From 172af0b6f222e44ab856c28452f4ef6ab85c3c97 Mon Sep 17 00:00:00 2001
From: frostedoyster <bigi.f@libero.it>
Date: Tue, 29 Oct 2024 19:05:39 +0100
Subject: [PATCH 1/6] Upgrade to `metatensor-torch` 0.6.0

---
 pyproject.toml                                             | 4 ++--
 .../experimental/alchemical_model/tests/test_exported.py   | 1 +
 .../alchemical_model/tests/test_functionality.py           | 1 +
 src/metatrain/experimental/pet/tests/test_exported.py      | 1 +
 src/metatrain/experimental/pet/tests/test_functionality.py | 3 +++
 .../experimental/soap_bpnn/tests/test_exported.py          | 1 +
 .../experimental/soap_bpnn/tests/test_functionality.py     | 5 +++++
 .../experimental/soap_bpnn/tests/test_torchscript.py       | 2 ++
 src/metatrain/utils/evaluate_model.py                      | 3 +++
 tests/utils/data/test_system_to_ase.py                     | 1 +
 tests/utils/data/test_target_writers.py                    | 1 +
 tests/utils/test_additive.py                               | 7 +++++++
 tests/utils/test_output_gradient.py                        | 6 ++++++
 tests/utils/test_per_atom.py                               | 3 +++
 tests/utils/test_transfer.py                               | 2 ++
 15 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 7e3f2d703..d78526820 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -11,8 +11,8 @@ 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-operations==0.3.0",
+    "metatensor-torch==0.6.0",
     "jsonschema",
     "omegaconf",
     "python-hostlist",
diff --git a/src/metatrain/experimental/alchemical_model/tests/test_exported.py b/src/metatrain/experimental/alchemical_model/tests/test_exported.py
index 891983693..8b06c2d74 100644
--- a/src/metatrain/experimental/alchemical_model/tests/test_exported.py
+++ b/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)
diff --git a/src/metatrain/experimental/alchemical_model/tests/test_functionality.py b/src/metatrain/experimental/alchemical_model/tests/test_functionality.py
index 7ee3331af..9f998f925 100644
--- a/src/metatrain/experimental/alchemical_model/tests/test_functionality.py
+++ b/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)
diff --git a/src/metatrain/experimental/pet/tests/test_exported.py b/src/metatrain/experimental/pet/tests/test_exported.py
index f67a15e4c..5ed9ecf6e 100644
--- a/src/metatrain/experimental/pet/tests/test_exported.py
+++ b/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)
diff --git a/src/metatrain/experimental/pet/tests/test_functionality.py b/src/metatrain/experimental/pet/tests/test_functionality.py
index ddf527603..35c44be0b 100644
--- a/src/metatrain/experimental/pet/tests/test_functionality.py
+++ b/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),
+        pbcs=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),
+        pbcs=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),
+        pbcs=torch.tensor([False, False, False]),
     )
     requested_neighbor_lists = get_requested_neighbor_lists(model)
     system = get_system_with_neighbor_lists(system, requested_neighbor_lists)
diff --git a/src/metatrain/experimental/soap_bpnn/tests/test_exported.py b/src/metatrain/experimental/soap_bpnn/tests/test_exported.py
index 63242161e..ff926aa60 100644
--- a/src/metatrain/experimental/soap_bpnn/tests/test_exported.py
+++ b/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),
+        pbcs=torch.tensor([False, False, False]),
     )
     requested_neighbor_lists = get_requested_neighbor_lists(exported)
     system = get_system_with_neighbor_lists(system, requested_neighbor_lists)
diff --git a/src/metatrain/experimental/soap_bpnn/tests/test_functionality.py b/src/metatrain/experimental/soap_bpnn/tests/test_functionality.py
index 25dd250b6..12bbe4c7a 100644
--- a/src/metatrain/experimental/soap_bpnn/tests/test_functionality.py
+++ b/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),
+        pbcs=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),
+        pbcs=torch.tensor([False, False, False]),
     )
 
     energy_monomer = model(
@@ -76,6 +78,7 @@ def test_prediction_subset_atoms():
             ],
         ),
         cell=torch.zeros(3, 3),
