Upgrade to metatensor-torch 0.6.0

pyproject.toml

@@ -10,9 +10,9 @@ authors = [{name = "metatrain developers"}]
 
 dependencies = [
     "ase < 3.23.0",
-    "metatensor-learn==0.2.3",
-    "metatensor-operations==0.2.3",
-    "metatensor-torch==0.5.5",
+    "metatensor-learn==0.3.0",
+    "metatensor-operations==0.3.0",
+    "metatensor-torch==0.6.0",
     "jsonschema",
     "omegaconf",
     "python-hostlist",
@@ -59,7 +59,7 @@ build-backend = "setuptools.build_meta"
 
 [project.optional-dependencies]
 soap-bpnn = [
-    "rascaline-torch @ git+https://github.com/luthaf/rascaline@d181b28#subdirectory=python/rascaline-torch",
+    "rascaline-torch @ git+https://github.com/luthaf/rascaline@5326b6e#subdirectory=python/rascaline-torch",
 ]
 alchemical-model = [
   "torch_alchemical @ git+https://github.com/abmazitov/torch_alchemical.git@51ff519",
@@ -68,7 +68,7 @@ pet = [
   "pet @ git+https://github.com/lab-cosmo/pet@7eddb2e",
 ]
 gap = [
-    "rascaline-torch @ git+https://github.com/luthaf/rascaline@d181b28#subdirectory=python/rascaline-torch",
+    "rascaline-torch @ git+https://github.com/luthaf/rascaline@5326b6e#subdirectory=python/rascaline-torch",
     "skmatter",
     "metatensor-learn",
     "scipy",

src/metatrain/experimental/alchemical_model/model.py

@@ -78,6 +78,7 @@ def requested_neighbor_lists(self) -> List[NeighborListOptions]:
             NeighborListOptions(
                 cutoff=self.cutoff,
                 full_list=True,
+                strict=True,
             )
         ]
 

src/metatrain/experimental/alchemical_model/tests/test_exported.py

@@ -33,6 +33,7 @@ def test_to(device, dtype):
         types=torch.tensor([6, 6]),
         positions=torch.tensor([[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]]),
         cell=torch.zeros(3, 3),
+        pbc=torch.tensor([False, False, False]),
     )
     requested_neighbor_lists = get_requested_neighbor_lists(exported)
     system = get_system_with_neighbor_lists(system, requested_neighbor_lists)

src/metatrain/experimental/alchemical_model/tests/test_functionality.py

@@ -27,6 +27,7 @@ def test_prediction_subset_elements():
         types=torch.tensor([6, 6]),
         positions=torch.tensor([[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]]),
         cell=torch.zeros(3, 3),
+        pbc=torch.tensor([False, False, False]),
     )
     requested_neighbor_lists = get_requested_neighbor_lists(model)
     system = get_system_with_neighbor_lists(system, requested_neighbor_lists)

src/metatrain/experimental/alchemical_model/tests/test_torch_alchemical_compatibility.py

@@ -29,6 +29,7 @@
 nl_options = NeighborListOptions(
     cutoff=5.0,
     full_list=True,
+    strict=True,
 )
 systems = [get_system_with_neighbor_lists(system, [nl_options]) for system in systems]
 

src/metatrain/experimental/pet/model.py

@@ -73,6 +73,7 @@ def requested_neighbor_lists(
             NeighborListOptions(
                 cutoff=self.cutoff,
                 full_list=True,
+                strict=True,
             )
         ]
 
@@ -111,9 +112,7 @@ def forward(
                 values=predictions,
             )
             if selected_atoms is not None:
-                block = metatensor.torch.slice_block(
-                    block, axis="samples", labels=selected_atoms
-                )
+                block = metatensor.torch.slice_block(block, "samples", selected_atoms)
             output_tmap = TensorMap(keys=empty_labels, blocks=[block])
             if not outputs[output_name].per_atom:
                 output_tmap = metatensor.torch.sum_over_samples(output_tmap, "atom")

src/metatrain/experimental/pet/tests/test_exported.py

@@ -61,6 +61,7 @@ def test_to(device):
         types=torch.tensor([6, 6]),
         positions=torch.tensor([[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]]),
         cell=torch.zeros(3, 3),
+        pbc=torch.tensor([False, False, False]),
     )
     requested_neighbor_lists = get_requested_neighbor_lists(exported)
     system = get_system_with_neighbor_lists(system, requested_neighbor_lists)

