StochasticForceInference 2.0.0__py3-none-any.whl

SFI/statefunc/params.py

@@ -0,0 +1,309 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Callable, Iterable, Tuple
+
+import equinox as eqx
+import jax.numpy as jnp
+
+
+# ---------------------------------------------------------------------
+# ParamSpec  - static description of one parameter block
+# ---------------------------------------------------------------------
+@dataclass(frozen=True, slots=True)
+class ParamSpec:
+    name: str  # <-- KEY for sharing
+    shape: Tuple[int, ...]
+    dtype: Any = jnp.float32
+    init: Callable | str = "zeros"  # PRNG->array or keyword
+    default: Any = None  # optional concrete value (scalar or array-like)
+
+    @property
+    def size(self) -> int:
+        from functools import reduce
+        from operator import mul
+
+        return reduce(mul, self.shape, 1)
+
+    def compatible_with(self, other: "ParamSpec") -> bool:
+        """Shareable iff shape and dtype match exactly."""
+        return (self.shape == other.shape) and (self.dtype == other.dtype)
+
+    def merged_with(self, other: "ParamSpec") -> "ParamSpec":
+        """
+        Return a single spec representing the shared parameter.
+        Requires compatibility; keeps `self.init` by default.
+        """
+        if not self.compatible_with(other):
+            raise ValueError(
+                f"ParamSpec mismatch for '{self.name}': {self.shape}/{self.dtype} vs {other.shape}/{other.dtype}"
+            )
+        # deterministically keep `self`, but salvage a default from `other`
+        if self.default is None and other.default is not None:
+            from dataclasses import replace
+
+            return replace(self, default=other.default)
+        return self
+
+
+class ParamSuite(eqx.Module):
+    """Immutable container holding a set of ``ParamSpec`` objects."""
+
+    specs: tuple[ParamSpec, ...] = eqx.field(static=True)
+    _lookup: dict[str, int] = eqx.field(static=True)  # name → idx
+
+    # ---------- construction ---------- #
+    def __init__(self, specs: Iterable[ParamSpec]):
+        specs = tuple(specs)
+        names = [ps.name for ps in specs]
+        if len(set(names)) != len(names):
+            dup = {n for n in names if names.count(n) > 1}
+            raise ValueError(f"Duplicate parameter names: {dup}")
+        object.__setattr__(self, "specs", specs)
+        object.__setattr__(self, "_lookup", {ps.name: i for i, ps in enumerate(specs)})
+
+    @classmethod
+    def from_specs(cls, *specs: ParamSpec) -> "ParamSuite":
+        return cls(specs)
+
+    # ---------- basic info ---------- #
+    def __iter__(self):
+        return iter(self.specs)
+
+    def __len__(self):
+        return len(self.specs)
+
+    def __getitem__(self, name: str) -> ParamSpec:
+        return self.specs[self._lookup[name]]
+
+    @property
+    def size(self) -> int:
+        return sum(ps.size for ps in self.specs)
+
+    # ---------- convenience constructors ---------- #
+    def zeros(self) -> dict[str, jnp.ndarray]:
+        """Return a parameter dict with all values initialized to zero."""
+        return {ps.name: jnp.zeros(ps.shape, dtype=ps.dtype) for ps in self.specs}
+
+    @property
+    def has_defaults(self) -> bool:
+        """True iff every spec in this suite carries a concrete ``default``."""
+        return len(self.specs) > 0 and all(ps.default is not None for ps in self.specs)
+
+    def defaults(self) -> dict[str, jnp.ndarray] | None:
+        """Return a parameter dict from spec ``default`` values, or None if any
+        spec has no default. Values are broadcast to the declared shape."""
+        if not self.has_defaults:
+            return None
+        out: dict[str, jnp.ndarray] = {}
+        for ps in self.specs:
+            arr = jnp.asarray(ps.default, dtype=ps.dtype)
+            if arr.shape != ps.shape:
+                arr = jnp.broadcast_to(arr, ps.shape)
+            out[ps.name] = arr
+        return out
+
+    # ---------- param ↔ vector helpers ---------- #
+    def materialize(
+        self,
+        vector: jnp.ndarray,
+        *,
+        dtype_overrides: dict[str, jnp.dtype] | None = None,
+    ):
+        if vector.ndim != 1 or vector.size != self.size:
+            raise ValueError("Flat vector has wrong length")
+        out, i = {}, 0
+        for ps in self.specs:
+            n = ps.size
+            arr = (
+                vector[i : i + n]
+                .reshape(ps.shape)
+                .astype(dtype_overrides.get(ps.name, ps.dtype) if dtype_overrides else ps.dtype)
+            )
+            out[ps.name] = arr
+            i += n
+        return out
+
+    def vectorize(self, tree: dict[str, jnp.ndarray]) -> jnp.ndarray:
+        parts = [jnp.ravel(tree[ps.name]).astype(ps.dtype) for ps in self.specs]
+        return jnp.concatenate(parts, axis=0)
+
+    # ---------- merging (parameter sharing) ---------- #
+    def merge(self, other: "ParamSuite | None") -> "ParamSuite":
+        """
+        Union with sharing-by-name:
+          - If a name appears in both suites and specs are compatible (shape/dtype),
+            they are **tied** (kept once).
+          - If incompatible → error.
+        """
+        if other is None:
+            return self
+        if not isinstance(other, ParamSuite):
+            raise TypeError(f"merge expects ParamSuite or None, got {type(other).__name__}")
+        if not self.specs:
+            return other
+        if not other.specs:
+            return self
+
+        left = {ps.name: ps for ps in self.specs}
+        right = {ps.name: ps for ps in other.specs}
+
+        out: dict[str, ParamSpec] = dict(left)  # copy
+        for name, ps_r in right.items():
+            if name in out:
+                ps_l = out[name]
+                out[name] = ps_l.merged_with(ps_r)  # validates compatibility
+            else:
+                out[name] = ps_r
+        return ParamSuite(out.values())
+
+    @classmethod
+    def merge_many(cls, *suites: "ParamSuite | None") -> "ParamSuite | None":
+        """Merge any number of suites, sharing parameters by name (shape/dtype must match)."""
+        merged: ParamSuite | None = None
+        for s in suites:
+            if s is None:
+                continue
+            merged = s if merged is None else merged.merge(s)
+        return merged
+
+    # ---------- PyTree protocol ---------- #
+    def tree_flatten(self):
+        children = ()  # no array leaves
+        aux = self.specs  # static
+        return children, aux
+
+    # ---------- universal parser ---------- #
+    @classmethod
+    def parse(cls, params) -> "ParamSuite | None":
+        """
+        Normalize various user-facing descriptions into a ParamSuite.
+
+        Accepts:
+
+        - ``None`` -- returns ``None``
+        - ``ParamSuite`` -- returned as-is
+        - ``dict[name -> array | shape]`` -- infer shape/dtype
+        - ``iterable[ParamSpec]`` -- from_specs
+        - ``iterable[str]`` -- scalar specs for each name
+
+        Shapes may be ``()``, ``(k,)``, ``(m, n, ...)`` or an integer k
+        (interpreted as ``(k,)``).
+        """
+        from collections.abc import Iterable as _Iterable
+
+        if params is None:
+            return None
+        if isinstance(params, ParamSuite):
+            return params
+        # dict path: values can be arrays (sample), shapes, ints, ParamSpec, or None→scalar
+        if isinstance(params, dict):
+            specs: list[ParamSpec] = []
+            for name, val in params.items():
+                if isinstance(val, ParamSpec):
+                    specs.append(val)
+                    continue
+                if hasattr(val, "shape"):
+                    shape = tuple(val.shape)
+                    dtype = getattr(val, "dtype", jnp.float32)
+                    specs.append(ParamSpec(name, shape, dtype=dtype))
+                    continue
+                if val is None:
+                    specs.append(ParamSpec(name, ()))
+                    continue
+                if isinstance(val, int):
+                    specs.append(ParamSpec(name, (int(val),)))
+                    continue
+                if isinstance(val, tuple) and all(isinstance(n, int) for n in val):
+                    specs.append(ParamSpec(name, tuple(int(n) for n in val)))
+                    continue
+                raise TypeError(f"ParamSuite.parse: unsupported dict value for '{name}': {type(val).__name__}")
+            return cls.from_specs(*specs)
+        # iterable path: ParamSpec or str
+        if isinstance(params, _Iterable):
+            items = list(params)
+            specs: list[ParamSpec] = []
+            for it in items:
+                if isinstance(it, ParamSpec):
+                    specs.append(it)
+                elif isinstance(it, str):
+                    specs.append(ParamSpec(it, ()))
+                else:
+                    raise TypeError(
+                        f"ParamSuite.parse: iterable must contain ParamSpec or str (got {type(it).__name__})"
+                    )
+            return cls.from_specs(*specs)
+        raise TypeError(f"ParamSuite.parse: unsupported params of type {type(params).__name__}")
+
+    @classmethod
+    def tree_unflatten(cls, aux, children):
+        return cls(aux)
+
+    def coerce(
+        self,
+        params: dict,
+        *,
+        allow_scalar_for_scalar: bool = True,
+        allow_scalar_to_len1: bool = True,
+        cast_dtype: bool = True,
+    ) -> dict[str, jnp.ndarray]:
+        """
+        Normalize a user param dict into JAX arrays matching this suite.
+
+        Rules:
+          - If spec.shape == (), accept Python scalars / 0-d arrays (if allowed).
+          - If spec.shape == (1,), optionally accept a scalar and expand to (1,).
+          - Otherwise, require exact shape; dtype is cast if `cast_dtype` is True.
+        Returns a NEW dict with normalized arrays.
+        """
+        out = {}
+        look = self._lookup
+        missing = [k for k in look if k not in params]
+        if missing:
+            raise KeyError(f"Missing params: {missing}")
+
+        for spec in self:
+            name, shape, dtype = spec.name, spec.shape, spec.dtype
+            val = params[name]
+
+            arr = jnp.asarray(val, dtype=dtype if cast_dtype else None)
+
+            # Handle scalar vs length-1 convenience
+            if shape == ():  # true scalar
+                # 0-d OK; (1,) OK if allow_scalar_for_scalar and allow squeezing
+                if arr.shape == ():
+                    pass
+                elif arr.shape == (1,) and allow_scalar_for_scalar:
+                    arr = jnp.reshape(arr, ())
+                else:
+                    raise TypeError(f"Param '{name}': expected scalar (), got {arr.shape}")
+            elif shape == (1,):
+                # accept scalar and expand to (1,) if asked
+                if arr.shape == ():
+                    if allow_scalar_to_len1:
+                        arr = jnp.reshape(arr, (1,))
+                    else:
+                        raise TypeError(f"Param '{name}': expected (1,), got scalar ()")
+                elif arr.shape == (1,):
+                    pass
+                else:
+                    raise TypeError(f"Param '{name}': expected {shape}, got {arr.shape}")
+            else:
+                if arr.shape != shape:
+                    raise TypeError(f"Param '{name}': expected {shape}, got {arr.shape}")
+
+            # Final dtype enforcement
+            if cast_dtype:
+                arr = arr.astype(dtype)
+            elif arr.dtype != dtype:
+                raise TypeError(f"Param '{name}': expected dtype {dtype}, got {arr.dtype}")
+
+            out[name] = arr
+        return out
+
+    def stamp(self, params_dict: dict) -> tuple:
+        """
+        Build a stable stamp for (shape, dtype) of each param in template order.
+        Caller is expected to pass an already-coerced dict.
+        """
+        return tuple((spec.name, params_dict[spec.name].shape, params_dict[spec.name].dtype) for spec in self)

