GitHub

dACE.jl

Directional atomic cluster expansion (dACE) learns rotation-equivariant, body-ordered acceleration rules from collective trajectory data. Each agent's acceleration is expanded in pure (self-interaction-free) basis functions of its own velocity and its neighbours' relative positions and velocities. Contributions from each body order can be read off the fitted model.

The 2D basis is equivariant under joint rotations. In 3D, the model uses the agent's heading frame, or a fixed vertical frame for systems with gravity.

Installation

Requires Julia 1.12 or later. The package is not yet registered. Install it directly from GitHub:

using Pkg
Pkg.add(url="https://github.com/TSGut/dACE.jl")

For experiment scripts and tracking-data setup, see Reproducing the paper.

Two dimensions

Build a local environment for each agent from positions and velocities at one time step, then evaluate a dACE basis for one agent:

using dACE

positions = [0.0 1.0 -1.0; 0.0 0.5 0.0]
velocities = [1.0 0.8 0.9; 0.0 0.1 -0.1]
envs = environments(positions, velocities; cutoff = 2.0)

basis = dACEBasis(maxdeg = (2, 3, 3), cutoff = 2.0)
features = evaluate(basis, envs[1])

For repeated evaluation, reuse a workspace and output buffer:

workspace = EvaluationWorkspace(basis)
features = zeros(ComplexF64, length(basis))
evaluate!(features, basis, envs[1], workspace)

A workspace belongs to one basis and should not be shared by concurrent tasks. Batch prediction and fitting manage their own reusable workspaces. Environment construction uses NeighbourLists.jl cell lists, retaining one minimum-image displacement per neighbour and the strict distance < cutoff condition in both 2D and 3D.

Three dimensions

In 3D a model has two parts, each a FrameBasis in a frame (e₁, e₂, e₃) of the agent: plane for the acceleration in the plane across e₃ and axis for the part along it. heading_bases uses the heading frame e₃ = v̂, vertical_bases the fixed frame e₃ = ẑ, for systems with a preferred vertical such as motion under gravity:

using dACE

positions = [0.0 1.0 -1.0; 0.0 0.5 0.0; 0.0 0.0 0.5]
velocities = [1.0 0.8 0.9; 0.0 0.1 -0.1; 0.0 0.0 0.1]
envs = environments(positions, velocities; cutoff = 2.0)

plane, axis = heading_bases(maxdeg = (2, 3), cutoff = 2.0)
features = (evaluate(plane, envs[1]), evaluate(axis, envs[1]))

The degree keywords per neighbour, in 2D and 3D:

keyworddACEBasis (2D)FrameBasis (3D)
maxdeg[n]total degree of the n-body termsthe same
maxmangular indices of the neighbour's position and velocitythe same, in the plane
maxldegree in the neighbour's relative speedthe same, in the plane
maxa—degree in the neighbour's direction along the axis
maxb—degree in its relative velocity along the axis

Fit from trajectories

For trajectory arrays with dimensions (dimension, agent, time), derive environments and acceleration targets, then fit:

envs, accs = trajectory_data(positions2d, velocities2d, dt; cutoff = 2.0)
model = fit(basis, envs, accs; reflections = true)
acceleration = predict(model, envs[1])
contributions = body_order_contributions(model, envs[1])   # [1]: 1-body, [2]: 2-body, …

envs3, accs3 = trajectory_data(positions3d, velocities3d, dt; cutoff = 2.0)
model3d = fit((plane, axis), envs3, accs3; reflections = true)

fit minimises the squared error plus a ridge penalty λ² Σₖ γₖ |c̃ₖ|² on the coefficients of the column-normalised basis.

  • Prior. The weights γ default to the algebraic smoothness prior of ACE, smoothness_prior(basis). prior = nothing gives equal weights.
  • Regularisation. By default λ maximises the evidence (marginal likelihood) of the fit, as Bayesian ridge regression does in ACE, from the same SVD that gives the fit, so it needs no validation data and no grid. A number, fit(...; λ = 1e-3), fixes it. model.λ is the value used.
  • Reflections. reflections = true restricts the coefficients to real values, which makes the model equivariant under reflections as well (in 3D, through planes containing the axis). Most collective-motion rules are. The default false allows chiral rules.
  • Solver. By default fit streams the design through a QR factorisation, one chunk of environments at a time, so memory depends on the basis size only and the condition number is that of the design, not its square. method = :normal_equations accumulates the normal equations instead, which is cheaper per environment but loses accuracy at weak regularisation and needs a numeric λ, since the evidence needs the QR factor.
  • Paths. fit_path returns the fits for a list of λ from a single pass over the data.