+        pbcs=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),
+        pbcs=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),
+        pbcs=torch.tensor([False, False, False]),
     )
 
     outputs = model(
diff --git a/src/metatrain/experimental/soap_bpnn/tests/test_torchscript.py b/src/metatrain/experimental/soap_bpnn/tests/test_torchscript.py
index 4d6b89898..9e41046a3 100644
--- a/src/metatrain/experimental/soap_bpnn/tests/test_torchscript.py
+++ b/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),
+        pbcs=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),
+        pbcs=torch.tensor([False, False, False]),
     )
     model(
         [system],
diff --git a/src/metatrain/utils/evaluate_model.py b/src/metatrain/utils/evaluate_model.py
index 48447b879..5b4500a21 100644
--- a/src/metatrain/utils/evaluate_model.py
+++ b/src/metatrain/utils/evaluate_model.py
@@ -278,6 +278,7 @@ def _prepare_system(
             positions=system.positions @ strain,
             cell=system.cell @ strain,
             types=system.types,
+            pbc=system.pbc,
         )
     else:
         if positions_grad:
@@ -285,6 +286,7 @@ def _prepare_system(
                 positions=system.positions.detach().clone().requires_grad_(True),
                 cell=system.cell,
                 types=system.types,
+                pbc=system.pbc,
             )
             strain = None
         else:
@@ -292,6 +294,7 @@ def _prepare_system(
                 positions=system.positions,
                 cell=system.cell,
                 types=system.types,
+                pbc=system.pbc,
             )
             strain = None
 
diff --git a/tests/utils/data/test_system_to_ase.py b/tests/utils/data/test_system_to_ase.py
index e6062b52d..3d66f7e7c 100644
--- a/tests/utils/data/test_system_to_ase.py
+++ b/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
diff --git a/tests/utils/data/test_target_writers.py b/tests/utils/data/test_target_writers.py
index 190d37fbe..525ccd446 100644
--- a/tests/utils/data/test_target_writers.py
+++ b/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)),
         ),
     ]
 
diff --git a/tests/utils/test_additive.py b/tests/utils/test_additive.py
index fd2179e5e..3daf3c980 100644
--- a/tests/utils/test_additive.py
+++ b/tests/utils/test_additive.py
@@ -33,6 +33,7 @@ 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(
@@ -40,6 +41,7 @@ def test_composition_model_train():
             ),
             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,6 +293,7 @@ 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(
@@ -296,6 +301,7 @@ def test_composition_model_missing_types():
             ),
             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]
diff --git a/tests/utils/test_output_gradient.py b/tests/utils/test_output_gradient.py
index 7bad9279e..513d43930 100644
--- a/tests/utils/test_output_gradient.py
+++ b/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)
     ]
diff --git a/tests/utils/test_per_atom.py b/tests/utils/test_per_atom.py
index 3f4f34945..0f84180e7 100644
--- a/tests/utils/test_per_atom.py
+++ b/tests/utils/test_per_atom.py
@@ -13,16 +13,19 @@ def test_average_by_num_atoms():
             positions=torch.tensor([[0.0, 0.0, 0.0]]),
             cell=torch.eye(3),
             types=torch.tensor([0]),
+            pbc=torch.tensor([True, True, True]),
         ),
         System(
             positions=torch.tensor([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]),
             cell=torch.eye(3),
             types=torch.tensor([0, 0]),
+            pbc=torch.tensor([True, True, True]),
         ),
         System(
             positions=torch.tensor([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]),
             cell=torch.eye(3),
             types=torch.tensor([0, 0, 0]),
+            pbc=torch.tensor([True, True, True]),
         ),
     ]
 
diff --git a/tests/utils/test_transfer.py b/tests/utils/test_transfer.py
index b669c0c29..283842779 100644
--- a/tests/utils/test_transfer.py
+++ b/tests/utils/test_transfer.py
@@ -14,6 +14,7 @@ def test_systems_and_targets_to_dtype():
         positions=torch.tensor([[1.0, 1.0, 1.0]]),
         cell=torch.tensor([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]]),
         types=torch.tensor([1]),
+        pbc=torch.tensor([True, True, True]),
     )
     targets = TensorMap(
         keys=Labels.single(),
@@ -39,6 +40,7 @@ def test_systems_and_targets_to_dtype_and_device():
         positions=torch.tensor([[1.0, 1.0, 1.0]]),
         cell=torch.tensor([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]]),
         types=torch.tensor([1]),
