Create a metatomic-torchsim package

HaoZeke · CompRhys · Luthaf · Luthaf · commit f53ce492e604 · 2026-03-17T14:34:25.000+01:00
It will hold the TorchSim &lt;=&gt; metatomic integration as a standalone
module.

Co-Authored-By: Rhys Goodall &lt;rhys.goodall@outlook.com&gt;
Co-Authored-By: Guillaume Fraux &lt;guillaume.fraux@epfl.ch&gt;
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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 <energy-output>` 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 <https://radical-ai.github.io/torch-sim/user/introduction.html#installation>`_.
+Only the :ref:`energy <energy-output>` 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
-<https://radical-ai.github.io/torch-sim/tutorials/metatomic_tutorial.html>`_,
-and generic documentation on torch-sim `there
-<radical-ai.github.io/torch-sim/>`_.
+.. 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
@@ -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
@@ -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).
 
 <!-- Possible sections for each package:
-
 ### Added
 
 ### Fixed
diff --git a/pyproject.toml b/pyproject.toml
@@ -67,6 +67,7 @@ filterwarnings = [
     "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:.*vesin.metatomic was only tested with metatomic.torch >=0.1.3,<0.2.*:UserWarning",
+    "ignore:`torch.jit.load` is deprecated. Please switch to `torch.export`.:DeprecationWarning",
 ]
 
 ### ======================================================================== ###
@@ -77,7 +78,7 @@ ignore = ["B018", "B904"]
 
 [tool.ruff.lint.isort]
 lines-after-imports = 2
-known-first-party = ["metatomic"]
+known-first-party = ["metatomic", "metatomic_torchsim", "metatomic_lj_test"]
 known-third-party = ["torch"]
 
 [tool.ruff.format]
@@ -88,4 +89,5 @@ docstring-code-format = true
 [tool.uv.pip]
 reinstall-package = [
     "metatomic-torch",
+    "metatomic-torchsim",
 ]
diff --git a/python/metatomic_torch/pyproject.toml b/python/metatomic_torch/pyproject.toml
@@ -58,8 +58,10 @@ 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.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",
 ]
diff --git a/python/metatomic_torch/tests/ase_calculator.py b/python/metatomic_torch/tests/ase_calculator.py
@@ -9,14 +9,14 @@
 import ase.calculators.lj
 import ase.md
 import ase.units
-import metatomic_lj_test
 import numpy as np
 import pytest
 import torch
 from ase.calculators.calculator import PropertyNotImplementedError
 from ase.md.velocitydistribution import MaxwellBoltzmannDistribution
 from metatensor.torch import Labels, TensorBlock, TensorMap
 
+import metatomic_lj_test
 from metatomic.torch import (
     AtomisticModel,
     ModelCapabilities,
diff --git a/python/metatomic_torchsim/AUTHORS b/python/metatomic_torchsim/AUTHORS
@@ -0,0 +1,4 @@
+Rhys Goodall
+Guillaume Fraux
+Filippo Bigi
+Rohit Goswami
diff --git a/python/metatomic_torchsim/Architecture.md b/python/metatomic_torchsim/Architecture.md
@@ -0,0 +1,71 @@
+# Architecture
+
+This page explains how `metatomic_torchsim.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/CHANGELOG.md b/python/metatomic_torchsim/CHANGELOG.md
@@ -0,0 +1,15 @@
+# Changelog
+
+All notable changes to metatomic-torchsim are documented here, following the
+[keep a changelog](https://keepachangelog.com/en/1.1.0/) format. This project
+follows [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
+
+<!-- Possible sections for each package:
+### Added
+
+### Fixed
+
+### Changed
+
+### Removed
+-->
diff --git a/python/metatomic_torchsim/MANIFEST.in 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
@@ -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
@@ -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
diff --git a/python/metatomic_torchsim/pyproject.toml b/python/metatomic_torchsim/pyproject.toml
diff --git a/python/metatomic_torchsim/setup.py b/python/metatomic_torchsim/setup.py
diff --git a/python/metatomic_torchsim/tests/__init__.py b/python/metatomic_torchsim/tests/__init__.py
diff --git a/python/metatomic_torchsim/tests/model_loading.py b/python/metatomic_torchsim/tests/model_loading.py
diff --git a/python/metatomic_torchsim/tests/torchsim.py b/python/metatomic_torchsim/tests/torchsim.py
diff --git a/setup.py b/setup.py
diff --git a/tox.ini b/tox.ini
