diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS new file mode 100644 index 00000000..611a1c95 --- /dev/null +++ b/.github/CODEOWNERS @@ -0,0 +1 @@ +/python/metatomic_torchsim/ @HaoZeke @CompRhys diff --git a/.github/workflows/torch-tests.yml b/.github/workflows/torch-tests.yml index 4aa30a9a..32ef1d82 100644 --- a/.github/workflows/torch-tests.yml +++ b/.github/workflows/torch-tests.yml @@ -70,6 +70,14 @@ jobs: METATOMIC_TESTS_TORCH_VERSION: ${{ matrix.torch-version }} METATOMIC_TESTS_NUMPY_VERSION_PIN: ${{ matrix.numpy-version-pin }} + - name: run metatomic-torchsim tests + if: matrix.python-version != '3.10' + run: tox -e torchsim-tests + env: + PIP_EXTRA_INDEX_URL: https://download.pytorch.org/whl/cpu + METATOMIC_TESTS_TORCH_VERSION: ${{ matrix.torch-version }} + METATOMIC_TESTS_NUMPY_VERSION_PIN: ${{ matrix.numpy-version-pin }} + - name: combine Python coverage files shell: bash run: | diff --git a/docs/src/engines/torch-sim.rst b/docs/src/engines/torch-sim.rst index 77b4edec..e7f4dfcf 100644 --- a/docs/src/engines/torch-sim.rst +++ b/docs/src/engines/torch-sim.rst @@ -8,24 +8,47 @@ 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 + +This pulls in ``torch-sim-atomistic`` and ``metatomic-torch`` as dependencies. + +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.io.atoms_to_state([atoms], model.device, 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] + +For more details, see the `metatomic-torchsim documentation +`_. diff --git a/pyproject.toml b/pyproject.toml index 1a67cd6a..f14c3ff1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -88,4 +88,5 @@ docstring-code-format = true [tool.uv.pip] reinstall-package = [ "metatomic-torch", + "metatomic-torchsim", ] diff --git a/python/metatomic_torch/metatomic/torch/model.py b/python/metatomic_torch/metatomic/torch/model.py index c2a9f528..b1af105f 100644 --- a/python/metatomic_torch/metatomic/torch/model.py +++ b/python/metatomic_torch/metatomic/torch/model.py @@ -549,7 +549,9 @@ def export(self, file: str, collect_extensions: Optional[str] = None): ) return self.save(file, collect_extensions) - def save(self, file: Union[str, Path], collect_extensions: Optional[str] = None): + def save( + self, file: Union[str, Path], collect_extensions: Optional[str] = None, **kwargs + ): """Save this model to a file that can then be loaded by simulation engine. The model will be saved with `requires_grad=False` for all parameters. diff --git a/python/metatomic_torchsim/AUTHORS b/python/metatomic_torchsim/AUTHORS new file mode 100644 index 00000000..0cbfac93 --- /dev/null +++ b/python/metatomic_torchsim/AUTHORS @@ -0,0 +1,4 @@ +Rhys Goodall +Guillaume Fraux +Filippo Bigi +Rohit Goswami diff --git a/python/metatomic_torchsim/CHANGELOG.md b/python/metatomic_torchsim/CHANGELOG.md new file mode 100644 index 00000000..f30a60fd --- /dev/null +++ b/python/metatomic_torchsim/CHANGELOG.md @@ -0,0 +1,5 @@ +# Changelog + +All notable changes to metatomic-torchsim will be documented in this file. + + diff --git a/python/metatomic_torchsim/MANIFEST.in b/python/metatomic_torchsim/MANIFEST.in new file mode 100644 index 00000000..12c08a00 --- /dev/null +++ b/python/metatomic_torchsim/MANIFEST.in @@ -0,0 +1,2 @@ +include pyproject.toml +include AUTHORS diff --git a/python/metatomic_torchsim/README.md b/python/metatomic_torchsim/README.md new file mode 100644 index 00000000..ce29a94d --- /dev/null +++ b/python/metatomic_torchsim/README.md @@ -0,0 +1,39 @@ +# 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 +``` + +To use metatrain checkpoints (`.ckpt` files) or the `pet-mad` shortcut: + +```bash +pip install metatomic-torchsim[metatrain] +``` + +## Usage + +```python +from metatomic.torchsim import MetatomicModel + +# From a saved .pt model +model = MetatomicModel("model.pt", device="cuda") + +# From a metatrain checkpoint (requires metatrain extra) +model = MetatomicModel("model.ckpt", device="cuda") + +# PET-MAD shortcut (requires metatrain extra) +model = MetatomicModel("pet-mad", device="cuda") + +# Use with TorchSim +output = model(sim_state) +energy = output["energy"] +forces = output["forces"] +stress = output["stress"] +``` diff --git a/python/metatomic_torchsim/docs/.gitignore b/python/metatomic_torchsim/docs/.gitignore new file mode 100644 index 00000000..69fa449d --- /dev/null +++ b/python/metatomic_torchsim/docs/.gitignore @@ -0,0 +1 @@ +_build/ diff --git a/python/metatomic_torchsim/docs/_static/.gitkeep b/python/metatomic_torchsim/docs/_static/.gitkeep new file mode 100644 index 00000000..e69de29b diff --git a/python/metatomic_torchsim/docs/_templates/partials/site-foot.html b/python/metatomic_torchsim/docs/_templates/partials/site-foot.html new file mode 100644 index 00000000..bbc7e8b4 --- /dev/null +++ b/python/metatomic_torchsim/docs/_templates/partials/site-foot.html @@ -0,0 +1,4 @@ +{% include "components/foot-copyright.html" %} +

