Extend atomic and system property support

.bumpversion.cfg

@@ -1,5 +1,5 @@
 [bumpversion]
-current_version = 0.1.1
+current_version = 0.1.2
 commit = True
 tag = False
 

Project.toml

@@ -1,25 +1,29 @@
 name = "ASEconvert"
 uuid = "3da9722f-58c2-4165-81be-b4d7253e8fd2"
 authors = ["Michael F. Herbst <info@michael-herbst.com>"]
-version = "0.1.1"
+version = "0.1.2"
 
 [deps]
 AtomsBase = "a963bdd2-2df7-4f54-a1ee-49d51e6be12a"
 CondaPkg = "992eb4ea-22a4-4c89-a5bb-47a3300528ab"
+PeriodicTable = "7b2266bf-644c-5ea3-82d8-af4bbd25a884"
 PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d"
 Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
+UnitfulAtomic = "a7773ee8-282e-5fa2-be4e-bd808c38a91a"
 
 [compat]
-julia = "1.6"
 AtomsBase = "0.2"
 CondaPkg = "0.2"
+PeriodicTable = "1"
 PythonCall = "0.9"
 Unitful = "1"
+UnitfulAtomic = "1"
+julia = "1.6"
 
 [extras]
-Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 ExtXYZ = "352459e4-ddd7-4360-8937-99dcb397b478"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
 test = ["Test", "ExtXYZ", "LinearAlgebra"]

README.md

@@ -19,7 +19,7 @@ using ASEconvert
 using DFTK
 
 # Make a silicon supercell using ASE
-atoms_ase = ase.build.bulk("Si") * (4, 1, 1)
+atoms_ase = ase.build.bulk("Si") * pytuple((4, 1, 1))
 
 # Convert to an AtomsBase-compatible structure
 atoms_ab = pyconvert(AbstractSystem, atoms_ase)

src/ASEconvert.jl

@@ -2,11 +2,13 @@ module ASEconvert
 using PythonCall
 using AtomsBase
 using Unitful
+using UnitfulAtomic
+import PeriodicTable
 
 export convert_ase
 export ase
-export pyconvert       # Reexport from PythonCall
-export AbstractSystem  # Reexport from AtomsBase
+export pyconvert, pytuple  # Reexport from PythonCall
+export AbstractSystem      # Reexport from AtomsBase
 
 """
 Global constant representing the `ase` python module available from Julia.
@@ -26,19 +28,32 @@ function ase_to_system(S::Type{<:AbstractSystem}, ase_atoms::Py)
     box = [pyconvert(Vector, ase_atoms.cell[i])u"Å" for i = 0:2]
 
     atnums     = pyconvert(Vector, ase_atoms.get_atomic_numbers())
+    atsyms     = pyconvert(Vector, ase_atoms.symbols)
+    atmasses   = pyconvert(Vector, ase_atoms.get_masses())
     positions  = pyconvert(Matrix, ase_atoms.get_positions())
     velocities = pyconvert(Matrix, ase_atoms.get_velocities())
     magmoms    = pyconvert(Vector, ase_atoms.get_initial_magnetic_moments())
-
-    atoms = [AtomsBase.Atom(atnums[i], positions[i, :]u"Å", velocities[i, :]u"Å/s";
-                            magnetic_moment=magmoms[i])
-             for i = 1:length(atnums)]
+    charges    = pyconvert(Vector, ase_atoms.get_initial_charges())
+    ase_info   = pyconvert(Dict{String,Any}, ase_atoms.info)
+
+    atoms = map(1:length(atnums)) do i
+        AtomsBase.Atom(atnums[i], positions[i, :]u"Å", velocities[i, :]u"Å/s";
+                       atomic_symbol=Symbol(atsyms[i]),
+                       atomic_number=atnums[i],
+                       atomic_mass=atmasses[i]u"u",
+                       magnetic_moment=magmoms[i],
+                       charge=charges[i]u"e_au")
+    end
 
     # Parse extra data in info struct
-    info = pyconvert(Dict{Symbol,Any}, ase_atoms.info)
-
-    # TODO We should have some mechanism to extract "official" ASE system keys not supported
-    #      in AtomsBase and put them into the system-level data.
+    info = Dict{Symbol, Any}()
+    for (k, v) in ase_info
+        if k == "charge"
+            info[Symbol(k)] = v * u"e_au"
+        else
+            info[Symbol(k)] = v
+        end
+    end
 
     bcs = [p ? Periodic() : DirichletZero() for p in pyconvert(Vector, ase_atoms.pbc)]
     PythonCall.pyconvert_return(atomic_system(atoms, box, bcs; info...))
@@ -49,14 +64,24 @@ end
 
 Convert a passed `system` (which satisfies the AtomsBase.AbstractSystem interface) to an
 `ase.Atoms` datastructure. Conversions to other ASE objects from equivalent Julia objects
-may be added in the future.
+may be added as additional methods in the future.
 """
 function convert_ase(system::AbstractSystem{D}) where {D}
     D != 3 && @warn "1D and 2D systems not yet fully supported."
 
     n_atoms = length(system)
     pbc     = map(isequal(Periodic()), boundary_conditions(system))
