Add torch-sim integration as a separate metatomic-torchsim package

.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

.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 }}

.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 }}

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 -------------------------------------------------
 

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:

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::
 

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

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",
 ]

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",
 ]

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,

python/metatomic_torchsim/AUTHORS

@@ -0,0 +1,4 @@
+Rhys Goodall
+Guillaume Fraux
+Filippo Bigi
+Rohit Goswami

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.

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
+-->

python/metatomic_torchsim/MANIFEST.in

@@ -0,0 +1,2 @@
+include AUTHORS
+include git_version_info

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).

python/metatomic_torchsim/metatomic_torchsim/__init__.py

@@ -0,0 +1,4 @@
+from ._model import MetatomicModel
+
+
+__all__ = ["MetatomicModel"]