onf-c2/src/onf/graph/core/nodes.py

"""The node table: one row per short, homogeneous run of consecutive demo frames.

Turns loaded demonstrations into g_nodes.npz artifacts.

Components:
    NodeTableBuilder: demos -> NodeTable
    NodeTableValidator: invariant enforcement
    NodeTableCodec: g_nodes.npz read/write
    NodeFeaturizer: table -> network features

WHY THIS SHAPE, why t_frac and phase are two different fields, why the feature
vector excludes task_id/stage, and why coarse qdot is an endpoint
difference rather than an average: see docs/technical/01-data-and-graph.md.
"""

from __future__ import annotations

import os
from dataclasses import dataclass, fields
from functools import lru_cache
from pathlib import Path
from typing import TYPE_CHECKING, Any, Final, Mapping, Sequence

import numpy as np
from numpy.typing import ArrayLike, NDArray

from onf.blend.ee_track import ACTION_DIM, POSE_DIM
from onf.graph.core import schema

if TYPE_CHECKING:
    import torch

__all__ = [
    "DemoRecord",
    "FeatureStats",
    "NodeFeaturizer",
    "NodeTable",
    "NodeTableBuilder",
    "NodeTableCodec",
    "NodeTableValidator",
    "Segment",
    "TimeEncoder",
    "coarsen_segments",
    "episode_phase",
    "psi",
]

_DTYPE: Final[dict[str, np.dtype]] = {
    "f4": np.dtype(np.float32),
    "u1": np.dtype(np.uint8),
    "i2": np.dtype(np.int16),
    "i4": np.dtype(np.int32),
}

_PER_NODE_FIELDS: Final[tuple[str, ...]] = (
    "q", "qdot", "grip", "stage", "t_idx", "t_frac", "t_raw",
    "n_members", "owner", "task_id", "phase",
)

_FINITE_FIELDS: Final[tuple[str, ...]] = (
    "q", "qdot", "t_frac", "phase", "q_raw", "qdot_raw", "ee_raw", "a_raw",
)

_PHASE_MAX_MIN: Final[float] = 0.99
_PHASE_MIN_MAX: Final[float] = 0.01
_SUMMARY_ATOL: Final[float] = 1e-4
_STD_FLOOR: Final[float] = 1e-6


# ==================================================================================
# Time / phase primitives
# ==================================================================================
def episode_phase(n_frames: int) -> NDArray[np.float32]:
    """Computes canonical phase over a full episode.

    This is schema.PHASE_CONV: the single definition every other notion of "how
    far along" must derive from. Never re-invent it.

    Args:
        n_frames: Episode length T.

    Returns:
        [T] f32 array spanning [0, 1].
    """
    n = int(n_frames)
    return (np.arange(n, dtype=np.float32) / max(1, n - 1)).astype(np.float32)


class TimeEncoder:
    """Sinusoidal encoder for a normalised time fraction."""

    def __init__(self, n_freq: int = schema.PSI_FREQS) -> None:
        """Initializes the frequency ladder.

        Args:
            n_freq: Number of octaves.

        Raises:
            ValueError: If n_freq is not positive.
        """
        if n_freq < 1:
            raise ValueError(f"n_freq must be >= 1, got {n_freq}")
        self._n_freq = int(n_freq)
        self._freqs = (2.0 ** np.arange(self._n_freq, dtype=np.float64)) * np.pi

    @property
    def width(self) -> int:
        """Returns the number of output columns."""
        return 2 * self._n_freq

    def encode(self, t_frac: ArrayLike) -> NDArray[np.float32]:
        """Encodes time fractions.

        Args:
            t_frac: [N] values, nominally in [0, 1].

        Returns:
            [N, 2 * n_freq] f32 array.
        """
        flat = np.asarray(t_frac, dtype=np.float32).reshape(-1)
        angles = flat[:, None].astype(np.float64) * self._freqs[None, :]
        out = np.empty((flat.shape[0], self.width), dtype=np.float32)
        out[:, 0::2] = np.sin(angles)
        out[:, 1::2] = np.cos(angles)
        return out


@lru_cache(maxsize=8)
def _time_encoder(n_freq: int) -> TimeEncoder:
    """Cached encoder for repeated feature calls.

    Args:
        n_freq: Number of octaves.

    Returns:
        A shared TimeEncoder instance.
    """
    return TimeEncoder(n_freq)


def psi(t_frac: ArrayLike, n_freq: int = schema.PSI_FREQS) -> NDArray[np.float32]:
    """Sinusoidal time encoding wrapper.

    Args:
        t_frac: [N] values in [0, 1].
        n_freq: Number of octaves.

    Returns:
        [N, 2 * n_freq] f32 array.
    """
    return _time_encoder(n_freq).encode(t_frac)