-    symbols = map(String, atomic_symbol(system))
+    numbers = atomic_number(system)
+    masses  = ustrip.(u"u", atomic_mass(system))
+
+    symbols_match = [
+        PeriodicTable.elements[atnum].symbol == string(atomic_symbol(system, i))
+        for (i, atnum) in enumerate(numbers)
+    ]
+    if !all(symbols_match)
+        @warn("Mismatch between atomic numbers and atomic symbols, which is not " *
+              "supported in ASE. Atomic numbers take preference.")
+    end
 
     cell = zeros(3, 3)
     for (i, v) in enumerate(bounding_box(system))
@@ -76,39 +101,29 @@ function convert_ase(system::AbstractSystem{D}) where {D}
         end
     end
 
-    # Map extra atomic properties
-    # Note: This is needed since ASE only support *some* atomic properties.
-    # In theory it is possible to support others using the system-level info dictionary,
-    # but this gets really messy and probably unreliable, so we don't do that.
-    #
-    # mapkeys maps from the atomic key in the AtomsBase convention to the ase.Atoms
-    # key and the default value (including units) in case an atom does not set this property.
-    mapkeys = Dict(:magnetic_moment => (:magmoms, 0.0), )
-    extra = Dict{Symbol,Any}()
-    for (key_atbase, (key_ase, atom_default)) in pairs(mapkeys)
-        if any(at -> hasproperty(at, key_atbase), system)
-            data_unit = unit(atom_default)
-            atpropdata = map(system) do atom
-                if hasproperty(atom, key_atbase)
-                    ustrip(data_unit, getproperty(atom, key_atbase))
-                else
-                    atom_default
-                end
-            end
-            extra[key_ase] = atpropdata
-        end
+    charges = nothing
+    if any(hasproperty(atom, :charge) for atom in system)
+        charges = [ustrip(u"e_au", atom.charge) for atom in system]
     end
 
+    magmoms = nothing
+    if any(hasproperty(atom, :magnetic_moment) for atom in system)
+        magmoms = [atom.magnetic_moment for atom in system]
+    end
+
+
     # Map extra system properties
     # TODO Probably we need some mechanism to map keys which happen to be "official"
-    #      ASE keys to their ASE counterparts.
+    #      AtomsBase keys to their ASE counterparts.
     info = Dict{String, Any}()
     if system isa FlexibleSystem
         # Here we can retrieve extra data
         # TODO not a good idea to directly access the field
         # TODO Implement and make use of a property interface on the system level
         for (k, v) in system.data
-            if v isa Quantity || (v isa AbstractArray && eltype(v) <: Quantity)
+            if k == :charge
+                info[string(k)] = ustrip(u"e_au", v)
+            elseif v isa Quantity || (v isa AbstractArray && eltype(v) <: Quantity)
                 @warn("Unitful quantities are not yet supported in convert_ase. " *
                       "Ignoring key $k")
             else
@@ -117,7 +132,22 @@ function convert_ase(system::AbstractSystem{D}) where {D}
         end
     end
 
-    ase.Atoms(; symbols, positions, cell, pbc, velocities, info, extra...)
+    # We don't map any extra properties, which are not available in ASE as this
+    # only causes a mess: ASE could do something to the atoms, but not taking
+    # care of the extra properties, thus rendering the extra properties invalid
+    # without the user noticing.
+    if first(system) isa Atom
+        # TODO not a good idea to directly access the field
+        # TODO Implement and make use of a property interface on the atom level
+        for k in keys(system[1].data)
+            if !(k in (:charge, :magnetic_moment))
+                @warn "Skipping atomic property $k, which is not supported in ASE."
+            end
+        end
+    end
+
+    ase.Atoms(; positions, numbers, masses, magmoms, charges,
+              cell, pbc, velocities, info)
 end
 
 # TODO Could have a convert_ase(Vector{AbstractSystem}) to make an ASE trajectory

test/common.jl

