Docstrings

struct EdgeEmbed

Wraps a layer that embeds an edge state into Vector to manage the reformatting of the embedding from a list of embedded states into a 3-dimensionsonal tensor that is aware of the graph structure.

Also implements an evaluate_ed wrapper, which is useful for computing jacobians.

struct EmbedDP

Embed a particle state into a vector. This is done by first applying a transform to the particle state into a number of SVector and then evaluating the basis (or other layer) on the transformed state.

This is basically a 3-stage Chain, but with additional logic, specifically the implementation of evaluate_ed allowing differentiation through and XState or NamedTuple input. This is e.g. needed for jacobians.

Implementation of nt_transform and dp_transform. For details see the docstring for those.

struct PooledSparseProduct : This implements a fused (tensor) product and pooling operation. Suppose we are given $N$ embeddings $\phi^{(i)}_{k_i}$ then the pooled sparse product generates feature vectors of the form

\[A_{k_1, \dots, k_N} = \sum_{j} \prod_{t = 1}^N \phi^{(t)}_{k_t}(x_j)\]

where $x_j$ are an list of inputs (multi-set).

The canonical example is

\[A_{nlm} = \sum_{j} R_{nl}(r_j, \mu_j) Y_l^m( \hat{\bm r}_j ),\]

where $R_{nl}$ is a radial embedding, possibly depending on a categorical variable $\mu_j$ and $Y_l^m$ are spherical harmonics.

Constructor

PooledSparseProduct(spec)

where spec is a list of $(k_1, \dots, k_N)$ tuples or vectors, or AbstractMatrix where each column specifies such a tuple.

struct SelectLinL <: AbstractLuxLayer

A Lux layer which acts as a simple linear layer, but using a categorical variable to select a weight matrix:

   P -> W[x] * P 

This layer is experimental and likely very inefficient.

struct SparseMatCSX

Sparse matrix format that is stored in both CSR and CSC, and can be transferred to GPU devices. Matrix multiplication via KernelAbstractions.jl for GPU support and uses whichever format (CSR, CSC) is most suitable for a given operation.

SparseSymmProd : sparse symmetric product with entries stored as tuples. Input is a vector A; each entry of the output vector AA is of the form

\[ {\bm A}_{i_1, \dots, i_N} = \prod_{t = 1}^N A_{i_t}.\]

Constructor

SparseSymmProd(spec)

where spec is a list of tuples or vectors, each of which specifies an AA basis function as described above. For example,

spec = [ (1,), (2,), (1,1), (1,2), (2,2), 
         (1,1,1), (1,1,2), (1,2,2), (2,2,2) ]
basis = SparseSymmProd(spec)         

defines a basis of 9 functions,

\[[ A_1, A_2, A_1^2, A_1 A_2, A_2^2, A_1^3, A_1^2 A_2, A_1 A_2^2, A_2^3 ]\]

struct TransSelSplines

A spline implementation that is intended exclusively to be used as an invariant embedding, e.g. radial basis.

x -> y -> s(y, i(x)) ->   R(x) = s(y) .* e(x) 
  ------> e(x) ------/

here y = trans(x), the selector chooses the index i(x) of the spline, e(x) is defined by the envelope.

Generate a transform with frozen parameters stored in a state; cf st_transform.

Takes a list of 𝔸 or 𝔹 specifications (many-body) in the form of

   [  [(n=., l=., m=.), (n=., l=., m=.)], ... ]

and converts it into a list of sorted unique (n=., l=.) named pairs, i.e the specification of the one-body basis.

autoYlmspec(mbspec)

takes a list of 𝔸 or 𝔹 specifications (many-body) and return the specification of the Ylm basis functions as a [ (l = ., m = .), ... ] list.

takes an nnll spec and generates a complete list of all possible nnllmm

evalcubic(t, fl, fr, gl, gr, h)

Evaluate cubic spline at position t in [0,1], given function values fl, fr and gradients gl, gr at the left and right endpoints.