def coarsen_segments(grip: ArrayLike, coarsen: int) -> NDArray[np.int32]:
    """Assigns coarse-node ids within one demo.

    Groups runs of up to coarsen consecutive frames. Cuts a group short when
    grip changes, so every node has exactly one gripper state. Deliberately does
    NOT cut on stage -- see docs/technical/01-data-and-graph.md.

    Args:
        grip: [T] u1 per-frame grasp flags.
        coarsen: Maximum frames folded into one node.

    Returns:
        [T] i4 node ids, non-decreasing.

    Raises:
        ValueError: If coarsen < 1.
    """
    if coarsen < 1:
        raise ValueError(f"coarsen must be >= 1, got {coarsen}")
    grip_arr = np.asarray(grip)
    n = len(grip_arr)
    if n == 0:
        return np.zeros(0, dtype=np.int32)

    idx = np.arange(n)
    changed = np.empty(n, dtype=bool)
    changed[0] = True
    changed[1:] = grip_arr[1:] != grip_arr[:-1]

    last_reset = np.maximum.accumulate(np.where(changed, idx, -1))
    pos_in_run = idx - last_reset
    is_new_node = (pos_in_run % coarsen) == 0
    return (np.cumsum(is_new_node) - 1).astype(np.int32)


# ==================================================================================
# Value objects
# ==================================================================================
def _as_f32_or_none(value: ArrayLike | None) -> NDArray[np.float32] | None:
    """Casts an optional array to f32, passing None through.

    Args:
        value: Array-like, or None.

    Returns:
        The f32 array, or None.
    """
    return None if value is None else np.asarray(value, dtype=np.float32)


@dataclass(slots=True)
class DemoRecord:
    """One demonstration trajectory.

    Attributes:
        q: [T, D] f32 joint positions.
        qdot: [T, D] f32 per-raw-frame velocities.
        grip: [T] u1, 1 = grasping.
        stage: [T] i2 release count.
        task_id: Index into task_names.
        ee: [T, 6] f32 end-effector pose, concat(position, unnormalized axis-angle orientation).
        a: [T, 7] f32 recorded action commands.
    """
    q: NDArray[np.float32]
    qdot: NDArray[np.float32]
    grip: NDArray[np.uint8]
    stage: NDArray[np.int16]
    task_id: int
    ee: NDArray[np.float32] | None = None
    a: NDArray[np.float32] | None = None

    def __post_init__(self) -> None:
        # Only a synthetic demo reaches here without the two: onf.graph.build.from_demos reads
        # both off the HDF5 and raises when either is absent.
        if self.ee is None:
            self.ee = np.zeros((len(self), POSE_DIM), dtype=np.float32)
        if self.a is None:
            self.a = np.zeros((len(self), ACTION_DIM), dtype=np.float32)

    def __len__(self) -> int:
        return int(self.q.shape[0])

    @classmethod
    def from_mapping(cls, payload: Mapping[str, Any]) -> DemoRecord:
        """Adapts a legacy per-demo dict to a DemoRecord.

        Args:
            payload: Mapping with keys q, qdot, grip, stage, task_id, and optionally ee and a.

        Returns:
            The DemoRecord instance.
        """
        return cls(
            q=np.asarray(payload["q"], dtype=np.float32),
            qdot=np.asarray(payload["qdot"], dtype=np.float32),
            grip=np.asarray(payload["grip"]),
            stage=np.asarray(payload["stage"]),
            task_id=int(payload["task_id"]),
            ee=_as_f32_or_none(payload.get("ee")),
            a=_as_f32_or_none(payload.get("a")),
        )


@dataclass(frozen=True, slots=True)
class Segment:
    """A contiguous run of raw frames decoded from one node.

    Unpacks as q, qdot -- the two-tuple shape every call site reads.

    Attributes:
        q: [k, D] f64 positions.
        qdot: [k, D] f64 velocities.
    """
    q: NDArray[np.float64]
    qdot: NDArray[np.float64]

    def __len__(self) -> int:
        return int(self.q.shape[0])

    def __iter__(self):
        """Yields q then qdot, so q, qdot = table.segment(...) works."""
        yield self.q
        yield self.qdot


@dataclass(frozen=True, slots=True)
class FeatureStats:
    """Normalisation statistics over the unbounded feature prefix.

    Unpacks as mean, std.

    Attributes:
        mean: [2D+1] f32 column means.
        std: [2D+1] f32 column stds, floored away from zero.
    """
    mean: NDArray[np.float32]
    std: NDArray[np.float32]

    def __iter__(self):
        """Yields mean then std, so mean, std = table.feature_stats() works."""
        yield self.mean
        yield self.std

    def normalize(self, features: NDArray[np.float32]) -> NDArray[np.float32]:
        """Applies statistics to the prefix columns of a feature block.

        Args:
            features: [N, F] raw features.

        Returns:
            [N, F] f32 with the prefix normalised.
        """
        out = np.array(features, dtype=np.float32, copy=True)
        width = self.mean.shape[0]
        out[:, :width] = (out[:, :width] - self.mean) / self.std
        return out


