README.md

---
license: mit
license_link: LICENSE
library_name: kups
tags:
  - chemistry
  - materials-science
  - molecular-dynamics
  - interatomic-potential
  - mlff
  - jax
  - mattersim
  - m3gnet
---

# MatterSim v1.0.0 — JAX build

This repository hosts JAX exports of [MatterSim v1.0.0](https://github.com/microsoft/mattersim) for use with [kUPS](https://github.com/cusp-ai-oss/kups), a JAX-native molecular-simulation toolkit. Each artefact is a self-contained `.zip` containing the serialized JAX computation graph, the original model parameters, and the minimal metadata needed to run inference.

**Important:** these are **re-exports, not retrainings.** Weights and architecture are the originals released by Microsoft Research. CuspAI's only contribution is converting the PyTorch reference implementation to JAX via [tojax](https://github.com/cusp-ai-oss/tojax). Every scientific claim, citation, and credit belongs to the original authors.

## Included models

| File | Upstream | License | Paper |
|------|----------|---------|-------|
| `mattersim-v1.0.0-1m.zip` (3.7 MB) | [MatterSim-v1.0.0-1M](https://github.com/microsoft/mattersim) | [MIT](LICENSE) | Yang et al. (2024), [arXiv:2405.04967](https://arxiv.org/abs/2405.04967) |
| `mattersim-v1.0.0-5m.zip` (18 MB) | [MatterSim-v1.0.0-5M](https://github.com/microsoft/mattersim) | [MIT](LICENSE) | Yang et al. (2024), [arXiv:2405.04967](https://arxiv.org/abs/2405.04967) |

Cutoff radius 5.0 Å (three-body cutoff 4.0 Å). Schema: `AtomGraphInput` (positions, atomic numbers, cell, pbc, edge index, cell offsets, batch index, charge, spin).

## Quick start

```sh
pip install kups[cuda]
```

```python
from huggingface_hub import hf_hub_download
from kups.potential.mliap.tojax import TojaxedMliap

path = hf_hub_download(
    repo_id="CuspAI/kUPS-mattersim-jax",
    filename="mattersim-v1.0.0-1m.zip",  # or mattersim-v1.0.0-5m.zip
)
model = TojaxedMliap.from_zip_file(path)
# model.cutoff, model.params, model.call(atom_graph_input) are now available.
```

kUPS ships CLI wrappers that take a YAML config pointing at one of these zips:

```sh
kups_relax_mlff --config relax_mattersim.yaml
```

Example configs live in the kUPS [examples/](https://github.com/cusp-ai-oss/kups/tree/main/examples) directory.

## What's in each `.zip`

- `model.jax` — JAX computation graph, serialized via `jax.export`.
- `params.msgpack` — parameters as a msgpack-encoded list of arrays.
- `metadata.json` — cutoff radius, three-body cutoff, and checkpoint name.
- `dtypes.json` — input dtypes for `AtomGraphInput`.

Exported with symbolic shapes (`--symbolic NSE`); accept variable atom, system, and edge counts without recompilation.

## Model details

**Upstream:** [microsoft/mattersim](https://github.com/microsoft/mattersim) · **Checkpoints:** `MatterSim-v1.0.0-1M.pth`, `MatterSim-v1.0.0-5M.pth`

MatterSim is a family of deep-learning interatomic potentials from Microsoft Research, designed for atomistic simulations of bulk materials across the periodic table, a wide range of temperatures, and pressures. The v1.0.0 checkpoints use the **M3GNet** architecture (a three-body message-passing graph network); the `1M` variant is a smaller, faster model and the `5M` variant trades speed for accuracy. Both are trained on data generated by the workflows described in the MatterSim manuscript.

**Original authors:** Han Yang, Chenxi Hu, Yichi Zhou, Xixian Liu, Yu Shi, Jielan Li, Guanzhi Li, Zekun Chen, Shuizhou Chen, Claudio Zeni, Matthew Horton, Robert Pinsler, Andrew Fowler, Daniel Zügner, Tian Xie, Jake Smith, Lixin Sun, Qian Wang, Lingyu Kong, Chang Liu, Hongxia Hao, Ziheng Lu. Microsoft Research AI4Science.

**Intended use and limitations:** MatterSim v1 is designed for atomistic simulations of **bulk** materials. Applications to surfaces, interfaces, or properties dominated by long-range interactions may be qualitatively reasonable but should not be used for quantitative analysis without fine-tuning. See the [upstream MODEL_CARD](https://github.com/microsoft/mattersim/blob/main/MODEL_CARD.md) and [documentation](https://microsoft.github.io/mattersim/) for authoritative guidance.

**Version reporting:** Per the upstream README, papers using these checkpoints should report the **exact** version — e.g. "MatterSim-v1.0.0-1M" — rather than the generic name "MatterSim", to ensure reproducibility.

**Citation:**

```bibtex
@article{yang2024mattersim,
  title         = {{MatterSim}: A Deep Learning Atomistic Model Across Elements,
                   Temperatures and Pressures},
  author        = {Yang, Han and Hu, Chenxi and Zhou, Yichi and Liu, Xixian and Shi, Yu
                   and Li, Jielan and Li, Guanzhi and Chen, Zekun and Chen, Shuizhou
                   and Zeni, Claudio and Horton, Matthew and Pinsler, Robert
                   and Fowler, Andrew and Z{\"u}gner, Daniel and Xie, Tian and Smith, Jake
                   and Sun, Lixin and Wang, Qian and Kong, Lingyu and Liu, Chang
                   and Hao, Hongxia and Lu, Ziheng},
  year          = {2024},
  eprint        = {2405.04967},
  archivePrefix = {arXiv},
  primaryClass  = {cond-mat.mtrl-sci},
  journal       = {arXiv preprint arXiv:2405.04967},
}
```

## Export pipeline and reproducibility

The archives were produced with the exporter in [tojax/examples/mlff/](https://github.com/cusp-ai-oss/tojax); one export per checkpoint:

```sh
for ckpt in mattersim-v1.0.0-1m mattersim-v1.0.0-5m; do
  uv run python export_mattersim.py --checkpoint ${ckpt}.pth \
      --output ${ckpt}.zip --symbolic NSE
done
```

tojax's export harness verifies numerical agreement with the PyTorch reference (default tolerances `rtol=1e-4`, `atol=1e-4`) before saving each archive.

## Changes from upstream

- **File format.** PyTorch `.pth` → JAX-exported `.zip` (graph + msgpack params).
- **Weights.** Unchanged, bit-for-bit, from the upstream checkpoint.
- **Architecture.** Translated operation-for-operation; no approximations or substitutions.
- **Numerics.** Verified within `rtol=1e-4, atol=1e-4` against the PyTorch reference during export.
- **Cutoffs and dtypes.** Preserved from upstream defaults (cutoff 5.0 Å, three-body cutoff 4.0 Å; float32 positions/cell/offsets, int64 indices, bool PBC).
- **Symbolic shapes.** Exports accept variable numbers of atoms, systems, and edges without recompilation.

## Attribution and license

These models exist because of the work of the MatterSim team at Microsoft Research. CuspAI's only contribution is the JAX export — we trained nothing, changed no weights, and designed none of the architecture. Please cite Yang et al. (2024) when using these checkpoints in research, and report the exact version name.

The archives are distributed under the **MIT License**; see [LICENSE](LICENSE) for the full text and [NOTICE](NOTICE) for attribution, modification, and trademark details. The kUPS / tojax tooling citations:

```bibtex
@software{kups2026,
  author = {{CuspAI}},
  title  = {kUPS},
  year   = {2026},
  url    = {https://github.com/cusp-ai-oss/kups},
}

@software{tojax2026,
  author = {{CuspAI}},
  title  = {tojax},
  year   = {2026},
  url    = {https://github.com/cusp-ai-oss/tojax},
}
```

## Contact

- Issues with the JAX export or with kUPS: [github.com/cusp-ai-oss/kups/issues](https://github.com/cusp-ai-oss/kups/issues)
- Scientific questions about MatterSim: please direct to the upstream authors via [microsoft/mattersim](https://github.com/microsoft/mattersim) or `ai4s-materials@microsoft.com`.