@@ -0,0 +1,114 @@
+using AtomsBase
+using Test
+using LinearAlgebra
+using Unitful
+using UnitfulAtomic
+
+"""
+Test whether two abstract systems are approximately the same. Certain atomic or system
+properties can be ignored during the comparison using the respective kwargs.
+"""
+function test_approx_eq(s::AbstractSystem, t::AbstractSystem;
+                        atol=1e-14, ignore_atprop=Symbol[], ignore_sysprop=Symbol[])
+    # TODO Introduce an == / ≈ method in AbstractSystem for this purpose
+
+    @test maximum(norm, position(s)     - position(t))     < atol * u"Å"
+    @test maximum(norm, bounding_box(s) - bounding_box(t)) < atol * u"Å"
+
+    if !(:velocity in ignore_atprop)
+        @test ismissing(velocity(s)) == ismissing(velocity(t))
+        if !ismissing(velocity(s)) && !ismissing(velocity(t))
+            @test maximum(norm, velocity(s) - velocity(t)) < atol * u"Å/s"
+        end
+    end
+
+    for method in (atomic_symbol, atomic_number, boundary_conditions)
+        @test method(s) == method(t)
+    end
+    @test maximum(abs, atomic_mass(s) - atomic_mass(t)) < atol * u"u"
+
+    extra_atomic_props = (:charge, :covalent_radius, :vdw_radius, :magnetic_moment)
+    for prop in extra_atomic_props
+        prop in ignore_atprop && continue
+        for (at_s, at_t) in zip(s, t)
+            @test hasproperty(at_s, prop) == hasproperty(at_t, prop)
+            if hasproperty(at_s, prop) && hasproperty(at_t, prop)
+                @test getproperty(at_s, prop) == getproperty(at_t, prop)
+            end
+        end
+    end
+
+    if s isa FlexibleSystem && t isa FlexibleSystem
+        for k in keys(s.data)
+            k in ignore_sysprop && continue
+            @test s.data[k] == t.data[k]
+        end
+    end
+end
+
+"""
+Setup a standard test system using some random data and supply the data to the caller.
+Extra atomic or system properties can be specified using `extra_atprop` and `extra_sysprop`
+and specific standard keys can be ignored using `drop_atprop` and `drop_sysprop`.
+"""
+function make_test_system(D=3; drop_atprop=Symbol[], drop_sysprop=Symbol[],
+                          extra_atprop=(; ), extra_sysprop=(; ), cellmatrix=:full)
+    # TODO Should be moved to AtomsBase
+    @assert D == 3
+    n_atoms = 5
+
+    # Generate some random data to store in Atoms
+    atprop = Dict{Symbol,Any}(
+        :position        => [randn(3) for _ = 1:n_atoms]u"Å",
+        :velocity        => [randn(3) for _ = 1:n_atoms]u"Å/s",
+        :atomic_symbol   => [:H, :H, :C, :N, :He],
+        :atomic_number   => [1, 1, 6, 7, 2],
+        :charge          => [2, 1, 3.0, -1.0, 0.0]u"e_au",
+        :atomic_mass     => 10rand(n_atoms)u"u",
+        :vdw_radius      => randn(n_atoms)u"Å",
+        :covalent_radius => randn(n_atoms)u"Å",
+        :magnetic_moment => [0.0, 0.0, 1.0, -1.0, 0.0],
+    )
+    sysprop = Dict{Symbol,Any}(
+        :extra_data   => 42,
+        :charge       => -1u"e_au",
+        :multiplicity => 2,
+    )
+
+    for prop in drop_atprop
+        pop!(atprop, prop)
+    end
+    for prop in drop_sysprop
+        pop!(sysprop, prop)
+    end
+    sysprop = merge(sysprop, pairs(extra_sysprop))
+    atprop  = merge(atprop,  pairs(extra_atprop))
+
+    atoms = map(1:n_atoms) do i
+        atargs = Dict(k => v[i] for (k, v) in pairs(atprop)
+                      if !(k in (:position, :velocity)))
+        if haskey(atprop, :velocity)
+            Atom(atprop[:atomic_symbol][i], atprop[:position][i], atprop[:velocity][i];
+                 atargs...)
+        else
+            Atom(atprop[:atomic_symbol][i], atprop[:position][i]; atargs...)
+        end
+    end
+    if cellmatrix == :upper_triangular
+        box = [[1.54732, -0.807289, -0.500870],
+               [    0.0, 0.4654985, 0.5615117],
+               [    0.0,       0.0, 0.7928950]]u"Å"
+    elseif cellmatrix == :lower_triangular
+        box = [[1.54732, 0.0, 0.0],
+               [-0.807289, 0.4654985, 0.0],
+               [-0.500870, 0.5615117, 0.7928950]]u"Å"
+    else
+        box = [[-1.50304, 0.850344, 0.717239],
+               [ 0.36113, 0.008144, 0.814712],
+               [ 0.06828, 0.381122, 0.129081]]u"Å"
+    end
+    bcs = [Periodic(), Periodic(), DirichletZero()]
+    system = atomic_system(atoms, box, bcs; sysprop...)
+
+    (; system, atoms, atprop=NamedTuple(atprop), sysprop=NamedTuple(sysprop), box, bcs)
+end