Build, fit and evaluate Atomic Cluster Expansion (ACE) interatomic
potentials in Python with JAX, on CPU or GPU. One
pip install gives all of these functions.
Documentation: https://acesuit.github.io/ace-jax/. It contains the installation, a quickstart, tutorials, how-to guides, and the CLI and API reference.
- Build a basis and fit it with one command.
aj fit --order 3 --max-degree 10 --train train.xyz --out fitbuilds the symmetry-adapted ACE basis and fits it. Thefit.yamlfile that it writes runs the fit again. - Bayesian fits. Linear ACE uses Bayesian linear regression. The fit
selects the energy, force and virial noise levels and the prior scale by
maximising the evidence, so you do not set weights by hand.
- An experimental hybrid ACE + Gaussian-process model
(
--m-per-species> 0) adds a calibrated set of uncertainty approximations (MAP, Laplace, Pathfinder, VI, NUTS). Most fits use the linear model, the default. - The linear model has POPS and ARD uncertainties.
- ARD gives a conformally calibrated per-atom force uncertainty
(
forces_std,forces_q, a 3x3forces_cov).aj calibraterecalibrates it on new labelled cells. - The fit can also learn the radial basis (
--learn-radial).
- An experimental hybrid ACE + Gaussian-process model
(
- Fast evaluation.
ACECalculatorandGPCalculatorare ASE calculators, fast enough for molecular dynamics. GP and ARD models also predictenergy_stdandforces_std.export_lammpsdeploys a model to LAMMPS with lammps-jax. On CPUs, linear ACE runs in LAMMPS as a C++ kernel (aj export --platform cpu). - PACE potentials. ace-jax loads, evaluates and writes pacemaker
.yacefiles. - Data in, data out. Training data is extended XYZ or a list of
ase.Atoms. ace-jax converts stress labels to virials.aj evalwrites predictions as extended XYZ, with an RMSE table for each configuration type.
pip install ace-jaxExtras: "ace-jax[cuda]" (CUDA 12 JAX), "ace-jax[cuda13]" (CUDA 13 JAX and nvcc
for the GPU kernel; NVIDIA driver 580 or later), "ace-jax[gp]" (the Pathfinder rung),
"ace-jax[native]" (C++ code on CPUs: a neighbour list and the kernel), "ace-jax[lammps]"
(LAMMPS export with lammps-jax) and "ace-jax[jaxmd]" (MD with JAX-MD).
To build a new basis, ace-jax needs the ace-jax-coupling wheel. This core
dependency is available for Linux x86_64 and aarch64, macOS arm64 and
Windows x64. On other platforms, fit and evaluate from an existing basis
file. See
Installation.
Fit a linear ACE model to labelled data, check it on a test set, and evaluate it:
aj fit --order 3 --max-degree 10 --train train.xyz --test test.xyz \
--e0 lsq --opt lbfgs --out fit
aj eval --model fit/model.npz --data test.xyz --out predictions.xyzUse the fitted model from Python:
import jax
jax.config.update("jax_enable_x64", True)
from ase.build import bulk
from ace_jax import ACECalculator
atoms = bulk("Si", "diamond", a=5.43, cubic=True)
atoms.calc = ACECalculator("fit/model.npz")
print(atoms.get_potential_energy(), atoms.get_forces())The Quickstart runs this on a
small silicon dataset. The
tutorials are notebooks that run
on a laptop CPU. Some of them come from the
MLIP School 2026. For coding
agents, skills/ace-jax/SKILL.md
is a compact usage guide.
Fits agree with the ACEfit design matrix and least-squares solution to 1e-8.
PACE evaluation agrees with the ML-PACE C++ code and python-ace. CI checks
both. The throughput benchmarks against LAMMPS ML-PACE are in
docs/dev/benchmarks.md.
See CONTRIBUTING.md for the development setup, tests and the reference-parity jobs. See the changelog for releases. ace-jax has the MIT licence and is part of ACEsuit.