src/metatrain/experimental/pet/tests/test_functionality.py

@@ -76,6 +76,7 @@ def test_prediction():
         types=torch.tensor([6, 6]),
         positions=torch.tensor([[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]]),
         cell=torch.zeros(3, 3),
+        pbc=torch.tensor([False, False, False]),
     )
     requested_neighbor_lists = get_requested_neighbor_lists(model)
     system = get_system_with_neighbor_lists(system, requested_neighbor_lists)
@@ -126,6 +127,7 @@ def test_per_atom_predictions_functionality():
         types=torch.tensor([6, 6]),
         positions=torch.tensor([[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]]),
         cell=torch.zeros(3, 3),
+        pbc=torch.tensor([False, False, False]),
     )
     requested_neighbor_lists = get_requested_neighbor_lists(model)
     system = get_system_with_neighbor_lists(system, requested_neighbor_lists)
@@ -177,6 +179,7 @@ def test_selected_atoms_functionality():
         types=torch.tensor([6, 6]),
         positions=torch.tensor([[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]]),
         cell=torch.zeros(3, 3),
+        pbc=torch.tensor([False, False, False]),
     )
     requested_neighbor_lists = get_requested_neighbor_lists(model)
     system = get_system_with_neighbor_lists(system, requested_neighbor_lists)

src/metatrain/experimental/pet/tests/test_pet_compatibility.py

@@ -59,7 +59,7 @@ def test_batch_dicts_compatibility(cutoff):
     structure = ase.io.read(DATASET_PATH)
     atomic_types = sorted(set(structure.numbers))
     system = systems_to_torch(structure)
-    options = NeighborListOptions(cutoff=cutoff, full_list=True)
+    options = NeighborListOptions(cutoff=cutoff, full_list=True, strict=True)
     system = get_system_with_neighbor_lists(system, [options])
 
     ARCHITECTURAL_HYPERS = Hypers(DEFAULT_HYPERS["model"])
@@ -121,7 +121,7 @@ def test_predictions_compatibility(cutoff):
     model.set_trained_model(raw_pet)
 
     system = systems_to_torch(structure)
-    options = NeighborListOptions(cutoff=cutoff, full_list=True)
+    options = NeighborListOptions(cutoff=cutoff, full_list=True, strict=True)
     system = get_system_with_neighbor_lists(system, [options])
 
     evaluation_options = ModelEvaluationOptions(

src/metatrain/experimental/soap_bpnn/tests/test_exported.py

@@ -33,6 +33,7 @@ def test_to(device, dtype):
         types=torch.tensor([6, 6]),
         positions=torch.tensor([[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]]),
         cell=torch.zeros(3, 3),
+        pbc=torch.tensor([False, False, False]),
     )
     requested_neighbor_lists = get_requested_neighbor_lists(exported)
     system = get_system_with_neighbor_lists(system, requested_neighbor_lists)

src/metatrain/experimental/soap_bpnn/tests/test_functionality.py

@@ -28,6 +28,7 @@ def test_prediction_subset_elements():
         types=torch.tensor([6, 6]),
         positions=torch.tensor([[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]]),
         cell=torch.zeros(3, 3),
+        pbc=torch.tensor([False, False, False]),
     )
     model(
         [system],
@@ -56,6 +57,7 @@ def test_prediction_subset_atoms():
             [[0.0, 0.0, 0.0], [0.0, 0.0, 1.0], [0.0, 0.0, 2.0]],
         ),
         cell=torch.zeros(3, 3),
+        pbc=torch.tensor([False, False, False]),
     )
 
     energy_monomer = model(
@@ -76,6 +78,7 @@ def test_prediction_subset_atoms():
             ],
         ),
         cell=torch.zeros(3, 3),
+        pbc=torch.tensor([False, False, False]),
     )
 
     selection_labels = metatensor.torch.Labels(
@@ -119,6 +122,7 @@ def test_output_last_layer_features():
             [[0.0, 0.0, 0.0], [0.0, 0.0, 1.0], [0.0, 0.0, 2.0], [0.0, 0.0, 3.0]],
         ),
         cell=torch.zeros(3, 3),
+        pbc=torch.tensor([False, False, False]),
     )
 
     # last-layer features per atom:
