Docstrings

abstract type AbstractP4MLBasis end

Annotates types that map a low-dimensional input, scalar or SVector, to a vector of scalars (feature vector, embedding, basis...). Can be used as a Lux layer.

AtomicOrbitals : a quantum-chemistry atomic-orbital basis whose functions are products ϕ_{n1,n2,l,m}(𝐫) = Pn_{n1}(r) * Dn_{n2,l}(r) * Ylm_{l,m}(𝐫) of a radial polynomial part Pn, a radial decay part Dn (a RadialDecay), and an angular part Ylm. The angular Ylm is a spherical/solid harmonics basis used purely through the ACEbase evaluate interface, so it carries no Polynomials4ML parameters or state.

Polynomials4ML ships no harmonics of its own — the Ylm is supplied by SpheriCart.jl. You therefore need to load SpheriCart (which activates the Polynomials4ML SpheriCart extension) before constructing or using an AtomicOrbitals basis:

import Polynomials4ML as P4ML
import SpheriCart                      # activates the P4ML SpheriCart extension
basis = P4ML._rand_gaussian_basis()    # a ready-made example basis

To assemble one directly, pass any SpheriCart harmonics basis as Ylm:

using SpheriCart: SolidHarmonics
basis = AtomicOrbitals(Pn, Dn, SolidHarmonics(L), spec, specidx)

Complex trigonometric polynomials up to degree N (inclusive). The basis is constructed in the order

[1, exp(im*θ), exp(-im*θ), exp(2im*θ), exp(-2im*θ), ..., 
                                exp(N*im*θ), exp(-N*im*θ) ]

where θ is input variable.

ChebBasis(N):

Chebyshev polynomials up to degree N-1 (inclusive). i.e basis with length N. The basis is ordered as

\[[1, x, 2x^2-1, 4x^3-3x, ..., 2xT_{N-1}(x)-T_{N-2}(x)]\]

where x is input variable.

The differences between ChebBasis and chebyshev_basis is that ChebBasis computes the basis on the fly when it is compiled and it does not store the recursion coefficients as in chebyshev_basis. There might be a small performance benefit from this.

Secondly, ChebBasis and chebyshev_basis use different normalization.

struct CubicSplines:

Statically typed cubic splines, compatible with P4ML type batched evaluation. For any P4ML basis with univariate input.

Standard Monomials basis. This should very rarely be used. Possibly useful in combination with a transformation of the inputs, e.g. exponential.

OrthPolyBasis1D3T: defines a basis of polynomials in terms of a 3-term recursion,

\[\begin{aligned} P_1(x) &= A_1 \\ P_2 &= A_2 x + B_2 \\ P_{n} &= (A_n x + B_n) P_{n-1}(x) + C_n P_{n-2}(x) \end{aligned}\]

Typically (but not necessarily) such bases are obtained by orthogonalizing the monomials with respect to a user-specified distribution, which can be either continuous or discrete but must have a density function. See also

• legendre_basis
• chebyshev_basis
• jacobi_basis

RTrigBasis(N):

Real trigonometric polynomials up to degree N (inclusive). The basis is ordered as

[1, cos(θ), sin(θ), cos(2θ), sin(2θ), ..., cos(Nθ), sin(Nθ) ]

where θ is input variable.

StaticBatch{N,T} : an auxiliary StaticArray type that is distinct from SVector{N,T}. It can be used to create a batch of inputs of static size N. It is in used to convert function calls with single inputs to function calls with a batch of inputs.

struct TransformedBasis

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

TransformedBasis(trans, basis)

The point of this structure is to provide a transformed basis that behaves exactly as all other P4ML bases.

Comments

• the "natural indices" will be the same as for basis
• _generate_input is not implemented for general input transforms; to

implement it for an input transform of type TIN one should monkey-patch

Polynomials4ML._generate_input(::TransformedBasis{TIN}, T::Type) where {TIN} = ...

• _valtype is implemented but unclear how well it behaves, might be necessary

to monkey-patch it as well

struct WrappedBasis

A wrapper type for a Lux layer l that behaves like a P4ML basis. The wrapper implements the necessary methods to make WrappedBasis a valid AbstractP4MLBasis. It assumes that the following calls are valid:

• l(x::T, ps, st) with T <: Number produces an AbstractVector{T}; i.e.
• l(X::AbstractVector{T}, ps, st) produces an B::AbstractMatrix{T} of

numbers, with B[i, :] == l(X[i], ps, st).

In particular it is assume that input and output types match. If this fails then the behaviour is undefined. (With the non-allocating interface this is likely unproblematic. Witht he allocating interface one could monkey-patch _valtype to get around this restriction.)