SFI/statefunc/psf.py

@@ -0,0 +1,112 @@
+"""PSF façade: parametric family of state functions."""
+
+import equinox as eqx
+import jax
+import jax.numpy as jnp
+
+from .nodes import BaseNode, DerivativeNode
+from .params import ParamSuite
+from .stateexpr import StateExpr
+
+
+# ============================================================================
+#  PSF  – parametric state-function family
+# ============================================================================
+class PSF(StateExpr):
+    """Parametric State-Function family `F(x; θ)`.
+
+    By default `drop_features=True`: when `n_features==1`, outputs **do not**
+    carry a trailing feature axis. `.d_theta()` forces `drop_features=False`.
+
+    Holds a `ParamSuite` template describing names, shapes, and dtypes of θ.
+    `__call__` evaluates `F` given a parameter dict matching the template.
+    Supports `.d_theta()` in addition to `.d_x()`/`.d_v()`.
+    """
+
+    template: ParamSuite = eqx.field(static=True)
+    drop_features: bool = eqx.field(static=True, default=True)
+    _validated_stamp: tuple | None = eqx.field(static=True, default=None, repr=False)
+
+    def __init__(self, root: BaseNode, *, drop_features: bool = True):
+        if root.param_suite is None:
+            raise ValueError("PSF root must carry parameters")
+        super().__init__(root)
+        object.__setattr__(self, "template", root.param_suite)
+        object.__setattr__(self, "drop_features", bool(drop_features))
+
+    def _coerce_and_check(self, params, extras):
+        # Normalize/validate via the suite
+        pnorm = self.template.coerce(params)
+        stamp = tuple(sorted((k, pnorm[k].shape, pnorm[k].dtype) for k in self.template._lookup))
+        if stamp != self._validated_stamp:
+            object.__setattr__(self, "_validated_stamp", stamp)
+        # Only presence for extras here; shape/broadcast is handled in node contracts
+        self._validate_extras_presence(extras)
+        return pnorm
+
+    def __call__(self, x, *, v=None, mask=None, extras=None, params=None):
+        if params is None:
+            if self.template.size == 0:
+                params = {}  # parameter-free PSF: auto-supply
+            else:
+                defaults = self.template.defaults()
+                if defaults is None:
+                    raise ValueError("PSF.__call__: params are required (template has parameters without defaults).")
+                params = defaults
+        params = self._coerce_and_check(params, extras)
+        y = self._caller(x, v, mask, extras, params)
+        if self.drop_features and self.n_features == 1:
+            return jnp.squeeze(y, axis=-1)
+        return y
+
+    def bind(self, params: dict[str, jax.Array] | None = None):
+        """Freeze parameter dict into an SF with normalized arrays.
+
+        If ``params is None``, fall back to spec defaults (``ParamSuite.defaults()``).
+        Raises if the template has parameters without defaults.
+        """
+        from .sf import SF
+
+        if params is None:
+            defaults = self.template.defaults()
+            if defaults is None:
+                raise ValueError("PSF.bind(): params are required (template has parameters without defaults).")
+            params = defaults
+        return SF(self, params, drop_features=self.drop_features)
+
+    def d_theta(self, *, mode: str = "auto"):
+        """Build an expression for the Jacobian w.r.t. parameters θ.
+
+        Shape effect
+        ------------
+        The final axis becomes `features × n_params_total`. Batch/pdepth/rank
+        prefixes are preserved exactly.
+
+        Notes
+        -----
+        The parameter PyTree is handled leafwise; each grad leaf is flattened over
+        its param part, then all leaves are concatenated along the final axis.
+        """
+        if self.root.param_suite is None:  # unreachable: enforced by __init__
+            raise AttributeError("Expression has no parameters to differentiate")
+        node = DerivativeNode(self.root, var="theta", mode=mode)
+        return PSF(node, drop_features=False)
+
+    # ------------- SciPy helpers ---------------------------------
+    @property
+    def labels(self):
+        """Basis labels from the underlying CoeffNode (if present)."""
+        from .nodes.ops.linear import CoeffNode
+
+        if isinstance(self.root, CoeffNode):
+            return self.root._basis_labels
+        _, labs, _ = self.root.flatten()
+        return labs
+
+    def flatten_params(self, params: dict[str, jax.Array]):
+        """Vectorize a parameter dict according to the template order."""
+        return self.template.vectorize(params)
+
+    def unflatten_params(self, vec: jax.Array):
+        """Materialize a parameter dict from a flat vector (inverse of `flatten_params`)."""
+        return self.template.materialize(vec)