@@ -190,6 +194,7 @@ def test_output_per_atom():
             [[0.0, 0.0, 0.0], [0.0, 0.0, 1.0], [0.0, 0.0, 2.0], [0.0, 0.0, 3.0]],
         ),
         cell=torch.zeros(3, 3),
+        pbc=torch.tensor([False, False, False]),
     )
 
     outputs = model(

src/metatrain/experimental/soap_bpnn/tests/test_torchscript.py

@@ -26,6 +26,7 @@ def test_torchscript():
             [[0.0, 0.0, 0.0], [0.0, 0.0, 1.0], [0.0, 0.0, 2.0], [0.0, 0.0, 3.0]]
         ),
         cell=torch.zeros(3, 3),
+        pbc=torch.tensor([False, False, False]),
     )
     model(
         [system],
@@ -52,6 +53,7 @@ def test_torchscript_with_identity():
             [[0.0, 0.0, 0.0], [0.0, 0.0, 1.0], [0.0, 0.0, 2.0], [0.0, 0.0, 3.0]]
         ),
         cell=torch.zeros(3, 3),
+        pbc=torch.tensor([False, False, False]),
     )
     model(
         [system],

src/metatrain/utils/additive/zbl.py

@@ -251,6 +251,7 @@ def requested_neighbor_lists(self) -> List[NeighborListOptions]:
             NeighborListOptions(
                 cutoff=self.cutoff_radius,
                 full_list=True,
+                strict=True,
             )
         ]
 

src/metatrain/utils/evaluate_model.py

@@ -278,20 +278,23 @@ def _prepare_system(
             positions=system.positions @ strain,
             cell=system.cell @ strain,
             types=system.types,
+            pbc=system.pbc,
         )
     else:
         if positions_grad:
             new_system = System(
                 positions=system.positions.detach().clone().requires_grad_(True),
                 cell=system.cell,
                 types=system.types,
+                pbc=system.pbc,
             )
             strain = None
         else:
             new_system = System(
                 positions=system.positions,
                 cell=system.cell,
                 types=system.types,
+                pbc=system.pbc,
             )
             strain = None
 

src/metatrain/utils/metrics.py

@@ -8,7 +8,7 @@ class RMSEAccumulator:
     """Accumulates the RMSE between predictions and targets for an arbitrary
     number of keys, each corresponding to one target."""
 
-    def __init__(self):
+    def __init__(self) -> None:
         """Initialize the accumulator."""
         self.information: Dict[str, Tuple[float, int]] = {}
 
@@ -91,7 +91,7 @@ class MAEAccumulator:
     """Accumulates the MAE between predictions and targets for an arbitrary
     number of keys, each corresponding to one target."""
 
-    def __init__(self):
+    def __init__(self) -> None:
         """Initialize the accumulator."""
         self.information: Dict[str, Tuple[float, int]] = {}
 

tests/utils/data/test_system_to_ase.py

@@ -11,6 +11,7 @@ def test_system_to_ase():
         positions=torch.tensor([[0.0, 0.0, 0.0], [1.0, 1.0, 1.0]]),
         types=torch.tensor([1, 8]),
         cell=torch.tensor([[10.0, 0.0, 0.0], [0.0, 10.0, 0.0], [0.0, 0.0, 10.0]]),
+        pbc=torch.tensor([True, True, True]),
     )
 
     # Convert the system to an ASE atoms object

tests/utils/data/test_target_writers.py

@@ -20,6 +20,7 @@ def systems_capabilities_predictions(
             types=torch.tensor([1, 1]),
             positions=torch.tensor([[0, 0, 0], [0, 0, 0.74]]),
             cell=cell,
+            pbc=torch.logical_not(torch.all(cell == 0, dim=1)),
         ),
     ]
 

tests/utils/test_additive.py

@@ -33,13 +33,15 @@ def test_composition_model_train():
             positions=torch.tensor([[0.0, 0.0, 0.0]], dtype=torch.float64),
             types=torch.tensor([8]),
             cell=torch.eye(3, dtype=torch.float64),
+            pbc=torch.tensor([True, True, True]),
         ),
         System(
             positions=torch.tensor(
                 [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], dtype=torch.float64
             ),
             types=torch.tensor([1, 1, 8]),
             cell=torch.eye(3, dtype=torch.float64),
+            pbc=torch.tensor([True, True, True]),
         ),
         System(
             positions=torch.tensor(
@@ -55,6 +57,7 @@ def test_composition_model_train():
             ),
             types=torch.tensor([1, 1, 8, 1, 1, 8]),
             cell=torch.eye(3, dtype=torch.float64),
+            pbc=torch.tensor([True, True, True]),
         ),
     ]
     energies = [1.0, 5.0, 10.0]