The forwardpass is computed via l(x, ps, st). Due to the above assumption, the optimal implementation of derivatives is forward-mode, hence evaluate_ed is implemented via ForwardDiff, and the rrule is provided by the standard P4ML interface.

cubsplwidthgrad(x, F, G, x0, x1, NX) -> (f, g)

Evaluate the cubic spline basis and its derivative at x. Outside [x0, x1] the value is Flat-extended and the derivative is zero.

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

Evaluate a Hermite cubic at t ∈ [0,1], given function values fl, fr and (interval-scaled) gradients gl, gr at the left and right endpoints.

evalcubic_d(t, fl, fr, gl, gr, h) -> (f, g)

Evaluate the Hermite cubic and its derivative w.r.t. the physical coordinate on an interval of width h. t is the local coordinate and fl, fr, gl, gr are the endpoint values / interval-scaled gradients. Value and t-derivative are computed analytically from the cubic f(t) = a3 t³ + a2 t² + a1 t + a0 (same coefficients as _eval_cubic); the t-derivative is then rescaled to the physical coordinate by 1/h.

evalcubspl(x, F, G, x0, x1, NX)

Evaluate the cubic spline basis at x, given the node value / gradient arrays F, G and the grid (x0, x1, NX).

generateinput(basis)

Returns a single randomly generated valid input for basis.

_gradtype(basis, x)

If the intention is that P, dP = evaluate_ed(basis, x) then then _gradtype(basis, x) should return etype(dP).

a fall-back method for initalparameters that all AbstractP4MLBasis should overload

splgrid(x, x0, x1, NX) -> (x, t, il, h)

Locate x on the uniform spline grid of NX nodes spanning [x0, x1]. Returns the clamped coordinate x (Flat boundary condition), the local coordinate t ∈ [0,1) within the interval, the 0-based left-node index il, and the grid spacing h. The cubic on that interval is then evaluated with _eval_cubic(t, F[il+1], F[il+2], h*G[il+1], h*G[il+2]).

GPU-safe: uses unsafe_trunc rather than floor(Int, ·).

_valtype(basis, x)

If the intention is that P = basis(x) where P is a Vector{T} then _valtype(basis, x) should return T.

Here, x can be a single input, a batch or a type. A new basis type TB only needs to implement _valtype(::TB, x::Type).

chebyshev_basis(N::Integer):

Constructs an OrthPolyBasis1D3T object representing a possibly rescaled version of the basis of Chebyshev polynomials of the first kind. N is the length of the basis, not the degree.

Careful: the normalisation may be non-standard.

construct_basis(RadialDecay, ζ_raw, D_raw, decay)

Construct a RadialDecay object from raw matrix data ζ_raw, D_raw and a DecayFunction(f, df) representing the decay form and its derivative. All input is converted to statically-sized SMatrix for efficiency.

index(basis, k) -> Integer

Given a "natural description" of a basis element return the index of that basis element in the computed vector of basis function values. For example, for Chebyshev polynomials, index(basis, n) returns n+1.

jacobi_basis(N::Integer, α::Real, β::Real):

Constructs an OrthPolyBasis1D3T object representing a possibly rescaled version of the basis of Jacobi polynomials Jαβ. N is the length of the basis, not the degree.

Careful: the normalisation may be non-standard.

legendre_basis(N::Integer):

Constructs an OrthPolyBasis1D3T object representing a possibly rescaled version of the basis of Legendre polynomials (L2 orthonormal on [-1, 1]). N is the length of the basis, not the degree.

Careful: the normalisation may be non-standard.

function orthpolybasis(...) : construct a univariate orthogonal polynomial basis with respect to some specified inner product. For the standard 3-term recursion polynomials, use legendre_basis, jacobi_basis or chebyshev_basis.

The orthpolybasis currently implements orthogonal polynomials for discrete weights with the following constructors:

orthpolybasis(N::Integer, W::DiscreteWeights{TW}; TX = Float64)
orthpolybasis(N::Integer, X::AbstractVector{<: Real}, 
              W::AbstractVector{<: Real}, normalizeW=false; kwargs...)

splinify(basis, x0, x1, NX; bspline=true)

Takes a P4ML basis with univariate input and constructs a cubic spline basis that interpolates the basis functions on a uniform grid with NX nodes. If bspline=true (default) the function values are first interpolated onto a B-spline representation to obtain C2,2 regularity of the splines.

x0, x1 are the left and right endpoints of the spline interval.

This is currently not exported and not part of the public interface. The interface can change in future releases.

struct LinL

A very basic linear layer that is compatible with the memory layout of P4ML.

_gensparse : function barrier for gensparse

gensparse(...) : 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]`
• tup2b = 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

• INT = Int : integer type to be used
• ordered = false : whether only ordered tuples are produced; ordered tuples

correspond to permutation-invariant basis functions