Buckets:

AsianPlayer
/

TriDrive

Files

xet

AsianPlayer/TriDrive / code /aide /tests /test_kv2.py

AsianPlayer

16 days ago

download

raw

15.1 kB

"""Guards for the K-v2 programme protocol side (docs/decisions.md 2026-09-14 "K-v2 programme — S0 pre-registration", protocol C23):
* `kv2_gate` S1 rule on synthetic per-clip arrays: survive / kill / HM guard / tie-with-fewer-components / ≤ 2 survivors;
* S1b continuation rule; S3 argmin + 0.10-px tie -> fewer parameters;
* MotionBERT-kp name-encoded variants build with the expected parameter counts, TRUNK_PREFIXES unchanged, foreign suffixes rejected;
* run-name / flag helpers shared by the launcher, the gate and the report;
* `dry_trunk_load` mechanics (REUSE on a matching checkpoint, NEED on a foreign one);
* `report_kv2_program` smoke on a tmp out-root with partial results (pending runs, conservative NOT ADOPTED line, append-only).
    python -m pytest aide/tests/test_kv2.py -q
"""
from __future__ import annotations

import os

os.environ.setdefault("OMP_NUM_THREADS", "4")

import glob  # noqa: E402
import json  # noqa: E402
import shutil  # noqa: E402
import sys  # noqa: E402
from pathlib import Path  # noqa: E402

import numpy as np  # noqa: E402
import pytest  # noqa: E402
import torch  # noqa: E402

torch.set_num_threads(4)
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
from aide.data import protocol as P  # noqa: E402
from aide.eval import kv2_gate as G  # noqa: E402
from aide.models import motionbert_kp  # noqa: E402

N_BOOT = 400


def _synthetic(seed=0, n=405, n_hm=40):
    rng = np.random.default_rng(seed)
    base = np.abs(rng.normal(55, 25, n))
    hm = np.zeros(n, bool)
    hm[rng.choice(n, n_hm, replace=False)] = True
    return rng, base, hm


def _cand(rng, base, shift, noise=0.05):
    return base + shift + rng.normal(0, noise, len(base))


def test_s1_rule_survive_kill_guard_tie_and_cap():
    rng, base, hm = _synthetic()
    cands = {                                  # cfg: (shift on all clips, extra shift on HM clips, components)
        "full": (-0.50, 0.0, 5),               # survives, mean −0.50
        "nohaz": (-0.45, 0.0, 3),              # within 0.10 of full with fewer components -> ranked FIRST
        "plain": (-0.31, 0.0, 4),              # survives but only 2 slots
        "khead": (-0.10, 0.0, 4),              # KILL (> −0.30)
        "q060": (-0.30, 0.0, 5),               # exactly at the threshold -> survives (≤)
        "hazaux": (-0.60, +6.0, 4),            # best All but HM guard: Δ HM > +3 with CI excl. 0 -> GUARD_KILL
    }
    rows = []
    for cfg, (sh, sh_hm, nc) in cands.items():
        c = _cand(rng, base, sh)
        c_hm = c[hm] + sh_hm
        rows.append(G.s1_row(cfg, G.kv2_run(cfg, 42), c, base, c_hm, base[hm], nc, N_BOOT, 0))
    st = {r["cfg"]: r["status"] for r in rows}
    assert st == {"full": "SURVIVE", "nohaz": "SURVIVE", "plain": "SURVIVE", "khead": "KILL", "q060": "SURVIVE", "hazaux": "GUARD_KILL"}
    r_haz = next(r for r in rows if r["cfg"] == "hazaux")
    assert r_haz["survive_all"] and r_haz["guard_kill"] and r_haz["d_hm"]["ci_lo"] > 0 and r_haz["d_hm"]["mean"] > 3.0
    assert abs(next(r for r in rows if r["cfg"] == "full")["d_all"]["mean"] + 0.50) < 0.02
    surv = G.s1_select(rows)
    assert surv == ["nohaz", "full"]           # tie rule (fewer components first), then the cap of 2 (plain / q060 excluded, hazaux guard-killed)
    # without the tie: a clear leader wins regardless of its component count
    rows2 = [r for r in rows if r["cfg"] != "nohaz"]
    assert G.s1_select(rows2) == ["full", "plain"]
    # a guard flagged but the All rule failed -> plain KILL, never GUARD_KILL
    c = _cand(rng, base, -0.05)
    r = G.s1_row("lam03", "x", c, base, c[hm] + 6.0, base[hm], 5, N_BOOT, 0)
    assert r["status"] == "KILL" and r["guard_kill"]
    assert G.s1_select([]) == []


