StochasticForceInference 2.0.0__py3-none-any.whl

SFI/statefunc/nodes/ops/linear.py

@@ -0,0 +1,153 @@
+# SFI/statefunc/nodes/ops/linear.py
+"""Linear algebra operator nodes: Dense and Coeff."""
+
+import equinox as eqx
+import jax.numpy as jnp
+
+from ...params import ParamSpec, ParamSuite
+from ..base import BaseNode, BaseOpNode
+from ..contract import _ContractMixin
+
+
+# ──────────────────────────────────────────────────────────────────────────────
+#  DenseNode  –  linear map on the feature axis of one child
+# ──────────────────────────────────────────────────────────────────────────────
+class DenseNode(BaseOpNode):
+    """
+    Affine map along the feature axis of a single child.
+
+    Parameters
+    ----------
+    n_out       : int                 - output feature dimension
+    weight      : str = "W"           - parameter, shape ``(n_in, n_out)``
+    bias        : str | None = "b"    - parameter, shape ``(n_out,)`` or None
+    """
+
+    weight_name: str = eqx.field(static=True)
+    bias_name: str | None = eqx.field(static=True)
+    _n_out: int = eqx.field(static=True)
+
+    # ------------------------------------------------------------------
+    def __init__(self, child: BaseNode, *, n_out: int, weight: str = "W", bias: str | None = "b"):
+        self.weight_name = weight
+        self.bias_name = bias
+        self._n_out = int(n_out)
+        super().__init__(child)
+
+        # own parameter template
+        W_spec = ParamSpec(weight, (child.n_features, n_out))
+        specs = [W_spec]
+        if bias is not None:
+            specs.append(ParamSpec(bias, (n_out,)))
+        own_suite = ParamSuite.from_specs(*specs)
+
+        merged = child.param_suite.merge(own_suite) if child.param_suite else own_suite
+        object.__setattr__(self, "param_suite", merged)
+
+    def with_children(self, new_children):
+        return DenseNode(
+            new_children[0],
+            n_out=self._n_out,
+            weight=self.weight_name,
+            bias=self.bias_name,
+        )
+
+    # ------------------------------------------------------------------
+    def _merge_static(self, children):
+        ch = children[0]
+        return _ContractMixin.inherit_contract(ch, n_features=self._n_out)
+
+    # ------------------------------------------------------------------
+    def _op(self, outs, *, params):
+        y_in = outs[0]
+        W = params[self.weight_name]
+        b = params[self.bias_name] if self.bias_name else None
+        y = jnp.tensordot(y_in, W, axes=[-1, 0])  # "...i,ij->...j"
+        return y + b if b is not None else y
+
+    # ------------------------------------------------------------------
+    #  Feature-wise flatten
+    # ------------------------------------------------------------------
+    def flatten(self):
+        """
+        Return (funcs, labels, descs) where
+
+        * ``funcs[j](...)``  calls the *DenseNode* itself and picks the
+          j-th feature slice → keeps maintenance trivial.
+        * Label is ``"dense{j}"`` as there's no practical way to propagate
+          previous labels through a Dense node.
+        * Descriptors are inherited from the **child** – the linear map
+          does not alter the mathematical identity, only its coefficients.
+        """
+
+        child_funcs, child_labels, child_descs = self.children[0].flatten()
+        n_out = self.n_features
+
+        # --- per-feature callables -----------------------------------
+        def _slice_feature(j):
+            # j is captured as default arg to avoid late binding
+            return lambda x, *, v=None, mask=None, extras=None, params=None, _j=j: (
+                self(x, v=v, mask=mask, extras=extras, params=params)[..., _j]
+            )
+
+        funcs = tuple(_slice_feature(j) for j in range(n_out))
+
+        # --- labels ---------------------------------------------------
+        labels = tuple(f"dense{j}" for j in range(n_out))
+
+        # --- descriptors (inherit) ------------------------------------
+        descs = tuple({"dense_of": tuple(child_descs), "index": j} for j in range(n_out))
+
+        return funcs, labels, descs
+
+
+# ──────────────────────────────────────────────────────────────────────────────
+#  CoeffNode – “Dense-with-θ-vector” (n_out = 1, bias = None)
+# ──────────────────────────────────────────────────────────────────────────────
+class CoeffNode(DenseNode):
+    """
+    Special-case of :class:`DenseNode` that stores its weights as a **vector**
+    θ ∈ ℝⁿ rather than a (n_in × 1) column matrix:
+
+        F(x; θ) = θ · f(x)           # dot on feature axis, returns scalar
+
+    The output therefore has a **singleton feature axis** (n_features = 1).
+
+    Parameters
+    ----------
+    child      : BaseNode
+        Basis/dictionary to be linearly combined.
+    coeff_key  : str, default ``"coeff"``
+        Parameter-dict name under which θ is stored.
+    """
+
+    _basis_labels: tuple[str, ...] = eqx.field(static=True, default=())
+
+    # ------------------------------------------------------------------
+    def __init__(self, child: BaseNode, *, coeff_key: str = "coeff"):
+        # call DenseNode constructor with n_out = 1, bias = None
+        super().__init__(child, n_out=1, weight=coeff_key, bias=None)
+
+        # Override the weight spec to a **vector** (n_in,) instead of matrix,
+        # with a default of ones so bare `Basis → .to_psf()` simulates out of
+        # the box (equivalent to "sum all features").
+        theta = ParamSpec(coeff_key, (child.n_features,), default=jnp.ones((child.n_features,)))
+        own = ParamSuite.from_specs(theta)
+        merged = child.param_suite.merge(own) if child.param_suite else own
+        object.__setattr__(self, "param_suite", merged)
+
+        # Preserve the original basis labels for downstream reporting
+        _, child_labels, _ = child.flatten()
+        object.__setattr__(self, "_basis_labels", tuple(child_labels))
+
+    def with_children(self, new_children):
+        return CoeffNode(new_children[0], coeff_key=self.weight_name)
+
+    # ------------------------------------------------------------------
+    # DenseNode._op expects W shaped (n_in, 1); we override to use a vector.
+    # ------------------------------------------------------------------
+    def _op(self, outs, *, params):
+        y_in = outs[0]  # (..., n_feat)
+        theta = params[self.weight_name]  # (n_feat,)
+        y = jnp.einsum("...b,b->...", y_in, theta)[..., None]  # (..., 1)
+        return y

