StochasticForceInference 2.0.0__py3-none-any.whl

SFI/bases/linear.py

@@ -0,0 +1,325 @@
+# SFI.bases.linear
+# =================
+# Lightweight linear utilities for building first-order bases in x and v.
+#
+# Contract assumed here:
+# - User functions are called on a SINGLE SAMPLE (no leading batch axes).
+# - For rank=VECTOR leaves, x and v arrive as shape (dim,).
+# - For rank=SCALAR leaves, the function may return () or (k,) where k=n_features.
+# - Feature axis must be present in the final leaf output; for n_features=1 on
+#   a scalar leaf, BasisLeaf will auto-insert it from a scalar return.
+#
+# pdepth=0 (non-interacting) everywhere in this helper module.
+
+from __future__ import annotations
+
+from typing import Optional, Sequence
+
+import jax.numpy as jnp
+
+from ..statefunc import Basis, Rank, make_basis
+from .monomials import monomials_degree  # degree-specific builder
+
+__all__ = [
+    "linear_basis",
+    "X",
+    "V",
+    "x_coordinate",
+    "x_coordinates",
+    "v_coordinate",
+    "v_coordinates",
+    "x_components",
+    "v_components",
+    "unit_axes",
+    "frame",
+]
+
+
+def linear_basis(dim: int, *, include_x: bool = True, include_v: bool = False):
+    """
+    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
+    -------
+    Basis
+        Rank-1 (vector) basis concatenating requested degree-1 monomials.
+    """
+    return monomials_degree(1, dim=dim, include_x=include_x, include_v=include_v)
+
+
+def X(dim: int, *, label: Optional[str] = None) -> Basis:
+    """
+    Identity in x with an explicit feature axis.
+
+    Input : x ∈ R^dim
+    Output: Y ∈ R^{dim×1}
+    """
+
+    def _eval(x):
+        return x[:, None]  # (dim, 1)
+
+    label = "x" if label is None else label
+    return make_basis(
+        func=_eval,
+        dim=dim,
+        rank=Rank.VECTOR,
+        n_features=1,
+        needs_v=False,
+        labels=[label],
+    )
+
+
+def V(dim: int, *, label: Optional[str] = None) -> Basis:
+    """
+    Identity in v with an explicit feature axis.
+
+    Input : v ∈ R^dim  (provided via keyword v=...)
+    Output: Y ∈ R^{dim×1}
+    """
+
+    def _eval(x, v):
+        return v[:, None]  # (dim, 1)
+
+    label = "v" if label is None else label
+    return make_basis(
+        func=_eval,
+        dim=dim,
+        rank=Rank.VECTOR,
+        n_features=1,
+        needs_v=True,
+        labels=[label],
+    )
+
+
+def x_coordinate(index: int, *, dim: int, label: Optional[str] = None) -> Basis:
+    """
+    Single x-coordinate as a scalar feature.
+
+    Input : x ∈ R^dim
+    Return: scalar (); BasisLeaf will auto-insert feature axis → (1,)
+    """
+
+    def _eval(x):
+        return x[index]  # ()
+
+    label = f"x{index}" if label is None else label
+    return make_basis(
+        func=_eval,
+        dim=dim,
+        rank=Rank.SCALAR,
+        n_features=1,
+        needs_v=False,
+        labels=[label],
+    )
+
+
+def field_component(index: int, *, n_fields: int, label: Optional[str] = None) -> Basis:
+    """Extract a single field component from an SPDE state vector.
+
+    Alias for :func:`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}]"``.
+    """
+    if label is None:
+        label = f"field[{index}]"
+    return x_coordinate(index, dim=n_fields, label=label)
+
+
+def x_coordinates(indices: Sequence[int], *, dim: int, labels: Optional[Sequence[str]] = None) -> Basis:
+    """
+    Multiple x-coordinates as scalar features.
+
+    Input : x ∈ R^dim
+    Output: y ∈ R^{k} with k=len(indices)
+    """
+    indices = jnp.array(indices)
+
+    def _eval(x):
+        return x[indices]  # (k,)
+
+    if labels is None:
+        labels = [f"x{i}" for i in indices]
+    elif len(labels) != len(indices):
+        raise ValueError("x_coordinates: labels length must match number of indices")
+
+    return make_basis(
+        func=_eval,
+        dim=dim,
+        rank=Rank.SCALAR,
+        n_features=len(indices),
+        needs_v=False,
+        labels=list(labels),
+    )
+
+
+def v_coordinate(index: int, *, dim: int, label: Optional[str] = None) -> Basis:
+    """
+    Single v-coordinate as a scalar feature.
+
+    Input : v ∈ R^dim (provided via keyword v=...)
+    Return: scalar (); BasisLeaf will auto-insert feature axis → (1,)
+    """
+
+    def _eval(x, v):
+        return v[index]  # ()
+
+    label = f"v{index}" if label is None else label
+    return make_basis(
+        func=_eval,
+        dim=dim,
+        rank=Rank.SCALAR,
+        n_features=1,
+        needs_v=True,
+        labels=[label],
+    )
+
+
+def v_coordinates(indices: Sequence[int], *, dim: int, labels: Optional[Sequence[str]] = None) -> Basis:
+    """
+    Multiple v-coordinates as scalar features.
+
+    Input : v ∈ R^dim (provided via keyword v=...)
+    Output: y ∈ R^{k} with k=len(indices)
+    """
+    indices = jnp.array(indices)
+
+    def _eval(x, v):
+        return v[indices]  # (k,)
+
+    if labels is None:
+        labels = [f"v{i}" for i in indices]
+    elif len(labels) != len(indices):
+        raise ValueError("v_coordinates: labels length must match number of indices")
+
+    return make_basis(
+        func=_eval,
+        dim=dim,
+        rank=Rank.SCALAR,
+        n_features=len(indices),
+        needs_v=True,
+        labels=list(labels),
+    )
+
+
+# ---------------------------------------------------------------------------
+# Component / axis unpackers
+# ---------------------------------------------------------------------------
+_DEFAULT_X_LABELS = ("x", "y", "z", "w")
+_DEFAULT_V_LABELS = ("vx", "vy", "vz", "vw")
+_DEFAULT_E_LABELS = ("ex", "ey", "ez", "ew")
+
+
+def _auto_labels(dim: int, defaults: Sequence[str], prefix: str) -> list[str]:
+    if dim <= len(defaults):
+        return [defaults[i] for i in range(dim)]
+    return [f"{prefix}{i}" for i in range(dim)]
+
+
+def x_components(dim: int, *, labels: Optional[Sequence[str]] = None) -> tuple[Basis, ...]:
+    """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.
+    """
+    if labels is None:
+        labels = _auto_labels(dim, _DEFAULT_X_LABELS, "x")
+    elif len(labels) != dim:
+        raise ValueError(f"x_components: labels length ({len(labels)}) must equal dim ({dim})")
+    return tuple(x_coordinate(i, dim=dim, label=labels[i]) for i in range(dim))
+
+
+def v_components(dim: int, *, labels: Optional[Sequence[str]] = None) -> tuple[Basis, ...]:
+    """Unpack scalar v-coordinate bases, one per axis.
+
+    >>> vx, vy, vz = v_components(3)
+    """
+    if labels is None:
+        labels = _auto_labels(dim, _DEFAULT_V_LABELS, "v")
+    elif len(labels) != dim:
+        raise ValueError(f"v_components: labels length ({len(labels)}) must equal dim ({dim})")
+    return tuple(v_coordinate(i, dim=dim, label=labels[i]) for i in range(dim))
+
+
+def unit_axes(dim: int, *, labels: Optional[Sequence[str]] = None) -> tuple[Basis, ...]:
+    """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.
+    """
+    # local import to avoid a circular import at module load
+    from .constants import unit_vector_basis
+
+    if labels is None:
+        labels = _auto_labels(dim, _DEFAULT_E_LABELS, "e")
+    elif len(labels) != dim:
+        raise ValueError(f"unit_axes: labels length ({len(labels)}) must equal dim ({dim})")
+
+    out = []
+    for i in range(dim):
+        e = unit_vector_basis(dim, axes=[i])
+        # Override the leaf's label with the friendly default. Bypass the
+        # Equinox freeze to patch labels after construction; fragile if
+        # Equinox changes its freeze mechanism — prefer constructing with
+        # labels set.
+        leaf = e.root
+        object.__setattr__(leaf, "labels", (labels[i],))
+        out.append(e)
+    return tuple(out)
+
+
+def frame(
+    dim: int,
+    *,
+    velocity: bool = False,
+    x_labels: Optional[Sequence[str]] = None,
+    v_labels: Optional[Sequence[str]] = None,
+    e_labels: Optional[Sequence[str]] = None,
+) -> tuple[Basis, ...]:
+    """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"])
+    """
+    from .constants import ones_basis
+
+    pieces: list[Basis] = [ones_basis(dim)]
+    pieces.extend(x_components(dim, labels=x_labels))
+    if velocity:
+        pieces.extend(v_components(dim, labels=v_labels))
+    pieces.extend(unit_axes(dim, labels=e_labels))
+    return tuple(pieces)