Choosing λ on held-out data instead

design_statistics streams one data set (a video, a trial, a group) into QR statistics once. Statistics pool with +, so cross-validation needs no second pass over the data. select_lambda picks the regularisation on validation statistics, and fit accepts statistics in place of data:

S = Dict(g => design_statistics(basis, envs[g], accs[g]) for g in groups)
train, validation = groups[1:end-1], groups[end:end]
choice = select_lambda(sum(S[g] for g in train), sum(S[g] for g in validation), basis;
                       λs = 10.0 .^ (-8:0.5:0), reflections = true)
model = fit(basis, sum(S[g] for g in groups); choice..., reflections = true)

choice holds λ and, for a basis with a self term, λself (see below). For the 3D pair (plane, axis), the statistics, the choice and the fit come per part.

Choosing the degree on held-out data

The total degrees maxdeg are chosen the same way. select_basis takes candidate bases, chooses λ for each on the validation data and returns the one with the lowest validation error. A higher degree resolves a smooth law better when the data constrain it and overfits when they do not. Passing a single basis fixes the degree.

ladder = [dACEBasis(maxdeg = d, maxl = 0, maxm = 1) for d in ((4, 8, 6), (4, 10, 8), (6, 12, 10))]
choice = select_basis(ladder, envs[train], accs[train], envs[val], accs[val]; λs = 10.0 .^ (-8:0.5:0))
model = fit(choice.basis, envs, accs; λ = choice.λ, λself = choice.λself)

choice.errors holds the validation error of every candidate.

Memory, speed maps and normalisation

