kups 0.1.1

A toolkit for building high-performance molecular simulations on JAX.

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 --config gcmc_co2_30box.yaml

Molecular Dynamics (Lennard-Jones)

cd examples
kups_md_lj --config 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 the docs locally

Always go through ./docs/scripts/build.sh. Do not run mkdocs or zensical directly — the script executes and converts docs/notebooks/*.ipynb to markdown, generates the API reference pages and a derived config (mkdocs.build.yml), and then invokes zensical.

./docs/scripts/build.sh                            # build into site/
./docs/scripts/build.sh --serve                    # live-reload server on http://127.0.0.1:8000
./docs/scripts/build.sh -f mkdocs.yml.dev --serve  # alternate config

Requires uv sync to have been run; the docs dependency group is pulled in by default (default-groups = "all" in pyproject.toml).

If --serve fails with OSError: [Errno 48] Address already in use, a previous server is still holding port 8000 — find and kill it:

lsof -iTCP:8000 -sTCP:LISTEN -n -P   # shows PID
kill <pid>

---

Citation

If you use kUPS in your research, please cite:

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

---

License

Apache License 2.0 — see LICENSE.