ACEpotentials.jl API

acefit!(rawdata, model; kwargs...)

provides a convenient interface to fitting the parameters of an ACE model. The data should be provided as a collection of AbstractSystem structures.

Keyword arguments:

• energy_key, force_key, virial_key specify

the label of the data to which the parameters will be fitted.

• weights specifies the regression weights, default is 30 for energy, 1 for forces and virials
• solver specifies the lsq solver, default is BLR (BayesianLinearRegression)
• smoothness specifies the smoothness prior, i.e. how strongly damped parameters corresponding to high polynomial degrees are; is 2.
• prior specifies a covariance of the prior, if nothing then a smoothness prior is used, using the smoothness parameter
• repulsion_restraint specifies whether to add artificial data to the training set that effectively introduces a restraints encouraging repulsion in the limit rij -> 0.
• restraint_weight specifies the weight of the repulsion restraint.
• export_lammps : path to a file to which the fitted potential will be exported in a LAMMPS compatible format (yace)
• export_json : path to a file to which the fitted potential will be exported in a JSON format, which can be read from Julia or Python

site_descriptors(system::AbstractSystem, model::ACEPotential;
                 domain, nlist)

Compute site descriptors for all atoms in system, returning them as a vector of vectors. If the optional kw argument domain is passed as a list of integers (atom indices), then only the site descriptors for those atoms are computed and returned. The neighbourlist nlist can be supplied optionally as a kw arg, otherwise it is recomputed.

function at_dimer(r, z1, z0) : generates a dimer with separation r and atomic numbers z1 and z0. (can also use symbols or strings)

function at_trimer(r1, r2, θ, z0, z1, z2) : generates a trimer with separations r1 and r2, angle θ and atomic numbers z0, z1 and z2 (can also use symbols or strings), where z0 is the species of the central atom, z1 at distance r1 and z2 at distance r2.

function atom_energy(IP, z0) : energy of an isolated atom

  copy_runfit(dest = pwd())

Copies the runfit.jl script and an example model parameter file to dest. If called from the destination directory, use

ACEpotentials.copy_runfit()

This is intended to setup a local project directory with the necessary scripts to run a fitting job.

  copy_tutorial(dest = pwd())

Converts the ACEpotential-Tutorial.jl and ACE+AtomsBase.jl literate notebooks to jupyter notebooks and copies them to the folder dest.

Generate a decohesion curve for testing the smoothness of a potential. Arguments:

• at0 : unit cell
• pot : potential implementing energy

Keyword Arguments:

• dim = 1 : dimension into which to expand
• mult = 10 : multiplicative factor for expanding the cell in dim direction
• aa = :auto : array of stretch values of the lattice parameter to use
• npoints = 100 : number of points to use in the stretch array (for auto aa)

function dimer_energy(pot, r, z1, z0) : energy of a dimer with separation r and atomic numbers z1 and z0 using the potential pot; subtracting the 1-body contributions.

dimers(potential, elements; kwargs...) : Generate a dictionary of dimer curves for a given potential.

• potential : potential to use to evaluate energy
• elements : list of chemical species, symbols for which the dimers are to be computed

The function returns a dictionary Ddim such that D[(s1, s2)] contains pairs or arrays (rr, E) which can be plotted plot(rr, E).

function get_adf(data::AbstractVector{<: Atoms}, r_cut; kwargs...) :

Angular distribution, i.e. list of angles in [0, π] between all pairs of bonds of length at most r_cut. Keyword arguments:

• skip = 3 : only consider every skipth atom in the dataset.
• maxsamples = 100_000 : maximum number of samples to return.

function get_rdf(data::AbstractVector{<: Atoms}, r_cut; kwargs...) :

Produce a list of r values that occur in the dataset, restricted to the cutoff radius r_cut. Keyword arguments:

• rescale = true : resample the data to account for volume scaling, i.e. a distance r will be kept with probability min(1, (r0/r)^2).
• r0 = :min : parameter for resampling. If :min then the minimum r occuring in the dataset is taken.
• maxsamples = 100_000 : maximum number of samples to return.

  make_model(model_dict::AbstractDict)

