From 65217b1d3c9f5449dd028c0b41b253026e3d3bae Mon Sep 17 00:00:00 2001 From: Rhys Goodall Date: Wed, 25 Mar 2026 11:28:54 -0400 Subject: [PATCH 1/6] fea: use upstream metatomic --- examples/tutorials/metatomic_tutorial.py | 2 +- pyproject.toml | 2 +- torch_sim/models/metatomic.py | 281 +---------------------- 3 files changed, 14 insertions(+), 271 deletions(-) diff --git a/examples/tutorials/metatomic_tutorial.py b/examples/tutorials/metatomic_tutorial.py index 4fdef426..ab657700 100644 --- a/examples/tutorials/metatomic_tutorial.py +++ b/examples/tutorials/metatomic_tutorial.py @@ -24,7 +24,7 @@ # %% from torch_sim.models.metatomic import MetatomicModel -model = MetatomicModel("pet-mad") +model = MetatomicModel("pet-mad") # type: ignore[arg-type] # %% [markdown] """ diff --git a/pyproject.toml b/pyproject.toml index 90de3c4a..18216e49 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -52,7 +52,7 @@ io = ["ase>=3.26", "phonopy>=2.37.0", "pymatgen>=2025.6.14"] symmetry = ["moyopy>=0.7.8"] mace = ["mace-torch>=0.3.15"] mattersim = ["mattersim>=0.1.2"] -metatomic = ["metatomic-torch>=0.1.3", "metatrain[pet]>=2025.12"] +metatomic = ["metatomic-torchsim>=0.0.1"] orb = ["orb-models>=0.6.0"] sevenn = ["sevenn[torchsim]>=0.12.1"] graphpes = ["graph-pes>=0.1", "mace-torch>=0.3.12"] diff --git a/torch_sim/models/metatomic.py b/torch_sim/models/metatomic.py index 758c50d6..f39e9e85 100644 --- a/torch_sim/models/metatomic.py +++ b/torch_sim/models/metatomic.py @@ -1,289 +1,32 @@ """Wrapper for metatomic-based models in TorchSim. -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 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 metatomic-torch package. +Re-exports the metatomic-torchsim package's TorchSim integration. """ import traceback import warnings -from pathlib import Path from typing import Any -import torch -import vesin.metatomic - -import torch_sim as ts -from torch_sim.models.interface import ModelInterface - try: - from metatomic.torch import ( - ModelEvaluationOptions, - ModelOutput, - System, - load_atomistic_model, + from metatomic_torchsim import MetatomicModel # pyright: ignore[reportMissingImports] +except ImportError as exc: + warnings.warn( + f"metatomic-torchsim import failed: {traceback.format_exc()}", stacklevel=2 ) - from metatrain.utils.io import load_model -except ImportError as exc: - warnings.warn(f"Metatomic import failed: {traceback.format_exc()}", stacklevel=2) + from torch_sim.models.interface import ModelInterface class MetatomicModel(ModelInterface): - """Metatomic model wrapper for torch-sim. - - This class is a placeholder for the MetatomicModel class. - It raises an ImportError if metatomic is not installed. - """ + """Placeholder when metatomic-torchsim is not installed.""" def __init__(self, err: ImportError = exc, *_args: Any, **_kwargs: Any) -> None: - """Dummy init for type checking.""" + """Raise the original ImportError.""" raise err + def forward(self, *_args: Any, **_kwargs: Any) -> Any: + """Unreachable — __init__ always raises.""" + raise NotImplementedError -class MetatomicModel(ModelInterface): - """Computes energies for a list of systems using a metatomic model. - - 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 metatomic's expected inputs. - - Attributes: - ... - """ - - def __init__( - self, - model: str | Path | None = None, - extensions_path: str | Path | None = None, - device: torch.device | str | None = None, - *, - check_consistency: bool = False, - compute_forces: bool = True, - compute_stress: bool = True, - ) -> None: - """Initialize the metatomic model for energy, force and stress calculations. - - Sets up a metatomic model for energy, force, and stress calculations within - the TorchSim framework. The model can be initialized with atomic numbers - and system indices, or these can be provided during the forward pass. - - Args: - 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". - extensions_path (str | Path | None): Optional, path to the folder containing - compiled extensions for the model. - device (torch.device | None): Device on which to run the model. If None, - defaults to "cuda" if available, otherwise "cpu". - check_consistency (bool): Whether to perform various consistency checks - during model evaluation. This should only be used in case of anomalous - behavior, as it can hurt performance significantly. - compute_forces (bool): Whether to compute forces. - compute_stress (bool): Whether to compute stresses. - - Raises: - TypeError: If model is neither a path nor "pet-mad". - """ - super().__init__() - - if model is None: - raise ValueError( - "A model path, or the name of a pre-defined model, must be provided. " - 'Currently only "pet-mad" is available as a pre-defined model.' - ) - - if model == "pet-mad": - path = "https://huggingface.co/lab-cosmo/pet-mad/resolve/v1.1.0/models/pet-mad-v1.1.0.ckpt" - self._model = load_model(path).export() - elif str(model).endswith(".ckpt"): - path = model - self._model = load_model(path).export() - elif str(model).endswith(".pt"): - path = model - self._model = load_atomistic_model(path, extensions_path) - else: - raise ValueError('Model must be a path to a .ckpt/.pt file, or "pet-mad".') - - if "energy" not in self._model.capabilities().outputs: - raise ValueError( - "This model does not support energy predictions. " - "The model must have an `energy` output to be used in TorchSim." - ) - - resolved_device: torch.device - if device is None: - resolved_device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - elif isinstance(device, str): - resolved_device = torch.device(device) - else: - resolved_device = device - self._device = resolved_device - if self._device.type not in self._model.capabilities().supported_devices: - raise ValueError( - f"Model does not support device {self._device}. Supported devices: " - f"{self._model.capabilities().supported_devices}. You might want to " - f"set the `device` argument to a supported device." - ) - - self._dtype = getattr(torch, self._model.capabilities().dtype) - self._model.to(self._device) - self._compute_forces = compute_forces - self._compute_stress = compute_stress - self._memory_scales_with = "n_atoms_x_density" # for the majority of models - self._check_consistency = check_consistency - self._requested_neighbor_lists = self._model.requested_neighbor_lists() - self._evaluation_options = ModelEvaluationOptions( - length_unit="angstrom", - outputs={"energy": ModelOutput(quantity="energy", unit="eV", per_atom=False)}, - ) - - def forward( # noqa: C901, PLR0915 - self, state: ts.SimState, **_kwargs: Any - ) -> dict[str, torch.Tensor]: - """Compute energies, forces, and stresses for the given atomic systems. - - Processes the provided state information and computes energies, forces, and - stresses using the underlying metatomic model. Handles batched calculations for - multiple systems as well as constructing the necessary neighbor lists. - - Args: - state (SimState): State object containing positions, cell, and other - system information. - **_kwargs: Unused; accepted for interface compatibility. - - Returns: - dict[str, torch.Tensor]: Computed properties: - - 'energy': System energies with shape [n_systems] - - 'forces': Atomic forces with shape [n_atoms, 3] if compute_forces=True - - 'stress': System stresses with shape [n_systems, 3, 3] if - compute_stress=True - """ - sim_state = state - - # Input validation is already done inside the forward method of the - # AtomisticModel class, so we don't need to do it again here. - - atomic_nums = sim_state.atomic_numbers - cell = sim_state.row_vector_cell - positions = sim_state.positions - - # 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, 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 - # caller code needs to call torch.autograd.grad or similar to compute them - # from the energy output - - # Process each system separately - systems: list[System] = [] - strains = [] - for sys_idx in range(len(cell)): - system_mask = sim_state.system_idx == sys_idx - system_positions = positions[system_mask] - system_cell = cell[sys_idx] - system_atomic_numbers = atomic_nums[system_mask] - - # Create a System object for this system - if self._compute_forces: - system_positions.requires_grad_() - if self._compute_stress: - strain = torch.eye( - 3, device=self._device, dtype=self._dtype, requires_grad=True - ) - system_positions = system_positions @ strain - system_cell = system_cell @ strain - strains.append(strain) - - systems.append( - System( - positions=system_positions, - types=system_atomic_numbers, - cell=system_cell, - pbc=sim_state.pbc, - ) - ) - - # Calculate the required neighbor list(s) for all the systems - - # move data to CPU because vesin only supports CPU for now - systems = [system.to(device="cpu") for system in systems] - vesin.metatomic.compute_requested_neighbors( - systems, system_length_unit="Angstrom", model=self._model - ) - # move back to the proper device - systems = [system.to(device=self.device) for system in systems] - - # Get model output - model_outputs = self._model( - systems=systems, - options=self._evaluation_options, - check_consistency=self._check_consistency, - ) - - results: dict[str, torch.Tensor] = {} - results["energy"] = model_outputs["energy"].block().values.squeeze(-1) - - # Compute forces and/or stresses if requested - tensors_for_autograd = [] - if self._compute_forces: - for system in systems: - tensors_for_autograd.append(system.positions) # noqa: PERF401 - if self._compute_stress: - for strain in strains: - tensors_for_autograd.append(strain) # noqa: PERF402 - - if self._compute_forces or self._compute_stress: - derivatives = torch.autograd.grad( - outputs=model_outputs["energy"].block().values, - inputs=tensors_for_autograd, - grad_outputs=torch.ones_like(model_outputs["energy"].block().values), - ) - else: - derivatives = [] - - results_by_system: dict[str, list[torch.Tensor]] = {} - if self._compute_forces and self._compute_stress: - results_by_system["forces"] = [-d for d in derivatives[: len(systems)]] - results_by_system["stress"] = [ - d / torch.abs(torch.det(system.cell.detach())) - for d, system in zip(derivatives[len(systems) :], systems, strict=False) - ] - elif self._compute_forces: - results_by_system["forces"] = [-d for d in derivatives] - elif self._compute_stress: - results_by_system["stress"] = [ - d / torch.abs(torch.det(system.cell.detach())) - for d, system in zip(derivatives, systems, strict=False) - ] - else: - pass - - # Concatenate/stack forces and stresses - if self._compute_forces: - if len(results_by_system["forces"]) > 0: - results["forces"] = torch.cat(results_by_system["forces"]) - else: - results["forces"] = torch.empty_like(positions) - if self._compute_stress: - if len(results_by_system["stress"]) > 0: - results["stress"] = torch.stack(results_by_system["stress"]) - else: - results["stress"] = torch.empty_like(cell) - return {k: v.detach() for k, v in results.items()} +__all__ = ["MetatomicModel"] From de293128ee14708a3f055a801b321692a7078f2d Mon Sep 17 00:00:00 2001 From: Rhys Goodall Date: Wed, 25 Mar 2026 12:07:43 -0400 Subject: [PATCH 2/6] clean: remove the tutorial here given upstream tutorials --- examples/tutorials/metatomic_tutorial.py | 69 ------------------------ 1 file changed, 69 deletions(-) delete mode 100644 examples/tutorials/metatomic_tutorial.py diff --git a/examples/tutorials/metatomic_tutorial.py b/examples/tutorials/metatomic_tutorial.py deleted file mode 100644 index ab657700..00000000 --- a/examples/tutorials/metatomic_tutorial.py +++ /dev/null @@ -1,69 +0,0 @@ -# %% -# /// script -# dependencies = [ -# "torch_sim_atomistic[metatomic]" -# ] -# /// - - -# %% [markdown] -""" -# 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 metatomic interface. - -## Loading the model - -Loading the model is simple: you simply need to specify the model name (in this case -"pet-mad"), as shown below. All other arguments are optional: for example, you could -specify the device. (If the device is not specified, like in this case, the optimal -device is chosen automatically.) -""" - -# %% -from torch_sim.models.metatomic import MetatomicModel - -model = MetatomicModel("pet-mad") # type: ignore[arg-type] - -# %% [markdown] -""" -## Using the model to run a molecular dynamics simulations - -Once the model is loaded, you can use it just like any other TorchSim model to run -simulations. Here, we show how to run a simple MD simulation consisting of an initial -NVT equilibration run followed by an NVE run. -""" -# %% -from ase.build import bulk -import torch_sim as ts - -atoms = bulk("Si", "diamond", a=5.43, cubic=True) - -equilibrated_state = ts.integrate( - system=atoms, - model=model, - integrator=ts.Integrator.nvt_langevin, - n_steps=100, - temperature=300, # K - timestep=0.001, # ps -) - -final_state = ts.integrate( - system=equilibrated_state, - model=model, - integrator=ts.Integrator.nve, - n_steps=100, - temperature=300, # K - timestep=0.001, # ps -) - -# %% [markdown] -""" -## Further steps - -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 -metatomic and PET-MAD. For more details on how to use TorchSim, you can refer to the -other tutorials in this section. -""" From 6b3c39a8bd555ed98d0e38e458b4bb43a38df683 Mon Sep 17 00:00:00 2001 From: Rhys Goodall Date: Thu, 2 Apr 2026 10:12:45 -0400 Subject: [PATCH 3/6] wip: broken model download of pet-mad --- pyproject.toml | 2 +- tests/models/test_metatomic.py | 26 +++++++++++++++++--------- 2 files changed, 18 insertions(+), 10 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 18216e49..ca722de6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -52,7 +52,7 @@ io = ["ase>=3.26", "phonopy>=2.37.0", "pymatgen>=2025.6.14"] symmetry = ["moyopy>=0.7.8"] mace = ["mace-torch>=0.3.15"] mattersim = ["mattersim>=0.1.2"] -metatomic = ["metatomic-torchsim>=0.0.1"] +metatomic = ["metatomic-torchsim>=0.1.1", "metatomic-ase>=0.1.0"] orb = ["orb-models>=0.6.0"] sevenn = ["sevenn[torchsim]>=0.12.1"] graphpes = ["graph-pes>=0.1", "mace-torch>=0.3.12"] diff --git a/tests/models/test_metatomic.py b/tests/models/test_metatomic.py index c42fa845..c0fbae7c 100644 --- a/tests/models/test_metatomic.py +++ b/tests/models/test_metatomic.py @@ -12,8 +12,9 @@ try: - from metatomic.torch import ase_calculator - from metatrain.utils.io import load_model + import requests + from metatomic.torch import AtomisticModel, load_atomistic_model + from metatomic_ase import MetatomicCalculator from torch_sim.models.metatomic import MetatomicModel except ImportError: @@ -24,18 +25,25 @@ @pytest.fixture -def metatomic_calculator(): - """Load a pretrained metatomic model for testing.""" +def metatomic_module(tmp_path): model_url = "https://huggingface.co/lab-cosmo/pet-mad/resolve/v1.1.0/models/pet-mad-v1.1.0.ckpt" - return ase_calculator.MetatomicCalculator( - model=load_model(model_url).export(), device=DEVICE - ) + model_path = tmp_path / "pet-mad-v1.1.0.ckpt" + response = requests.get(model_url) + response.raise_for_status() + model_path.write_bytes(response.content) + return load_atomistic_model(model_path) + + +@pytest.fixture +def metatomic_calculator(metatomic_module: AtomisticModel) -> MetatomicCalculator: + """Load a pretrained metatomic model for testing.""" + return MetatomicCalculator(model=metatomic_module, device=DEVICE) @pytest.fixture -def metatomic_model() -> MetatomicModel: +def metatomic_model(metatomic_module: AtomisticModel) -> MetatomicModel: """Create an MetatomicModel wrapper for the pretrained model.""" - return MetatomicModel(model="pet-mad", device=DEVICE) + return MetatomicModel(model=metatomic_module, device=DEVICE) def test_metatomic_initialization() -> None: From 34402403f402f924f28f57c27f77df55c0977f38 Mon Sep 17 00:00:00 2001 From: Rhys Goodall Date: Thu, 2 Apr 2026 10:23:40 -0400 Subject: [PATCH 4/6] just use upet utils as that's what most users will want? --- pyproject.toml | 2 +- tests/models/test_metatomic.py | 21 +++++---------------- 2 files changed, 6 insertions(+), 17 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index ca722de6..0090b57f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -52,7 +52,7 @@ io = ["ase>=3.26", "phonopy>=2.37.0", "pymatgen>=2025.6.14"] symmetry = ["moyopy>=0.7.8"] mace = ["mace-torch>=0.3.15"] mattersim = ["mattersim>=0.1.2"] -metatomic = ["metatomic-torchsim>=0.1.1", "metatomic-ase>=0.1.0"] +metatomic = ["metatomic-torchsim>=0.1.1", "metatomic-ase>=0.1.0", "upet>=0.2.0"] orb = ["orb-models>=0.6.0"] sevenn = ["sevenn[torchsim]>=0.12.1"] graphpes = ["graph-pes>=0.1", "mace-torch>=0.3.12"] diff --git a/tests/models/test_metatomic.py b/tests/models/test_metatomic.py index c0fbae7c..c853f55a 100644 --- a/tests/models/test_metatomic.py +++ b/tests/models/test_metatomic.py @@ -12,9 +12,9 @@ try: - import requests - from metatomic.torch import AtomisticModel, load_atomistic_model + from metatomic.torch import AtomisticModel from metatomic_ase import MetatomicCalculator + from upet import get_upet from torch_sim.models.metatomic import MetatomicModel except ImportError: @@ -25,33 +25,22 @@ @pytest.fixture -def metatomic_module(tmp_path): - model_url = "https://huggingface.co/lab-cosmo/pet-mad/resolve/v1.1.0/models/pet-mad-v1.1.0.ckpt" - model_path = tmp_path / "pet-mad-v1.1.0.ckpt" - response = requests.get(model_url) - response.raise_for_status() - model_path.write_bytes(response.content) - return load_atomistic_model(model_path) +def metatomic_module() -> AtomisticModel: + return get_upet(model="pet-mad") @pytest.fixture def metatomic_calculator(metatomic_module: AtomisticModel) -> MetatomicCalculator: - """Load a pretrained metatomic model for testing.""" return MetatomicCalculator(model=metatomic_module, device=DEVICE) @pytest.fixture def metatomic_model(metatomic_module: AtomisticModel) -> MetatomicModel: - """Create an MetatomicModel wrapper for the pretrained model.""" return MetatomicModel(model=metatomic_module, device=DEVICE) def test_metatomic_initialization() -> None: - """Test that the metatomic model initializes correctly.""" - model = MetatomicModel( - model="pet-mad", - device=DEVICE, - ) + model = MetatomicModel(model=get_upet(model="pet-mad"), device=DEVICE) assert model.device == DEVICE assert model.dtype == torch.float32 From 5920087df31cbf7d54d3170e7158c524c5f49bdf Mon Sep 17 00:00:00 2001 From: Rhys Goodall Date: Thu, 2 Apr 2026 10:51:06 -0400 Subject: [PATCH 5/6] require torch optional dep in alchemiops --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 0090b57f..f083769c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -30,7 +30,7 @@ dependencies = [ "h5py>=3.15.1", "numpy>=1.26,<3; python_version < '3.13'", "numpy>=2.3.2,<3; python_version >= '3.13'", - "nvalchemi-toolkit-ops>=0.3.0", + "nvalchemi-toolkit-ops[torch]>=0.3.0", "tables>=3.11.1", "torch>=2", "tqdm>=4.67", From c1e173ed515c662027067f557edab825af5a4898 Mon Sep 17 00:00:00 2001 From: Rhys Goodall Date: Thu, 2 Apr 2026 11:07:06 -0400 Subject: [PATCH 6/6] bump orb for alchemiops 0.3.0 --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index f083769c..3f60469b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -53,7 +53,7 @@ symmetry = ["moyopy>=0.7.8"] mace = ["mace-torch>=0.3.15"] mattersim = ["mattersim>=0.1.2"] metatomic = ["metatomic-torchsim>=0.1.1", "metatomic-ase>=0.1.0", "upet>=0.2.0"] -orb = ["orb-models>=0.6.0"] +orb = ["orb-models>=0.6.2"] sevenn = ["sevenn[torchsim]>=0.12.1"] graphpes = ["graph-pes>=0.1", "mace-torch>=0.3.12"] nequip = ["nequip>=0.17.0"]