Load and write PACE `.yace` models¶

ace-jax reads pacemaker .yace potential files and evaluates them in JAX. It uses the same calculator, CLI and LAMMPS export as for its own models. The results agree with the ML-PACE C++ code, python-ace and LAMMPS.

Evaluate a `.yace`¶

import jax
jax.config.update("jax_enable_x64", True)
from ace_jax import ACECalculator

atoms.calc = ACECalculator("model.yace")
atoms.get_potential_energy(); atoms.get_forces()

From the shell:

aj eval --model model.yace --data test.xyz --out predictions.xyz

Supported features¶

Feature	Supported
radial basis	ChebExpCos, ChebPow, ChebLinear, SBessel
embedding	FinnisSinclair, FinnisSinclairShiftedScaled
inner cutoff	`density`, `distance`, `zbl`

If a file uses a different feature, ace-jax gives an error when it loads the file. It does not evaluate the file incorrectly.

Load, modify and write back¶

ace_jax.load returns a PACEModel, its metadata and the parsed file. The model is a JAX pytree of arrays. Thus you can change it with equinox tools. write_yace writes it to a file:

import ace_jax as aj
from ace_jax.eval import write_yace

pm, meta, spec = aj.load("model.yace")       # PACEModel, meta dict, parsed spec
write_yace(pm, "copy.yace", spec=spec)       # numbers from pm, layout from spec

write_yace takes the numeric values from the model. It takes all other data (the function layout, the element names, the file structure) from spec. Thus it never makes the basis layout again. If you write a model without changes, it gives the same results as the original.

The aj export command writes the same file:

aj export --model model.yace --out copy.yace

Write an ACE model as a .yace file¶

write_yace also writes a linear ACE model (an .npz model) as a .yace file. The file gives the same energies, forces and stresses as the .npz model.

import ace_jax as aj
from ace_jax.eval import write_yace

model, meta, _ = aj.load("model.npz")         # a linear ACE model
info = write_yace(model, "model.yace")        # returns the sizes of the file
pm, meta, spec = aj.load("model.yace")        # ace-jax reads the file back

The aj export command writes the same file:

aj export --model model.npz --out model.yace

The file keeps the radial functions as tables of values. The table spacing is 0.001 Å. Set a different spacing with delta=. The cutoff must be a whole number of table spacings.

Warning

ML-PACE reads these tables only with its tabulated-radial addition. That addition is not yet in ML-PACE releases. Without it, LAMMPS refuses the file. ace-jax reads the file without it. For the status of the addition, see ML-PACE pull request 31.

Note

A model with more than one element gives a larger file. Each product of the model becomes one product for each combination of neighbour elements.

Limits¶

Use .yace models only for evaluation and export. aj fit and the radial learner need an ACE .npz model.
The return value of aj.load(path) depends on the file extension: .yace gives (PACEModel, meta, spec), .npz gives (ACEModel, meta, arrays).
For an SBessel radial with nradbase >= 12, ace-jax selects a matrix form at load time (PACEModel.sbessel_form == "matmul"). The values agree with the recurrence to round-off.

Load and write PACE .yace models¶

Evaluate a .yace¶