@@ -200,6 +203,7 @@ def test_composition_model_torchscript(tmpdir):
         positions=torch.tensor([[0.0, 0.0, 0.0]], dtype=torch.float64),
         types=torch.tensor([8]),
         cell=torch.eye(3, dtype=torch.float64),
+        pbc=torch.tensor([True, True, True]),
     )
 
     composition_model = CompositionModel(
@@ -289,13 +293,15 @@ def test_composition_model_missing_types():
             positions=torch.tensor([[0.0, 0.0, 0.0]], dtype=torch.float64),
             types=torch.tensor([8]),
             cell=torch.eye(3, dtype=torch.float64),
+            pbc=torch.tensor([True, True, True]),
         ),
         System(
             positions=torch.tensor(
                 [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], dtype=torch.float64
             ),
             types=torch.tensor([1, 1, 8]),
             cell=torch.eye(3, dtype=torch.float64),
+            pbc=torch.tensor([True, True, True]),
         ),
         System(
             positions=torch.tensor(
@@ -311,6 +317,7 @@ def test_composition_model_missing_types():
             ),
             types=torch.tensor([1, 1, 8, 1, 1, 8]),
             cell=torch.eye(3, dtype=torch.float64),
+            pbc=torch.tensor([True, True, True]),
         ),
     ]
     energies = [1.0, 5.0, 10.0]

tests/utils/test_neighbor_list.py

@@ -14,9 +14,9 @@ def test_attach_neighbor_lists():
     systems = read_systems_ase(filename)
 
     requested_neighbor_lists = [
-        NeighborListOptions(cutoff=4.0, full_list=True),
-        NeighborListOptions(cutoff=5.0, full_list=False),
-        NeighborListOptions(cutoff=6.0, full_list=True),
+        NeighborListOptions(cutoff=4.0, full_list=True, strict=True),
+        NeighborListOptions(cutoff=5.0, full_list=False, strict=True),
+        NeighborListOptions(cutoff=6.0, full_list=True, strict=True),
     ]
 
     new_system = get_system_with_neighbor_lists(systems[0], requested_neighbor_lists)
@@ -25,7 +25,7 @@ def test_attach_neighbor_lists():
     assert requested_neighbor_lists[1] in new_system.known_neighbor_lists()
     assert requested_neighbor_lists[2] in new_system.known_neighbor_lists()
 
-    extraneous_nl = NeighborListOptions(cutoff=5.0, full_list=True)
+    extraneous_nl = NeighborListOptions(cutoff=5.0, full_list=True, strict=True)
     assert extraneous_nl not in new_system.known_neighbor_lists()
 
     for nl_options in new_system.known_neighbor_lists():

tests/utils/test_output_gradient.py

@@ -32,6 +32,7 @@ def test_forces(is_training):
             positions=system.positions.requires_grad_(True),
             cell=system.cell,
             types=system.types,
+            pbc=system.pbc,
         )
         for system in systems
     ]
@@ -50,6 +51,7 @@ def test_forces(is_training):
             positions=system.positions.requires_grad_(True),
             cell=system.cell,
             types=system.types,
+            pbc=system.pbc,
         )
         for system in systems
     ]
@@ -96,6 +98,7 @@ def test_virial(is_training):
             positions=system.positions @ strain,
             cell=system.cell @ strain,
             types=system.types,
+            pbc=system.pbc,
         )
         for system, strain in zip(systems, strains)
     ]
@@ -121,6 +124,7 @@ def test_virial(is_training):
             positions=system.positions @ strain,
             cell=system.cell @ strain,
             types=system.types,
+            pbc=system.pbc,
         )
         for system, strain in zip(systems, strains)
     ]
@@ -170,6 +174,7 @@ def test_both(is_training):
             positions=system.positions @ strain,
             cell=system.cell @ strain,
             types=system.types,
+            pbc=system.pbc,
         )
         for system, strain in zip(systems, strains)
     ]
@@ -193,6 +198,7 @@ def test_both(is_training):
             positions=system.positions @ strain,
             cell=system.cell @ strain,
             types=system.types,
+            pbc=system.pbc,
         )
         for system, strain in zip(systems, strains)
     ]