Bases API¶

This page documents the public API of the SFI.bases package. For usage examples, the construction decision tree, canonical patterns (Lorenz, harmonic UD, limit cycle, double-well), and guidance on choosing a basis, see Building bases.

Reference card¶

All builders are importable from SFI.bases directly (from SFI.bases import monomials_up_to):

Builder	One-liner
`monomials_up_to(order, dim, rank=...)`	Polynomials up to total order; `rank` lifts to vector/tensor
`monomials_degree(degree, dim, ...)`	Polynomials of exact total degree
`ones_basis(dim)`	Scalar constant 1
`unit_vector_basis(dim, axes)`	Cartesian unit vectors (rank-1)
`identity_matrix_basis(dim)`	Isotropic \(\delta_{ij}\) (rank-2)
`symmetric_matrix_basis(dim)`	All symmetric unit matrices (rank-2)
`constant_array(A)`	Wrap a constant array as one feature
`X(dim)` / `V(dim)`	The position / velocity vector itself (rank-1)
`x_coordinate(i, dim)`	Single coordinate \(x_i\) (rank-0)
`linear_basis(dim)`	Coordinate-extraction identity
`x_components()` / `v_components()` / `unit_axes()` / `frame`	Compositional per-axis primitives (see Building bases)
`named_scalar(s)`	Parametric scalar PSFs with optional defaults
`SFI.bases.pairs.*`	Pair-interaction builders and kernel families (below)
`SFI.bases.spde.*`	PDE differential operators (SPDE API reference, experimental)

Two ways to build a basis¶

Both yield a Basis and are freely combinable with +, *, and & (see Building bases).

Factory path (this page). Pre-built builders return a complete Basis in one call: monomials_up_to(), monomials_degree(), linear_basis(), the kernel families under SFI.bases.pairs, and the SPDE operators under SFI.bases.spde. Right tool when the structure of the basis is already standard (polynomials, pair potentials, finite-difference operators).
Compositional path — primitives x_components(), v_components(), unit_axes(), frame, named_scalar(s) plus the algebra operators + / * / &. Right tool when the basis follows the symmetry of the problem rather than a generic family (e.g. ex * x[0] + ey * x[1] for a 2D radial force).

Monomials¶

SFI.bases.monomials_up_to(order, *, dim, include_constant=True, include_x=True, include_v=False, rank='scalar')[source]

Concatenate degree-wise monomial bases for degrees 0..order (ascending).

[Basis functions] Multivariate polynomial basis

\[\begin{split}f_\\alpha(x) = \\prod_{k=1}^{d} x_k^{\\alpha_k}, \\qquad |\\alpha| \\le \\texttt{order}\end{split}\]

Full polynomial dictionary up to a given total degree, optionally including velocity monomials and lifted to vector or matrix rank.

Parameters:

order (int) – Maximum total polynomial degree.
dim (int) – Spatial dimension.
include_constant (bool) – If False, skip degree-0 (constant) term.
include_x (bool) – Which variables to include.
include_v (bool) – Which variables to include.
rank (str) –
Output tensor rank. See monomials_degree() for allowed values: 'scalar', 'vector', 'matrix' / 'symmetric_matrix', 'identity_matrix'.

For force inference, use rank='vector' (the most common choice).

SFI.bases.monomials_degree(degree, *, dim, include_x=True, include_v=False, rank='scalar')[source]

All monomials of exact total degree in x and/or v.

Parameters:

degree (int) – Exact total polynomial degree.
dim (int) – Spatial dimension.
include_x (bool) – Which variables to include.
include_v (bool) – Which variables to include.
rank (str) – Output rank. 'scalar' (default) returns a scalar Basis with F features. 'vector' lifts to rank-1 via Cartesian product with unit_vector_basis(dim) (F × dim features). 'matrix' / 'symmetric_matrix' lifts to rank-2 via symmetric_matrix_basis(dim) (F × dim(dim+1)/2 features). 'identity_matrix' lifts via identity_matrix_basis(dim) (F × 1 features, isotropic).

Structural constants¶

SFI.bases.ones_basis(dim, pdepth=0)[source]

Parameters:

dim (int)
pdepth (int)

Return type:

Basis

SFI.bases.unit_vector_basis(dim, axes=None)[source]

Parameters:

dim (int)
axes (Sequence[int] | None)

Return type:

Basis

SFI.bases.identity_matrix_basis(dim, pdepth=0)[source]

Parameters:

dim (int)
pdepth (int)

Return type:

Basis

SFI.bases.symmetric_matrix_basis(dim, pdepth=0)[source]

Constant symmetric-matrix templates spanning the space of real symmetric dim × dim matrices.

For dim=d there are d(d+1)/2 features: one per upper-triangle entry (i,j) with i <= j. Each feature is a (dim, dim) matrix: S_{(i,j)} = δ_{ia}δ_{jb} + δ_{ib}δ_{ja} (so the off-diagonal templates equal 1 in both symmetric slots, diagonal templates equal 1 on the diagonal).

Rank is MATRIX (2), and the output shape is (dim, dim, F).

Parameters:

dim (int)
pdepth (int)

Return type:

Basis

SFI.bases.constant_array(A, *, label='const', descriptor='constant-array', as_sf=True)[source]

Constant basis/sf with a single feature whose value is a fixed tensor A of shape (dim,)*rank (rank inferred from A.ndim, dim from A.shape[0]).

Errors if A is not a hypercube tensor (all axes same length).
Broadcasts over batch/particles: output shape is (*x.shape[:-1], (dim,)*rank, 1).

Parameters:

label (str)
descriptor (str)

Linear shorthand¶

SFI.bases.X(dim, *, label=None)[source]

Identity in x with an explicit feature axis.

Input : x ∈ R^dim Output: Y ∈ R^{dim×1}

Parameters:

dim (int)
label (str | None)

Return type:

Basis

SFI.bases.V(dim, *, label=None)[source]

Identity in v with an explicit feature axis.

Input : v ∈ R^dim (provided via keyword v=…) Output: Y ∈ R^{dim×1}

Parameters:

dim (int)
label (str | None)

Return type:

Basis

SFI.bases.x_coordinate(index, *, dim, label=None)[source]

Single x-coordinate as a scalar feature.

Input : x ∈ R^dim Return: scalar (); BasisLeaf will auto-insert feature axis → (1,)

Parameters:

index (int)
dim (int)
label (str | None)

Return type:

Basis

SFI.bases.linear_basis(dim, *, include_x=True, include_v=False)[source]

Degree-1 monomial basis in (x, v).

Parameters:

dim (int) – Spatial dimension.
include_x (bool) – Include linear x terms.
include_v (bool) – Include linear v terms.

Returns:

Rank-1 (vector) basis concatenating requested degree-1 monomials.

Return type:

Basis

Pair-interaction bases¶

SFI.bases.pairs.scalar_pair_basis(kernels, *, dim, box=None, spatial_dims=None, labels=None)[source]

Build a rank-0 pair Interactor with scalar radial-kernel features.

Each feature is \(\phi_\alpha(r_{ij})\) — the raw kernel value without the directional \(\hat{r}\) factor.

Use this for energy-like quantities, as radial weights for angular coupling, or as building blocks for tensor pair features composed via the * operator.

Parameters:

kernels (Sequence[tuple[Callable, str]]) – Same as radial_pair_basis().
dim (int) – Same as radial_pair_basis().
box (Any) – Same as radial_pair_basis().
spatial_dims (slice | Sequence[int] | None) – Same as radial_pair_basis().
labels (Sequence[str] | None) – Same as radial_pair_basis().

Returns:

Rank-0 (scalar) interactor with K=2.

Return type:

Interactor

SFI.bases.pairs.radial_pair_basis(kernels, *, dim, box=None, spatial_dims=None, embed_dim=None, embed_axes=None, labels=None)[source]

Build a rank-1 pair Interactor with radial-kernel features.

Each feature is \(\phi_\alpha(r_{ij})\,\hat{\mathbf{r}}_{ij}\) where \(r_{ij}\) is the pairwise distance (optionally with PBC).

[Basis functions] Radial pair interaction basis

\[f_\alpha(\mathbf{r}_{ij}) = \phi_\alpha(r_{ij})\;\hat{\mathbf{r}}_{ij}\]

Scalar radial kernel \(\phi_\alpha\) times the unit displacement vector. Available kernel families: exponential-polynomial, Gaussian, power-law, and compactly supported.

Parameters:

kernels (list of (callable, str)) – 1-D kernel functions and their labels, as returned by exp_poly_kernels(), gaussian_kernels(), etc.
dim (int) – Full state-vector dimension per particle.
box (array, "extras", or None) – PBC box lengths. None (default) = free-space, no periodic boundaries. Pass an array for a static box captured in the closure, or "extras" to read the box from extras["box"] at evaluation time. The box is applied over spatial_dims only.
spatial_dims (slice or index array, optional) – Which axes of the state vector are spatial coordinates (the ones over which distances are computed and the output vector is defined). Default: slice(None) (all axes are spatial).
embed_dim (int, optional) – If the output should be embedded into a larger vector (e.g., the displacement lives in 2D but the state vector is 3D with an angle), set embed_dim to the output vector size. Spatial components are placed at embed_axes indices; remaining indices are zero.
embed_axes (sequence of int, optional) – Indices into the embed_dim-length output where spatial components are placed. Required when embed_dim is not None.
labels (sequence of str, optional) – Override labels (one per kernel).

Returns:

Rank-1 (vector) interactor with K=2, n_features=len(kernels). Call .dispatch_pairs(...) to stream over neighbour lists.

Return type:

Interactor

SFI.bases.pairs.dyadic_pair_basis(kernels, *, dim, box=None, spatial_dims=None, labels=None)[source]

Build a rank-2 (tensor) pair Interactor: \(\phi(r)\,\hat{r}\otimes\hat{r}\).

Each feature is the outer product of the unit displacement vector with itself, weighted by a radial kernel. Useful for directional diffusion tensors, nematic order parameters, etc.

Parameters:

kernels (Sequence[tuple[Callable, str]]) – Same as radial_pair_basis().
dim (int) – Same as radial_pair_basis().
box (Any) – Same as radial_pair_basis().
spatial_dims (slice | Sequence[int] | None) – Same as radial_pair_basis().
labels (Sequence[str] | None) – Same as radial_pair_basis().

Returns:

Rank-2 (matrix) interactor with K=2.

Return type:

Interactor

SFI.bases.pairs.angular_pair_basis(kernels, coupling_fn, *, dim, angle_index, output_index=None, box=None, spatial_dims=None, coupling_label='g', labels=None)[source]

Build a rank-1 pair Interactor for orientation coupling.

Each feature computes \(\phi_\alpha(r_{ij})\,g(\theta_j - \theta_i)\) and embeds the result along output_index in a dim-d output vector.

Parameters:

kernels (list of (callable, str)) – Radial weight functions (same format as other kernel factories).
coupling_fn (callable) – Scalar function of the angle difference, e.g. jnp.sin for alignment, lambda a: jnp.cos(2*a) for nematic coupling.
dim (int) – Full state-vector dimension.
angle_index (int) – Index of the angle coordinate in the state vector.
output_index (int, optional) – Index along which the coupled output is placed. Defaults to angle_index.
box (Any) – PBC and spatial-dimension controls (same as radial_pair_basis()).
spatial_dims (slice | Sequence[int] | None) – PBC and spatial-dimension controls (same as radial_pair_basis()).
coupling_label (str) – Short label for the coupling (appears in feature labels).
labels (list of str, optional) – Override labels.

Returns:

Rank-1 (vector) interactor with K=2.

Return type:

Interactor

Kernel families¶

Each family parametrises a different radial profile for pair potentials. Choose the one that best matches the expected interaction range and short-distance behaviour:

gaussian_kernels() — smooth bumps \(\phi_\sigma(r) = \exp(-r^2/2\sigma^2)\); good general-purpose default for soft, finite-range interactions.
exp_poly_kernels() — \(\phi_{k,L}(r) = r^k \exp(-r/L)\); exponentially-decaying with a polynomial prefactor (e.g. screened forces).
power_kernels() — pure power laws \(\phi_k(r) = r^k\) (no decay); use with care, requires a finite neighbour cutoff.
compact_kernels() — compactly supported \(r^k (1 - r/r_c)^2\) for \(r < r_c\); identically zero past r_c, ideal with a neighbour list.

SFI.bases.pairs.gaussian_kernels(sigmas)[source]

Radial kernels \(\phi_\sigma(r) = \exp(-r^2 / 2\sigma^2)\).

Parameters:: sigmas (sequence of float) – Gaussian widths.
Return type:: list of (callable, str)

SFI.bases.pairs.exp_poly_kernels(degrees, lengths)[source]

Radial kernels \(\phi_{k,L}(r) = r^k \exp(-r/L)\).

Parameters:

degrees (sequence of int) – Polynomial degrees k.
lengths (sequence of float) – Exponential decay lengths L.

Returns:

Each entry is (phi, label) where phi(r) -> scalar.

Return type:

list of (callable, str)

SFI.bases.pairs.power_kernels(degrees)[source]

Radial kernels \(\phi_k(r) = r^k\).

Parameters:: degrees (sequence of int)
Return type:: list of (callable, str)

SFI.bases.pairs.compact_kernels(degrees, cutoff)[source]

Compactly-supported kernels \(r^k (1 - r/r_c)^2\) for \(r < r_c\).

Parameters:

degrees (sequence of int)
cutoff (float) – Support radius r_c.

Return type:

list of (callable, str)

Pair utilities¶

SFI.bases.pairs.heading_vector(dim, angle_index, *, spatial_axes=None)[source]

Single-particle heading vector from an angle coordinate.

Returns a rank-1 Basis whose single feature is the unit vector \((\cos\theta, \sin\theta)\) embedded in a dim-d vector, with the cosine and sine placed at spatial_axes[0] and spatial_axes[1] respectively.

Parameters:

dim (int) – State-vector dimension.
angle_index (int) – Index of the angle coordinate \(\theta\).
spatial_axes ((int, int), optional) – Indices for (cos θ, sin θ) in the output. Default: (0, 1) — i.e. the first two axes.

Returns:

Rank-1, 1-feature heading-vector basis.

Return type:

Basis

SFI.bases.pairs.pbc_displacement(xj, xi, box)[source]

Minimum-image displacement xj - xi under periodic boundaries.

Works in any dimension. All inputs are plain JAX arrays of shape (d,) (or broadcastable).

Parameters:

xj (array, shape (d,)) – Positions (or sub-positions) of two particles.
xi (array, shape (d,)) – Positions (or sub-positions) of two particles.
box (array, shape (d,)) – Box lengths along each axis.

Returns:

dx – xj - xi folded via minimum-image convention.

Return type:

array, shape (d,)

SFI.bases.pairs.wrap_angle(a)[source]: Wrap angle(s) to (-π, π].

SPDE operators¶

Composable spatial differential operators for grid-based field inference are documented in a dedicated section: