StochasticForceInference 2.0.0__py3-none-any.whl

SFI/statefunc/layout/_sectors.py

@@ -0,0 +1,166 @@
+"""Sector types for the structured dimensions layer.
+
+Defines :class:`ScalarSector`, :class:`VectorSector`,
+:class:`SymTensorSector`, :class:`TensorSector` and the convenience
+type alias :data:`Sector`.
+"""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+
+# =====================================================================
+# Sector types
+# =====================================================================
+
+
+@dataclass(frozen=True, slots=True)
+class ScalarSector:
+    """A single scalar field occupying one data index.
+
+    ``sdims = ()``, ``n_data = 1``.
+    """
+
+    indices: tuple[int, ...]
+
+    def __post_init__(self) -> None:
+        object.__setattr__(self, "indices", tuple(self.indices))
+        if len(self.indices) != 1:
+            raise ValueError(f"ScalarSector requires exactly 1 index, got {len(self.indices)}: {self.indices}")
+
+    @property
+    def sdims(self) -> tuple[int, ...]:
+        return ()
+
+    @property
+    def n_data(self) -> int:
+        return 1
+
+
+@dataclass(frozen=True, slots=True)
+class VectorSector:
+    """A vector field with *sdim* components.
+
+    ``sdims = (sdim,)``, ``n_data = sdim``.
+
+    Parameters
+    ----------
+    spatial : bool
+        When ``True``, declares that this sector's components correspond
+        to spatial coordinates.  Layouts with ``ndim`` validate that
+        ``sdim == ndim``.
+    """
+
+    indices: tuple[int, ...]
+    sdim: int
+    spatial: bool = False
+
+    def __post_init__(self) -> None:
+        object.__setattr__(self, "indices", tuple(self.indices))
+        if self.sdim < 1:
+            raise ValueError(f"sdim must be >= 1, got {self.sdim}")
+        if len(self.indices) != self.sdim:
+            raise ValueError(f"VectorSector(sdim={self.sdim}) requires {self.sdim} indices, got {len(self.indices)}")
+
+    @property
+    def sdims(self) -> tuple[int, ...]:
+        return (self.sdim,)
+
+    @property
+    def n_data(self) -> int:
+        return self.sdim
+
+
+@dataclass(frozen=True, slots=True)
+class SymTensorSector:
+    """A symmetric tensor field, Voigt-packed in data space.
+
+    ``sdims = (sdim, sdim)``.
+
+    When ``traceless=False`` (default), ``n_data = sdim * (sdim + 1) // 2``
+    (full Voigt packing).  When ``traceless=True``,
+    ``n_data = sdim * (sdim + 1) // 2 - 1`` (the trace degree of freedom
+    is removed: the diagonal is constrained so that
+    :math:`Q_{00} + Q_{11} + \\ldots = 0`).
+
+    Parameters
+    ----------
+    traceless : bool
+        When ``True``, the tensor is both symmetric *and* traceless.
+        Only the independent components are stored in data space;
+        the last diagonal entry is reconstructed as minus the sum
+        of the other diagonal entries.
+    """
+
+    indices: tuple[int, ...]
+    sdim: int
+    traceless: bool = False
+
+    def __post_init__(self) -> None:
+        object.__setattr__(self, "indices", tuple(self.indices))
+        expected = self._n_independent
+        if len(self.indices) != expected:
+            kind = "traceless symmetric" if self.traceless else "symmetric"
+            raise ValueError(
+                f"SymTensorSector(sdim={self.sdim}, traceless={self.traceless}) "
+                f"requires {expected} indices ({kind} packing), "
+                f"got {len(self.indices)}"
+            )
+
+    @property
+    def _n_independent(self) -> int:
+        """Number of independent components stored in data space."""
+        n_voigt = self.sdim * (self.sdim + 1) // 2
+        return n_voigt - 1 if self.traceless else n_voigt
+
+    @property
+    def sdims(self) -> tuple[int, ...]:
+        return (self.sdim, self.sdim)
+
+    @property
+    def n_data(self) -> int:
+        return self._n_independent
+
+    @property
+    def voigt_pairs(self) -> list[tuple[int, int]]:
+        """Upper-triangle ``(i, j)`` pairs matching the index order.
+
+        When ``traceless=True`` the last diagonal pair
+        ``(sdim-1, sdim-1)`` is omitted (it is reconstructed from the
+        other diagonal entries).
+        """
+        pairs: list[tuple[int, int]] = []
+        for i in range(self.sdim):
+            for j in range(i, self.sdim):
+                pairs.append((i, j))
+        if self.traceless:
+            # Remove the last diagonal entry (sdim-1, sdim-1)
+            pairs = [p for p in pairs if p != (self.sdim - 1, self.sdim - 1)]
+        return pairs
+
+
+@dataclass(frozen=True, slots=True)
+class TensorSector:
+    """A general tensor field with arbitrary ``sdims``.
+
+    ``n_data = prod(sdims)``.
+    """
+
+    indices: tuple[int, ...]
+    sdims: tuple[int, ...]
+
+    def __post_init__(self) -> None:
+        object.__setattr__(self, "indices", tuple(self.indices))
+        object.__setattr__(self, "sdims", tuple(self.sdims))
+        expected = math.prod(self.sdims)
+        if len(self.indices) != expected:
+            raise ValueError(f"TensorSector(sdims={self.sdims}) requires {expected} indices, got {len(self.indices)}")
+
+    @property
+    def n_data(self) -> int:
+        return math.prod(self.sdims)
+
+
+# Convenience type alias
+Sector = ScalarSector | VectorSector | SymTensorSector | TensorSector