def test_s1b_rule():
    ok, why = G.s1b_decide(-0.40, -0.15)
    assert ok and why.startswith("CONTINUE")           # Δ123 < 0 and mean −0.275 ≤ −0.25
    assert not G.s1b_decide(-0.60, +0.05)[0]           # Δ123 ≥ 0 -> killed even though the mean passes
    assert not G.s1b_decide(-0.30, -0.10)[0]           # mean −0.20 > −0.25 -> killed
    ok, why = G.s1b_decide(-0.25, -0.25)
    assert ok and "val-killed" not in why
    assert "val-killed at seed 123" in G.s1b_decide(-0.30, -0.10)[1]


def test_s3_argmin_and_tie_to_fewer_params():
    rows = [{"cfg": "default", "val_all": 55.55, "n_params": 813_088},
            {"cfg": "d192", "val_all": 55.50, "n_params": 1_809_440},      # within 0.10 of the best -> fewer params wins
            {"cfg": "L4", "val_all": 55.70, "n_params": 1_077_664}]
    assert G.s3_select(rows) == "default"
    rows.append({"cfg": "do0.1", "val_all": 55.30, "n_params": 813_088})   # clear leader (> 0.10 better than default)
    assert G.s3_select(rows) == "do0.1"
    rows.append({"cfg": "wd0.1", "val_all": 55.25, "n_params": 813_088})   # ties do0.1 (same params) -> lower val
    assert G.s3_select(rows) == "wd0.1"
    assert G.s3_select([{"cfg": "x", "val_all": None, "n_params": 1}]) is None
    assert G.s3_select([]) is None


def test_mb_name_variants_build_with_expected_params_and_same_trunk_prefixes():
    expected = {"motionbert_kp": 813_088, "motionbert_kp_2b": 548_512, "motionbert_kp_d192": 1_809_440,
                "motionbert_kp_L4": 1_077_664, "motionbert_kp_d192_L4": 2_402_912}
    ref = motionbert_kp.MotionBERTKP.TRUNK_PREFIXES
    assert ref == ("inp.", "joint_emb", "time_emb", "blocks.", "ln_f.")
    for name, n in expected.items():
        m = motionbert_kp.build(name, dropout=0.0)
        assert sum(p.numel() for p in m.parameters()) == n, name
        assert m.TRUNK_PREFIXES == ref
        x = torch.rand(2, P.T_OBS, P.J, 3) * torch.tensor([P.W, P.H, 1.0])
        with torch.no_grad():
            y = m.eval()(x)
        assert y.shape == (2, P.T_PRED, P.J, 2) and torch.equal(y, x[:, -1:, :, :2].expand_as(y))      # zero-velocity at init
    assert motionbert_kp.parse_variant("motionbert_kp_d192_L4") == {"d": 192, "blocks": 4}
    for bad in ("motionbert_kp_do0.1", "motionbert_kp_lr1e-4", "motionbert_kp_d192_d128", "motionbert_kpx", "stgcn_kp"):
        with pytest.raises(ValueError):
            motionbert_kp.build(bad)
    with pytest.raises(ValueError):
        motionbert_kp.build("motionbert_kp_d192", d=128)                   # width given twice
    # every MB grid model name builds through the same dispatcher train_kp uses (name.startswith("motionbert_kp"))
    for cfg, (model, _) in G.MB_GRID.items():
        motionbert_kp.build(model, dropout=0.0)


