From 02a52b3e203d1a6164aaf7d93e4c07706c09438b Mon Sep 17 00:00:00 2001 From: Rohit Goswami Date: Tue, 17 Mar 2026 10:52:41 +0100 Subject: [PATCH] Create a metatomic-torchsim package It will hold the TorchSim <=> metatomic integration as a standalone module. Co-Authored-By: Rhys Goodall Co-Authored-By: Guillaume Fraux --- .github/workflows/build-wheels.yml | 13 + .github/workflows/torch-tests.yml | 2 +- .github/workflows/torchsim-tests.yml | 60 +++ docs/src/conf.py | 9 +- docs/src/engines/torch-sim.rst | 46 ++- docs/src/index.rst | 2 +- metatomic-torch/CHANGELOG.md | 1 - pyproject.toml | 4 +- python/metatomic_torch/pyproject.toml | 2 + .../metatomic_torch/tests/ase_calculator.py | 2 +- python/metatomic_torchsim/AUTHORS | 4 + python/metatomic_torchsim/Architecture.md | 71 ++++ python/metatomic_torchsim/CHANGELOG.md | 15 + python/metatomic_torchsim/MANIFEST.in | 2 + python/metatomic_torchsim/README.md | 33 ++ .../metatomic_torchsim/__init__.py | 4 + .../metatomic_torchsim/_model.py | 390 ++++++++++++++++++ python/metatomic_torchsim/pyproject.toml | 61 +++ python/metatomic_torchsim/setup.py | 114 +++++ python/metatomic_torchsim/tests/__init__.py | 0 .../metatomic_torchsim/tests/model_loading.py | 56 +++ python/metatomic_torchsim/tests/torchsim.py | 281 +++++++++++++ setup.py | 7 + tox.ini | 48 ++- 24 files changed, 1206 insertions(+), 21 deletions(-) create mode 100644 .github/workflows/torchsim-tests.yml create mode 100644 python/metatomic_torchsim/AUTHORS create mode 100644 python/metatomic_torchsim/Architecture.md create mode 100644 python/metatomic_torchsim/CHANGELOG.md create mode 100644 python/metatomic_torchsim/MANIFEST.in create mode 100644 python/metatomic_torchsim/README.md create mode 100644 python/metatomic_torchsim/metatomic_torchsim/__init__.py create mode 100644 python/metatomic_torchsim/metatomic_torchsim/_model.py create mode 100644 python/metatomic_torchsim/pyproject.toml create mode 100644 python/metatomic_torchsim/setup.py create mode 100644 python/metatomic_torchsim/tests/__init__.py create mode 100644 python/metatomic_torchsim/tests/model_loading.py create mode 100644 python/metatomic_torchsim/tests/torchsim.py diff --git a/.github/workflows/build-wheels.yml b/.github/workflows/build-wheels.yml index 9f6d8f984..42379bedd 100644 --- a/.github/workflows/build-wheels.yml +++ b/.github/workflows/build-wheels.yml @@ -207,6 +207,9 @@ jobs: - name: build metatomic sdist and wheel run: python -m build . --outdir=dist/ + - name: build metatomic_torchsim sdist and wheel + run: python -m build python/metatomic_torchsim --outdir=dist/ + - name: check sdist and wheels with twine run: twine check dist/*.tar.gz dist/*.whl @@ -279,6 +282,16 @@ jobs: env: GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + - name: upload to GitHub release (metatomic-torchsim) + if: startsWith(github.ref, 'refs/tags/metatomic-torchsim-v') + uses: softprops/action-gh-release@v2 + with: + files: | + wheels/metatomic_torchsim-* + prerelease: ${{ contains(github.ref, '-rc') }} + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + test-build-external: # This checks building the wheels with external libraries. This setup is # mainly used for the conda packages metatensor-*-python, which use the diff --git a/.github/workflows/torch-tests.yml b/.github/workflows/torch-tests.yml index 4aa30a9ae..3a6286bac 100644 --- a/.github/workflows/torch-tests.yml +++ b/.github/workflows/torch-tests.yml @@ -64,7 +64,7 @@ jobs: python -m pip install tox coverage - name: run tests - run: tox + run: tox -e lint,torch-tests,torch-tests-cxx,torch-install-tests-cxx,docs-tests env: PIP_EXTRA_INDEX_URL: https://download.pytorch.org/whl/cpu METATOMIC_TESTS_TORCH_VERSION: ${{ matrix.torch-version }} diff --git a/.github/workflows/torchsim-tests.yml b/.github/workflows/torchsim-tests.yml new file mode 100644 index 000000000..a6953f4fc --- /dev/null +++ b/.github/workflows/torchsim-tests.yml @@ -0,0 +1,60 @@ +name: TorchSim integration + +on: + push: + branches: [main] + pull_request: + # Check all PR + +concurrency: + group: torchsim-tests-${{ github.ref }} + cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} + +jobs: + tests: + runs-on: ubuntu-24.04 + name: tests + steps: + - uses: actions/checkout@v6 + with: + fetch-depth: 0 + + - name: setup Python + uses: actions/setup-python@v6 + with: + python-version: "3.13" + + - name: Setup sccache + uses: mozilla-actions/sccache-action@v0.0.9 + with: + version: "v0.10.0" + + - name: Setup sccache environnement variables + run: | + echo "SCCACHE_GHA_ENABLED=true" >> $GITHUB_ENV + echo "RUSTC_WRAPPER=sccache" >> $GITHUB_ENV + echo "CMAKE_C_COMPILER_LAUNCHER=sccache" >> $GITHUB_ENV + echo "CMAKE_CXX_COMPILER_LAUNCHER=sccache" >> $GITHUB_ENV + + - name: install tests dependencies + run: | + python -m pip install --upgrade pip + python -m pip install tox coverage + + - name: run tests + run: tox -e torchsim-tests + env: + PIP_EXTRA_INDEX_URL: https://download.pytorch.org/whl/cpu + + - name: combine Python coverage files + shell: bash + run: | + coverage combine .tox/*/.coverage + coverage xml + + - name: upload to codecov.io + uses: codecov/codecov-action@v5 + with: + fail_ci_if_error: true + files: coverage.xml + token: ${{ secrets.CODECOV_TOKEN }} diff --git a/docs/src/conf.py b/docs/src/conf.py index 74ff15ead..c053c0a1d 100644 --- a/docs/src/conf.py +++ b/docs/src/conf.py @@ -171,16 +171,17 @@ def setup(app): "python": ("https://docs.python.org/3", None), "numpy": ("https://numpy.org/doc/stable/", None), "torch": ("https://docs.pytorch.org/docs/stable/", None), - "featomic": ("https://metatensor.github.io/featomic/latest/", None), + "featomic": ("http://docs.metatensor.org/featomic/latest/", None), "metatensor": ("https://docs.metatensor.org/latest/", None), "ase": ("https://ase-lib.org/", None), "matplotlib": ("https://matplotlib.org/stable/", None), + "torch_sim": ("https://torchsim.github.io/torch-sim/", None), } # sitemap/SEO settings -html_baseurl = "https://docs.metatensor.org/metatomic/latest/" # prefix for the sitemap -sitemap_url_scheme = "{link}" # avoids language settings -html_extra_path = ["robots.txt"] # extra files to move +html_baseurl = "https://docs.metatensor.org/metatomic/latest/" # prefix for the sitemap +sitemap_url_scheme = "{link}" # avoids language settings +html_extra_path = ["robots.txt"] # extra files to move # -- Options for HTML output ------------------------------------------------- diff --git a/docs/src/engines/torch-sim.rst b/docs/src/engines/torch-sim.rst index 77b4edec5..4932994d9 100644 --- a/docs/src/engines/torch-sim.rst +++ b/docs/src/engines/torch-sim.rst @@ -8,24 +8,48 @@ torch-sim * - Official website - How is metatomic supported? - * - https://radical-ai.github.io/torch-sim/ - - In the official version + * - https://torchsim.github.io/torch-sim/ + - Via the ``metatomic-torchsim`` package -Supported model outputs +How to install the code ^^^^^^^^^^^^^^^^^^^^^^^ -Only the :ref:`energy ` output is supported. +Install the integration package from PyPI: -How to install the code +.. code-block:: bash + + pip install metatomic-torchsim + +For the full TorchSim documentation, see https://torchsim.github.io/torch-sim/. + +Supported model outputs ^^^^^^^^^^^^^^^^^^^^^^^ -The code is available in the ``torch-sim`` package, see the corresponding -`installation instructions `_. +Only the :ref:`energy ` output is supported. Forces and stresses +are derived via autograd. How to use the code ^^^^^^^^^^^^^^^^^^^ -You can find the documentation for metatomic models in torch-sim `here -`_, -and generic documentation on torch-sim `there -`_. +.. code-block:: python + + import ase.build + import torch_sim as ts + from metatomic_torchsim import MetatomicModel + + model = MetatomicModel("model.pt", device="cpu") + + atoms = ase.build.bulk("Si", "diamond", a=5.43, cubic=True) + sim_state = ts.initialize_state(atoms, device=model.device, dtype=model.dtype) + + results = model(sim_state) + print(results["energy"]) # shape [1] + print(results["forces"]) # shape [n_atoms, 3] + print(results["stress"]) # shape [1, 3, 3] + +API documentation +----------------- + +.. autoclass:: metatomic_torchsim.MetatomicModel + :show-inheritance: + :members: diff --git a/docs/src/index.rst b/docs/src/index.rst index d41b2501a..7f85a55aa 100644 --- a/docs/src/index.rst +++ b/docs/src/index.rst @@ -17,7 +17,7 @@ This library focuses on exporting and importing fully working, already trained models. If you want to train existing architectures with new data or re-use existing trained models, look into the metatrain_ project instead. -.. _metatrain: https://github.com/lab-cosmo/metatrain +.. _metatrain: https://github.com/metatensor/metatrain .. grid:: diff --git a/metatomic-torch/CHANGELOG.md b/metatomic-torch/CHANGELOG.md index 00947737a..45c02fd0f 100644 --- a/metatomic-torch/CHANGELOG.md +++ b/metatomic-torch/CHANGELOG.md @@ -5,7 +5,6 @@ a changelog](https://keepachangelog.com/en/1.1.0/) format. This project follows [Semantic Versioning](https://semver.org/spec/v2.0.0.html). diff --git a/python/metatomic_torchsim/MANIFEST.in b/python/metatomic_torchsim/MANIFEST.in new file mode 100644 index 000000000..576b01ef3 --- /dev/null +++ b/python/metatomic_torchsim/MANIFEST.in @@ -0,0 +1,2 @@ +include AUTHORS +include git_version_info diff --git a/python/metatomic_torchsim/README.md b/python/metatomic_torchsim/README.md new file mode 100644 index 000000000..275dbb411 --- /dev/null +++ b/python/metatomic_torchsim/README.md @@ -0,0 +1,33 @@ +# metatomic-torchsim + +TorchSim integration for metatomic atomistic models. + +Wraps metatomic models as TorchSim `ModelInterface` instances, enabling their +use in TorchSim molecular dynamics and other simulation workflows. + +## Installation + +```bash +pip install metatomic-torchsim +``` + +For universal potential models, see +[upet](https://github.com/lab-cosmo/upet). + +## Usage + +```python +from metatomic_torchsim import MetatomicModel + +# From a saved .pt model +model = MetatomicModel("model.pt", device="cuda") + +# Use with TorchSim +output = model(sim_state) +energy = output["energy"] +forces = output["forces"] +stress = output["stress"] +``` + +For full documentation, see the +[torch-sim engine page](https://docs.metatensor.org/metatomic/latest/engines/torch-sim.html). diff --git a/python/metatomic_torchsim/metatomic_torchsim/__init__.py b/python/metatomic_torchsim/metatomic_torchsim/__init__.py new file mode 100644 index 000000000..3950290f9 --- /dev/null +++ b/python/metatomic_torchsim/metatomic_torchsim/__init__.py @@ -0,0 +1,4 @@ +from ._model import MetatomicModel + + +__all__ = ["MetatomicModel"] diff --git a/python/metatomic_torchsim/metatomic_torchsim/_model.py b/python/metatomic_torchsim/metatomic_torchsim/_model.py new file mode 100644 index 000000000..6612c4612 --- /dev/null +++ b/python/metatomic_torchsim/metatomic_torchsim/_model.py @@ -0,0 +1,390 @@ +"""TorchSim wrapper for metatomic atomistic models. + +Adapts metatomic models to the TorchSim ModelInterface protocol, allowing them to +be used within the torch-sim simulation framework for MD and other simulations. + +Supports batched computations for multiple systems simultaneously, computing +energies, forces, and stresses via autograd. +""" + +import logging +import os +import pathlib +from typing import Dict, List, Optional, Union + +import torch +import vesin.metatomic +from metatensor.torch import Labels, TensorBlock + +from metatomic.torch import ( + AtomisticModel, + ModelEvaluationOptions, + ModelOutput, + NeighborListOptions, + System, + load_atomistic_model, + pick_device, +) + + +try: + from nvalchemiops.torch.neighbors import neighbor_list as nvalchemi_neighbor_list + + HAS_NVALCHEMIOPS = True +except ImportError: + HAS_NVALCHEMIOPS = False + + +try: + import torch_sim as ts + from torch_sim.models.interface import ModelInterface +except ImportError as e: + raise ImportError( + "the torch_sim package is required for metatomic-torchsim: " + "pip install torch-sim-atomistic" + ) from e + + +FilePath = Union[str, bytes, pathlib.PurePath] + +LOGGER = logging.getLogger(__name__) + +STR_TO_DTYPE = { + "float32": torch.float32, + "float64": torch.float64, +} + + +class MetatomicModel(ModelInterface): + """TorchSim wrapper for metatomic atomistic models. + + Wraps a metatomic model to compute energies, forces, and stresses within the + TorchSim framework. Handles the translation between TorchSim's batched + ``SimState`` and metatomic's list-of-``System`` convention, and uses autograd + for force/stress derivatives. + + Neighbor lists are computed with vesin, or with nvalchemiops on CUDA when + available and the model requests full neighbor lists. + """ + + def __init__( + self, + model: Union[FilePath, AtomisticModel, "torch.jit.RecursiveScriptModule"], + *, + extensions_directory: Optional[FilePath] = None, + device: Optional[Union[torch.device, str]] = None, + check_consistency: bool = False, + compute_forces: bool = True, + compute_stress: bool = True, + ) -> None: + """ + :param model: Model to use. Accepts a file path to a ``.pt`` saved + model, a Python :py:class:`AtomisticModel` instance, or a + TorchScript :py:class:`torch.jit.RecursiveScriptModule`. + :param extensions_directory: Directory containing compiled TorchScript + extensions required by the model, if any. + :param device: Torch device for evaluation. When ``None``, the best + device is selected from the model's ``supported_devices``. + :param check_consistency: Run consistency checks during model evaluation. + Useful for debugging but hurts performance. + :param compute_forces: Compute atomic forces via autograd. + :param compute_stress: Compute stress tensors via the strain trick. + """ + super().__init__() + + self._check_consistency = check_consistency + + # Load the model, following the same patterns as ase_calculator.py + if isinstance(model, (str, bytes, pathlib.PurePath)): + model_path = str(model) + if not os.path.exists(model_path): + raise ValueError(f"given model path '{model_path}' does not exist") + model = load_atomistic_model( + model_path, extensions_directory=extensions_directory + ) + elif isinstance(model, torch.jit.RecursiveScriptModule): + if model.original_name != "AtomisticModel": + raise TypeError( + "torch model must be 'AtomisticModel', " + f"got '{model.original_name}' instead" + ) + elif isinstance(model, AtomisticModel): + pass + else: + raise TypeError(f"unknown type for model: {type(model)}") + + capabilities = model.capabilities() + + # Resolve device + if device is not None: + if isinstance(device, str): + device = torch.device(device) + self._device = device + else: + self._device = torch.device( + pick_device(capabilities.supported_devices, None) + ) + + # Resolve dtype from model capabilities + if capabilities.dtype in STR_TO_DTYPE: + self._dtype = STR_TO_DTYPE[capabilities.dtype] + else: + raise ValueError( + f"unexpected dtype in model capabilities: {capabilities.dtype}" + ) + + if "energy" not in capabilities.outputs: + raise ValueError( + "model does not have an 'energy' output. " + "Only models with energy outputs can be used with TorchSim." + ) + + self._model = model.to(device=self._device) + self._compute_forces = compute_forces + self._compute_stress = compute_stress + + 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(self, state: "ts.SimState") -> Dict[str, torch.Tensor]: + """Compute energies, forces, and stresses for the given simulation state. + + :param state: TorchSim simulation state + + :returns: Dictionary with ``"energy"`` (shape ``[n_systems]``), + ``"forces"`` (shape ``[n_atoms, 3]``, if ``compute_forces``), and + ``"stress"`` (shape ``[n_systems, 3, 3]``, if ``compute_stress``). + """ + positions = state.positions + cell = state.row_vector_cell + atomic_nums = state.atomic_numbers + + if positions.dtype != self._dtype: + raise TypeError( + f"positions dtype {positions.dtype} does not match " + f"model dtype {self._dtype}" + ) + + # Build per-system System objects. Metatomic expects a list of System + # rather than a single batched graph. + systems: List[System] = [] + strains: List[torch.Tensor] = [] + n_systems = len(cell) + + for sys_idx in range(n_systems): + mask = state.system_idx == sys_idx + sys_positions = positions[mask] + sys_cell = cell[sys_idx] + sys_types = atomic_nums[mask] + + if self._compute_forces: + sys_positions = sys_positions.detach().requires_grad_(True) + + if self._compute_stress: + strain = torch.eye( + 3, + device=self._device, + dtype=self._dtype, + requires_grad=True, + ) + sys_positions = sys_positions @ strain + sys_cell = sys_cell @ strain + strains.append(strain) + + systems.append( + System( + positions=sys_positions, + types=sys_types, + cell=sys_cell, + pbc=state.pbc, + ) + ) + + # Compute neighbor lists + systems = _compute_requested_neighbors( + systems=systems, + requested_options=self._requested_neighbor_lists, + check_consistency=self._check_consistency, + ) + + # Run the model + model_outputs = self._model( + systems=systems, + options=self._evaluation_options, + check_consistency=self._check_consistency, + ) + + energy_values = model_outputs["energy"].block().values + + results: Dict[str, torch.Tensor] = {} + results["energy"] = energy_values.detach().squeeze(-1) + + # Compute forces and/or stresses via autograd + if self._compute_forces or self._compute_stress: + grad_inputs: List[torch.Tensor] = [] + if self._compute_forces: + for system in systems: + grad_inputs.append(system.positions) + if self._compute_stress: + grad_inputs.extend(strains) + + grads = torch.autograd.grad( + outputs=energy_values, + inputs=grad_inputs, + grad_outputs=torch.ones_like(energy_values), + ) + + if self._compute_forces and self._compute_stress: + n_sys = len(systems) + force_grads = grads[:n_sys] + stress_grads = grads[n_sys:] + elif self._compute_forces: + force_grads = grads + stress_grads = () + else: + force_grads = () + stress_grads = grads + + if self._compute_forces: + results["forces"] = torch.cat([-g for g in force_grads]) + + if self._compute_stress: + results["stress"] = torch.stack( + [ + g / torch.abs(torch.det(system.cell.detach())) + for g, system in zip(stress_grads, systems, strict=False) + ] + ) + + return results + + +# -- Neighbor list helpers (shared with ase_calculator.py patterns) ---------- + + +def _compute_requested_neighbors( + systems: List[System], + requested_options: List[NeighborListOptions], + check_consistency: bool = False, +) -> List[System]: + """Compute all neighbor lists requested by the model and store them in the systems. + + Uses nvalchemiops for full neighbor lists on CUDA when available, vesin otherwise. + """ + can_use_nvalchemi = HAS_NVALCHEMIOPS and all( + system.device.type == "cuda" for system in systems + ) + + if can_use_nvalchemi: + full_nl_options = [] + half_nl_options = [] + for options in requested_options: + if options.full_list: + full_nl_options.append(options) + else: + half_nl_options.append(options) + + systems = _compute_requested_neighbors_nvalchemi( + systems=systems, + requested_options=full_nl_options, + ) + systems = _compute_requested_neighbors_vesin( + systems=systems, + requested_options=half_nl_options, + check_consistency=check_consistency, + ) + else: + systems = _compute_requested_neighbors_vesin( + systems=systems, + requested_options=requested_options, + check_consistency=check_consistency, + ) + + return systems + + +def _compute_requested_neighbors_vesin( + systems: List[System], + requested_options: List[NeighborListOptions], + check_consistency: bool = False, +) -> List[System]: + """Compute neighbor lists using vesin.""" + system_devices = [] + moved_systems = [] + for system in systems: + system_devices.append(system.device) + if system.device.type not in ["cpu", "cuda"]: + moved_systems.append(system.to(device="cpu")) + else: + moved_systems.append(system) + + vesin.metatomic.compute_requested_neighbors_from_options( + systems=moved_systems, + system_length_unit="angstrom", + options=requested_options, + check_consistency=check_consistency, + ) + + systems = [] + for system, device in zip(moved_systems, system_devices, strict=True): + systems.append(system.to(device=device)) + + return systems + + +def _compute_requested_neighbors_nvalchemi( + systems: List[System], + requested_options: List[NeighborListOptions], +) -> List[System]: + """Compute full neighbor lists on CUDA using nvalchemiops.""" + for options in requested_options: + assert options.full_list + for system in systems: + assert system.device.type == "cuda" + + edge_index, _, S = nvalchemi_neighbor_list( + system.positions, + options.engine_cutoff("angstrom"), + cell=system.cell, + pbc=system.pbc, + return_neighbor_list=True, + ) + D = ( + system.positions[edge_index[1]] + - system.positions[edge_index[0]] + + S.to(system.cell.dtype) @ system.cell + ) + P = edge_index.T + + neighbors = TensorBlock( + D.reshape(-1, 3, 1), + samples=Labels( + names=[ + "first_atom", + "second_atom", + "cell_shift_a", + "cell_shift_b", + "cell_shift_c", + ], + values=torch.hstack([P, S]), + ), + components=[ + Labels( + "xyz", + torch.tensor([[0], [1], [2]], device=system.device), + ) + ], + properties=Labels( + "distance", + torch.tensor([[0]], device=system.device), + ), + ) + system.add_neighbor_list(options, neighbors) + + return systems diff --git a/python/metatomic_torchsim/pyproject.toml b/python/metatomic_torchsim/pyproject.toml new file mode 100644 index 000000000..a1252ef88 --- /dev/null +++ b/python/metatomic_torchsim/pyproject.toml @@ -0,0 +1,61 @@ +[project] +name = "metatomic-torchsim" +dynamic = ["version", "authors"] +requires-python = ">=3.10" + +readme = "README.md" +license = "BSD-3-Clause" +description = "TorchSim integration for metatomic models" + +keywords = ["machine learning", "molecular modeling", "torch", "torchsim"] +classifiers = [ + "Development Status :: 3 - Alpha", + "Intended Audience :: Science/Research", + "Operating System :: POSIX", + "Operating System :: MacOS :: MacOS X", + "Operating System :: Microsoft :: Windows", + "Programming Language :: Python", + "Programming Language :: Python :: 3", + "Topic :: Scientific/Engineering", + "Topic :: Scientific/Engineering :: Chemistry", + "Topic :: Scientific/Engineering :: Physics", + "Topic :: Software Development :: Libraries", + "Topic :: Software Development :: Libraries :: Python Modules", +] + +dependencies = [ + "metatomic-torch >=0.1.11,<0.2", + "torch-sim-atomistic >=0.5,<0.6", +] + +[project.urls] +homepage = "https://docs.metatensor.org/metatomic/latest/engines/torch-sim.html" +documentation = "https://docs.metatensor.org/metatomic/latest/engines/torch-sim.html" +repository = "https://github.com/metatensor/metatomic/tree/main/python/metatomic_torchsim" + +### ======================================================================== ### + +[build-system] +requires = [ + "setuptools >=77", + "packaging >=23", +] +build-backend = "setuptools.build_meta" + +### ======================================================================== ### + +[tool.pytest.ini_options] +python_files = ["*.py"] +testpaths = ["tests"] +filterwarnings = [ + "error", + # TorchScript deprecation warnings + "ignore:`torch.jit.script` is deprecated. Please switch to `torch.compile` or `torch.export`:DeprecationWarning", + "ignore:`torch.jit.script_method` is deprecated. Please switch to `torch.compile` or `torch.export`:DeprecationWarning", + "ignore:`torch.jit.save` is deprecated. Please switch to `torch.export`:DeprecationWarning", + "ignore:`torch.jit.load` is deprecated. Please switch to `torch.export`:DeprecationWarning", + # There is a circular dependency between metatomic-torch and vesin.metatomic + "ignore:.*vesin.metatomic was only tested with metatomic.torch >=0.1.3,<0.2.*:UserWarning", + # This comes from inside TorchSim + "ignore:The 'nvalchemiops.neighborlist' module has been renamed to 'nvalchemiops.neighbors':DeprecationWarning", +] diff --git a/python/metatomic_torchsim/setup.py b/python/metatomic_torchsim/setup.py new file mode 100644 index 000000000..9d542f8d1 --- /dev/null +++ b/python/metatomic_torchsim/setup.py @@ -0,0 +1,114 @@ +import os +import subprocess +import sys + +import packaging.version +from setuptools import setup +from setuptools.command.sdist import sdist + + +ROOT = os.path.realpath(os.path.dirname(__file__)) + +METATOMIC_TORCHSIM_VERSION = "0.1.0" + + +class sdist_generate_data(sdist): + """ + Create a sdist with an additional generated files: + - `git_version_info` + - `metatomic-torch-cxx-*.tar.gz` + """ + + def run(self): + n_commits, git_hash = git_version_info() + with open("git_version_info", "w") as fd: + fd.write(f"{n_commits}\n{git_hash}\n") + + # run original sdist + super().run() + + os.unlink("git_version_info") + + +def git_version_info(): + """ + If git is available and we are building from a checkout, get the number of commits + since the last tag & full hash of the code. Otherwise, this always returns (0, ""). + """ + TAG_PREFIX = "metatomic-torchsim-v" + + if os.path.exists("git_version_info"): + # we are building from a sdist, without git available, but the git + # version was recorded in the `git_version_info` file + with open("git_version_info") as fd: + n_commits = int(fd.readline().strip()) + git_hash = fd.readline().strip() + else: + script = os.path.join(ROOT, "..", "..", "scripts", "git-version-info.py") + assert os.path.exists(script) + + output = subprocess.run( + [sys.executable, script, TAG_PREFIX], + stderr=subprocess.PIPE, + stdout=subprocess.PIPE, + encoding="utf8", + ) + + if output.returncode != 0: + raise Exception( + "failed to get git version info.\n" + f"stdout: {output.stdout}\n" + f"stderr: {output.stderr}\n" + ) + elif output.stderr: + print(output.stderr, file=sys.stderr) + n_commits = 0 + git_hash = "" + else: + lines = output.stdout.splitlines() + n_commits = int(lines[0].strip()) + git_hash = lines[1].strip() + + return n_commits, git_hash + + +def create_version_number(version): + version = packaging.version.parse(version) + + n_commits, git_hash = git_version_info() + + if n_commits != 0: + # if we have commits since the last tag, this mean we are in a pre-release of + # the next version. So we increase either the minor version number or the + # release candidate number (if we are closing up on a release) + if version.pre is not None: + assert version.pre[0] == "rc" + pre = ("rc", version.pre[1] + 1) + release = version.release + else: + major, minor, patch = version.release + release = (major, minor + 1, 0) + pre = None + + # this is using a private API which is intended to become public soon: + # https://github.com/pypa/packaging/pull/698. In the mean time we'll + # use this + version._version = version._version._replace(release=release) + version._version = version._version._replace(pre=pre) + version._version = version._version._replace(dev=("dev", n_commits)) + version._version = version._version._replace(local=(git_hash,)) + + return str(version) + + +if __name__ == "__main__": + with open(os.path.join(ROOT, "AUTHORS")) as fd: + authors = fd.read().splitlines() + + setup( + version=create_version_number(METATOMIC_TORCHSIM_VERSION), + author=", ".join(authors), + cmdclass={ + "sdist": sdist_generate_data, + }, + ) diff --git a/python/metatomic_torchsim/tests/__init__.py b/python/metatomic_torchsim/tests/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/python/metatomic_torchsim/tests/model_loading.py b/python/metatomic_torchsim/tests/model_loading.py new file mode 100644 index 000000000..503f9111d --- /dev/null +++ b/python/metatomic_torchsim/tests/model_loading.py @@ -0,0 +1,56 @@ +"""Tests for MetatomicModel loading paths.""" + +import pytest +import torch + +import metatomic_lj_test +from metatomic_torchsim import MetatomicModel + + +DEVICE = torch.device("cpu") + + +@pytest.fixture +def lj_model(): + return metatomic_lj_test.lennard_jones_model( + atomic_type=28, + cutoff=5.0, + sigma=1.5808, + epsilon=0.1729, + length_unit="Angstrom", + energy_unit="eV", + with_extension=False, + ) + + +def test_load_from_pt_file(lj_model, tmp_path): + """Model loads from a saved .pt file.""" + pt_path = tmp_path / "test_model.pt" + lj_model.save(str(pt_path)) + + model = MetatomicModel(model=str(pt_path), device=DEVICE) + assert model.device == DEVICE + + +def test_nonexistent_path_raises_valueerror(): + """ValueError raised for a path that does not exist.""" + with pytest.raises(ValueError, match="does not exist"): + MetatomicModel(model="/non/existent/path.pt", device=DEVICE) + + +def test_wrong_model_type_raises_typeerror(): + """TypeError raised when passing an unsupported type.""" + with pytest.raises(TypeError, match="unknown type for model"): + MetatomicModel(model=42, device=DEVICE) + + +def test_non_atomisticmodel_scriptmodule_raises_typeerror(): + """TypeError raised for a ScriptModule that is not AtomisticModel.""" + + class Dummy(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x + + dummy_scripted = torch.jit.script(Dummy()) + with pytest.raises(TypeError, match="must be 'AtomisticModel'"): + MetatomicModel(model=dummy_scripted, device=DEVICE) diff --git a/python/metatomic_torchsim/tests/torchsim.py b/python/metatomic_torchsim/tests/torchsim.py new file mode 100644 index 000000000..83428dd93 --- /dev/null +++ b/python/metatomic_torchsim/tests/torchsim.py @@ -0,0 +1,281 @@ +"""Tests for the MetatomicModel TorchSim wrapper. + +Uses the metatomic-lj-test model so that tests run without +downloading large model files. +""" + +import numpy as np +import pytest +import torch +import torch_sim as ts + +import metatomic_lj_test +from metatomic_torchsim import MetatomicModel + + +CUTOFF = 5.0 +SIGMA = 1.5808 +EPSILON = 0.1729 + +DEVICE = torch.device("cpu") +DTYPE = torch.float64 + + +@pytest.fixture +def lj_model(): + return metatomic_lj_test.lennard_jones_model( + atomic_type=28, + cutoff=CUTOFF, + sigma=SIGMA, + epsilon=EPSILON, + length_unit="Angstrom", + energy_unit="eV", + with_extension=False, + ) + + +@pytest.fixture +def ni_atoms(): + """Create a small perturbed Ni FCC supercell.""" + import ase.build + + np.random.seed(0xDEADBEEF) + atoms = ase.build.make_supercell( + ase.build.bulk("Ni", "fcc", a=3.6, cubic=True), 2 * np.eye(3) + ) + atoms.positions += 0.2 * np.random.rand(*atoms.positions.shape) + return atoms + + +@pytest.fixture +def metatomic_model(lj_model): + return MetatomicModel(model=lj_model, device=DEVICE) + + +def test_initialization(lj_model): + """MetatomicModel initializes with correct device and dtype.""" + model = MetatomicModel(model=lj_model, device=DEVICE) + assert model.device == DEVICE + assert model.dtype == DTYPE + assert model.compute_forces is True + assert model.compute_stress is True + + +def test_initialization_no_forces(lj_model): + """Can disable force computation.""" + model = MetatomicModel(model=lj_model, device=DEVICE, compute_forces=False) + assert model.compute_forces is False + assert model.compute_stress is True + + +def test_forward_returns_energy(metatomic_model, ni_atoms): + """Forward pass returns energy with correct shape.""" + sim_state = ts.io.atoms_to_state([ni_atoms], DEVICE, DTYPE) + output = metatomic_model(sim_state) + + assert "energy" in output + assert output["energy"].shape == (1,) + assert output["energy"].dtype == DTYPE + + +def test_forward_returns_forces(metatomic_model, ni_atoms): + """Forward pass returns forces with correct shape.""" + sim_state = ts.io.atoms_to_state([ni_atoms], DEVICE, DTYPE) + output = metatomic_model(sim_state) + + assert "forces" in output + n_atoms = len(ni_atoms) + assert output["forces"].shape == (n_atoms, 3) + assert output["forces"].dtype == DTYPE + + +def test_forward_returns_stress(metatomic_model, ni_atoms): + """Forward pass returns stress with correct shape.""" + sim_state = ts.io.atoms_to_state([ni_atoms], DEVICE, DTYPE) + output = metatomic_model(sim_state) + + assert "stress" in output + assert output["stress"].shape == (1, 3, 3) + assert output["stress"].dtype == DTYPE + + +def test_forward_no_stress(lj_model, ni_atoms): + """Stress is not returned when compute_stress=False.""" + model = MetatomicModel(model=lj_model, device=DEVICE, compute_stress=False) + sim_state = ts.io.atoms_to_state([ni_atoms], DEVICE, DTYPE) + output = model(sim_state) + + assert "energy" in output + assert "forces" in output + assert "stress" not in output + + +def test_forward_no_forces(lj_model, ni_atoms): + """Forces are not returned when compute_forces=False.""" + model = MetatomicModel(model=lj_model, device=DEVICE, compute_forces=False) + sim_state = ts.io.atoms_to_state([ni_atoms], DEVICE, DTYPE) + output = model(sim_state) + + assert "energy" in output + assert "forces" not in output + assert "stress" in output + + +@pytest.fixture +def ni_atoms_2(): + """Create a second Ni supercell (same size, different lattice parameter).""" + import ase.build + + np.random.seed(0xCAFEBABE) + atoms = ase.build.make_supercell( + ase.build.bulk("Ni", "fcc", a=3.5, cubic=True), 2 * np.eye(3) + ) + atoms.positions += 0.1 * np.random.rand(*atoms.positions.shape) + return atoms + + +def test_batched_forward(metatomic_model, ni_atoms, ni_atoms_2): + """Forward pass handles batched systems correctly.""" + sim_state = ts.io.atoms_to_state([ni_atoms, ni_atoms_2], DEVICE, DTYPE) + output = metatomic_model(sim_state) + + assert output["energy"].shape == (2,) + n_total = len(ni_atoms) + len(ni_atoms_2) + assert output["forces"].shape == (n_total, 3) + assert output["stress"].shape == (2, 3, 3) + + +def test_energy_consistency_single_vs_batch(metatomic_model, ni_atoms, ni_atoms_2): + """Energy from single system matches the corresponding entry in a batch.""" + + # single + state_1 = ts.io.atoms_to_state([ni_atoms], DEVICE, DTYPE) + out_1 = metatomic_model(state_1) + + state_2 = ts.io.atoms_to_state([ni_atoms_2], DEVICE, DTYPE) + out_2 = metatomic_model(state_2) + + # batch + state_batch = ts.io.atoms_to_state([ni_atoms, ni_atoms_2], DEVICE, DTYPE) + out_batch = metatomic_model(state_batch) + + torch.testing.assert_close(out_1["energy"], out_batch["energy"][:1]) + torch.testing.assert_close(out_2["energy"], out_batch["energy"][1:]) + + +def test_forces_sum_to_zero(metatomic_model, ni_atoms): + """Net force on the system should be approximately zero (Newton's 3rd law).""" + sim_state = ts.io.atoms_to_state([ni_atoms], DEVICE, DTYPE) + output = metatomic_model(sim_state) + + net_force = output["forces"].sum(dim=0) + torch.testing.assert_close( + net_force, torch.zeros(3, dtype=DTYPE), atol=1e-6, rtol=0 + ) + + +def test_validate_model_outputs(metatomic_model): + """Model passes TorchSim's validate_model_outputs check.""" + try: + from torch_sim.models.interface import validate_model_outputs + except ImportError: + pytest.skip("validate_model_outputs not available in this torch-sim version") + + # validate_model_outputs creates its own test systems (Si diamond + Fe FCC). + # Our LJ model only knows atomic_type=28 (Ni), but the validator uses Si (14) + # and Fe (26). So we skip if the validator would fail for type reasons. + try: + validate_model_outputs(metatomic_model, DEVICE, DTYPE) + except Exception as exc: + if "atomic type" in str(exc).lower() or "species" in str(exc).lower(): + pytest.skip(f"LJ test model does not support Si/Fe types: {exc}") + raise + + +def test_wrong_dtype_raises(metatomic_model, ni_atoms): + """TypeError raised when positions have wrong dtype.""" + sim_state = ts.io.atoms_to_state([ni_atoms], DEVICE, torch.float32) + with pytest.raises(TypeError, match="dtype"): + metatomic_model(sim_state) + + +def test_single_atom_system(lj_model): + """Model handles a single-atom system.""" + import ase + + atoms = ase.Atoms( + symbols=["Ni"], + positions=[[0.0, 0.0, 0.0]], + cell=[10.0, 10.0, 10.0], + pbc=True, + ) + model = MetatomicModel(model=lj_model, device=DEVICE) + sim_state = ts.io.atoms_to_state([atoms], DEVICE, DTYPE) + output = model(sim_state) + + assert output["energy"].shape == (1,) + assert output["forces"].shape == (1, 3) + assert output["stress"].shape == (1, 3, 3) + + +def test_energy_only_mode(lj_model, ni_atoms): + """Model returns only energy when forces and stress are disabled.""" + model = MetatomicModel( + model=lj_model, device=DEVICE, compute_forces=False, compute_stress=False + ) + sim_state = ts.io.atoms_to_state([ni_atoms], DEVICE, DTYPE) + output = model(sim_state) + + assert "energy" in output + assert "forces" not in output + assert "stress" not in output + + +def test_check_consistency_mode(lj_model, ni_atoms): + """Model runs with consistency checking enabled.""" + model = MetatomicModel(model=lj_model, device=DEVICE, check_consistency=True) + sim_state = ts.io.atoms_to_state([ni_atoms], DEVICE, DTYPE) + output = model(sim_state) + + assert "energy" in output + assert "forces" in output + assert "stress" in output + + +def test_forces_match_finite_difference(lj_model, ni_atoms): + """Autograd forces match finite-difference gradient of energy.""" + delta = 1e-4 + model = MetatomicModel(model=lj_model, device=DEVICE, compute_stress=False) + sim_state = ts.io.atoms_to_state([ni_atoms], DEVICE, DTYPE) + output = model(sim_state) + autograd_forces = output["forces"] + + for i in range(3): + for j in range(3): + atoms_plus = ni_atoms.copy() + atoms_minus = ni_atoms.copy() + atoms_plus.positions[i, j] += delta + atoms_minus.positions[i, j] -= delta + + state_plus = ts.io.atoms_to_state([atoms_plus], DEVICE, DTYPE) + state_minus = ts.io.atoms_to_state([atoms_minus], DEVICE, DTYPE) + + e_plus = model(state_plus)["energy"][0] + e_minus = model(state_minus)["energy"][0] + + numerical_force = -(e_plus - e_minus) / (2 * delta) + torch.testing.assert_close( + autograd_forces[i, j], + numerical_force, + atol=1e-4, + rtol=0, + ) + + +def test_stress_is_symmetric(metatomic_model, ni_atoms): + """Stress tensor is symmetric.""" + sim_state = ts.io.atoms_to_state([ni_atoms], DEVICE, DTYPE) + output = metatomic_model(sim_state) + stress = output["stress"] + + torch.testing.assert_close(stress, stress.transpose(-2, -1), atol=1e-10, rtol=0) diff --git a/setup.py b/setup.py index 2aa9f2beb..b38f1ba0e 100644 --- a/setup.py +++ b/setup.py @@ -13,6 +13,8 @@ # when packaging a sdist for release, we should never use local dependencies METATOMIC_NO_LOCAL_DEPS = os.environ.get("METATOMIC_NO_LOCAL_DEPS", "0") == "1" + METATOMIC_TORCHSIM = os.path.join(ROOT, "python", "metatomic_torchsim") + if not METATOMIC_NO_LOCAL_DEPS and os.path.exists(METATOMIC_TORCH): # we are building from a git checkout extras_require["torch"] = f"metatomic-torch @ file://{METATOMIC_TORCH}" @@ -20,6 +22,11 @@ # we are building from a sdist/installing from a wheel extras_require["torch"] = "metatomic-torch" + if not METATOMIC_NO_LOCAL_DEPS and os.path.exists(METATOMIC_TORCHSIM): + extras_require["torchsim"] = f"metatomic-torchsim @ file://{METATOMIC_TORCHSIM}" + else: + extras_require["torchsim"] = "metatomic-torchsim" + setup( author=", ".join(open(os.path.join(ROOT, "AUTHORS")).read().splitlines()), extras_require=extras_require, diff --git a/tox.ini b/tox.ini index 3d2fbdabf..cdcd65179 100644 --- a/tox.ini +++ b/tox.ini @@ -10,7 +10,7 @@ envlist = torch-tests-cxx torch-install-tests-cxx docs-tests - + torchsim-tests [testenv] passenv = * @@ -164,6 +164,49 @@ commands = pytest --cov={env_site_packages_dir}/metatomic --cov-report= --import-mode=append {posargs} +[testenv:build-metatomic-torchsim] +passenv = * +setenv = + PYTHONPATH= + +description = + Build the metatomic-torchsim wheel for testing +deps = + {[testenv]metatensor_deps} + torch=={env:METATOMIC_TESTS_TORCH_VERSION:2.10}.* + +commands = + pip wheel python/metatomic_torchsim {[testenv]build_single_wheel} --wheel-dir {envtmpdir}/dist + + +[testenv:torchsim-tests] +description = Run the tests of the metatomic-torchsim Python package +package = external +package_env = build-metatomic-torch +deps = + {[testenv]testing_deps} + {[testenv]metatensor_deps} + torch=={env:METATOMIC_TESTS_TORCH_VERSION:2.10}.* + numpy {env:METATOMIC_TESTS_NUMPY_VERSION_PIN} + vesin + ase + torch-sim-atomistic + # for metatensor-lj-test + setuptools-scm + cmake + +changedir = python/metatomic_torchsim +commands = + # install metatomic-torchsim + pip install {[testenv]build_single_wheel} {toxinidir}/python/metatomic_torchsim + + # use the reference LJ implementation for tests + pip install {[testenv]build_single_wheel} git+https://github.com/metatensor/lj-test@f7401a8 + + pytest --cov={env_site_packages_dir}/metatomic_torchsim --cov-report= --import-mode=append {posargs} + + + [testenv:docs-tests] description = Run the doctests defined in any metatomic package deps = @@ -229,10 +272,13 @@ deps = # required for autodoc torch=={env:METATOMIC_TESTS_TORCH_VERSION:2.10}.* + torch-sim-atomistic >=0.5,<0.6 # required for examples ase chemiscope commands = + pip install {[testenv]build_single_wheel} {toxinidir}/python/metatomic_torchsim + sphinx-build -d docs/build/doctrees -W -b html docs/src docs/build/html