From 41acc63896dfe5e02ab6d20200f9705683df566d Mon Sep 17 00:00:00 2001
From: Rashid Rafeek <rashidrafeek@gmail.com>
Date: Thu, 26 Dec 2024 19:38:28 +0530
Subject: [PATCH] Add pbc_shortest_vector function (#11)

---
 Project.toml        |  2 +-
 src/AtomsToolbox.jl |  2 +-
 src/pbc_utils.jl    | 37 +++++++++++++++++++++++++++++++++++++
 src/piracy.jl       |  2 +-
 test/Project.toml   |  1 +
 test/runtests.jl    | 35 +++++++++++++++++++++++++++++++++++
 6 files changed, 76 insertions(+), 3 deletions(-)

diff --git a/Project.toml b/Project.toml
index c7af703..851c26f 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "AtomsToolbox"
 uuid = "f386ec8f-7408-45ea-aa1c-e80969af901e"
 authors = ["Rashid Rafeek <rashidrafeek@gmail.com> and contributors"]
-version = "0.1.2"
+version = "0.1.3"
 
 [deps]
 AtomsBase = "a963bdd2-2df7-4f54-a1ee-49d51e6be12a"
diff --git a/src/AtomsToolbox.jl b/src/AtomsToolbox.jl
index 62945ea..a0b959f 100644
--- a/src/AtomsToolbox.jl
+++ b/src/AtomsToolbox.jl
@@ -18,7 +18,7 @@ using AtomsBase: AtomsBase,
 using Unitful: Unitful, ustrip, @u_str, unit
 import Graphs
 using LinearAlgebra: LinearAlgebra, ⋅, det, norm, lu, ×
-using StaticArrays: StaticMatrix, Size
+using StaticArrays: StaticMatrix, SVector, Size
 import Base: angle, sort # To extend for AbstractSystem 
 
 # Some type piracy to deal with other packages
diff --git a/src/pbc_utils.jl b/src/pbc_utils.jl
index 30169b6..ded2324 100644
--- a/src/pbc_utils.jl
+++ b/src/pbc_utils.jl
@@ -5,6 +5,43 @@ const images_view = permutedims(reduce(hcat,
                    [[a,b,c] for (a,b,c) in Iterators.product(-1:1, -1:1, -1:1)][:])
                )
 
+function pbc_shortest_vector(system::AbstractSystem, at1::Int, at2::Int)
+    pos1 = position(system, at1)
+    pos2 = position(system, at2)
+
+    return pbc_shortest_vector(system, pos1, pos2)
+end
+function pbc_shortest_vector(system::AbstractSystem, pos1::AbstractVector, pos2::AbstractVector)
+    diff = pos2 - pos1
+    cell = cell_matrix(system)
+
+    return pbc_shortest_vector(diff, cell)
+end
+function pbc_shortest_vector(diff::AbstractVector, cell::AbstractMatrix)
+    cell_inv = inv(cell)
+    # Convert to fractional coordinates
+    frac = cell_inv * diff
+    
+    # For each component, find the image that minimizes the distance
+    # Check -1, 0, +1 translations in each direction
+    min_dist = Inf*oneunit(eltype(diff))
+    min_image = diff
+    
+    for i in -1:1, j in -1:1, k in -1:1
+        shift = SVector(i, j, k)
+        trial_frac = frac + shift
+        trial_cart = cell * trial_frac
+        # dist = trial_cart ⋅ trial_cart  # squared distance
+        dist = norm(trial_cart)  # distance
+        if dist < min_dist
+            min_dist = dist
+            min_image = trial_cart
+        end
+    end
+    
+    return min_image
+end
+
 function dot_2d_mod(a,b)
     I = size(a)[1]
     J = size(b)[2]
diff --git a/src/piracy.jl b/src/piracy.jl
index d004a0d..0395e63 100644
--- a/src/piracy.jl
+++ b/src/piracy.jl
@@ -10,6 +10,6 @@ function Base.inv(x::StaticMatrix{N,M,T}) where {N,M,T <: Unitful.AbstractQuanti
 end 
 
 # These were removed in the AtomsBase 0.4 update
-Base.position(sys::AbstractSystem) = position.(sys, :)
+Base.position(sys::AbstractSystem) = position.(Ref(sys), 1:length(sys))
 AtomsBase.atomic_symbol(sys::AbstractSystem) = atomic_symbol.(species(sys, :))
 AtomsBase.atomic_number(sys::AbstractSystem) = atomic_number.(species(sys, :))
diff --git a/test/Project.toml b/test/Project.toml
index fd02eb6..085f66b 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -1,5 +1,6 @@
 [deps]
 AtomsBase = "a963bdd2-2df7-4f54-a1ee-49d51e6be12a"
 AtomsIO = "1692102d-eeb4-4df9-807b-c9517f998d44"
+StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
diff --git a/test/runtests.jl b/test/runtests.jl
index ffc47a2..0bba965 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -2,6 +2,7 @@ using Test
 using AtomsToolbox
 using AtomsBase, Unitful
 using AtomsIO
+using StaticArrays: SVector
 
 @testset "AtomsToolbox.jl" begin
     # Si primitive crystal 
@@ -68,4 +69,38 @@ using AtomsIO
         end
     end
 
+    @testset "PBC" begin
+        # Define a simple cubic cell with side length 3 Å.
+        box = (
+            [3.0, 0.0, 0.0]u"Å", 
+            [0.0, 3.0, 0.0]u"Å", 
+            [0.0, 0.0, 3.0]u"Å"
+        )
+
+        # Create a small system with two atoms:
+        #    - First at [0, 0, 0] Å
+        #    - Second at [2.5, 2.5, 2.5] Å
+        # We expect pbc_shortest_vector to produce [-0.5, -0.5, -0.5] Å 
+        # as the minimal image distance from the first to second atom.
+        atom1 = Atom(:X, [0.0, 0.0, 0.0]u"Å")
+        atom2 = Atom(:X, [2.5, 2.5, 2.5]u"Å")
+
+        # Build a FlexibleSystem (from AtomsBase) with periodic boundary conditions.
+        test_sys = FlexibleSystem(
+            [atom1, atom2]; 
+            cell_vectors = box, 
+            periodicity = (true, true, true)
+        )
+
+        @testset "Two-atom cubic system" begin
+            vec = AtomsToolbox.pbc_shortest_vector(test_sys, 1, 2)
+            @test vec ≈ SVector(-0.5, -0.5, -0.5)u"Å" atol=1e-10u"Å"
+
+            # Also test the version that accepts explicit positions:
+            pos1 = [0.0, 0.0, 0.0]u"Å"
+            pos2 = [2.5, 2.5, 2.5]u"Å"
+            vec2 = AtomsToolbox.pbc_shortest_vector(test_sys, pos1, pos2)
+            @test vec2 ≈ SVector(-0.5, -0.5, -0.5)u"Å" atol=1e-10u"Å"
+        end
+    end
 end