User-facing script to generate a model from a dictionary. See documentation for details.

  save_model(model, filename; kwargs...)

save model constructor, model parameters, and other information to a JSON file.

• model : the model to be saved
• filename : the name of the file to which the model will be saved
• model_spec : the arguments used to construct the model; without this the model cannot be reconstructed unless the original script is available
• errors : the fitting / test errors computed during the fitting
• verbose : print information about the saving process
• save_project : save Project.toml and Manifest.toml for reproducibility

function trimer_energy(IP, r1, r2, θ, z0, z1, z2) : computes the energy of a trimer, subtracting the 2-body and 1-body contributions.

trimers(potential, elements, r1, r2; kwargs...) : Generate a dictionary of trimer curves for a given potential.

• potential : potential to use to evaluate energy
• elements : list of chemical species, symbols for which the trimers are to be computed
• r1, r2 : distance between the central atom and the first, second neighbour

The function returns a dictionary Dtri such that D[(s1, s2, s3)] contains pairs or arrays (θ, E) which can be plotted plot(θ, E).

The pair basis radial envelope implemented in ACE1.jl

Maps the transform trans to the standardized interval [-1, 1]

mutable struct OneBody{T}

this should not normally be constructed by a user, but instead E0 should be passed to the relevant model constructor, which will construct it.

Takes an object and converts it to an SMatrix{NZ, NZ} via the following rules:

• if obj is already an SMatrix{NZ, NZ} then it just return obj
• if obj is an AbstractMatrix and size(obj) == (NZ, NZ) then it converts it to an SMatrix{NZ, NZ} with the same entries.
• otherwise it generates an SMatrix{NZ, NZ} filled with the value obj.

Helper function to check if any signed combination of mm satisfies |sum| <= L. Adapted from EquivariantModels.jl for migration to EquivariantTensors.jl.

_rpe_filter_real(L::Integer)

Filter for real spherical harmonics basis with equivariance level L. Replacement for EquivariantModels.RPEfilterreal.

Checks:

1. m-quantum number compatibility (via lazy signed m-set)
2. Parity condition: sum(l) + L must be even
3. Special case: for L=0 with single element, l must be 0

function agnesi_transform: constructs a generalized agnesi transform.

trans = agnesi_transform(r0, p, q)

with q >= p. This generates an AnalyticTransform object that implements

\[ x(r) = \frac{1}{1 + a (r/r_0)^q / (1 + (r/r0)^(q-p))}\]

with default a chosen such that $|x'(r)|$ is maximised at $r = r_0$. But a may also be specified directly as a keyword argument.

The transform satisfies

\[ x(r) \sim \frac{1}{1 + a (r/r_0)^p} \quad \text{as} \quad r \to 0 \quad \text{and} \quad x(r) \sim \frac{1}{1 + a (r/r_0)^p} \quad \text{as} r \to \infty.\]

As default parameters we recommend p = 2, q = 4 and the defaults for a.

Get the specification of the BBbasis as a list (Vector) of vectors of @NamedTuple{n::Int, l::Int}.

Parameters

• model : an ACEModel

Set the radial weights as they would be in a linear ACE model.

This is one of the most important functions to generate an ACE model with sparse AA basis. It generates the AA basis specification as a list (Vector) of vectors of @NamedTuple{n::Int, l::Int, m::Int}.

Parameters

• order : maximum correlation order
• r_spec : radial basis specification in the format Vector{@NamedTuple{a::Int64, b::Int64}}
• max_level : maximum level of the basis, either a single scalar, or an iterable (one for each order)
• level : a function that computes the level of a basis element; see e.g. TotalDegree and EuclideanDegree

ETACEPotential

AtomsCalculators-compatible calculator wrapping an ETACE model. This is a type alias for WrappedSiteCalculator{<:ETACE, PS, ST}.

Access underlying components via:

• calc.model - The ETACE model
• calc.ps - Model parameters
• calc.st - Model state
• calc.rcut - Cutoff radius in Ångström
• calc.co_ps - Committee parameters (optional)

Example

