ace-jax¶
ace-jax builds, fits and evaluates Atomic Cluster Expansion (ACE)
interatomic potentials in Python with JAX. One
pip install gives all of these functions. It also evaluates pacemaker PACE .yace potentials.
- Fit from data with one command.
aj fit --order 3 --max-degree 10 ...builds the ACE basis for the species in your data and fits it. The fit is a Bayesian linear regression: the evidence sets the weights and the regularisation. Each fit writes afit.yamlfile that runs it again. - Build a basis on its own with
aj basis. You can then save it, share it or fit it again. - Learn the radial basis by variable projection. For deployment, the learned radials become splines, so they are as fast as a stock model.
- Evaluate with
aj eval, or in Python with the ASE calculatorACECalculator. - Deploy to LAMMPS with lammps-jax.
pip install ace-jax
aj fit --order 3 --max-degree 10 \
--train train.xyz --test test.xyz \
--e0 lsq --m-per-species 0 --opt lbfgs \
--out fit
aj eval --model fit/model.npz --data test.xyz --out predictions.xyz
The Quickstart runs these commands on a small silicon dataset. The tutorials are marimo notebooks. You can run them on your computer or open them in molab.
Where to go next¶
| If you want to | Read |
|---|---|
| install ace-jax and its extras | Installation |
| run fit → eval once from the shell | Quickstart |
| work through a fit in a notebook | Tutorials |
| understand the basis, E0, the prior and the fit | Concepts |
use a model in ASE or LAMMPS, or load a .yace |
How-to guides: ASE, LAMMPS, PACE |
reproduce or vary a fit from its fit.yaml |
Run files |
| learn the radial basis | Learn the radial basis |
| put calibrated error bars on forces | Per-atom force uncertainty |
| look up every CLI flag | CLI reference |
| look up a Python function | Python API |
| correct a problem | FAQ and troubleshooting |
ace-jax is part of the ACEsuit organisation. Its licence is MIT; see Licence and citation.