Metadata-Version: 2.5
Name: ace-jax
Version: 0.2.2
Summary: Fit and evaluate ACE interatomic potentials in Python/JAX
Project-URL: Homepage, https://github.com/ACEsuit/ace-jax
Project-URL: Documentation, https://acesuit.github.io/ace-jax/
Project-URL: Source, https://github.com/ACEsuit/ace-jax
Project-URL: Issues, https://github.com/ACEsuit/ace-jax/issues
Project-URL: Changelog, https://github.com/ACEsuit/ace-jax/blob/main/CHANGELOG.md
Author: ACEsuit contributors
License-Expression: MIT
License-File: LICENSE
Keywords: ace,chemistry,gaussian-process,interatomic-potential,jax,molecular-dynamics
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Programming Language :: Python :: 3.14
Classifier: Topic :: Scientific/Engineering :: Chemistry
Classifier: Topic :: Scientific/Engineering :: Physics
Requires-Python: >=3.11
Requires-Dist: ace-jax-coupling==0.2.0; (sys_platform == 'linux' and (platform_machine == 'x86_64' or platform_machine == 'aarch64')) or (sys_platform == 'darwin' and platform_machine == 'arm64') or (sys_platform == 'win32' and platform_machine == 'AMD64')
Requires-Dist: ase-extxyz>=0.1.2
Requires-Dist: ase>=3.22
Requires-Dist: equinox>=0.11
Requires-Dist: extxyz>=0.4.6
Requires-Dist: jax>=0.10.1
Requires-Dist: lineax>=0.1.1
Requires-Dist: numpy>=1.24
Requires-Dist: numpyro>=0.16
Requires-Dist: optax>=0.2
Requires-Dist: pyyaml>=6
Requires-Dist: scipy>=1.10
Provides-Extra: cuda
Requires-Dist: jax[cuda12]>=0.10.1; extra == 'cuda'
Provides-Extra: fast-neighbours
Requires-Dist: matscipy-neighbours>=1.0; extra == 'fast-neighbours'
Provides-Extra: gp
Requires-Dist: blackjax>=1.6; extra == 'gp'
Description-Content-Type: text/markdown

# ace-jax

[![PyPI](https://img.shields.io/pypi/v/ace-jax)](https://pypi.org/project/ace-jax/)
[![Docs](https://img.shields.io/github/actions/workflow/status/ACEsuit/ace-jax/docs.yml?branch=main&label=docs)](https://acesuit.github.io/ace-jax/)
[![Tests](https://img.shields.io/github/actions/workflow/status/ACEsuit/ace-jax/test.yml?branch=main&label=tests)](https://github.com/ACEsuit/ace-jax/actions/workflows/test.yml)
[![License: MIT](https://img.shields.io/badge/license-MIT-blue)](https://github.com/ACEsuit/ace-jax/blob/main/LICENSE)

Build, fit and evaluate **Atomic Cluster Expansion (ACE)** interatomic
potentials in Python with [JAX](https://docs.jax.dev), 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.

## What it does

- **Build a basis and fit it with one command.** `aj fit --order 3
  --max-degree 10 --train train.xyz --out fit` builds the symmetry-adapted
  ACE basis and fits it. The `fit.yaml` file 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.
    - A hybrid ACE + Gaussian-process model adds a calibrated set of
      uncertainty approximations (MAP, Laplace, Pathfinder, VI, NUTS).
    - The linear model has POPS and ARD uncertainties.
    - ARD gives a conformally calibrated per-atom force uncertainty
      (`forces_std`, `forces_q`, a 3x3 `forces_cov`). `aj calibrate`
      recalibrates it on new labelled cells.
    - The fit can also learn the radial basis (`--learn-radial`).
- **Fast evaluation.** `ACECalculator` and `GPCalculator` are ASE
  calculators, fast enough for molecular dynamics. GP and ARD models also
  predict `energy_std` and `forces_std`. `export_lammps` deploys a model to
  LAMMPS with [lammps-jax](https://github.com/abhijeetgangan/lammps-jax).
- **PACE potentials.** ace-jax loads, evaluates and writes pacemaker
  `.yace` files.
- **Data in, data out.** Training data is extended XYZ or a list of
  `ase.Atoms`. ace-jax converts stress labels to virials. `aj eval` writes
  predictions as extended XYZ, with an RMSE table for each configuration
  type.

## Install

```bash
pip install ace-jax
```

Extras: `"ace-jax[cuda]"` (CUDA 12 JAX), `"ace-jax[gp]"` (the Pathfinder rung)
and `"ace-jax[fast-neighbours]"` (a C++ neighbour list).

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](https://acesuit.github.io/ace-jax/installation/).

## Quickstart

Fit a linear ACE model to labelled data, check it on a test set, and evaluate it:

```bash
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
```

Use the fitted model from Python:

```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](https://acesuit.github.io/ace-jax/quickstart/) runs this on a
small silicon dataset. The
[tutorials](https://acesuit.github.io/ace-jax/tutorials/) are notebooks that run
on a laptop CPU. Some of them come from the
[MLIP School 2026](https://mlipschool.uk/). For coding
agents, [`skills/ace-jax/SKILL.md`](https://github.com/ACEsuit/ace-jax/blob/main/skills/ace-jax/SKILL.md)
is a compact usage guide.

## Performance and validation

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`](https://github.com/ACEsuit/ace-jax/blob/main/docs/dev/benchmarks.md).

## Contributing

See [CONTRIBUTING.md](https://github.com/ACEsuit/ace-jax/blob/main/CONTRIBUTING.md)
for the development setup, tests and the reference-parity jobs. See the
[changelog](https://github.com/ACEsuit/ace-jax/blob/main/CHANGELOG.md) for
releases. ace-jax has the MIT licence and is part of [ACEsuit](https://github.com/ACEsuit).
