Merge pull request #12 from JuliaMolSim/soa

SoA example

sandbox/aos_vs_soa.jl

@@ -0,0 +1,17 @@
+using AtomsBase
+using StaticArrays
+using Unitful
+
+box = [[1, 0, 0], [0, 1, 0], [0, 0, 1]]u"m"
+bcs = [Periodic(), Periodic(), Periodic()]
+positions = [0.25 0.25 0.25; 0.75 0.75 0.75]u"m"
+elements = [ChemicalElement(:C), ChemicalElement(:C)]
+
+atom1 = SimpleAtom(SVector{3}(positions[1,:]),elements[1])
+atom2 = SimpleAtom(SVector{3}(positions[2,:]),elements[2])
+
+aos = FlexibleSystem(box, bcs, [atom1, atom2])
+soa = FastSystem(box, bcs, positions, elements)
+
+# And now we can ask questions like...
+soa .== aos

src/AtomsBase.jl

@@ -1,7 +1,8 @@
 module AtomsBase
 
 include("interface.jl")
-include("implementation_simple.jl")
+include("implementation_soa.jl")
+include("implementation_aos.jl")
 
 # Write your package code here.
 

src/implementation_aos.jl

@@ -0,0 +1,35 @@
+# Example implementation using as few function definitions as possible
+#
+using StaticArrays
+
+export FlexibleSystem
+
+# TODO Switch order of type arguments?
+struct FlexibleSystem{D,ET<:AbstractElement,AT<:AbstractParticle{ET}} <: AbstractSystem{D,ET,AT}
+    box::SVector{D,<:SVector{D,<:Unitful.Length}}
+    boundary_conditions::SVector{D,<:BoundaryCondition}
+    particles::Vector{AT}
+end
+# convenience constructor where we don't have to preconstruct all the static stuff...
+function FlexibleSystem(
+    box::Vector{Vector{L}},
+    boundary_conditions::Vector{BC},
+    particles::Vector{AT},
+) where {BC<:BoundaryCondition,L<:Unitful.Length,AT<:AbstractParticle}
+    D = length(box)
+    if !all(length.(box) .== D)
+        throw(ArgumentError("box must have D vectors of length D"))
+    end
+    FlexibleSystem(SVector{D,SVector{D,L}}(box), SVector{D,BoundaryCondition}(boundary_conditions), particles)
+end
+
+bounding_box(sys::FlexibleSystem) = sys.box
+boundary_conditions(sys::FlexibleSystem) = sys.boundary_conditions
+
+function Base.show(io::IO, ::MIME"text/plain", sys::FlexibleSystem)
+    print(io, "FlexibleSystem with ", length(sys), " particles")
+end
+
+Base.size(sys::FlexibleSystem) = size(sys.particles)
+Base.length(sys::FlexibleSystem) = length(sys.particles)
+Base.getindex(sys::FlexibleSystem, i::Int) = getindex(sys.particles, i)

src/implementation_soa.jl

@@ -0,0 +1,61 @@
+# Example implementation using as few function definitions as possible
+#
+using StaticArrays
+
+export FastSystem
+
+struct FastSystem{D,ET<:AbstractElement,AT<:AbstractParticle{ET},L<:Unitful.Length} <:
+       AbstractSystem{D,ET,AT}
+    box::SVector{D,SVector{D,L}}
+    boundary_conditions::SVector{D,BoundaryCondition}
+    positions::Vector{SVector{D,L}}
+    elements::Vector{ET}
+    # janky inner constructor that we need for some reason
+    FastSystem(box, boundary_conditions, positions, elements) =
+        new{length(boundary_conditions),eltype(elements),SimpleAtom,eltype(eltype(positions))}(
+            box,
+            boundary_conditions,
+            positions,
+            elements,
+        )
+end
+
+# convenience constructor where we don't have to preconstruct all the static stuff...
+function FastSystem(
+    box::AbstractVector{Vector{L}},
+    boundary_conditions::AbstractVector{BC},
+    positions::AbstractMatrix{M},
+    elements::AbstractVector{ET},
+) where {L<:Unitful.Length,BC<:BoundaryCondition,M<:Unitful.Length,ET<:AbstractElement}
+    N = length(elements)
+    D = length(box)
+    if !all(length.(box) .== D)
+        throw(ArgumentError("box must have D vectors of length D"))
+    end
+    if !(size(positions, 1) == N)
+        throw(ArgumentError("box must have D vectors of length D"))
+    end
+    if !(size(positions, 2) == D)
+        throw(ArgumentError("box must have D vectors of length D"))
+    end
+    FastSystem(
+        SVector{D,SVector{D,L}}(box),
+        SVector{D,BoundaryCondition}(boundary_conditions),
+        Vector{SVector{D,eltype(positions)}}([positions[i, :] for i = 1:N]),
+        elements,
+    )
+end
+
+function Base.show(io::IO, ::MIME"text/plain", sys::FastSystem)
+    print(io, "FastSystem with ", length(sys), " particles")
+end
+
+bounding_box(sys::FastSystem) = sys.box
+boundary_conditions(sys::FastSystem) = sys.boundary_conditions
+
+# Base.size(sys::FastSystem) = size(sys.particles)
+Base.length(sys::FastSystem{D,ET,AT}) where {D,ET,AT} = length(sys.elements)
+
+# first piece of trickiness: can't do a totally abstract dispatch here because we need to know the signature of the constructor for AT
+Base.getindex(sys::FastSystem{D,ET,SimpleAtom}, i::Int) where {D,ET} =
+    SimpleAtom{D}(sys.positions[i], sys.elements[i])

