onf-c2/scripts/build_sentinel_artifacts.py

#!/usr/bin/env python
"""scripts/build_sentinel_artifacts.py — the demo-only builder for g_track.npz, the sentinel's
fitted belief-filter kernel.

Reads a built graph dir (g_nodes.npz/g_edges.npz/g_head.npz), NEVER a simulator and NEVER an
SR outcome, and produces:

    <graph_dir>/g_track.npz              -- onf.graph.run.track.TrackParams (pi/beta/leak +
                                             basin_r/basin_h), graph_hash-stamped like g_head.npz.
    <graph_dir>/sentinel_artifacts.json  -- the provenance record this run's numbers are cited from: every
                                             fitted scalar, basin geometry summary, graph_hash, head
                                             filename, and the exact split rule used.

R1 (this run) is DERIVED-ONLY: the retriever loads schema.HEAD_NPZ (the existing head, never
fine-tuned for the sentinel), so nothing downstream of this script's output has seen a single SR
outcome.

THE SPLIT RULE (declared here, not tuned)
    nodes.owner is a demo id in [0, n_demos). Every owner with owner % HELD_OUT_STRIDE == 0
    is HELD OUT (never touched by the kernel fit) and used ONLY for this script's own Step-2 sanity
    probes; every other owner is a FIT owner, used ONLY to fit pi/beta/leak. The two sets are disjoint
    by construction (a stride partition, not a random split -- no RNG, exactly reproducible); with
    owners assigned in contiguous per-task blocks (as this repo's graph builder does), the stride
    spreads the held-out demos evenly across every task rather than clumping them in one.
    FIT_CLEAN_STRIDE further subsamples the FIT owners' CLEAN sequences (still a stride, still
    deterministic) purely to bound wall-clock; the drift/cross_strand sequences use the full FIT pool as
    their sampling base (see onf.graph.build.tracker_fit.build_likelihood_sequences).

    NOTE (a real limitation, not fixable by this script alone): neither split is verified disjoint from
    whatever demos actually trained g_head.npz -- the graph-build pipeline in this repo does not
    currently record a train/held-out demo split for the GNN itself. The FIT/HELD-OUT disjointness this
    script enforces is therefore a guarantee about the KERNEL fit relative to the probe set, not about
    either relative to the head's own training set. Recorded in sentinel_artifacts.json
    as split.head_training_disjointness == "not verified (no recorded head train/held-out split)" so
    nobody mistakes this for a stronger claim than it is.

WHY A MIXED (clean + drift + cross_strand) KERNEL FIT, NOT held_out_owners= ALONE
    fit_transition_kernel's own docstring (2026-07-30 retirement note): fitting beta/leak
    against CLEAN-ONLY sequences answers "does sibling-mixing help when the arm never leaves its strand"
    (answer: no, beta=0) rather than the actual deployment question. This script therefore builds a
    MIXED fit set (clean over a strided FIT subsample + drift + cross_strand, all drawn only from FIT
    owners) and passes it via sequences=, per that docstring's own recommendation. If beta still
    comes out 0.0 on this mix, that is reported verbatim in both the console output and
    sentinel_artifacts.json -- it would directly contradict schema.py's "sibling mixing is MANDATORY"
    claim, and this script does not paper over that.

Usage:
    .venv-core/bin/python scripts/build_sentinel_artifacts.py outputs/long/frozen_learn1/artifacts \\
        --suite long --device cuda
"""

from __future__ import annotations

import argparse
import dataclasses
import json
import sys
import time
from pathlib import Path

import numpy as np

REPO_ROOT = Path(__file__).resolve().parents[1]
if str(REPO_ROOT / "src") not in sys.path:
    sys.path.insert(0, str(REPO_ROOT / "src"))

from onf.config import GraphConfig, SentinelConfig  # noqa: E402
from onf.graph.core import corpus, schema  # noqa: E402
from onf.graph.core.edges import EdgeSet  # noqa: E402
from onf.graph.core.nodes import NodeTable  # noqa: E402
from onf.graph.report import StageLogger  # noqa: E402
from onf.graph.run.retrieve import GraphRetriever  # noqa: E402
from onf.graph.run.track import GraphTracker, TrackParams, TransitionKernel  # noqa: E402
from onf.graph.run.params import TrackerCalibration  # noqa: E402
from onf.graph.build.tracker_fit import build_likelihood_sequences, fit_transition_kernel  # noqa: E402

