SFI.bases.pairs module

SFI.bases.pairs

Generic building blocks for multi-particle (pair-interaction) systems.

This module provides:

• PBC utilities: minimum-image displacement in arbitrary dimension.

• Kernel families: pre-built 1-D radial kernel functions.

• Radial / scalar pair bases: build Interactor objects from kernel families, ready for dispatch over neighbor lists.

• Angular coupling bases: weighted orientation-coupling interactors.

• Heading vector: single-particle heading vector from an angle coordinate.

• Tensor pair features: rank-2 dyadic basis for diffusion tensors, nematic order, etc.

SFI.bases.pairs.angle_coupling(coupling_fn, *, dim, angle_index, output_index=None, label='g')

Scalar orientation coupling embedded as a rank-1 Interactor.

Computes coupling_fn(θ_j − θ_i) and places the result along output_index in a dim-d output vector.

Parameters:

• coupling_fn (callable) – Scalar function of the wrapped angle difference, e.g. jnp.sin.

• dim (int) – Full state-vector dimension.

• angle_index (int) – Index of the angle coordinate in the state vector.

• output_index (int, optional) – Output axis for the coupling value. Defaults to angle_index.

• label (str) – Short label used in feature names.

Returns:

Rank-1, 1-feature 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)

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

SFI.bases.pairs.compact_kernels(degrees, cutoff)

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)

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

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.exp_poly_kernels(degrees, lengths)

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

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.heading_vector(dim, angle_index, *, spatial_axes=None)

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.pair_direction(*, dim, box=None, spatial_dims=None, embed_dim=None, embed_axes=None)

Unit displacement vector \(\hat{r}_{ij}\) as a rank-1 Interactor.

Returns a single-feature, rank-1 pair Interactor whose output is the unit vector pointing from particle i to particle j.

Parameters:

• dim (int) – Full state-vector dimension per particle.

• box (array, "extras", or None) – PBC box lengths (same semantics as radial_pair_basis()).

• spatial_dims (slice or index array, optional) – Which axes are spatial coordinates.

• embed_dim (int, optional) – Embed into a larger output vector (e.g. 2-D displacement in 3-D state).

• embed_axes (sequence of int, optional) – Indices for embedding (required when embed_dim is set).

Returns:

Rank-1, 1-feature interactor with K=2.

Return type:

Interactor

SFI.bases.pairs.pair_velocity_difference(*, dim)

Velocity difference \(\mathbf{v}_j - \mathbf{v}_i\) as a rank-1 Interactor.

Returns a single-feature, rank-1 pair Interactor whose output is the velocity difference between neighbor and focal particle. Designed for composition with scalar pair Interactors via the * operator, e.g.:

scalar_pair_basis(kernels, dim=d) * pair_velocity_difference(dim=d)

Parameters:

dim (int) – State-vector dimension per particle.

Returns:

Rank-1, 1-feature interactor with K=2, needs_v=True.

Return type:

Interactor

SFI.bases.pairs.parametric_radial_kernel(kernel_fn, *, params, dim, box=None, spatial_dims=None)

Parametric scalar radial kernel as a rank-0 Interactor.

Wraps a user-supplied function kernel_fn(r, params) into a rank-0 pair Interactor with learnable parameters.

Parameters:

• kernel_fn (callable) – kernel_fn(r, params) -> scalar where r is the inter-particle distance and params is a dict of JAX arrays.

• params (dict) – Parameter specification passed to make_interactor(), e.g. {"eps": (), "R0": ()} for two scalar parameters.

• dim (int) – Full state-vector dimension.

• box (array, "extras", or None) – PBC box (same semantics as radial_pair_basis()).

• spatial_dims (slice or index array, optional) – Which axes are spatial coordinates.

Returns:

Rank-0, 1-feature parametric interactor with K=2.

Return type:

Interactor

SFI.bases.pairs.particle_heading(which, *, dim, angle_index, spatial_axes=None)

Heading vector of one particle in a pair, as a rank-1 Interactor.

Returns \((\cos\theta, \sin\theta)\) of the selected particle (which=0 for the focal particle, which=1 for the neighbor), embedded in a dim-d output vector.

Parameters:

• which (int) – 0 for the focal particle, 1 for the neighbor.

• dim (int) – Full state-vector dimension.

• angle_index (int) – Index of the angle coordinate.

• spatial_axes ((int, int), optional) – Indices for (cos θ, sin θ) in the output. Default: (0, 1).

Returns:

Rank-1, 1-feature interactor with K=2.

Return type:

Interactor

SFI.bases.pairs.particle_velocity(which, *, dim)

Velocity of one particle in a pair, as a rank-1 Interactor.

Returns the velocity of either the focal particle (which=0) or the neighbor (which=1) as a rank-1 Interactor. Designed for composition with scalar pair Interactors via the * operator, e.g.:

scalar_pair_basis(kernels, dim=d) * particle_velocity(which=1, dim=d)

Parameters:

• which (int) – 0 for the focal particle’s velocity, 1 for the neighbor’s.

• dim (int) – State-vector dimension per particle.

Returns:

Rank-1, 1-feature interactor with K=2, needs_v=True.

Return type:

Interactor

SFI.bases.pairs.pbc_displacement(xj, xi, box)

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

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

Parameters:

degrees (sequence of int)

Return type:

list of (callable, str)

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

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.scalar_pair_basis(kernels, *, dim, box=None, spatial_dims=None, labels=None)

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.vision_gate(gate_fn, *, dim, angle_index, box=None, spatial_dims=None)

Scalar vision-cone gate as a rank-0 Interactor.

Computes the bearing angle \(\delta\) from the focal particle’s heading to the displacement toward the neighbor, then returns gate_fn(δ). This makes interactions nonreciprocal: the gate value depends on whether the neighbor is “in front” or “behind” the focal particle.

Parameters:

• gate_fn (callable) – Scalar function of the bearing angle, e.g. lambda d: (1 + jnp.cos(d)) / 2 for a cosine vision cone.

• dim (int) – Full state-vector dimension.

• angle_index (int) – Index of the angle coordinate.

• box (array, "extras", or None) – PBC box (same semantics as radial_pair_basis()).

• spatial_dims (slice or index array, optional) – Which axes are spatial coordinates (for the displacement). Must select exactly 2 dimensions — bearing angle is 2-D only.

Returns:

Rank-0, 1-feature interactor with K=2.

Return type:

Interactor

SFI.bases.pairs.wrap_angle(a)

Wrap angle(s) to (-π, π].