CuspAI/kUPS-mattersim-jax

MatterSim v1.0.0 — JAX build

This repository hosts JAX exports of MatterSim v1.0.0 for use with 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. 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	MIT	Yang et al. (2024), arXiv:2405.04967
mattersim-v1.0.0-5m.zip (18 MB)	MatterSim-v1.0.0-5M	MIT	Yang et al. (2024), arXiv: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

pip install kups[cuda]

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:

kups_relax_mlff --config relax_mattersim.yaml

Example configs live in the kUPS 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 · 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 and documentation 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:

@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/; one export per checkpoint:

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 for the full text and NOTICE for attribution, modification, and trademark details. The kUPS / tojax tooling citations:

@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
• Scientific questions about MatterSim: please direct to the upstream authors via microsoft/mattersim or ai4s-materials@microsoft.com.