src/interface.jl

@@ -2,13 +2,20 @@ using Unitful
 using UnitfulAtomic
 using PeriodicTable
 using StaticArrays
+import Base.position
 
 export AbstractElement, AbstractParticle, AbstractAtom, AbstractSystem, AbstractAtomicSystem
-export ChemicalElement
+export ChemicalElement, SimpleAtom
 export BoundaryCondition, DirichletZero, Periodic
-export atomic_mass, atomic_number, atomic_symbol,
-    bounding_box, element, position, velocity,
-    boundary_conditions, periodic_dims
+export atomic_mass,
+    atomic_number,
+    atomic_symbol,
+    bounding_box,
+    element,
+    position,
+    velocity,
+    boundary_conditions,
+    periodic_dims
 export atomic_property, has_atomic_property, atomic_propertynames
 export n_dimensions
 
@@ -18,13 +25,14 @@ struct ChemicalElement <: AbstractElement
     data::PeriodicTable.Element
 end
 
-ChemicalElement(symbol::Union{Symbol,Integer,AbstractString}) = ChemicalElement(PeriodicTable.elements[symbol])
+ChemicalElement(symbol::Union{Symbol,Integer,AbstractString}) =
+    ChemicalElement(PeriodicTable.elements[symbol])
 Base.show(io::IO, elem::ChemicalElement) = print(io, "Element(", atomic_symbol(elem), ")")
 
 # These are always only read-only ... and allow look-up into a database
 atomic_symbol(el::ChemicalElement) = el.data.symbol
 atomic_number(el::ChemicalElement) = el.data.number
-atomic_mass(el::ChemicalElement)   = el.data.atomic_mass
+atomic_mass(el::ChemicalElement) = el.data.atomic_mass
 
 
 
@@ -36,8 +44,8 @@ atomic_mass(el::ChemicalElement)   = el.data.atomic_mass
 # IdType:  Type used to identify the particle
 #
 abstract type AbstractParticle{ET<:AbstractElement} end
-velocity(::AbstractParticle)::AbstractVector{<: Unitful.Velocity} = missing
-position(::AbstractParticle)::AbstractVector{<: Unitful.Length}   = error("Implement me")
+velocity(::AbstractParticle)::AbstractVector{<:Unitful.Velocity} = missing
+position(::AbstractParticle)::AbstractVector{<:Unitful.Length} = error("Implement me")
 (element(::AbstractParticle{ET})::ET) where {ET<:AbstractElement} = error("Implement me")
 
 
@@ -57,64 +65,85 @@ element(::AbstractAtom)::ChemicalElement = error("Implement me")
 # implementations, therefore these interfaces are forwarded from the Element object.
 atomic_symbol(atom::AbstractAtom) = atomic_symbol(element(atom))
 atomic_number(atom::AbstractAtom) = atomic_number(element(atom))
-atomic_mass(atom::AbstractAtom)   = atomic_mass(element(atom))
+atomic_mass(atom::AbstractAtom) = atomic_mass(element(atom))
 
 # Custom atomic properties:
-atomic_property(::AbstractAtom, ::Symbol, default=missing) = default
-has_atomic_property(atom::AbstractAtom, property::Symbol) = !ismissing(atomic_property(atom, property))
+atomic_property(::AbstractAtom, ::Symbol, default = missing) = default
+has_atomic_property(atom::AbstractAtom, property::Symbol) =
+    !ismissing(atomic_property(atom, property))
 atomic_propertynames(::AbstractAtom) = Symbol[]
 
 #
 # Identifier for boundary conditions per dimension
 #
 abstract type BoundaryCondition end
 struct DirichletZero <: BoundaryCondition end  # Dirichlet zero boundary (i.e. molecular context)
