From a34d094b238514871f081973f2009c2ccc71d7c7 Mon Sep 17 00:00:00 2001
From: ACEsuit <christohortner@gmail.com>
Date: Tue, 21 May 2024 23:00:58 -0700
Subject: [PATCH] sketching out an atomsbase interface

---
 Project.toml              |  7 ++--
 src/DecoratedParticles.jl |  1 +
 src/structures.jl         | 67 +++++++++++++++++++++++++++++++++++++++
 test/test_atomsbase.jl    | 41 ++++++++++++++++++++++++
 4 files changed, 114 insertions(+), 2 deletions(-)
 create mode 100644 test/test_atomsbase.jl

diff --git a/Project.toml b/Project.toml
index 966fd03..f0d4716 100644
--- a/Project.toml
+++ b/Project.toml
@@ -4,19 +4,22 @@ authors = ["Christoph Ortner <christophortner@gmail.com> and contributors"]
 version = "0.0.2"
 
 [deps]
+AtomsBase = "a963bdd2-2df7-4f54-a1ee-49d51e6be12a"
+AtomsBuilder = "f5cc8831-eeb7-4288-8d9f-d6c1ddb77004"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 NamedTupleTools = "d9ec5142-1e00-5aa0-9d6a-321866360f50"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 
 [compat]
-StaticArrays = "1"
 ChainRulesCore = "1"
 NamedTupleTools = "0.14"
+StaticArrays = "1"
 julia = "1.9"
 
 [extras]
+AtomsBuilder = "f5cc8831-eeb7-4288-8d9f-d6c1ddb77004"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test"]
+test = ["Test", "AtomsBuilder"]
diff --git a/src/DecoratedParticles.jl b/src/DecoratedParticles.jl
index 886c2e0..36f7e40 100644
--- a/src/DecoratedParticles.jl
+++ b/src/DecoratedParticles.jl
@@ -5,5 +5,6 @@ include("states.jl")
 
 include("show.jl")
 
+include("structures.jl")
 
 end
diff --git a/src/structures.jl b/src/structures.jl
index 8b13789..6e320d9 100644
--- a/src/structures.jl
+++ b/src/structures.jl
@@ -1 +1,68 @@
 
+import AtomsBase 
+import AtomsBase: AbstractSystem, ChemicalElement, 
+                  position, velocity, atomic_mass, atomic_number, 
+                  atomic_symbol
+
+# --------------------------------------------------- 
+# 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)) = :m
+symbol(::typeof(atomic_symbol)) = :Z
+
+position(atom::PState) = atom.𝐫 
+velocity(atom::PState) = atom.𝐯
+atomic_mass(atom::PState) = atom.m
+atomic_symbol(atom::PState) = atom.Z   # this one I'm not sure about 
+
+atomic_number(atom::PState) = atomic_number(atomic_symbol(atom))
+
+""" 
+Generate an atom with the given properties. 
+"""
+atom(at; properties = (position, atomic_mass, atomic_symbol)) = 
+      PState((; [symbol(p) => p(at) for p in properties]...))
+
+# --------------------------------------------------- 
+
+mutable struct AosSystem{D, TCELL, TPART} <: AbstractSystem{D} 
+   cell::TCELL 
+   particles::Vector{TPART}
+   # -------- 
+   meta::Dict{String, Any}
+end
+
+
+function AosSystem(sys::AbstractSystem; 
+                   properties = (position, atomic_mass, atomic_symbol), )
+
+   X = [ atom(sys[i]; properties = properties) for i = 1:length(sys) ]
+   cell = AtomsBase.get_cell(sys)
+   D = AtomsBase.n_dimensions(cell)
+   return AosSystem{D, typeof(cell), eltype(X)}(cell, X, Dict{String, Any}())
+end
+
+
+# ---------------------------------------------------
+# implementing the interface 
+
+Base.length(at::AosSystem) = length(at.particles)
+
+Base.getindex(at::AosSystem, i::Int) = at.particles[i]
+
+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]
+end
+
+AtomsBase.get_cell(at::AosSystem) = at.cell
+
+AtomsBase.n_dimensions(at::AosSystem) = AtomsBase.n_dimensions(at.cell)
+AtomsBase.bounding_box(at::AosSystem) = AtomsBase.bounding_box(at.cell)
+AtomsBase.boundary_conditions(at::AosSystem) = AtomsBase.boundary_conditions(at.cell)
+AtomsBase.periodicity(at::AosSystem) = AtomsBase.periodicity(at.cell)
+
diff --git a/test/test_atomsbase.jl b/test/test_atomsbase.jl
new file mode 100644
index 0000000..9db321e
--- /dev/null
+++ b/test/test_atomsbase.jl
@@ -0,0 +1,41 @@
+
+using DecoratedParticles, AtomsBase, StaticArrays, Unitful, Test 
+using AtomsBase: ChemicalElement, Atom 
+using AtomsBuilder: bulk, rattle! 
+DP = DecoratedParticles
+
+## 
+#generate an atom and check that the accessors work
+
+x = PState(𝐫 = SA[1.0, 2.0, 3.0], 𝐯 = SA[0.1, 0.2, 0.3], 
+           m = 1.0, Z = ChemicalElement(6) )
+display(x)           
+@test position(x) == x.𝐫
+@test velocity(x) == x.𝐯
+@test atomic_mass(x) == x.m
+@test atomic_symbol(x) == x.Z
+@test atomic_number(x) == 6
+
+## 
+#convert an Atom 
+
+at = Atom(6, SA[1.0, 2.0, 3.0]u"Å"; atomic_mass = 1.0u"u")
+x = DP.atom(at; properties = (position, atomic_mass, atomic_symbol))
+display(x)
+@test x.𝐫 == position(x) == position(at)
+@test x.m == atomic_mass(x) == atomic_mass(at)
+@test x.Z == atomic_symbol(x) == atomic_symbol(at)
+
+
+## 
+# convert an entire system 
+
+sys = rattle!(bulk(:Si, cubic=true) * 2, 0.1);
+aos = DP.AosSystem(sys);
+
+aos[1]
+aos[1, position]
+aos[1, atomic_mass]
+atomic_mass(aos, 1)
+
+get_cell(aos)
\ No newline at end of file