dependency bump: metatensor 0.9, metatomic 0.1.14, featomic 0.7.4

.github/workflows/architecture-tests.yml

@@ -22,7 +22,7 @@ jobs:
           - pet
           - phace
           - soap-bpnn
-          - dpa3       
+          - dpa3
 
     runs-on: ubuntu-24.04
 

.github/workflows/tests.yml

@@ -54,7 +54,7 @@ jobs:
       env:
         # Use the CPU only version of torch when building/running the code
         PIP_EXTRA_INDEX_URL: https://download.pytorch.org/whl/cpu
-        # Temporary workaround: avoid the deprecated hf_xet.download_files() path 
+        # Temporary workaround: avoid the deprecated hf_xet.download_files() path
         # until huggingface_hub switches to the new Xet API.
         HF_HUB_DISABLE_XET: 1
         HUGGINGFACE_TOKEN_METATRAIN: ${{ secrets.HUGGINGFACE_TOKEN }}

pyproject.toml

@@ -12,18 +12,18 @@ dependencies = [
     "ase",
     "huggingface_hub",
     "numpy",
-    "metatensor-learn >=0.4.0,<0.5",
-    "metatensor-operations >=0.4.0,<0.5",
-    "metatensor-torch >=0.8.2,<0.9",
-    "metatomic-torch >=0.1.8,<0.2",
-    "metatomic-ase >=0.1.0,<0.2",
+    "metatensor-learn >=0.5.0,<0.6",
+    "metatensor-operations >=0.5.0,<0.6",
+    "metatensor-torch >=0.9.0,<0.10",
+    "metatomic-torch >=0.1.14,<0.2",
+    "metatomic-ase >=0.1.1,<0.2",
     "jsonschema",
     "pydantic >= 2.12", # Because of passing arguments to validator functions
     "typing_extensions", # For python < 3.12 typing.TypedDict has limited functionality
     "omegaconf >= 2.3.0",
     "python-hostlist",
     "tqdm",
-    "vesin >= 0.5.2" # Because of https://github.com/Luthaf/vesin/pull/119,
+    "vesin >= 0.5.2", # Because of https://github.com/Luthaf/vesin/pull/119,
 ]
 
 keywords = ["machine learning", "molecular modeling"]
@@ -66,11 +66,12 @@ build-backend = "setuptools.build_meta"
 [project.optional-dependencies]
 soap-bpnn = [
     "torch-spex>=0.1,<0.2",
+    'sphericart-torch == 1.0.8 ; platform_system == "Windows"', # for some reason sphericart-torch 1.0.9 can't load on windows
     "wigners",
 ]
 pet = []
 gap = [
-    "featomic-torch >=0.7,<0.8",
+    "featomic-torch >=0.7.4,<0.8",
     "skmatter",
     "scipy",
 ]
@@ -180,6 +181,8 @@ filterwarnings = [
     "ignore:`torch.jit.load` is deprecated. Please switch to `torch.export`:DeprecationWarning",
     "ignore:`torch.jit.trace_method` is deprecated. Please switch to `torch.compile` or `torch.export`:DeprecationWarning",
     "ignore:`torch.jit.trace` is deprecated. Please switch to `torch.compile` or `torch.export`:DeprecationWarning",
+    "ignore:`torch.jit.interface` is deprecated. Please use `torch.compile` instead.:DeprecationWarning",
+    "ignore:`torch.jit.set_fusion_strategy` is deprecated. Please use `torch.compile` instead.:DeprecationWarning",
     # PyTorch does not want these, but mypy requires them
     "ignore:The TorchScript type system doesn't support instance-level annotations on empty non-base types",
     # Sphericart double backward warning
@@ -196,7 +199,19 @@ filterwarnings = [
     # MACE warning with newer versions of pytorch (because they use e3nn==0.4.4)
     "ignore:Environment variable TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD detected, since the`weights_only` argument was not explicitly passed:UserWarning",
     # vesin.metatomic deprecation
-    "ignore:`compute_requested_neighbors_from_options` is deprecated and will be removed in a future version:UserWarning"
+    "ignore:`compute_requested_neighbors_from_options` is deprecated and will be removed in a future version:UserWarning",
+    # metatomic 0.1.12 deprecations
+    "ignore:the 'features' output name is deprecated:UserWarning",
+    "ignore:the 'features' quantity is deprecated:UserWarning",
+    "ignore:the 'features' quantity is deprecated:DeprecationWarning",
+    "ignore:the 'non_conservative_forces' output name is deprecated:UserWarning",
+    "ignore:the 'momenta' quantity is deprecated:DeprecationWarning",
+    "ignore:the 'positions' quantity is deprecated:DeprecationWarning",
+    "ignore:the 'masses' quantity is deprecated:DeprecationWarning",
+    "ignore:the 'velocities' quantity is deprecated:DeprecationWarning",
+    "ignore:the 'charges' quantity is deprecated:DeprecationWarning",
+    "ignore:`per_atom` is deprecated:DeprecationWarning",
+    "ignore:ModelOutput.quantity is deprecated:DeprecationWarning"
 ]
 addopts = ["-p", "mtt_plugin"]
 pythonpath = "src/metatrain/utils/testing"

src/metatrain/experimental/flashmd/model.py

@@ -588,7 +588,7 @@ def forward(
                                     )
                                 )
                             else:
-                                output_blocks.append(b)
+                                output_blocks.append(b.copy(deep=False))
                         return_dict[name] = TensorMap(
                             return_dict[name].keys, output_blocks
                         )

src/metatrain/experimental/flashmd_symplectic/model.py

@@ -643,7 +643,7 @@ def forward(
                                     )
                                 )
                             else:
-                                output_blocks.append(b)
+                                output_blocks.append(b.copy(deep=False))
                         return_dict[name] = TensorMap(
                             return_dict[name].keys, output_blocks
                         )

src/metatrain/experimental/mace/model.py

@@ -528,7 +528,7 @@ def add_additive_contributions(
                             )
                         )
                     else:
-                        output_blocks.append(b)
+                        output_blocks.append(b.copy(deep=False))
                 values[name] = TensorMap(values[name].keys, output_blocks)
 
     def supported_outputs(self) -> Dict[str, ModelOutput]:

src/metatrain/experimental/phace/model.py

@@ -513,7 +513,7 @@ def forward(
                                 )
                             )
                         else:
-                            output_blocks.append(b)
+                            output_blocks.append(b.copy(deep=False))
                     return_dict[name] = TensorMap(return_dict[name].keys, output_blocks)
 
             # For atomic basis targets, sparsify to create blocks with "atom_type"

src/metatrain/gap/model.py