@dataclass(slots=True)
class _DemoNodes:
    """Coarse-node arrays for a single demo."""
    q: NDArray[np.float32]
    qdot: NDArray[np.float32]
    grip: NDArray[Any]
    stage: NDArray[Any]
    t_idx: NDArray[np.int32]
    t_frac: NDArray[np.float32]
    t_raw: NDArray[np.int32]
    n_members: NDArray[np.int16]
    owner: NDArray[np.int32]
    task_id: NDArray[np.int16]
    phase: NDArray[np.float32]

    def __len__(self) -> int:
        return int(self.q.shape[0])


# ==================================================================================
# The table
# ==================================================================================
@dataclass(eq=False)
class NodeTable:
    """One row per coarse node, plus every raw frame.

    eq=False: numpy arrays don't support the dataclass-generated ==.
    """

    # ---- per-node fields (length V) ----
    q: NDArray[np.float32]
    qdot: NDArray[np.float32]
    grip: NDArray[np.uint8]
    stage: NDArray[np.int16]
    t_idx: NDArray[np.int32]
    t_frac: NDArray[np.float32]
    t_raw: NDArray[np.int32]
    n_members: NDArray[np.int16]
    owner: NDArray[np.int32]
    task_id: NDArray[np.int16]
    phase: NDArray[np.float32]

    # ---- CSR-style per-demo pointers (length n_demos + 1) ----
    demo_ptr: NDArray[np.int32]
    raw_ptr: NDArray[np.int32]

    # ---- every raw frame ----
    # ee_raw is [N_raw, 6] concat(position, unnormalized axis-angle) as RECORDED, i.e. world-frame;
    # the blend reads the base-frame ee_base property instead, and this stays as the ground truth
    # the FK behind it is checked against. a_raw is [N_raw, 7] commands.
    q_raw: NDArray[np.float32]
    qdot_raw: NDArray[np.float32]
    ee_raw: NDArray[np.float32]
    a_raw: NDArray[np.float32]

    # ---- metadata ----
    task_names: list[str]
    coarsen: int

    def __post_init__(self) -> None:
        self._ee_base: NDArray[np.float32] | None = None

    def __len__(self) -> int:
        return int(self.q.shape[0])

    @property
    def ee_base(self) -> NDArray[np.float32]:
        """[N_raw, 6] f32 concat(position, unnormalized axis-angle) in the ROBOT BASE frame.

        The frame every reference segment the blend follows is stated in. ee_raw is world-frame,
        and LIBERO parks the base somewhere different in each scene -- kitchen z 0.912, living room
        z 0.420, study x -0.750 -- so two world poses from different scenes differ by up to 0.67 m
        of pure scene offset. A retrieval that crosses a scene boundary then makes the tracking
        chunk's row 0 command ~46 action units toward another room, against a simulator that clips
        at 1. In the base frame the offset does not exist and the same retrieval degrades to the
        joint-space error it actually is.

        Derived, not stored: it is a pure function of q_raw, so an artifact column would be a second
        source of truth that goes stale the moment the URDF changes, and every g_nodes.npz on disk
        would need rebuilding to gain it. Computed by forward kinematics on first access and cached
        -- [138090, 6] f32 = 3.3 MB and 1.8 s on the long suite, paid once per loaded table.

        Returns:
            The pose table.

        Raises:
            ValueError: q_raw is not a 7-DoF Panda configuration. Assign a table instead when it is
                some other robot.
            ImportError: pytorch_kinematics is not installed in this interpreter.
        """
        if self._ee_base is None:
            from onf.blend.kinematics import PandaKinematics

            self.ee_base = PandaKinematics().pose(np.asarray(self.q_raw, dtype=np.float64))
        return self._ee_base

    @ee_base.setter
    def ee_base(self, pose: ArrayLike) -> None:
        """Supply the base-frame table directly, for a corpus the Panda FK above does not describe.

        Args:
            pose: [N_raw, 6] base-frame poses.

        Raises:
            ValueError: pose is not [N_raw, 6].
        """
        arr = np.asarray(pose, dtype=np.float32)
        if arr.shape != (self.n_raw, POSE_DIM):
            raise ValueError(
                f"ee_base must be [Nraw,{POSE_DIM}]=[{self.n_raw},{POSE_DIM}], got {arr.shape}"
            )
        self._ee_base = arr

    @property
    def dim(self) -> int:
        return int(self.q.shape[1])

    @property
    def n_demos(self) -> int:
        return int(self.demo_ptr.shape[0] - 1)

    @property
    def n_tasks(self) -> int:
        return len(self.task_names)

    @property
    def n_raw(self) -> int:
        return int(self.q_raw.shape[0])

    @classmethod
    def from_demos(
        cls,
        demos: Sequence[DemoRecord | Mapping[str, Any]],
        task_names: Sequence[str],
        coarsen: int = schema.COARSEN,
    ) -> NodeTable:
        """Builds a validated table from loaded demonstrations.

        Args:
            demos: Records, or legacy per-demo mappings.
            task_names: Task name per task_id.
            coarsen: Maximum raw frames per node.

        Returns:
            The built NodeTable.
        """
        return NodeTableBuilder(coarsen=coarsen).build(demos, task_names)

    def save(self, path: str | os.PathLike[str]) -> Path:
        """Writes g_nodes.npz.

        Args:
            path: Destination file or parent directory.

        Returns:
            The written path.
        """
        return NodeTableCodec().save(self, path)

    @classmethod
    def load(cls, path: str | os.PathLike[str], *, strict: bool = True) -> NodeTable:
        """Reads g_nodes.npz.

        Args:
            path: Source file or parent directory.
            strict: Verify the stamped phase_conv and run full validation.

        Returns:
            The loaded table.
        """
        return NodeTableCodec().load(path, strict=strict)

    def validate(self) -> None:
        """Enforces every structural invariant."""
        NodeTableValidator(self).run()

    def features(self) -> NDArray[np.float32]:
        """Generates raw, unnormalised network features.

        Returns:
            [V, 2D+1+2*PSI_FREQS] f32 array.
        """
        return NodeFeaturizer().encode(self)

    def feature_stats(self) -> FeatureStats:
        """Calculates statistics over the unbounded prefix.

        Returns:
            Mean and floored std.
        """
        return NodeFeaturizer().stats(self)

    def raw_span(self, node: int, k: int) -> tuple[int, int]:
        """Resolves the raw-frame slice a k-frame segment starting at a node may read.

        Clipped at the demo's raw end -- never reaches past raw_ptr[owner+1] into the
        next demo. The shortfall is what the callers edge-pad.

        Args:
            node: Node index.
            k: Number of frames requested.

        Returns:
            (start, end) global raw-frame offsets, with end - start <= k.

        Raises:
            ValueError: If k < 1 or t_raw[node] falls outside its demo's range.
        """
        if k < 1:
            raise ValueError(f"k must be >= 1, got {k}")
        owner = int(self.owner[node])
        start = int(self.t_raw[node])
        demo_start = int(self.raw_ptr[owner])
        demo_end = int(self.raw_ptr[owner + 1])
        if not demo_start <= start < demo_end:
            raise ValueError(
                f"t_raw[{node}]={start} outside demo {owner}'s raw range "
                f"[{demo_start}, {demo_end})"
            )
        return start, min(start + k, demo_end)

    def segment(self, node: int, k: int) -> Segment:
        """Decodes k raw frames starting at a node.

        Clipped at the demo's raw end and edge-padded up to exactly k frames --
        never reads past raw_ptr[owner+1] into the next demo.

        Args:
            node: Node index.
            k: Number of frames to return.

        Returns:
            The decoded Segment.

        Raises:
            ValueError: If k < 1 or t_raw[node] falls outside its demo's range.
        """
        start, end = self.raw_span(node, k)
        return Segment(
            q=_edge_pad(self.q_raw[start:end].astype(np.float64), k),
            qdot=_edge_pad(self.qdot_raw[start:end].astype(np.float64), k),
        )

    def ee_segment(self, node: int, k: int) -> NDArray[np.float64]:
        """Decodes k raw end-effector poses starting at a node, in the ROBOT BASE frame.

        Same clipping and edge-padding as segment, off the same raw_ptr boundaries.

        Args:
            node: Node index.
            k: Number of frames to return.

        Returns:
            [k, 6] f64 concat(position, unnormalized axis-angle orientation) off ee_base.

        Raises:
            ValueError: If k < 1 or t_raw[node] falls outside its demo's range.
        """
        start, end = self.raw_span(node, k)
        return _edge_pad(self.ee_base[start:end].astype(np.float64), k)

    def demo_nodes(self, owner: int) -> NDArray[np.int64]:
        """Retrieves node indices belonging to one demo.

        Args:
            owner: Demo id.

        Returns:
            Indices in t_idx order.
        """
        return np.arange(
            int(self.demo_ptr[owner]), int(self.demo_ptr[owner + 1]), dtype=np.int64
        )

    def node_at_raw(self, owner: int, raw_idx: int) -> int:
        """Finds the node on owner containing raw_idx.

        Args:
            owner: Demo id.
            raw_idx: Global raw-frame offset.

        Returns:
            Node index.
        """
        return int(self.node_at_raw_many(owner, raw_idx).reshape(()))

    def node_at_raw_many(
        self, owner: ArrayLike, raw_idx: ArrayLike
    ) -> NDArray[np.int64]:
        """Vectorised form of node_at_raw.

        t_raw is a GLOBAL raw offset, monotone over the WHOLE table, so ONE
        searchsorted answers every row at once; owner only clips the result
        back into its own node range. This is what makes
        onf.graph.run.kernel.TransitionKernel.build affordable at V in the
        tens of thousands.

        Args:
            owner: Demo ids, broadcastable against raw_idx.
            raw_idx: Global raw-frame offsets.

        Returns:
            Node indices with the broadcast shape.
        """
        owner_arr, raw_arr = np.broadcast_arrays(
            np.asarray(owner, dtype=np.int64), np.asarray(raw_idx, dtype=np.int64)
        )
        t_raw = self.t_raw.astype(np.int64, copy=False)
        demo_ptr = self.demo_ptr.astype(np.int64, copy=False)

        flat_owner = owner_arr.reshape(-1)
        pos = np.searchsorted(t_raw, raw_arr.reshape(-1), side="right") - 1
        lo = demo_ptr[flat_owner]
        hi = demo_ptr[flat_owner + 1] - 1
        return np.clip(pos, lo, hi).reshape(owner_arr.shape).astype(np.int64)

    def as_torch(self, device: str | None = None) -> dict[str, torch.Tensor]:
        """Returns every schema array as a torch tensor.

        f4 columns become float32; id columns become int64. Torch is imported
        lazily, so this module stays numpy-only until someone calls this.

        Args:
            device: Target device, or None for CPU.

        Returns:
            Mapping of field names to tensors.
        """
        torch_mod = _import_torch()
        out: dict[str, torch.Tensor] = {}
        for key, code in schema.NODES_KEYS.items():
            arr = getattr(self, key)
            cast = arr.astype(np.float32) if code == "f4" else arr.astype(np.int64)
            tensor = torch_mod.as_tensor(cast)
            out[key] = tensor.to(device) if device is not None else tensor
        return out

    # Historical name; EdgeSet.torch still spells it this way.
    torch = as_torch