def test_run_names_flags_and_grid_sizes():
    assert set(G.KV2_CONFIGS) == {"full", "nospatial", "khead", "nohaz", "hazaux", "plain", "q060", "q090", "lam03"}
    assert G.kv2_run("full", 42) == "cwm3d_v2_s42_v2"
    assert G.kv2_run("nospatial", 123) == "cwm3d_v2_nospatial_s123_v2"
    assert G.kv2_run("nospatial", 123, pre=True) == "cwm3d_v2_nospatial_batonpre_trunkp123_s123_v2"
    assert G.kv2_run("full", 789, pre=True) == "cwm3d_v2_batonpre_trunkp123_s789_v2"
    assert G.mb_run("default", 42) == "motionbert_kp_s42_v2"
    assert G.mb_run("default", 789, pre=True) == "motionbert_kp_batonpre_trunkp123_s789_v2"
    assert G.mb_run("d192_L4_do0.1", 42) == "motionbert_kp_d192_L4_do0.1_s42_v2"
    assert G.k_run(1011) == "cwm3d_kp_s1011_v2" and G.k_run(42, pre=True) == "cwm3d_kp_batonpre_trunkp123_s42_v2"
    assert len(G.MB_GRID) == 10 and len(G.MB_NEW) == 9 and "default" not in G.MB_NEW
    assert G.flags("kv2", "plain") == ["--model", "cwm3d_v2", "--no-physics", "--w-nll", "0"]
    assert G.flags("kv2", "hazaux") == ["--model", "cwm3d_v2", "--haz-mod", "0"]
    assert G.flags("mb", "d192_L4_do0.1") == ["--model", "motionbert_kp_d192_L4", "--dropout", "0.1"]
    assert G.flags("mb", "lr1e-4") == ["--model", "motionbert_kp", "--lr", "1e-4"]
    assert G.mb_changes_dl("d192") and G.mb_changes_dl("d192_L4") and not G.mb_changes_dl("wd0.1") and not G.mb_changes_dl("default")
    assert G.mb_pretrain_run("d192") == "motionbert_kp_d192_batonpre_pretrain_s123"
    assert G.kv2_pretrain_run("nospatial") == "cwm3d_v2_nospatial_batonpre_pretrain_s123"          # own trunk architecture (no spatial / pool)
    assert G.kv2_pretrain_run("lam03") == G.kv2_pretrain_run("full") == "cwm3d_v2_batonpre_pretrain_s123"
    # component counts as pre-registered
    assert {c: n for c, (_, n) in G.KV2_CONFIGS.items()} == {"full": 5, "q060": 5, "q090": 5, "lam03": 5, "nospatial": 4, "khead": 4, "hazaux": 4, "plain": 4, "nohaz": 3}
    # the trainer flags used by the plain / K-side configs exist in train_kp today
    src = (Path(G.ROOT) / "aide" / "train" / "train_kp.py").read_text()
    assert '"--no-physics"' in src and '"--w-nll"' in src and '"--init-parts"' in src


def test_dry_trunk_load_reuse_vs_need(tmp_path):
    torch.manual_seed(0)
    src = motionbert_kp.build("motionbert_kp", dropout=0.0)
    with torch.no_grad():
        for p in src.parameters():
            p.add_(0.05 * torch.randn_like(p))
    ck = tmp_path / "mb.pt"
    torch.save({"model": src.state_dict(), "epoch": 3}, ck)
    torch.manual_seed(1)
    dst = motionbert_kp.build("motionbert_kp", dropout=0.0)
    rep = G.dry_trunk_load(dst, ck)
    assert rep["decision"] == "REUSE_K_CKPT" and rep["trunk_complete"] is True and rep["n_skipped_shape"] == 0 and rep["n_missing_trunk"] == 0
    assert rep["n_loaded"] == rep["n_trunk_tensors"] == sum(k.startswith(tuple(dst.TRUNK_PREFIXES)) for k in dst.state_dict())
    # a wider model against the same checkpoint: every trunk tensor is shape-mismatched or missing -> NEED
    torch.manual_seed(2)
    wide = motionbert_kp.build("motionbert_kp_d192", dropout=0.0)
    rep2 = G.dry_trunk_load(wide, ck)
    assert rep2["decision"] == "NEED_V2_PRETRAIN" and rep2["trunk_complete"] is False and (rep2["n_skipped_shape"] > 0 or rep2["n_missing_trunk"] > 0)


# ----------------------------------------------------------------------------- report smoke (needs the AIDE cache + one finished K run)
def _copy_run(runs_dst: Path, src_name: str, dst_name: str, shift_px: float = 0.0, seed: int = 0, haz: bool = False):
    """Copy an existing run's outputs into a tmp runs dir under a new name; optionally move every prediction by `shift_px` along x
    (a synthetic 'different model') and add a fake haz_P array; rewrite the run json name / t_start_utc."""
    src = P.RESULTS / "runs" / src_name
    dst = runs_dst / dst_name
    dst.mkdir(parents=True)
    rng = np.random.default_rng(seed)
    for split in ("val", "test"):
        z = dict(np.load(src / f"pred_{split}.npz", allow_pickle=False))
        if shift_px:
            z["pred"] = (z["pred"] + rng.normal(shift_px, 1.0, z["pred"].shape).astype(np.float32))
        if haz:
            z["haz_P"] = rng.uniform(0, 1, (len(z["pred"]), 4, 5)).astype(np.float32)
        np.savez(dst / f"pred_{split}.npz", **z)
        shutil.copy(src / f"metrics_{split}.json", dst / f"metrics_{split}.json")
    js = sorted(glob.glob(str(src / f"run_{src_name}_2*.json")))[-1]
    rec = json.load(open(js))
    rec["name"] = dst_name
    rec["t_start_utc"] = "2026-09-15T01:00:00Z"
    (dst / f"run_{dst_name}_20260915T010000Z.json").write_text(json.dumps(rec, default=str))


