ace-jax 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- ace_jax-0.1.0/.gitignore +25 -0
- ace_jax-0.1.0/LICENSE +21 -0
- ace_jax-0.1.0/PKG-INFO +135 -0
- ace_jax-0.1.0/README.md +92 -0
- ace_jax-0.1.0/pyproject.toml +145 -0
- ace_jax-0.1.0/src/ace_jax/__init__.py +28 -0
- ace_jax-0.1.0/src/ace_jax/basis/__init__.py +15 -0
- ace_jax-0.1.0/src/ace_jax/basis/coupling.py +274 -0
- ace_jax-0.1.0/src/ace_jax/basis/embedding.py +134 -0
- ace_jax-0.1.0/src/ace_jax/basis/export.py +159 -0
- ace_jax-0.1.0/src/ace_jax/basis/model.py +417 -0
- ace_jax-0.1.0/src/ace_jax/basis/prior.py +143 -0
- ace_jax-0.1.0/src/ace_jax/basis/radial_ace1.py +151 -0
- ace_jax-0.1.0/src/ace_jax/basis/radial_init.py +331 -0
- ace_jax-0.1.0/src/ace_jax/basis/spec.py +191 -0
- ace_jax-0.1.0/src/ace_jax/calc/__init__.py +9 -0
- ace_jax-0.1.0/src/ace_jax/calc/gp.py +87 -0
- ace_jax-0.1.0/src/ace_jax/calc/point.py +540 -0
- ace_jax-0.1.0/src/ace_jax/calc/skin.py +214 -0
- ace_jax-0.1.0/src/ace_jax/cli.py +410 -0
- ace_jax-0.1.0/src/ace_jax/eval/__init__.py +24 -0
- ace_jax-0.1.0/src/ace_jax/eval/api.py +82 -0
- ace_jax-0.1.0/src/ace_jax/eval/edge_model.py +359 -0
- ace_jax-0.1.0/src/ace_jax/eval/harmonics.py +107 -0
- ace_jax-0.1.0/src/ace_jax/eval/io.py +157 -0
- ace_jax-0.1.0/src/ace_jax/eval/model.py +778 -0
- ace_jax-0.1.0/src/ace_jax/eval/nlist.py +269 -0
- ace_jax-0.1.0/src/ace_jax/eval/pace_build.py +105 -0
- ace_jax-0.1.0/src/ace_jax/eval/pace_io.py +203 -0
- ace_jax-0.1.0/src/ace_jax/eval/pace_model.py +245 -0
- ace_jax-0.1.0/src/ace_jax/eval/pace_radial.py +244 -0
- ace_jax-0.1.0/src/ace_jax/eval/radial.py +133 -0
- ace_jax-0.1.0/src/ace_jax/eval/splinify.py +362 -0
- ace_jax-0.1.0/src/ace_jax/export/__init__.py +1 -0
- ace_jax-0.1.0/src/ace_jax/export/lammps.py +379 -0
- ace_jax-0.1.0/src/ace_jax/fit/__init__.py +2 -0
- ace_jax-0.1.0/src/ace_jax/fit/ard.py +494 -0
- ace_jax-0.1.0/src/ace_jax/fit/baseline.py +68 -0
- ace_jax-0.1.0/src/ace_jax/fit/block_lbfgs.py +182 -0
- ace_jax-0.1.0/src/ace_jax/fit/data.py +272 -0
- ace_jax-0.1.0/src/ace_jax/fit/density.py +76 -0
- ace_jax-0.1.0/src/ace_jax/fit/embedding.py +41 -0
- ace_jax-0.1.0/src/ace_jax/fit/feature.py +42 -0
- ace_jax-0.1.0/src/ace_jax/fit/fit_api.py +18 -0
- ace_jax-0.1.0/src/ace_jax/fit/hostcache.py +173 -0
- ace_jax-0.1.0/src/ace_jax/fit/hypers.py +54 -0
- ace_jax-0.1.0/src/ace_jax/fit/inducing.py +162 -0
- ace_jax-0.1.0/src/ace_jax/fit/kernels.py +106 -0
- ace_jax-0.1.0/src/ace_jax/fit/ladder.py +244 -0
- ace_jax-0.1.0/src/ace_jax/fit/loo.py +63 -0
- ace_jax-0.1.0/src/ace_jax/fit/metrics.py +60 -0
- ace_jax-0.1.0/src/ace_jax/fit/multistart.py +69 -0
- ace_jax-0.1.0/src/ace_jax/fit/objective.py +182 -0
- ace_jax-0.1.0/src/ace_jax/fit/paramset.py +156 -0
- ace_jax-0.1.0/src/ace_jax/fit/pipeline/__init__.py +9 -0
- ace_jax-0.1.0/src/ace_jax/fit/pipeline/config.py +124 -0
- ace_jax-0.1.0/src/ace_jax/fit/pipeline/data.py +211 -0
- ace_jax-0.1.0/src/ace_jax/fit/pipeline/export.py +147 -0
- ace_jax-0.1.0/src/ace_jax/fit/pipeline/lstsq.py +40 -0
- ace_jax-0.1.0/src/ace_jax/fit/pipeline/mapfit.py +73 -0
- ace_jax-0.1.0/src/ace_jax/fit/pipeline/objective.py +75 -0
- ace_jax-0.1.0/src/ace_jax/fit/pipeline/outputs.py +112 -0
- ace_jax-0.1.0/src/ace_jax/fit/pipeline/predict.py +144 -0
- ace_jax-0.1.0/src/ace_jax/fit/pipeline/problem.py +46 -0
- ace_jax-0.1.0/src/ace_jax/fit/pipeline/radials.py +84 -0
- ace_jax-0.1.0/src/ace_jax/fit/pipeline/run.py +82 -0
- ace_jax-0.1.0/src/ace_jax/fit/pipeline/rungs.py +50 -0
- ace_jax-0.1.0/src/ace_jax/fit/pops.py +212 -0
- ace_jax-0.1.0/src/ace_jax/fit/predict.py +597 -0
- ace_jax-0.1.0/src/ace_jax/fit/radial_learn.py +623 -0
- ace_jax-0.1.0/src/ace_jax/fit/radial_model.py +246 -0
- ace_jax-0.1.0/src/ace_jax/fit/report.py +48 -0
- ace_jax-0.1.0/src/ace_jax/fit/rows.py +209 -0
- ace_jax-0.1.0/src/ace_jax/fit/sharding.py +30 -0
- ace_jax-0.1.0/src/ace_jax/fit/solve.py +196 -0
- ace_jax-0.1.0/src/ace_jax/fit/stats.py +531 -0
- ace_jax-0.1.0/src/ace_jax/fit/summary.py +60 -0
- ace_jax-0.1.0/src/ace_jax/fit/varpro.py +49 -0
- ace_jax-0.1.0/src/ace_jax/fit/weights.py +139 -0
- ace_jax-0.1.0/src/ace_jax/fit/xyz.py +84 -0
- ace_jax-0.1.0/src/ace_jax/fit/zbl.py +52 -0
- ace_jax-0.1.0/src/ace_jax/runfile.py +177 -0
- ace_jax-0.1.0/src/ace_jax/tutorials/__init__.py +2 -0
- ace_jax-0.1.0/src/ace_jax/tutorials/labels.py +181 -0
- ace_jax-0.1.0/src/ace_jax/tutorials/structures.py +84 -0
ace_jax-0.1.0/.gitignore
ADDED
|
@@ -0,0 +1,25 @@
|
|
|
1
|
+
__pycache__/
|
|
2
|
+
*.pyc
|
|
3
|
+
.venv/
|
|
4
|
+
.pytest_cache/
|
|
5
|
+
dist/
|
|
6
|
+
build/
|
|
7
|
+
*.egg-info/
|
|
8
|
+
.uv/
|
|
9
|
+
.jax_cache/
|
|
10
|
+
pace_ref/build/
|
|
11
|
+
pace_ref/.venv/
|
|
12
|
+
bench/scaling/models/
|
|
13
|
+
coupling/build/
|
|
14
|
+
!coupling/julia/build/
|
|
15
|
+
dist-linux/
|
|
16
|
+
# marimo session caches (written by `marimo export`/`edit`); also the docs build output
|
|
17
|
+
__marimo__/
|
|
18
|
+
site/
|
|
19
|
+
# tutorial pages rendered from the notebooks by docs/build_tutorials.py
|
|
20
|
+
docs/user/tutorials/first-fit.md
|
|
21
|
+
docs/user/tutorials/learned-radials.md
|
|
22
|
+
docs/user/tutorials/dataset-and-properties.md
|
|
23
|
+
docs/user/tutorials/basis-and-evidence.md
|
|
24
|
+
docs/user/tutorials/truth-about-the-truth.md
|
|
25
|
+
docs/user/tutorials/*_files/
|
ace_jax-0.1.0/LICENSE
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2022 Christoph Ortner <christophortner0@gmail.com> and contributors
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
ace_jax-0.1.0/PKG-INFO
ADDED
|
@@ -0,0 +1,135 @@
|
|
|
1
|
+
Metadata-Version: 2.5
|
|
2
|
+
Name: ace-jax
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Fit and evaluate ACE interatomic potentials in Python/JAX
|
|
5
|
+
Project-URL: Homepage, https://github.com/ACEsuit/ace-jax
|
|
6
|
+
Project-URL: Documentation, https://acesuit.github.io/ace-jax/
|
|
7
|
+
Project-URL: Source, https://github.com/ACEsuit/ace-jax
|
|
8
|
+
Project-URL: Issues, https://github.com/ACEsuit/ace-jax/issues
|
|
9
|
+
Project-URL: Changelog, https://github.com/ACEsuit/ace-jax/blob/main/CHANGELOG.md
|
|
10
|
+
Author: ACEsuit contributors
|
|
11
|
+
License-Expression: MIT
|
|
12
|
+
License-File: LICENSE
|
|
13
|
+
Keywords: ace,chemistry,gaussian-process,interatomic-potential,jax,molecular-dynamics
|
|
14
|
+
Classifier: Development Status :: 3 - Alpha
|
|
15
|
+
Classifier: Intended Audience :: Science/Research
|
|
16
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
17
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
18
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
19
|
+
Classifier: Programming Language :: Python :: 3.13
|
|
20
|
+
Classifier: Programming Language :: Python :: 3.14
|
|
21
|
+
Classifier: Topic :: Scientific/Engineering :: Chemistry
|
|
22
|
+
Classifier: Topic :: Scientific/Engineering :: Physics
|
|
23
|
+
Requires-Python: >=3.11
|
|
24
|
+
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')
|
|
25
|
+
Requires-Dist: ase-extxyz>=0.1.2
|
|
26
|
+
Requires-Dist: ase>=3.22
|
|
27
|
+
Requires-Dist: equinox>=0.11
|
|
28
|
+
Requires-Dist: extxyz>=0.4.6
|
|
29
|
+
Requires-Dist: jax>=0.10.1
|
|
30
|
+
Requires-Dist: lineax>=0.1.1
|
|
31
|
+
Requires-Dist: numpy>=1.24
|
|
32
|
+
Requires-Dist: numpyro>=0.16
|
|
33
|
+
Requires-Dist: optax>=0.2
|
|
34
|
+
Requires-Dist: pyyaml>=6
|
|
35
|
+
Requires-Dist: scipy>=1.10
|
|
36
|
+
Provides-Extra: cuda
|
|
37
|
+
Requires-Dist: jax[cuda12]>=0.10.1; extra == 'cuda'
|
|
38
|
+
Provides-Extra: fast-neighbours
|
|
39
|
+
Requires-Dist: matscipy-neighbours>=1.0; extra == 'fast-neighbours'
|
|
40
|
+
Provides-Extra: gp
|
|
41
|
+
Requires-Dist: blackjax>=1.6; extra == 'gp'
|
|
42
|
+
Description-Content-Type: text/markdown
|
|
43
|
+
|
|
44
|
+
# ace-jax
|
|
45
|
+
|
|
46
|
+
[](https://pypi.org/project/ace-jax/)
|
|
47
|
+
[](https://acesuit.github.io/ace-jax/)
|
|
48
|
+
[](https://github.com/ACEsuit/ace-jax/actions/workflows/test.yml)
|
|
49
|
+
[](https://github.com/ACEsuit/ace-jax/blob/main/LICENSE)
|
|
50
|
+
|
|
51
|
+
Build, fit and evaluate **Atomic Cluster Expansion (ACE)** interatomic
|
|
52
|
+
potentials in Python with [JAX](https://docs.jax.dev), on CPU or GPU, all from
|
|
53
|
+
one `pip install`.
|
|
54
|
+
|
|
55
|
+
**Documentation: <https://acesuit.github.io/ace-jax/>**: installation,
|
|
56
|
+
quickstart, tutorials, how-to guides, and the CLI and API reference.
|
|
57
|
+
|
|
58
|
+
## What it does
|
|
59
|
+
|
|
60
|
+
- **Build a basis and fit it in one command.** `aj fit --order 3 --max-degree 10
|
|
61
|
+
--train train.xyz --out fit` builds the symmetry-adapted ACE basis and fits it.
|
|
62
|
+
The `fit.yaml` it writes reproduces the run.
|
|
63
|
+
- **Bayesian fits.** Linear ACE is fitted by Bayesian linear regression, with the
|
|
64
|
+
energy, force and virial noise levels and the prior scale chosen by maximising
|
|
65
|
+
the evidence: no hand-tuned weights. A hybrid ACE + Gaussian-process arm adds
|
|
66
|
+
a calibrated uncertainty ladder (MAP, Laplace, Pathfinder, VI, NUTS), and the
|
|
67
|
+
linear model has POPS and ARD uncertainties. The radial basis can be learned
|
|
68
|
+
as part of the fit (`--learn-radial`).
|
|
69
|
+
- **Fast evaluation.** `ACECalculator` and `GPCalculator` are ASE calculators,
|
|
70
|
+
fast enough for molecular dynamics, with predicted `energy_std` and
|
|
71
|
+
`forces_std` for GP and ARD models. `export_lammps` deploys a model to LAMMPS
|
|
72
|
+
through [lammps-jax](https://github.com/abhijeetgangan/lammps-jax).
|
|
73
|
+
- **PACE potentials.** pacemaker `.yace` files load, evaluate and write back.
|
|
74
|
+
- **Data in, data out.** Training data is extended XYZ or a list of `ase.Atoms`
|
|
75
|
+
(stress labels are converted to virials). `aj eval` writes predictions as
|
|
76
|
+
extended XYZ, with an RMSE table per config type.
|
|
77
|
+
|
|
78
|
+
## Install
|
|
79
|
+
|
|
80
|
+
```bash
|
|
81
|
+
pip install ace-jax
|
|
82
|
+
```
|
|
83
|
+
|
|
84
|
+
Extras: `"ace-jax[cuda]"` (CUDA 12 JAX), `"ace-jax[gp]"` (the Pathfinder rung)
|
|
85
|
+
and `"ace-jax[fast-neighbours]"` (a C++ neighbour list). Building new bases needs
|
|
86
|
+
the `ace-jax-coupling` wheel, a core dependency available for Linux x86_64 and
|
|
87
|
+
aarch64, macOS arm64 and Windows x64; elsewhere, fit and evaluate from an
|
|
88
|
+
existing basis file. See
|
|
89
|
+
[Installation](https://acesuit.github.io/ace-jax/installation/).
|
|
90
|
+
|
|
91
|
+
## Quickstart
|
|
92
|
+
|
|
93
|
+
Fit a linear ACE model to labelled data, check it on a test set, and evaluate it:
|
|
94
|
+
|
|
95
|
+
```bash
|
|
96
|
+
aj fit --order 3 --max-degree 10 --train train.xyz --test test.xyz \
|
|
97
|
+
--e0 lsq --m-per-species 0 --opt lbfgs --out fit
|
|
98
|
+
aj eval --model fit/model.npz --data test.xyz --out predictions.xyz
|
|
99
|
+
```
|
|
100
|
+
|
|
101
|
+
Use the fitted model from Python:
|
|
102
|
+
|
|
103
|
+
```python
|
|
104
|
+
import jax
|
|
105
|
+
jax.config.update("jax_enable_x64", True)
|
|
106
|
+
|
|
107
|
+
from ase.build import bulk
|
|
108
|
+
from ace_jax import ACECalculator
|
|
109
|
+
|
|
110
|
+
atoms = bulk("Si", "diamond", a=5.43, cubic=True)
|
|
111
|
+
atoms.calc = ACECalculator("fit/model.npz")
|
|
112
|
+
print(atoms.get_potential_energy(), atoms.get_forces())
|
|
113
|
+
```
|
|
114
|
+
|
|
115
|
+
The [Quickstart](https://acesuit.github.io/ace-jax/quickstart/) runs this on a
|
|
116
|
+
small silicon data set, and the
|
|
117
|
+
[tutorials](https://acesuit.github.io/ace-jax/tutorials/) are notebooks that run
|
|
118
|
+
on a laptop CPU, including ones adapted from the
|
|
119
|
+
[MLIP school 2026](https://github.com/ACEsuit/MLIP-school-2026). For coding
|
|
120
|
+
agents, [`skills/ace-jax/SKILL.md`](https://github.com/ACEsuit/ace-jax/blob/main/skills/ace-jax/SKILL.md)
|
|
121
|
+
is a compact usage guide.
|
|
122
|
+
|
|
123
|
+
## Performance and validation
|
|
124
|
+
|
|
125
|
+
Fits reproduce ACEfit's design matrix and least-squares solve to 1e-8, and PACE
|
|
126
|
+
evaluation matches the ML-PACE C++ code and python-ace; CI checks both. The
|
|
127
|
+
throughput benchmarks against LAMMPS ML-PACE are in
|
|
128
|
+
[`docs/dev/benchmarks.md`](https://github.com/ACEsuit/ace-jax/blob/main/docs/dev/benchmarks.md).
|
|
129
|
+
|
|
130
|
+
## Contributing
|
|
131
|
+
|
|
132
|
+
See [CONTRIBUTING.md](https://github.com/ACEsuit/ace-jax/blob/main/CONTRIBUTING.md)
|
|
133
|
+
for the development setup, tests and the reference-parity jobs, and the
|
|
134
|
+
[changelog](https://github.com/ACEsuit/ace-jax/blob/main/CHANGELOG.md) for
|
|
135
|
+
releases. ace-jax is MIT-licensed and part of [ACEsuit](https://github.com/ACEsuit).
|
ace_jax-0.1.0/README.md
ADDED
|
@@ -0,0 +1,92 @@
|
|
|
1
|
+
# ace-jax
|
|
2
|
+
|
|
3
|
+
[](https://pypi.org/project/ace-jax/)
|
|
4
|
+
[](https://acesuit.github.io/ace-jax/)
|
|
5
|
+
[](https://github.com/ACEsuit/ace-jax/actions/workflows/test.yml)
|
|
6
|
+
[](https://github.com/ACEsuit/ace-jax/blob/main/LICENSE)
|
|
7
|
+
|
|
8
|
+
Build, fit and evaluate **Atomic Cluster Expansion (ACE)** interatomic
|
|
9
|
+
potentials in Python with [JAX](https://docs.jax.dev), on CPU or GPU, all from
|
|
10
|
+
one `pip install`.
|
|
11
|
+
|
|
12
|
+
**Documentation: <https://acesuit.github.io/ace-jax/>**: installation,
|
|
13
|
+
quickstart, tutorials, how-to guides, and the CLI and API reference.
|
|
14
|
+
|
|
15
|
+
## What it does
|
|
16
|
+
|
|
17
|
+
- **Build a basis and fit it in one command.** `aj fit --order 3 --max-degree 10
|
|
18
|
+
--train train.xyz --out fit` builds the symmetry-adapted ACE basis and fits it.
|
|
19
|
+
The `fit.yaml` it writes reproduces the run.
|
|
20
|
+
- **Bayesian fits.** Linear ACE is fitted by Bayesian linear regression, with the
|
|
21
|
+
energy, force and virial noise levels and the prior scale chosen by maximising
|
|
22
|
+
the evidence: no hand-tuned weights. A hybrid ACE + Gaussian-process arm adds
|
|
23
|
+
a calibrated uncertainty ladder (MAP, Laplace, Pathfinder, VI, NUTS), and the
|
|
24
|
+
linear model has POPS and ARD uncertainties. The radial basis can be learned
|
|
25
|
+
as part of the fit (`--learn-radial`).
|
|
26
|
+
- **Fast evaluation.** `ACECalculator` and `GPCalculator` are ASE calculators,
|
|
27
|
+
fast enough for molecular dynamics, with predicted `energy_std` and
|
|
28
|
+
`forces_std` for GP and ARD models. `export_lammps` deploys a model to LAMMPS
|
|
29
|
+
through [lammps-jax](https://github.com/abhijeetgangan/lammps-jax).
|
|
30
|
+
- **PACE potentials.** pacemaker `.yace` files load, evaluate and write back.
|
|
31
|
+
- **Data in, data out.** Training data is extended XYZ or a list of `ase.Atoms`
|
|
32
|
+
(stress labels are converted to virials). `aj eval` writes predictions as
|
|
33
|
+
extended XYZ, with an RMSE table per config type.
|
|
34
|
+
|
|
35
|
+
## Install
|
|
36
|
+
|
|
37
|
+
```bash
|
|
38
|
+
pip install ace-jax
|
|
39
|
+
```
|
|
40
|
+
|
|
41
|
+
Extras: `"ace-jax[cuda]"` (CUDA 12 JAX), `"ace-jax[gp]"` (the Pathfinder rung)
|
|
42
|
+
and `"ace-jax[fast-neighbours]"` (a C++ neighbour list). Building new bases needs
|
|
43
|
+
the `ace-jax-coupling` wheel, a core dependency available for Linux x86_64 and
|
|
44
|
+
aarch64, macOS arm64 and Windows x64; elsewhere, fit and evaluate from an
|
|
45
|
+
existing basis file. See
|
|
46
|
+
[Installation](https://acesuit.github.io/ace-jax/installation/).
|
|
47
|
+
|
|
48
|
+
## Quickstart
|
|
49
|
+
|
|
50
|
+
Fit a linear ACE model to labelled data, check it on a test set, and evaluate it:
|
|
51
|
+
|
|
52
|
+
```bash
|
|
53
|
+
aj fit --order 3 --max-degree 10 --train train.xyz --test test.xyz \
|
|
54
|
+
--e0 lsq --m-per-species 0 --opt lbfgs --out fit
|
|
55
|
+
aj eval --model fit/model.npz --data test.xyz --out predictions.xyz
|
|
56
|
+
```
|
|
57
|
+
|
|
58
|
+
Use the fitted model from Python:
|
|
59
|
+
|
|
60
|
+
```python
|
|
61
|
+
import jax
|
|
62
|
+
jax.config.update("jax_enable_x64", True)
|
|
63
|
+
|
|
64
|
+
from ase.build import bulk
|
|
65
|
+
from ace_jax import ACECalculator
|
|
66
|
+
|
|
67
|
+
atoms = bulk("Si", "diamond", a=5.43, cubic=True)
|
|
68
|
+
atoms.calc = ACECalculator("fit/model.npz")
|
|
69
|
+
print(atoms.get_potential_energy(), atoms.get_forces())
|
|
70
|
+
```
|
|
71
|
+
|
|
72
|
+
The [Quickstart](https://acesuit.github.io/ace-jax/quickstart/) runs this on a
|
|
73
|
+
small silicon data set, and the
|
|
74
|
+
[tutorials](https://acesuit.github.io/ace-jax/tutorials/) are notebooks that run
|
|
75
|
+
on a laptop CPU, including ones adapted from the
|
|
76
|
+
[MLIP school 2026](https://github.com/ACEsuit/MLIP-school-2026). For coding
|
|
77
|
+
agents, [`skills/ace-jax/SKILL.md`](https://github.com/ACEsuit/ace-jax/blob/main/skills/ace-jax/SKILL.md)
|
|
78
|
+
is a compact usage guide.
|
|
79
|
+
|
|
80
|
+
## Performance and validation
|
|
81
|
+
|
|
82
|
+
Fits reproduce ACEfit's design matrix and least-squares solve to 1e-8, and PACE
|
|
83
|
+
evaluation matches the ML-PACE C++ code and python-ace; CI checks both. The
|
|
84
|
+
throughput benchmarks against LAMMPS ML-PACE are in
|
|
85
|
+
[`docs/dev/benchmarks.md`](https://github.com/ACEsuit/ace-jax/blob/main/docs/dev/benchmarks.md).
|
|
86
|
+
|
|
87
|
+
## Contributing
|
|
88
|
+
|
|
89
|
+
See [CONTRIBUTING.md](https://github.com/ACEsuit/ace-jax/blob/main/CONTRIBUTING.md)
|
|
90
|
+
for the development setup, tests and the reference-parity jobs, and the
|
|
91
|
+
[changelog](https://github.com/ACEsuit/ace-jax/blob/main/CHANGELOG.md) for
|
|
92
|
+
releases. ace-jax is MIT-licensed and part of [ACEsuit](https://github.com/ACEsuit).
|
|
@@ -0,0 +1,145 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["hatchling"]
|
|
3
|
+
build-backend = "hatchling.build"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "ace-jax"
|
|
7
|
+
version = "0.1.0"
|
|
8
|
+
description = "Fit and evaluate ACE interatomic potentials in Python/JAX"
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
license = "MIT"
|
|
11
|
+
license-files = ["LICENSE"]
|
|
12
|
+
requires-python = ">=3.11"
|
|
13
|
+
authors = [{ name = "ACEsuit contributors" }]
|
|
14
|
+
keywords = ["ace", "interatomic-potential", "jax", "molecular-dynamics", "chemistry", "gaussian-process"]
|
|
15
|
+
classifiers = [
|
|
16
|
+
"Development Status :: 3 - Alpha",
|
|
17
|
+
"Intended Audience :: Science/Research",
|
|
18
|
+
"License :: OSI Approved :: MIT License",
|
|
19
|
+
"Programming Language :: Python :: 3.11",
|
|
20
|
+
"Programming Language :: Python :: 3.12",
|
|
21
|
+
"Programming Language :: Python :: 3.13",
|
|
22
|
+
"Programming Language :: Python :: 3.14",
|
|
23
|
+
"Topic :: Scientific/Engineering :: Chemistry",
|
|
24
|
+
"Topic :: Scientific/Engineering :: Physics",
|
|
25
|
+
]
|
|
26
|
+
dependencies = [
|
|
27
|
+
"jax>=0.10.1",
|
|
28
|
+
"equinox>=0.11",
|
|
29
|
+
"numpy>=1.24",
|
|
30
|
+
"scipy>=1.10",
|
|
31
|
+
"lineax>=0.1.1",
|
|
32
|
+
"ase>=3.22", # core: neighbour-list fallback (eval) + calculators + element data
|
|
33
|
+
"pyyaml>=6", # .yace (PACE) model files
|
|
34
|
+
# training/eval data reader (libAtoms C parser, binary wheels): labels come back
|
|
35
|
+
# as written, independent of ASE's calculator-key handling (see fit/xyz.py)
|
|
36
|
+
"extxyz>=0.4.6",
|
|
37
|
+
# aj eval writes its extxyz output through the ASE plugin (create_calc=False keeps every
|
|
38
|
+
# label under its own name); 0.1.2 / extxyz 0.4.6 fixed a transposed Lattice on write
|
|
39
|
+
"ase-extxyz>=0.1.2",
|
|
40
|
+
# the fit pipeline: the theta-MAP (numpyro SVI) and the L-BFGS / Adam optimisers (optax)
|
|
41
|
+
"numpyro>=0.16",
|
|
42
|
+
"optax>=0.2",
|
|
43
|
+
# coupling tables for new basis shapes (a compiled library, nothing to install
|
|
44
|
+
# besides the wheel); platform wheels only -- elsewhere fit from an existing .npz
|
|
45
|
+
"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')",
|
|
46
|
+
]
|
|
47
|
+
|
|
48
|
+
[project.optional-dependencies]
|
|
49
|
+
# The pathfinder rung of the hyperparameter UQ ladder (everything else in `aj fit`,
|
|
50
|
+
# GP arm included, needs only the core dependencies).
|
|
51
|
+
gp = ["blackjax>=1.6"]
|
|
52
|
+
cuda = ["jax[cuda12]>=0.10.1"]
|
|
53
|
+
# Fast neighbour lists; ace-jax uses them when present and falls back to ASE's
|
|
54
|
+
# list otherwise, so they are intentionally NOT core deps. matscipy-neighbours
|
|
55
|
+
# (preferred: OpenMP CPU, native dense neighbour_matrix) ships as a source
|
|
56
|
+
# distribution: installing it compiles C++ (CMake + a C++17 compiler). For its
|
|
57
|
+
# GPU backend (graph built on the device, handed to JAX via DLPack):
|
|
58
|
+
# pip install -C cmake.define.ENABLE_CUDA=ON \
|
|
59
|
+
# -C cmake.define.CMAKE_CUDA_ARCHITECTURES=<sm> matscipy-neighbours
|
|
60
|
+
fast-neighbours = ["matscipy-neighbours>=1.0"]
|
|
61
|
+
|
|
62
|
+
[project.scripts]
|
|
63
|
+
ace-jax = "ace_jax.cli:main"
|
|
64
|
+
aj = "ace_jax.cli:main" # short alias
|
|
65
|
+
|
|
66
|
+
[project.urls]
|
|
67
|
+
Homepage = "https://github.com/ACEsuit/ace-jax"
|
|
68
|
+
Documentation = "https://acesuit.github.io/ace-jax/"
|
|
69
|
+
Source = "https://github.com/ACEsuit/ace-jax"
|
|
70
|
+
Issues = "https://github.com/ACEsuit/ace-jax/issues"
|
|
71
|
+
Changelog = "https://github.com/ACEsuit/ace-jax/blob/main/CHANGELOG.md"
|
|
72
|
+
|
|
73
|
+
[dependency-groups]
|
|
74
|
+
# mirrors the `gp` extra so a plain `uv run pytest` can run the GP ladder tests
|
|
75
|
+
dev = [
|
|
76
|
+
"pytest>=8", "pytest-split>=0.8", "pytest-xdist>=3.5", "numpyro>=0.16", "optax>=0.2", "blackjax>=1.6",
|
|
77
|
+
# lint + hooks; keep ruff in step with the rev in .pre-commit-config.yaml
|
|
78
|
+
"ruff==0.16.9", "pre-commit>=4",
|
|
79
|
+
]
|
|
80
|
+
|
|
81
|
+
[tool.hatch.build.targets.wheel]
|
|
82
|
+
packages = ["src/ace_jax"]
|
|
83
|
+
|
|
84
|
+
[tool.hatch.build.targets.sdist]
|
|
85
|
+
include = ["/src/ace_jax", "/README.md", "/LICENSE"] # anchored: unanchored names match nested READMEs too
|
|
86
|
+
|
|
87
|
+
[tool.pytest.ini_options]
|
|
88
|
+
pythonpath = ["src", "tests"]
|
|
89
|
+
testpaths = ["tests"]
|
|
90
|
+
# `slow` (real-model ladder, bit-exact goldens) is opt-in: `pytest -m slow`
|
|
91
|
+
addopts = "--durations=15 -m 'not slow'"
|
|
92
|
+
markers = [
|
|
93
|
+
"slow: heavy integration (real-model MCMC ladder, bit-exact driver goldens); excluded by default -- run with `pytest -m slow`",
|
|
94
|
+
]
|
|
95
|
+
|
|
96
|
+
[tool.ruff]
|
|
97
|
+
# ~99.9% of src/tests lines are <= 110 already; no formatter is enforced
|
|
98
|
+
# (`ruff format` would rewrite the deliberately dense one-line style).
|
|
99
|
+
line-length = 110
|
|
100
|
+
target-version = "py311" # requires-python >= 3.11
|
|
101
|
+
# Not linted: reference generators / experiments that are not part of the package.
|
|
102
|
+
extend-exclude = ["pace_ref", "julia", "docs", "fixtures", ".worktrees"]
|
|
103
|
+
|
|
104
|
+
[tool.ruff.lint]
|
|
105
|
+
# pycodestyle, pyflakes, bugbear, pyupgrade; NPY guards NumPy-2 removals.
|
|
106
|
+
# No "I" (isort): ruff's isort always splits `import a, b` and re-wraps dense
|
|
107
|
+
# parenthesised imports one-per-line (74 files, +270/-138), and reorders package
|
|
108
|
+
# __init__ imports; that is a restyle, not a fix.
|
|
109
|
+
select = ["E", "F", "W", "B", "UP", "NPY"]
|
|
110
|
+
ignore = [
|
|
111
|
+
# House style, not bugs: dense `a = ...; b = ...` lines, one-line `if x: y`,
|
|
112
|
+
# short math names (l, I, O), named lambdas, `import os, sys`.
|
|
113
|
+
"E701", "E702", "E741", "E731", "E401",
|
|
114
|
+
# Imports after jax.config.update("jax_enable_x64", ...) / importorskip are deliberate.
|
|
115
|
+
"E402",
|
|
116
|
+
# Long lines are tolerated (no formatter); line-length documents the target.
|
|
117
|
+
"E501",
|
|
118
|
+
# zip(strict=True) would change behaviour on length mismatch; opt in case by case.
|
|
119
|
+
"B905",
|
|
120
|
+
]
|
|
121
|
+
|
|
122
|
+
|
|
123
|
+
[tool.ruff.lint.per-file-ignores]
|
|
124
|
+
# B023 (closure over a loop variable): every site in these files was reviewed and
|
|
125
|
+
# is a false positive -- the closure is consumed in the same iteration
|
|
126
|
+
# (`jax.tree.map(lambda a: a[i], ds)`, in-loop jit timing). Listed per file so
|
|
127
|
+
# the rule still guards new files.
|
|
128
|
+
"bench/learn_radial/compare_plots.py" = ["B023"]
|
|
129
|
+
"bench/learn_radial/plots.py" = ["B023"]
|
|
130
|
+
"bench/learn_radial/rmse.py" = ["B023"]
|
|
131
|
+
"bench/learn_radial/rmse_npz.py" = ["B023"]
|
|
132
|
+
"bench/learn_radial/tradeoff.py" = ["B023"]
|
|
133
|
+
"src/ace_jax/fit/hostcache.py" = ["B023"]
|
|
134
|
+
"src/ace_jax/fit/pipeline/predict.py" = ["B023"]
|
|
135
|
+
"src/ace_jax/fit/predict.py" = ["B023"]
|
|
136
|
+
"src/ace_jax/fit/radial_model.py" = ["B023"]
|
|
137
|
+
"src/ace_jax/fit/solve.py" = ["B023"]
|
|
138
|
+
"src/ace_jax/fit/stats.py" = ["B023"]
|
|
139
|
+
"tests/test_ard.py" = ["B023"]
|
|
140
|
+
"tests/test_gp_learn_radial.py" = ["B023"]
|
|
141
|
+
"tests/test_gp_pops_paper.py" = ["B023"]
|
|
142
|
+
"tests/test_gp_rows.py" = ["B023"]
|
|
143
|
+
"tests/test_gp_stats.py" = ["B023"]
|
|
144
|
+
"tests/test_radial_learn_flags.py" = ["B023"]
|
|
145
|
+
"tests/test_rows_chunked.py" = ["B023"]
|
|
@@ -0,0 +1,28 @@
|
|
|
1
|
+
"""ace-jax: fit and evaluate ACE interatomic potentials in Python/JAX.
|
|
2
|
+
|
|
3
|
+
Build a basis (`ace_jax.basis`), fit it (`ace_jax.fit.pipeline`, or `aj fit`) and
|
|
4
|
+
evaluate it (`ACECalculator`, `export_lammps`). Documentation:
|
|
5
|
+
https://acesuit.github.io/ace-jax/
|
|
6
|
+
"""
|
|
7
|
+
import importlib.metadata as _md
|
|
8
|
+
|
|
9
|
+
try:
|
|
10
|
+
__version__ = _md.version("ace-jax")
|
|
11
|
+
except _md.PackageNotFoundError: # a source tree that was never installed
|
|
12
|
+
__version__ = "0+unknown"
|
|
13
|
+
from .eval import (ACEModel, load, highest_precision, site_descriptors, species_indices,
|
|
14
|
+
sparse_graph, dense_graph)
|
|
15
|
+
|
|
16
|
+
__all__ = ["__version__", "ACEModel", "load", "highest_precision", "site_descriptors", "species_indices",
|
|
17
|
+
"sparse_graph", "dense_graph", "ACECalculator", "GPCalculator"]
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
def __getattr__(name):
|
|
21
|
+
# ase-dependent calculators kept out of the core import path
|
|
22
|
+
if name == "ACECalculator":
|
|
23
|
+
from .calc.point import ACECalculator
|
|
24
|
+
return ACECalculator
|
|
25
|
+
if name == "GPCalculator":
|
|
26
|
+
from .calc.gp import GPCalculator
|
|
27
|
+
return GPCalculator
|
|
28
|
+
raise AttributeError(name)
|
|
@@ -0,0 +1,15 @@
|
|
|
1
|
+
"""Build ACE basis definitions from Python (`BasisSpec`, `build_basis`).
|
|
2
|
+
|
|
3
|
+
Scope now: the SO(3) coupling table (A2B + aa_spec), generated by the
|
|
4
|
+
compiled EquivariantTensors library (ace-jax-coupling) -- see docs/dev/coupling-etshim-spec.md --
|
|
5
|
+
and the algebraic smoothness prior (prior.py, pure numpy, parity-tested against
|
|
6
|
+
julia/smoothness_reference.jl). Radials / pair basis / embedding stay exported
|
|
7
|
+
(separate follow-up).
|
|
8
|
+
|
|
9
|
+
prior.py is deliberately numpy-only so the eval/fit path can import it as a
|
|
10
|
+
fallback for exports that lack gamma (cli.py and the bench driver go through
|
|
11
|
+
prior.prior_diagonal, which rebuilds it from the export's meta["nnll"]); the
|
|
12
|
+
coupling shim is the only piece that
|
|
13
|
+
needs the compiled ace-jax-coupling library (a core dependency on supported platforms;
|
|
14
|
+
elsewhere `BasisUnavailable`). Refitting on an existing model never touches it.
|
|
15
|
+
"""
|