SFI.bases.constants module

SFI.bases.constants.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.constants.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.constants.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.constants.identity_matrix_basis(dim, pdepth=0)[source]
Parameters:
  • dim (int)

  • pdepth (int)

Return type:

Basis

SFI.bases.constants.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.constants.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.constants.ones_basis(dim, pdepth=0)[source]
Parameters:
  • dim (int)

  • pdepth (int)

Return type:

Basis

SFI.bases.constants.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.constants.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.constants.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.constants.unit_vector_basis(dim, axes=None)[source]
Parameters:
  • dim (int)

  • axes (Sequence[int] | None)

Return type:

Basis