SFI/bases/monomials.py

@@ -0,0 +1,218 @@
+import functools
+import itertools
+
+import jax.numpy as jnp
+
+from ..statefunc import Basis, Rank, make_basis
+
+_RANK_ALIASES = {
+    "scalar": None,
+    "vector": "vector",
+    "matrix": "symmetric",
+    "symmetric_matrix": "symmetric",
+    "identity_matrix": "identity",
+}
+
+
+def _lift(B: Basis, rank: str, dim: int):
+    """Lift a scalar Basis to higher rank via Cartesian product with a structural basis."""
+    import warnings
+
+    from .constants import identity_matrix_basis, symmetric_matrix_basis, unit_vector_basis
+
+    if rank not in _RANK_ALIASES:
+        raise ValueError(f"Unknown rank={rank!r}. Choose from: {list(_RANK_ALIASES.keys())}")
+    mode = _RANK_ALIASES[rank]
+    if mode is None:
+        return B
+    # Suppress the Cartesian-product warning — this is intentional lifting.
+    with warnings.catch_warnings():
+        warnings.filterwarnings("ignore", ".*Cartesian product.*")
+        if mode == "vector":
+            return B * unit_vector_basis(dim)
+        if mode == "identity":
+            return B * identity_matrix_basis(dim)
+        if mode == "symmetric":
+            return B * symmetric_matrix_basis(dim)
+    raise AssertionError("unreachable")  # pragma: no cover
+
+
+@functools.lru_cache(maxsize=64)
+def _exponent_mats(dim: int, degree: int, include_x: bool, include_v: bool):
+    """Return (ex_x, ex_v, labels) for all monomials of exact total degree."""
+    if not include_x and not include_v:
+        raise ValueError("At least one of include_x / include_v must be True.")
+    L = (dim if include_x else 0) + (dim if include_v else 0)
+
+    # stars-and-bars enumeration for total degree == degree
+    if degree == 0:
+        gammas = [tuple(0 for _ in range(L))]
+    else:
+        gammas = []
+        for stars in itertools.combinations_with_replacement(range(L), degree):
+            arr = [0] * L
+            for s in stars:
+                arr[s] += 1
+            gammas.append(tuple(arr))
+
+    ex_x, ex_v, labels = [], [], []
+    for g in gammas:
+        alpha = g[:dim] if include_x else (0,) * dim
+        beta = g[-dim:] if include_v else (0,) * dim
+
+        parts = []
+        for k, p in enumerate(alpha):
+            if p:
+                parts.append(f"x{k}^{p}" if p > 1 else f"x{k}")
+        for k, p in enumerate(beta):
+            if p:
+                parts.append(f"v{k}^{p}" if p > 1 else f"v{k}")
+        lab = "1" if not parts else "·".join(parts)
+
+        ex_x.append(alpha)
+        ex_v.append(beta)
+        labels.append(lab)
+
+    ex_x = jnp.array(ex_x, dtype=jnp.int32)  # (F, dim)
+    ex_v = jnp.array(ex_v, dtype=jnp.int32)  # (F, dim)
+    return ex_x, ex_v, tuple(labels)
+
+
+def monomials_degree(
+    degree: int,
+    *,
+    dim: int,
+    include_x: bool = True,
+    include_v: bool = False,
+    rank: str = "scalar",
+):
+    """
+    All monomials of **exact** total degree in x and/or v.
+
+    Parameters
+    ----------
+    degree : int
+        Exact total polynomial degree.
+    dim : int
+        Spatial dimension.
+    include_x, 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).
+    """
+    ex_x, ex_v, labels = _exponent_mats(dim, degree, include_x, include_v)
+    F = int(ex_x.shape[0])
+    use_x, use_v = bool(include_x), bool(include_v)
+
+    # Precompute static stuff outside eval (captured in the closure)
+    Dx = int(ex_x.max()) if use_x else 0  # max degree for x
+    Dv = int(ex_v.max()) if use_v else 0  # max degree for v
+    Ix = ex_x.T if use_x else None  # (dim, F) gather indices
+    Iv = ex_v.T if use_v else None  # (dim, F)
+
+    def _power_table_fixed(z, D):
+        """
+        z: (dim,) -> (dim, D+1) = [1, z, z^2, ..., z^D]
+
+        Important: avoid jnp.power with float exponents because its JVP/VJP uses log(z)
+        and yields NaNs at z=0 even for integer powers (e.g. 0^0 pathways).
+        This iterative construction is polynomial and differentiable everywhere.
+        """
+        dim = z.shape[0]
+        one = jnp.ones((dim, 1), dtype=z.dtype)
+        if D == 0:
+            return one
+        # powers[:, k] = z^k
+        powers = [one, z[:, None]]
+        for _ in range(2, D + 1):
+            powers.append(powers[-1] * z[:, None])
+        return jnp.concatenate(powers, axis=1)
+
+    if use_x and use_v:
+
+        def _eval(x, *, v=None, mask=None):
+            # x,v: (dim,) ; returns (F,)
+            if v is None:
+                v = jnp.zeros_like(x)
+            Tx = _power_table_fixed(x, Dx)  # (dim, Dx+1)
+            Tv = _power_table_fixed(v, Dv)  # (dim, Dv+1)
+            vx = jnp.take_along_axis(Tx, Ix, axis=1)  # (dim, F)
+            vv = jnp.take_along_axis(Tv, Iv, axis=1)  # (dim, F)
+            return jnp.prod(vx * vv, axis=0)  # (F,)
+    elif use_x:
+
+        def _eval(x, *, v=None, mask=None):
+            Tx = _power_table_fixed(x, Dx)  # (dim, Dx+1)
+            vx = jnp.take_along_axis(Tx, Ix, axis=1)  # (dim, F)
+            return jnp.prod(vx, axis=0)  # (F,)
+    else:  # use_v only
+
+        def _eval(x, *, v=None, mask=None):
+            Tv = _power_table_fixed(v if v is not None else jnp.zeros_like(x), Dv)
+            vv = jnp.take_along_axis(Tv, Iv, axis=1)  # (dim, F)
+            return jnp.prod(vv, axis=0)  # (F,)
+
+    B = make_basis(
+        func=_eval,
+        dim=dim,
+        rank=Rank.SCALAR,
+        n_features=F,
+        needs_v=use_v,
+        labels=list(labels),
+    )
+    return _lift(B, rank, dim)
+
+
+def monomials_up_to(
+    order: int,
+    *,
+    dim: int,
+    include_constant: bool = True,
+    include_x: bool = True,
+    include_v: bool = False,
+    rank: str = "scalar",
+):
+    r"""
+    Concatenate degree-wise monomial bases for degrees 0..order (ascending).
+
+    .. physics:: Multivariate polynomial basis
+       :label: polynomial-basis
+       :category: Basis functions
+
+       .. math::
+
+          f_\\alpha(x) = \\prod_{k=1}^{d} x_k^{\\alpha_k},
+          \\qquad |\\alpha| \\le \\texttt{order}
+
+       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, include_v : bool
+        Which variables to include.
+    rank : str
+        Output tensor rank.  See :func:`monomials_degree` for allowed
+        values: ``'scalar'``, ``'vector'``, ``'matrix'`` /
+        ``'symmetric_matrix'``, ``'identity_matrix'``.
+
+        For force inference, use ``rank='vector'`` (the most common choice).
+    """
+    start = 0 if include_constant else 1
+    if start > order:
+        raise ValueError("No monomials to include: set include_constant=True or order>=1.")
+    Bs = [monomials_degree(d, dim=dim, include_x=include_x, include_v=include_v) for d in range(start, order + 1)]
+    B = Basis.stack(Bs)
+    return _lift(B, rank, dim)