+        pbc=torch.tensor([True, True, True]),
     )
     targets = TensorMap(
         keys=Labels.single(),

From 0f3bfd7b19d98ea3a9be62d13d10f6ba55a18e6b Mon Sep 17 00:00:00 2001
From: frostedoyster <bigi.f@libero.it>
Date: Wed, 30 Oct 2024 15:35:41 +0100
Subject: [PATCH 2/6] Upgrade `metatensor-learn`

---
 pyproject.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pyproject.toml b/pyproject.toml
index d78526820..6e09fc511 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -10,7 +10,7 @@ authors = [{name = "metatrain developers"}]
 
 dependencies = [
     "ase < 3.23.0",
-    "metatensor-learn==0.2.3",
+    "metatensor-learn==0.3.0",
     "metatensor-operations==0.3.0",
     "metatensor-torch==0.6.0",
     "jsonschema",

From 0e8e2f71bd0dcf3c1bad534e0ae74270c84c5cb3 Mon Sep 17 00:00:00 2001
From: frostedoyster <bigi.f@libero.it>
Date: Wed, 30 Oct 2024 17:30:22 +0100
Subject: [PATCH 3/6] Update strict NL

---
 pyproject.toml                                            | 4 ++--
 src/metatrain/experimental/alchemical_model/model.py      | 1 +
 .../tests/test_torch_alchemical_compatibility.py          | 1 +
 src/metatrain/experimental/pet/model.py                   | 1 +
 .../experimental/pet/tests/test_pet_compatibility.py      | 4 ++--
 src/metatrain/utils/additive/zbl.py                       | 1 +
 src/metatrain/utils/metrics.py                            | 4 ++--
 tests/utils/test_neighbor_list.py                         | 8 ++++----
 8 files changed, 14 insertions(+), 10 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 6e09fc511..f75f0c018 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -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@b70b19e#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@b70b19e#subdirectory=python/rascaline-torch",
     "skmatter",
     "metatensor-learn",
     "scipy",
diff --git a/src/metatrain/experimental/alchemical_model/model.py b/src/metatrain/experimental/alchemical_model/model.py
index 8ffbd8dc7..4e985dc75 100644
--- a/src/metatrain/experimental/alchemical_model/model.py
+++ b/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,
             )
         ]
 
diff --git a/src/metatrain/experimental/alchemical_model/tests/test_torch_alchemical_compatibility.py b/src/metatrain/experimental/alchemical_model/tests/test_torch_alchemical_compatibility.py
index 03b7ef1df..7b560e786 100644
--- a/src/metatrain/experimental/alchemical_model/tests/test_torch_alchemical_compatibility.py
+++ b/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]
 
diff --git a/src/metatrain/experimental/pet/model.py b/src/metatrain/experimental/pet/model.py
index bf567dee1..512dc319b 100644
--- a/src/metatrain/experimental/pet/model.py
+++ b/src/metatrain/experimental/pet/model.py
@@ -73,6 +73,7 @@ def requested_neighbor_lists(
             NeighborListOptions(
                 cutoff=self.cutoff,
                 full_list=True,
+                strict=True,
             )
         ]
 
diff --git a/src/metatrain/experimental/pet/tests/test_pet_compatibility.py b/src/metatrain/experimental/pet/tests/test_pet_compatibility.py
index 04ecae1bc..4a2bdd71e 100644
--- a/src/metatrain/experimental/pet/tests/test_pet_compatibility.py
+++ b/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(
diff --git a/src/metatrain/utils/additive/zbl.py b/src/metatrain/utils/additive/zbl.py
index edf2ec7b2..cec2ee96e 100644
--- a/src/metatrain/utils/additive/zbl.py
+++ b/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,
             )
         ]
 
diff --git a/src/metatrain/utils/metrics.py b/src/metatrain/utils/metrics.py
index 357b567d0..29ace5104 100644
--- a/src/metatrain/utils/metrics.py
+++ b/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]] = {}
 
diff --git a/tests/utils/test_neighbor_list.py b/tests/utils/test_neighbor_list.py
index 9c35e05d5..768114d6c 100644
--- a/tests/utils/test_neighbor_list.py
+++ b/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():

From d632dc2f68a6b92b96e55874de2b2cbe7a8892d4 Mon Sep 17 00:00:00 2001
From: frostedoyster <bigi.f@libero.it>
Date: Wed, 30 Oct 2024 17:39:08 +0100
Subject: [PATCH 4/6] Fix PBCs

---
 .../experimental/pet/tests/test_functionality.py       |  6 +++---
 .../experimental/soap_bpnn/tests/test_exported.py      |  2 +-
 .../experimental/soap_bpnn/tests/test_functionality.py | 10 +++++-----
 .../experimental/soap_bpnn/tests/test_torchscript.py   |  4 ++--
 4 files changed, 11 insertions(+), 11 deletions(-)

diff --git a/src/metatrain/experimental/pet/tests/test_functionality.py b/src/metatrain/experimental/pet/tests/test_functionality.py
index 35c44be0b..b0f5771ea 100644
--- a/src/metatrain/experimental/pet/tests/test_functionality.py
+++ b/src/metatrain/experimental/pet/tests/test_functionality.py
@@ -76,7 +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),
-        pbcs=torch.tensor([False, False, False]),
+        pbc=torch.tensor([False, False, False]),
     )
     requested_neighbor_lists = get_requested_neighbor_lists(model)
     system = get_system_with_neighbor_lists(system, requested_neighbor_lists)
@@ -127,7 +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),
-        pbcs=torch.tensor([False, False, False]),
+        pbc=torch.tensor([False, False, False]),
     )
     requested_neighbor_lists = get_requested_neighbor_lists(model)
     system = get_system_with_neighbor_lists(system, requested_neighbor_lists)
