From 0cf096ad40db2867280150c34edf35bcf92ee9dc Mon Sep 17 00:00:00 2001
From: Yeonghun Kang <62047140+Yeonghun1675@users.noreply.github.com>
Date: Fri, 24 May 2024 19:05:38 +0900
Subject: [PATCH] Update potential.pyx

---
 fast_grid/libs/potential.pyx | 93 +++++++++++++++++++++++-------------
 1 file changed, 59 insertions(+), 34 deletions(-)

diff --git a/fast_grid/libs/potential.pyx b/fast_grid/libs/potential.pyx
index 13e1aa7..7d7d45a 100644
--- a/fast_grid/libs/potential.pyx
+++ b/fast_grid/libs/potential.pyx
@@ -8,31 +8,6 @@ cimport numpy as np
 from libc.math cimport exp
 
 
-def minimum_image_triclinic(np.ndarray[np.float64_t, ndim=1] dx, np.ndarray[np.float64_t, ndim=2] box):
-    cdef int ix, iy, iz
-    cdef float rx, ry0, ry1, rz0, rz1, rz2, dsq
-    cdef float dsq_min = np.finfo(np.float64).max
-    cdef np.ndarray[np.float64_t, ndim=1] dx_min = np.zeros(3, dtype=np.float64)
-
-    for ix in range(-1, 2):
-        rx = dx[0] + box[0, 0] * ix
-        for iy in range(-1, 2):
-            ry0 = rx + box[1, 0] * iy
-            ry1 = dx[1] + box[1, 1] * iy
-            for iz in range(-1, 2):
-                rz0 = ry0 + box[2, 0] * iz
-                rz1 = ry1 + box[2, 1] * iz
-                rz2 = dx[2] + box[2, 2] * iz
-                dsq = rz0 * rz0 + rz1 * rz1 + rz2 * rz2
-                if dsq < dsq_min:
-                    dsq_min = dsq
-                    dx_min[0] = rz0
-                    dx_min[1] = rz1
-                    dx_min[2] = rz2
-
-    dx[:] = dx_min
-
-
 @cython.boundscheck(False)
 @cython.wraparound(False)
 @cython.cdivision(True)
@@ -48,6 +23,11 @@ def lj_potential_cython(np.ndarray[np.float64_t, ndim=2] pos1,
     cdef double r2, lj6, lj12, inv_r2, inv_r6, inv_r12, e, s, s6, s12, energy
     cdef double threshold = 1e-10
     cdef double cutoff_squared = cutoff * cutoff
+
+    cdef int ix, iy, iz
+    cdef float rx, ry0, ry1, rz0, rz1, rz2, dsq
+    cdef float dsq_min = np.finfo(np.float64).max
+    cdef np.ndarray[np.float64_t, ndim=1] dx_min = np.zeros(3, dtype=np.float64)
     
     cdef np.ndarray[np.float64_t, ndim=1] energy_grid = np.zeros((n), dtype=np.float64)
     cdef np.ndarray[np.float64_t, ndim=1] diff = np.zeros(3, dtype=np.float64)
@@ -58,11 +38,31 @@ def lj_potential_cython(np.ndarray[np.float64_t, ndim=2] pos1,
             diff[0] = pos2[j, 0] - pos1[i, 0]
             diff[1] = pos2[j, 1] - pos1[i, 1]
             diff[2] = pos2[j, 2] - pos1[i, 2]
-            
-            with gil:
-                minimum_image_triclinic(diff, box)
-            r2 = diff[0] * diff[0] + diff[1] * diff[1] + diff[2] * diff[2]
 
+            dsq_min = 10000.
+            dx_min[0] = 0
+            dx_min[1] = 0
+            dx_min[2] = 0
+
+            # minimum_image_triclinic
+            for ix in range(-1, 2):
+                rx = diff[0] + box[0, 0] * ix
+                for iy in range(-1, 2):
+                    ry0 = rx + box[1, 0] * iy
+                    ry1 = diff[1] + box[1, 1] * iy
+                    for iz in range(-1, 2):
+                        rz0 = ry0 + box[2, 0] * iz
+                        rz1 = ry1 + box[2, 1] * iz
+                        rz2 = diff[2] + box[2, 2] * iz
+                        dsq = rz0 * rz0 + rz1 * rz1 + rz2 * rz2
+                        if dsq < dsq_min:
+                            dsq_min = dsq
+                            dx_min[0] = rz0
+                            dx_min[1] = rz1
+                            dx_min[2] = rz2
+
+            r2 = dx_min[0] * dx_min[0] + dx_min[1] * dx_min[1] + dx_min[2] * dx_min[2]
+            
             if r2 < cutoff_squared and r2 > threshold:
                 e = epsilon[j]
                 s = sigma[j]
@@ -97,6 +97,11 @@ def gaussian_cython(np.ndarray[np.float64_t, ndim=2] pos1,
     cdef double threshold = 1e-10
     cdef double width_squared = width * width
     cdef double cutoff_squared = cutoff * cutoff
+
+    cdef int ix, iy, iz
+    cdef float rx, ry0, ry1, rz0, rz1, rz2, dsq
+    cdef float dsq_min = np.finfo(np.float64).max
+    cdef np.ndarray[np.float64_t, ndim=1] dx_min = np.zeros(3, dtype=np.float64)
     
     cdef np.ndarray[np.float64_t, ndim=1] energy_grid = np.zeros((n), dtype=np.float64)
     cdef np.ndarray[np.float64_t, ndim=1] diff = np.zeros(3, dtype=np.float64)
@@ -107,14 +112,34 @@ def gaussian_cython(np.ndarray[np.float64_t, ndim=2] pos1,
             diff[0] = pos2[j, 0] - pos1[i, 0]
             diff[1] = pos2[j, 1] - pos1[i, 1]
             diff[2] = pos2[j, 2] - pos1[i, 2]
-            
-            with gil:
-                minimum_image_triclinic(diff, box)
-            r2 = diff[0] * diff[0] + diff[1] * diff[1] + diff[2] * diff[2]
+
+            dsq_min = 10000.
+            dx_min[0] = 0
+            dx_min[1] = 0
+            dx_min[2] = 0
+
+            # minimum_image_triclinic
+            for ix in range(-1, 2):
+                rx = diff[0] + box[0, 0] * ix
+                for iy in range(-1, 2):
+                    ry0 = rx + box[1, 0] * iy
+                    ry1 = diff[1] + box[1, 1] * iy
+                    for iz in range(-1, 2):
+                        rz0 = ry0 + box[2, 0] * iz
+                        rz1 = ry1 + box[2, 1] * iz
+                        rz2 = diff[2] + box[2, 2] * iz
+                        dsq = rz0 * rz0 + rz1 * rz1 + rz2 * rz2
+                        if dsq < dsq_min:
+                            dsq_min = dsq
+                            dx_min[0] = rz0
+                            dx_min[1] = rz1
+                            dx_min[2] = rz2
+
+            r2 = dx_min[0] * dx_min[0] + dx_min[1] * dx_min[1] + dx_min[2] * dx_min[2]
 
             if r2 < cutoff_squared and r2 > threshold:
                 energy += height * exp(-1 * r2 / width_squared)
         
         energy_grid[i] += energy
     
-    return energy_grid
\ No newline at end of file
+    return energy_grid