SFI/statefunc/nodes/ops/mapn.py

@@ -0,0 +1,138 @@
+# SFI/statefunc/nodes/ops/mapn.py
+"""MapN: broadcast and apply N-ary function across children."""
+
+from typing import Callable
+
+import equinox as eqx
+
+from ...params import ParamSuite
+from ..base import BaseNode, BaseOpNode
+
+
+# ──────────────────────────────────────────────────────────────────────────────
+#  MapNNode – elementwise map of N children (with optional parameters)
+# ──────────────────────────────────────────────────────────────────────────────
+class MapNNode(BaseOpNode):
+    """
+    Apply an arbitrary **element-wise** callable to *N* child bases.
+
+    Parameters
+    ----------
+    mapper : Callable
+        A function that will be called as::
+
+            mapper(*arrays[, theta]) -> array
+
+        The *arrays* (one per child) have **identical shapes**, including feature axis.
+        If a ``mapper_param_suite`` is supplied, a θ-dict with **only** those names
+        is appended as the *last positional argument*; otherwise no parameter payload
+        is passed to the mapper.
+
+    mapper_param_suite : ParamSuite | None
+        Parameter template **owned by the mapper itself** (not the children).
+
+    Contract rules
+    --------------
+    - Children must match in (rank, n_features). ``dim`` is unified if concrete.
+    - Particle depth follows broadcast rules: you may mix (pdepth=0, particles_input=False)
+      with (pdepth=1, True). The output uses the latter, and non-P children treat P as part
+      of the batch prefix.
+    - The mapper **must** return a tensor with the same last-axis size (``n_features``),
+      otherwise ``_assert_outputs`` will trigger.
+
+    Notes
+    -----
+    This node's full ``param_suite`` is the **shared-by-name** merge of
+    ``(children's suites) ∪ (mapper_param_suite)``.
+
+    Reusing the same name across children/mapper **ties** that parameter provided
+    the shape/dtype match exactly; incompatible duplicates raise an error.
+    """
+
+    CONTRACT_MODE = "map"
+
+    mapper: Callable = eqx.field(static=True)
+    mapper_param_suite: ParamSuite | None = eqx.field(static=True)
+    label_fn: Callable[[str], str] | None = eqx.field(static=True, default=None, repr=False)
+
+    # ------------------------------------------------------------------
+    def __init__(
+        self,
+        mapper: Callable,
+        *children: BaseNode,
+        mapper_param_suite: ParamSuite | None = None,
+        label_fn: Callable[[str], str] | None = None,
+    ):
+        object.__setattr__(self, "mapper", mapper)
+        object.__setattr__(self, "mapper_param_suite", mapper_param_suite)
+        object.__setattr__(self, "label_fn", label_fn)
+
+        # Build base composite first: this sets `param_suite` from children
+        super().__init__(*children)
+
+        # Now incorporate the mapper's own params using shared-by-name merge
+        # (allows parameter tying across children and mapper when specs match).
+        merged = ParamSuite.merge_many(self.param_suite, self.mapper_param_suite)
+        object.__setattr__(self, "param_suite", merged)
+
+    def with_children(self, new_children):
+        return MapNNode(
+            self.mapper,
+            *new_children,
+            mapper_param_suite=self.mapper_param_suite,
+            label_fn=self.label_fn,
+        )
+
+    # ------------------------------------------------------------------
+    def _op(self, outs, *, params):
+        """
+        Combine child outputs element-wise via `mapper`.
+
+        The `params` dict passed here is the full suite; we slice out only the
+        mapper's subset (if any) and append it as the final positional argument.
+        """
+        if self.mapper_param_suite:
+            theta = {n: params[n] for n in self.mapper_param_suite._lookup}
+            return self.mapper(*outs, theta)
+        return self.mapper(*outs)
+
+    # ------------------------------------------------------------------
+    #  Feature-aware flatten  (N children, element-wise map)
+    # ------------------------------------------------------------------
+    def flatten(self):
+        """
+        Build (funcs, labels, descs) where *n_features* matches children exactly.
+
+        - funcs[i](x, ...) := mapper( child0_i(x), child1_i(x), …, θ? )
+        - labels are produced by `label_fn` if given, otherwise
+          `"{mapper_name}(" + ",".join(child_labels) + ")"`.
+        - descs  = tuple of children’s desc[i]  (one per child)  for traceability.
+        """
+        # ---- gather children ----
+        child_flat = [ch.flatten() for ch in self.children]
+        funcs_lists, labels_lists, descs_lists = zip(*child_flat)
+        n_feat = len(labels_lists[0])  # already validated in _merge_static
+
+        # --- per-feature callables -----------------------------------
+        def _slice_feature(j):
+            # j is captured as default arg to avoid late binding
+            return lambda x, *, v=None, mask=None, extras=None, params=None, _j=j: (
+                self(x, v=v, mask=mask, extras=extras, params=params)[..., _j]
+            )
+
+        funcs = tuple(_slice_feature(j) for j in range(n_feat))
+
+        # ---- labels ----
+        if self.label_fn is not None:
+            labels = tuple(self.label_fn(*(lablist[i] for lablist in labels_lists)) for i in range(n_feat))
+        else:
+            mname = self.mapper.__name__ if getattr(self.mapper, "__name__", "") else "map"
+            labels = tuple(f"{mname}(" + ",".join(lab_tuple) + ")" for lab_tuple in zip(*labels_lists))
+
+        # ---- descriptors ----
+        descs = [
+            tuple(dl[i] for dl in descs_lists)  # one tuple per child
+            for i in range(n_feat)
+        ]
+
+        return funcs, labels, descs