# ---- declared, deterministic split (see module docstring) ------------------------------------------
HELD_OUT_STRIDE = 5     # owner % 5 == 0 -> held out (100/500 on long); everything else -> FIT
FIT_CLEAN_STRIDE = 3    # further stride the FIT owners for the CLEAN half of the mixed kernel fit
N_DRIFT_SEQ = 80        # drift-kind sequences drawn from the FIT pool
N_CROSS_STRAND_SEQ = 80  # cross_strand-kind sequences drawn from the FIT pool
DRIFT_SEQ_SEED = 0

# ---- Step-2 sanity probe knobs (debugging only) -------------------------------------------------
N_PROBE_DEMOS = 40          # held-out demos sampled (strided) for the stall probe
BURNIN_CHECKS_MIN = 3        # minimum burn-in checks before the probe/branch point
PROBE_PHASE_FRAC = 0.5       # probe at the raw-frame index closest to 50% through the demo -- mid-episode,
                              # where schema.py measures the largest clean-vs-drift intrinsic spread
N_STALL_CHECKS = 6           # checks held at the last frame (+ qdot-floor noise) after the branch point
STALL_QDOT_FLOOR_NOISE = 1e-4  # rad, tiny per-step jitter so a held window isn't bit-identical every check


def _quantiles(x: np.ndarray) -> dict:
    if x.size == 0:
        return {"n": 0, "p1": None, "p5": None, "p50": None, "p95": None, "p99": None,
                "min": None, "max": None, "mean": None, "std": None}
    return {
        "n": int(x.size),
        "p1": float(np.quantile(x, 0.01)), "p5": float(np.quantile(x, 0.05)),
        "p50": float(np.quantile(x, 0.50)), "p95": float(np.quantile(x, 0.95)),
        "p99": float(np.quantile(x, 0.99)),
        "min": float(x.min()), "max": float(x.max()),
        "mean": float(x.mean()), "std": float(x.std()),
    }


def _split_owners(n_demos: int, stride: int) -> tuple[np.ndarray, np.ndarray]:
    owners = np.arange(n_demos, dtype=np.int64)
    held_out = owners[owners % stride == 0]
    fit = owners[owners % stride != 0]
    return fit, held_out


def _build_mixed_fit_sequences(
    retriever, fit_owners: np.ndarray, *, check_every: int, log,
    fit_clean_stride: int, n_drift: int, n_cross_strand: int,
) -> list:
    clean_owners = fit_owners[::fit_clean_stride]
    seqs = build_likelihood_sequences(retriever, clean_owners, kind="clean", check_every=check_every)
    log.step("fit-seq clean", f"{len(seqs)} sequences from {len(clean_owners)} fit owners "
              f"(stride {fit_clean_stride})", n_clean=len(seqs))
    d = build_likelihood_sequences(
        retriever, fit_owners, kind="drift", n_sequences=n_drift, check_every=check_every,
        seed=DRIFT_SEQ_SEED,
    )
    seqs.extend(d)
    log.step("fit-seq drift", f"{len(d)} sequences", n_drift=len(d))
    x = build_likelihood_sequences(
        retriever, fit_owners, kind="cross_strand", n_sequences=n_cross_strand, check_every=check_every,
        seed=DRIFT_SEQ_SEED,
    )
    seqs.extend(x)
    log.step("fit-seq cross_strand", f"{len(x)} sequences", n_cross_strand=len(x))
    return seqs


