Store arbitrary data in Flexible System (#53)

Co-authored-by: Jan Habscheid <Jan.Habscheid@rwth-aachen.de>
Co-authored-by: Michael F. Herbst <info@michael-herbst.com>

src/atom.jl

@@ -34,7 +34,7 @@ end
 
 function Atom(identifier::AtomId,
               position::AbstractVector{L},
-              velocity::AbstractVector{V}=zeros(3)u"bohr/s";
+              velocity::AbstractVector{V}=zeros(length(position))u"bohr/s";
               atomic_symbol=Symbol(elements[identifier].symbol),
               atomic_number=elements[identifier].number,
               atomic_mass::M=elements[identifier].atomic_mass,
@@ -66,27 +66,28 @@ end
 #
 # Special high-level functions to construct atomic systems
 #
-atomic_system(atoms::AbstractVector{<:Atom}, box, bcs) = FlexibleSystem(atoms, box, bcs)
-atomic_system(atoms::AbstractVector, box, bcs) = FlexibleSystem(convert.(Atom, atoms), box, bcs)
+atomic_system(atoms::AbstractVector{<:Atom}, box, bcs; kwargs...) = FlexibleSystem(atoms, box, bcs; kwargs...)
+atomic_system(atoms::AbstractVector, box, bcs; kwargs...) = FlexibleSystem(convert.(Atom, atoms), box, bcs; kwargs...)
 
-function isolated_system(atoms::AbstractVector{<:Atom})
+
+function isolated_system(atoms::AbstractVector{<:Atom}; kwargs...)
     # Use dummy box and boundary conditions
     D = n_dimensions(first(atoms))
-    atomic_system(atoms, infinite_box(D), fill(DirichletZero(), D))
+    atomic_system(atoms, infinite_box(D), fill(DirichletZero(), D); kwargs...)
 end
-isolated_system(atoms::AbstractVector) = isolated_system(convert.(Atom, atoms))
+isolated_system(atoms::AbstractVector; kwargs...) = isolated_system(convert.(Atom, atoms); kwargs...)
 
 function periodic_system(atoms::AbstractVector,
-                         box::AbstractVector{<:AbstractVector},
-                         boundary_conditions=fill(Periodic(), length(box));
+                         box::AbstractVector{<:AbstractVector};
                          fractional=false, kwargs...)
+    boundary_conditions=fill(Periodic(), length(box))
     lattice = hcat(box...)
-    !fractional && return atomic_system(atoms, box, boundary_conditions)
+    !fractional && return atomic_system(atoms, box, boundary_conditions; kwargs...)
 
     parse_fractional(atom::Atom) = atom
     function parse_fractional(atom::Pair)::Atom
         id, pos_fractional = atom
         Atom(id, lattice * pos_fractional)
     end
-    atomic_system(parse_fractional.(atoms), box, boundary_conditions)
+    atomic_system(parse_fractional.(atoms), box, boundary_conditions; kwargs...)
 end

src/flexible_system.jl

@@ -3,32 +3,42 @@
 #
 export FlexibleSystem
 
+
 struct FlexibleSystem{D, S, L<:Unitful.Length} <: AbstractSystem{D}
     particles::AbstractVector{S}
     box::SVector{D, SVector{D, L}}
     boundary_conditions::SVector{D, BoundaryCondition}
+    data::Dict{Symbol, Any}  # Store arbitrary data about the atom.
 end
 
+Base.hasproperty(system::FlexibleSystem, x::Symbol) = hasfield(FlexibleSystem, x) || haskey(system.data, x)
+Base.getproperty(system::FlexibleSystem, x::Symbol) = hasfield(FlexibleSystem, x) ? getfield(system, x) : getindex(system.data, x)
+
 function FlexibleSystem(
     particles::AbstractVector{S},
     box::AbstractVector{<:AbstractVector{L}},
-    boundary_conditions::AbstractVector{BC}
+    boundary_conditions::AbstractVector{BC};
+    kwargs...
 ) where {BC<:BoundaryCondition, L<:Unitful.Length, S}
     D = length(box)
     if !all(length.(box) .== D)
         throw(ArgumentError("Box must have D vectors of length D"))
     end
-    FlexibleSystem{D, S, L}(particles, box, boundary_conditions)
+    FlexibleSystem{D, S, L}(convert.(Atom, particles), box, boundary_conditions, Dict(kwargs...))
+
 end
 
 # Update constructor
 function FlexibleSystem(system::AbstractSystem;
                         particles=nothing, atoms=nothing,
                         box=bounding_box(system),
-                        boundary_conditions=boundary_conditions(system))
+                        boundary_conditions=boundary_conditions(system),
+                        kwargs...)
     particles = something(particles, atoms, collect(system))
-    FlexibleSystem(particles, box, boundary_conditions)
+    extra = system isa FlexibleSystem ? system.data : (; )
+    FlexibleSystem(particles, box, boundary_conditions; extra..., kwargs...)
 end
+FlexibleSystem(;system::FlexibleSystem, kwargs...) = FlexibleSystem(system; kwargs...)
 
 function Base.show(io::IO, system::FlexibleSystem)
     print(io, "FlexibleSystem")

test/atom.jl

@@ -5,6 +5,9 @@ using Test
 
 @testset "atomic systems" begin
     @testset "Atom construction" begin
+        at2D = Atom(:Si, zeros(2) * u"m", extradata=41)
+        @test n_dimensions(at2D) == 2
+
         at = Atom(:Si, zeros(3) * u"m", extradata=42)
 
         @test n_dimensions(at) == 3
@@ -36,15 +39,18 @@ using Test
     @testset "flexible atomic systems" begin
         box = [[10, 0.0, 0.0], [0.0, 5, 0.0], [0.0, 0.0, 7]]u"Å"
         bcs = [Periodic(), DirichletZero(), DirichletZero()]
+        dic = Dict{String, Any}("extradata_dic"=>"44")
         atoms = [:Si => [0.0, 1.0, 1.5]u"Å",
                  :C  => [0.0, 0.8, 1.7]u"Å",
                  Atom(:H, zeros(3) * u"Å", ones(3) * u"bohr/s")]
-        system = atomic_system(atoms, box, bcs)
+        system = atomic_system(atoms, box, bcs, extradata=45; dic)
         @test length(system) == 3
         @test atomic_symbol(system) == [:Si, :C, :H]
         @test boundary_conditions(system) == [Periodic(), DirichletZero(), DirichletZero()]
         @test position(system) == [[0.0, 1.0, 1.5], [0.0, 0.8, 1.7], [0.0, 0.0, 0.0]]u"Å"
         @test velocity(system) == [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [1.0, 1.0, 1.0]]u"bohr/s"
+        @test system.extradata == 45
+        @test system.dic["extradata_dic"] == "44"
 
         # Test update constructor
         newatoms  = [system[1], system[2]]
@@ -57,28 +63,34 @@ using Test
     end
 
     @testset "isolated_system" begin
+        dic = Dict{String, Any}("extradata_dic"=>"47")
         system = isolated_system([:Si => [0.0, 1.0, 1.5]u"Å",
                                   :C  => [0.0, 0.8, 1.7]u"Å",
-                                  Atom(:H, zeros(3) * u"Å")])
+                                  Atom(:H, zeros(3) * u"Å")], extradata=46; dic)
         @test length(system) == 3
         @test atomic_symbol(system) == [:Si, :C, :H]
         @test boundary_conditions(system) == [DirichletZero(), DirichletZero(), DirichletZero()]
         @test position(system) == [[0.0, 1.0, 1.5], [0.0, 0.8, 1.7], [0.0, 0.0, 0.0]]u"Å"
         @test velocity(system) == [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]u"bohr/s"
         @test bounding_box(system) == infinite_box(3)
+        @test system.extradata == 46
+        @test system.dic["extradata_dic"] == "47"
     end
 
     @testset "periodic_system" begin
         box = [[10, 0.0, 0.0], [0.0, 5, 0.0], [0.0, 0.0, 7]]u"Å"
+        dic = Dict{String, Any}("extradata_dic"=>"49")
         atoms = [:Si => [0.0, -0.125, 0.0],
                  :C  => [0.125, 0.0, 0.0],
                  Atom(:H, zeros(3) * u"Å")]
-        system = periodic_system(atoms, box; fractional=true)
+        system = periodic_system(atoms, box, extradata=48; fractional=true, dic)
 
         @test length(system) == 3
         @test atomic_symbol(system) == [:Si, :C, :H]
         @test boundary_conditions(system) == [Periodic(), Periodic(), Periodic()]
         @test position(system) == [[0.0, -0.625, 0.0], [1.25, 0.0, 0.0], [0.0, 0.0, 0.0]]u"Å"
+        @test system.extradata == 48
+        @test system.dic["extradata_dic"] == "49"
     end