Command-line reference¶
ace-jax and its short alias aj are the same program. The documentation
build makes this page from the --help output of the program. Thus the page
agrees with the installed version.
The commands, in the order of a typical workflow:
aj fitfits a model to labelled extxyz data. It builds the basis from--order/--max-degree, or uses a saved basis from--model. It writes the fitted model, the metrics and afit.yamlrun file to--out.aj evalevaluates a fitted model (ACE.npz,gp_model.npzor PACE.yace) on an extxyz file.aj calibraterecalibrates the per-atom force uncertainty of a--uq ardfit on labelled target-regime data.aj basisbuilds a basis and saves it on its own.
Each command exits with status 0 if it is successful. The default label
keys are energy, forces and virial. For other names, use
--energy-key, --force-key and --virial-key.
usage: ace-jax [-h] [--version] {fit,eval,calibrate,basis} ...
Fit and evaluate ACE models in JAX (short alias: aj)
positional arguments:
{fit,eval,calibrate,basis}
fit fit the hybrid linear-ACE + residual GP (--m-per-species 0 =
linear-only fit)
eval evaluate a model on a dataset: RMSE table vs the labels,
predictions to extxyz
calibrate recalibrate the per-group conformal scales of an ARD posterior
on a labelled set (per-group replace by default)
basis author a new ACE basis: a frozen model (seeded radial init)
saved as .npz
options:
-h, --help show this help message and exit
--version show program's version number and exit
aj fit¶
The options fall into groups:
- data:
--train/--testfiles, or one--datafile divided by a seeded permutation (--ntrain,--ntest,--test-start).--oodadds a second test set.--weightssets weights for each configuration type. - basis: the options under "basis" below build the basis in the fit.
Alternatively,
--modeluses a saved basis; then--r0is necessary. - model:
--m-per-species 0is linear ACE. A positive value adds a Gaussian-process arm, with that number of inducing sites for each species. The default (500) is the GP arm. - hyperparameters:
--opt,--map-steps,--map-restarts,--map-polish,--strict,--init,--noise.--noise sharedlearns one noise scale for all weighted rows, so--weightssets the balance of E, F and V (as in ACEpotentials).--rungsother thanmapadds approximations of the hyperparameter posterior, at a much higher cost. - uncertainty:
--uq blr(default),popsorardfor the linear model. The--ard-*and--force-shapeoptions configureard(Per-atom force uncertainty). - run file:
--config fit.yaml. Command-line flags override the file (Run files).
usage: ace-jax fit [-h] [--config CONFIG] [--model MODEL]
[--train TRAIN | --data DATA] [--test TEST] [--ood OOD]
[--ntrain NTRAIN] [--ntest NTEST] [--test-start TEST_START]
[--energy-key ENERGY_KEY] [--force-key FORCE_KEY]
[--virial-key VIRIAL_KEY] [--stress-key STRESS_KEY]
[--weights WEIGHTS] [--e0 {model,lsq,prefit}] [--baseline BASELINE]
[--configs-per-batch CONFIGS_PER_BATCH]
[--batch-pack {auto,on,off}] [--m-per-species M_PER_SPECIES]
[--kernel {cosine,matern32}] [--no-bump]
[--density {none,pair,pca}] [--pca-d PCA_D] [--embedding EMBEDDING]
[--objective {lml,loo}] [--lml {device,host-cache}]
[--lml-solver {qr,cholesky}] [--opt {adam,lbfgs}]
[--map-polish {auto,on,off}] [--strict]
[--noise {per-quantity,shared}] [--solver {evidence,lstsq}]
[--map-restarts MAP_RESTARTS] [--init INIT] [--rungs RUNGS]
[--n-draws N_DRAWS] [--laplace {svi,fd}] [--map-steps MAP_STEPS]
[--vi-steps VI_STEPS] [--nuts-warmup NUTS_WARMUP]
[--nuts-samples NUTS_SAMPLES] [--nuts-chains NUTS_CHAINS] [--r0 R0]
[--uq {blr,pops,ard}] [--ard-mode {joint,sequential}]
[--ard-variance {sandwich,kappa}] [--ard-val-frac ARD_VAL_FRAC]
[--force-shape {iso,aniso}] [--ard-coverage ARD_COVERAGE]
[--ard-groups {distortion,none}]
[--ard-cluster-size ARD_CLUSTER_SIZE] [--ard-press {exact,block}]
[--ard-transfer {exponent,sqrt,none}] [--ard-n-min ARD_N_MIN]
[--no-ard-support] [--ard-support-features {raw,normalised}]
[--learn-radial] [--radial-n-q RADIAL_N_Q]
[--radial-steps RADIAL_STEPS] [--radial-lam-grid RADIAL_LAM_GRID]
[--radial-val-frac RADIAL_VAL_FRAC] [--pops-ridge POPS_RIDGE]
[--seed SEED] [--out OUT] [--model-draws MODEL_DRAWS]
[--no-save-model] [--devices DEVICES] [--elements ELEMENTS]
[--order ORDER] [--max-degree MAX_DEGREE] [--wL WL] [--rcut RCUT]
[--rin RIN] [--radial-mode {onehot,glorot_normal,zero}]
[--pair-mode PAIR_MODE] [--no-gamma] [--no-coupling-cache]
[--coupling-cache-dir COUPLING_CACHE_DIR]
[--basis-embedding BASIS_EMBEDDING] [--d-max D_MAX] [--maxl MAXL]
[--reduction {pca,truncate}]
options:
-h, --help show this help message and exit
--config CONFIG a fit.yaml run file (flag names as keys, the basis in a
`basis:` block); command-line flags override it
--model MODEL an ACE basis/model .npz (or give --order/--max-degree to build
the basis)
--train TRAIN training extxyz (with --test, or tested on itself)
--data DATA one extxyz split by a seeded permutation
(--ntrain/--ntest/--test-start)
--test TEST
--ood OOD extra out-of-distribution test extxyz
--ntrain NTRAIN
--ntest NTEST
--test-start TEST_START
--energy-key ENERGY_KEY
--force-key FORCE_KEY
--virial-key VIRIAL_KEY
--stress-key STRESS_KEY
stress label (eV/A^3, 3x3 or Voigt-6) used as the virial,
virial = -stress * volume, when a config has no virial label
--weights WEIGHTS JSON: an ACEfit weights dict {"default":
{"E":..,"F":..,"V":..}, <config_type>: ..} or a list of weight
factors [{"Structural": {}}, {"ConfigType": {...}}]
--e0 {model,lsq,prefit}
per-species E0: the model's (default); lsq fits it jointly
with the readout (a wide prior around a least-squares start);
prefit fixes it at least squares on the training energies
before the fit
--baseline BASELINE dimer_mean.npz: fit the residual to this pair mean
--configs-per-batch CONFIGS_PER_BATCH
--batch-pack {auto,on,off}
size-aware batching: pack configs into batches by an atom
budget (the largest config) instead of fixed groups of
--configs-per-batch; auto = only when the fixed layout would
pad badly (mixed bulk + big cells)
--m-per-species M_PER_SPECIES
--kernel {cosine,matern32}
--no-bump
--density {none,pair,pca}
GP feature map: full descriptor, pair densities, or a PCA view
(--pca-d)
--pca-d PCA_D
--embedding EMBEDDING
MACE element table (JSON): frozen species coregionalization
--objective {lml,loo}
--lml {device,host-cache}
host-cache: cache the linear design rows in host RAM (GP,
pair|pca, L-BFGS, map only)
--lml-solver {qr,cholesky}
linear arm: the evidence and posterior by QR (stable; default)
or by a Cholesky of the Gram (faster at large basis, but loses
accuracy when the fit nearly interpolates its data). The GP
arm always uses the Cholesky
--opt {adam,lbfgs} hyperparameter MAP optimiser: lbfgs (default) = bounded
L-BFGS-B, then (linear arm, see --map-polish) a Newton polish
to a stationary point; adam = numpyro SVI for --map-steps
steps, which can stop far from the optimum on large data
--map-polish {auto,on,off}
Newton polish of the L-BFGS MAP to a stationary point, with
the exact Hessian built one Hessian-vector product per free
hyperparameter (~2.3x the gradient's memory). auto (default):
on for the linear arm, off for the GP arm, where each HVP is a
forward-over-reverse pass through the streamed objective (one
gradient can take a minute, a polish hours at large scale)
--strict fail, instead of warning, when the MAP ends away from a
stationary point (a predicted Newton gain above 1e-3 nats)
--noise {per-quantity,shared}
noise hyperparameters: per-quantity (default) learns sigma_E,
sigma_F and sigma_V separately, and at the optimum each
cancels its quantity's --weights; shared learns ONE sigma for
every weighted row (ACEpotentials' BLR), so the E:F:V weights
set the balance. Use shared whenever the weights are meant to
set it (an ACEpotentials weights dict, e.g. E 30 / F 1 / V 1):
per-quantity lets the many force rows push sigma_E up and
underweight the energies
--solver {evidence,lstsq}
lstsq: plain weighted least squares with no prior (teaching;
overfits a large basis)
--map-restarts MAP_RESTARTS
L-BFGS multi-start (best log-posterior)
--init INIT theta_map.json to start the MAP from
--rungs RUNGS comma-separated from map,laplace,pathfinder,vi,nuts (default
map; the others add hyperparameter draws and cost far more
than the MAP)
--n-draws N_DRAWS
--laplace {svi,fd}
--map-steps MAP_STEPS
MAP iterations: L-BFGS-B iterations (at most 4x as many
evaluations), or Adam steps
--vi-steps VI_STEPS
--nuts-warmup NUTS_WARMUP
--nuts-samples NUTS_SAMPLES
--nuts-chains NUTS_CHAINS
--r0 R0 typical nearest-neighbour distance (A); centres the GP
hyperprior (default when building the basis: its mean bond
length)
--uq {blr,pops,ard} pops/ard: linear arm (--m-per-species 0); ard = ARD posterior
with a calibrated per-atom forces_std (see --ard-variance),
writes posterior.npz
--ard-mode {joint,sequential}
joint: noise + ARD scales by evidence; sequential: ARD only,
one Gram (low memory)
--ard-variance {sandwich,kappa}
ARD force-uncertainty shape: sandwich (default) = delete-one-
cluster PRESS jackknife (misspecification-robust); kappa = the
posterior A^-1 shape. Both get the per-group scales
--ard-val-frac ARD_VAL_FRAC
stratified train fraction held out to score the per-group
force scales (rms factor and conformal quantile) with the
hold-out posterior
--force-shape {iso,aniso}
ard: anisotropic (Mahalanobis, default) or isotropic
(|e|/sqrt(v/3), spherical radius) conformal force scores
--ard-coverage ARD_COVERAGE
ard: nominal coverage 1 - alpha of the per-group conformal
quantile
--ard-groups {distortion,none}
ard: conformal groups = distortion bands x [z = z*] (8), or
none (2)
--ard-cluster-size ARD_CLUSTER_SIZE
ard: sandwich block side in units of r_cut ('inf': whole
configurations)
--ard-press {exact,block}
ard: PRESS correction of the jackknife scores (exact: per-
cluster solve; block: block approximation)
--ard-transfer {exponent,sqrt,none}
ard: carry the hold-out scales to the served posterior by
(N/N_fit)^beta -- exponent: beta fitted per run from a second,
smaller hold-out fit (clipped to [0, 1/2]); sqrt: beta = 1/2;
none: beta = 0
--ard-n-min ARD_N_MIN
ard: groups with fewer configurations borrow a neighbouring
group's scales
--no-ard-support ard: skip the covariate-shift support reference
--ard-support-features {raw,normalised}
ard: support-reference features -- raw descriptors, or
normalised (unit-norm descriptor plus log-norm channels per
body order, which flag atoms losing neighbours)
--learn-radial learn the tensor radials (VarPro, held-out gate) before the
fit; the saved model is marked radial_learned and splined at
deploy time
--radial-n-q RADIAL_N_Q
tensor-radial polynomial span after widening (with --learn-
radial)
--radial-steps RADIAL_STEPS
L-BFGS steps per roughness weight (with --learn-radial)
--radial-lam-grid RADIAL_LAM_GRID
comma-separated relative roughness weights the gate picks
among, alongside the initial radials (with --learn-radial)
--radial-val-frac RADIAL_VAL_FRAC
fraction of the training configs held out for the gate (with
--learn-radial)
--pops-ridge POPS_RIDGE
--seed SEED
--out OUT
--model-draws MODEL_DRAWS
GP arm: hyperparameter draws stored in gp_model.npz (1 = the
MAP; more = evenly spaced draws of the last rung, each adding
a (Dt, Dt) factor)
--no-save-model skip writing the fitted model (model.npz linear / gp_model.npz
GP)
--devices DEVICES shard sufficient statistics over this many devices
(experimental)
basis (built in memory; instead of --model):
--elements ELEMENTS comma-separated Z numbers or symbols (default: the species in
the data)
--order ORDER correlation order
--max-degree MAX_DEGREE
TotalDegree level bound
--wL WL
--rcut RCUT cutoff (default 5.5; with an embedding, 2.5 x mean bond
length)
--rin RIN
--radial-mode {onehot,glorot_normal,zero}
initial tensor radials: onehot (R_n = P_n, the frozen-fit
default) or seeded glorot_normal mixtures (a random start,
e.g. for learned radials)
--pair-mode PAIR_MODE
--no-gamma skip the smoothness prior
--no-coupling-cache always compute the coupling instead of using the per-shape
cache
--coupling-cache-dir COUPLING_CACHE_DIR
coupling cache directory (default: $ACEJAX_COUPLING_CACHE or
~/.cache/ace-jax/coupling)
--basis-embedding BASIS_EMBEDDING
frozen element embedding of the basis: a JSON table {Z, emb}
or 'identity' (builds ace_embedding_model: ace1-compatible,
factorised radial)
--d-max D_MAX cap on per-order channel widths (default lossless)
--maxl MAXL
--reduction {pca,truncate}
aj eval¶
With --posterior (a posterior.npz from aj fit --uq ard), --per-atom
writes the per-atom force uncertainty, and --support adds the support flag
(Per-atom force uncertainty).
usage: ace-jax eval [-h] --model MODEL --data DATA [--energy-key ENERGY_KEY]
[--force-key FORCE_KEY] [--virial-key VIRIAL_KEY]
[--stress-key STRESS_KEY] [--out OUT] [--prefix PREFIX]
[--posterior POSTERIOR] [--per-atom PER_ATOM] [--support]
[--shape-path {rows,committee}] [--no-deriv-dtc]
options:
-h, --help show this help message and exit
--model MODEL
--data DATA
--energy-key ENERGY_KEY
--force-key FORCE_KEY
--virial-key VIRIAL_KEY
--stress-key STRESS_KEY
stress label (eV/A^3, 3x3 or Voigt-6) used as the virial,
virial = -stress * volume, when a config has no virial label
--out OUT extxyz to write: the input structures, every label kept, plus
the predictions (<prefix>energy, <prefix>forces,
<prefix>stress, and *_std for UQ models)
--prefix PREFIX name prefix of the predicted keys (default ace_)
--posterior POSTERIOR
posterior.npz from `fit --uq ard`: adds <prefix>forces_std
--per-atom PER_ATOM with --posterior: extxyz of the served per-atom arrays
(forces_pred, forces_std, forces_q, forces_group; forces_cov,
forces_q_mahal for an aniso posterior)
--support with --posterior --per-atom: add support_ok and support_q
(covariate-shift support)
--shape-path {rows,committee}
ard: evaluate the uncertainty shape from the whole cell's
force design rows (rows) or as the forces of an r-output
linear ACE (committee: same values, memory O(N r) not O(N L))
--no-deriv-dtc gp_model.npz only: SoR-only forces_std, without the
derivative-DTC term (whose whole-cell (n, K, d, 3) arrays may
not fit for a big cell)
aj calibrate¶
Calculates the group scales of an ARD posterior again, on labelled configurations that were not in the training set. It writes a new posterior. The model does not change. See Recalibrate on target data.
usage: ace-jax calibrate [-h] --model MODEL --posterior POSTERIOR --data DATA
[--energy-key ENERGY_KEY] [--force-key FORCE_KEY]
[--virial-key VIRIAL_KEY] [--shape-path {rows,committee}]
[--append | --replace] --out OUT
options:
-h, --help show this help message and exit
--model MODEL
--posterior POSTERIOR
--data DATA labelled extxyz U (needs forces under --force-key)
--energy-key ENERGY_KEY
--force-key FORCE_KEY
--virial-key VIRIAL_KEY
--shape-path {rows,committee}
ard: evaluate the uncertainty shape from the whole cell's
force design rows (rows) or as the forces of an r-output
linear ACE (committee: same values, memory O(N r) not O(N L))
--append pool U with the stored T_val scores in every group
--replace use U only in every group
--out OUT
aj basis¶
usage: ace-jax basis [-h] --elements ELEMENTS --order ORDER --max-degree MAX_DEGREE
[--wL WL] [--rcut RCUT] [--rin RIN]
[--radial-mode {onehot,glorot_normal,zero}]
[--pair-mode PAIR_MODE] [--no-gamma] [--no-coupling-cache]
[--coupling-cache-dir COUPLING_CACHE_DIR] [--embedding EMBEDDING]
[--d-max D_MAX] [--maxl MAXL] [--reduction {pca,truncate}]
[--seed SEED] --out OUT [--config CONFIG]
options:
-h, --help show this help message and exit
--elements ELEMENTS comma-separated Z numbers or symbols
--order ORDER correlation order
--max-degree MAX_DEGREE
TotalDegree level bound
--wL WL
--rcut RCUT cutoff (default 5.5; with an embedding, 2.5 x mean bond
length)
--rin RIN
--radial-mode {onehot,glorot_normal,zero}
initial tensor radials: onehot (R_n = P_n, the frozen-fit
default) or seeded glorot_normal mixtures (a random start,
e.g. for learned radials)
--pair-mode PAIR_MODE
--no-gamma skip the smoothness prior
--no-coupling-cache always compute the coupling instead of using the per-shape
cache
--coupling-cache-dir COUPLING_CACHE_DIR
coupling cache directory (default: $ACEJAX_COUPLING_CACHE or
~/.cache/ace-jax/coupling)
--embedding EMBEDDING
frozen element embedding of the basis: a JSON table {Z, emb}
or 'identity' (builds ace_embedding_model: ace1-compatible,
factorised radial)
--d-max D_MAX cap on per-order channel widths (default lossless)
--maxl MAXL
--reduction {pca,truncate}
--seed SEED
--out OUT
--config CONFIG a fit.yaml: its `basis:` block (and `seed`); command-line
flags override it