Buckets:
| """cwm3/tests/test_bank_causality.py — the token bank never lets an anchor see a feature window that ends after the anchor. | |
| Rule under test (data/token_bank.py, derived from s12_wm/extract_crop.py): a bank row with clip END time t_row (window | |
| [t_row-2, t_row]; rel = linspace(-2, 0, 16) so the last decoded frame is exactly t_row) is usable for an anchor at T iff | |
| t_row <= T; slot j targets floor(2T)/2 - (15-j)*0.5. Checks, on the materialised files: | |
| (a) the extraction code still has rel = np.linspace(-CLIP_SPAN, 0.0, CLIP_FRAMES) with CLIP_SPAN = 2.0 (window ends at t) | |
| (b) every gathered bank row satisfies t_bank <= t_anchor + 1e-6 and belongs to the anchor's own route | |
| (c) slots are 0.5 s apart and the last slot lags the anchor by < 0.5 s | |
| (d) the anchor_index_for() builder refuses a bank row placed 0.1 s after the anchor (synthetic) | |
| (e) the E-branch attention mask is causal across slots (no query slot attends to a later slot; the readout sees all) | |
| Run: cd cwm3 && python -m pytest tests/test_bank_causality.py -q | |
| """ | |
| from __future__ import annotations | |
| import re | |
| import sys | |
| from pathlib import Path | |
| import numpy as np | |
| import pytest | |
| CWM3 = Path(__file__).resolve().parents[1] | |
| EXPLORE = CWM3.parent | |
| sys.path.insert(0, str(CWM3 / "data")) | |
| sys.path.insert(0, str(CWM3)) | |
| import token_bank as tb # noqa: E402 | |
| BANK = CWM3 / "results" / "banks" / "front_tok.npz" | |
| S12R = EXPLORE / "s12_wm" / "results" | |
| have_bank = BANK.exists() and (CWM3 / "results" / "banks" / "anchor_index_manual_val.npy").exists() | |
| def test_extraction_window_ends_at_t(): | |
| src = (EXPLORE / "s12_wm" / "extract_crop.py").read_text() | |
| assert re.search(r"CLIP_FRAMES,\s*CLIP_SPAN,\s*STRIDE\s*=\s*16,\s*2\.0,\s*0\.5", src) | |
| assert "rel = np.linspace(-CLIP_SPAN, 0.0, CLIP_FRAMES)" in src | |
| assert "np.round((t + rel) * 20)" in src # frames at (t + rel) * 20 -> last frame exactly at t | |
| assert "clip 覆盖 [ts-2, ts]" in src | |
| def test_gathered_rows_are_causal_and_same_route(regime, split): | |
| with np.load(BANK, allow_pickle=True) as z: | |
| bank_t, names, off = z["t"].astype(np.float64), [str(u) for u in z["uid_index"]], z["route_off"] | |
| idx = np.load(CWM3 / "results" / "banks" / f"anchor_index_{regime}_{split}.npy") | |
| with np.load(S12R / f"targets_{regime}_{split}.npz", allow_pickle=True) as z: | |
| uid, t = z["uid"], z["t"].astype(np.float64) | |
| assert idx.shape == (len(t), 16) | |
| ok = idx >= 0 | |
| assert ok[:, -1].mean() > 0.95 | |
| # (b) causal | |
| lag = t[:, None] - bank_t[np.where(ok, idx, 0)] | |
| assert np.all(lag[ok] >= -1e-6), "a gathered bank row ends after its anchor" | |
| assert lag[ok].max() <= 7.9 + 1e-6 | |
| # (b) same route: bank row's segment == the anchor's uid | |
| seg = np.searchsorted(off, idx[ok], side="right") - 1 | |
| route_of = {u: i for i, u in enumerate(names)} | |
| r_anchor = np.array([route_of.get(u, -1) for u in uid]) | |
| assert not ok[r_anchor < 0].any(), "anchors of routes absent from the bank must have no gathered rows" | |
| want = r_anchor[:, None].repeat(16, 1)[ok] | |
| assert np.array_equal(seg, want) | |
| # (c) spacing 0.5 s where both neighbours exist; last slot lag < 0.5 s | |
| both = ok[:, 1:] & ok[:, :-1] | |
| d = bank_t[idx[:, 1:][both]] - bank_t[idx[:, :-1][both]] | |
| assert np.allclose(d, 0.5, atol=1e-3) | |
| assert np.all(lag[:, -1][ok[:, -1]] < 0.5) | |
| def test_builder_rejects_future_row(): | |
| bank_t = np.array([9.5, 10.0, 10.5], np.float32) | |
| idx = tb.anchor_index_for(bank_t, {"r": 0}, np.array([0, 3]), np.array(["r", "r"]), np.array([10.2, 10.6])) | |
| # anchor 10.2 -> last slot 10.0 (row 1), not 10.5; anchor 10.6 -> 10.5 (row 2) | |
| assert idx[0, -1] == 1 and idx[1, -1] == 2 and idx[0, -2] == 0 | |
| assert bank_t[idx[0, -1]] <= 10.2 and bank_t[idx[1, -1]] <= 10.6 | |
| # a bank row that is 0.1 s after the anchor can never be selected for it | |
| bank_t2 = np.array([10.3], np.float32) | |
| idx2 = tb.anchor_index_for(bank_t2, {"r": 0}, np.array([0, 1]), np.array(["r"]), np.array([10.2])) | |
| assert (idx2 < 0).all() | |
| def test_env_attention_mask_is_causal(): | |
| torch = pytest.importorskip("torch") | |
| from models.branches import EnvBranch | |
| for tokens, n_slots in (("region", 16), ("region", 8), ("pooled", 16)): | |
| m = EnvBranch(tokens=tokens, n_slots=n_slots) | |
| R = m.n_reg | |
| mask = m.attn_mask.cpu().numpy() # True = blocked | |
| L = n_slots * R | |
| q_slot = np.arange(L) // R | |
| for q in range(L): | |
| for k in range(L): | |
| assert mask[q, k] == (q_slot[k] > q_slot[q]) | |
| assert not mask[-1].any(), "readout must see every slot" | |
| assert mask[:-1, -1].all(), "no token may attend to the readout" | |
| # ----------------------------------------------------------------------------- 2026-09-10 extension (latent-quality study, probes/) | |
| # Stored-feature shift tests: the features an anchor receives must not change when every window that ends AFTER the anchor is removed. | |
| # (f) token bank : re-gather from a bank truncated to rows with t_row <= T -> identical slots and identical stored PCA-64 tokens | |
| # (g) PCA streams: the stored pooled Vf row (s5 clips ending in [t-2, t]) is reproduced from the s5 library using ONLY clips with ts <= t, | |
| # and adding the first clip with ts > t changes it -> the pooled stream never reads a window that ends after t (on the s5 label axis; | |
| # the s5 label axis itself drifts from the frame/20 axis on some routes — quantified in results/tables/latent_audit.md A.2, not a rule error) | |
| # (h) future F : F(t+k) is the future_pca projection of the s5 clip ending at the last ts <= t+k, i.e. the window [t+k-2, t+k]; this documents | |
| # that F(t+1) overlaps the history by 1 s and F(t+2) touches it (relevant to how P2 restricts to non-overlapping horizons k >= 3) | |
| PROBES = CWM3 / "results" / "probes" | |
| have_probe_feats = (PROBES / "feats_manual_val.npz").exists() and (PROBES / "raw_manual_val.npz").exists() | |
| S5F = EXPLORE / "s5_features" / "vjepa2g384" / "front" | |
| def _bank_memmap(key): | |
| sys.path.insert(0, str(CWM3 / "probes")) | |
| import importlib.util as ilu | |
| spec = ilu.spec_from_file_location("probes_common", CWM3 / "probes" / "common.py") | |
| mod = ilu.module_from_spec(spec) | |
| spec.loader.exec_module(mod) | |
| return mod.npz_memmap(BANK, key) | |
| def test_stored_bank_tokens_invariant_to_removing_future_windows(): | |
| with np.load(BANK, allow_pickle=True) as z: | |
| bank_t, names, off = z["t"].astype(np.float64), [str(u) for u in z["uid_index"]], z["route_off"] | |
| tok = _bank_memmap("tok") | |
| idx_full = np.load(CWM3 / "results" / "banks" / "anchor_index_manual_val.npy") | |
| with np.load(S12R / "targets_manual_val.npz", allow_pickle=True) as z: | |
| uid, t = z["uid"], z["t"].astype(np.float64) | |
| rng = np.random.default_rng(0) | |
| route_of = {u: i for i, u in enumerate(names)} | |
| for i in rng.choice(len(t), 40, replace=False): | |
| r = route_of.get(uid[i]) | |
| if r is None: | |
| continue | |
| lo, hi = off[r], off[r + 1] | |
| keep = bank_t[lo:hi] <= t[i] + 1e-6 # remove every window ending after the anchor | |
| bt_trunc = bank_t[lo:hi][keep] | |
| idx_trunc = tb.anchor_index_for(bt_trunc, {uid[i]: 0}, np.array([0, len(bt_trunc)]), np.array([uid[i]]), np.array([t[i]]))[0] | |
| full = idx_full[i] | |
| okf, okt = full >= 0, idx_trunc >= 0 | |
| assert np.array_equal(okf, okt), "slot availability changed when future windows were removed" | |
| assert np.allclose(bank_t[full[okf]], bt_trunc[idx_trunc[okt]]), "slot times changed when future windows were removed" | |
| # the stored PCA-64 tokens of the gathered rows are, by construction, rows with t_row <= T: shifting the anchor later by 0.2 s | |
| # (T + 0.2, same 0.5 s grid cell when floor(2T) unchanged) gathers exactly the same rows | |
| if np.floor(2 * (t[i] + 0.2)) == np.floor(2 * t[i]): | |
| idx_shift = tb.anchor_index_for(bank_t, route_of, off, np.array([uid[i]]), np.array([t[i] + 0.2]))[0] | |
| assert np.array_equal(idx_shift, full) | |
| assert np.array_equal(np.asarray(tok[full[okf]]), np.asarray(tok[idx_shift[okf]])) | |
| def test_pca_stream_pool_uses_only_clips_ending_at_or_before_t(): | |
| fe = np.load(PROBES / "feats_manual_val.npz", allow_pickle=True) | |
| raw = np.load(PROBES / "raw_manual_val.npz", allow_pickle=True) | |
| Vf, uid, t = fe["Vf4224"], raw["uid"], raw["t"].astype(np.float64) | |
| have = np.flatnonzero(np.isfinite(Vf[:, 0])) | |
| rng = np.random.default_rng(0) | |
| n_checked = 0 | |
| for i in rng.choice(have, 120, replace=False): | |
| p = S5F / f"{uid[i]}.npz" | |
| if not p.exists(): | |
| continue | |
| with np.load(p) as z: | |
| ts, F = z["timestamps"].astype(np.float64), z["features"].astype(np.float32) | |
| lo = np.searchsorted(ts, t[i] - 2.0 - 1e-6) | |
| hi = np.searchsorted(ts, t[i] + 1e-6, side="right") | |
| if hi - lo < 3: | |
| continue | |
| assert ts[hi - 1] <= t[i] + 1e-6 | |
| W = F[lo:hi] | |
| rec = np.concatenate([W.mean(0), W[-1], W[-1] - W[0]]).astype(np.float16) | |
| assert np.array_equal(rec, Vf[i]), "stored pooled Vf != recomputation from clips ending <= t" | |
| if hi < len(ts): # adding the first clip that ends after t must change the pooled vector | |
| W2 = F[lo:hi + 1] | |
| rec2 = np.concatenate([W2.mean(0), W2[-1], W2[-1] - W2[0]]).astype(np.float16) | |
| assert not np.array_equal(rec2, Vf[i]) | |
| n_checked += 1 | |
| assert n_checked >= 50 | |
| def test_future_latent_window_ends_at_t_plus_k(): | |
| raw = np.load(PROBES / "raw_manual_val.npz", allow_pickle=True) | |
| fp = np.load(S12R / "future_pca.npz") | |
| mu, W = fp["mu"].astype(np.float32), fp["W"].astype(np.float32) | |
| uid, t, F, Fok = raw["uid"], raw["t"].astype(np.float64), raw["F"], raw["Fok"].astype(bool) | |
| rng = np.random.default_rng(1) | |
| n_checked, ends = 0, {k: [] for k in range(1, 6)} | |
| for i in rng.choice(len(t), 80, replace=False): | |
| p = S5F / f"{uid[i]}.npz" | |
| if not p.exists(): | |
| continue | |
| with np.load(p) as z: | |
| ts, Fs = z["timestamps"].astype(np.float64), z["features"].astype(np.float32) | |
| for k in range(1, 6): | |
| jj = np.searchsorted(ts, t[i] + k + 1e-6) - 1 | |
| ok = jj >= 0 and abs(ts[jj] - (t[i] + k)) <= 0.6 | |
| assert ok == bool(Fok[i, k - 1]) | |
| if ok: | |
| rec = ((Fs[jj] - mu) @ W).astype(np.float16) | |
| assert np.allclose(rec.astype(np.float32), F[i, k - 1].astype(np.float32), rtol=2e-2, atol=2e-2) | |
| ends[k].append(ts[jj] - t[i]) | |
| n_checked += 1 | |
| assert n_checked >= 100 | |
| # window END relative to t lies in (k-0.5, k]; window START = END - 2 -> k=1 overlaps [t-1, t], k=2 starts at >= t-0.5 | |
| for k in range(1, 6): | |
| e = np.array(ends[k]) | |
| assert np.all(e <= k + 1e-6) and np.all(e > k - 0.5 - 1e-6) | |
Xet Storage Details
- Size:
- 11.5 kB
- Xet hash:
- a4dd02d9dcc5c1f9ec9f83532b9bb76739b308befd8f86776dc93bad08d63e93
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.