Update to the latest vesin release

Author: Luthaf

python/metatomic_ase/setup.py

@@ -122,7 +122,7 @@ def create_version_number(version):
         # No dependency on ASE itself until this package is no longer a direct
         # dependency of metatomic-torch
         # "ase >=3.22.0",
-        "vesin >=0.5.2,<0.6",
+        "vesin >=0.5.5,<0.6",
     ]
 
     # when packaging a sdist for release, we should never use local dependencies

python/metatomic_ase/src/metatomic_ase/_calculator.py

@@ -21,7 +21,7 @@
     pick_output,
 )
 
-from ._neighbors import _compute_requested_neighbors
+from ._neighbors import AllNeighborsCalculator
 
 
 import ase  # isort: skip
@@ -317,6 +317,11 @@ def __init__(
                     "be positive"
                 )
 
+        self._nl_calculators = AllNeighborsCalculator(
+            requested_options=self._model.requested_neighbor_lists(),
+            check_consistency=check_consistency,
+        )
+
         # We do our own check to verify if a property is implemented in `calculate()`,
         # so we pretend to be able to compute all properties ASE knows about.
         self.implemented_properties = ALL_ASE_PROPERTIES
@@ -398,11 +403,7 @@ def run_model(
             systems.append(system)
 
         # Compute the neighbors lists requested by the model
-        input_systems = _compute_requested_neighbors(
-            systems=systems,
-            requested_options=self._model.requested_neighbor_lists(),
-            check_consistency=self.parameters["check_consistency"],
-        )
+        input_systems = self._nl_calculators.compute(systems=systems)
 
         available_outputs = self._model.capabilities().outputs
         for key in outputs:
@@ -538,11 +539,7 @@ def calculate(
         with record_function("MetatomicCalculator::compute_neighbors"):
             # convert from ase.Atoms to metatomic.torch.System
             system = System(types, positions, cell, pbc)
-            input_system = _compute_requested_neighbors(
-                systems=[system],
-                requested_options=self._model.requested_neighbor_lists(),
-                check_consistency=self.parameters["check_consistency"],
-            )[0]
+            input_system = self._nl_calculators.compute(systems=[system])[0]
 
         with record_function("MetatomicCalculator::get_model_inputs"):
             for name, option in self._model.requested_inputs().items():
@@ -721,11 +718,7 @@ def compute_energy(
             systems.append(system)
 
         # Compute the neighbors lists requested by the model
-        input_systems = _compute_requested_neighbors(
-            systems=systems,
-            requested_options=self._model.requested_neighbor_lists(),
-            check_consistency=self.parameters["check_consistency"],
-        )
+        input_systems = self._nl_calculators.compute(systems=systems)
 
         predictions = self._model(
             systems=input_systems,

python/metatomic_ase/src/metatomic_ase/_neighbors.py

@@ -31,57 +31,68 @@
     HAS_NVALCHEMIOPS = False
 
 
-def _compute_requested_neighbors(
-    systems: List[System],
-    requested_options: List[NeighborListOptions],
-    check_consistency=False,
-) -> List[System]:
-    """
-    Compute all neighbor lists requested by ``model`` and store them inside the systems.
-    """
-    can_use_nvalchemi = HAS_NVALCHEMIOPS and all(
-        system.device.type == "cuda" for system in systems
-    )
-
-    if can_use_nvalchemi:
-        full_nl_options = []
-        half_nl_options = []
-        for options in requested_options:
-            if options.full_list:
-                full_nl_options.append(options)
-            else:
-                half_nl_options.append(options)
-
-        # Do the full neighbor lists with nvalchemi, and the rest with vesin
-        systems = _compute_requested_neighbors_nvalchemi(
-            systems=systems,
-            requested_options=full_nl_options,
-        )
-        systems = _compute_requested_neighbors_vesin(
-            systems=systems,
-            requested_options=half_nl_options,
-            check_consistency=check_consistency,
+class AllNeighborsCalculator:
+    def __init__(
+        self,
+        requested_options: List[NeighborListOptions],
+        check_consistency=False,
+    ):
+        self.check_consistency = check_consistency
+        self._full_nl_options = [
+            options for options in requested_options if options.full_list
+        ]
+        self._full_vesin_calculators = [
+            vesin.metatomic.NeighborList(
+                options=options,
+                length_unit="angstrom",
+                check_consistency=check_consistency,
+            )
+            for options in requested_options
+            if options.full_list
+        ]
+        self._half_vesin_calculators = [
+            vesin.metatomic.NeighborList(
+                options=options,
+                length_unit="angstrom",
+                check_consistency=check_consistency,
+            )
+            for options in requested_options
+            if not options.full_list
+        ]
+
+    def compute(self, systems: List[System]) -> List[System]:
+        assert isinstance(systems, list)
+        assert isinstance(systems[0], torch.ScriptObject)
+
+        can_use_nvalchemi = HAS_NVALCHEMIOPS and all(
+            system.device.type == "cuda" for system in systems
         )
-    else:
+
+        if can_use_nvalchemi:
+            # Do the full neighbor lists with nvalchemi
+            systems = _compute_requested_neighbors_nvalchemi(
+                systems=systems,
+                requested_options=self._full_nl_options,
+            )
+        else:
+            systems = _compute_requested_neighbors_vesin(
+                systems=systems,
+                calculators=self._full_vesin_calculators,
+            )
+
+        # always compute the half neighbor lists with vesin
         systems = _compute_requested_neighbors_vesin(
             systems=systems,
-            requested_options=requested_options,
-            check_consistency=check_consistency,
+            calculators=self._half_vesin_calculators,
         )
 
-    return systems
+        return systems
 
 
 def _compute_requested_neighbors_vesin(
     systems: List[System],
-    requested_options: List[NeighborListOptions],
-    check_consistency=False,
+    calculators: List[vesin.metatomic.NeighborList],
 ) -> List[System]:
-    """
-    Compute all neighbor lists requested by ``model`` and store them inside the systems,
-    using vesin.
-    """
-
     system_devices = []
     moved_systems = []
     for system in systems:
@@ -91,12 +102,13 @@ def _compute_requested_neighbors_vesin(
         else:
             moved_systems.append(system)
 
-    vesin.metatomic.compute_requested_neighbors_from_options(
-        systems=moved_systems,
-        system_length_unit="angstrom",
-        options=requested_options,
-        check_consistency=check_consistency,
-    )
+    for calculator in calculators:
+        calculator.add_neighbor_list(
+            systems=moved_systems,
+            # if we have more than one system, we can no keep the data as a reference
+            # to memory allocated in the calculator and we need to make a copy
+            copy=len(systems) > 1,
+        )
 
     systems = []
     for system, device in zip(moved_systems, system_devices, strict=True):
@@ -142,6 +154,7 @@ def _compute_requested_neighbors_nvalchemi(systems, requested_options):
                         "cell_shift_c",
                     ],
                     values=torch.hstack([P, S]),
+                    assume_unique=True,
                 ),
                 components=[
                     Labels("xyz", torch.tensor([[0], [1], [2]], device=system.device))

python/metatomic_torchsim/metatomic_torchsim/_model.py

@@ -28,7 +28,7 @@
     pick_output,
 )
 
-from ._neighbors import _compute_requested_neighbors
+from ._neighbors import AllNeighborsCalculator
 
 
 try:
@@ -249,7 +249,6 @@ def __init__(
                     "be positive"
                 )
 
-        self._requested_neighbor_lists = self._model.requested_neighbor_lists()
         self._requested_inputs = self._model.requested_inputs()
         if len(self._requested_inputs) != 0:
             raise ValueError(
@@ -283,6 +282,11 @@ def __init__(
             outputs=run_outputs,
         )
 
+        self._nl_calculators = AllNeighborsCalculator(
+            requested_options=self._model.requested_neighbor_lists(),
+            check_consistency=check_consistency,
+        )
+
         self.additional_outputs: Dict[str, TensorMap] = {}
         """
         Additional outputs computed by :py:meth:`forward` are stored here.
@@ -355,11 +359,7 @@ def forward(self, state: "ts.SimState") -> Dict[str, torch.Tensor]:
             )
 
         # Compute neighbor lists
-        systems = _compute_requested_neighbors(
-            systems=systems,
-            requested_options=self._requested_neighbor_lists,
-            check_consistency=self._check_consistency,
-        )
+        systems = self._nl_calculators.compute(systems=systems)
 
         # Run the model (evaluation options precomputed in __init__)
         model_outputs = self._model(