A toolkit for building high-performance molecular simulations on JAX

kUPS provides composable, differentiable primitives — samplers, potentials, and propagators — with hardware acceleration on CPU, GPU, and TPU.

Documentation | Quick Start | Features | Examples

---

Installation

Standard Installation	pip install kups
GPU Support	pip install kups[cuda]
Development	git clone https://github.com/cusp-ai-oss/kups.git cd kups uv sync

Quick Start

Tip

The repository includes example applications built with kUPS in the examples/ directory.

Monte Carlo Simulation (GCMC)

cd examples
kups_mcmc_rigid mcmc_rigid.yaml

Molecular Dynamics (Lennard-Jones)

cd examples
kups_md_lj md_lj_argon_nvt.yaml

Features

Simulation Methods Monte Carlo — NVT and GCMC ensembles with translation, rotation, reinsertion, and exchange moves Molecular Dynamics — NVE, NVT, NPT ensembles Geometry Optimization — FIRE and L-BFGS relaxation	Force Fields & Potentials Lennard-Jones potential Coulomb interactions (Ewald summation) Harmonic bonds and angles Morse potential MACE and UMA ML force fields
Core Capabilities Composable — shared propagator interface; methods and potentials snap together freely Batched — run thousands of independent simulations as vectorized computations	Performance & Integration GPU-native — JIT-compiled on CPU, GPU, and TPU with no code changes Differentiable — full automatic differentiation via JAX PyTorch interop — bring PyTorch models into JAX via Tojax

Documentation

Full documentation is available at cusp-ai-oss.github.io/kUPS.

Building docs locally

To build the docs locally run ./docs/scripts/build.sh, which executes and renders all documentation notebooks and generates API pages in markdown.

./docs/scripts/build.sh           # build into site/
./docs/scripts/build.sh --serve   # serve with live updates on http://127.0.0.1:8000

---

Citation

If you use kUPS in your research, please cite:

@software{kups2026,
  author = {Gao, Nicholas
    and K{\"o}hler, Jonas
    and Hanke, Felix
    and Ramanan, Anita
    and Moubarak, Elias
    and Morrow, Joe
    and de Haan, Pim
    and Openshaw, Hannah
    and Welling, Max
    and CuspAI Team},
  title = {kUPS - a universal particle simulation toolkit},
  year = {2026},
  url = {https://github.com/cusp-ai-oss/kups}
}

---

License

Apache License 2.0 — see LICENSE.