Buckets:
| """cwm3/probes/vjepa2_temporal_probe.py — P1b probes on the re-extracted V-JEPA2 windows (results/probes/vjepa2_temporal/): | |
| does the WITHIN-WINDOW time dimension (8 tubelet groups) carry road information that the cache's time-mean discards? | |
| Representations of the same window (n <= 2,000; small -> 5-fold device-grouped CV, ridge with an inner alpha grid selected on the training | |
| folds' out-of-fold AP, out-of-fold scores pooled, device-cluster bootstrap CI): | |
| mean1408 time-mean, spatial-mean clip token (= the cache's `feats` / raw1408) | |
| region_pca time-mean 16x1408 (= the cache's `region`) -> PCA-k on the training folds | |
| timegroups_pca 8 time groups x spatial mean (8x1408) -> PCA-k | |
| full_pca 8x16x1408 -> PCA-k | |
| motion1408 last time group - first time group (spatial mean) [pure within-window change] | |
| tstd1408 std over the 8 time groups (spatial mean) [within-window variability] | |
| mean+motion concat(mean1408, motion1408) | |
| Targets: current road pressure (P0 per family + any), future onsets within 3 s (P per family + any, brake/gas, hand_leave/head_turn), vEgo bins. | |
| -> results/tables/latent_vjepa2_temporal.md. No openpilot DMS field is read. | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import os | |
| import sys | |
| import time | |
| from pathlib import Path | |
| os.environ.setdefault("OMP_NUM_THREADS", "4") | |
| import numpy as np # noqa: E402 | |
| sys.path.insert(0, str(Path(__file__).resolve().parent)) | |
| import common as C # noqa: E402 | |
| OUT = C.OUT / "vjepa2_temporal" | |
| K_PCA = 128 | |
| ALPHAS = C.RIDGE_ALPHAS | |
| def load(): | |
| z = np.load(OUT / "windows.npz", allow_pickle=True) | |
| tt = np.asarray(np.load(str(z["tt_file"]), mmap_mode="r")) # f16 [N,8,16,1408] in RAM (0.7 GB): random memmap reads were IO-bound | |
| done = z["done"].astype(bool) | |
| meta = {k: z[k] for k in ("regime", "split", "row", "uid", "dongle", "t_anchor", "t_row", "enriched")} | |
| return tt, z["g"].astype(np.float32), z["cos_check"], done, meta | |
| def targets_for(meta, done): | |
| """re-derive the probe targets of every selected window from its raw_{regime}_{split} row.""" | |
| n = len(meta["row"]) | |
| T = {} | |
| for regime in C.REGIMES: | |
| for split in C.SPLITS: | |
| m = (meta["regime"] == regime) & (meta["split"] == split) & done | |
| if not m.any(): | |
| continue | |
| raw = C.load_raw(regime, split) | |
| rows = meta["row"][m].astype(int) | |
| sub = {k: (v[rows] if isinstance(v, np.ndarray) and v.ndim >= 1 and len(v) == len(raw["t"]) else v) for k, v in raw.items()} | |
| Tt = C.build_targets(sub) | |
| for k, v in Tt.items(): | |
| if k not in T: | |
| T[k] = {"y": np.zeros(n, bool), "m": np.zeros(n, bool), "kind": v["kind"], "family": v["family"]} | |
| T[k]["y"][m] = v["y"] | |
| T[k]["m"][m] = v["m"] | |
| # pooled "any" targets | |
| for pref, kind in (("cur_P0_", "cur"), ("fut_P_", "fut")): | |
| ks = [k for k in T if k.startswith(pref)] | |
| y = np.zeros(n, bool); m = np.zeros(n, bool) | |
| for k in ks: | |
| y |= T[k]["y"] & T[k]["m"]; m |= T[k]["m"] | |
| T[pref + "any"] = {"y": y, "m": m, "kind": kind, "family": "pressure"} | |
| return T | |
| def pca_fit(X, k): | |
| import torch | |
| dev = C.torch_device() | |
| Xt = torch.as_tensor(np.asarray(X, np.float32), device=dev) | |
| mu = Xt.mean(0) | |
| U, S, V = torch.svd_lowrank(Xt - mu, q=min(k + 16, Xt.shape[0] - 1, Xt.shape[1] - 1), niter=4) | |
| return mu.cpu().numpy(), V[:, :k].cpu().numpy() | |
| _PCA_CACHE = {} | |
| def device_folds(groups, k_folds=5, seed=0): | |
| """one seeded device -> fold assignment shared by every target and representation (so in-fold PCAs can be cached).""" | |
| uniq = np.unique(groups) | |
| rng = np.random.default_rng(seed) | |
| return {d: i % k_folds for i, d in enumerate(rng.permutation(uniq))} | |
| def cv_ridge(X_fn, y, m, groups, k_folds=5, seed=0, pca_k=None, fold_of_dev=None, sel=None, cache_key=None): | |
| """X_fn(rows) -> matrix for those rows. PCA (if pca_k) is fitted on ALL windows of the regime set `sel` whose device is outside the test | |
| fold (never on the test fold) and cached per (cache_key, fold). -> out-of-fold scores [n] (NaN outside m).""" | |
| import torch | |
| rows = np.flatnonzero(m) | |
| if fold_of_dev is None: | |
| fold_of_dev = device_folds(groups[rows], k_folds, seed) | |
| folds = np.array([fold_of_dev[d] for d in groups[rows]]) | |
| all_folds = np.array([fold_of_dev[d] for d in groups]) | |
| scores = np.full(len(y), np.nan, np.float32) | |
| dev = C.torch_device() | |
| for f in range(k_folds): | |
| tr, te = rows[folds != f], rows[folds == f] | |
| if len(te) == 0 or y[tr].sum() < 5 or (~y[tr]).sum() < 5: | |
| continue | |
| Xtr, Xte = X_fn(tr), X_fn(te) | |
| if pca_k is not None and Xtr.shape[1] > pca_k: | |
| key = (cache_key, f) | |
| if key not in _PCA_CACHE: | |
| fit_rows = np.flatnonzero((sel if sel is not None else m) & (all_folds != f)) | |
| _PCA_CACHE[key] = pca_fit(X_fn(fit_rows), pca_k) | |
| mu, W = _PCA_CACHE[key] | |
| Xtr, Xte = (Xtr - mu) @ W, (Xte - mu) @ W | |
| std = C.Std(Xtr) | |
| Ztr, Zte = std(Xtr), std(Xte) | |
| # dual (kernel) ridge: n <= 1.6k rows << d -> one n x n eigendecomposition per fold serves every alpha (identical solution to the primal | |
| # closed form of common.ridge_fit_predict: lambda = alpha * trace(G)/d, centred labels, bias column) | |
| best_a, best_ap = 0, -1 | |
| inner = np.array([fold_of_dev[d] % 3 for d in groups[tr]]) | |
| aps = np.full(len(ALPHAS), np.nan) | |
| for g in range(3): | |
| itr, ite = inner != g, inner == g | |
| if y[tr][itr].sum() < 3 or y[tr][ite].sum() < 3 or ite.sum() < 5: | |
| continue | |
| S = ridge_dual_scores(Ztr[itr], y[tr][itr].astype(np.float32), Ztr[ite], ALPHAS, dev) | |
| a_ = np.array([C.ap_weighted(y[tr][ite], S[a]) for a in range(len(ALPHAS))]) | |
| aps = np.nanmean(np.stack([aps, a_]), 0) if np.isfinite(aps).any() else a_ | |
| best_a = int(np.nanargmax(aps)) if np.isfinite(aps).any() else 2 | |
| scores[te] = ridge_dual_scores(Ztr, y[tr].astype(np.float32), Zte, (ALPHAS[best_a],), dev)[0] | |
| return scores | |
| def ridge_dual_scores(Ztr, ytr, Zev, alphas, dev): | |
| """kernel form of the closed-form ridge in common.ridge_fit_predict -> scores [len(alphas), n_ev].""" | |
| import torch | |
| Xt = torch.as_tensor(Ztr, device=dev) | |
| Xe = torch.as_tensor(Zev, device=dev) | |
| n, d = Xt.shape | |
| Xb = torch.cat([Xt, torch.ones(n, 1, device=dev)], 1) | |
| Xeb = torch.cat([Xe, torch.ones(len(Xe), 1, device=dev)], 1) | |
| yv = torch.as_tensor(ytr, device=dev) | |
| yv = yv - yv.mean() | |
| K = (Xb @ Xb.T).double() | |
| tr_g = torch.trace(K) / d # trace(X^T X) = trace(X X^T) | |
| evals, V = torch.linalg.eigh(K) | |
| Vy = V.T @ yv.double() | |
| Kev = (Xeb @ Xb.T).double() # [n_ev, n] | |
| out = [] | |
| for al in alphas: | |
| dual = V @ (Vy / (evals + al * tr_g)) | |
| out.append((Kev @ dual).float().cpu().numpy()) | |
| return np.stack(out) | |
| def main(): | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--n-boot", type=int, default=500) | |
| ap.add_argument("--pca-k", type=int, default=K_PCA) | |
| ap.add_argument("--max-minutes", type=float, default=0, help="exit 3 (resumable) once this wall budget is used; 0 = no limit") | |
| a = ap.parse_args() | |
| t0 = time.time() | |
| tt, g, cos_check, done, meta = load() | |
| T = targets_for(meta, done) | |
| n = len(done) | |
| print(f"windows {done.sum()}/{n}; cos_check median {np.nanmedian(cos_check[done]):.4f} p01 {np.nanpercentile(cos_check[done], 1):.4f}") | |
| mean1408 = g.mean(1) | |
| motion = g[:, -1] - g[:, 0] | |
| tstd = g.std(1) | |
| tg = g.reshape(n, -1) | |
| reps = {"mean1408": (lambda r: mean1408[r], None), "region_pca": (lambda r: np.asarray(tt[r]).astype(np.float32).mean(1).reshape(len(r), -1), a.pca_k), | |
| "timegroups_pca": (lambda r: tg[r], a.pca_k), "full_pca": (lambda r: np.asarray(tt[r]).astype(np.float32).reshape(len(r), -1), a.pca_k), | |
| "motion1408": (lambda r: motion[r], None), "tstd1408": (lambda r: tstd[r], None), "mean+motion": (lambda r: np.concatenate([mean1408[r], motion[r]], 1), None)} | |
| import json | |
| cells_p = OUT / "probe_cells.jsonl" # incremental, resumable (the host has killed long jobs for memory) | |
| rows_out = [json.loads(l) for l in cells_p.read_text().splitlines() if l.strip()] if cells_p.exists() else [] | |
| have = {(c["regime"], c["target"]) for c in rows_out} | |
| fold_of_dev = device_folds(meta["dongle"][done]) | |
| for regime in ("both", "manual", "engaged"): | |
| sel = done if regime == "both" else (done & (meta["regime"] == regime)) | |
| _PCA_CACHE.clear() | |
| for tn, tv in T.items(): | |
| m = tv["m"] & sel | |
| if m.sum() < 60 or tv["y"][m].sum() < 15 or (~tv["y"][m]).sum() < 15: | |
| continue | |
| if (regime, tn) in have: | |
| continue | |
| if a.max_minutes and (time.time() - t0) > a.max_minutes * 60: | |
| print(f"time budget reached after {len(rows_out)} cells; exit 3 (resumable)", flush=True) | |
| sys.exit(3) | |
| cell = {"regime": regime, "target": tn, "kind": tv["kind"], "n": int(m.sum()), "n_pos": int(tv["y"][m].sum())} | |
| for rn, (fn, k) in reps.items(): | |
| s = cv_ridge(fn, tv["y"], m, meta["dongle"], pca_k=k, fold_of_dev=fold_of_dev, sel=sel, cache_key=(rn, regime)) | |
| ok = np.isfinite(s) | |
| if ok.sum() < 30 or tv["y"][ok].sum() < 5: | |
| cell[rn] = None | |
| continue | |
| r = C.cluster_boot(tv["y"][ok], s[ok], meta["dongle"][ok], n_boot=a.n_boot) | |
| cell[rn] = (round(r["ap"], 4), round(r["ci"][0], 4), round(r["ci"][1], 4)) | |
| rows_out.append(cell) | |
| with open(cells_p, "a") as fh: | |
| fh.write(json.dumps(cell) + "\n") | |
| print(regime, tn, {k: (v[0] if isinstance(v, tuple) else v) for k, v in cell.items() if k in reps}, flush=True) | |
| C.write_json(OUT / "probe_latest.json", {"cells": rows_out, "n_windows": int(done.sum()), "pca_k": a.pca_k, "cos_check_median": float(np.nanmedian(cos_check[done])), "created_utc": C.stamp()}) | |
| L = ["# P1b — within-window time dimension of V-JEPA2 (re-extracted windows)", "", | |
| f"Generated {C.stamp()} by `probes/vjepa2_temporal_probe.py` on {int(done.sum())} re-extracted 2 s windows (`results/probes/vjepa2_temporal/`, selection balanced over " | |
| f"regime x split, device-capped, half enriched with road pressure). Time-mean of the new tensor vs the cached `region`: cos median {np.nanmedian(cos_check[done]):.4f}, " | |
| f"p01 {np.nanpercentile(cos_check[done], 1):.4f} (extraction reproduces the cache). Probe: ridge, 5-fold device-grouped CV with inner alpha selection, PCA-{a.pca_k} fitted " | |
| "inside the training folds; PR-AUC of pooled out-of-fold scores with device-cluster bootstrap 95 % CI. Small n -> wide CIs; read the ordering, not the third decimal.", "", | |
| C.md_table(["regime", "target", "n / pos"] + list(reps), [[c["regime"], c["target"], f"{c['n']} / {c['n_pos']}"] + [(f"{v[0]:.3f} [{v[1]:.3f}, {v[2]:.3f}]" if v else "–") for v in (c.get(rn) for rn in reps)] for c in rows_out])] | |
| (C.TABLES / "latent_vjepa2_temporal.md").write_text("\n".join(L) + "\n") | |
| print(f"-> {C.TABLES / 'latent_vjepa2_temporal.md'} ({time.time() - t0:.0f}s)") | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 11.6 kB
- Xet hash:
- dc21422fa74771225cd450525eb6e97063cd2b3162fa7e24f03833b6042454ec
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.