def _import_torch() -> Any:
    """Lazily imports torch."""
    import torch as torch_mod
    return torch_mod


# ==================================================================================
# Builder
# ==================================================================================
class NodeTableBuilder:
    """Coarsens loaded demonstrations into a NodeTable."""

    def __init__(self, *, coarsen: int = schema.COARSEN, validate: bool = True) -> None:
        """Configures the builder.

        Args:
            coarsen: Maximum raw frames per node.
            validate: Run NodeTableValidator on the result.

        Raises:
            ValueError: If coarsen < 1.
        """
        if coarsen < 1:
            raise ValueError(f"coarsen must be >= 1, got {coarsen}")
        self._coarsen = int(coarsen)
        self._validate = validate

    def build(
        self,
        demos: Sequence[DemoRecord | Mapping[str, Any]],
        task_names: Sequence[str],
    ) -> NodeTable:
        """Builds a table.

        Args:
            demos: Records, or legacy per-demo mappings.
            task_names: Task name per task_id.

        Returns:
            The built NodeTable.

        Raises:
            ValueError: If demos is empty.
        """
        records = [
            d if isinstance(d, DemoRecord) else DemoRecord.from_mapping(d)
            for d in demos
        ]
        if not records:
            raise ValueError("demos must be non-empty")

        chunks: list[_DemoNodes] = []
        q_raw: list[NDArray[np.float32]] = []
        qdot_raw: list[NDArray[np.float32]] = []
        ee_raw: list[NDArray[np.float32]] = []
        a_raw: list[NDArray[np.float32]] = []
        demo_ptr: list[int] = [0]
        raw_ptr: list[int] = [0]
        raw_offset = 0

        for demo_id, record in enumerate(records):
            chunk = self._coarsen_demo(record, demo_id=demo_id, raw_offset=raw_offset)
            chunks.append(chunk)
            q_raw.append(record.q)
            qdot_raw.append(record.qdot)
            ee_raw.append(record.ee)
            a_raw.append(record.a)
            demo_ptr.append(demo_ptr[-1] + len(chunk))
            raw_ptr.append(raw_ptr[-1] + len(record))
            raw_offset += len(record)

        table = NodeTable(
            **self._concat_chunks(chunks),
            demo_ptr=np.asarray(demo_ptr, dtype=_DTYPE["i4"]),
            raw_ptr=np.asarray(raw_ptr, dtype=_DTYPE["i4"]),
            q_raw=_stack(q_raw, "f4"),
            qdot_raw=_stack(qdot_raw, "f4"),
            ee_raw=_stack(ee_raw, "f4"),
            a_raw=_stack(a_raw, "f4"),
            task_names=list(task_names),
            coarsen=self._coarsen,
        )
        if self._validate:
            table.validate()
        return table

    def _coarsen_demo(
        self, record: DemoRecord, *, demo_id: int, raw_offset: int
    ) -> _DemoNodes:
        """Coarsens one demo into per-node arrays.

        Args:
            record: The demonstration.
            demo_id: Global demo id.
            raw_offset: Running raw-frame offset.

        Returns:
            The demo's node chunk.
        """
        q = np.asarray(record.q, dtype=np.float32)
        n_frames = len(q)

        node_id = coarsen_segments(record.grip, self._coarsen)
        n_nodes = int(node_id[-1]) + 1
        bounds = np.searchsorted(node_id, np.arange(n_nodes + 1))
        starts = bounds[:-1]
        ends = bounds[1:] - 1

        phase = episode_phase(n_frames)
        return _DemoNodes(
            q=q[starts],
            # Coarse-time DISPLACEMENT, never the mean of per-step deltas.
            qdot=q[ends] - q[starts],
            grip=np.asarray(record.grip)[starts],
            stage=np.asarray(record.stage)[starts],
            t_idx=np.arange(n_nodes, dtype=np.int32),
            t_frac=(np.arange(n_nodes) / max(1, n_nodes - 1)).astype(np.float32),
            t_raw=starts.astype(np.int32) + raw_offset,
            n_members=(ends - starts + 1).astype(np.int16),
            owner=np.full(n_nodes, demo_id, dtype=np.int32),
            task_id=np.full(n_nodes, int(record.task_id), dtype=np.int16),
            phase=phase[starts],
        )

    @staticmethod
    def _concat_chunks(chunks: Sequence[_DemoNodes]) -> dict[str, NDArray[Any]]:
        """Concatenates per-demo chunks field-wise.

        Args:
            chunks: Per-demo node arrays.

        Returns:
            Mapping of field names to table-wide arrays.
        """
        codes = {name: schema.NODES_KEYS[name] for name in _PER_NODE_FIELDS}
        return {
            name: _stack([getattr(c, name) for c in chunks], codes[name])
            for name in (f.name for f in fields(_DemoNodes))
        }


