Upstream metatomic TorchSim model

.github/CODEOWNERS

@@ -0,0 +1 @@
+/python/metatomic_torchsim/ @HaoZeke @CompRhys

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

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 <energy-output>` 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 <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.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
+<https://docs.metatensor.org/metatomic/latest/torchsim/>`_.

pyproject.toml

@@ -88,4 +88,5 @@ docstring-code-format = true
 [tool.uv.pip]
 reinstall-package = [
     "metatomic-torch",
+    "metatomic-torchsim",
 ]

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.

python/metatomic_torchsim/AUTHORS

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

python/metatomic_torchsim/CHANGELOG.md

@@ -0,0 +1,5 @@
+# Changelog
+
+All notable changes to metatomic-torchsim will be documented in this file.
+
+<!-- towncrier release notes start -->

python/metatomic_torchsim/MANIFEST.in

@@ -0,0 +1,2 @@
+include pyproject.toml
+include AUTHORS

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"]
+```

python/metatomic_torchsim/docs/.gitignore

@@ -0,0 +1 @@
+_build/

python/metatomic_torchsim/docs/_templates/partials/site-foot.html

@@ -0,0 +1,4 @@
+{% include "components/foot-copyright.html" %}
+<p class="site-foot-extra">
+  Part of the <a href="https://docs.metatensor.org">Metatensor ecosystem</a>.
+</p>

python/metatomic_torchsim/docs/changelog.md

@@ -0,0 +1,5 @@
+# Changelog
+
+```{include} ../CHANGELOG.md
+:start-after: "# Changelog"
+```

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

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.