Packaged Simulations¶

kUPS ships with several ready-to-use simulation applications as CLI tools. Each is a thin layer built on the core primitives (propagators, potentials, lenses, tables) and serves as both a useful tool and a reference implementation. All commands take a YAML configuration file via and use nanoargs for argument parsing, so any configuration value can also be overridden from the command line. Example configurations are provided in the examples/ directory and should be run from there.

Molecular Dynamics¶

Run molecular dynamics trajectories in the NVE, NVT, or NPT ensemble with kups_md. The force field is selected by the potential.backend field:

cd examples
kups_md md_lj_argon_nvt.yaml
kups_md md_lj_argon_nve.yaml
kups_md md_mace.yaml
kups_md md_orb.yaml

Ensembles and integrators:

• NVE — velocity Verlet. Constant energy, useful for validating energy conservation.
• NVT — Langevin thermostat (BAOAB splitting) or canonical sampling via velocity rescaling (CSVR). Constant temperature.
• NPT — CSVR thermostat with stochastic cell rescaling barostat (csvr_npt), or fully-flexible-cell BAOAB NPT Langevin (baoab_npt_langevin, Gao–Fang–Wang JCP 2016). Constant temperature and pressure.

All integrators are built from the same composable propagator primitives described in the Propagators tutorial.

Geometry Optimization¶

Relax atomic positions (and optionally lattice vectors) to a local energy minimum with kups_relax. The force field is selected by the potential.backend field:

cd examples
kups_relax relax_mace.yaml
kups_relax relax_orb.yaml
kups_relax relax_torch_mace.yaml   # native MACE .model checkpoint
kups_relax relax_torch_uma.yaml    # native UMA .pt checkpoint (fairchem)

Optimizers:

• FIRE — fast inertial relaxation engine. Adaptive timestep, robust for rough energy landscapes.
• L-BFGS — limited-memory quasi-Newton method. Fast convergence near the minimum.
• Any Optax optimizer (Adam, SGD, etc.) can be plugged in via the same interface.

Relaxation converges when the maximum force on any atom drops below a configurable tolerance.

Grand-Canonical Monte Carlo (GCMC)¶

Simulate adsorption of rigid molecules in a host framework at constant chemical potential, volume, and temperature (μVT ensemble).

cd examples
kups_mcmc_rigid mcmc_rigid.yaml

Monte Carlo moves:

• Translation — displace a molecule by a random vector.
• Rotation — rotate a molecule about its center of mass.
• Reinsertion — delete a molecule and reinsert it at a random position and orientation.
• Exchange — insert or delete a molecule based on the chemical potential (fugacity computed via the Peng-Robinson equation of state).

Move probabilities and step sizes are configurable. The simulation supports multiple adsorbate species (CO₂, CH₄, H₂O, N₂, etc.) with pre-defined molecular geometries.

Widom Test-Particle Insertion¶

Compute the excess chemical potential, Henry coefficient, and isosteric heat of adsorption for a rigid adsorbate in a host framework. Ghost insertions accumulate Boltzmann factors \(\exp(-\beta\Delta U)\) alongside an optional NVT displacement chain over real adsorbates.

cd examples
kups_mcmc_widom mcmc_widom.yaml

Per-cycle WidomStatistics snapshots are written to the HDF5 output file. analyze_widom_file reduces them post-hoc into block-averaged \(\mu^\mathrm{ex}\), \(K_H\), and \(q_\mathrm{st}\) with standard errors (Vlugt 2008 eq. 16, \(N=0\)).

Machine-learning Force Fields¶

CuspAI publishes JAX exports of MACE and Orb on the Hugging Face Hub — one repository per model so each retains its upstream license:

These are re-exports (via Tojax), not retrainings — weights and architectures are unchanged from upstream.

Meta's UMA model is not redistributed by CuspAI. Two routes to run it with kUPS: (1) download the PyTorch checkpoint from Hugging Face and run it natively via kups_relax (potential.backend: uma); or (2) port it to JAX using Tojax following the instructions here and run via kups_relax (potential.backend: tojax).

Any potential.model_path: field accepts either an hf://<owner>/<repo>/<filename> URI (fetched via huggingface_hub.hf_hub_download and cached on first use) or a local filesystem path to a Tojax-exported .zip:

potential:
  backend: tojax
  # Remote (HF Hub, requires pip install kups[hf])
  model_path: hf://CuspAI/kUPS-mace-jax/mace-mpa-0-medium_32.zip
  # ...or local (anything readable by TojaxedMliap.from_zip_file)
  # model_path: ./my_model.zip

The hf:// scheme requires the optional huggingface_hub dependency: pip install kups[hf]. Local paths work without it.

Direct PyTorch Bridge¶

The mace and uma backends (for both kups_relax and kups_md) load native PyTorch MLFF checkpoints (no Tojax conversion) and run them from JAX via a DLPack-based bridge. Functionality across the two paths is identical (forces, stress, optimisation, HDF5 trajectory output) — the differences are about how the model is plumbed in.

Pros

• No JAX-traceability requirement. Any torch.nn.Module plugs in as-is: you write one adapter forward that consumes the universal AtomGraphInput dict and you're done. The model never has to be expressible in pure JAX ops or jaxprs.
• Run upstream checkpoints unmodified (MACE .model, UMA .pt) — no export step, no risk of conversion drift, easy to diff against the reference implementation.
• Full upstream feature surface — every inference toggle the original library exposes is still accessible. For UMA that means tunable InferenceSettings and all five dataset heads (omat/omol/oc20/odac/omc) selected per-call, without re-exporting one zip per head.
• Custom CUDA kernels (e.g. cuequivariance for MACE/UMA) execute on the original PyTorch path, no JAX reimplementation required.
• UMA turbo is faster than its Tojax export. inference_settings: turbo merges the mixture-of-linear-experts at lazy-init, beating the JAX-exported equivalent on the same model.

Cons

• Per-call JAX↔PyTorch hand-off. Zero-copy via DLPack but a serial dispatch — adds latency on every potential evaluation.
• MACE is slower than its Tojax export. For MACE specifically the JAX build wins; use the bridge mainly for unconverted checkpoints or reference diffs.
• Extra runtime dependency — you install the model's own PyTorch package (mace-torch, fairchem-core, …) alongside kUPS.
• Python-version constraints from upstream propagate (e.g. fairchem-core>=2.0 caps at Python ≤3.13).

Backends are selected via the discriminated potential: union in the YAML config:

# MACE — native .model checkpoint
potential:
  backend: mace
  model_path: ./mace-mpa-0-medium.model
  device: cuda
  dtype: float32           # or float64

# UMA — native .pt checkpoint via fairchem-core
potential:
  backend: uma
  model_path: ./uma-s-1p2.pt
  device: cuda
  task_name: omat          # one of: omat | omol | oc20 | odac | omc
  inference_settings: default   # or "turbo" (compiled MOLE merge — same composition every call)

Optional extras install the model's own PyTorch package alongside kUPS:

pip install kups[mace]   # mace-torch
pip install kups[uma]    # fairchem-core (≥2.0; Python ≤3.13)

Under the hood both backends share the universal TorchMliap container — adding a third backend means writing one torch.nn.Module that consumes the AtomGraphInput schema and returns {"energy", "position_gradients", "cell_gradients"}. See kups.potential.mliap.torch.mace and kups.potential.mliap.torch.uma for the adapter pattern.