metatensor models have been renamed to metatomic models (#205)

Author: Luthaf

.github/workflows/docs.yml

@@ -34,7 +34,7 @@ jobs:
           uv pip install ".[test,docs]" --system
 
       - name: Install extras for tutorial generation
-        run: uv pip install ".[graphpes,mace,metatensor]" --system
+        run: uv pip install ".[graphpes,mace,metatomic]" --system
 
       - name: Copy tutorials
         run: |

.github/workflows/test.yml

@@ -49,7 +49,7 @@ jobs:
             --ignore=tests/models/test_orb.py \
             --ignore=tests/models/test_sevennet.py \
             --ignore=tests/models/test_mattersim.py \
-            --ignore=tests/models/test_metatensor.py \
+            --ignore=tests/models/test_metatomic.py \
             --ignore=tests/test_optimizers_vs_ase.py \
 
       - name: Upload coverage to Codecov
@@ -73,7 +73,7 @@ jobs:
           - { name: mace, test_path: "tests/test_elastic.py" }
           - { name: mace, test_path: "tests/test_optimizers_vs_ase.py" }
           - { name: mattersim, test_path: "tests/models/test_mattersim.py" }
-          - { name: metatensor, test_path: "tests/models/test_metatensor.py" }
+          - { name: metatomic, test_path: "tests/models/test_metatomic.py" }
           - { name: orb, test_path: "tests/models/test_orb.py" }
           - { name: sevenn, test_path: "tests/models/test_sevennet.py" }
     runs-on: ${{ matrix.os }}

README.md

@@ -14,7 +14,7 @@ era. By rewriting the core primitives of atomistic simulation in Pytorch, it all
 orders of magnitude acceleration of popular machine learning potentials.
 
 * Automatic batching and GPU memory management allowing significant simulation speedup
-* Support for MACE, Fairchem, SevenNet, ORB, MatterSim, graph-pes, and metatensor MLIP models
+* Support for MACE, Fairchem, SevenNet, ORB, MatterSim, graph-pes, and metatomic MLIP models
 * Support for classical lennard jones, morse, and soft-sphere potentials
 * Molecular dynamics integration schemes like NVE, NVT Langevin, and NPT Langevin
 * Relaxation of atomic positions and cell with gradient descent and FIRE

docs/conf.py

@@ -66,7 +66,7 @@
     "fairchem",
     "mace",
     "mattersim",
-    "metatensor",
+    "metatomic",
     "orb",
     "sevennet",
     "graphpes",

docs/tutorials/index.rst

@@ -19,4 +19,4 @@ versions of the tutorials can also be found in the `torch-sim /examples/tutorial
     low_level_tutorial
     hybrid_swap_tutorial
     using_graphpes_tutorial
-    metatensor_tutorial
+    metatomic_tutorial

examples/tutorials/metatensor_tutorial.py renamed to examples/tutorials/metatomic_tutorial.py

@@ -3,19 +3,19 @@
 #   <summary>Dependencies</summary>
 # /// script
 # dependencies = [
-#     "metatrain[pet]==2025.6",
-#     "metatensor-torch==0.7.6"
+#     "metatrain[pet]==2025.7",
+#     "metatomic-torch>=0.1.1,<0.2"
 # ]
 # ///
 # </details>
 
 
 # %% [markdown]
 """
-# Using the PET-MAD model with metatensor
+# Using the PET-MAD model with TorchSim
 
 This tutorial explains how to use the PET-MAD model (https://arxiv.org/abs/2503.14118)
-via TorchSim's metatensor interface.
+via TorchSim's metatomic interface.
 
 ## Loading the model
 
@@ -26,9 +26,9 @@
 """
 
 # %%
-from torch_sim.models.metatensor import MetatensorModel
+from torch_sim.models.metatomic import MetatomicModel
 
-model = MetatensorModel("pet-mad")
+model = MetatomicModel("pet-mad")
 
 # %% [markdown]
 """
@@ -68,6 +68,6 @@
 
 Of course, in reality, you would want to run the simulation for much longer, probably
 save trajectories, and much more. However, this is all you need to get started with
-metatensor and PET-MAD. For more details on how to use TorchSim, you can refer to the
+metatomic and PET-MAD. For more details on how to use TorchSim, you can refer to the
 other tutorials in this section.
 """

pyproject.toml

@@ -32,8 +32,8 @@ dependencies = [
     "tables>=3.10.2",
     "torch>=2",
     "tqdm>=4.67",
-    "vesin-torch>=0.3.1",
-    "vesin>=0.3.1",
+    "vesin-torch>=0.3.7",
+    "vesin>=0.3.7",
 ]
 
 [project.optional-dependencies]
@@ -48,7 +48,7 @@ test = [
 io = ["ase>=3.24", "phonopy>=2.37.0", "pymatgen>=2024.11.3"]
 mace = ["mace-torch>=0.3.12"]
 mattersim = ["mattersim>=0.1.2"]
-metatensor = ["metatensor-torch==0.7.6", "metatrain[pet]==2025.6"]
+metatomic = ["metatomic-torch>=0.1.1,<0.2", "metatrain[pet]==2025.7"]
 orb = ["orb-models>=0.5.2"]
 sevenn = ["sevenn>=0.11.0"]
 graphpes = ["graph-pes>=0.0.34", "mace-torch>=0.3.12"]

tests/models/test_metatensor.py

@@ -0,0 +1,66 @@
+import pytest
+import torch
+
+from tests.models.conftest import (
+    consistency_test_simstate_fixtures,
+    make_model_calculator_consistency_test,
+    make_validate_model_outputs_test,
+)
+
+
+try:
+    from metatomic.torch import ase_calculator
+    from metatrain.utils.io import load_model
+
+    from torch_sim.models.metatomic import MetatomicModel
+except ImportError:
+    pytest.skip("metatomic not installed", allow_module_level=True)
+
+
+@pytest.fixture
+def dtype() -> torch.dtype:
+    """Fixture to provide the default dtype for testing."""
+    return torch.float32
+
+
+@pytest.fixture
+def metatomic_calculator(device: torch.device):
+    """Load a pretrained metatomic model for testing."""
+    return ase_calculator.MetatomicCalculator(
+        model=load_model(
+            "https://huggingface.co/lab-cosmo/pet-mad/resolve/main/models/pet-mad-latest.ckpt"
+        ).export(),
+        device=device,
+    )
+
+
+@pytest.fixture
+def metatomic_model(device: torch.device) -> MetatomicModel:
+    """Create an MetatomicModel wrapper for the pretrained model."""
+    return MetatomicModel(
+        model="pet-mad",
+        device=device,
+    )
+
+
+def test_metatomic_initialization(device: torch.device) -> None:
+    """Test that the metatomic model initializes correctly."""
+    model = MetatomicModel(
+        model="pet-mad",
+        device=device,
+    )
+    assert model.device == device
+    assert model.dtype == torch.float32
+
+
+test_metatomic_consistency = make_model_calculator_consistency_test(
+    test_name="metatomic",
+    model_fixture_name="metatomic_model",
+    calculator_fixture_name="metatomic_calculator",
+    sim_state_names=consistency_test_simstate_fixtures,
+    energy_atol=5e-5,
+)
+
+test_metatomic_model_outputs = make_validate_model_outputs_test(
+    model_fixture_name="metatomic_model",
+)

tests/models/test_metatomic.py

@@ -1,14 +1,14 @@
-"""Wrapper for metatensor-based models in TorchSim.
+"""Wrapper for metatomic-based models in TorchSim.
 
-This module provides a TorchSim wrapper of metatensor models for computing
+This module provides a TorchSim wrapper of metatomic models for computing
 energies, forces, and stresses for atomistic systems, including batched computations
 for multiple systems simultaneously.
 
-The MetatensorModel class adapts metatensor models to the ModelInterface protocol,
+The MetatomicModel class adapts metatomic models to the ModelInterface protocol,
 allowing them to be used within the broader torch_sim simulation framework.
 
 Notes:
-    This module depends on the metatensor-torch package.
+    This module depends on the metatomic-torch package.
 """
 
 import traceback
@@ -17,15 +17,15 @@
 from typing import Any
 
 import torch
-import vesin.torch.metatensor
+import vesin.metatomic
 
 import torch_sim as ts
 from torch_sim.models.interface import ModelInterface
 from torch_sim.typing import StateDict
 
 
 try:
-    from metatensor.torch.atomistic import (
+    from metatomic.torch import (
         ModelEvaluationOptions,
         ModelOutput,
         System,
@@ -34,27 +34,27 @@
     from metatrain.utils.io import load_model
 
 except ImportError as exc:
-    warnings.warn(f"Metatensor import failed: {traceback.format_exc()}", stacklevel=2)
+    warnings.warn(f"Metatomic import failed: {traceback.format_exc()}", stacklevel=2)
 
-    class MetatensorModel(torch.nn.Module, ModelInterface):
-        """Metatensor model wrapper for torch_sim.
+    class MetatomicModel(torch.nn.Module, ModelInterface):
+        """Metatomic model wrapper for torch_sim.
 
-        This class is a placeholder for the MetatensorModel class.
-        It raises an ImportError if metatensor is not installed.
+        This class is a placeholder for the MetatomicModel class.
+        It raises an ImportError if metatomic is not installed.
         """
 
         def __init__(self, err: ImportError = exc, *_args: Any, **_kwargs: Any) -> None:
             """Dummy init for type checking."""
             raise err
 
 
-class MetatensorModel(torch.nn.Module, ModelInterface):
-    """Computes energies for a list of systems using a metatensor model.
+class MetatomicModel(torch.nn.Module, ModelInterface):
+    """Computes energies for a list of systems using a metatomic model.
 
-    This class wraps a metatensor model to compute energies, forces, and stresses for
+    This class wraps a metatomic model to compute energies, forces, and stresses for
     atomic systems within the TorchSim framework. It supports batched calculations
     for multiple systems and handles the necessary transformations between
-    TorchSim's data structures and metatensor's expected inputs.
+    TorchSim's data structures and metatomic's expected inputs.
 
     Attributes:
         ...
@@ -70,14 +70,14 @@ def __init__(
         compute_forces: bool = True,
         compute_stress: bool = True,
     ) -> None:
-        """Initialize the metatensor model for energy, force and stress calculations.
+        """Initialize the metatomic model for energy, force and stress calculations.
 
-        Sets up a metatensor model for energy, force, and stress calculations within
+        Sets up a metatomic model for energy, force, and stress calculations within
         the TorchSim framework. The model can be initialized with atomic numbers
         and batch indices, or these can be provided during the forward pass.
 
         Args:
-            model (str | Path | None): Path to the metatensor model file or a
+            model (str | Path | None): Path to the metatomic model file or a
                 pre-defined model name. Currently only "pet-mad"
                 (https://arxiv.org/abs/2503.14118) is supported as a pre-defined model.
                 If None, defaults to "pet-mad".
@@ -155,7 +155,7 @@ def forward(  # noqa: C901, PLR0915
         """Compute energies, forces, and stresses for the given atomic systems.
 
         Processes the provided state information and computes energies, forces, and
-        stresses using the underlying metatensor model. Handles batched calculations for
+        stresses using the underlying metatomic model. Handles batched calculations for
         multiple systems as well as constructing the necessary neighbor lists.
 
         Args:
@@ -175,21 +175,21 @@ def forward(  # noqa: C901, PLR0915
             state = ts.SimState(**state, masses=torch.ones_like(state["positions"]))
 
         # Input validation is already done inside the forward method of the
-        # MetatensorAtomisticModel class, so we don't need to do it again here.
+        # AtomisticModel class, so we don't need to do it again here.
 
         atomic_numbers = state.atomic_numbers
         cell = state.row_vector_cell
         positions = state.positions
         pbc = state.pbc
 
-        # Check dtype (metatensor models require a specific input dtype)
+        # Check dtype (metatomic models require a specific input dtype)
         if positions.dtype != self._dtype:
             raise TypeError(
                 f"Positions dtype {positions.dtype} does not match model dtype "
                 f"{self._dtype}"
             )
 
-        # Compared to other models, metatensor models have two peculiarities:
+        # Compared to other models, metatomic models have two peculiarities:
         # - different structures are fed to the models separately as a list of System
         #   objects, and not as a single graph-like batch
         # - the model does not compute forces and stresses itself, but rather the
@@ -232,7 +232,7 @@ def forward(  # noqa: C901, PLR0915
 
         # move data to CPU because vesin only supports CPU for now
         systems = [system.to(device="cpu") for system in systems]
-        vesin.torch.metatensor.compute_requested_neighbors(
+        vesin.metatomic.compute_requested_neighbors(
             systems, system_length_unit="Angstrom", model=self._model
         )
         # move back to the proper device