API reference

dACE

Directional atomic cluster expansion: pure body-ordered, rotation-equivariant expansion of an agent's acceleration in its own velocity and its neighbours' relative positions and velocities.

ARModel(model, ρ)

aᵢ(t) = F(envᵢ(t)) + ηᵢ(t), ηᵢ(t) = ρ ηᵢ(t - Δ) + ξ: force law F = model and persistent individual noise (from fit_ar1). predict(m, env): the force law; predict(m, env, lagged_env, lagged_acc): one step ahead given the environment and acceleration a step Δ earlier.

CompositeBasis(interaction, selfterm)

The columns of interaction followed by those of another self model selfterm, both evaluated on the same environments, fitted and predicted as one linear model (fit, design_statistics, predict). selfterm must implement length, evaluate!(out, selfterm, env, work), evaluation_workspaces(selfterm) and smoothness_prior(selfterm; p), and basis_history(selfterm) if it needs past states (for rollout); dACE's own self term is part of dACEBasis(; history) and needs no composite.

Environment(v, r, u; past_v = [], past_vn = [], dt = NaN)

v = vᵢ; r[j] = xⱼ - xᵢ, u[j] = vⱼ - vᵢ for neighbours within the cutoff (2D). For bases with history (see History): past_v[k] = vᵢ(t - kΔ), the agent's past velocities, past_vn[j] = vⱼ(t - Δ), each neighbour's velocity one step earlier (empty: none), dt = Δ; filled by trajectory_data(...; history). A neighbour at r[j] = 0 has no direction and is an error (coincident positions, e.g. merged tracking identities, have to be removed beforehand).

Environment3(v, r, u; past_v = [], past_vn = [], dt = NaN)

v = vᵢ; r[j] = xⱼ - xᵢ, u[j] = vⱼ - vᵢ for neighbours within the cutoff (3D); past_v, past_vn, dt: past velocities of the agent and (one step earlier) of each neighbour, as for Environment. A neighbour at r[j] = 0 is an error, as in 2D.

EvaluationWorkspace(basis)

Reusable storage for evaluate!(out, basis, env, workspace). Construct one for each basis and concurrent task; a workspace must not be shared by simultaneous evaluations.

FrameBasis(; maxdeg = (4, 8, 6), charge = 1, cutoff = 5.0, rscale = cutoff / 2,
           vscale = 2.0, ascale = 1.0, maxm = 1, maxl = 0, maxa = 2, maxb = 1, vertical = false,
           history = nothing, inputmap = :linear, normaliser = nothing)

3D basis in a frame (e₁, e₂, e₃): the heading frame e₃ = v̂ᵢ (default), or with vertical = true the fixed frame e₃ = ẑ (symmetry: rotations about the vertical only, e.g. under gravity). charge = 1: the part of the acceleration in the plane a₁ + i a₂; charge = 0: the part along the axis, a₃ (built in pairs by heading_bases and vertical_bases). Centres Pₚ(2|ṽ| - 1) (heading frame) or Pₚ(2|ṽ⊥|² - 1) P_q(ṽ₃) ṽ⊥^m of charge m (vertical frame).

Degrees, compared with dACEBasis: maxdeg[n] is the total degree of the n-body terms, as in 2D. Per neighbour, maxm bounds the angular indices in the plane (as in 2D), maxl the degree in the in-plane relative speed (as maxl in 2D), maxa the degree in the axial direction of the neighbour (t = r̂ ⋅ e₃), maxb the degree in its axial relative velocity (ũ₃), the two directions 2D does not have. history, ascale, inputmap, normaliser: as for dACEBasis; past states enter in the frame (history centres ã⊥^k ã₃^j, neighbour factors b̃⊥^k b̃₃^j, |k| + j ≤ interaction_degree; the self term in the frame of the basis).

FrameModel(plane, c_plane, axis, c_axis, λ = NaN, λself = λ)

aᵢ = Re(a⊥) e₁ + Im(a⊥) e₂ + a∥ e₃ in the frame of the bases, a⊥ = Σ c_plane B_plane, a∥ = Re Σ c_axis B_axis. λ, λself: the regularisation of the fit, one per part (plane = ..., axis = ...).

History(; selfterm = nothing, interactions = (;), use = :velocity,
        selfterm_degree = (8, 4), selfterm_maxm = 2, interaction_degree = 1)