@pytest.mark.skipif(not P.CACHE_NPZ.exists() or not (P.RESULTS / "runs" / "cwm3d_kp_s42_v2" / "pred_val.npz").exists(),
                    reason="AIDE cache / reference run not available")
def test_report_smoke_with_partial_results(tmp_path):
    from aide.eval import report_kv2_program as R
    runs = tmp_path / "runs"
    for s in (42, 123):
        _copy_run(runs, f"cwm3d_kp_s{s}_v2", G.k_run(s))                                # K baselines (copied, not retrained)
        _copy_run(runs, f"motionbert_kp_s{s}_v2", G.mb_run("default", s))
    _copy_run(runs, "cwm3d_kp_s42_v2", G.kv2_run("full", 42), shift_px=0.0, seed=1, haz=True)   # identical predictions -> Δ = 0 -> KILL
    _copy_run(runs, "cwm3d_kp_s42_v2", G.kv2_run("nospatial", 42), shift_px=+3.0, seed=2)      # worse
    _copy_run(runs, "motionbert_kp_s42_v2", G.mb_run("do0.1", 42), shift_px=+2.0, seed=3)
    kv2 = tmp_path / "kv2"
    kv2.mkdir()
    (kv2 / "S0_UTC.txt").write_text("2026-09-14T22:51:07Z\n")
    # gate on partial results: incomplete, pending rows, no survivors
    dec = G.stage_s1(runs, kv2, N_BOOT)
    assert dec["complete"] is False and dec["survivors"] == []
    st = {r["cfg"]: r["status"] for r in dec["rows"]}
    assert st["full"] == "KILL" and st["nospatial"] == "KILL" and st["khead"] == "PENDING" and sum(v == "PENDING" for v in st.values()) == 7
    assert abs(next(r for r in dec["rows"] if r["cfg"] == "full")["d_all"]["mean"]) < 1e-9
    assert (kv2 / "s1_decision.json").exists() and len(list(kv2.glob("s1_decision_*.json"))) == 1
    j = json.load(open(kv2 / "s1_decision.json"))
    assert set(("utc", "rule", "rows", "survivors")) <= set(j)
    dec3 = G.stage_s3(runs, kv2, N_BOOT)
    assert dec3["complete"] is False and dec3["mb_best"] == "default"           # do0.1 is +2 px worse
    assert (kv2 / "s3_decision.json").exists()
    assert G.stage_s1b("full", runs, kv2, N_BOOT)["decision"] == "PENDING"     # seed 123 K-v2 missing
    # cli read helper
    assert G.main(["read", "s1_decision", "complete", "--out", str(kv2)]) == 0
    # report on the tmp root: runs gracefully, pending everywhere, conservative adoption line, append-only on the second call
    tables = tmp_path / "tables"
    assert R.main(["--out-root", str(tmp_path), "--n-boot", str(N_BOOT)]) == 0
    md = (tables / "kv2_program.md").read_text(encoding="utf-8")
    for sec in ("## 0. Pre-registration", "## 1. S1 val screen", "### 1.1 S1b", "## 2. S2", "## 3. S3", "## 4. Primary endpoints", "### 4.1 Paper-wording", "## 5. GPU minutes", "## 6. Provenance"):
        assert sec in md, sec
    assert "**ADOPTION A1: NOT ADOPTED**" in md and "**ADOPTION A2: NOT ADOPTED**" in md and "**ADOPTION A1: ADOPTED**" not in md   # exact bold marker only
    assert "pending" in md and "PENDING" in md
    assert "S0 UTC: **2026-09-14T22:51:07Z**" in md and "VIOLATION" not in md                  # copied runs start after S0
    assert "K-v2 programme — S0 pre-registration" in md                                         # the row is quoted from docs/decisions.md
    assert (tables / "kv2_program.json").exists() and len(json.load(open(tables / "kv2_program.json"))) == 1
    assert list(tables.glob("run_report_kv2_program_2*.json"))
    n1 = md.count("## 0. Pre-registration")
    assert R.main(["--out-root", str(tmp_path), "--n-boot", str(N_BOOT)]) == 0
    md2 = (tables / "kv2_program.md").read_text(encoding="utf-8")
    assert md2.startswith(md) and md2.count("## 0. Pre-registration") == n1 + 1 and "appended regeneration" in md2
    assert len(json.load(open(tables / "kv2_program.json"))) == 2
    # AUC helper: perfect separation -> 1, inverted -> 0, ties -> 0.5
    assert R.rank_auc([0.1, 0.2, 0.9, 0.8], [False, False, True, True]) == 1.0
    assert R.rank_auc([0.9, 0.8, 0.1, 0.2], [False, False, True, True]) == 0.0
    assert R.rank_auc([0.5, 0.5, 0.5, 0.5], [False, False, True, True]) == 0.5