# ======================================================================================================
# Step 2 -- bank sanity probes (debugging only)
# ======================================================================================================
def _probe_owners(held_out: np.ndarray, n: int) -> np.ndarray:
    if held_out.size <= n:
        return held_out
    stride = max(1, held_out.size // n)
    return held_out[::stride][:n]


def _burnin_steps(traj_raw: np.ndarray, t_probe_raw: int, check_every: int) -> list[int]:
    steps = list(range(0, t_probe_raw, check_every))
    if not steps or steps[-1] != t_probe_raw:
        steps.append(t_probe_raw)
    return steps


def _run_stall_branch(tracker, traj_raw: np.ndarray, steps: list[int], hist: int, *,
                       stall_checks: int, rng: np.random.RandomState):
    """Replay steps as burn-in, then hold the last raw frame fixed (+ tiny qdot-floor noise) for
    stall_checks more checks and return those TrackerSteps -- the stall-mass (w0) probe."""
    tracker.reset()
    n_raw = traj_raw.shape[0]
    for t in steps:
        hi = min(t + 1, n_raw)
        lo = max(0, hi - hist)
        tracker.step(traj_raw[lo:hi].astype(np.float64).copy())
    held_frame = traj_raw[min(steps[-1], n_raw - 1)].astype(np.float64)
    out = []
    for _ in range(stall_checks):
        noise = rng.randn(hist, held_frame.shape[0]) * STALL_QDOT_FLOOR_NOISE
        out.append(tracker.step(np.tile(held_frame, (hist, 1)) + noise))
    return out


def run_bank_sanity(retriever, held_out_owners: np.ndarray, tracker_factory, *, check_every: int,
                     hist: int, log) -> dict:
    nodes = retriever.nodes
    probe_owners = _probe_owners(held_out_owners, N_PROBE_DEMOS)

    w0_clean_vals: list[float] = []
    w0_stall_vals: list[float] = []

    for j in probe_owners:
        j = int(j)
        demo_start, demo_end = int(nodes.raw_ptr[j]), int(nodes.raw_ptr[j + 1])
        traj = nodes.q_raw[demo_start:demo_end]
        n_raw = traj.shape[0]
        if n_raw < hist + check_every * BURNIN_CHECKS_MIN:
            continue
        t_probe = int(min(n_raw - 1, max(hist, round(PROBE_PHASE_FRAC * n_raw))))
        steps = _burnin_steps(traj, t_probe, check_every)
        if len(steps) < BURNIN_CHECKS_MIN:
            continue
        rng = np.random.RandomState(seed=j)

        # ---- stall mass: w0 continuing clean vs w0 while stalled ----
        clean_cont_steps = list(steps)
        n_stall = 0
        while len(clean_cont_steps) < len(steps) + N_STALL_CHECKS and clean_cont_steps[-1] + check_every < n_raw:
            clean_cont_steps.append(clean_cont_steps[-1] + check_every)
            n_stall += 1
        if n_stall > 0:
            trk = tracker_factory()
            trk.reset()
            w0s = []
            for i, t in enumerate(clean_cont_steps):
                hi = min(t + 1, n_raw)
                lo = max(0, hi - hist)
                out = trk.step(traj[lo:hi].astype(np.float64))
                if i >= len(steps) and not np.isnan(out.w0):
                    w0s.append(out.w0)
            w0_clean_vals.extend(w0s)

            stall_out = _run_stall_branch(tracker_factory(), traj, steps, hist,
                                           stall_checks=n_stall, rng=rng)
            w0_stall_vals.extend(o.w0 for o in stall_out if not np.isnan(o.w0))

    log.step("sanity probes", f"{len(probe_owners)} held-out demos probed", n_probe_demos=len(probe_owners))

    result = {
        "n_probe_demos": int(len(probe_owners)),
        "w0_clean": _quantiles(np.asarray(w0_clean_vals, dtype=np.float64)),
        "w0_stall": _quantiles(np.asarray(w0_stall_vals, dtype=np.float64)),
    }
    return result


# ======================================================================================================
# main
# ======================================================================================================
def main(argv: list[str] | None = None) -> int:
    p = argparse.ArgumentParser(
        description="Fit onf.graph.run.track's kernel + basin geometry from a built graph's "
        "DEMONSTRATION corpus alone (no simulator, no SR outcome) and write g_track.npz + "
        "sentinel_artifacts.json into graph_dir.",
    )
    p.add_argument("graph_dir", help="dir holding g_nodes.npz/g_edges.npz/g_head.npz (e.g. "
                    "outputs/long/frozen_learn1/artifacts)")
    p.add_argument("--suite", default=None, help="suite label for outputs/<suite>/... progress "
                    "bookkeeping only (default: inferred from graph_dir's grandparent dir name)")
    p.add_argument("--device", default=None, help="torch device (default: auto)")
    p.add_argument("--head-npz", default=schema.HEAD_NPZ, help=f"head artifact to load the retriever "
                    f"from (default: schema.HEAD_NPZ={schema.HEAD_NPZ!r} -- R1 is derived-only, the "
                    f"existing head, never a sentinel-fine-tuned one)")
    p.add_argument("--held-out-stride", type=int, default=HELD_OUT_STRIDE)
    p.add_argument("--fit-clean-stride", type=int, default=FIT_CLEAN_STRIDE)
    p.add_argument("--n-drift", type=int, default=N_DRIFT_SEQ)
    p.add_argument("--n-cross-strand", type=int, default=N_CROSS_STRAND_SEQ)
    p.add_argument("--n-sweeps", type=int, default=3, help="fit_transition_kernel's coordinate-ascent "
                    "sweep count (default: that function's own default)")
    p.add_argument("--skip-sanity", action="store_true", help="skip the Step-2 bank sanity probes")
    p.add_argument("--out", default=None, help="StageLogger bookkeeping root (default: outputs/)")
    p.add_argument("--quiet", action="store_true")
    args = p.parse_args(argv)

    held_out_stride = int(args.held_out_stride)
    fit_clean_stride = int(args.fit_clean_stride)
    n_drift = int(args.n_drift)
    n_cross_strand = int(args.n_cross_strand)

    graph_dir = Path(args.graph_dir)
    suite = args.suite or (graph_dir.resolve().parents[1].name if len(graph_dir.resolve().parents) > 1 else "unknown")

    config = {
        "command": "build_sentinel_artifacts", "graph_dir": str(graph_dir), "suite": suite,
        "device": args.device, "head_npz": args.head_npz, "held_out_stride": held_out_stride,
        "fit_clean_stride": fit_clean_stride, "n_drift": n_drift, "n_cross_strand": n_cross_strand,
        "n_sweeps": args.n_sweeps, "skip_sanity": bool(args.skip_sanity),
    }
    print("===== build_sentinel_artifacts: resolved config =====")
    print(json.dumps(config, indent=2))

    check_every = SentinelConfig().check_every
    hist = schema.HIST_H

    with StageLogger(
        "sentinel_fit", suite, root=args.out, config=config, n_stages=8, quiet=args.quiet,
        update_latest=False,
    ) as log:
        t_start = time.time()

        log.step("load graph", str(graph_dir))
        nodes = NodeTable.load(graph_dir)
        edges = EdgeSet.load(graph_dir / schema.EDGES_NPZ)
        edges.validate(nodes)
        gcfg = GraphConfig(hist=hist, device=(args.device or ""))
        retriever = GraphRetriever.load(graph_dir, cfg=gcfg, device=(args.device or None), head_npz=args.head_npz)
        g_hash = retriever.graph_hash          # graph_hash v2: geometry + constants + this head's C2
        log.log(f"[sentinel_fit] V={len(nodes):,} E={len(edges.src):,} n_demos={nodes.n_demos} "
                f"n_tasks={nodes.n_tasks} graph_hash={g_hash[:16]}... device={retriever.device}")

        log.step("basin geometry", "corpus.measure_basin_geometry")
        basin = corpus.measure_basin_geometry(nodes)
        log.log(f"[sentinel_fit] basin_r: min={basin['basin_r'].min():.4f} "
                f"median={np.median(basin['basin_r']):.4f} max={basin['basin_r'].max():.4f}  "
                f"basin_h: min={basin['basin_h'].min():.4f} median={np.median(basin['basin_h']):.4f} "
                f"max={basin['basin_h'].max():.4f}  n_cells_backfilled={basin['n_cells_backfilled']} "
                f"/ {basin['n_tasks'] * basin['nbins']}  spacing={basin['spacing']:.4f}")

        log.step("split owners", f"stride={held_out_stride}")
        fit_owners, held_out_owners = _split_owners(nodes.n_demos, held_out_stride)
        assert set(fit_owners.tolist()).isdisjoint(set(held_out_owners.tolist()))
        log.log(f"[sentinel_fit] fit_owners={len(fit_owners)}  held_out_owners={len(held_out_owners)}")

        log.step("build fit sequences", "clean(strided) + drift + cross_strand, from FIT owners only")
        fit_sequences = _build_mixed_fit_sequences(
            retriever, fit_owners, check_every=check_every, log=log, fit_clean_stride=fit_clean_stride,
            n_drift=n_drift, n_cross_strand=n_cross_strand,
        )
        log.log(f"[sentinel_fit] total fit sequences={len(fit_sequences)}")

        log.step("fit_transition_kernel", f"{len(fit_sequences)} sequences, n_sweeps={args.n_sweeps}")
        t0 = time.time()
        fit = fit_transition_kernel(
            retriever, sequences=fit_sequences, check_every=check_every, n_sweeps=args.n_sweeps,
        )
        fit_elapsed = time.time() - t0
        beta_is_zero = bool(fit.beta == 0.0)
        log.log(f"[sentinel_fit] FITTED pi={np.round(fit.pi, 4).tolist()} beta={fit.beta:.4f} "
                f"leak={fit.leak:.4f}  log_evidence={fit.log_evidence:.2f} "
                f"(uniform_pi baseline={fit.log_evidence_uniform_pi:.2f}, "
                f"delta={fit.log_evidence - fit.log_evidence_uniform_pi:+.2f})  "
                f"n_checks={fit.n_checks}  fit_elapsed={fit_elapsed:.1f}s")
        if beta_is_zero:
            log.log("[sentinel_fit] *** WARNING: fitted beta == 0.0 -- sibling mixing did NOT help on "
                     "this mixed (clean+drift+cross_strand) fit set, which CONTRADICTS schema.py's "
                     "'sibling mixing is MANDATORY' claim. Reported verbatim, not papered over. ***")

        kernel = TransitionKernel.build(nodes, edges, pi=fit.pi, beta=fit.beta, device=retriever.device)

        params = TrackParams(
            pi=fit.pi, beta=float(fit.beta), leak=float(fit.leak),
            basin_r=basin["basin_r"], basin_h=basin["basin_h"],
        )

        log.step("save g_track.npz", str(graph_dir))
        track_path = params.save(graph_dir, nodes, edges)
        log.log(f"[sentinel_fit] wrote {track_path}")

        sanity = None
        if not args.skip_sanity:
            log.step("Step-2 sanity probes", f"n_probe_demos<= {N_PROBE_DEMOS}")

            def tracker_factory():
                return GraphTracker(retriever, kernel, calib=TrackerCalibration(
                    leak=params.leak, basin_r=params.basin_r, basin_h=params.basin_h,
                ))

            t0 = time.time()
            sanity = run_bank_sanity(
                retriever, held_out_owners, tracker_factory, check_every=check_every, hist=hist, log=log,
            )
            sanity_elapsed = time.time() - t0
            log.log(f"[sentinel_fit] sanity probes elapsed={sanity_elapsed:.1f}s")
            wc, ws = sanity["w0_clean"], sanity["w0_stall"]
            log.log(f"[sentinel_fit] w0 median: clean={wc['p50']}  stall={ws['p50']}")
        else:
            log.step("Step-2 sanity probes", "SKIPPED (--skip-sanity)")

        # ---- provenance record --------------------------------------------------------------------
        artifact_record = {
            "graph_dir": str(graph_dir), "graph_hash": g_hash, "head_npz": args.head_npz,
            "n_nodes": len(nodes), "n_edges": int(len(edges.src)), "n_demos": int(nodes.n_demos),
            "n_tasks": int(nodes.n_tasks), "device": str(retriever.device),
            "split": {
                "rule": f"owner % {held_out_stride} == 0 -> held_out, else -> fit (deterministic stride "
                        "partition of nodes.owner, no RNG)",
                "held_out_stride": held_out_stride, "fit_clean_stride": fit_clean_stride,
                "n_fit_owners": int(len(fit_owners)), "n_held_out_owners": int(len(held_out_owners)),
                "disjoint_fit_vs_held_out": True,
                "head_training_disjointness": "not verified (no recorded head train/held-out split in "
                                               "this repo's graph-build pipeline)",
            },
            "fit_sequences": {
                "n_total": len(fit_sequences), "n_clean": len(fit_sequences) - n_drift - n_cross_strand,
                "n_drift": n_drift, "n_cross_strand": n_cross_strand, "seed": DRIFT_SEQ_SEED,
            },
            "fit_transition_kernel": {
                "pi": fit.pi.tolist(), "advances": list(schema.TRACK_ADVANCE_SET), "beta": float(fit.beta),
                "beta_is_zero": beta_is_zero, "leak": float(fit.leak),
                "log_evidence": float(fit.log_evidence),
                "log_evidence_uniform_pi_baseline": float(fit.log_evidence_uniform_pi),
                "log_evidence_delta_vs_uniform_pi": float(fit.log_evidence - fit.log_evidence_uniform_pi),
                "n_checks": int(fit.n_checks), "n_sweeps": int(args.n_sweeps), "elapsed_s": fit_elapsed,
            },
            "basin_geometry": {
                "radius_quantile": basin["radius_quantile"], "bandwidth_quantile": basin["bandwidth_quantile"],
                "min_demos": basin["min_demos"], "spacing": basin["spacing"],
                "n_tasks": basin["n_tasks"], "nbins": basin["nbins"],
                "n_cells_backfilled": basin["n_cells_backfilled"],
                "n_cells_total": basin["n_tasks"] * basin["nbins"],
                "basin_r": {"min": float(basin["basin_r"].min()), "median": float(np.median(basin["basin_r"])),
                             "max": float(basin["basin_r"].max())},
                "basin_h": {"min": float(basin["basin_h"].min()), "median": float(np.median(basin["basin_h"])),
                             "max": float(basin["basin_h"].max())},
            },
            "g_track_npz_path": track_path,
            "step2_bank_sanity_probes": sanity,
            "elapsed_s_total": time.time() - t_start,
        }
        record_path = graph_dir / "sentinel_artifacts.json"
        with open(record_path, "w") as fh:
            json.dump(artifact_record, fh, indent=2)
        log.log(f"[sentinel_fit] wrote {record_path}")
        log.metric(**{k: v for k, v in artifact_record.items() if k not in ("step2_bank_sanity_probes",)})

    print(json.dumps(artifact_record, indent=2))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())