Docstrings

To be used as part of TransformedBasis as follows:

y = transform(x) 
P = basis(y) * envelope(x, y)

Warning: P may be modified in-place!

If a particle x is represented as a NamedTuple, e.g., x = (r = SA[...], Z = 13) then an NTtransform can be used to embed this named tuple into ℝ in a differentiable way, e.g.,

x = (𝐫 = randn(StaticVector{3, Float64}), Z = rand(10:50))
r0 = Float64[ ... ]  # list of r0 values for rescaling r 
trans = NTtransform(x -> 1 / (1 + norm(x.𝐫)/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 again 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.

struct ParallelEmbed : basically a variation on Lux.Parallel, but makes some assumptions on what the individual layers to be evaluated in parallel do to enable some simple optimizations.

The assumption is that each layer is a "basis", i.e. a mapping x -> B(x), where x is low-dimensional and B(x) is moderate-dimensional. This means that the optimal differentiation is in forward-mode. This is exploited by implementing the functions evaluate and evaluate_ed. For implementing models that use ParallelEmbed a forward-pass that also differentiates (e.g. energy and forces) is implemented with evaluate_ed.

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 TransformedBasis

Basically a three-stage chain, consisting of an input transformation, and basis evaluation. Constructor:

TransformedBasis(; transin, basis, transout)

defaults for transin and transout are IDtrans(), i.e. identity transformations.

The layer performs the following chain:

y = evaluate(transin, x, ...)
P = evaluate(basis, y, ...) 
B = evaluate(transout, P, x, ...)

• It is assumed that transout only utilizes the categorical variables stored in

x but not the continuous variables. This means that when differentiating, one only needs to differentiate B(P(y), x) with respect to y but not with respect to x.

• it is also assumed that basis has no parameters

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

convert readable AA_spec into the internal representation of the AA basis

convert readable A_spec into the internal representation of the A basis

  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

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.