@@ -243,9 +243,10 @@ def forward(
         for idx_key in range(len(self._species_labels)):
             key = self._species_labels.entry(idx_key)
             if soap_features.keys.position(key) is not None:
-                new_blocks.append(soap_features.block(key))
+                block = soap_features.block(key)
+                new_blocks.append(block.copy(deep=False))
             else:
-                new_blocks.append(dummyblock)
+                new_blocks.append(dummyblock.copy(deep=False))
         soap_features = TensorMap(keys=self._species_labels, blocks=new_blocks)
         soap_features = soap_features.keys_to_samples("center_type")
         # here, we move to properties to use metatensor operations to aggregate
@@ -254,7 +255,10 @@ def forward(
         soap_features = soap_features.keys_to_properties(
             ["neighbor_1_type", "neighbor_2_type"]
         )
-        soap_features = TensorMap(self._keys, soap_features.blocks())
+        soap_features = TensorMap(
+            self._keys,
+            [block.copy(deep=False) for block in soap_features.blocks()],
+        )
         output_key = list(outputs.keys())[0]
         energies = self._subset_of_regressors_torch(soap_features)
         return_dict = {output_key: energies}

src/metatrain/soap_bpnn/model.py

@@ -865,7 +865,7 @@ def forward(
                                 )
                             )
                         else:
-                            output_blocks.append(b)
+                            output_blocks.append(b.copy(deep=False))
                     return_dict[name] = TensorMap(return_dict[name].keys, output_blocks)
 
             # For atomic basis targets, sparsify to create blocks with "atom_type"

src/metatrain/utils/additive/remove.py

@@ -3,7 +3,7 @@
 
 import metatensor.torch as mts
 import torch
-from metatensor.torch import Labels, TensorBlock, TensorMap
+from metatensor.torch import Labels, TensorMap
 from metatensor.torch.operations._add import _add_block_block
 from metatensor.torch.operations._multiply import _multiply_block_constant
 from metatomic.torch import System
@@ -116,14 +116,7 @@ def remove_additive(
                 )
 
             else:
-                new_target_blocks.append(
-                    TensorBlock(
-                        values=block.values,
-                        samples=block.samples,
-                        components=block.components,
-                        properties=block.properties,
-                    )
-                )
+                new_target_blocks.append(block.copy(deep=False))
 
         targets[target_key] = TensorMap(
             keys=targets[target_key].keys,

src/metatrain/utils/data/atomic_basis_helpers.py

@@ -129,7 +129,7 @@ def _densify_per_atom_atomic_basis_target(
                 properties=existing_block.properties,
             )
         else:
-            block = layout_block.copy()
+            block = layout_block.copy(deep=False)
             assert len(block.samples) == 0
         blocks.append(block)
 

src/metatrain/utils/data/target_info.py

@@ -758,7 +758,10 @@ def is_auxiliary_output(name: str) -> bool:
     :return: `True` if the target is an auxiliary output, `False` otherwise.
     """
     is_auxiliary = (
-        name == "features" or name == "energy_ensemble" or name.startswith("mtt::aux::")
+        name == "features"
+        or name == "feature"
+        or name == "energy_ensemble"
+        or name.startswith("mtt::aux::")
     )
     return is_auxiliary
 

src/metatrain/utils/evaluate_model.py

@@ -62,7 +62,7 @@ def evaluate_model(
     energy_targets_that_require_strain_gradients = []
     for target_name in targets.keys():
         # Check if the target is an energy:
-        if model_outputs[target_name].quantity == "energy":
+        if targets[target_name].quantity == "energy":
             energy_targets.append(target_name)
             if isinstance(targets[target_name], TargetInfo):
                 # Check if the energy requires gradients:
@@ -111,7 +111,7 @@ def evaluate_model(
                 destroy_graph=(index == len(energy_targets_with_gradients) - 1),
             )
             old_energy_tensor_map = model_outputs[energy_target]
-            new_block = old_energy_tensor_map.block().copy()
+            new_block = old_energy_tensor_map.block().copy(deep=False)
             new_block.add_gradient(
                 "positions", _position_gradients_to_block(gradients[: len(systems)])
             )
@@ -132,7 +132,7 @@ def evaluate_model(
                 destroy_graph=(index == len(energy_targets_with_gradients) - 1),
             )
             old_energy_tensor_map = model_outputs[energy_target]
-            new_block = old_energy_tensor_map.block().copy()
+            new_block = old_energy_tensor_map.block().copy(deep=False)
             new_block.add_gradient("positions", _position_gradients_to_block(gradients))
             new_energy_tensor_map = TensorMap(
                 keys=old_energy_tensor_map.keys,
@@ -147,7 +147,7 @@ def evaluate_model(
                 destroy_graph=(index == len(energy_targets_with_gradients) - 1),
             )
             old_energy_tensor_map = model_outputs[energy_target]
-            new_block = old_energy_tensor_map.block().copy()
+            new_block = old_energy_tensor_map.block().copy(deep=False)
             new_block.add_gradient("strain", _strain_gradients_to_block(gradients))
             new_energy_tensor_map = TensorMap(
                 keys=old_energy_tensor_map.keys,

src/metatrain/utils/per_atom.py

@@ -57,7 +57,7 @@ def divide_by_num_atoms(tensor_map: TensorMap, num_atoms: torch.Tensor) -> Tenso
     blocks = []
     for block in tensor_map.blocks():
         if "atom" in block.samples.names:
-            new_block = block.copy()
+            new_block = block.copy(deep=False)
         else:
             values = block.values / num_atoms.view(
                 -1, *[1] * (len(block.values.shape) - 1)
@@ -70,7 +70,7 @@ def divide_by_num_atoms(tensor_map: TensorMap, num_atoms: torch.Tensor) -> Tenso
             )
             for gradient_name, gradient in block.gradients():
                 if "atom" in gradient.samples.names:
-                    new_gradient = gradient.copy()
+                    new_gradient = gradient.copy(deep=False)
                 else:
                     values = gradient.values / num_atoms.view(
                         -1, *[1] * (len(gradient.values.shape) - 1)

src/metatrain/utils/scaler/_base_scaler.py

@@ -831,11 +831,11 @@ def _set_fixed_weights(
 
         self.scales[target_name] = TensorMap(
             self.Y2[target_name].keys.to(device=block.values.device),
-            [block],
+            [block.copy(deep=False)],
         )
         self.per_target_scales[target_name] = TensorMap(
             self.Y2[target_name].keys.to(device=block.values.device),
-            [block],
+            [block.copy(deep=False)],
         )
 
     def _sync_device_dtype(self, device: torch.device, dtype: torch.dtype) -> None:

src/metatrain/utils/testing/output.py

@@ -627,7 +627,7 @@ def test_output_features(
 
         features_output_options = ModelOutput(
             quantity="",
-            unit="unitless",
+            unit="",
             per_atom=per_atom,
         )
         model = model.to(system.positions.dtype)
@@ -705,7 +705,7 @@ def test_output_last_layer_features(
         # last-layer features per atom:
         ll_output_options = ModelOutput(
             quantity="",
-            unit="unitless",
+            unit="",
             per_atom=per_atom,
         )
         model = model.to(system.positions.dtype)

tests/utils/data/test_readers.py

@@ -53,7 +53,8 @@ def test_read_unknonw_library():
 def test_unsupported_target_name():
     conf = {"free_energy": {"quantity": "energy"}}
     with pytest.raises(
-        ValueError, match="Invalid name for model output: 'free_energy'"
+        ValueError,
+        match="invalid model output name 'free_energy': this is not a known quantity",
     ):
         read_targets(OmegaConf.create(conf))
 

tests/utils/test_metrics.py

@@ -189,15 +189,15 @@ def test_per_block(accumulator_class, tensor_map_with_grad_1, tensor_map_with_gr
     tensor_1 = TensorMap(
         keys=Labels.range("label_name", 2),
         blocks=[
-            tensor_map_with_grad_1.block().copy(),
-            tensor_map_with_grad_1.block().copy(),
+            tensor_map_with_grad_1.block().copy(deep=False),
+            tensor_map_with_grad_1.block().copy(deep=False),
         ],
     )
     tensor_2 = TensorMap(
         keys=Labels.range("label_name", 2),
         blocks=[
-            tensor_map_with_grad_2.block().copy(),
-            tensor_map_with_grad_2.block().copy(),
+            tensor_map_with_grad_2.block().copy(deep=False),
+            tensor_map_with_grad_2.block().copy(deep=False),
         ],
     )