The past states each part of a model sees, in sampling steps (one step: the environments' dt, e.g. K video frames; 1:4 means 1 to 4 steps ago). Anything left out sees no past; the current velocities of the agent and its neighbours are always used. Passed to a basis as history.

• selfterm = steps: the self term, the agent's own past velocities at steps, in a temporal body-order expansion in the frame of the current heading (functions of one past state, total degree ≤ selfterm_degree[1], and products over two, ≤ selfterm_degree[2]; angular index ≤ selfterm_maxm). It is part of the 1-body term and acts whether or not there are neighbours.
• interactions = (own = 1, neighbours = 1): the reactions to neighbours (the 2-body and higher terms) also see the agent's own state one step ago (own: the reaction may depend on what the agent was just doing), and each neighbour's state one step ago (neighbours), polynomials of degree ≤ interaction_degree. One step only.
• use = :velocity (default), :acceleration ((v(t - (k - 1)Δ) - v(t - kΔ)) / Δ, with v(t) the current velocity) or :both: what a past state is. Velocities are scaled by the basis vscale, accelerations by the basis ascale. At one step, given the current velocity, a past velocity and acceleration carry the same information, so in interactions :both means the acceleration.

Fill the environments with trajectory_data(...; history) or environments(...; history, past_velocities, dt).

RolloutDivergence(step)

Thrown by rollout when the velocities become non-finite at step: the model has left the region where it predicts finite accelerations (typically states unlike its training data).

dACEBasis(; maxdeg = (2, 8, 6), cutoff = 5.0, rscale = cutoff / 2, vscale = 2.0, ascale = 1.0,
          maxl = 2, maxm = 2, inner_envelope = 0, filter = (ψ, φs) -> true, pure = true,
          history = nothing, inputmap = :linear, charge = 1, normaliser = nothing)

Basis of total charge charge (1: the acceleration aₓ + i a_y; 0: invariants). maxdeg[n]: total polynomial degree of the n-body terms, so length(maxdeg) is the highest body order ((4, 6, 4): 1-, 2- and 3-body). rscale: length scale of the radial map y(ρ) (small values resolve short distances finely). maxl, maxm bound the degree in the neighbour's relative speed (l) and the angular indices |mr|, |mu| of its position and velocity. inner_envelope = q: radial envelope fc(ρ) ((1 - y(ρ))/2)^q, vanishing at ρ = 0 for q > 0. pure = false: self-interacting products ψ ∏ₜ Σⱼ φ_{kₜ}(xⱼ) on the same index set. filter(ψ, φs): keeps the functions with centre ψ and one-particle factors φs for which it is true (socialforce_basis). history: a History (which past states each part of the model sees; nothing: none); its self term follows the interaction functions as part of the 1-body term. inputmap: map from velocities (scale vscale) and accelerations (scale ascale) to the scaled inputs of the polynomials; :tanh keeps all magnitudes below 1, where the polynomials are resolved evenly (see scaled). normaliser = (c, p): every 2-body and higher function is divided by (c + S₀)^p, S₀ = Σⱼ fc(ρⱼ) (fixed crowding normalisation; p = 1: averaging over neighbours, p = 1/2: Finnis–Sinclair-like).

dACEModel(basis, c, λ = NaN, λself = λ)

aₓ + i a_y = Σₖ cₖ Bₖ(env), c ∈ ℂᴷ, for a dACEBasis (with or without history) or a CompositeBasis. λ, λself: the regularisation of the fit (NaN when not from fit).

body_order_contributions(model, env)

Contributions of the 1-, 2-, 3-body … terms ([n]: the n-body term, the self term in the 1-body term); they sum to predict(model, env).

body_orders(basis)

Body order of every column of basis: 1 for the 1-body terms (including the self term), 2 for the pair terms, ….

design_statistics(basis, envs, accs; chunksize = 2000) -> QRStatistics

Streamed QR statistics of the design of basis and the accelerations accs (one data set, e.g. one video or trial). Statistics of separate data sets pool with +, so a cross-validation needs one pass over each data set: fit(basis, S_train + S_validation; λ...) after select_lambda(S_train, S_validation, basis; λs). For the 3D bases (plane, axis): a pair of statistics (FrameStatistics).

environments(pos, vel; cutoff, extent = nothing, periodic = false, history = nothing, past_velocities = [], dt = NaN)

Environments of all agents of one snapshot (Environment in 2D, Environment3 in 3D); minimum image in a box extent if periodic. For a basis with a History (history = basis.history): past_velocities[k] are the velocities of all agents k sampling steps of dt earlier (same layout as vel, at least history_length(history) of them); each environment keeps the agent's past velocities that the history uses, and its neighbours' velocities one step earlier if it uses interactions = (neighbours = 1,).

evaluate!(out, basis, env, workspace)

Evaluate into out using an EvaluationWorkspace constructed for basis. The output and workspace are overwritten; the basis and environment are unchanged. evaluate!(workspace, basis, env) evaluates into the workspace's own output vector and returns it.

evaluate(basis, env)

All basis functions at env (purified).

evaluate(basis, P::Purification, env)

All basis functions at env via P.

fit((plane, axis), envs, accs; λ = :evidence, λself = λ, prior = smoothness_prior.((plane, axis)),
    reflections = false, self_fit = :staged, method = :qr, chunksize = 2000)
fit((plane, axis), S::FrameStatistics; λ, λself, prior, reflections, self_fit)

fit of the parts of accs in the plane and along the axis of each agent's frame, as the 2D fit; prior: one weight vector per part. λ, λself: :evidence (default, chosen per part), a number, or one per part (plane = ..., axis = ...) (as returned by select_lambda). reflections = true: real coefficients, the model also equivariant under reflections through planes containing the axis.

fit(basis, envs, accs; λ = :evidence, λself = λ, prior = smoothness_prior(basis), reflections = false,
    self_fit = :staged, method = :qr, chunksize = 2000)
fit(basis, S::QRStatistics; λ, λself, prior, reflections, self_fit)
fit(basis, envs, accs; solver)

min Σₛ |aₛ - Σₖ cₖ Bₖ(envₛ)|² + λ² Σₖ γₖ |c̃ₖ|² (c̃: column-normalised basis, γ = prior: by default the smoothness prior, prior = nothing for equal weights). λ = :evidence (default): the regularisation that maximises the evidence (marginal likelihood) of the fit, as Bayesian ridge regression in ACE, stage by stage for a staged fit, without validation data (method = :qr); a number fixes it. The model records the λ and λself it was fitted with (m.λ). reflections = true: c ∈ ℝᴷ, i.e. the model is also equivariant under reflections (most collective-motion models are; the default false allows chiral rules). method = :qr (default): streamed QR of the design, one chunk of environments at a time, so that memory does not grow with the data and the condition number is that of the design; method = :normal_equations: the normal equations accumulated over chunks, which square the condition number but are cheaper to accumulate. From streamed statistics S (design_statistics, pooled with +): the same fit without the data, e.g. after select_lambda. With an ACEfit solver: a direct solve of the full design matrix (joint only, reflections honoured, the regularisation that of the solver: λ and λself must be omitted, prior is not used). A CompositeBasis is fitted by the streamed QR only (no method, no solver).

self_fit (for a basis with a History): :staged (default) fits the interaction functions first, alone, with λ, then the self term on their residual with λself, then the reactions to neighbours that see the agent's own past (interactions = (own = ...,)) on what is left, with λ. A self term with velocity history can absorb interaction signal (the two are not identified), so the staged fit keeps the interactions of the fit without any past. :self_first: the self term first, then the interactions, then the own terms (the self term absorbs whatever interaction signal it can). :joint: all columns in one least-squares problem, as ACE fits a potential.

fit_ar1(basis, envs, accs, lagged_envs, lagged_accs; λ = :evidence, prior = smoothness_prior(basis),
        reflections = false, method = :qr, chunksize = 2000) -> ARModel

Diagnostic, not for interaction laws: the force law fitted jointly with persistent individual noise η(t) = ρ η(t - Δ) + ξ by Cochrane–Orcutt iteration (lagged_envs, lagged_accs: the same agents one step Δ earlier, as from trajectory_data(...; lagged = true)); all columns jointly. The fitted ρ shows whether residuals persist. With tracked positions the previous finite-difference acceleration shares frames with the target, so the one-step prediction of an ARModel partly predicts tracking error, and as ρ → 1 the force law is not identified: judge the force law (predict(m.model, env)), not the one-step prediction. λ = :evidence (default, method = :qr): each Cochrane–Orcutt fit chooses λ by the evidence of its quasi-differenced system; m.model.λ is the last one.

fit_path(basis, envs, accs; λs, λself = nothing, prior = smoothness_prior(basis), reflections = false,
         self_fit = :staged, method = :qr, chunksize = 2000)

fit for each λ ∈ λs from one pass over the data: one streamed QR, then one SVD of its triangular factor for a joint fit, or one SVD of the first stage of a staged fit and a ridge solve of the later stages per λ (method = :qr), or one accumulation of the normal equations. Staged fits use λself for the self term (default: the same λ), and select_lambda picks both on validation data.

heading_bases(; kwargs...) -> (plane, axis)

The two FrameBasis of a 3D model in the heading frame: plane (charge 1) for the part of the acceleration across the heading, axis (charge 0) for the part along it.

max_body_order(basis)

The highest body order of basis: length(maxdeg), unless filter removed every function of that order.

predict(model, env)
predict(model, envs::AbstractVector)

The acceleration a fitted model (dACEModel, FrameModel or ARModel) predicts for one environment, as an SVector, or for each of a vector of environments (with one reusable workspace for the batch). For an ARModel, predict(model, env, lagged_env, lagged_acc) adds the persistent noise one step ahead (see ARModel).

reflect(env::Environment3)

Reflection (x, y, z) ↦ (x, -y, z) of all positions and (current and past) velocities.

reflect(env)

Reflection (x, y) ↦ (x, -y) of all positions and (current and past) velocities.

rmse(model, envs, accs; relative = false)

RMS acceleration error, relative to the RMS of accs if relative.

rollout(model, pos0, vel0, dt, nsteps; extent = nothing, periodic = false, past_velocities = nothing) -> (pos, vel)

Synchronous explicit Euler: vᵢ += dt aᵢ, xᵢ += dt vᵢ. Models with a History see the velocities of the previous steps (step dt, so dt must be the sampling step of the training data); past_velocities[k]: the velocities k steps before the start, as for environments (default: vel0, i.e. no past acceleration). Throws a RolloutDivergence if the velocities become non-finite.

rotate(env, α)

Joint rotation of all positions and (current and past) velocities by α.

rotate(env::Environment3, Q)

Joint rotation by Q.

select_basis(bases, envs, accs, val_envs, val_accs; λs, prior = smoothness_prior, reflections = false,
             self_fit = :staged) -> (; basis, λ, λself, index, errors, statistics)

The candidate basis (for instance a ladder of total degrees, dACEBasis(maxdeg = d, ...) for several d) with the lowest relative error on the validation data, each candidate with its own regularisation chosen on the same validation data by select_lambda. prior is a function of the basis (default smoothness_prior) or nothing. errors are the validation errors of all candidates. When the candidates differ only in degree, the design is built once, for the largest. statistics are those of the chosen basis on training and validation data together, for the refit on all data: fit(basis, statistics; λ, λself). Passing a single candidate fixes the basis.

select_lambda(S, Sval, basis; λs, prior = smoothness_prior(basis), reflections = false, self_fit = :staged)
    -> (; λ, λself)

The regularisation that minimises the error on validation statistics Sval of the fit on S, for any self_fit: joint, one λ over the path (λself = λ); staged, stage by stage (select_staged).

smoothness_prior(basis; p = 4)

γₖ = 1 + deg(ψ)ᵖ + Σₜ deg(φₜ)ᵖ for Bₖ = ψ ∏ₜ φₜ (algebraic prior as in ACEpotentials), for a dACEBasis, a FrameBasis or a CompositeBasis.

socialforce_basis(; maxdeg = (4, 8), kwargs...)

aᵢ = s(|vᵢ|) vᵢ + Σⱼ p(ρᵢⱼ) r̂ᵢⱼ + Σⱼ q(ρᵢⱼ, |uᵢⱼ|) uᵢⱼ as a restricted dACE basis.

trajectory_data(pos, vel, dt; cutoff, extent, periodic, times, history = nothing, lagged = false)
    -> (envs, accs) or, with `lagged`, (envs, accs, lagged_envs, lagged_accs)

Environments and aᵢ = (vᵢ(t+1) - vᵢ(t)) / dt for the frames t ∈ times of pos, vel (dt apart). history (a History, e.g. basis.history): the environments carry the past velocities that it uses (frames without enough past are skipped). lagged: also the environments and accelerations one frame earlier, t - 1 (for the AR(1) diagnostic, fit_ar1).

vertical_bases(; kwargs...) -> (plane, axis)

heading_bases in the fixed frame e₃ = ẑ: plane for the horizontal acceleration a_x + i a_y, axis for the vertical a_z.