Buckets:
| """cwm3/tests/test_overlap_head.py — F8 overlap head (models/overlap_head.py, train/train_overlap_head.py, eval/overlap_eval.py; R2 WP4 batch 2, 2026-09-17; | |
| pre-registration docs/decisions.md 2026-09-17T19:32:07Z). | |
| (a) IDENTITY (pre-registered): the joint head replaced by the product distribution joint_from_marginals(p_hands_off, p_head, p_D) reproduces | |
| q_k == p_U · p_D EXACTLY, with p_U = 1 − (1 − p_hands_off)(1 − p_head) (manual) and p_U = p_head (engaged); the 8-way law sums to 1 and its marginals | |
| are the inputs. | |
| (b) masked joint NLL: fully observed = the categorical NLL of the state; one unobserved variable = −log of the marginal over it; nothing observed = 0 and | |
| excluded from the mask; the loss is finite and differentiable. | |
| (c) strict future: build_targets maps horizon k (>= 1) to index k−1 of the future arrays (a state change at step k* appears at index k*−1 only); poisoning | |
| the anchor (S0 / P0) leaves the targets bitwise unchanged; poisoning the future leaves the anchor features unchanged. | |
| (d) shapes / params: OverlapHead forward -> [B, 15, 8]; joint arm < 80 k parameters; margonly arm ignores z (use_z False) and consumes zk [B, 15, 16]. | |
| (e) realised_overlap: y = any_k (U ∧ D) over fully observed steps only; J_from_q = 0.2 Σ q. | |
| (f) alert_timing_eval.overlap_spec: `name=DIR` -> ov_<name>, bare DIR -> ov_J (backward compatible). | |
| (g) overlap_eval.paired: Brier / AUROC / PR-AUC on identical inputs give Δ = 0 with a degenerate CI; the tie-aware AUROC equals the mid-rank formula. | |
| Synthetic data only (never a result). Run: cd cwm3 && python -m pytest tests/test_overlap_head.py -q | |
| """ | |
| from __future__ import annotations | |
| import os | |
| import sys | |
| from pathlib import Path | |
| os.environ.setdefault("OMP_NUM_THREADS", "4") | |
| import numpy as np # noqa: E402 | |
| import torch # noqa: E402 | |
| torch.set_num_threads(4) | |
| CWM3 = Path(__file__).resolve().parents[1] | |
| for p in (CWM3, CWM3 / "data", CWM3 / "train", CWM3 / "common", CWM3 / "eval", CWM3.parent / "s12_wm", CWM3.parent / "cwm2" / "eval"): | |
| if str(p) not in sys.path: | |
| sys.path.insert(0, str(p)) | |
| from models import overlap_head as oh # noqa: E402 | |
| def test_product_identity(): | |
| torch.manual_seed(0) | |
| B, Kh = 64, oh.K | |
| p_ho, p_ht, p_D = torch.rand(B, Kh, dtype=torch.float64), torch.rand(B, Kh, dtype=torch.float64), torch.rand(B, Kh, dtype=torch.float64) | |
| P8 = oh.joint_from_marginals(p_ho, p_ht, p_D) | |
| assert P8.shape == (B, Kh, 8) | |
| assert torch.allclose(P8.sum(-1), torch.ones(B, Kh, dtype=torch.float64), atol=1e-12) | |
| mg = oh.marginals_from_joint(P8) | |
| assert torch.allclose(mg["p_hands_off"], p_ho, atol=1e-12) and torch.allclose(mg["p_head_turned"], p_ht, atol=1e-12) and torch.allclose(mg["p_D"], p_D, atol=1e-12) | |
| pU = 1 - (1 - p_ho) * (1 - p_ht) | |
| q_manual = oh.overlap_from_joint(P8, use_hands=True) | |
| q_engaged = oh.overlap_from_joint(P8, use_hands=False) | |
| assert torch.allclose(q_manual, pU * p_D, atol=1e-12), "identity test: q != p_U p_D under the product distribution" | |
| assert torch.allclose(q_engaged, p_ht * p_D, atol=1e-12) | |
| assert torch.allclose(mg["p_U"], pU, atol=1e-12) | |
| # degenerate corners | |
| P8 = oh.joint_from_marginals(torch.zeros(1, 1, dtype=torch.float64), torch.ones(1, 1, dtype=torch.float64), torch.ones(1, 1, dtype=torch.float64)) | |
| assert float(oh.overlap_from_joint(P8, True)) == 1.0 and float(oh.overlap_from_joint(P8, False)) == 1.0 | |
| def test_masked_joint_nll_forms(): | |
| torch.manual_seed(1) | |
| logits = torch.randn(5, 3, 8, requires_grad=True) | |
| logp = torch.log_softmax(logits, -1) | |
| y3 = torch.tensor([[[1, 0, 1]] * 3] * 5, dtype=torch.uint8) | |
| # fully observed: -log p_s with s = 4*1 + 2*0 + 1 = 5 | |
| obs = torch.ones(5, 3, 3, dtype=torch.bool) | |
| nll, m = oh.masked_joint_nll(logits, y3, obs) | |
| assert m.all() and torch.allclose(nll, -logp[..., 5], atol=1e-6) | |
| # D unobserved: -log (p_4 + p_5) | |
| obs2 = obs.clone(); obs2[..., 2] = False | |
| nll2, m2 = oh.masked_joint_nll(logits, y3, obs2) | |
| assert m2.all() and torch.allclose(nll2, -torch.logsumexp(logp[..., [4, 5]], -1), atol=1e-6) | |
| # hands unobserved: states with head=0, D=1 -> s in {1, 5} | |
| obs3 = obs.clone(); obs3[..., 0] = False | |
| nll3, _ = oh.masked_joint_nll(logits, y3, obs3) | |
| assert torch.allclose(nll3, -torch.logsumexp(logp[..., [1, 5]], -1), atol=1e-6) | |
| # nothing observed -> 0, excluded | |
| obs4 = torch.zeros_like(obs); obs4[0, 0] = True | |
| nll4, m4 = oh.masked_joint_nll(logits, y3, obs4) | |
| assert int(m4.sum()) == 1 and float(nll4[1:].abs().sum()) == 0.0 | |
| loss, parts = oh.overlap_loss(logits, y3, obs2) | |
| assert torch.isfinite(loss) | |
| loss.backward() | |
| assert torch.isfinite(logits.grad).all() | |
| def test_strict_future_targets_and_anchor_independence(): | |
| import train_overlap_head as TO | |
| rng = np.random.default_rng(0) | |
| N, H = 50, 25 | |
| S = np.zeros((N, H, 2), np.float32); S[:, :, 0] = 1.0 # hands on, head not turned everywhere | |
| kstar = 7 | |
| S[0, kstar - 1, 0] = 0.0 # hands leave at horizon k* = 7 -> index 6 | |
| S[1, :, 1] = np.nan # head unobserved for row 1 | |
| P = np.zeros((N, H, 6), bool); Pv = np.ones((N, H, 6), bool); M = np.ones((N, H), bool); ok = np.ones(N, bool) | |
| P[2, 9, 3] = True # demand at horizon 10 -> index 9 | |
| Pv[3, 4] = False # D unobserved at horizon 5 for row 3 | |
| y3, obs3 = oh.build_targets(S, P, Pv, M, ok) | |
| assert y3.shape == (N, oh.K, 3) and obs3.shape == (N, oh.K, 3) | |
| assert y3[0, :, 0].tolist() == [1 if k == kstar - 1 else 0 for k in range(oh.K)] | |
| assert not obs3[1, :, 1].any() and obs3[1, :, 0].all() | |
| assert y3[2, :, 2].tolist() == [1 if k == 9 else 0 for k in range(oh.K)] | |
| assert not obs3[3, 4, 2] and obs3[3, 5, 2] | |
| # anchor poisoning leaves the targets untouched; future poisoning leaves the anchor features untouched | |
| tg = {"S0": rng.normal(size=(N, 2)).astype(np.float32), "P0": rng.random((N, 6)) < 0.2, "P0valid": rng.random((N, 6)) < 0.9} | |
| A = TO.anchor_features(tg) | |
| assert A.shape == (N, oh.D_ANCHOR) | |
| S_poison = S.copy(); S_poison[:, :, :] = np.nan | |
| A2 = TO.anchor_features(tg) # anchor features do not read S at all | |
| assert np.array_equal(A, A2) | |
| y3b, obs3b = oh.build_targets(S, P, Pv, M, ok) | |
| tg2 = {"S0": tg["S0"] + 5, "P0": ~tg["P0"], "P0valid": tg["P0valid"]} | |
| assert np.array_equal(y3, y3b) and np.array_equal(obs3, obs3b) and not np.array_equal(TO.anchor_features(tg2), A) | |
| # realised overlap (e): row 4 has hands off at k=3 and demand at k=3 -> y; row 5 has them at different steps -> no | |
| S2 = S.copy(); P2 = P.copy() | |
| S2[4, 2, 0] = 0.0; P2[4, 2, 0] = True | |
| S2[5, 2, 0] = 0.0; P2[5, 6, 0] = True | |
| y3c, obs3c = oh.build_targets(S2, P2, Pv, M, ok) | |
| y, full = oh.realised_overlap(y3c, obs3c, use_hands=True) | |
| assert bool(y[4]) and not bool(y[5]) and full[4] and not full[3] | |
| yh, _ = oh.realised_overlap(y3c, obs3c, use_hands=False) | |
| assert not bool(yh[4]) | |
| q = torch.full((2, oh.K), 0.5) | |
| assert torch.allclose(oh.J_from_q(q), torch.full((2,), 0.2 * 0.5 * oh.K)) | |
| def test_head_shapes_params_and_margonly(): | |
| m = oh.OverlapHead(d_z=384, d=96) | |
| n = oh.n_params(m) | |
| assert 30_000 < n < 80_000, n | |
| z, a, mode = torch.randn(7, 384), torch.rand(7, oh.D_ANCHOR), torch.randint(0, 2, (7,)) | |
| out = m(z, a, mode) | |
| assert out.shape == (7, oh.K, 8) and torch.isfinite(out).all() | |
| mm = oh.OverlapHead(d_z=384, d=96, d_k=16, use_z=False) | |
| zk = torch.randn(7, oh.K, 16) | |
| o1 = mm(torch.zeros(7, 1), a, mode, zk) | |
| o2 = mm(torch.randn(7, 1), a, mode, zk) | |
| assert torch.allclose(o1, o2) and o1.shape == (7, oh.K, 8) | |
| def test_alert_timing_overlap_spec(): | |
| import alert_timing_eval as AT | |
| assert AT.overlap_spec("OV_joint_s0") == ("ov_J", "OV_joint_s0") | |
| assert AT.overlap_spec("joint=OV_joint_s0") == ("ov_joint", "OV_joint_s0") | |
| assert AT.overlap_spec("prod=OV_prod_s{seed}".format(seed=2)) == ("ov_prod", "OV_prod_s2") | |
| def test_overlap_eval_paired_identity_and_tie_aware_auroc(): | |
| import overlap_eval as OE | |
| rng = np.random.default_rng(3) | |
| N = 2000 | |
| y = rng.random(N) < 0.1 | |
| s = np.round(rng.random(N), 1) # many ties | |
| dongle = np.array([f"d{i % 17}" for i in range(N)]) | |
| boot = OE.Boot(dongle, 50, 0) | |
| for metric in ("brier", "auroc", "pr_auc"): | |
| r = OE.paired(y, s, s, boot, metric) | |
| assert r["delta"] == 0.0 and r["ci_d"] == (0.0, 0.0) and r["n_dev"] == 17 | |
| # tie-aware AUROC == mid-rank formula (Mann-Whitney with ties) | |
| prep = OE._auroc_prep(y, s) | |
| a = OE._auroc_w(prep, np.ones(N)) | |
| from scipy.stats import rankdata | |
| ranks = rankdata(s) | |
| n_pos, n_neg = y.sum(), (~y).sum() | |
| mw = (ranks[y].sum() - n_pos * (n_pos + 1) / 2) / (n_pos * n_neg) | |
| assert abs(a - mw) < 1e-9, (a, mw) | |
Xet Storage Details
- Size:
- 9.21 kB
- Xet hash:
- d04864dedf185c3e950c890d15cf2e5624d277c4a2adfb9e32093668fc1bb41f
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.