Skip to content

kUPS

kUPS is a toolkit for building high-performance molecular simulations on JAX. It provides composable, differentiable primitives — samplers, potentials, and propagators — with hardware acceleration on CPU, GPU, and TPU.

Installation

pip install kups

For GPU support (Linux only):

pip install kups[cuda]

Quick Start

The repository includes example applications built with kUPS. To try them, run from the examples/ directory:

cd examples
kups_mcmc_rigid mcmc_rigid.yaml

cd examples
kups_md_lj md_lj_argon_nvt.yaml

See the examples/ directory for more configurations.

Features

  • Composable — every operation is a propagator with a shared interface; methods and potentials snap together freely
  • 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 — Lennard-Jones, Coulomb (Ewald summation), harmonic bonds/angles, Morse, MACE, UMA
  • Differentiable — full automatic differentiation through simulations via JAX
  • Batched — run thousands of independent simulations as a single vectorized computation
  • GPU-native — JIT-compiled on CPU, GPU, and TPU with no code changes
  • PyTorch interop — bring any PyTorch model into JAX via Tojax

License

Apache License 2.0 — see LICENSE.