calc = ETACEPotential(et_model, ps, st, 5.5)
E = potential_energy(sys, calc)

ETOneBodyPotential

AtomsCalculators-compatible calculator wrapping an ETOneBody model. This is a type alias for WrappedSiteCalculator{<:ETOneBody, PS, ST}.

ETOneBody has no learnable parameters, so training assembly returns empty results:

• length_basis(calc) returns 0
• energy_forces_virial_basis(sys, calc) returns empty arrays
• Forces and virial are always zero (energy only depends on atom types)

Example

et_onebody = one_body(Dict(:Si => -0.846), x -> x.z)
_, st = Lux.setup(rng, et_onebody)
calc = ETOneBodyPotential(et_onebody, nothing, st, 3.0)
E = potential_energy(sys, calc)

ETPairPotential

AtomsCalculators-compatible calculator wrapping an ETPairModel. This is a type alias for WrappedSiteCalculator{<:ETPairModel, PS, ST}.

Supports training assembly functions:

• length_basis(calc) - Total linear parameters
• energy_forces_virial_basis(sys, calc) - Full EFV design row
• potential_energy_basis(sys, calc) - Energy design row
• get_linear_parameters(calc) / set_linear_parameters!(calc, θ)

Example

et_pair = convertpair(model)
ps, st = Lux.setup(rng, et_pair)
calc = ETPairPotential(et_pair, ps, st, 5.5)
E = potential_energy(sys, calc)

struct EnvRBranchL

An auxiliary layer that is basically a branch layer needed to build radial bases, with additional evaluate_ed functionality, needed for Jacobians.

StackedCalculator{N, C<:Tuple}

Combines multiple AtomsCalculators by summing their energy, forces, and virial. Each calculator in the tuple must implement the AtomsCalculators interface.

This allows combining site-based calculators (via WrappedSiteCalculator) with calculators that don't have site decompositions (e.g., Coulomb, dispersion).

The implementation uses compile-time loop unrolling for efficiency when the number of calculators is small and known at compile time.

Example

# Wrap site energy models
E0_calc = WrappedSiteCalculator(E0Model(Dict(:Si => -0.846)))
ace_calc = WrappedSiteCalculator(WrappedETACE(et_model, ps, st, 5.5))

# Stack them (could also add Coulomb, dispersion, etc.)
calc = StackedCalculator((E0_calc, ace_calc))

E = potential_energy(sys, calc)
F = forces(sys, calc)

Fields

• calcs::Tuple - Tuple of N calculators implementing AtomsCalculators interface

WrappedSiteCalculator{M, PS, ST}

Wraps any ETACE-pattern model (ETACE, ETPairModel, ETOneBody) and provides the AtomsCalculators interface.

All wrapped models must implement the ETACE interface:

• model(G, ps, st) → (site_energies, st)
• site_grads(model, G, ps, st) → edge gradients

Mutable to allow parameter updates during training.

Example

# With ETACE model
calc = WrappedSiteCalculator(et_model, ps, st, 5.5)

# With ETOneBody (upstream)
et_onebody = ETM.one_body(Dict(:Si => -0.846), x -> x.z)
_, onebody_st = Lux.setup(rng, et_onebody)
calc = WrappedSiteCalculator(et_onebody, nothing, onebody_st, 3.0)

E = potential_energy(sys, calc)
F = forces(sys, calc)

Fields

• model - ETACE-pattern model (ETACE, ETPairModel, or ETOneBody)
• ps - Model parameters (can be nothing for ETOneBody)
• st - Model state
• rcut::Float64 - Cutoff radius for graph construction (Å)
• co_ps - Optional committee parameters for uncertainty quantification

convert2et_full(model, ps, st; rng=default_rng()) -> StackedCalculator

Convert a complete ACE model (E0 + Pair + Many-body) to an ETACE-based StackedCalculator. This creates a calculator that combines:

1. ETOneBody - reference energies per species
2. ETPairModel - pair potential
3. ETACE - many-body ACE potential

The returned StackedCalculator is fully compatible with AtomsCalculators and can be used for energy, forces, and virial evaluation.