-struct Periodic      <: BoundaryCondition end  # Periodic BCs
+struct Periodic <: BoundaryCondition end  # Periodic BCs
 
 
 #
 # The system type
 #     Again readonly.
 #
 
-abstract type AbstractSystem{D, ET<:AbstractElement, AT<:AbstractParticle{ET}} end
-(bounding_box(::AbstractSystem{D})::SVector{D, SVector{D, <:Unitful.Length}}) where {D} = error("Implement me")
-(boundary_conditions(::AbstractSystem{D})::SVector{D,BoundaryCondition}) where {D} = error("Implement me")
+abstract type AbstractSystem{D,ET<:AbstractElement,AT<:AbstractParticle{ET}} end
+(bounding_box(::AbstractSystem{D})::SVector{D,SVector{D,<:Unitful.Length}}) where {D} =
+    error("Implement me")
+(boundary_conditions(::AbstractSystem{D})::SVector{D,BoundaryCondition}) where {D} =
+    error("Implement me")
 
-get_periodic(sys::AbstractSystem) = [isa(bc, Periodic) for bc in get_boundary_conditions(sys)]
+get_periodic(sys::AbstractSystem) =
+    [isa(bc, Periodic) for bc in get_boundary_conditions(sys)]
 
 # Note: Can't use ndims, because that is ndims(sys) == 1 (because of AbstractVector interface)
 n_dimensions(::AbstractSystem{D}) where {D} = D
 
 
 # indexing interface
-Base.getindex(::AbstractSystem, ::Int)  = error("Implement me")
-Base.size(::AbstractSystem)             = error("Implement me")
-Base.length(::AbstractSystem)           = error("Implement me")
+Base.getindex(::AbstractSystem, ::Int) = error("Implement me")
+Base.size(::AbstractSystem) = error("Implement me")
+Base.length(::AbstractSystem) = error("Implement me")
 Base.setindex!(::AbstractSystem, ::Int) = error("AbstractSystem objects are not mutable.")
 Base.firstindex(::AbstractSystem) = 1
 Base.lastindex(s::AbstractSystem) = length(s)
 Base.iterate(S::AbstractSystem, i::Int=1) = (1 <= i <= length(S)) ? (@inbounds S[i], i+1) : nothing
 
+# iteration interface, needed for default broadcast dispatches below to work
+Base.iterate(sys::AbstractSystem{D,ET,AT}, state = firstindex(sys)) where {D,ET,AT} =
+    state > length(sys) ? nothing : (sys[state], state + 1)
+
 # TODO Support similar, push, ...
 
 # Some implementations might prefer to store data in the System as a flat list and
 # expose Atoms as a view. Therefore these functions are needed. Of course this code
 # should be autogenerated later on ...
 position(sys::AbstractSystem) = position.(sys)    # in Cartesian coordinates!
 velocity(sys::AbstractSystem) = velocity.(sys)    # in Cartesian coordinates!
-element(sys::AbstractSystem)  = element.(sys)
+element(sys::AbstractSystem) = element.(sys)
 
 #
 # Extra stuff only for Systems composed of atoms
 #
 const AbstractAtomicSystem{D,AT<:AbstractAtom} = AbstractSystem{D,ChemicalElement,AT}
 atomic_symbol(sys::AbstractAtomicSystem) = atomic_symbol.(sys)
 atomic_number(sys::AbstractAtomicSystem) = atomic_number.(sys)
-atomic_mass(sys::AbstractAtomicSystem)   = atomic_mass.(sys)
-atomic_property(sys::AbstractAtomicSystem, property::Symbol)::Vector{Any} = atomic_property.(sys, property)
+atomic_mass(sys::AbstractAtomicSystem) = atomic_mass.(sys)
+atomic_property(sys::AbstractAtomicSystem, property::Symbol)::Vector{Any} =
+    atomic_property.(sys, property)
 atomic_propertiesnames(sys::AbstractAtomicSystem) = unique(sort(atomic_propertynames.(sys)))
 
+struct SimpleAtom{D} <: AbstractAtom
+    position::SVector{D,<:Unitful.Length}
+    element::ChemicalElement
+end
+SimpleAtom(position, element) = SimpleAtom{length(position)}(position, element)
+position(atom::SimpleAtom) = atom.position
+element(atom::SimpleAtom) = atom.element
+
+function SimpleAtom(position, symbol::Union{Integer,AbstractString,Symbol,AbstractVector})
+    SimpleAtom(position, ChemicalElement(symbol))
+end
+
 # Just to make testing a little easier for now
 function Base.show(io::IO, ::MIME"text/plain", part::AbstractParticle)
     print(io, "Particle(", element(part), ") @ ", position(part))