Xet Storage Details

Size:: 15.1 kB
Xet hash:: df00b0abaf6fc8222a97c19265646f74b90778176893c75907be910b4c232856

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.

	"""Guards for the K-v2 programme protocol side (docs/decisions.md 2026-09-14 "K-v2 programme — S0 pre-registration", protocol C23):

	* `kv2_gate` S1 rule on synthetic per-clip arrays: survive / kill / HM guard / tie-with-fewer-components / ≤ 2 survivors;
	* S1b continuation rule; S3 argmin + 0.10-px tie -> fewer parameters;
	* MotionBERT-kp name-encoded variants build with the expected parameter counts, TRUNK_PREFIXES unchanged, foreign suffixes rejected;
	* run-name / flag helpers shared by the launcher, the gate and the report;
	* `dry_trunk_load` mechanics (REUSE on a matching checkpoint, NEED on a foreign one);
	* `report_kv2_program` smoke on a tmp out-root with partial results (pending runs, conservative NOT ADOPTED line, append-only).

	python -m pytest aide/tests/test_kv2.py -q
	"""
	from __future__ import annotations

	import os

	os.environ.setdefault("OMP_NUM_THREADS", "4")

	import glob # noqa: E402
	import json # noqa: E402
	import shutil # noqa: E402
	import sys # noqa: E402
	from pathlib import Path # noqa: E402

	import numpy as np # noqa: E402
	import pytest # noqa: E402
	import torch # noqa: E402

	torch.set_num_threads(4)
	sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
	from aide.data import protocol as P # noqa: E402
	from aide.eval import kv2_gate as G # noqa: E402
	from aide.models import motionbert_kp # noqa: E402

	N_BOOT = 400


	def _synthetic(seed=0, n=405, n_hm=40):
	rng = np.random.default_rng(seed)
	base = np.abs(rng.normal(55, 25, n))
	hm = np.zeros(n, bool)
	hm[rng.choice(n, n_hm, replace=False)] = True
	return rng, base, hm


	def _cand(rng, base, shift, noise=0.05):
	return base + shift + rng.normal(0, noise, len(base))


	def test_s1_rule_survive_kill_guard_tie_and_cap():
	rng, base, hm = _synthetic()
	cands = { # cfg: (shift on all clips, extra shift on HM clips, components)
	"full": (-0.50, 0.0, 5), # survives, mean −0.50
	"nohaz": (-0.45, 0.0, 3), # within 0.10 of full with fewer components -> ranked FIRST
	"plain": (-0.31, 0.0, 4), # survives but only 2 slots
	"khead": (-0.10, 0.0, 4), # KILL (> −0.30)
	"q060": (-0.30, 0.0, 5), # exactly at the threshold -> survives (≤)
	"hazaux": (-0.60, +6.0, 4), # best All but HM guard: Δ HM > +3 with CI excl. 0 -> GUARD_KILL
	}
	rows = []
	for cfg, (sh, sh_hm, nc) in cands.items():
	c = _cand(rng, base, sh)
	c_hm = c[hm] + sh_hm
	rows.append(G.s1_row(cfg, G.kv2_run(cfg, 42), c, base, c_hm, base[hm], nc, N_BOOT, 0))
	st = {r["cfg"]: r["status"] for r in rows}
	assert st == {"full": "SURVIVE", "nohaz": "SURVIVE", "plain": "SURVIVE", "khead": "KILL", "q060": "SURVIVE", "hazaux": "GUARD_KILL"}
	r_haz = next(r for r in rows if r["cfg"] == "hazaux")
	assert r_haz["survive_all"] and r_haz["guard_kill"] and r_haz["d_hm"]["ci_lo"] > 0 and r_haz["d_hm"]["mean"] > 3.0
	assert abs(next(r for r in rows if r["cfg"] == "full")["d_all"]["mean"] + 0.50) < 0.02
	surv = G.s1_select(rows)
	assert surv == ["nohaz", "full"] # tie rule (fewer components first), then the cap of 2 (plain / q060 excluded, hazaux guard-killed)
	# without the tie: a clear leader wins regardless of its component count
	rows2 = [r for r in rows if r["cfg"] != "nohaz"]
	assert G.s1_select(rows2) == ["full", "plain"]
	# a guard flagged but the All rule failed -> plain KILL, never GUARD_KILL
	c = _cand(rng, base, -0.05)
	r = G.s1_row("lam03", "x", c, base, c[hm] + 6.0, base[hm], 5, N_BOOT, 0)
	assert r["status"] == "KILL" and r["guard_kill"]
	assert G.s1_select([]) == []


	def test_s1b_rule():
	ok, why = G.s1b_decide(-0.40, -0.15)
	assert ok and why.startswith("CONTINUE") # Δ123 < 0 and mean −0.275 ≤ −0.25
	assert not G.s1b_decide(-0.60, +0.05)[0] # Δ123 ≥ 0 -> killed even though the mean passes
	assert not G.s1b_decide(-0.30, -0.10)[0] # mean −0.20 > −0.25 -> killed
	ok, why = G.s1b_decide(-0.25, -0.25)
	assert ok and "val-killed" not in why
	assert "val-killed at seed 123" in G.s1b_decide(-0.30, -0.10)[1]


	def test_s3_argmin_and_tie_to_fewer_params():
	rows = [{"cfg": "default", "val_all": 55.55, "n_params": 813_088},
	{"cfg": "d192", "val_all": 55.50, "n_params": 1_809_440}, # within 0.10 of the best -> fewer params wins
	{"cfg": "L4", "val_all": 55.70, "n_params": 1_077_664}]
	assert G.s3_select(rows) == "default"
	rows.append({"cfg": "do0.1", "val_all": 55.30, "n_params": 813_088}) # clear leader (> 0.10 better than default)
	assert G.s3_select(rows) == "do0.1"
	rows.append({"cfg": "wd0.1", "val_all": 55.25, "n_params": 813_088}) # ties do0.1 (same params) -> lower val
	assert G.s3_select(rows) == "wd0.1"
	assert G.s3_select([{"cfg": "x", "val_all": None, "n_params": 1}]) is None
	assert G.s3_select([]) is None


	def test_mb_name_variants_build_with_expected_params_and_same_trunk_prefixes():
	expected = {"motionbert_kp": 813_088, "motionbert_kp_2b": 548_512, "motionbert_kp_d192": 1_809_440,
	"motionbert_kp_L4": 1_077_664, "motionbert_kp_d192_L4": 2_402_912}
	ref = motionbert_kp.MotionBERTKP.TRUNK_PREFIXES
	assert ref == ("inp.", "joint_emb", "time_emb", "blocks.", "ln_f.")
	for name, n in expected.items():
	m = motionbert_kp.build(name, dropout=0.0)
	assert sum(p.numel() for p in m.parameters()) == n, name
	assert m.TRUNK_PREFIXES == ref
	x = torch.rand(2, P.T_OBS, P.J, 3) * torch.tensor([P.W, P.H, 1.0])
	with torch.no_grad():
	y = m.eval()(x)
	assert y.shape == (2, P.T_PRED, P.J, 2) and torch.equal(y, x[:, -1:, :, :2].expand_as(y)) # zero-velocity at init
	assert motionbert_kp.parse_variant("motionbert_kp_d192_L4") == {"d": 192, "blocks": 4}
	for bad in ("motionbert_kp_do0.1", "motionbert_kp_lr1e-4", "motionbert_kp_d192_d128", "motionbert_kpx", "stgcn_kp"):
	with pytest.raises(ValueError):
	motionbert_kp.build(bad)
	with pytest.raises(ValueError):
	motionbert_kp.build("motionbert_kp_d192", d=128) # width given twice
	# every MB grid model name builds through the same dispatcher train_kp uses (name.startswith("motionbert_kp"))
	for cfg, (model, _) in G.MB_GRID.items():
	motionbert_kp.build(model, dropout=0.0)


	def test_run_names_flags_and_grid_sizes():
	assert set(G.KV2_CONFIGS) == {"full", "nospatial", "khead", "nohaz", "hazaux", "plain", "q060", "q090", "lam03"}
	assert G.kv2_run("full", 42) == "cwm3d_v2_s42_v2"
	assert G.kv2_run("nospatial", 123) == "cwm3d_v2_nospatial_s123_v2"
	assert G.kv2_run("nospatial", 123, pre=True) == "cwm3d_v2_nospatial_batonpre_trunkp123_s123_v2"
	assert G.kv2_run("full", 789, pre=True) == "cwm3d_v2_batonpre_trunkp123_s789_v2"
	assert G.mb_run("default", 42) == "motionbert_kp_s42_v2"
	assert G.mb_run("default", 789, pre=True) == "motionbert_kp_batonpre_trunkp123_s789_v2"
	assert G.mb_run("d192_L4_do0.1", 42) == "motionbert_kp_d192_L4_do0.1_s42_v2"
	assert G.k_run(1011) == "cwm3d_kp_s1011_v2" and G.k_run(42, pre=True) == "cwm3d_kp_batonpre_trunkp123_s42_v2"
	assert len(G.MB_GRID) == 10 and len(G.MB_NEW) == 9 and "default" not in G.MB_NEW
	assert G.flags("kv2", "plain") == ["--model", "cwm3d_v2", "--no-physics", "--w-nll", "0"]
	assert G.flags("kv2", "hazaux") == ["--model", "cwm3d_v2", "--haz-mod", "0"]
	assert G.flags("mb", "d192_L4_do0.1") == ["--model", "motionbert_kp_d192_L4", "--dropout", "0.1"]
	assert G.flags("mb", "lr1e-4") == ["--model", "motionbert_kp", "--lr", "1e-4"]
	assert G.mb_changes_dl("d192") and G.mb_changes_dl("d192_L4") and not G.mb_changes_dl("wd0.1") and not G.mb_changes_dl("default")
	assert G.mb_pretrain_run("d192") == "motionbert_kp_d192_batonpre_pretrain_s123"
	assert G.kv2_pretrain_run("nospatial") == "cwm3d_v2_nospatial_batonpre_pretrain_s123" # own trunk architecture (no spatial / pool)
	assert G.kv2_pretrain_run("lam03") == G.kv2_pretrain_run("full") == "cwm3d_v2_batonpre_pretrain_s123"
	# component counts as pre-registered
	assert {c: n for c, (_, n) in G.KV2_CONFIGS.items()} == {"full": 5, "q060": 5, "q090": 5, "lam03": 5, "nospatial": 4, "khead": 4, "hazaux": 4, "plain": 4, "nohaz": 3}
	# the trainer flags used by the plain / K-side configs exist in train_kp today
	src = (Path(G.ROOT) / "aide" / "train" / "train_kp.py").read_text()
	assert '"--no-physics"' in src and '"--w-nll"' in src and '"--init-parts"' in src


	def test_dry_trunk_load_reuse_vs_need(tmp_path):
	torch.manual_seed(0)
	src = motionbert_kp.build("motionbert_kp", dropout=0.0)
	with torch.no_grad():
	for p in src.parameters():
	p.add_(0.05 * torch.randn_like(p))
	ck = tmp_path / "mb.pt"
	torch.save({"model": src.state_dict(), "epoch": 3}, ck)
	torch.manual_seed(1)
	dst = motionbert_kp.build("motionbert_kp", dropout=0.0)
	rep = G.dry_trunk_load(dst, ck)
	assert rep["decision"] == "REUSE_K_CKPT" and rep["trunk_complete"] is True and rep["n_skipped_shape"] == 0 and rep["n_missing_trunk"] == 0
	assert rep["n_loaded"] == rep["n_trunk_tensors"] == sum(k.startswith(tuple(dst.TRUNK_PREFIXES)) for k in dst.state_dict())
	# a wider model against the same checkpoint: every trunk tensor is shape-mismatched or missing -> NEED
	torch.manual_seed(2)
	wide = motionbert_kp.build("motionbert_kp_d192", dropout=0.0)
	rep2 = G.dry_trunk_load(wide, ck)
	assert rep2["decision"] == "NEED_V2_PRETRAIN" and rep2["trunk_complete"] is False and (rep2["n_skipped_shape"] > 0 or rep2["n_missing_trunk"] > 0)


	# ----------------------------------------------------------------------------- report smoke (needs the AIDE cache + one finished K run)
	def _copy_run(runs_dst: Path, src_name: str, dst_name: str, shift_px: float = 0.0, seed: int = 0, haz: bool = False):
	"""Copy an existing run's outputs into a tmp runs dir under a new name; optionally move every prediction by `shift_px` along x
	(a synthetic 'different model') and add a fake haz_P array; rewrite the run json name / t_start_utc."""
	src = P.RESULTS / "runs" / src_name
	dst = runs_dst / dst_name
	dst.mkdir(parents=True)
	rng = np.random.default_rng(seed)
	for split in ("val", "test"):
	z = dict(np.load(src / f"pred_{split}.npz", allow_pickle=False))
	if shift_px:
	z["pred"] = (z["pred"] + rng.normal(shift_px, 1.0, z["pred"].shape).astype(np.float32))
	if haz:
	z["haz_P"] = rng.uniform(0, 1, (len(z["pred"]), 4, 5)).astype(np.float32)
	np.savez(dst / f"pred_{split}.npz", **z)
	shutil.copy(src / f"metrics_{split}.json", dst / f"metrics_{split}.json")
	js = sorted(glob.glob(str(src / f"run_{src_name}_2*.json")))[-1]
	rec = json.load(open(js))
	rec["name"] = dst_name
	rec["t_start_utc"] = "2026-09-15T01:00:00Z"
	(dst / f"run_{dst_name}_20260915T010000Z.json").write_text(json.dumps(rec, default=str))


	@pytest.mark.skipif(not P.CACHE_NPZ.exists() or not (P.RESULTS / "runs" / "cwm3d_kp_s42_v2" / "pred_val.npz").exists(),
	reason="AIDE cache / reference run not available")
	def test_report_smoke_with_partial_results(tmp_path):
	from aide.eval import report_kv2_program as R
	runs = tmp_path / "runs"
	for s in (42, 123):
	_copy_run(runs, f"cwm3d_kp_s{s}_v2", G.k_run(s)) # K baselines (copied, not retrained)
	_copy_run(runs, f"motionbert_kp_s{s}_v2", G.mb_run("default", s))
	_copy_run(runs, "cwm3d_kp_s42_v2", G.kv2_run("full", 42), shift_px=0.0, seed=1, haz=True) # identical predictions -> Δ = 0 -> KILL
	_copy_run(runs, "cwm3d_kp_s42_v2", G.kv2_run("nospatial", 42), shift_px=+3.0, seed=2) # worse
	_copy_run(runs, "motionbert_kp_s42_v2", G.mb_run("do0.1", 42), shift_px=+2.0, seed=3)
	kv2 = tmp_path / "kv2"
	kv2.mkdir()
	(kv2 / "S0_UTC.txt").write_text("2026-09-14T22:51:07Z\n")
	# gate on partial results: incomplete, pending rows, no survivors
	dec = G.stage_s1(runs, kv2, N_BOOT)
	assert dec["complete"] is False and dec["survivors"] == []
	st = {r["cfg"]: r["status"] for r in dec["rows"]}
	assert st["full"] == "KILL" and st["nospatial"] == "KILL" and st["khead"] == "PENDING" and sum(v == "PENDING" for v in st.values()) == 7
	assert abs(next(r for r in dec["rows"] if r["cfg"] == "full")["d_all"]["mean"]) < 1e-9
	assert (kv2 / "s1_decision.json").exists() and len(list(kv2.glob("s1_decision_*.json"))) == 1
	j = json.load(open(kv2 / "s1_decision.json"))
	assert set(("utc", "rule", "rows", "survivors")) <= set(j)
	dec3 = G.stage_s3(runs, kv2, N_BOOT)
	assert dec3["complete"] is False and dec3["mb_best"] == "default" # do0.1 is +2 px worse
	assert (kv2 / "s3_decision.json").exists()
	assert G.stage_s1b("full", runs, kv2, N_BOOT)["decision"] == "PENDING" # seed 123 K-v2 missing
	# cli read helper
	assert G.main(["read", "s1_decision", "complete", "--out", str(kv2)]) == 0
	# report on the tmp root: runs gracefully, pending everywhere, conservative adoption line, append-only on the second call
	tables = tmp_path / "tables"
	assert R.main(["--out-root", str(tmp_path), "--n-boot", str(N_BOOT)]) == 0
	md = (tables / "kv2_program.md").read_text(encoding="utf-8")
	for sec in ("## 0. Pre-registration", "## 1. S1 val screen", "### 1.1 S1b", "## 2. S2", "## 3. S3", "## 4. Primary endpoints", "### 4.1 Paper-wording", "## 5. GPU minutes", "## 6. Provenance"):
	assert sec in md, sec
	assert "ADOPTION A1: NOT ADOPTED" in md and "ADOPTION A2: NOT ADOPTED" in md and "ADOPTION A1: ADOPTED" not in md # exact bold marker only
	assert "pending" in md and "PENDING" in md
	assert "S0 UTC: 2026-09-14T22:51:07Z" in md and "VIOLATION" not in md # copied runs start after S0
	assert "K-v2 programme — S0 pre-registration" in md # the row is quoted from docs/decisions.md
	assert (tables / "kv2_program.json").exists() and len(json.load(open(tables / "kv2_program.json"))) == 1
	assert list(tables.glob("run_report_kv2_program_2*.json"))
	n1 = md.count("## 0. Pre-registration")
	assert R.main(["--out-root", str(tmp_path), "--n-boot", str(N_BOOT)]) == 0
	md2 = (tables / "kv2_program.md").read_text(encoding="utf-8")
	assert md2.startswith(md) and md2.count("## 0. Pre-registration") == n1 + 1 and "appended regeneration" in md2
	assert len(json.load(open(tables / "kv2_program.json"))) == 2
	# AUC helper: perfect separation -> 1, inverted -> 0, ties -> 0.5
	assert R.rank_auc([0.1, 0.2, 0.9, 0.8], [False, False, True, True]) == 1.0
	assert R.rank_auc([0.9, 0.8, 0.1, 0.2], [False, False, True, True]) == 0.0
	assert R.rank_auc([0.5, 0.5, 0.5, 0.5], [False, False, True, True]) == 0.5