Speed up neighborlists

molmod · Oct 27, 2024 · 43b5713 · 43b5713
1 parent c062027
commit 43b5713
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/.gitignore b/.gitignore
@@ -4,3 +4,4 @@ __pycache__
 src/tinyff/_version.py
 dist
 build
+tmp
diff --git a/src/tinyff/neighborlist.py b/src/tinyff/neighborlist.py
@@ -26,6 +26,7 @@
 
 __all__ = ("NLIST_DTYPE", "NBuildSimple", "NBuildCellLists")
 
+
 NLIST_DTYPE = [
     # First atoms.
     ("iatom0", int),
@@ -189,7 +190,9 @@ def _apply_mic(deltas: NDArray[float], cell_lengths: NDArray[float]):
     deltas /= cell_lengths
     deltas -= np.round(deltas)
     deltas *= cell_lengths
-    return np.linalg.norm(deltas, axis=1)
+    dist = np.einsum("ij,ij->i", deltas, deltas)
+    np.sqrt(dist, out=dist)
+    return dist
 
 
 @attrs.define
-Original file line number
+Diff line change
@@ Expand Up / @@ -4,3 +4,4 @@ __pycache__ @@
     src/tinyff/_version.py
     dist
     build
+    tmp