add matscipy neighborlist option (#1)

---------

Co-authored-by: Linux-cpp-lisp <[email protected]>

CHANGELOG.md

@@ -23,6 +23,7 @@ Most recent change on the bottom.
 - `nequip-deploy --using-dataset` to support data-dependent deployment steps
 - Support for Gaussian Mixture Model uncertainty quantification (https://doi.org/10.1063/5.0136574)
 - `nequip-deploy build --checkpoint` and `--override` to avoid many largely duplicated YAML files
+- matscipy neighborlist support enabled with `NEQUIP_MATSCIPY_NL` environment variable
 
 ### Changed
 - Always require explicit `seed`

nequip/data/AtomicData.py

@@ -5,7 +5,7 @@
 
 import warnings
 from copy import deepcopy
-from typing import Union, Tuple, Dict, Optional, List, Set, Sequence
+from typing import Union, Tuple, Dict, Optional, List, Set, Sequence, Final
 from collections.abc import Mapping
 import os
 
@@ -705,6 +705,13 @@ def without_nodes(self, which_nodes):
 assert _ERROR_ON_NO_EDGES in ("true", "false")
 _ERROR_ON_NO_EDGES = _ERROR_ON_NO_EDGES == "true"
 
+_NEQUIP_MATSCIPY_NL: Final[bool] = os.environ.get("NEQUIP_MATSCIPY_NL", "false").lower()
+assert _NEQUIP_MATSCIPY_NL in ("true", "false")
+_NEQUIP_MATSCIPY_NL = _NEQUIP_MATSCIPY_NL == "true"
+
+if _NEQUIP_MATSCIPY_NL:
+    import matscipy.neighbours
+
 
 def neighbor_list_and_relative_vec(
     pos,
@@ -783,15 +790,25 @@ def neighbor_list_and_relative_vec(
     # ASE dependent part
     temp_cell = ase.geometry.complete_cell(temp_cell)
 
-    first_idex, second_idex, shifts = ase.neighborlist.primitive_neighbor_list(
-        "ijS",
-        pbc,
-        temp_cell,
-        temp_pos,
-        cutoff=float(r_max),
-        self_interaction=strict_self_interaction,  # we want edges from atom to itself in different periodic images!
-        use_scaled_positions=False,
-    )
+    if _NEQUIP_MATSCIPY_NL:
+        assert strict_self_interaction and not self_interaction
+        first_idex, second_idex, shifts = matscipy.neighbours.neighbour_list(
+            "ijS",
+            pbc=pbc,
+            cell=temp_cell,
+            positions=temp_pos,
+            cutoff=float(r_max),
+        )
+    else:
+        first_idex, second_idex, shifts = ase.neighborlist.primitive_neighbor_list(
+            "ijS",
+            pbc,
+            temp_cell,
+            temp_pos,
+            cutoff=float(r_max),
+            self_interaction=strict_self_interaction,  # we want edges from atom to itself in different periodic images!
+            use_scaled_positions=False,
+        )
 
     # Eliminate true self-edges that don't cross periodic boundaries
     if not self_interaction: