SFI.bases package¶

SFI.bases¶

Library of ready-made basis builders on top of SFI.statefunc.

Submodules¶

monomials : scalar (or lifted) monomial families in x and/or v.
constants : structural bases: ones, unit vectors, identity/symmetric matrices.
linear : coordinate-extraction helpers (X, V, x_coordinate, …).
pairs : pair-interaction toolkit: radial kernels, PBC, heading vectors.
spde : composable spatial operators (Laplacian, Gradient, Divergence, Curl, …) on regular grids.

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(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.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)

SFI.bases.dataset_indicator(n_datasets, *, dim=None)[source]¶

Rank-0 Basis of n_datasets one-hot features 1{dataset == d}.

The linear-estimator route to per-dataset coefficients: multiply a feature by the indicator and concatenate, and each dataset gets an independent linear coefficient for that feature (the Gram is block-diagonal across datasets), PASTIS-prunable like any feature.

B = B_shared & (dataset_indicator(n) * X(dim))

Reads the reserved extras["dataset_index"] injected by TrajectoryCollection.

Parameters:

n_datasets (int)
dim (int | None)

SFI.bases.extra_scalar(name, *, dim=None, label=None)[source]¶

Rank-0, 1-feature Basis whose value is read from extras[name].

The compositional symbol for data-carried quantities: an extra (per-experiment constant, a time-dependent drive delivered per frame, a per-particle property) becomes an expression that composes with x_components / unit_axes / named_scalar through the usual algebra — and, being a parameter-free Basis, slots directly into linear-estimator dictionaries.

Parameters:

name (str) – Extras key to read; also the default feature label.
dim (int or None) – Spatial dimensionality; None (default) lets it be inferred at first call (the value is independent of x).
label (str or None) – Optional human-readable feature label (defaults to name).

Examples

>>> from SFI.bases import X, extra_scalar
>>> k_t = extra_scalar("k_drive")          # delivered by the dataset
>>> B = X(dim=2) & (k_t * X(dim=2))        # static + driven trap terms
>>> # simulate: OverdampedProcess(F=B, theta_F=jnp.array([-1.0, -1.0]))
>>> # infer:    inf.infer_force_linear(B)

Notes

At inference time the trajectory layer materializes extras[name] per frame (slicing TimeSeriesExtra values); in simulation, set_extras accepts the same time-dependent forms.

SFI.bases.field_component(index, *, n_fields, label=None)[source]¶

Extract a single field component from an SPDE state vector.

Alias for x_coordinate() with SPDE-oriented naming.

Parameters:

index (int) – Zero-based index of the field component to extract.
n_fields (int) – Total number of field components per grid site (= dim).
label (str, optional) – Human-readable label; defaults to "field[{index}]".

Return type:

Basis

SFI.bases.frame(dim, *, velocity=False, x_labels=None, v_labels=None, e_labels=None)[source]¶

Default compositional frame: constant 1 + coordinate scalars + unit axes.

Overdamped (velocity=False):

one, *x_components(dim), *unit_axes(dim)

Underdamped (velocity=True):

one, *x_components(dim), *v_components(dim), *unit_axes(dim)

Examples

>>> one, x, y, z, ex, ey, ez = frame(3)
>>> one, x, y, z, vx, vy, vz, ex, ey, ez = frame(3, velocity=True)

Custom labels (useful for dim > 4):

>>> bundle = frame(5, x_labels=["q0","q1","q2","q3","q4"])

Parameters:

dim (int)
velocity (bool)
x_labels (Sequence[str] | None)
v_labels (Sequence[str] | None)
e_labels (Sequence[str] | None)

Return type:

tuple[Basis, …]

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

Parameters:

dim (int)
pdepth (int)

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

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).

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.named_scalar(name, default=None, *, dim=None, label=None)[source]¶

Rank-0, 1-feature PSF whose value is a single named scalar parameter.

The returned PSF carries a single ParamSpec with shape (), optional default, and label label or name.

Parameters:

name (str) – Parameter name; also the default feature label.
default (scalar or None) – Optional default value. When set, the PSF can be evaluated, bound, or passed to a simulator without explicit params.
dim (int or None) – Spatial dimensionality; None (default) lets it be inferred at first call (the value is independent of x).
label (str or None) – Optional human-readable feature label (defaults to name).

Examples

>>> sigma = named_scalar("sigma", default=20.0)
>>> sigma()                       # uses default
Array(20., dtype=float32)
>>> sigma(params={"sigma": 30.})  # explicit override
Array(30., dtype=float32)

SFI.bases.named_scalars(*args, **kwargs)[source]¶

Unpack named scalar parameters, one PSF per name.

Two equivalent call styles, both with deterministic ordering:

Positional names (no defaults):

sigma, rho, beta = named_scalars("sigma", "rho", "beta")

Keyword names with defaults (Python preserves call-site order):

sigma, rho, beta = named_scalars(sigma=20.0, rho=8.0, beta=2.0)

Returns:: One PSF per name, in the order given.
Return type:: tuple[PSF, …]

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

Parameters:

dim (int)
pdepth (int)

Return type:

Basis

SFI.bases.per_dataset_scalar(name, n_datasets, default=None, *, dim=None)[source]¶

Rank-0, 1-feature PSF whose value is dataset-specific.

Carries a parameter array of shape (n_datasets,) and reads the entry of the current dataset through the reserved extras["dataset_index"] (injected automatically by TrajectoryCollection). Compose it with shared named_scalar() terms to fit pooled multi-experiment models where part of the parameters is experiment-specific and the rest is shared.

Parameters:

name (str) – Parameter name (one entry per dataset).
n_datasets (int) – Number of datasets in the collection the model will be fit on.
default (scalar or array of shape (n_datasets,), optional) – Optional default value(s); a scalar is broadcast.
dim (int or None) – Spatial dimensionality (None → inferred).

Notes

Indexed parameter access is nonlinear in the bookkeeping sense, so the parametric estimators fit models containing this primitive on the L-BFGS path. For the linear estimators, use the one-hot route instead: dataset_indicator().

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.time_fourier(n_modes, period=None, *, dim=None, label=None)[source]¶

Rank-0 time-Fourier dictionary read from the reserved time extra.

Emits 1 + 2 * n_modes parameter-free features

\[\bigl\{\,1,\; \cos(k\omega t),\; \sin(k\omega t)\,\bigr\}_{k=1}^{n_\text{modes}}, \qquad \omega = 2\pi / P,\]

evaluated at each frame’s absolute time t — the auto-injected time extra (see TrajectoryDataset.build_extras()), so no bookkeeping is required. Tensor it with a spatial basis to learn an unknown time-dependent force field by expansion: time_fourier(4) * X(dim=1) recovers a time-varying stiffness \(k(t)\), and time_fourier(4) * unit_vector_basis(1) a moving trap centre. Sparse selection (PASTIS) keeps only the harmonics the data support.

Parameters:

n_modes (int) – Number of harmonics; produces 1 + 2 * n_modes features.
period (float or None) – Fundamental period \(P\). If None (default), it defaults to the full trajectory duration (read from the auto-injected duration extra), i.e. the fundamental frequency is the inverse of the total observation time.
dim (int or None) – Spatial dimensionality; None lets it be inferred (the value is independent of x).
label (str or None) – Optional label prefix for the features.

Examples

>>> from SFI.bases import X, time_fourier
>>> B = time_fourier(4) * X(dim=1)        # learn k(t) over the trajectory
>>> inf.infer_force_linear(B)

SFI.bases.unit_axes(dim, *, labels=None)[source]¶

Unpack unit-vector bases (one per spatial axis).

>>> ex, ey, ez = unit_axes(3)

Each returned basis is rank-1 with a single feature carrying the unit vector along that axis. Labels default to ("ex", "ey", "ez", "ew") for dim <= 4 and ("e0", "e1", ...) otherwise.

Parameters:

dim (int)
labels (Sequence[str] | None)

Return type:

tuple[Basis, …]

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

Parameters:

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

Return type:

Basis

SFI.bases.v_components(dim, *, labels=None)[source]¶

Unpack scalar v-coordinate bases, one per axis.

>>> vx, vy, vz = v_components(3)

Parameters:

dim (int)
labels (Sequence[str] | None)

Return type:

tuple[Basis, …]

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

Single v-coordinate as a scalar feature.

Input : v ∈ R^dim (provided via keyword v=…) Return: scalar (); BasisLeaf will auto-insert feature axis → (1,)

Parameters:

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

Return type:

Basis

SFI.bases.v_coordinates(indices, *, dim, labels=None)[source]¶

Multiple v-coordinates as scalar features.

Input : v ∈ R^dim (provided via keyword v=…) Output: y ∈ R^{k} with k=len(indices)

Parameters:

indices (Sequence[int])
dim (int)
labels (Sequence[str] | None)

Return type:

Basis

SFI.bases.x_components(dim, *, labels=None)[source]¶

Unpack scalar x-coordinate bases, one per axis.

>>> x, y, z = x_components(3)

Each returned basis is rank-0 with one feature. Labels default to ("x", "y", "z", "w") for dim <= 4 and ("x0", "x1", ...) otherwise.

Parameters:

dim (int)
labels (Sequence[str] | None)

Return type:

tuple[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.x_coordinates(indices, *, dim, labels=None)[source]¶

Multiple x-coordinates as scalar features.

Input : x ∈ R^dim Output: y ∈ R^{k} with k=len(indices)

Parameters:

indices (Sequence[int])
dim (int)
labels (Sequence[str] | None)

Return type:

Basis