diff --git a/src/atomsbase.jl b/src/atomsbase.jl index aea44b2..dc77922 100644 --- a/src/atomsbase.jl +++ b/src/atomsbase.jl @@ -14,31 +14,35 @@ export get_cell, AosSystem, SoaSystem, ChemicalElement # an `Atom` is now just a `PState`, so we define # accessors for the PState fields with canonical names. -symbol(::typeof(position)) = :𝐫 -symbol(::typeof(velocity)) = :𝐯 -symbol(::typeof(atomic_mass)) = :𝑚 -symbol(::typeof(atomic_symbol)) = :𝑍 - -const _atom_syms = (𝐫 = position, - 𝐯 = velocity, - 𝑚 = atomic_mass, - 𝑍 = atomic_symbol) - -position(atom::PState) = atom.𝐫 -velocity(atom::PState) = atom.𝐯 -atomic_mass(atom::PState) = atom.𝑚 -atomic_symbol(atom::PState) = atom.𝑍 # this one I'm not sure about +# list of properties that DP knows about: +_list_of_properties = [ + (:𝐫, :position), + (:𝐯, :velocity), + (:𝑚, :atomic_mass), + (:𝑍, :atomic_symbol), + (:𝑞, :charge), + (:μ, :dipole), + (:𝐩, :momentum), + (:𝐸, :energy), + (:𝑀, :mass), + ] + +for (sym, name) in _list_of_properties + @eval $name(x::XState) = x.$sym + @eval symbol(::typeof($name)) = $(Meta.quot(sym)) +end atomic_number(atom::PState) = atomic_number(atomic_symbol(atom)) -_post(a) = a -_post(sym::Symbol) = ChemicalElement(sym) +_post(p, a) = a +_post(::typeof(atomic_symbol), sym::Symbol) = ChemicalElement(sym) +_post(::typeof(atomic_symbol), z::Integer) = ChemicalElement(z) """ Generate an atom with the given properties. """ atom(at; properties = (position, atomic_mass, atomic_symbol)) = - PState((; [symbol(p) => _post(p(at)) for p in properties]...)) + PState((; [symbol(p) => _post(p, p(at)) for p in properties]...)) # --------------------------------------------------- # Array of Structs System @@ -68,10 +72,10 @@ Base.length(at::AosSystem) = length(at.particles) Base.getindex(at::AosSystem, i::Int) = at.particles[i] Base.getindex(at::AosSystem, inds::AbstractVector) = at.particles[inds] -for f in (:position, :velocity, :atomic_mass, :atomic_symbol) - @eval $f(sys::AosSystem) = [ $f(x) for x in sys.particles ] - @eval $f(sys::AosSystem, i::Integer) = $f(sys.particles[i]) - @eval $f(sys::AosSystem, inds::AbstractVector) = [$f(sys.particles[i]) for i in inds] +for (sym, name) in [ _list_of_properties; [(:ignore, :atomic_number) ] ] + @eval $name(sys::AosSystem) = [ $name(x) for x in sys.particles ] + @eval $name(sys::AosSystem, i::Integer) = $name(sys.particles[i]) + @eval $name(sys::AosSystem, inds::AbstractVector) = [$name(sys.particles[i]) for i in inds] end # AtomsBase. @@ -95,7 +99,7 @@ end function SoaSystem(sys::AbstractSystem; properties = (position, atomic_mass, atomic_symbol), ) - arrays = (; [symbol(p) => _post.(p(sys)) for p in properties]... ) + arrays = (; [symbol(p) => _post.(p, p(sys)) for p in properties]... ) cell = get_cell(sys) D = AtomsBase.n_dimensions(cell) return SoaSystem{D, typeof(cell), typeof(arrays)}( @@ -130,13 +134,14 @@ end Base.getindex(sys::SoaSystem, inds::AbstractVector{<: Integer}) = [ sys[i] for i in inds ] -for f in (:position, :velocity, :atomic_mass, :atomic_symbol) - @eval $f(sys::SoaSystem) = getfield(sys.arrays, symbol($f)) - @eval $f(sys::SoaSystem, i::Integer) = getfield(sys.arrays, symbol($f))[i] - @eval $f(sys::SoaSystem, inds::AbstractVector) = getfield(sys.arrays, symbol($f))[inds] + +for (sym, name) in [ _list_of_properties; [(:ignore, :atomic_number) ] ] + @eval $name(sys::SoaSystem) = copy(sys.arrays.$sym) + @eval $name(sys::SoaSystem, i::Integer) = sys.arrays.$sym[i] + @eval $name(sys::SoaSystem, inds::AbstractVector) = sys.arrays.$sym[inds] end -# AtomsBase. + get_cell(at::SoaSystem) = at.cell for f in (:n_dimensions, :bounding_box, :boundary_conditions, :periodicity) @@ -149,37 +154,49 @@ end # --------------------------------------------------------------- # Extension of the AtomsBase interface with setter functions -export set_position, - set_position!, - set_positions!, - set_bounding_box! - -set_position(x::PState, 𝐫::SVector) = setproperty(x, :𝐫, 𝐫) - -function set_position!(sys::AosSystem, i::Integer, 𝐫::SVector) - xi = sys.particles[i] - sys.particles[i] = set_position(xi, 𝐫) - return nothing -end - -function set_positions!(sys::AosSystem, R::AbstractVector{<: SVector}) - for i = 1:length(sys) - set_position!(sys, i, R[i]) +for (sym, name) in _list_of_properties + set_name = Symbol("set_$name") + set_name_ip = Symbol("set_$(name)!") + set_names_ip = Symbol("set_$(name)s!") + @eval export $set_name, $set_name_ip, $set_names_ip + # e.g. set_position(x::PState, 𝐫) = setproperty(x, :𝐫, 𝐫) + @eval $set_name(x::PState, t) = setproperty(x, $(Meta.quot(sym)), t) + # e.g. set_position!(sys, i, 𝐫) + @eval begin + function $set_name_ip(sys::AosSystem, i::Integer, t) + xi = sys.particles[i] + sys.particles[i] = $set_name(xi, t) + return nothing + end + end + @eval begin + function $set_name_ip(sys::SoaSystem, i::Integer, t) + sys.arrays.$sym[i] = t + return nothing + end + end + # e.g. set_positions!(sys, X) + @eval begin + function $set_names_ip(sys::AosSystem, Rs) + for i = 1:length(sys) + $set_name_ip(sys, i, Rs[i]) + end + end + return nothing + end + @eval begin + function $set_names_ip(sys::SoaSystem, R::AbstractVector) + copy!(sys.arrays.$sym, R) + return nothing + end end - return nothing end +# todo: set_atomic_number & co -function set_position!(sys::SoaSystem, i::Integer, 𝐫::SVector) - sys.arrays.𝐫[i] = 𝐫 - return nothing -end - -function set_positions!(sys::SoaSystem, R::AbstractVector{<: SVector}) - copy!(sys.arrays.𝐫, R) - return nothing -end +# cell-level setter +export set_bounding_box! function set_bounding_box!(sys::Union{SoaSystem, AosSystem}, bb) cell = PCell(bb, periodicity(sys)) diff --git a/src/states.jl b/src/states.jl index 37d1e10..52e46fe 100644 --- a/src/states.jl +++ b/src/states.jl @@ -9,7 +9,7 @@ import Base: *, +, -, zero, rand, randn, show, promote_rule, rtoldefault, import LinearAlgebra: norm, promote_leaf_eltypes -export PState, VState, vstate_type +export PState, VState, vstate_type, setproperty abstract type XState{NT <: NamedTuple} end diff --git a/test/test_atomsbase.jl b/test/test_atomsbase.jl index d17df6e..419a4d6 100644 --- a/test/test_atomsbase.jl +++ b/test/test_atomsbase.jl @@ -18,6 +18,7 @@ display(x) @test atomic_symbol(x) == x.𝑍 @test atomic_number(x) == 6 + ## #convert an Atom @@ -27,7 +28,9 @@ display(x) @test x.𝐫 == position(x) == position(at) @test x.𝑚 == atomic_mass(x) == atomic_mass(at) @test x.𝑍 == atomic_symbol(x) == atomic_symbol(at) - +@test DP.symbol(position) == :𝐫 +@test DP.symbol(atomic_mass) == :𝑚 +@test DP.symbol(atomic_symbol) == :𝑍 ## # convert an entire system