@@ -179,7 +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),
-        pbcs=torch.tensor([False, False, False]),
+        pbc=torch.tensor([False, False, False]),
     )
     requested_neighbor_lists = get_requested_neighbor_lists(model)
     system = get_system_with_neighbor_lists(system, requested_neighbor_lists)
diff --git a/src/metatrain/experimental/soap_bpnn/tests/test_exported.py b/src/metatrain/experimental/soap_bpnn/tests/test_exported.py
index ff926aa60..349e89cfd 100644
--- a/src/metatrain/experimental/soap_bpnn/tests/test_exported.py
+++ b/src/metatrain/experimental/soap_bpnn/tests/test_exported.py
@@ -33,7 +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),
-        pbcs=torch.tensor([False, False, False]),
+        pbc=torch.tensor([False, False, False]),
     )
     requested_neighbor_lists = get_requested_neighbor_lists(exported)
     system = get_system_with_neighbor_lists(system, requested_neighbor_lists)
diff --git a/src/metatrain/experimental/soap_bpnn/tests/test_functionality.py b/src/metatrain/experimental/soap_bpnn/tests/test_functionality.py
index 12bbe4c7a..b1fb4f1d3 100644
--- a/src/metatrain/experimental/soap_bpnn/tests/test_functionality.py
+++ b/src/metatrain/experimental/soap_bpnn/tests/test_functionality.py
@@ -28,7 +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),
-        pbcs=torch.tensor([False, False, False]),
+        pbc=torch.tensor([False, False, False]),
     )
     model(
         [system],
@@ -57,7 +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),
-        pbcs=torch.tensor([False, False, False]),
+        pbc=torch.tensor([False, False, False]),
     )
 
     energy_monomer = model(
@@ -78,7 +78,7 @@ def test_prediction_subset_atoms():
             ],
         ),
         cell=torch.zeros(3, 3),
-        pbcs=torch.tensor([False, False, False]),
+        pbc=torch.tensor([False, False, False]),
     )
 
     selection_labels = metatensor.torch.Labels(
@@ -122,7 +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),
-        pbcs=torch.tensor([False, False, False]),
+        pbc=torch.tensor([False, False, False]),
     )
 
     # last-layer features per atom:
@@ -194,7 +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),
-        pbcs=torch.tensor([False, False, False]),
+        pbc=torch.tensor([False, False, False]),
     )
 
     outputs = model(
diff --git a/src/metatrain/experimental/soap_bpnn/tests/test_torchscript.py b/src/metatrain/experimental/soap_bpnn/tests/test_torchscript.py
index 9e41046a3..6be11f582 100644
--- a/src/metatrain/experimental/soap_bpnn/tests/test_torchscript.py
+++ b/src/metatrain/experimental/soap_bpnn/tests/test_torchscript.py
@@ -26,7 +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),
-        pbcs=torch.tensor([False, False, False]),
+        pbc=torch.tensor([False, False, False]),
     )
     model(
         [system],
@@ -53,7 +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),
-        pbcs=torch.tensor([False, False, False]),
+        pbc=torch.tensor([False, False, False]),
     )
     model(
         [system],

From cf62ee7b32f7ca5ee2f8cf1d5f49bf77386ec3f5 Mon Sep 17 00:00:00 2001
From: frostedoyster <bigi.f@libero.it>
Date: Fri, 1 Nov 2024 12:33:06 +0100
Subject: [PATCH 5/6] Upgrade rascaline-torch

---
 pyproject.toml | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index f75f0c018..bec22e930 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -59,7 +59,7 @@ build-backend = "setuptools.build_meta"
 
 [project.optional-dependencies]
 soap-bpnn = [
-    "rascaline-torch @ git+https://github.com/luthaf/rascaline@b70b19e#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@b70b19e#subdirectory=python/rascaline-torch",
+    "rascaline-torch @ git+https://github.com/luthaf/rascaline@5326b6e#subdirectory=python/rascaline-torch",
     "skmatter",
     "metatensor-learn",
     "scipy",

From 0ebcc9aba44ff12aa06b90b35725e7f8ae018796 Mon Sep 17 00:00:00 2001
From: frostedoyster <bigi.f@libero.it>
Date: Sat, 2 Nov 2024 01:34:31 +0100
Subject: [PATCH 6/6] Fix `slice` argument name

---
 src/metatrain/experimental/pet/model.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/src/metatrain/experimental/pet/model.py b/src/metatrain/experimental/pet/model.py
index 512dc319b..d0d21d73f 100644
--- a/src/metatrain/experimental/pet/model.py
+++ b/src/metatrain/experimental/pet/model.py
@@ -112,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")