Arguments

• model: ACE model (from ACEpotentials.Models)
• ps: Model parameters (from Lux.setup)
• st: Model state (from Lux.setup)
• rng: Random number generator (default: default_rng())

Returns

• StackedCalculator combining ETOneBody, ETPairModel, and ETACE

Example

model = ace_model(elements=[:Si], order=3, totaldegree=8)
ps, st = Lux.setup(rng, model)
# ... fit model ...
calc = convert2et_full(model, ps, st)
E = potential_energy(sys, calc)

copy_ace_params!(et_ps, ps, model)

Copy many-body (ACE) parameters from ACE model format to ETACE format.

copy_pair_params!(et_ps, ps, model)

Copy pair potential parameters from ACE model format to ETPairModel format. Based on parameter mapping from test/etmodels/test_etpair.jl.

get_linear_parameters(calc::StackedCalculator)

Get concatenated linear parameters from all stacked calculators.

get_linear_parameters(calc::ETACEPotential)

Extract the linear parameters (readout weights) as a flat vector. Parameters are ordered as: [W[1,:,1]; W[1,:,2]; ... ; W[1,:,nspecies]]

get_linear_parameters(calc::ETOneBodyPotential)

Return empty vector - ETOneBody has no learnable parameters.

get_linear_parameters(calc::ETPairPotential)

Extract the linear parameters (readout weights) as a flat vector.

one_body(D::Dict, catfun)

Create a one-body energy model that assigns to each atom an energy based on a categorical variable that is extracted from the atom state via catfun. The dictionary D contains category-value pairs. The one-body energy assigned to an atom with state x is D[catfun(x)].

set_linear_parameters!(calc::StackedCalculator, θ::AbstractVector)

Set linear parameters for all stacked calculators from concatenated vector.

set_linear_parameters!(calc::ETACEPotential, θ::AbstractVector)

Set the linear parameters (readout weights) from a flat vector.

set_linear_parameters!(calc::ETOneBodyPotential, θ::AbstractVector)

No-op for ETOneBody (no learnable parameters).

set_linear_parameters!(calc::ETPairPotential, θ::AbstractVector)

Set the linear parameters (readout weights) from a flat vector.

energy_forces_virial_basis(sys::AbstractSystem, calc::StackedCalculator)

Compute concatenated basis for all stacked calculators.

energy_forces_virial_basis(sys::AbstractSystem, calc::ETACEPotential)

Compute the basis functions for energy, forces, and virial. Returns a named tuple with:

• energy::Vector{Float64} - length = length_basis(calc)
• forces::Matrix{SVector{3,Float64}} - size = (natoms, length_basis)
• virial::Vector{SMatrix{3,3,Float64}} - length = length_basis(calc)

The linear combination of basis values with parameters gives: E = dot(energy, params) F = forces * params V = sum(params .* virial)

energy_forces_virial_basis(sys::AbstractSystem, calc::ETOneBodyPotential)

Return empty arrays - ETOneBody has no learnable parameters.

energy_forces_virial_basis(sys::AbstractSystem, calc::ETPairPotential)

Compute the basis functions for energy, forces, and virial for pair potential.

length_basis(calc::StackedCalculator)

Return total number of linear parameters across all stacked calculators.

length_basis(calc::ETACEPotential)

Return the number of linear parameters in the model (nbasis * nspecies).

length_basis(calc::ETOneBodyPotential)

Return 0 - ETOneBody has no learnable linear parameters.

length_basis(calc::ETPairPotential)

Return the number of linear parameters in the pair model (nbasis * nspecies).

potential_energy_basis(sys::AbstractSystem, calc::StackedCalculator)

Compute concatenated energy basis for all stacked calculators.

potential_energy_basis(sys::AbstractSystem, calc::ETACEPotential)

Compute only the energy basis (faster when forces/virial not needed).

potential_energy_basis(sys::AbstractSystem, calc::ETOneBodyPotential)

Return empty array - ETOneBody has no learnable parameters.

potential_energy_basis(sys::AbstractSystem, calc::ETPairPotential)

Compute only the energy basis for pair potential.