def _stack(parts: Sequence[NDArray[Any]], code: str) -> NDArray[Any]:
    """Concatenates along axis 0 and casts to a schema dtype."""
    return np.concatenate(parts, axis=0).astype(_DTYPE[code])


def _edge_pad(block: NDArray[np.float64], k: int) -> NDArray[np.float64]:
    """Repeats a block's last row until it has exactly k rows.

    Args:
        block: [n, W] f64 rows, n >= 1.
        k: Target row count; n >= k returns block unchanged.

    Returns:
        [k, W] f64 rows.
    """
    pad = k - block.shape[0]
    if pad <= 0:
        return block
    return np.concatenate([block, np.repeat(block[-1:], pad, axis=0)], axis=0)


# ==================================================================================
# Validator
# ==================================================================================
class NodeTableValidator:
    """Enforces the structural invariants of a NodeTable."""

    def __init__(self, table: NodeTable) -> None:
        """Binds a table.

        Args:
            table: The table to check.
        """
        self._table = table
        self._n_nodes = len(table)
        self._n_demos = table.n_demos
        self._n_raw = table.n_raw

    def run(self) -> None:
        """Runs every default check in dependency order.

        Raises:
            TypeError: On a dtype mismatch.
            ValueError: On any shape, layout, or coarsening violation.
        """
        self.check_dtypes()
        self.check_shapes()
        self.check_finite()
        self.check_domains()
        if self._n_nodes == 0:
            return
        self.check_phase_convention()
        self.check_owner_layout()
        self.check_raw_layout()
        self.check_coarsening()

    def check_dtypes(self) -> None:
        """Verifies schema dtypes."""
        for key, code in schema.NODES_KEYS.items():
            dtype = getattr(self._table, key).dtype
            if dtype != _DTYPE[code]:
                raise TypeError(
                    f"{key}: dtype {dtype} != schema.NODES_KEYS[{key!r}]={_DTYPE[code]}"
                )

    def check_shapes(self) -> None:
        """Verifies array shapes."""
        table, n_nodes, dim = self._table, self._n_nodes, self._table.dim
        for key in _PER_NODE_FIELDS:
            length = getattr(table, key).shape[0]
            if length != n_nodes:
                raise ValueError(f"{key}: length {length} != V={n_nodes}")
        if table.q.shape != (n_nodes, dim) or table.qdot.shape != (n_nodes, dim):
            raise ValueError(f"q/qdot must be [V,D]=[{n_nodes},{dim}]")
        if table.q_raw.shape != (self._n_raw, dim) or table.qdot_raw.shape != (self._n_raw, dim):
            raise ValueError(f"q_raw/qdot_raw must be [Nraw,D]=[{self._n_raw},{dim}]")
        if table.ee_raw.shape != (self._n_raw, POSE_DIM):
            raise ValueError(f"ee_raw must be [Nraw,{POSE_DIM}]=[{self._n_raw},{POSE_DIM}]")
        if table.a_raw.shape != (self._n_raw, ACTION_DIM):
            raise ValueError(f"a_raw must be [Nraw,{ACTION_DIM}]=[{self._n_raw},{ACTION_DIM}]")
        expected_ptr = (self._n_demos + 1,)
        if table.demo_ptr.shape != expected_ptr or table.raw_ptr.shape != expected_ptr:
            raise ValueError("demo_ptr/raw_ptr must have length n_demos+1")

    def check_finite(self) -> None:
        """Verifies float fields contain no NaN or inf."""
        for key in _FINITE_FIELDS:
            if not np.all(np.isfinite(getattr(self._table, key))):
                raise ValueError(f"{key} contains NaN/inf")

    def check_domains(self) -> None:
        """Verifies categorical values."""
        if self._n_nodes == 0:
            return
        grip = self._table.grip
        if not np.all((grip == 0) | (grip == 1)):
            raise ValueError("grip must be 0/1")
        if not np.all(self._table.n_members >= 1):
            raise ValueError("n_members must be >= 1")

    def check_phase_convention(self) -> None:
        """Verifies global phase spans [0, 1].

        A violation is the signature of a FOREIGN convention: truncated at the grasp
        cut (q_flow.npz-style, max 0.436 on long) or divided by T instead of T-1
        (q_manifold_full.npz-style). Rebuild via episode_phase.
        """
        phase = self._table.phase
        p_max, p_min = float(phase.max()), float(phase.min())
        if not (p_max > _PHASE_MAX_MIN and p_min < _PHASE_MIN_MAX):
            raise ValueError(
                f"phase convention violated: table-wide max={p_max:.4f} min={p_min:.4f}, "
                f"expected max>{_PHASE_MAX_MIN} and min<{_PHASE_MIN_MAX}"
            )

    def check_phase_per_demo(self) -> None:
        """Verifies phase spans [0, 1] within EVERY demo.

        Not part of run -- an opt-in, stricter check.
        """
        table = self._table
        for demo_id in range(self._n_demos):
            lo, hi = int(table.demo_ptr[demo_id]), int(table.demo_ptr[demo_id + 1])
            if hi - lo < 2:
                continue
            span = table.phase[lo:hi]
            if not (float(span.max()) > _PHASE_MAX_MIN and float(span.min()) < _PHASE_MIN_MAX):
                raise ValueError(f"demo {demo_id}: phase span does not reach endpoints")

    def check_owner_layout(self) -> None:
        """Verifies owner contiguity."""
        table = self._table
        if not np.all(np.diff(table.owner) >= 0):
            raise ValueError("owner must be non-decreasing")
        node_counts = np.bincount(table.owner, minlength=self._n_demos)
        if (
            table.demo_ptr[0] != 0
            or table.demo_ptr[-1] != self._n_nodes
            or not np.array_equal(np.diff(table.demo_ptr), node_counts)
        ):
            raise ValueError("demo_ptr inconsistent with bincount(owner)")

        expected_t_idx = np.arange(self._n_nodes) - table.demo_ptr[table.owner]
        if not np.array_equal(table.t_idx.astype(np.int64), expected_t_idx.astype(np.int64)):
            raise ValueError("t_idx is not contiguous within demo")

    def check_raw_layout(self) -> None:
        """Verifies raw pointers.

        The last check is the worst failure mode: a t_raw outside its own demo
        makes NodeTable.segment silently decode a DIFFERENT demonstration.
        """
        table = self._table
        raw_counts = np.bincount(
            table.owner, weights=table.n_members.astype(np.int64), minlength=self._n_demos,
        ).astype(np.int64)
        if (
            table.raw_ptr[0] != 0
            or table.raw_ptr[-1] != self._n_raw
            or not np.array_equal(np.diff(table.raw_ptr).astype(np.int64), raw_counts)
        ):
            raise ValueError("raw_ptr inconsistent with n_members sum")

        if self._n_nodes > 1:
            same_owner = table.owner[:-1] == table.owner[1:]
            contiguous = table.t_raw[1:] == (table.t_raw[:-1] + table.n_members[:-1])
            if not np.all(~same_owner | contiguous):
                raise ValueError("t_raw is not contiguous within a demo")

        raw_start = table.raw_ptr[table.owner]
        raw_end = table.raw_ptr[table.owner + 1]
        if not (
            np.all(table.t_raw >= raw_start) and np.all(table.t_raw + table.n_members <= raw_end)
        ):
            raise ValueError("t_raw falls outside its own demo's raw range")

    def check_coarsening(self) -> None:
        """Verifies summaries agree with raw frames.

        Catches exactly the "coarse qdot == mean of per-step deltas" bug, which
        understates node motion by roughly coarsen x.
        """
        table = self._table
        last_idx = table.t_raw + table.n_members.astype(np.int32) - 1
        if not np.allclose(table.q, table.q_raw[table.t_raw], atol=_SUMMARY_ATOL):
            raise ValueError("q does not match q_raw[t_raw]")

        expected_qdot = table.q_raw[last_idx] - table.q_raw[table.t_raw]
        if not np.allclose(table.qdot, expected_qdot, atol=_SUMMARY_ATOL):
            raise ValueError("qdot must be the NET displacement across node's raw span")


