From e51ff5076b510d4b759593b385021c7106d91a88 Mon Sep 17 00:00:00 2001 From: Guillaume Fraux Date: Fri, 30 May 2025 14:57:13 +0200 Subject: [PATCH] metatensor models have been renamed to metatomic models --- .github/workflows/docs.yml | 2 +- .github/workflows/test.yml | 4 +- README.md | 2 +- docs/conf.py | 2 +- docs/tutorials/index.rst | 2 +- ...nsor_tutorial.py => metatomic_tutorial.py} | 14 ++-- pyproject.toml | 6 +- tests/models/test_metatensor.py | 66 ------------------- tests/models/test_metatomic.py | 66 +++++++++++++++++++ .../models/{metatensor.py => metatomic.py} | 46 ++++++------- 10 files changed, 105 insertions(+), 105 deletions(-) rename examples/tutorials/{metatensor_tutorial.py => metatomic_tutorial.py} (82%) delete mode 100644 tests/models/test_metatensor.py create mode 100644 tests/models/test_metatomic.py rename torch_sim/models/{metatensor.py => metatomic.py} (87%) diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index af5da9ae..03afc521 100644 --- a/.github/workflows/docs.yml +++ b/.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: | diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index dc560663..f08a83c1 100644 --- a/.github/workflows/test.yml +++ b/.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 }} diff --git a/README.md b/README.md index 598ff824..6f096ed4 100644 --- a/README.md +++ b/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 diff --git a/docs/conf.py b/docs/conf.py index bd46a1d6..acb69117 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -66,7 +66,7 @@ "fairchem", "mace", "mattersim", - "metatensor", + "metatomic", "orb", "sevennet", "graphpes", diff --git a/docs/tutorials/index.rst b/docs/tutorials/index.rst index 3517bff5..8dad9bc2 100644 --- a/docs/tutorials/index.rst +++ b/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 diff --git a/examples/tutorials/metatensor_tutorial.py b/examples/tutorials/metatomic_tutorial.py similarity index 82% rename from examples/tutorials/metatensor_tutorial.py rename to examples/tutorials/metatomic_tutorial.py index af9585a7..284ff2bf 100644 --- a/examples/tutorials/metatensor_tutorial.py +++ b/examples/tutorials/metatomic_tutorial.py @@ -3,8 +3,8 @@ # Dependencies # /// script # dependencies = [ -# "metatrain[pet]==2025.6", -# "metatensor-torch==0.7.6" +# "metatrain[pet]==2025.7", +# "metatomic-torch>=0.1.1,<0.2" # ] # /// # @@ -12,10 +12,10 @@ # %% [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. """ diff --git a/pyproject.toml b/pyproject.toml index 61b18e1e..28e812f8 100644 --- a/pyproject.toml +++ b/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"] diff --git a/tests/models/test_metatensor.py b/tests/models/test_metatensor.py deleted file mode 100644 index 49f4fdb4..00000000 --- a/tests/models/test_metatensor.py +++ /dev/null @@ -1,66 +0,0 @@ -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 metatensor.torch.atomistic import ase_calculator - from metatrain.utils.io import load_model - - from torch_sim.models.metatensor import MetatensorModel -except ImportError: - pytest.skip("Metatensor 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 metatensor_calculator(device: torch.device): - """Load a pretrained metatensor model for testing.""" - return ase_calculator.MetatensorCalculator( - model=load_model( - "https://huggingface.co/lab-cosmo/pet-mad/resolve/main/models/pet-mad-latest.ckpt" - ).export(), - device=device, - ) - - -@pytest.fixture -def metatensor_model(device: torch.device) -> MetatensorModel: - """Create an MetatensorModel wrapper for the pretrained model.""" - return MetatensorModel( - model="pet-mad", - device=device, - ) - - -def test_metatensor_initialization(device: torch.device) -> None: - """Test that the metatensor model initializes correctly.""" - model = MetatensorModel( - model="pet-mad", - device=device, - ) - assert model.device == device - assert model.dtype == torch.float32 - - -test_mattersim_consistency = make_model_calculator_consistency_test( - test_name="metatensor", - model_fixture_name="metatensor_model", - calculator_fixture_name="metatensor_calculator", - sim_state_names=consistency_test_simstate_fixtures, - energy_atol=5e-5, -) - -test_mattersim_model_outputs = make_validate_model_outputs_test( - model_fixture_name="metatensor_model", -) diff --git a/tests/models/test_metatomic.py b/tests/models/test_metatomic.py new file mode 100644 index 00000000..98bad1ff --- /dev/null +++ b/tests/models/test_metatomic.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", +) diff --git a/torch_sim/models/metatensor.py b/torch_sim/models/metatomic.py similarity index 87% rename from torch_sim/models/metatensor.py rename to torch_sim/models/metatomic.py index f728e55f..31d26966 100644 --- a/torch_sim/models/metatensor.py +++ b/torch_sim/models/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,7 +17,7 @@ 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 @@ -25,7 +25,7 @@ try: - from metatensor.torch.atomistic import ( + from metatomic.torch import ( ModelEvaluationOptions, ModelOutput, System, @@ -34,13 +34,13 @@ 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: @@ -48,13 +48,13 @@ def __init__(self, err: ImportError = exc, *_args: Any, **_kwargs: Any) -> None: 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