SFI/statefunc/memhint.py

@@ -0,0 +1,222 @@
+# SFI/statefunc/memhint.py
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Iterable, Optional, Protocol
+
+import jax.numpy as jnp
+
+# ──────────────────────────────────────────────────────────────────────────────
+# Core types
+# ──────────────────────────────────────────────────────────────────────────────
+
+
+@dataclass(frozen=True)
+class MemHint:
+    """
+    Conservative memory footprint per SINGLE sample.
+    - per_sample_bytes scales with the number of samples.
+    - persistent_bytes is model state / CSR / weights, does NOT scale with samples.
+    """
+
+    per_sample_bytes: int = 0
+    persistent_bytes: int = 0
+
+    def __add__(self, other: "MemHint") -> "MemHint":
+        return MemHint(
+            self.per_sample_bytes + other.per_sample_bytes,
+            self.persistent_bytes + other.persistent_bytes,
+        )
+
+    def scaled(self, k: int) -> "MemHint":
+        """Scale only the transient part (useful if you ever convert to total bytes)."""
+        return MemHint(self.per_sample_bytes * int(k), self.persistent_bytes)
+
+
+@dataclass(frozen=True)
+class SampleMeta:
+    """
+    Optional single-sample context to refine estimates.
+
+    P: number of particles in ONE sample (for nodes with particle axes, pdepth>0).
+    K: arity for interaction gathers when fixed/known (pairs=2, etc.).
+    has_v: whether a velocity block will be present.
+    has_mask: whether a boolean mask participates in the call.
+    """
+
+    P: Optional[int] = None
+    K: Optional[int] = None
+    has_v: Optional[bool] = None
+    has_mask: Optional[bool] = None
+
+    @staticmethod
+    def from_arrays(x=None, v=None, mask=None) -> "SampleMeta":
+        """
+        Try to recover P from provided single-sample arrays without allocating anything.
+        Heuristics:
+
+          - If x has at least 2 dims and looks like (..., dim) with a leading particle axis,
+            we guess P from the axis before `dim`. If ambiguous, we leave P=None.
+          - v and mask only toggle flags; they don't change P.
+        """
+        P = None
+        if x is not None and hasattr(x, "shape"):
+            # super conservative: if it looks like (..., P, dim) with small dim
+            if x.ndim >= 2:
+                dim_guess = int(x.shape[-1])
+                if 1 <= dim_guess <= 16:
+                    # take the previous axis as P
+                    P = int(x.shape[-2])
+        return SampleMeta(
+            P=P,
+            K=None,
+            has_v=(v is not None),
+            has_mask=(mask is not None),
+        )
+
+
+# ──────────────────────────────────────────────────────────────────────────────
+# Helpers
+# ──────────────────────────────────────────────────────────────────────────────
+
+
+def itemsize_of(dtype) -> int:
+    """Return dtype itemsize as an int; defaults to float32 when dtype is None."""
+    return int(jnp.dtype(jnp.float32 if dtype is None else dtype).itemsize)
+
+
+class _HasContract(Protocol):
+    # Minimal surface we read from nodes (already present on your BaseNode/ops).
+    rank: int
+    dim: Optional[int]
+    pdepth: int
+    n_features: int
+
+
+def resolve_P(particle_size: Optional[int], sample) -> Optional[int]:
+    """Pick P from explicit particle_size, else from SampleMeta or array, else None."""
+    if particle_size is not None:
+        return int(particle_size)
+    if sample is None:
+        return None
+
+    # Accept either a SampleMeta or an array-like (single-sample x)
+    if isinstance(sample, SampleMeta):
+        sm = sample
+    else:
+        # best-effort duck-typing: treat it as x and try to read P
+        try:
+            sm = SampleMeta.from_arrays(x=sample)
+        except Exception:
+            sm = None
+
+    return None if (sm is None or sm.P is None) else int(sm.P)
+
+
+def output_elems_per_sample(node: _HasContract, *, particle_size: Optional[int]) -> int:
+    """
+    Count output elements of a node for ONE sample from its static contract.
+
+    Output shape suffix is ``(*particle_axes, *rank_axes, n_features)``
+    with ``particle_axes = (P,) * pdepth`` and ``rank_axes = (dim,) * rank``.
+
+    If ``particle_size`` is None, treat ``P = 1`` (safe lower bound for
+    batch picking).
+    """
+    n_features = int(getattr(node, "n_features", 0) or 0)
+    if n_features <= 0:
+        return 0
+
+    dim = getattr(node, "dim", None)
+    rank = int(getattr(node, "rank", 0) or 0)
+    if dim is None and rank > 0:
+        return 0  # cannot infer rank block without dim
+
+    pdepth = int(getattr(node, "pdepth", 0) or 0)
+    P = 1 if particle_size is None else int(particle_size)
+
+    elems = (dim or 1) ** rank
+    if pdepth:
+        elems *= P**pdepth
+    elems *= n_features
+    return int(elems)
+
+
+def output_bytes_per_sample(node: _HasContract, *, dtype, particle_size: Optional[int]) -> int:
+    """Translate element count to bytes using dtype itemsize."""
+    return output_elems_per_sample(node, particle_size=particle_size) * itemsize_of(dtype)
+
+
+def default_leaf_hint(node: _HasContract, *, dtype, particle_size: Optional[int], mode: str) -> MemHint:
+    """
+    Default for leaf-like nodes: count only the output buffer per sample.
+    Composite nodes will also include their children.
+    """
+    return MemHint(per_sample_bytes=output_bytes_per_sample(node, dtype=dtype, particle_size=particle_size))
+
+
+def broadcast_extra_bytes_for_children(*, children: Iterable, dtype, particle_size: Optional[int]) -> int:
+    """
+    When children have different particle depths, lower-depth outputs are broadcast
+    to match the max. Materializing that broadcast costs memory. We conservatively
+    account for an additional slab equal to (P^Δ - 1) times the child's OUTPUT bytes.
+    """
+    # target = max pdepth among children
+    target = 0
+    pd = []
+    for ch in children:
+        p = int(getattr(ch, "pdepth", 0) or 0)
+        pd.append(p)
+        if p > target:
+            target = p
+    if target == 0:
+        return 0
+
+    extra = 0
+    P = 1 if particle_size is None else int(particle_size)
+    for ch in children:
+        p = int(getattr(ch, "pdepth", 0) or 0)
+        if p < target:
+            delta = target - p
+            # only the OUTPUT slab needs broadcasting; not the child's internal working set
+            base_out = output_bytes_per_sample(ch, dtype=dtype, particle_size=particle_size)
+            factor = P**delta - 1
+            if factor > 0 and base_out > 0:
+                extra += base_out * factor
+    return int(extra)
+
+
+def default_op_hint(
+    node: _HasContract,
+    *,
+    children: Iterable,
+    dtype,
+    particle_size: Optional[int],
+    mode: str,
+) -> MemHint:
+    """
+    Default for composite ops: sum(children.hint) + my output buffer + broadcast overhead.
+    We assume child outputs coexist while the op constructs its result.
+    """
+    hint = MemHint()
+    for ch in children:
+        hint = hint + ch.memory_hint(dtype=dtype, particle_size=particle_size, mode=mode)
+    # add broadcast materialization if children pdepth differ
+    hint = MemHint(
+        hint.per_sample_bytes
+        + broadcast_extra_bytes_for_children(children=children, dtype=dtype, particle_size=particle_size),
+        hint.persistent_bytes,
+    )
+    # add this op's own output
+    return hint + default_leaf_hint(node, dtype=dtype, particle_size=particle_size, mode=mode)
+
+
+def inflate_for_grad(hint: MemHint, *, factor: float = 2.0) -> MemHint:
+    """
+    Blanket inflation for gradient-mode nodes to account for tangents/tapes.
+    Keep it simple and conservative; adjust locally if you get better bounds.
+    """
+    return MemHint(
+        per_sample_bytes=int(hint.per_sample_bytes * factor),
+        persistent_bytes=hint.persistent_bytes,
+    )