# ==================================================================================
# Codec
# ==================================================================================
def _resolve_npz(path: str | os.PathLike[str]) -> Path:
    """Resolves a directory-or-file path to the g_nodes.npz file."""
    p = Path(path)
    return p if p.suffix == ".npz" else p / schema.NODES_NPZ


def _suite_hint(src: Path) -> str:
    """Guesses the suite an artifact belongs to, for a rebuild instruction.

    Args:
        src: Path to a g_nodes.npz, nominally under outputs/<suite>/<run>/artifacts/.

    Returns:
        The suite name, or a "<suite>" placeholder when the layout says nothing.
    """
    parts = src.parts
    if "outputs" in parts:
        index = parts.index("outputs") + 1
        if index < len(parts):
            return parts[index]
    return "<suite>"


class NodeTableCodec:
    """Reads and writes the NodeTable artifact."""

    def save(self, table: NodeTable, path: str | os.PathLike[str]) -> Path:
        """Writes a table.

        Args:
            table: Table to persist.
            path: Destination file or parent directory.

        Returns:
            The written path.
        """
        out = _resolve_npz(path)
        out.parent.mkdir(parents=True, exist_ok=True)
        payload: dict[str, Any] = {k: getattr(table, k) for k in schema.NODES_KEYS}
        payload["task_names"] = np.array(table.task_names)
        payload["coarsen"] = np.array(table.coarsen)
        payload["phase_conv"] = np.array(schema.PHASE_CONV)
        np.savez_compressed(os.fspath(out), **payload)
        return out

    def load(self, path: str | os.PathLike[str], *, strict: bool = True) -> NodeTable:
        """Reads a table.

        Args:
            path: Source file or parent directory.
            strict: Verify the stamped phase_conv and run full validation.

        Returns:
            The loaded table.

        Raises:
            ValueError: If the artifact predates a schema column, or if strict and its
                phase convention is foreign.
        """
        src = _resolve_npz(path)
        with np.load(os.fspath(src)) as archive:
            self._check_columns(archive.files, src)
            arrays = {
                key: np.asarray(archive[key]).astype(_DTYPE[code])
                for key, code in schema.NODES_KEYS.items()
            }
            task_names = [str(x) for x in archive["task_names"]]
            coarsen = int(archive["coarsen"])
            stamp = str(archive["phase_conv"]) if "phase_conv" in archive.files else ""

        table = NodeTable(**arrays, task_names=task_names, coarsen=coarsen)
        if strict:
            self._check_stamp(stamp, src)
            table.validate()
        return table

    @staticmethod
    def _check_columns(present: Sequence[str], src: Path) -> None:
        """Refuses an artifact written before one of the schema's columns existed.

        Zero-filling a missing ee_raw instead would leave the blend commanding a dead
        end-effector track, with nothing anywhere reporting an error.

        Args:
            present: Column names found in the archive.
            src: Source path, for the error message.

        Raises:
            ValueError: If any schema.NODES_KEYS column is absent.
        """
        missing = [key for key in schema.NODES_KEYS if key not in present]
        if missing:
            raise ValueError(
                f"{src}: g_nodes.npz has no {', '.join(missing)} column -- it predates the "
                f"current schema. Rebuild it with: "
                f"python -m onf.graph build --suite {_suite_hint(src)}"
            )

    @staticmethod
    def _check_stamp(stamp: str, src: Path) -> None:
        """Compares an artifact's phase-convention stamp against the schema.

        Args:
            stamp: The stamp read off the archive, or "" if absent.
            src: Source path, for the error message.

        Raises:
            ValueError: If the stamp is missing or foreign.
        """
        if stamp != schema.PHASE_CONV:
            raise ValueError(
                f"{src}: phase_conv={stamp or '<missing>'!r} != "
                f"schema.PHASE_CONV={schema.PHASE_CONV!r}. This file was not built by "
                "NodeTable.save -- it is likely a q_flow.npz- or q_manifold_full.npz-style "
                "legacy artifact. Rebuild via NodeTable.from_demos."
            )