convert readable AA_spec into the internal representation of the AA basis

convert readable A_spec into the internal representation of the A basis

splgrid(y, x0, x1, NX)

Compute the local spline grid information needed for evalcubic

agnesi_params(pcut, pin, rin, req, rcut)

Precompute the parameters for the generalized Agnesi transform, to be used with eval_agnesi. See docs for ??? for details.

agnesi_transform(...)

Construct a generalized Agnesi transform, a layer that implements the operation

\[ y(x) = b_0 + \frac{b_1}{1 + a s^q / (1 + s^(q-p))}\]

where

\[ s = \frac{r - r_{\rm in}}{r_{\rm eq} - r_{\rm in}}, \]

with default b_0, b_1 such that $r \in [r_{\rm in}, r_{\rm cut}]$ is mapped to [-1, 1] and a such that $|dy/dr|$ is maximised at $r = r_{\rm eq}$.

The transform is constructed to satisfy (assum $r_{\rm in} = 0$)

\[ y \sim \frac{1}{1 + a (r/r_{\rm eq})^q} \quad \text{as} \quad r \to 0 \quad \text{and} \quad y \sim (r/r_{\rm eq})^{-p} \quad \text{as} r \to \infty.\]

The values for $p, q, r_{\rm in}, r_{\rm eq}, r_{\rm cut}$ need to be specified. Typical defaults for $p,q$ are p = 2, q = 4.

cat2idx(categories, a)

maps a -> i such that categories[i] == a. Basically the same as findfirst, similar performance but seems to avoid an allocation.

catcat2idx(categories1, categories2, a1, a2) catcat2idx(categories, a1, a2)

same as cat2idx but for pairs of categorical variables, mapping to a linear index.

catcat2idx_sym(categories, a1, a2)

Same as catcat2idx but for symmetric pairs, i.e., (a1, a2) and (a2, a1) map to the same index.

function nttransform(f::Function) function nttransform(f::Function, refstate::NamedTuple)