SFI/statefunc/nodes/ops/reshape_rank.py

@@ -0,0 +1,158 @@
+# SFI/statefunc/nodes/ops/reshape_rank.py
+"""Lossless reshape between spatial (rank) axes and the feature axis."""
+
+from functools import reduce
+from operator import mul
+
+import equinox as eqx
+import jax.numpy as jnp
+
+from ..base import BaseNode, BaseOpNode
+from ..contract import Rank, _ContractMixin
+
+
+def _rank_axis_sizes(node: _ContractMixin) -> tuple[int, ...]:
+    """Return per-axis sizes for rank axes: sdims if present, else (dim,)^rank."""
+    if node.sdims is not None:
+        return node.sdims
+    if node.dim is None:
+        raise ValueError("ReshapeRankNode requires known axis sizes (dim or sdims); got dim=None, sdims=None.")
+    return (node.dim,) * int(node.rank)
+
+
+class ReshapeRankNode(BaseOpNode):
+    """Lossless reshape that moves data between rank (spatial) axes and features.
+
+    Given a child with output layout::
+
+        batch · (dim,)^source_rank · n_features
+
+    this node reshapes to::
+
+        batch · (dim,)^target_rank · n_features'
+
+    where  ``n_features' = n_features x dim^(source_rank - target_rank)``
+    (folding when *target_rank < source_rank*) or
+    ``n_features' = n_features / dim^(target_rank - source_rank)``
+    (unfolding when *target_rank > source_rank*).
+
+    The reshape uses C-order so that folding followed by unfolding back
+    to the original rank is the identity — the pair is invertible.
+
+    Parameters
+    ----------
+    child : BaseNode
+        Source node (must have known ``dim``).
+    target_rank : int or Rank
+        Desired output rank.
+
+    Notes
+    -----
+    * Folding merges the *innermost* spatial axes into the feature axis.
+    * Unfolding splits the feature axis to create new *innermost* spatial axes.
+    * The child's ``dim`` must be known (not ``None``).
+    * For unfolding, ``n_features`` must be divisible by ``dim^Δrank``.
+    """
+
+    _target_rank: Rank = eqx.field(static=True)
+
+    def __init__(self, child: BaseNode, *, target_rank: int | Rank):
+        object.__setattr__(self, "_target_rank", Rank(int(target_rank)))
+
+        if child.dim is None and child.sdims is None:
+            raise ValueError("ReshapeRankNode requires a child with known dim or sdims; got dim=None, sdims=None.")
+        src = int(child.rank)
+        tgt = int(self._target_rank)
+
+        if tgt > src:
+            # Unfolding: need to split the feature axis. Determine how many
+            # elements need to be carved out. For uniform dim this is dim^(tgt-src).
+            # For sdims, the target rank axes beyond the source's sdims are unknown
+            # at this point — we require uniform dim for unfolding (sdims=None or
+            # all dims equal).
+            if child.sdims is not None:
+                raise ValueError(
+                    "ReshapeRankNode: unfolding with non-uniform sdims is not "
+                    "supported. Use explicit reshape or einsum instead."
+                )
+            assert child.dim is not None  # guaranteed: sdims is None and guard on L62 ensures one is non-None
+            factor = child.dim ** (tgt - src)
+            if child.n_features % factor != 0:
+                raise ValueError(
+                    f"Cannot unfold {child.n_features} features into rank-{tgt} "
+                    f"with dim={child.dim}: {child.n_features} is not divisible "
+                    f"by dim^{tgt - src} = {factor}."
+                )
+
+        super().__init__(child)
+
+    # ------------------------------------------------------------------
+    def _merge_static(self, children):
+        ch = children[0]
+        src = int(ch.rank)
+        tgt = int(self._target_rank)
+        axis_sizes = _rank_axis_sizes(ch)
+
+        if tgt < src:
+            # Folding: innermost (src − tgt) rank axes → features
+            factor = reduce(mul, axis_sizes[tgt:], 1)
+            n_out = ch.n_features * factor
+            out_sdims = axis_sizes[:tgt] if ch.sdims is not None else None
+        elif tgt > src:
+            # Unfolding: features → new innermost rank axes (uniform dim only)
+            dim = ch.dim
+            factor = dim ** (tgt - src)
+            n_out = ch.n_features // factor
+            # Uniform unfolding: new axes are all dim-sized
+            out_sdims = axis_sizes + (dim,) * (tgt - src) if ch.sdims is not None else None
+        else:
+            n_out = ch.n_features
+            out_sdims = ch.sdims
+
+        return _ContractMixin.inherit_contract(ch, rank=self._target_rank, n_features=n_out, sdims=out_sdims)
+
+    # ------------------------------------------------------------------
+    def _op(self, outs, *, params):
+        y = outs[0]
+        src = int(self.children[0].rank)
+        tgt = int(self._target_rank)
+
+        if tgt == src:
+            return y
+
+        if tgt < src:
+            # Fold innermost (src − tgt) spatial axes + feature axis into one.
+            n_fold = src - tgt
+            new_shape = y.shape[: -(n_fold + 1)] + (-1,)
+            return jnp.reshape(y, new_shape)
+        else:
+            # Unfold feature axis into (tgt − src) new spatial axes + features.
+            n_unfold = tgt - src
+            dim = self.children[0].dim
+            n_feat_out = self.n_features
+            new_shape = y.shape[:-1] + (dim,) * n_unfold + (n_feat_out,)
+            return jnp.reshape(y, new_shape)
+
+    # ------------------------------------------------------------------
+    def flatten(self):
+        child_funcs, child_labels, child_descs = self.children[0].flatten()
+
+        n_out = self.n_features
+
+        def _slice_feature(j):
+            return lambda x, *, v=None, mask=None, extras=None, params=None, _j=j: (
+                self(x, v=v, mask=mask, extras=extras, params=params)[..., _j]
+            )
+
+        src = int(self.children[0].rank)
+        tgt = int(self._target_rank)
+        tag = "fold" if tgt < src else "unfold"
+
+        funcs = [_slice_feature(j) for j in range(n_out)]
+        labels = [f"{tag}_{j}" for j in range(n_out)]
+        descs = [{f"{tag}_of": [d for d in child_descs], "index": j} for j in range(n_out)]
+        return funcs, labels, descs
+
+    # ------------------------------------------------------------------
+    def _tree_id(self):
+        return hash(("ReshapeRankNode", int(self._target_rank), self.children[0]._tree_id()))