SFI/statefunc/nodes/__init__.py

@@ -0,0 +1,56 @@
+# SFI/statefunc/nodes/__init__.py
+"""
+Developer-facing node classes.
+
+Most users should import from `SFI.statefunc` (Basis/PSF/SF).
+This module exposes the node types for advanced composition.
+"""
+
+from .base import BaseNode, BaseOpNode
+from .contract import Rank
+from .leaf import InteractionLeaf, SimpleLeaf
+from .ops import (
+    CoeffNode,
+    ConcatNode,
+    DenseNode,
+    DerivativeNode,
+    EinsumNode,
+    MapNNode,
+    ReshapeRankNode,
+    SliceFeaturesNode,
+)
+
+__all__ = [
+    "BaseNode",
+    "BaseOpNode",
+    "SimpleLeaf",
+    "ConcatNode",
+    "MapNNode",
+    "EinsumNode",
+    "DenseNode",
+    "CoeffNode",
+    "SliceFeaturesNode",
+    "DerivativeNode",
+    "ReshapeRankNode",
+    "Rank",
+]
+
+
+from .interactions import (
+    AutoPairs,
+    FromExtrasPairsCSR,
+    HyperCSR,
+    HyperFixed,
+    InteractionDispatcher,
+    PairsCSR,
+)
+
+__all__ += [
+    "InteractionDispatcher",
+    "PairsCSR",
+    "HyperFixed",
+    "HyperCSR",
+    "AutoPairs",
+    "FromExtrasPairsCSR",
+    "InteractionLeaf",
+]