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

"""Contract for the demonstration-graph subsystem.

Configuration dataclasses, the relation vocabulary, and the on-disk artifact layout.
Design rationale lives in docs/technical/01-data-and-graph.md.

TWO VIEWS, ONE SOURCE OF TRUTH
    The dataclasses below are the definition: they carry the validation and the
    grouping. The flat COARSEN/NODES_KEYS/RELATIONS/... names at the
    bottom are DERIVED from DEFAULTS and exist because 424 call sites across
    44 modules read them. Never assign a literal to a flat name -- change the
    dataclass default and the flat name follows.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from types import MappingProxyType
from typing import ClassVar, Final, Mapping

import numpy as np
from numpy.typing import ArrayLike

__all__ = [
    "NpzSchema",
    "ArtifactSchemas",
    "RelationVocab",
    "GripperConfig",
    "GraphBuildConfig",
    "NetworkConfig",
    "RetrievalConfig",
    "ReadoutConfig",
    "TrackingConfig",
    "TrainingConfig",
    "GraphContract",
    "DEFAULTS",
    "grip_flag",
]


# =================================================================================================
# I/O schema
# =================================================================================================
@dataclass(frozen=True, slots=True)
class NpzSchema:
    """Filename and column dtypes for one on-disk npz artifact.

    Args:
        filename: Basename of the artifact within the graph directory.
        keys: Column name -> numpy dtype string.
        hash_stamped: Whether loaders must verify a graph_hash stamp and refuse on mismatch.
    """

    filename: str
    keys: Mapping[str, str] = field(default_factory=dict)
    hash_stamped: bool = False

    def __post_init__(self) -> None:
        if not self.filename:
            raise ValueError("NpzSchema.filename must be non-empty")
        for name, dtype in self.keys.items():
            try:
                np.dtype(dtype)
            except TypeError as exc:
                raise ValueError(f"{self.filename}: key {name!r} has bad dtype {dtype!r}") from exc
        object.__setattr__(self, "keys", MappingProxyType(dict(self.keys)))


class ArtifactSchemas:
    """Namespace of the on-disk npz artifact schemas. Not instantiated."""

    # q_raw/qdot_raw/raw_ptr keep every raw frame so segment length k stays a runtime choice.
    NODES: Final[NpzSchema] = NpzSchema(
        filename="g_nodes.npz",
        keys={
            "q": "f4", "qdot": "f4", "grip": "u1", "stage": "i2", "t_idx": "i4", "t_frac": "f4",
            "t_raw": "i4", "n_members": "i2", "owner": "i4", "task_id": "i2", "phase": "f4",
            "demo_ptr": "i4", "q_raw": "f4", "qdot_raw": "f4", "raw_ptr": "i4",
            "ee_raw": "f4", "a_raw": "f4",
        },
    )

    # Sorted by dst so indptr is a valid CSR offset array and index_add_ collisions coalesce.
    EDGES: Final[NpzSchema] = NpzSchema(
        filename="g_edges.npz",
        keys={"src": "i4", "dst": "i4", "rel": "i1", "w": "f4", "indptr": "i4", "log_idf": "f4"},
    )

    HEAD: Final[NpzSchema] = NpzSchema(filename="g_head.npz", keys={}, hash_stamped=True)

    # Hash stamp is mandatory: the advance operator is node-id-positional, so a foreign file
    # advances the belief along the wrong strands with no shape error and no exception.
    TRACK: Final[NpzSchema] = NpzSchema(
        filename="g_track.npz",
        keys={
            "pi": "f8", "beta": "f8", "leak": "f8", "delta_star": "f8",
            "s_ref_per_joint": "f8",
            "basin_r": "f8",           # [n_tasks, nbins_align] certified radius
            "basin_h": "f8",           # [n_tasks, nbins_align] KDE bandwidth
        },
        hash_stamped=True,
    )

    def __init__(self) -> None:
        raise TypeError("ArtifactSchemas is a namespace and must not be instantiated")


# =================================================================================================
# Relation vocabulary
# =================================================================================================
@dataclass(frozen=True, slots=True)
class RelationVocab:
    """Edge-relation alphabet. Order is load-bearing: rel indexes into relations.

    Args:
        dilations: Temporal hop lengths in coarse node steps.
        rel_dtype: Numpy dtype string of the rel edge column, used for a width check.
    """

    # Dilation is load-bearing, not an optimisation: depth L reaches only L hops, while a
    # pre-grasp approach is ~80 raw frames. Changing this invalidates every g_edges.npz.
    dilations: tuple[int, ...] = (1, 2, 4, 8, 16)
    rel_dtype: str = ArtifactSchemas.EDGES.keys["rel"]

    @property
    def next_relations(self) -> tuple[str, ...]:
        """Forward-along-demo relation names, dilated.

        Returns:
            Relation names in dilation order.
        """
        return tuple(f"next{d}" for d in self.dilations)

    @property
    def prev_relations(self) -> tuple[str, ...]:
        """Backward-along-demo relation names, dilated.

        Returns:
            Relation names in dilation order.
        """
        return tuple(f"prev{d}" for d in self.dilations)

    @property
    def relations(self) -> tuple[str, ...]:
        """Full relation alphabet in rel-index order.

        Returns:
            All relation names. sibling is kNN within a task; align crosses tasks at a
            shared phase bin. There is deliberately no stage relation.
        """
        return (*self.next_relations, *self.prev_relations, "sibling", "align")

    @property
    def rel_index(self) -> Mapping[str, int]:
        """Relation name to its rel column value.

        Returns:
            Read-only name -> index mapping.
        """
        return MappingProxyType({r: i for i, r in enumerate(self.relations)})

    @property
    def n_relations(self) -> int:
        """Size of the relation alphabet.

        Returns:
            Number of distinct relations.
        """
        return len(self.relations)

    def __post_init__(self) -> None:
        if not self.dilations:
            raise ValueError("RelationVocab.dilations must be non-empty")
        if any(d <= 0 for d in self.dilations):
            raise ValueError(f"dilations must be positive, got {self.dilations}")
        if len(set(self.dilations)) != len(self.dilations):
            raise ValueError(f"dilations contains duplicates: {self.dilations}")

        names = self.relations
        if len(set(names)) != len(names):
            raise ValueError(f"relation alphabet contains duplicates: {names}")

        rel_max = int(np.iinfo(np.dtype(self.rel_dtype)).max)
        if len(names) > rel_max + 1:
            raise ValueError(
                f"{len(names)} relations do not fit dtype {self.rel_dtype!r} (max {rel_max})"
            )


# =================================================================================================
# Configs
# =================================================================================================
@dataclass(frozen=True, slots=True)
class GripperConfig:
    """Gripper-state thresholding.

    Args:
        open_thresh: Mean absolute finger qpos at or below which the gripper counts as closed.
    """

    # Single source of truth; SentinelConfig.gripper_open_thresh must import this, not redeclare it.
    open_thresh: float = 0.035

    def __post_init__(self) -> None:
        if self.open_thresh <= 0.0:
            raise ValueError(f"open_thresh={self.open_thresh} must be positive")


@dataclass(frozen=True, slots=True)
class GraphBuildConfig:
    """Node-table and edge-construction parameters.

    Args:
        coarsen: Raw frames merged into one node; splits at every gripper-state change.
        psi_freqs: Sinusoidal time-encoding octaves; yields 2 * psi_freqs feature slots.
        k_sibling: kNN degree for the sibling relation.
        k_align: kNN degree for the align relation.
        nbins_align: Phase bins for the align bucket. Exported as NBINS_ALIGN, which
            every other reader of the binning (edges, loss, metrics, the sentinel's phase gate)
            must agree on -- see onf.graph.core.binning.phase_bins.
        cdist_bs: Row-block size for batched torch.cdist.
        kernel_bw: Joint-space kernel bandwidth.
        geo_scale: Radians; divisor in f = ||q - nearest node|| / geo_scale.
        phase_conv: Phase convention stamped into the npz and re-checked on load.
    """

    coarsen: int = 5
    psi_freqs: int = 4
    k_sibling: int = 8
    k_align: int = 4
    nbins_align: int = 20
    cdist_bs: int = 2048
    kernel_bw: float = 0.15
    # Calibrated so median f over leave-one-demo-out node distances matches the trained field's
    # median f on the same graph. Used only under ONF_CLEANLINESS=geo.
    geo_scale: float = 0.5303
    # Full-episode convention: phase[0] == 0 and phase[-1] == 1 exactly. Deliberately NOT the
    # truncated-at-grasp convention used by q_flow.npz -- align bins phase across tasks.
    phase_conv: str = "arange(T)/(T-1) full"

    def __post_init__(self) -> None:
        positive_ints = {
            "coarsen": self.coarsen, "psi_freqs": self.psi_freqs,
            "k_sibling": self.k_sibling, "k_align": self.k_align,
            "nbins_align": self.nbins_align, "cdist_bs": self.cdist_bs,
        }
        for name, value in positive_ints.items():
            if value <= 0:
                raise ValueError(f"{name}={value} must be a positive integer")
        for name, value in (("kernel_bw", self.kernel_bw), ("geo_scale", self.geo_scale)):
            if value <= 0.0:
                raise ValueError(f"{name}={value} must be positive")
        if not self.phase_conv:
            raise ValueError("phase_conv must be non-empty")


@dataclass(frozen=True, slots=True)
class NetworkConfig:
    """Message-passing network hyperparameters.

    Args:
        hidden: Hidden width.
        layers: Message-passing depth; reach is layers * max(dilations) * coarsen raw frames.
        agg: Neighbour aggregation, "sum" or "logsumexp".
        agg_temp: Temperature for logsumexp aggregation.
    """

    hidden: int = 64
    # 3 gives ~240 raw frames of reach, just under the long suite median demo length (p50 263).
    # Deeper dilutes: from a 256-node seed set, 6 layers reach 100% of a long graph's nodes.
    layers: int = 3
    agg: str = "sum"          # logsumexp with m = -cost recovers DTW's soft-min
    agg_temp: float = 1.0

    # ClassVar, not a field: a slots dataclass would otherwise put this in __init__.
    _VALID_AGG: ClassVar[frozenset[str]] = frozenset({"sum", "logsumexp"})

    def __post_init__(self) -> None:
        if self.hidden <= 0:
            raise ValueError(f"hidden={self.hidden} must be positive")
        if self.layers <= 0:
            raise ValueError(f"layers={self.layers} must be positive")
        if self.agg not in self._VALID_AGG:
            raise ValueError(f"agg={self.agg!r} must be one of {sorted(self._VALID_AGG)}")
        if self.agg_temp <= 0.0:
            raise ValueError(f"agg_temp={self.agg_temp} must be positive")


@dataclass(frozen=True, slots=True)
class RetrievalConfig:
    """Query construction and seeding.

    Args:
        hist_h: Query history window length W, one policy action chunk.
        seg_k: Reference-segment length; a segment restores velocity, a single point does not.
        seed_vel_w: Weight on the velocity-direction factor in the seeding kernel; 0 disables.
        seed_topk: Only the top-K nodes receive a nonzero initial hidden state.
        advance: Ground truth sits this many raw frames ahead; 0 labels the window's own end.
    """

    hist_h: int = 8
    seg_k: int = 8
    # Drift is nearly a pure translation of the query window, corrupting q but not qdot: at the
    # deployment-median 0.30 rad displacement, phase localisation is 4x better by velocity.
    seed_vel_w: float = 2.0
    seed_topk: int = 256
    # The head answers "where am I NOW"; the fitted transition kernel owns the lookahead. Any
    # advance > 0 is added to the kernel's own E[a] (~9.6 raw frames), overshooting the chunk.
    advance: int = 0

    def __post_init__(self) -> None:
        for name, value in (
            ("hist_h", self.hist_h), ("seg_k", self.seg_k), ("seed_topk", self.seed_topk),
        ):
            if value <= 0:
                raise ValueError(f"{name}={value} must be positive")
        if self.advance < 0:
            raise ValueError(f"advance={self.advance} must be non-negative")
        if self.seed_vel_w < 0.0:
            raise ValueError(f"seed_vel_w={self.seed_vel_w} must be non-negative")


@dataclass(frozen=True, slots=True)
class ReadoutConfig:
    """Readout arms and basin geometry.

    Args:
        arms: Available readout aggregations. All consume the same p(v|Q).
        basin_radius_quantile: Quantile of LODO NN distances defining the certified radius.
        basin_bandwidth_quantile: Quantile defining the KDE bandwidth / stop-short margin.
        basin_min_demos: Cells with fewer distinct demos inherit the per-task median r/h.
    """

    arms: tuple[str, ...] = ("euc_raw", "basin")
    basin_radius_quantile: float = 0.95
    basin_bandwidth_quantile: float = 0.50
    basin_min_demos: int = 2

    def __post_init__(self) -> None:
        if not self.arms:
            raise ValueError("ReadoutConfig.arms must be non-empty")
        if len(set(self.arms)) != len(self.arms):
            raise ValueError(f"arms contains duplicates: {self.arms}")
        for name, value in (
            ("basin_radius_quantile", self.basin_radius_quantile),
            ("basin_bandwidth_quantile", self.basin_bandwidth_quantile),
        ):
            if not 0.0 < value < 1.0:
                raise ValueError(f"{name}={value} must lie in (0, 1)")
        if self.basin_min_demos < 1:
            raise ValueError(f"basin_min_demos={self.basin_min_demos} must be at least 1")


@dataclass(frozen=True, slots=True)
class TrackingConfig:
    """Belief filter over the graph.

    Args:
        advance_set: Raw-frame advances mixed by the transition kernel; must contain 0.
        kernel_relations: Relations used by the dynamics kernel.
        belief_topk: Belief entries kept when handing a readout to host memory.
    """

    advance_set: tuple[int, ...] = (0, 1, 2, 4, 6, 8, 12, 16)
    # Deliberately excludes align and prev^d.
    kernel_relations: tuple[str, ...] = (
        "next1", "next2", "next4", "next8", "next16", "sibling",
    )
    belief_topk: int = 2048

    def __post_init__(self) -> None:
        if 0 not in self.advance_set:
            raise ValueError("advance_set must contain the mandatory a=0 stall/self-loop")
        if any(a < 0 for a in self.advance_set):
            raise ValueError(f"advance_set must be non-negative, got {self.advance_set}")
        if not self.kernel_relations:
            raise ValueError("kernel_relations must be non-empty")
        if self.belief_topk <= 0:
            raise ValueError(f"belief_topk={self.belief_topk} must be positive")


@dataclass(frozen=True, slots=True)
class TrainingConfig:
    """Objective weights for p(v|Q) = p(T|Q) * p(strand | T, Q).

    Phase is the primary signal; node identity is auxiliary and only weakly learnable.

    Args:
        phase_ce_w: Weight of the cross-entropy term on the phase-bin marginal.
        phase_expect_w: Weight of the phase-expectation regression term.
        node_w: Weight of the auxiliary node-identity BCE+rank term.
        phase_bin_smooth_w: Phase bins smoothed onto the phase-CE target, plus or minus.
        phase_bin_smooth_t: Smoothing decay constant in bins, as exp(-|db| / t).
        where_phase_band: Phase band defining the multi-positive candidate set.
        where_move_temp: Softmax temperature for inverse-movement-cost weighting.
        where_true_bonus: Multiplicative boost on the true continuation's weight.
    """

    phase_ce_w: float = 1.0
    phase_expect_w: float = 0.5
    node_w: float = 0.1
    phase_bin_smooth_w: float = 1.0
    phase_bin_smooth_t: float = 1.0
    where_phase_band: float = 0.05
    where_move_temp: float = 0.1
    where_true_bonus: float = 2.0

    def __post_init__(self) -> None:
        for name, value in (
            ("phase_ce_w", self.phase_ce_w),
            ("phase_expect_w", self.phase_expect_w),
            ("node_w", self.node_w),
            ("phase_bin_smooth_w", self.phase_bin_smooth_w),
        ):
            if value < 0.0:
                raise ValueError(f"{name}={value} must be non-negative")
        for name, value in (
            ("phase_bin_smooth_t", self.phase_bin_smooth_t),
            ("where_move_temp", self.where_move_temp),
            ("where_true_bonus", self.where_true_bonus),
        ):
            if value <= 0.0:
                raise ValueError(f"{name}={value} must be positive")
        if not 0.0 < self.where_phase_band < 1.0:
            raise ValueError(f"where_phase_band={self.where_phase_band} must lie in (0, 1)")


@dataclass(frozen=True, slots=True)
class GraphContract:
    """Composed configuration for the demonstration-graph subsystem.

    Args:
        vocab: Relation alphabet.
        gripper: Gripper thresholding.
        build: Node-table and edge construction.
        network: Message-passing hyperparameters.
        retrieval: Query construction and seeding.
        readout: Readout arms and basin geometry.
        tracking: Sequential test and belief filter.
        training: Objective weights.
    """

    vocab: RelationVocab = field(default_factory=RelationVocab)
    gripper: GripperConfig = field(default_factory=GripperConfig)
    build: GraphBuildConfig = field(default_factory=GraphBuildConfig)
    network: NetworkConfig = field(default_factory=NetworkConfig)
    retrieval: RetrievalConfig = field(default_factory=RetrievalConfig)
    readout: ReadoutConfig = field(default_factory=ReadoutConfig)
    tracking: TrackingConfig = field(default_factory=TrackingConfig)
    training: TrainingConfig = field(default_factory=TrainingConfig)

    @property
    def reach_raw_frames(self) -> int:
        """Raw-frame reach of the message-passing stack.

        Returns:
            layers * max(dilations) * coarsen.
        """
        return self.network.layers * max(self.vocab.dilations) * self.build.coarsen

    def __post_init__(self) -> None:
        unknown = set(self.tracking.kernel_relations) - set(self.vocab.relations)
        if unknown:
            raise ValueError(f"kernel_relations names unknown relations: {sorted(unknown)}")


# Convenience default. Prefer injecting a GraphContract explicitly in new code.
DEFAULTS: Final[GraphContract] = GraphContract()


# =================================================================================================
# Utilities
# =================================================================================================
def grip_flag(gripper_qpos: ArrayLike | None, thresh: float | None = None) -> float:
    """Map gripper joint positions to the graph's grip channel.

    Polarity is load-bearing: 1.0 means CLOSED, matching NodeTable.grip. Callers broadcast a
    single value across the whole window, matching how training builds it.

    Args:
        gripper_qpos: Finger joint positions, any array-like shape, or None if unavailable.
        thresh: Closed-gripper threshold. Defaults to DEFAULTS.gripper.open_thresh.

    Returns:
        1.0 if the gripper is closed (grasping), else 0.0. None yields 0.0.
    """
    if gripper_qpos is None:
        return 0.0
    limit = DEFAULTS.gripper.open_thresh if thresh is None else thresh
    mean_abs = float(np.abs(np.asarray(gripper_qpos, np.float64)).mean())
    return 1.0 if mean_abs <= limit else 0.0


# =================================================================================================
# Flat view -- DERIVED from DEFAULTS, never assigned a literal
# =================================================================================================
# 424 call sites across 44 modules read these names. They are aliases, not a second
# source of truth: every one of them reads through DEFAULTS or ArtifactSchemas, so a
# dataclass default and its flat name cannot drift apart.

# ---- relation vocabulary ----
DILATIONS: Final = DEFAULTS.vocab.dilations
RELATIONS: Final = DEFAULTS.vocab.relations
REL_INDEX: Final = DEFAULTS.vocab.rel_index
N_RELATIONS: Final = DEFAULTS.vocab.n_relations

# ---- node-table construction ----
COARSEN: Final = DEFAULTS.build.coarsen
PSI_FREQS: Final = DEFAULTS.build.psi_freqs
GRIP_OPEN_THR: Final = DEFAULTS.gripper.open_thresh
PHASE_CONV: Final = DEFAULTS.build.phase_conv

# ---- edge construction ----
K_SIBLING: Final = DEFAULTS.build.k_sibling
K_ALIGN: Final = DEFAULTS.build.k_align
NBINS_ALIGN: Final = DEFAULTS.build.nbins_align
CDIST_BS: Final = DEFAULTS.build.cdist_bs
KERNEL_BW: Final = DEFAULTS.build.kernel_bw
GEO_SCALE: Final = DEFAULTS.build.geo_scale

# ---- query / retrieval ----
HIST_H: Final = DEFAULTS.retrieval.hist_h
SEG_K: Final = DEFAULTS.retrieval.seg_k
SEED_VEL_W: Final = DEFAULTS.retrieval.seed_vel_w
SEED_TOPK: Final = DEFAULTS.retrieval.seed_topk
ADVANCE: Final = DEFAULTS.retrieval.advance

# ---- network ----
HIDDEN: Final = DEFAULTS.network.hidden
LAYERS: Final = DEFAULTS.network.layers
AGG: Final = DEFAULTS.network.agg
AGG_TEMP: Final = DEFAULTS.network.agg_temp

# ---- readout arms / basin geometry ----
READOUT_ARMS: Final = DEFAULTS.readout.arms
BASIN_RADIUS_QUANTILE: Final = DEFAULTS.readout.basin_radius_quantile
BASIN_BANDWIDTH_QUANTILE: Final = DEFAULTS.readout.basin_bandwidth_quantile
BASIN_MIN_DEMOS: Final = DEFAULTS.readout.basin_min_demos

# ---- WHERE: belief filter ----
TRACK_ADVANCE_SET: Final = DEFAULTS.tracking.advance_set
TRACK_KERNEL_RELATIONS: Final = DEFAULTS.tracking.kernel_relations
TRACK_BELIEF_TOPK: Final = DEFAULTS.tracking.belief_topk

# ---- objective weights ----
PHASE_CE_W: Final = DEFAULTS.training.phase_ce_w
PHASE_EXPECT_W: Final = DEFAULTS.training.phase_expect_w
NODE_W: Final = DEFAULTS.training.node_w
PHASE_BIN_SMOOTH_W: Final = DEFAULTS.training.phase_bin_smooth_w
PHASE_BIN_SMOOTH_T: Final = DEFAULTS.training.phase_bin_smooth_t
WHERE_PHASE_BAND: Final = DEFAULTS.training.where_phase_band
WHERE_MOVE_TEMP: Final = DEFAULTS.training.where_move_temp
WHERE_TRUE_BONUS: Final = DEFAULTS.training.where_true_bonus

# ---- artifact layout ----
NODES_NPZ: Final = ArtifactSchemas.NODES.filename
EDGES_NPZ: Final = ArtifactSchemas.EDGES.filename
HEAD_NPZ: Final = ArtifactSchemas.HEAD.filename
TRACK_NPZ: Final = ArtifactSchemas.TRACK.filename
NODES_KEYS: Final = ArtifactSchemas.NODES.keys
EDGES_KEYS: Final = ArtifactSchemas.EDGES.keys
TRACK_KEYS: Final = ArtifactSchemas.TRACK.keys

__all__ += [
    "DILATIONS", "RELATIONS", "REL_INDEX", "N_RELATIONS",
    "COARSEN", "PSI_FREQS", "GRIP_OPEN_THR", "PHASE_CONV",
    "K_SIBLING", "K_ALIGN", "NBINS_ALIGN", "CDIST_BS", "KERNEL_BW", "GEO_SCALE",
    "HIST_H", "SEG_K", "SEED_VEL_W", "SEED_TOPK", "ADVANCE",
    "HIDDEN", "LAYERS", "AGG", "AGG_TEMP",
    "READOUT_ARMS", "BASIN_RADIUS_QUANTILE", "BASIN_BANDWIDTH_QUANTILE", "BASIN_MIN_DEMOS",
    "TRACK_ADVANCE_SET", "TRACK_KERNEL_RELATIONS", "TRACK_BELIEF_TOPK",
    "PHASE_CE_W", "PHASE_EXPECT_W", "NODE_W", "PHASE_BIN_SMOOTH_W", "PHASE_BIN_SMOOTH_T",
    "WHERE_PHASE_BAND", "WHERE_MOVE_TEMP", "WHERE_TRUE_BONUS",
    "NODES_NPZ", "EDGES_NPZ", "HEAD_NPZ", "TRACK_NPZ",
    "NODES_KEYS", "EDGES_KEYS", "TRACK_KEYS",
]