Animal trajectories often need more than positions and velocities. All options keep the basis exactly equivariant and are available in 2D and 3D:

  • history = History(...): which past velocities each part of the model sees, in sampling steps (one step is the environments' dt, so 1:4 means 1 to 4 steps ago). Anything left out sees no past, and the current velocities are always used.

    part of the modelkeywordsees
    self termselfterm = 1:4the agent's own velocities at those steps
    reactions to neighboursinteractions = (own = 1,)also the agent's own state one step ago
    reactions to neighboursinteractions = (neighbours = 1,)also each neighbour's state one step ago
    • selfterm: a temporal body-order expansion of the agent's own past in the frame of its current heading (functions of one past state and products of two), part of the 1-body term, sized by selfterm_degree and selfterm_maxm. It acts whether or not there are neighbours. Animals and drones carry much of their behaviour here (speed regulation, burst-and-coast swimming, turning bouts), which atoms, without self-propulsion, do not need.
    • own: the reaction to a neighbour may depend on what the agent was just doing (speeding up, slowing down, turning). It acts only through neighbours, and in data where agents nearly always have neighbours it is hard to separate from the self term.
    • use = :velocity (default), :acceleration or :both: what a past state is, for all parts. Past velocities are scaled by the basis vscale, past accelerations by the basis ascale.

    Environments carry past velocities only, from frames up to the current one, so a past acceleration is a difference of velocities and shares no position with a finite-difference target ahead of t (a past acceleration computed from the previous target would leak tracking noise into the prediction). Fill them with trajectory_data(...; history = basis.history), or for one snapshot with environments(...; history, past_velocities, dt). CompositeBasis(interaction, selfterm) combines an interaction basis with another self model.

  • inputmap = :tanh: velocities and accelerations are mapped to tanh(|x| / scale) x / |x| before the polynomials, which resolves heavy-tailed speeds (bursts and glides) and keeps tracking glitches bounded. The default is the linear x / scale.

  • normaliser = (c, p): the 2-body and higher terms are divided by (c + S₀)^p, S₀ = Σⱼ fc(ρⱼ), a fixed normalisation by crowding. Use it to test whether a response is normalised (averaged over neighbours), not as a default: it helps on rules that average over neighbours but biases rules that sum over them, most for many-body terms. Compare fits with and without it on held-out data.

  • inner_envelope = q: radial functions vanish at distance 0 like ρ^q.

using dACE

positions = [0.0 1.0 -1.0; 0.0 0.5 0.0]
velocities = [1.0 0.8 0.9; 0.0 0.1 -0.1]
past_velocities = [[0.9 0.8 1.0; 0.1 0.0 -0.2], [0.8 0.9 1.0; 0.2 0.1 -0.2]]   # one and two steps earlier

history = History(selfterm = 1:2, interactions = (neighbours = 1,), use = :both)
basis = dACEBasis(maxdeg = (4, 4), cutoff = 2.0, history = history, ascale = 2.0, inputmap = :tanh)
envs = environments(positions, velocities; cutoff = 2.0, history, past_velocities, dt = 0.1)
features = evaluate(basis, envs[1])

Fitting a basis with a self term

A self term with velocity history and the interaction terms are not identified against each other: the self term can absorb part of the interaction signal, which raises the one-step prediction while the interactions get worse. fit therefore takes self_fit:

self_fitwhat is fitteduse
:staged (default)the interaction functions alone (regularisation λ), then the self term on their residual (λself, default λ), then the own terms on what is left (λ)interaction laws: the interactions are those of the fit without any past
:self_firstthe self term alone (λself), then the interactions, then the own terms (λ)for comparison, since the self term takes whatever interaction signal it can, so expect worse interactions
:jointall columns in one least-squares problem, as ACE fits a potentialwhen only the prediction matters, or the self term is known not to absorb interactions

self_fit applies to dACEBasis, the 3D bases and CompositeBasis (whose interaction part goes first), to fits from data or from statistics, with method = :qr or :normal_equations, and in fit_path (λself, default the same λ as the interactions). select_lambda chooses λ and λself stage by stage. Fitting the own terms last protects the plain interactions but recovers an own-state-dependent response only partly when the two overlap. :joint recovers it fully, at the cost of that protection. A basis with only one kind of column fits the same either way. An ACEfit solver fits jointly only and needs self_fit = :joint.

accs = [[0.1, 0.0], [0.0, -0.1], [0.05, 0.05]]
model = fit(basis, envs, accs; λ = 1e-3, λself = 1e-4)                    # staged
model_joint = fit(basis, envs, accs; λ = 1e-3, self_fit = :joint)         # ACE-like
dACEModel{dACEBasis{Polynomials4ML.OrthPolyBasis1D3T{5, Float64}, Polynomials4ML.OrthPolyBasis1D3T{3, Float64}, dACE.Purification{Polynomials4ML.OrthPolyBasis1D3T{4, Float64}, Polynomials4ML.OrthPolyBasis1D3T{2, Float64}}, History, dACE.SelfTerm}}(dACEBasis{Polynomials4ML.OrthPolyBasis1D3T{5, Float64}, Polynomials4ML.OrthPolyBasis1D3T{3, Float64}, dACE.Purification{Polynomials4ML.OrthPolyBasis1D3T{4, Float64}, Polynomials4ML.OrthPolyBasis1D3T{2, Float64}}, History, dACE.SelfTerm}(2.0, 1.0, 2.0, 2.0, :tanh, History(1:2, nothing, 1:1, :both, (8, 4), 2, 1), dACE.SelfTerm(2, [(:velocity, 1), (:acceleration, 1), (:velocity, 2), (:acceleration, 2)], [2.0, 2.0, 2.0, 2.0], 2.0, :tanh, dACE.SelfFunction[dACE.SelfFunction(0, 0, [(1, 0, -2, 0)]), dACE.SelfFunction(1, 0, [(1, 0, -2, 0)]), dACE.SelfFunction(2, 0, [(1, 0, -2, 0)]), dACE.SelfFunction(3, 0, [(1, 0, -2, 0)]), dACE.SelfFunction(4, 0, [(1, 0, -2, 0)]), dACE.SelfFunction(5, 0, [(1, 0, -2, 0)]), dACE.SelfFunction(6, 0, [(1, 0, -2, 0)]), dACE.SelfFunction(0, 0, [(1, 0, -1, 0)]), dACE.SelfFunction(1, 0, [(1, 0, -1, 0)]), dACE.SelfFunction(2, 0, [(1, 0, -1, 0)])  …  dACE.SelfFunction(0, 0, [(3, 2, 0, 0), (4, 1, 0, 0)]), dACE.SelfFunction(1, 0, [(3, 2, 0, 0), (4, 1, 0, 0)]), dACE.SelfFunction(0, 0, [(3, 2, 0, 0), (4, 1, 1, 0)]), dACE.SelfFunction(0, 0, [(3, 2, 0, 0), (4, 2, 0, 0)]), dACE.SelfFunction(0, 0, [(3, 2, 1, 0), (4, 0, -1, 0)]), dACE.SelfFunction(0, 0, [(3, 2, 1, 0), (4, 0, 1, 0)]), dACE.SelfFunction(0, 0, [(3, 2, 1, 0), (4, 1, 0, 0)]), dACE.SelfFunction(0, 0, [(3, 3, 0, 0), (4, 0, -1, 0)]), dACE.SelfFunction(0, 0, [(3, 3, 0, 0), (4, 0, 1, 0)]), dACE.SelfFunction(0, 0, [(3, 3, 0, 0), (4, 1, 0, 0)])], 8), 0, 1, nothing, OrthPolyBasis1D3T(maxn = 5), OrthPolyBasis1D3T(maxn = 3), dACE.OneParticle[dACE.OneParticle(0, 0, 0, 0, 0, 0), dACE.OneParticle(0, 0, -1, 0, 0, 0), dACE.OneParticle(0, 0, 0, -1, 0, 0), dACE.OneParticle(0, 0, 0, 0, 0, -1), dACE.OneParticle(0, 0, 0, 0, 0, 1), dACE.OneParticle(0, 0, 0, 1, 0, 0), dACE.OneParticle(0, 0, 1, 0, 0, 0), dACE.OneParticle(1, 0, 0, 0, 0, 0), dACE.OneParticle(0, 0, -2, 0, 0, 0), dACE.OneParticle(0, 0, -1, -1, 0, 0)  …  dACE.OneParticle(2, 0, 1, 1, 0, 0), dACE.OneParticle(2, 0, 2, 0, 0, 0), dACE.OneParticle(2, 1, 0, 0, 0, 0), dACE.OneParticle(3, 0, -1, 0, 0, 0), dACE.OneParticle(3, 0, 0, -1, 0, 0), dACE.OneParticle(3, 0, 0, 0, 0, -1), dACE.OneParticle(3, 0, 0, 0, 0, 1), dACE.OneParticle(3, 0, 0, 1, 0, 0), dACE.OneParticle(3, 0, 1, 0, 0, 0), dACE.OneParticle(4, 0, 0, 0, 0, 0)], dACE.Centre[dACE.Centre(0, -4, 0, 0), dACE.Centre(0, -3, 0, 0), dACE.Centre(0, -2, 0, 0), dACE.Centre(0, -1, 0, 0), dACE.Centre(0, 0, 0, 0), dACE.Centre(0, 1, 0, 0), dACE.Centre(0, 2, 0, 0), dACE.Centre(0, 3, 0, 0), dACE.Centre(0, 4, 0, 0), dACE.Centre(1, -2, 0, 0), dACE.Centre(1, -1, 0, 0), dACE.Centre(1, 0, 0, 0), dACE.Centre(1, 1, 0, 0), dACE.Centre(1, 2, 0, 0), dACE.Centre(2, 0, 0, 0)], [[1], [2], [3], [4], [5], [6], [7], [8], [12], [13]  …  [140], [141], [144], [145], [151], [152], [153], [174], [175], [176]], dACE.BasisFunction[dACE.BasisFunction(6, Int64[], dACE.ProductTerm[dACE.ProductTerm(1.0, Int64[])]), dACE.BasisFunction(13, Int64[], dACE.ProductTerm[dACE.ProductTerm(1.0, Int64[])]), dACE.BasisFunction(6, [1], dACE.ProductTerm[dACE.ProductTerm(1.0, [1])]), dACE.BasisFunction(13, [1], dACE.ProductTerm[dACE.ProductTerm(1.0, [1])]), dACE.BasisFunction(7, [2], dACE.ProductTerm[dACE.ProductTerm(1.0, [2])]), dACE.BasisFunction(7, [3], dACE.ProductTerm[dACE.ProductTerm(1.0, [3])]), dACE.BasisFunction(7, [4], dACE.ProductTerm[dACE.ProductTerm(1.0, [4])]), dACE.BasisFunction(5, [5], dACE.ProductTerm[dACE.ProductTerm(1.0, [5])]), dACE.BasisFunction(12, [5], dACE.ProductTerm[dACE.ProductTerm(1.0, [5])]), dACE.BasisFunction(5, [6], dACE.ProductTerm[dACE.ProductTerm(1.0, [6])])  …  dACE.BasisFunction(5, [140], dACE.ProductTerm[dACE.ProductTerm(1.0, [57])]), dACE.BasisFunction(5, [141], dACE.ProductTerm[dACE.ProductTerm(1.0, [58])]), dACE.BasisFunction(5, [144], dACE.ProductTerm[dACE.ProductTerm(1.0, [59])]), dACE.BasisFunction(5, [145], dACE.ProductTerm[dACE.ProductTerm(1.0, [60])]), dACE.BasisFunction(5, [151], dACE.ProductTerm[dACE.ProductTerm(1.0, [61])]), dACE.BasisFunction(5, [152], dACE.ProductTerm[dACE.ProductTerm(1.0, [62])]), dACE.BasisFunction(5, [153], dACE.ProductTerm[dACE.ProductTerm(1.0, [63])]), dACE.BasisFunction(5, [174], dACE.ProductTerm[dACE.ProductTerm(1.0, [64])]), dACE.BasisFunction(5, [175], dACE.ProductTerm[dACE.ProductTerm(1.0, [65])]), dACE.BasisFunction(5, [176], dACE.ProductTerm[dACE.ProductTerm(1.0, [66])])], dACE.Purification{Polynomials4ML.OrthPolyBasis1D3T{4, Float64}, Polynomials4ML.OrthPolyBasis1D3T{2, Float64}}([(1, 0, 0, 0, 0, 0, 0), (1, 0, 0, -1, 0, 0, 0), (1, 0, 0, 0, -1, 0, 0), (1, 0, 0, 0, 0, 0, -1), (1, 0, 0, 0, 0, 0, 1), (1, 0, 0, 0, 1, 0, 0), (1, 0, 0, 1, 0, 0, 0), (1, 1, 0, 0, 0, 0, 0), (1, 0, 0, -1, 0, 0, 1), (1, 0, 0, -1, 1, 0, 0)  …  (1, 1, 0, 1, -1, 0, 1), (1, 1, 0, 1, 1, 0, -1), (1, 1, 0, 2, -1, 0, 0), (1, 1, 0, 2, 0, 0, -1), (1, 1, 1, 0, 0, 0, 1), (1, 1, 1, 0, 1, 0, 0), (1, 1, 1, 1, 0, 0, 0), (1, 3, 0, 0, 0, 0, 1), (1, 3, 0, 0, 1, 0, 0), (1, 3, 0, 1, 0, 0, 0)], sparse([1, 2, 3, 4, 5, 6, 7, 8, 9, 10  …  57, 58, 59, 60, 61, 62, 63, 64, 65, 66], [1, 2, 3, 4, 5, 6, 7, 8, 9, 10  …  57, 58, 59, 60, 61, 62, 63, 64, 65, 66], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0000000000000004, 1.0, 1.0  …  1.0000000000000004, 1.0000000000000004, 1.0000000000000004, 1.0000000000000004, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], 66, 66), 3, 1, 2, 2, 1, 0, 1, OrthPolyBasis1D3T(maxn = 4), OrthPolyBasis1D3T(maxn = 2))), ComplexF64[0.008064813641394792 - 0.000936955791435848im, -0.0001564538584899391 + 4.1436375624081834e-5im, 0.02851936860234543 + 0.0010692853946715633im, -0.0006311108903328435 + 1.0570155876304989e-5im, -0.0003207014202784482 + 0.028064688657057604im, -0.15372586426572693 - 0.005449032417099244im, -0.0053379064473057115 + 0.012333005708235024im, -0.004292997584593657 - 0.016171026995009725im, 0.0002643615619241793 + 0.0016946278238395918im, -0.019925953138052892 - 0.09890635501004631im  …  9.517813364522534e-5 + 0.00038516575202421236im, -0.00027841192444442607 - 0.0008775900128685485im, -0.0007771912161576538 - 0.000636140036575897im, -1.7497862627750127e-5 - 0.00012198869074122512im, 1.4420506724822297e-5 - 0.00014169306043316794im, -0.00044545229238925635 - 0.0003037630132546427im, 5.193030943565792e-5 + 0.0002754972927963128im, 6.442438432563888e-5 + 4.559943662739539e-5im, 0.00026914203847915963 + 0.00022945125179065523im, 2.4566295775264325e-5 - 0.0005294468461882176im], 0.001, 0.001)

Rollouts

rollout(model, pos0, vel0, dt, nsteps; past_velocities) integrates a fitted 2D or 3D model forward with explicit Euler steps. Models with a History carry their past velocities along (dt must be the sampling step of the training data).

AR(1) diagnostic

fit_ar1(basis, envs, accs, lagged_envs, lagged_accs) fits the force law jointly with persistent individual noise η(t) = ρ η(t - Δ) + ξ and returns an ARModel (2D only, with the lagged environments and accelerations one step earlier from trajectory_data(...; lagged = true)). The fitted ρ shows whether residuals persist. This is currentl mainly intended for diagnostic use. On tracked positions the previous finite-difference acceleration often shares frames with the target, so the one-step prediction predict(model, env, lagged_env, lagged_acc) partly predicts tracking error, and as ρ → 1 the force law is no longer identified. Evaluate the force law, predict(model.model, env).

See the API reference for all exported functions and options.