From 4ae1272810578ffbe659591f2a49de08f19828d1 Mon Sep 17 00:00:00 2001
From: JuliaMolSim <christohortner@gmail.com>
Date: Wed, 22 May 2024 10:01:41 -0700
Subject: [PATCH] cleanup, union, deleteat!

---
 Project.toml      |  4 +++-
 src/utils.jl      | 61 +++++++++++++++++++++++++----------------------
 test/test_bulk.jl | 13 +++++++++-
 3 files changed, 47 insertions(+), 31 deletions(-)

diff --git a/Project.toml b/Project.toml
index bc7a6dd..07cb144 100644
--- a/Project.toml
+++ b/Project.toml
@@ -6,6 +6,7 @@ version = "0.0.2"
 AtomsBase = "a963bdd2-2df7-4f54-a1ee-49d51e6be12a"
 JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 
@@ -13,13 +14,14 @@ Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 AtomsBase = "0.3.5"
 JSON = "0.21.4"
 LinearAlgebra = "1.9.0, 1.10.0"
+Random = "1.9.0, 1.10.0"
 StaticArrays = "1.9"
 Unitful = "1.19"
 julia = "1.9.0, 1.10.0"
 
 [extras]
-Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 JuLIP = "945c410c-986d-556a-acb1-167a618e0462"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
 test = ["Test", "JuLIP"]
diff --git a/src/utils.jl b/src/utils.jl
index 5734a9a..9730808 100644
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -1,7 +1,7 @@
 
 using AtomsBase 
 using AtomsBase: Atom, FlexibleSystem, Periodic
-using Unitful: unit, ustrip 
+using Unitful: unit, ustrip, Quantity
 using LinearAlgebra: norm 
 
 export rattle!, 
@@ -93,50 +93,53 @@ import Base.*
 
 
 """
-`rattle!(at, r::Float64) -> at`
+```
+rattle!(sys, r::Union{AbstractFloat, Quantity}) -> at
+```
 
 Randomly perturbs the atom positions within a ball of radius `r`. The perturbation 
 is uniform in angular component, and uniform in radial component. (Note this is 
 not the same as choosing them uniform in cartesian coordinates!). 
+
+If `r` is unitless, then the unit of the system is applied. 
 """
-function rattle!(at::FlexibleSystem, r::AbstractFloat)
+function rattle!(at::FlexibleSystem, r::Quantity)
    for i = 1:length(at.particles)
       p = at.particles[i]
       𝐫ᵢ = p.position 
       T = typeof(ustrip(𝐫ᵢ[1]))
       ui = randn(Vec3{T})
-      p_new = _set_position(p, 𝐫ᵢ + r * ui / norm(ui) * unit(𝐫ᵢ[1]))
+      p_new = _set_position(p, 𝐫ᵢ + r * ui / norm(ui))
       at.particles[i] = p_new
    end
    return at
 end
 
+rattle!(sys::FlexibleSystem, r::AbstractFloat) = 
+      rattle!(sys, r * unit(position(sys)[1][1]))
 
 
-# union(at1::Atoms, at2::Atoms) =
-#    Atoms( X = union(at1.X, at2.X),
-#           P = union(at1.P, at2.P),
-#           M = union(at1.M, at2.M),
-#           Z = union(at1.Z, at2.Z),
-#           cell = cell(at1),
-#           pbc = pbc(at1) )
-
-
-
-# """
-# `deleteat!(at::Atoms, n) -> at`:
-
-# returns the same atoms object `at`, but with the atom(s) specified by `n`
-# removed.
-# """
-# function Base.deleteat!(at::Atoms, n)
-#    deleteat!(at.X, n)
-#    deleteat!(at.P, n)
-#    deleteat!(at.M, n)
-#    deleteat!(at.Z, n)
-#    update_data!(at, Inf)
-#    JuLIP.reset_clamp!(at)
-#    return at
-# end
+"""
+```
+union(sys1::FlexibleSystem, sys2::FlexibleSystem)
+```
+takes the union of two particle systems provided their cells are identical. 
+"""
+function union(sys1::FlexibleSystem, sys2::FlexibleSystem) 
+   @assert boundary_conditions(sys1) == boundary_conditions(sys2)
+   @assert bounding_box(sys1) == bounding_box(sys2)
+   return FlexibleSystem(union(sys1.particles, sys2.particles),  
+                        bounding_box(at), 
+                        boundary_conditions(at) )
+end
 
+"""
+`deleteat!(sys::FlexibleSystem, n) -> sys`:
 
+returns the same `FlexibleSystem`` object `sys`, but with the atom(s) specified by `n`
+removed.
+"""
+function Base.deleteat!(sys::FlexibleSystem, n)
+   deleteat!(sys.particles, n)
+   return sys
+end 
\ No newline at end of file
diff --git a/test/test_bulk.jl b/test/test_bulk.jl
index dc3cf12..c84aa39 100644
--- a/test/test_bulk.jl
+++ b/test/test_bulk.jl
@@ -1,5 +1,5 @@
 
-using AtomsBuilder, Test, AtomsBase, Unitful
+using AtomsBuilder, Test, AtomsBase, Unitful, Random
 import JuLIP
 
 ##
@@ -39,6 +39,7 @@ end
 ##
 
 # not sure how to write a test for this, but at least it should execute
+sys0 = rattle!(bulk(:C, cubic=true) * (2,3,4), 0.1u"Å")
 sys1 = rattle!(bulk(:C, cubic=true) * (2,3,4), 0.1)
 sys2 = rattle!(bulk(:C) * (2,3,4), 0.1)
 
@@ -54,3 +55,13 @@ Znew = copy(Z); Znew[3:5:end] .= zO
 sys4 = set_elements(sys3, Znew)
 @test all(atomic_number(sys4) .== Znew)
 
+pold = deepcopy(sys4.particles) 
+deleteat!(sys4, 1:5)
+@test position.(pold[6:end]) == position(sys4)
+@test atomic_mass.(pold[6:end]) == atomic_mass(sys4)
+@test atomic_number.(pold[6:end]) == atomic_number(sys4)
+
+
+
+
+