# ==================================================================================
# Featurizer
# ==================================================================================
class NodeFeaturizer:
    """Builds network features from a NodeTable.

    Features are [q, qdot, grip, psi(t_frac)] ONLY. Never add task_id (it
    would make the input layer suite-shaped) or stage (diagnostic-only).
    """

    def __init__(self, n_freq: int = schema.PSI_FREQS) -> None:
        """Configures the featurizer.

        Args:
            n_freq: Octaves in the time encoding.
        """
        self._encoder = _time_encoder(n_freq)

    def width(self, dim: int) -> int:
        """Returns total feature width for a configuration dimension."""
        return 2 * dim + 1 + self._encoder.width

    def encode(self, table: NodeTable) -> NDArray[np.float32]:
        """Builds raw, unnormalised features.

        Args:
            table: Source table.

        Returns:
            [V, 2D+1+2*n_freq] f32 array.
        """
        return np.concatenate(
            [self._prefix(table), self._encoder.encode(table.t_frac)], axis=1
        ).astype(np.float32)

    def stats(self, table: NodeTable) -> FeatureStats:
        """Computes normalisation statistics over the unbounded prefix.

        The psi columns are already bounded in [-1, 1] and are excluded.

        Args:
            table: Source table.

        Returns:
            Mean and floored std.
        """
        prefix = self._prefix(table)
        std = prefix.std(axis=0)
        std = np.where(std < _STD_FLOOR, 1.0, std)
        return FeatureStats(
            mean=prefix.mean(axis=0).astype(np.float32), std=std.astype(np.float32)
        )

    @staticmethod
    def _prefix(table: NodeTable) -> NDArray[np.float32]:
        """Returns the [q, qdot, grip] block needing normalisation."""
        grip = table.grip.astype(np.float32)[:, None]
        return np.concatenate([table.q, table.qdot, grip], axis=1).astype(np.float32)