If a particle x is represented as a PState (DecoratedParticles.jl) or a NamedTuple, e.g.,x = PState(r = SA[...], Z = 13), then annt_transform` generates a type that incorporates broadcasting and differentiation. For example

x = (𝐫 = randn(StaticVector{3, Float64}), Z = rand(10:50))
refstate = (; r0 = SA[ ... ])    # list of r0 values for rescaling r 
trans = dp_transform( (x, st) -> 1 / (1 + norm(x.𝐫)/ st.r0[x.Z]))

We can then evaluate and differenitate

y, _ = evaluate(trans, x, ps, st)
(y, dy), _ = evaluate_ed(trans, x, ps, st)

Here, dy is a VState or a named-tuple with the derivative w.r.t. x.𝐫 stored as dy.𝐫. The derivative w.r.t. Z is not taken because Z is a categorical variable.

The transform is assumed to be parameter-free, or with frozen parameters stored in the refstate. If a transform has no parameters at all then one can generate it via

trans = dp_transform( x -> 1 / (1 + norm(x.𝐫)))

eval_agnesi(r::Real, params)

Evaluate the generalized Agnesi transform at distance r, with parameters provided produced by agnesi_params.

  invmap(a::AbstractVector)

Returns a structure that makes looking up the index of an element in a vector convenient and fast. Assumes that elements of a are unique.

inva = invmap(a) 
inva[a[i]] == i  # true for all i

reshape_embedding(P, ii, jj, nnodes, maxneigs)

Takes a Nedges x Nfeat matrix and writes it into a 3-dimensional array of size (maxneigs, nnodes, Nfeat) where each column corresponds to a node. The "missing" neighbours are filled with zeros.

revreshapeembedding(P3, ii, jj, nnodes, maxneigs) -> P

Reverse operation for reshape_embedding. P3 is of shape (maxneigs, nnodes, nfeatures), and this gets written into P which is of shape (nedges, nfeatures) and then returned.

setproduct(A)

Assumes the A is a length-N collection of collections. It returns a P = Matrix{T} where each P[i, :] is a vector of length N with P[i, j] ∈ A[j]. The number of columns of P is the number of such products, i.e. prod(length.(A)).

In constrast with Iterators.product this implementation is type-stable for a priori unknown N.

sparseequivarianttensor(L, mbspec, Rnlspec, Ylm_spec, basis)

sparse_product(...) : utility function to generate high-dimensional sparse grids which are downsets. All arguments are keyword arguments (with defaults):

• NU : maximum correlation order
• minvv = 0 : minvv[i] gives the minimum value forvv[i]`
• maxvv = Inf : maxvv[i] gives the minimum value forvv[i]`
• tup2bb = vv -> vv :
• admissible = _ -> false : determines whether a tuple belongs to the downset
• filter = _ -> true : a callable object that returns true of tuple is to be kept and

false otherwise (whether or not it is part of the downset!) This is used, e.g. to enfore conditions such as ∑ lₐ = even or |∑ mₐ| ≦ M

• ordered = false : whether only ordered tuples are produced; ordered tuples

correspond to permutation-invariant basis functions

Generate a transform with frozen parameters stored in a state. The advantage of this layer over WrappedFunction is that the state can be moved to a device or converted to different floating point formats.

symidx(j1, j2, n)

Computes the index in a flattened symmetric matrix of size n x n, given the row/column indices j1, j2 (1-based). E.g. if n == 3 then the mapping is as follows:

   [ 1 2 3
     2 4 5
     3 5 6 ]

symmetrisationmatrix(L, mbspec; prune, kwargs...) -> 𝔸2𝔹, 𝔸_spec

Generates the symmetrization operator for a sparse ACE basis. The basis is

specified via the input mb_spec, which is a list of basis function

specifications of the form

mb_spec = [ [(n=0, l=0), (n=1, l=0)], [(n=1, l=1), (n=2, l=1)], ... ] 

i.e. a Vector{Vector{NL}}. where NL = @NamedTuple{n::Int, l::Int}.

The parameter L determines the order of the ouput, e.g. L=0 for an invariant scalar, L = 1 for a vector, and so forth.

The output is given in terms of a sparse matrix 𝔸2𝔹 in CCS format and a specification of the 𝔸 basis as a Vector{Vector{NLM}} where NLM = @NamedTuple{n::Int, l::Int, m::Int}.

struct TYVec2CartMat

transformation from a real Y vector to a cartesian matrix; accepts as input either an SVector{3}(i.e. the Y_1^m) or aSYYVector(containing both the Y_0^0 and Y_1^m); the output is aSMatrix{3,3}`.

struct TYVec2CartVec

transformation from a real Y vector to a cartesian vector; accepts as input either an SVector{3}(i.e. the Y_1^m) or aSYYVector(containing both the Y_0^0 and Y_1^m); the output is aSVector{3}`.

struct TYVec2YMat

transformation from a real Y vector to a spherical matrix

Dfromangles(l, θ, basis)

Here, l::Integer and θ a 3-element vector or tuple, basis must be either real or complex. Output is a Wigner-D matrix such that y ∘ Q = D * y with y real/complex spherical harmonics.

produces a rotation Q and Wigner-D matrix D such that y ∘ Q = D * y with y real spherical harmonics.

O3.coupling_coeffs(L, ll, nn; PI, basis)
O3.coupling_coeffs(L, ll; PI, basis)

Compute coupling coefficients for the spherical harmonics basis, where

• L must be an Integer;
• ll, nn must be vectors or tuples of Integer of the same length.
• PI: whether or not the coupled basis is permutation-invariant (or the

corresponding tensor symmetric); default is true when nn is provided and false when nn is not provided.

• basis: which basis is being coupled, default is complex, alternative

choice is real, which is compatible with the SpheriCart.jl convention.