onf-c2/scripts/capture_golden.py

"""capture_golden.py — build tests/golden/parity.json, the bit-identical regression net that
proves the two LIVE inference paths still produce IDENTICAL output after the upcoming deletion
campaign:

  1. t=0 ENTRY  -- onf.graph.run.retrieve.GraphRetriever, queried with a zero-velocity window
     (q0 repeated hist times).
  2. t>0 SENTINEL -- onf.graph.run.track.GraphTracker, driven step-by-step over a short
     synthetic joint trajectory sliced from one real demo's raw frames.

Everything below is DETERMINISTIC: one seeded numpy.random.RandomState (never the global
np.random state, never wall-clock/PID-derived anything) drives every draw, device="cpu"
throughout, and every float that goes into the golden file is round-tripped via float.hex() so
tests/test_parity.py can assert exact (==, not approx) equality against a fresh rerun of
this exact recipe.

Run (regenerates tests/golden/parity.json in place):

    PYTHONPATH=src .venv-core/bin/python scripts/capture_golden.py

Both halves read ONLY the real trained artifacts at outputs/long/latest/artifacts/ (never a
synthetic toy graph -- the toy fixtures in tests/test_retrieve.py exist for fast, artifact-free
unit tests of the *mechanics*; this script exists to pin the *actual* deployed numbers). If
g_track.npz is absent (as it is in this checkout), the sentinel half still runs -- GraphTracker.
load falls back to an unfit transition kernel -- but this is recorded in meta so a future re-run
against a fitted graph is not silently compared against a stale, unfit-kernel golden file.
"""

from __future__ import annotations

import hashlib
import json
from pathlib import Path

import numpy as np

from onf.config import GraphConfig
from onf.graph.core import schema
from onf.graph.run.retrieve import GraphRetriever
from onf.graph.run.track import GraphTracker

REPO_ROOT = Path(__file__).resolve().parents[1]
ARTIFACT_DIR = REPO_ROOT / "outputs" / "long" / "latest" / "artifacts"
GOLDEN_PATH = REPO_ROOT / "tests" / "golden" / "parity.json"

SEED = 20260805          # fixed, arbitrary -- never re-drawn, never wall-clock derived
N_ENTRY = 64
N_TRACK_STEPS = 64
TRACK_STRIDE = 2          # raw-frame stride between consecutive tracker checks

# t=0 ENTRY config -- the entry-band recipe pinned here as literal kwargs, not read from the
# environment, so this script has no dependency on which env vars a particular eval mode happens to set.
ENTRY_CFG_KWARGS = dict(entry_band=0.05, readout="euc_raw", move_cost=False, topm=32, hist=8, device="cpu")

# t>0 SENTINEL config -- GraphTracker._finalize always hand-rolls its own argmax reduction regardless
# of cfg.readout, so this is deliberately a separate config from the entry one above.
TRACK_CFG_KWARGS = dict(hist=8, device="cpu")


# The four artifact files the recorded numbers depend on. outputs/ is a symlink into a tree
# shared with sibling worktrees, so any of them can be rewritten by another branch; hashing them
# lets test_parity say "the artifact changed" instead of blaming the code.
ARTIFACT_FILES = ("g_nodes.npz", "g_edges.npz", "g_head.npz", "g_track.npz")


def artifact_sha256() -> dict:
    """{filename: sha256} for each present artifact in ARTIFACT_DIR; absent files are omitted."""
    out = {}
    for name in ARTIFACT_FILES:
        p = ARTIFACT_DIR / name
        if p.exists():
            out[name] = hashlib.sha256(p.read_bytes()).hexdigest()
    return out


def _hexf(x) -> str:
    """One float -> its exact hex round-trip representation."""
    return float(x).hex()


def _hex_list(arr) -> list:
    return [_hexf(v) for v in np.asarray(arr, dtype=np.float64).ravel().tolist()]