+ Part of the Metatensor ecosystem. +

diff --git a/python/metatomic_torchsim/docs/changelog.md b/python/metatomic_torchsim/docs/changelog.md new file mode 100644 index 00000000..98320037 --- /dev/null +++ b/python/metatomic_torchsim/docs/changelog.md @@ -0,0 +1,5 @@ +# Changelog + +```{include} ../CHANGELOG.md +:start-after: "# Changelog" +``` diff --git a/python/metatomic_torchsim/docs/conf.py b/python/metatomic_torchsim/docs/conf.py new file mode 100644 index 00000000..0ca6c5b0 --- /dev/null +++ b/python/metatomic_torchsim/docs/conf.py @@ -0,0 +1,127 @@ +"""Sphinx configuration for metatomic-torchsim documentation.""" + +from datetime import datetime + + +project = "metatomic-torchsim" +author = "the metatomic developers" +copyright = f"{datetime.now().date().year}, {author}" + +# -- General configuration --------------------------------------------------- + +extensions = [ + "myst_parser", + "sphinx.ext.autodoc", + "sphinx.ext.viewcode", + "sphinx.ext.intersphinx", + "autoapi.extension", +] + +templates_path = ["_templates"] +exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"] +suppress_warnings = ["autoapi.python_import_resolution"] + +# -- AutoAPI ----------------------------------------------------------------- + +autoapi_dirs = ["../metatomic/torchsim"] +autoapi_options = [ + "members", + "undoc-members", + "show-inheritance", + "show-module-summary", + "special-members", + "imported-members", +] +autoapi_python_class_content = "both" +autoapi_member_order = "bysource" +autoapi_keep_files = False +autoapi_add_toctree_entry = False + +# -- MyST -------------------------------------------------------------------- + +myst_enable_extensions = ["colon_fence", "fieldlist"] + +# -- Autodoc mocking --------------------------------------------------------- + +autodoc_mock_imports = [ + "torch", + "torch_sim", + "metatensor", + "metatomic", + "vesin", + "nvalchemiops", +] + +# -- Intersphinx ------------------------------------------------------------- + +intersphinx_mapping = { + "python": ("https://docs.python.org/3", None), + "torch": ("https://docs.pytorch.org/docs/stable/", None), + "metatomic": ("https://docs.metatensor.org/metatomic/latest/", None), + "metatensor": ("https://docs.metatensor.org/latest/", None), + "ase": ("https://ase-lib.org/", None), +} + +# -- HTML output ------------------------------------------------------------- + +html_theme = "shibuya" +html_static_path = ["_static"] + +html_context = { + "source_type": "github", + "source_user": "metatensor", + "source_repo": "metatomic", + "source_version": "main", + "source_docs_path": "/python/metatomic_torchsim/docs/", +} + +html_theme_options = { + "github_url": "https://github.com/metatensor/metatomic", + "accent_color": "teal", + "dark_code": True, + "globaltoc_expand_depth": 1, + "nav_links": [ + { + "title": "Ecosystem", + "children": [ + { + "title": "metatensor", + "url": "https://docs.metatensor.org", + "summary": "Data storage for atomistic ML", + "external": True, + }, + { + "title": "metatomic", + "url": "https://docs.metatensor.org/metatomic/latest/", + "summary": "Atomistic models with metatensor", + "external": True, + }, + { + "title": "metatrain", + "url": "https://metatensor.github.io/metatrain/latest/", + "summary": "Training framework for atomistic ML", + "external": True, + }, + { + "title": "torch-sim", + "url": "https://torchsim.github.io/torch-sim/", + "summary": "Differentiable molecular dynamics in PyTorch", + "external": True, + }, + ], + }, + { + "title": "PyPI", + "url": "https://pypi.org/project/metatomic-torchsim/", + "external": True, + }, + ], +} + +html_sidebars = { + "**": [ + "sidebars/localtoc.html", + "sidebars/repo-stats.html", + "sidebars/edit-this-page.html", + ], +} diff --git a/python/metatomic_torchsim/docs/explanation/architecture.md b/python/metatomic_torchsim/docs/explanation/architecture.md new file mode 100644 index 00000000..900a10a0 --- /dev/null +++ b/python/metatomic_torchsim/docs/explanation/architecture.md @@ -0,0 +1,77 @@ +# Architecture + +This page explains how `MetatomicModel` bridges TorchSim and metatomic. + +## SimState vs list of System + +TorchSim represents a simulation as a single batched `SimState` containing all +atoms from all systems, with a `system_idx` tensor tracking ownership. +Metatomic expects a `list[System]` where each `System` holds one periodic +structure. + +`MetatomicModel.forward` converts between these representations: + +1. Split the batched positions and atomic numbers by `system_idx` +2. Create one `System` per sub-structure with its own cell +3. Call the model on the list of systems +4. Concatenate results back into batched tensors + +## Forces via autograd + +Metatomic models typically output only total energies. Forces are computed as +the negative gradient of the energy with respect to atomic positions: + +``` +F_i = -dE/dr_i +``` + +Before calling the model, each system's positions are detached and set to +`requires_grad_(True)`. After the forward pass, `torch.autograd.grad` computes +the derivatives. + +## Stress via the strain trick + +Stress is computed using the Knuth strain trick. An identity strain tensor +(3x3, `requires_grad=True`) is applied to both positions and cell vectors: + +``` +r' = r @ strain +h' = h @ strain +``` + +The stress per system is then: + +``` +sigma = (1/V) * dE/d(strain) +``` + +where V is the cell volume. This gives the full 3x3 stress tensor without +finite differences. + +## Neighbor lists + +Models specify what neighbor lists they need via +`model.requested_neighbor_lists()`, which returns a list of +`NeighborListOptions` (cutoff radius, full vs half list). + +The wrapper computes these using: + +- **vesin**: Default backend for both CPU and GPU. Handles half and full + neighbor lists. Systems on non-CPU/CUDA devices are temporarily moved to CPU + for the computation. +- **nvalchemiops**: Used automatically on CUDA for full neighbor lists when + installed. Keeps everything on GPU, avoiding host-device transfers. + +The decision happens per-call in `_compute_requested_neighbors`: if all systems +are on CUDA and nvalchemiops is available, full-list requests go through +nvalchemi while half-list requests still use vesin. + +## Why a separate package + +metatomic-torchsim has its own versioning, release schedule, and dependency set +(`torch-sim-atomistic`). Keeping it separate from metatomic-torch avoids +forcing a torch-sim dependency on users who only need the ASE calculator or +other integrations. + +The package is pure Python with no compiled extensions, making it lightweight +to install. diff --git a/python/metatomic_torchsim/docs/howto/batched_simulations.md b/python/metatomic_torchsim/docs/howto/batched_simulations.md new file mode 100644 index 00000000..a0ef542b --- /dev/null +++ b/python/metatomic_torchsim/docs/howto/batched_simulations.md @@ -0,0 +1,68 @@ +# Batched simulations + +TorchSim supports batching multiple systems into a single `SimState` for +efficient parallel evaluation on GPU. `MetatomicModel` handles this +transparently. + +## Creating a batched state + +Pass a list of ASE `Atoms` objects to `atoms_to_state`: + +```python +import ase.build +import torch_sim as ts +from metatomic.torchsim import MetatomicModel + +model = MetatomicModel("model.pt", device="cpu") + +atoms_list = [ + ase.build.bulk("Cu", "fcc", a=3.6, cubic=True), + ase.build.bulk("Ni", "fcc", a=3.52, cubic=True), + ase.build.bulk("Al", "fcc", a=4.05, cubic=True), +] + +sim_state = ts.io.atoms_to_state(atoms_list, model.device, model.dtype) +``` + +## Evaluating the batch + +A single forward call evaluates all systems: + +```python +results = model(sim_state) +``` + +The output shapes reflect the batch: + +- `results["energy"]` has shape `[3]` (one energy per system) +- `results["forces"]` has shape `[n_total_atoms, 3]` (all atoms concatenated) +- `results["stress"]` has shape `[3, 3, 3]` (one 3x3 tensor per system) + +## How system_idx works + +`SimState` tracks which atom belongs to which system via the `system_idx` +tensor. For three 4-atom systems, `system_idx` looks like: + +``` +[0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2] +``` + +`MetatomicModel.forward` uses this to split the batched positions and types +into per-system `System` objects before calling the underlying model. + +## Batch consistency + +Energies computed in a batch match those computed individually. This is +guaranteed because each system gets its own neighbor list and independent +evaluation. The existing test `test_energy_consistency_single_vs_batch` +validates this property. + +## Performance considerations + +Batching is most beneficial on GPU, where the neighbor list computation and +model forward pass can run in parallel across systems. On CPU, the speedup +comes from reduced Python overhead (one call instead of N). + +For very large systems or many small ones, adjust the batch size to fit in GPU +memory. TorchSim does not impose a maximum batch size, but each system gets its +own neighbor list, so memory scales with the sum of per-system sizes. diff --git a/python/metatomic_torchsim/docs/howto/model_loading.md b/python/metatomic_torchsim/docs/howto/model_loading.md new file mode 100644 index 00000000..b739c6e3 --- /dev/null +++ b/python/metatomic_torchsim/docs/howto/model_loading.md @@ -0,0 +1,86 @@ +# Loading models + +`MetatomicModel` accepts several input formats. Each section below shows one +loading pattern. + +## From a saved `.pt` file + +The most common case. Pass the path to a TorchScript-exported metatomic model: + +```python +from metatomic.torchsim import MetatomicModel + +model = MetatomicModel("path/to/model.pt", device="cpu") +``` + +The file must exist and contain a valid `AtomisticModel`. A `ValueError` is +raised if the path does not exist. + +## From a metatrain checkpoint + +Pass a `.ckpt` path to load a metatrain checkpoint directly. This requires the +`metatrain` package: + +```python +model = MetatomicModel("path/to/checkpoint.ckpt") +``` + +The checkpoint is exported to an `AtomisticModel` internally. + +## PET-MAD shortcut + +The string `"pet-mad"` downloads and loads the PET-MAD universal model: + +```python +model = MetatomicModel("pet-mad") +``` + +This also requires `metatrain` to be installed. The model weights are fetched +from HuggingFace on first use. + +## From a Python AtomisticModel + +If you already have an `AtomisticModel` instance (for example, built +programmatically): + +```python +from metatomic.torch import AtomisticModel + +atomistic_model = build_my_model() # returns AtomisticModel +model = MetatomicModel(atomistic_model, device="cuda") +``` + +## From a TorchScript RecursiveScriptModule + +If you have a scripted model loaded via `torch.jit.load`: + +```python +import torch + +scripted = torch.jit.load("model.pt") +model = MetatomicModel(scripted, device="cpu") +``` + +The script module must have `original_name == "AtomisticModel"`. Otherwise a +`TypeError` is raised. + +## Selecting a device + +By default, `MetatomicModel` picks the best device from the model's +`supported_devices`. Override with the `device` parameter: + +```python +model = MetatomicModel("model.pt", device="cuda:0") +``` + +## Extensions directory + +Some models require compiled TorchScript extensions. Point to their location +with `extensions_directory`: + +```python +model = MetatomicModel( + "model.pt", + extensions_directory="path/to/extensions/", +) +``` diff --git a/python/metatomic_torchsim/docs/index.md b/python/metatomic_torchsim/docs/index.md new file mode 100644 index 00000000..e82574a7 --- /dev/null +++ b/python/metatomic_torchsim/docs/index.md @@ -0,0 +1,63 @@ +# metatomic-torchsim + +```{toctree} +:hidden: + +tutorials/getting_started +``` + +```{toctree} +:hidden: +:caption: How-to Guides + +howto/model_loading +howto/batched_simulations +``` + +```{toctree} +:hidden: +:caption: Understanding + +explanation/architecture +``` + +```{toctree} +:hidden: +:caption: Reference + +autoapi/metatomic/torchsim/index +changelog +``` + +**metatomic-torchsim** adapts [metatomic](https://docs.metatensor.org/metatomic/latest/) +atomistic models for use with [TorchSim](https://radical-ai.github.io/torch-sim/), +a differentiable molecular dynamics framework built on PyTorch. + +## Features + +- Run any metatomic-compatible model (PET-MAD, MACE, etc.) inside TorchSim + simulations +- Compute energies, forces, and stresses via autograd +- Batch multiple systems in a single forward pass +- GPU-accelerated neighbor lists via nvalchemiops when available + +## Quick install + +```bash +pip install metatomic-torchsim +``` + +## Minimal example + +```python +from metatomic.torchsim import MetatomicModel +import torch_sim as ts + +model = MetatomicModel("model.pt", device="cpu") +sim_state = ts.io.atoms_to_state([atoms], model.device, model.dtype) +results = model(sim_state) + +print(results["energy"]) # shape [n_systems] +print(results["forces"]) # shape [n_atoms, 3] +print(results["stress"]) # shape [n_systems, 3, 3] +``` diff --git a/python/metatomic_torchsim/docs/newsfragments/.gitkeep b/python/metatomic_torchsim/docs/newsfragments/.gitkeep new file mode 100644 index 00000000..e69de29b diff --git a/python/metatomic_torchsim/docs/tutorials/getting_started.md b/python/metatomic_torchsim/docs/tutorials/getting_started.md new file mode 100644 index 00000000..b20f23a7 --- /dev/null +++ b/python/metatomic_torchsim/docs/tutorials/getting_started.md @@ -0,0 +1,84 @@ +# Getting started + +This tutorial walks through running a short NVE molecular dynamics simulation +with a metatomic model and TorchSim. + +## Prerequisites + +Install the package and its dependencies: + +```bash +pip install metatomic-torchsim +``` + +You also need a saved metatomic model file (`.pt`). If you have a metatrain +checkpoint (`.ckpt`), install metatrain as well: + +```bash +pip install metatrain +``` + +## Load the model + +```python +from metatomic.torchsim import MetatomicModel + +model = MetatomicModel("path/to/model.pt", device="cpu") +``` + +The wrapper detects the model's dtype and supported devices automatically. Pass +`device="cuda"` to run on GPU. + +## Build a simulation state + +TorchSim works with `SimState` objects. Convert ASE `Atoms` using +`torch_sim.io.atoms_to_state`: + +```python +import ase.build +import torch_sim as ts + +atoms = ase.build.bulk("Si", "diamond", a=5.43, cubic=True) +sim_state = ts.io.atoms_to_state([atoms], model.device, model.dtype) +``` + +## Evaluate the model + +Call the model on the simulation state to get energies, forces, and stresses: + +```python +results = model(sim_state) + +print("Energy:", results["energy"]) # shape [1] +print("Forces:", results["forces"]) # shape [n_atoms, 3] +print("Stress:", results["stress"]) # shape [1, 3, 3] +``` + +## Run NVE dynamics + +Use TorchSim's Velocity Verlet integrator: + +```python +from torch_sim.integrators import VelocityVerletIntegrator + +integrator = VelocityVerletIntegrator( + model=model, + state=sim_state, + dt=1.0, # femtoseconds +) + +for step in range(100): + sim_state = integrator.step(sim_state) + if step % 10 == 0: + energy = model(sim_state)["energy"].item() + print(f"Step {step:3d} E = {energy:.4f} eV") +``` + +The total energy should remain approximately constant in an NVE simulation, +which serves as a basic sanity check for your model. + +## Next steps + +- {doc}`/howto/model_loading` covers all supported input formats +- {doc}`/howto/batched_simulations` explains running multiple systems at once +- {doc}`/explanation/architecture` describes the internals diff --git a/python/metatomic_torchsim/metatomic/__init__.py b/python/metatomic_torchsim/metatomic/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/python/metatomic_torchsim/metatomic/torchsim/__init__.py b/python/metatomic_torchsim/metatomic/torchsim/__init__.py new file mode 100644 index 00000000..5452e18c --- /dev/null +++ b/python/metatomic_torchsim/metatomic/torchsim/__init__.py @@ -0,0 +1,4 @@ +from ._calculator import MetatomicModel + + +__all__ = ["MetatomicModel"] diff --git a/python/metatomic_torchsim/metatomic/torchsim/_calculator.py b/python/metatomic_torchsim/metatomic/torchsim/_calculator.py new file mode 100644 index 00000000..23901a5b --- /dev/null +++ b/python/metatomic_torchsim/metatomic/torchsim/_calculator.py @@ -0,0 +1,413 @@ +"""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.neighborlist 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 exc: + raise ImportError( + "torch-sim is required for metatomic-torchsim: pip install torch-sim-atomistic" + ) from exc + + +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: + """Initialize the metatomic model wrapper. + + :param model: Model to use. Accepts a file path to a ``.pt`` saved + model, a ``.ckpt`` metatrain checkpoint (requires ``metatrain``), the + string ``"pet-mad"`` (shortcut for the PET-MAD model, requires + ``metatrain``), 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) and model == "pet-mad": + model = self._load_metatrain_model( + "https://huggingface.co/lab-cosmo/pet-mad/resolve/v1.1.0/" + "models/pet-mad-v1.1.0.ckpt" + ) + elif isinstance(model, (str, bytes, pathlib.PurePath)): + model_path = str(model) + if model_path.endswith(".ckpt"): + model = self._load_metatrain_model(model_path) + else: + 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._memory_scales_with = "n_atoms_x_density" + 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) + }, + ) + + @staticmethod + def _load_metatrain_model(path: str) -> AtomisticModel: + """Load a metatrain checkpoint and export it as an AtomisticModel.""" + try: + from metatrain.utils.io import load_model + except ImportError as exc: + raise ImportError( + "metatrain is required to load .ckpt files or use the 'pet-mad' " + "shortcut: pip install metatrain" + ) from exc + + return load_model(path).export() + + 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 00000000..33860d6c --- /dev/null +++ b/python/metatomic_torchsim/pyproject.toml @@ -0,0 +1,146 @@ +[project] +name = "metatomic-torchsim" +dynamic = ["version"] +requires-python = ">=3.10" + +readme = "README.md" +license = "BSD-3-Clause" +description = "TorchSim integration for metatomic atomistic models" + +authors = [ + {name = "Rhys Goodall"}, + {name = "Guillaume Fraux"}, + {name = "Filippo Bigi"}, + {name = "Rohit Goswami"}, +] + +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", + "torch-sim-atomistic", +] + +[project.optional-dependencies] +metatrain = ["metatrain"] + +[project.urls] +homepage = "https://docs.metatensor.org/metatomic/" +documentation = "https://docs.metatensor.org/metatomic/" +repository = "https://github.com/metatensor/metatomic" + +### ======================================================================== ### + +[build-system] +requires = [ + "hatchling >=1.21", + "hatch-vcs >=0.4", +] +build-backend = "hatchling.build" + +[tool.hatch.version] +source = "vcs" +fallback-version = "0.0.0.dev0" + +[tool.hatch.version.raw-options] +root = "../.." +git_describe_command = ["git", "describe", "--dirty", "--tags", "--long", "--match", "metatomic-torchsim-v*"] +tag_regex = "^metatomic-torchsim-v(?P[vV]?\\d+(?:\\.\\d+)*(?:[._-]?\\w+)*)$" + +[tool.hatch.build.targets.wheel] +packages = ["metatomic"] + +### ======================================================================== ### + +[tool.towncrier] +directory = "docs/newsfragments" +filename = "CHANGELOG.md" +title_format = "## metatomic-torchsim v{version} ({project_date})" +underlines = ["", "", ""] +issue_format = "[#{issue}](https://github.com/metatensor/metatomic/issues/{issue})" + +[[tool.towncrier.type]] +directory = "added" +name = "Added" +showcontent = true + +[[tool.towncrier.type]] +directory = "changed" +name = "Changed" +showcontent = true + +[[tool.towncrier.type]] +directory = "fixed" +name = "Fixed" +showcontent = true + +[[tool.towncrier.type]] +directory = "removed" +name = "Removed" +showcontent = true + +[[tool.towncrier.type]] +directory = "deprecated" +name = "Deprecated" +showcontent = true + +[[tool.towncrier.type]] +directory = "security" +name = "Security" +showcontent = true + +[[tool.towncrier.type]] +directory = "dev" +name = "Developer" +showcontent = true + +[[tool.towncrier.type]] +directory = "misc" +name = "Miscellaneous" +showcontent = true + +### ======================================================================== ### + +[tool.tbump] + +[tool.tbump.version] +current = "0.0.0" +regex = ''' + (?P\d+) + \. + (?P\d+) + \. + (?P\d+) + ''' + +[tool.tbump.git] +push_remote = "origin" +tag_template = "metatomic-torchsim-v{new_version}" +message_template = "Release metatomic-torchsim v{new_version}" + +### ======================================================================== ### + +[tool.pytest.ini_options] +python_files = ["*.py"] +testpaths = ["tests"] +filterwarnings = [ + "error", + "ignore:`torch.jit.script` 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", + "ignore:.*vesin.metatomic was only tested with metatomic.torch >=0.1.3,<0.2.*:UserWarning", +] diff --git a/python/metatomic_torchsim/tests/__init__.py b/python/metatomic_torchsim/tests/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/python/metatomic_torchsim/tests/test_model_loading.py b/python/metatomic_torchsim/tests/test_model_loading.py new file mode 100644 index 00000000..3e14a135 --- /dev/null +++ b/python/metatomic_torchsim/tests/test_model_loading.py @@ -0,0 +1,64 @@ +"""Tests for MetatomicModel loading paths.""" + +import pytest +import torch + + +ts = pytest.importorskip("torch_sim") + +import metatomic_lj_test # noqa: E402 + +from metatomic.torchsim import MetatomicModel # noqa: E402 + + +DEVICE = torch.device("cpu") + + +def _make_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, + ) + + +@pytest.fixture +def lj_model(): + return _make_lj_model() + + +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/test_torchsim.py b/python/metatomic_torchsim/tests/test_torchsim.py new file mode 100644 index 00000000..fd98433e --- /dev/null +++ b/python/metatomic_torchsim/tests/test_torchsim.py @@ -0,0 +1,293 @@ +"""Tests for the MetatomicModel TorchSim wrapper. + +Uses the metatomic-lj-test model so that tests run without needing metatrain or +downloading large model files. +""" + +import numpy as np +import pytest +import torch + + +ts = pytest.importorskip("torch_sim") + +import metatomic_lj_test # noqa: E402 + +from metatomic.torchsim import MetatomicModel # noqa: E402 + + +CUTOFF = 5.0 +SIGMA = 1.5808 +EPSILON = 0.1729 + +DEVICE = torch.device("cpu") +DTYPE = torch.float64 + + +def _make_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, + ) + + +def _make_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 lj_model(): + return _make_lj_model() + + +@pytest.fixture +def ni_atoms(): + return _make_ni_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 + + +def _make_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): + """Forward pass handles batched systems correctly.""" + atoms_2 = _make_ni_atoms_2() + sim_state = ts.io.atoms_to_state([ni_atoms, atoms_2], DEVICE, DTYPE) + output = metatomic_model(sim_state) + + assert output["energy"].shape == (2,) + n_total = len(ni_atoms) + len(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): + """Energy from single system matches the corresponding entry in a batch.""" + atoms_2 = _make_ni_atoms_2() + + # 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([atoms_2], DEVICE, DTYPE) + out_2 = metatomic_model(state_2) + + # batch + state_batch = ts.io.atoms_to_state([ni_atoms, 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 2aa9f2be..b38f1ba0 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 cd0a281d..5282b07b 100644 --- a/tox.ini +++ b/tox.ini @@ -162,6 +162,68 @@ 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 = + hatchling >=1.21 + hatch-vcs >=0.4 + {[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 + hatchling >=1.21 + hatch-vcs >=0.4 + # for metatensor-lj-test + setuptools-scm + cmake + +changedir = python/metatomic_torchsim +commands = + # install metatomic-torchsim (pure Python, hatchling backend) + 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 --cov-report= --import-mode=append {posargs} + + +[testenv:torchsim-docs] +description = Build the metatomic-torchsim documentation +package = skip +deps = + shibuya + sphinx >=9 + sphinx-autoapi >=3.6 + myst-parser + {toxinidir}/python/metatomic_torchsim + +commands = + sphinx-build -W -b html \ + {toxinidir}/python/metatomic_torchsim/docs \ + {toxinidir}/python/metatomic_torchsim/docs/_build + + [testenv:docs-tests] description = Run the doctests defined in any metatomic package deps =