SFI/statefunc/nodes/ops/slice.py

@@ -0,0 +1,83 @@
+# SFI/statefunc/nodes/ops/slice.py
+"""Slice features along the feature axis."""
+
+from typing import Sequence
+
+import equinox as eqx
+
+from ..base import BaseNode, BaseOpNode
+from ..contract import _ContractMixin
+
+
+# ---------------------------------------------------------------------
+# Helper node – slice the **feature axis**
+# ---------------------------------------------------------------------
+class SliceFeaturesNode(BaseOpNode):
+    """Select or reorder *feature* entries of a child basis.
+
+    The node simply indexes the last axis (feature axis) of its child
+    while preserving every spatial axis.  Supported *idx* forms are
+
+    * ``int``               - keep a single feature → ``n_features = 1``
+    * ``slice``             - standard Python slice
+    * ``Sequence[int]``     - fancy indexing/ordering
+    * ``Sequence[bool]``    - boolean mask (length must equal ``n_features``)
+
+    Notes
+    -----
+    - Rank, dim, pdepth & needs_v are *identical* to the child - only
+      ``n_features`` changes.
+    - No parameters or extras are consumed, so the node is deterministic.
+    """
+
+    idx: int | slice | Sequence[int] | Sequence[bool] = eqx.field(static=True)
+
+    # ---------------- constructor ----------------
+    def __init__(self, child: BaseNode, idx):
+        object.__setattr__(self, "idx", idx)
+        super().__init__(child)
+        object.__setattr__(self, "param_suite", child.param_suite)
+
+    # ------------- static contract merge ----------
+    def _merge_static(self, children):
+        ch = children[0]
+        n_out: int
+        if isinstance(self.idx, int):
+            n_out = 1
+        elif isinstance(self.idx, slice):
+            n_out = len(range(*self.idx.indices(ch.n_features)))
+        else:  # fancy or boolean indexing
+            n_out = len(self.idx)
+        return _ContractMixin.inherit_contract(ch, n_features=n_out)
+
+    # ---------------- combine ---------------------
+    def _op(self, outs, *, params):
+        y = outs[0]
+        if isinstance(self.idx, int):
+            # Keep the feature axis (size 1), including for negative indices.
+            return y[..., [self.idx]]
+        return y[..., self.idx]
+
+    # ---------------- flatten (label-aware) --------
+    def flatten(self):
+        """Select the labels/funcs/descriptors matching the sliced features."""
+        child_funcs, child_labels, child_descs = self.children[0].flatten()
+
+        # Resolve idx to a concrete list of integer indices
+        n_child = len(child_funcs)
+        if isinstance(self.idx, int):
+            indices = [self.idx % n_child]
+        elif isinstance(self.idx, slice):
+            indices = list(range(*self.idx.indices(n_child)))
+        else:
+            seq = list(self.idx)
+            if seq and isinstance(seq[0], bool):
+                # boolean mask
+                indices = [i for i, b in enumerate(seq) if b]
+            else:
+                indices = [int(i) % n_child for i in seq]
+
+        funcs = [child_funcs[i] for i in indices]
+        labels = [child_labels[i] for i in indices]
+        descs = [child_descs[i] for i in indices]
+        return funcs, labels, descs