# ======================================================================================================
# 1. t=0 ENTRY capture
# ======================================================================================================
def capture_entry(retriever: GraphRetriever, rng: np.random.RandomState) -> list[dict]:
    """64 t=0 queries: fixed strided node indices, each node's own q perturbed by the seeded RNG,
    then wrapped into a zero-velocity window (q0 repeated cfg.hist times)."""
    nodes = retriever.nodes
    v = len(nodes)
    stride = max(v // N_ENTRY, 1)
    node_idx = [min(i * stride, v - 1) for i in range(N_ENTRY)]

    records = []
    for i in node_idx:
        q0 = nodes.q[i].astype(np.float64)
        q_pert = q0 + rng.normal(scale=0.02, size=q0.shape)
        q_hist = np.repeat(q_pert[None, :], retriever.cfg.hist, axis=0)   # zero velocity, by construction
        result = retriever.retrieve(q_hist)
        records.append({
            "node_idx": int(i),
            "q_star_hex": _hex_list(result.q_star),
            "depth_hex": _hexf(result.depth),
            "reached": int(result.info["reached"]),   # the surviving candidate pool, post every mask
        })
    return records


# ======================================================================================================
# 2. t>0 SENTINEL capture
# ======================================================================================================
def capture_track(tracker: GraphTracker, rng: np.random.RandomState) -> dict:
    """64 t>0 steps: a short synthetic trajectory built by walking consecutive raw q frames of one
    (seed-chosen) demo, fed to onf.graph.run.track.GraphTracker one check at a time -- this is
    onf.graph.run.track's own "clean" likelihood-sequence regime
    (onf.graph.build.tracker_fit.build_likelihood_sequences's kind="clean"), just driven through the
    live step() call instead of the batch fitting helper.

    Each step also records GraphTracker.where_target, the WHERE half of the sentinel: it is the
    only route into onf.graph.run.readout's basin arm and the only consumer of basin_r/
    basin_h, so without it that whole path -- and every argument of GraphTracker's constructor
    that feeds it -- ships with no bit-exact net. It draws no RNG and cannot perturb the step()
    records above it: TransitionKernel.push allocates its output and never mutates the belief.
    """
    nodes = tracker.retriever.nodes
    j = int(rng.randint(0, nodes.n_demos))
    demo_start, demo_end = int(nodes.raw_ptr[j]), int(nodes.raw_ptr[j + 1])
    traj = nodes.q_raw[demo_start:demo_end].astype(np.float64)

    hist = tracker.retriever.cfg.hist
    needed = hist + (N_TRACK_STEPS - 1) * TRACK_STRIDE
    if traj.shape[0] < needed:
        raise RuntimeError(
            f"capture_track: demo {j} has only {traj.shape[0]} raw frames, need >= {needed} "
            f"(hist={hist}, stride={TRACK_STRIDE}, steps={N_TRACK_STEPS}) -- pick a different seed/demo"
        )

    tracker.reset()
    records = []
    for s in range(N_TRACK_STEPS):
        start = s * TRACK_STRIDE
        q_hist = traj[start:start + hist]
        step = tracker.step(q_hist)
        b = step.belief_view
        where = tracker.where_target(q_hist[-1])            # WHERE half; q_now = the most recent frame
        records.append({
            "step": s,
            "belief_entropy_hex": _hexf(step.belief_entropy),
            "q_star_hex": _hex_list(b.q_star),
            "depth_hex": _hexf(b.depth),
            "phase_hat_hex": _hexf(step.phase_hat),
            # Both shipped arms are free-form (node is None); -1 is the value this field was captured
            # with when the golden file was frozen, so the reader maps None back to it.
            "where_node": -1 if where.node is None else int(where.node),
            "where_q_star_hex": _hex_list(where.q_star),
            "where_depth_hex": _hexf(where.depth),
            "where_conf_hex": _hexf(where.conf),
            "where_off_manifold_hex": _hexf(where.off_manifold),
        })
    return {"demo_owner": j, "records": records}


# ======================================================================================================
# entry point
# ======================================================================================================
def build_golden() -> dict:
    head_path = ARTIFACT_DIR / schema.HEAD_NPZ
    graph_hash_hex = None
    if head_path.exists():
        raw = np.load(str(head_path))
        if "graph_hash" in raw.files:
            graph_hash_hex = str(raw["graph_hash"])

    track_npz_present = (ARTIFACT_DIR / schema.TRACK_NPZ).exists()

    meta = {
        "artifact_dir": str(ARTIFACT_DIR.relative_to(REPO_ROOT)),
        "graph_hash": graph_hash_hex,
        "seed": SEED,
        "n_entry": N_ENTRY,
        "n_track_steps": N_TRACK_STEPS,
        "track_stride": TRACK_STRIDE,
        "entry_cfg": ENTRY_CFG_KWARGS,
        "track_cfg": TRACK_CFG_KWARGS,
        "track_npz_present": track_npz_present,
        "artifact_sha256": artifact_sha256(),
    }

    rng = np.random.RandomState(SEED)

    entry_cfg = GraphConfig(**ENTRY_CFG_KWARGS)
    entry_retriever = GraphRetriever.load(ARTIFACT_DIR, cfg=entry_cfg, device="cpu")
    entry_records = capture_entry(entry_retriever, rng)

    golden = {"meta": meta, "entry": entry_records}

    if track_npz_present is False:
        print(
            f"capture_golden: {ARTIFACT_DIR / schema.TRACK_NPZ} is ABSENT -- skipping the t>0 sentinel "
            "half (GraphTracker still loads with an unfit transition kernel, but there is no point "
            "pinning numbers against a kernel this checkout does not actually ship). meta.track_npz_"
            "present=False records this."
        )
        golden["track"] = None
    else:
        track_cfg = GraphConfig(**TRACK_CFG_KWARGS)
        tracker = GraphTracker.load(ARTIFACT_DIR, cfg=track_cfg, device="cpu")
        golden["track"] = capture_track(tracker, rng)

    return golden


def main() -> None:
    if not ARTIFACT_DIR.exists():
        raise SystemExit(f"capture_golden: {ARTIFACT_DIR} does not exist -- nothing to capture against")
    golden = build_golden()
    GOLDEN_PATH.parent.mkdir(parents=True, exist_ok=True)
    with open(GOLDEN_PATH, "w") as f:
        json.dump(golden, f, indent=2, sort_keys=True)
        f.write("\n")
    n_track = 0 if golden["track"] is None else len(golden["track"]["records"])
    print(f"wrote {GOLDEN_PATH} -- {len(golden['entry'])} entry records, {n_track} track records")


if __name__ == "__main__":
    main()