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.
Files changed (85) hide show
  1. ace_jax-0.1.0/.gitignore +25 -0
  2. ace_jax-0.1.0/LICENSE +21 -0
  3. ace_jax-0.1.0/PKG-INFO +135 -0
  4. ace_jax-0.1.0/README.md +92 -0
  5. ace_jax-0.1.0/pyproject.toml +145 -0
  6. ace_jax-0.1.0/src/ace_jax/__init__.py +28 -0
  7. ace_jax-0.1.0/src/ace_jax/basis/__init__.py +15 -0
  8. ace_jax-0.1.0/src/ace_jax/basis/coupling.py +274 -0
  9. ace_jax-0.1.0/src/ace_jax/basis/embedding.py +134 -0
  10. ace_jax-0.1.0/src/ace_jax/basis/export.py +159 -0
  11. ace_jax-0.1.0/src/ace_jax/basis/model.py +417 -0
  12. ace_jax-0.1.0/src/ace_jax/basis/prior.py +143 -0
  13. ace_jax-0.1.0/src/ace_jax/basis/radial_ace1.py +151 -0
  14. ace_jax-0.1.0/src/ace_jax/basis/radial_init.py +331 -0
  15. ace_jax-0.1.0/src/ace_jax/basis/spec.py +191 -0
  16. ace_jax-0.1.0/src/ace_jax/calc/__init__.py +9 -0
  17. ace_jax-0.1.0/src/ace_jax/calc/gp.py +87 -0
  18. ace_jax-0.1.0/src/ace_jax/calc/point.py +540 -0
  19. ace_jax-0.1.0/src/ace_jax/calc/skin.py +214 -0
  20. ace_jax-0.1.0/src/ace_jax/cli.py +410 -0
  21. ace_jax-0.1.0/src/ace_jax/eval/__init__.py +24 -0
  22. ace_jax-0.1.0/src/ace_jax/eval/api.py +82 -0
  23. ace_jax-0.1.0/src/ace_jax/eval/edge_model.py +359 -0
  24. ace_jax-0.1.0/src/ace_jax/eval/harmonics.py +107 -0
  25. ace_jax-0.1.0/src/ace_jax/eval/io.py +157 -0
  26. ace_jax-0.1.0/src/ace_jax/eval/model.py +778 -0
  27. ace_jax-0.1.0/src/ace_jax/eval/nlist.py +269 -0
  28. ace_jax-0.1.0/src/ace_jax/eval/pace_build.py +105 -0
  29. ace_jax-0.1.0/src/ace_jax/eval/pace_io.py +203 -0
  30. ace_jax-0.1.0/src/ace_jax/eval/pace_model.py +245 -0
  31. ace_jax-0.1.0/src/ace_jax/eval/pace_radial.py +244 -0
  32. ace_jax-0.1.0/src/ace_jax/eval/radial.py +133 -0
  33. ace_jax-0.1.0/src/ace_jax/eval/splinify.py +362 -0
  34. ace_jax-0.1.0/src/ace_jax/export/__init__.py +1 -0
  35. ace_jax-0.1.0/src/ace_jax/export/lammps.py +379 -0
  36. ace_jax-0.1.0/src/ace_jax/fit/__init__.py +2 -0
  37. ace_jax-0.1.0/src/ace_jax/fit/ard.py +494 -0
  38. ace_jax-0.1.0/src/ace_jax/fit/baseline.py +68 -0
  39. ace_jax-0.1.0/src/ace_jax/fit/block_lbfgs.py +182 -0
  40. ace_jax-0.1.0/src/ace_jax/fit/data.py +272 -0
  41. ace_jax-0.1.0/src/ace_jax/fit/density.py +76 -0
  42. ace_jax-0.1.0/src/ace_jax/fit/embedding.py +41 -0
  43. ace_jax-0.1.0/src/ace_jax/fit/feature.py +42 -0
  44. ace_jax-0.1.0/src/ace_jax/fit/fit_api.py +18 -0
  45. ace_jax-0.1.0/src/ace_jax/fit/hostcache.py +173 -0
  46. ace_jax-0.1.0/src/ace_jax/fit/hypers.py +54 -0
  47. ace_jax-0.1.0/src/ace_jax/fit/inducing.py +162 -0
  48. ace_jax-0.1.0/src/ace_jax/fit/kernels.py +106 -0
  49. ace_jax-0.1.0/src/ace_jax/fit/ladder.py +244 -0
  50. ace_jax-0.1.0/src/ace_jax/fit/loo.py +63 -0
  51. ace_jax-0.1.0/src/ace_jax/fit/metrics.py +60 -0
  52. ace_jax-0.1.0/src/ace_jax/fit/multistart.py +69 -0
  53. ace_jax-0.1.0/src/ace_jax/fit/objective.py +182 -0
  54. ace_jax-0.1.0/src/ace_jax/fit/paramset.py +156 -0
  55. ace_jax-0.1.0/src/ace_jax/fit/pipeline/__init__.py +9 -0
  56. ace_jax-0.1.0/src/ace_jax/fit/pipeline/config.py +124 -0
  57. ace_jax-0.1.0/src/ace_jax/fit/pipeline/data.py +211 -0
  58. ace_jax-0.1.0/src/ace_jax/fit/pipeline/export.py +147 -0
  59. ace_jax-0.1.0/src/ace_jax/fit/pipeline/lstsq.py +40 -0
  60. ace_jax-0.1.0/src/ace_jax/fit/pipeline/mapfit.py +73 -0
  61. ace_jax-0.1.0/src/ace_jax/fit/pipeline/objective.py +75 -0
  62. ace_jax-0.1.0/src/ace_jax/fit/pipeline/outputs.py +112 -0
  63. ace_jax-0.1.0/src/ace_jax/fit/pipeline/predict.py +144 -0
  64. ace_jax-0.1.0/src/ace_jax/fit/pipeline/problem.py +46 -0
  65. ace_jax-0.1.0/src/ace_jax/fit/pipeline/radials.py +84 -0
  66. ace_jax-0.1.0/src/ace_jax/fit/pipeline/run.py +82 -0
  67. ace_jax-0.1.0/src/ace_jax/fit/pipeline/rungs.py +50 -0
  68. ace_jax-0.1.0/src/ace_jax/fit/pops.py +212 -0
  69. ace_jax-0.1.0/src/ace_jax/fit/predict.py +597 -0
  70. ace_jax-0.1.0/src/ace_jax/fit/radial_learn.py +623 -0
  71. ace_jax-0.1.0/src/ace_jax/fit/radial_model.py +246 -0
  72. ace_jax-0.1.0/src/ace_jax/fit/report.py +48 -0
  73. ace_jax-0.1.0/src/ace_jax/fit/rows.py +209 -0
  74. ace_jax-0.1.0/src/ace_jax/fit/sharding.py +30 -0
  75. ace_jax-0.1.0/src/ace_jax/fit/solve.py +196 -0
  76. ace_jax-0.1.0/src/ace_jax/fit/stats.py +531 -0
  77. ace_jax-0.1.0/src/ace_jax/fit/summary.py +60 -0
  78. ace_jax-0.1.0/src/ace_jax/fit/varpro.py +49 -0
  79. ace_jax-0.1.0/src/ace_jax/fit/weights.py +139 -0
  80. ace_jax-0.1.0/src/ace_jax/fit/xyz.py +84 -0
  81. ace_jax-0.1.0/src/ace_jax/fit/zbl.py +52 -0
  82. ace_jax-0.1.0/src/ace_jax/runfile.py +177 -0
  83. ace_jax-0.1.0/src/ace_jax/tutorials/__init__.py +2 -0
  84. ace_jax-0.1.0/src/ace_jax/tutorials/labels.py +181 -0
  85. ace_jax-0.1.0/src/ace_jax/tutorials/structures.py +84 -0
@@ -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
+ [![PyPI](https://img.shields.io/pypi/v/ace-jax)](https://pypi.org/project/ace-jax/)
47
+ [![Docs](https://img.shields.io/github/actions/workflow/status/ACEsuit/ace-jax/docs.yml?branch=main&label=docs)](https://acesuit.github.io/ace-jax/)
48
+ [![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)
49
+ [![License: MIT](https://img.shields.io/badge/license-MIT-blue)](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).
@@ -0,0 +1,92 @@
1
+ # ace-jax
2
+
3
+ [![PyPI](https://img.shields.io/pypi/v/ace-jax)](https://pypi.org/project/ace-jax/)
4
+ [![Docs](https://img.shields.io/github/actions/workflow/status/ACEsuit/ace-jax/docs.yml?branch=main&label=docs)](https://acesuit.github.io/ace-jax/)
5
+ [![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)
6
+ [![License: MIT](https://img.shields.io/badge/license-MIT-blue)](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
+ """