SFI/statefunc/sf.py

@@ -0,0 +1,104 @@
+"""SF façade: state function with fixed parameters."""
+
+import equinox as eqx
+import jax
+import jax.numpy as jnp
+
+from .nodes import BaseNode
+from .psf import PSF
+from .stateexpr import StateExpr, _specialize_node
+
+
+# ============================================================================
+#  SF  – bound (θ-fixed) state-function
+# ============================================================================
+class SF(StateExpr):
+    """State-Function with θ fixed (a thin wrapper over the PSF’s root).
+
+    Behaves like a `Basis` for evaluation purposes (no `.d_theta()`), but you
+    can still build `.d_x()` / `.d_v()` expressions.
+
+    Feature axis handling mirrors the parent `PSF`:
+    if `drop_features=True` and `n_features==1`, the final axis is removed.
+    """
+
+    params: dict[str, jax.Array] = eqx.field(static=False, repr=False)
+    _psf: PSF = eqx.field(static=True, repr=False)
+    drop_features: bool = eqx.field(static=True, default=True)
+
+    def __init__(
+        self,
+        psf: PSF,
+        params: dict[str, jax.Array],
+        *,
+        drop_features: bool | None = None,
+    ):
+        super().__init__(psf.root)
+        params = psf.template.coerce(params)
+        object.__setattr__(self, "params", params)
+        object.__setattr__(self, "_psf", psf)
+        object.__setattr__(
+            self,
+            "drop_features",
+            psf.drop_features if drop_features is None else bool(drop_features),
+        )
+
+    def __call__(self, x, *, v=None, mask=None, extras=None):
+        """Evaluate the bound function on a **batched** input."""
+        self._validate_extras_presence(extras)
+
+        y = self._caller(x, v, mask, extras, self.params)
+
+        if self.drop_features and self.n_features == 1:
+            return jnp.squeeze(y, axis=-1)
+        return y
+
+    @property
+    def labels(self):
+        """Basis labels propagated from the parent PSF."""
+        return self._psf.labels
+
+    def _with_node(self, new_root: BaseNode):
+        new_psf = PSF(new_root, drop_features=self.drop_features)
+        return SF(new_psf, self.params, drop_features=self.drop_features)
+
+    def specialize(self, *, dataset: int) -> "SF":
+        """Specialize a *bound* function at condition ``dataset``.
+
+        Rewrites the graph (folding ``dataset_index``-reading leaves) and
+        projects the bound parameter values onto the shrunken template: a
+        per-condition spec whose shape loses a leading axis is sliced at
+        ``dataset``; shared specs are kept verbatim.
+        """
+        k = int(dataset)
+        new_root = _specialize_node(self.root, k)
+        new_psf = PSF(new_root, drop_features=self.drop_features)
+        new_params = _project_params(self.params, self._psf.template, new_psf.template, k)
+        return SF(new_psf, new_params, drop_features=self.drop_features)
+
+
+def _project_params(old_params, old_template, new_template, k: int) -> dict:
+    """Map a bound parameter dict onto a specialized template.
+
+    For each spec in ``new_template``: keep the old value when the shape is
+    unchanged; slice the leading axis at ``k`` when specialization dropped it
+    (the per-condition case, e.g. ``(K,) -> ()``).
+    """
+    out: dict = {}
+    for spec in new_template:
+        name = spec.name
+        if name not in old_params:
+            continue  # genuinely new param (none today); leave to template defaults
+        old_val = old_params[name]
+        old_shape = tuple(old_template[name].shape)
+        new_shape = tuple(spec.shape)
+        if old_shape == new_shape:
+            out[name] = old_val
+        elif len(old_shape) == len(new_shape) + 1 and old_shape[1:] == new_shape:
+            out[name] = old_val[k]  # per-condition slice
+        else:
+            raise ValueError(
+                f"Cannot project param {name!r} from shape {old_shape} to "
+                f"{new_shape} during specialize(dataset={k})."
+            )
+    return out