← Walkthroughevals/neuroprobe/readout.py
Source snapshot · source SHA-256 d9119b6ba95d
1"""Frozen-feature linear ridge for the three Neuroprobe-Lite regimes.23Features use training-only standardization. Each fit selects regularization4on validation AUROC over the fixed 25-point grid, keeping the smallest lambda5on ties, and reports the separate test half. WS retains both fold scores.6"""78from __future__ import annotations910import argparse11import contextlib12import glob13import json14import multiprocessing as mp15import os16import pickle17import time1819import numpy as np20import torch2122UNKNOWN_REGION_ID = 742324_PH: dict = {}2526@contextlib.contextmanager27def _timed(name):28 t = time.perf_counter()29 try:30 yield31 finally:32 _PH[name] = _PH.get(name, 0.0) + (time.perf_counter() - t)3334def _phase_report(label) -> None:35 tot = sum(_PH.values()) or 1.036 print(f"[phase] {label} — wall budget by phase (total {tot / 60:.1f} min):", flush=True)37 for k, v in sorted(_PH.items(), key=lambda kv: -kv[1]):38 print(f"[phase] {k:22s} {v / 60:7.2f} min {100 * v / tot:5.1f}%", flush=True)3940BOARD_TASKS = (41 "onset", "speech", "volume", "delta_volume", "pitch", "word_index",42 "word_gap", "gpt2_surprisal", "word_head_pos", "word_part_speech",43 "word_length", "global_flow", "local_flow", "frame_brightness", "face_num",44)4546CS_TRAIN_ANCHOR = (2, 4)4748CS_TEST_CELLS = ((1, 1), (1, 2), (3, 0), (3, 1), (4, 0), (4, 1),49 (7, 0), (7, 1), (10, 0), (10, 1))5051LITE_SESSIONS = ((1, 1), (1, 2), (2, 0), (2, 4), (3, 0), (3, 1),52 (4, 0), (4, 1), (7, 0), (7, 1), (10, 0), (10, 1))5354ENCODERS = ("enc0", "enc3", "enc6", "enc9", "enc12") # region-mean (electrodes pooled at encode)5556ELEC_TAPS = ("enc0_elec", "enc12_elec") # per-electrode: enc0_elec = depth-0 parity FLOOR for enc12_elec5758ALL_TAPS = ELEC_TAPS + ENCODERS # universe for --taps validation5960WS_TAPS = ELEC_TAPS6162CS_TAPS = ENCODERS6364CSESSION_TAPS = ELEC_TAPS6566CSESSION_CELLS = LITE_SESSIONS # every Lite session is a cross-session test cell, sibling-trained6768def _sibling(cell):69 """The other Lite trial of the SAME subject — upstream cross_session's train trial70 (train_test_splits.py:146: the one other NEUROPROBE_LITE_SUBJECT_TRIALS entry for the subject)."""71 s, _ = cell72 sibs = [c for c in LITE_SESSIONS if c[0] == s and c != cell]73 assert len(sibs) == 1, f"subject {s} must have exactly one sibling Lite trial, got {sibs}"74 return sibs[0]7576_ELEC_LABELS_SIDECAR: dict | None = None7778LAM_MULTS = tuple(np.logspace(-4.0, 4.0, 25))7980def auroc(scores, labels) -> float:81 """Verbatim from online_probe.auroc. NaN if the eval half is single-class."""82 from sklearn.metrics import roc_auc_score8384 y = (np.asarray(labels) > 0).astype(int)85 if y.min() == y.max():86 return float("nan")87 return float(roc_auc_score(y, np.asarray(scores)))8889def _finite(y: np.ndarray, rows: np.ndarray) -> np.ndarray:90 r = np.asarray(rows, dtype=np.int64)91 return r[np.isfinite(y[r])]9293def _standardize(z_tr, others):94 """Per-feature z-score on TRAIN stats only (never fit on val/test). σ=0 → 1.9596 This is the canonical frozen-FM linear probe: the MAE/MoCo-v3/DINO lineage puts a97 BatchNorm WITHOUT affine before the linear head, and BN at eval uses statistics98 accumulated over probe TRAINING — i.e. train-set stats, which in CS means the ANCHOR's.99 """100 with _timed("standardize"):101 mu = z_tr.mean(axis=0)102 sd = z_tr.std(axis=0)103 sd[sd == 0] = 1.0104 return (z_tr - mu) / sd, [(z - mu) / sd for z in others]105106_STD_BLOCK = 1024107108def _standardize_inplace(z_tr, others, blk=_STD_BLOCK):109 """Apply training mean/std in place, blocking columns to bound temporary memory.110111 Reductions stay on axis 0. Avoid width-one blocks, whose NumPy summation112 order differs; tests check exact agreement with whole-array standardization.113 """114 with _timed("standardize"):115 d = z_tr.shape[1]116 edges = list(range(0, d, blk)) + [d]117 if len(edges) > 2 and edges[-1] - edges[-2] == 1:118 edges.pop(-2) # avoid a width-one trailing block119 for lo, hi in zip(edges, edges[1:]):120 b = z_tr[:, lo:hi]121 mu = b.mean(axis=0)122 sd = b.std(axis=0)123 sd[sd == 0] = 1.0124 b -= mu125 b /= sd126 for z in others:127 zb = z[:, lo:hi]128 zb -= mu129 zb /= sd130 return z_tr, others131132def _have(rec, tap) -> bool:133 return tap in rec["feats"]134135def _validate_taps(taps) -> None:136 for tap in taps:137 if tap not in ALL_TAPS:138 raise SystemExit(f"unknown tap {tap!r}; choose from {ALL_TAPS}")139140def _is_elec(tap) -> bool:141 return tap in ELEC_TAPS142143def _feat(rec, enc, rows, col_idx=None) -> np.ndarray:144 """Gather rows and aligned contact/region columns, then flatten to fp32."""145 with _timed("gather_fp16"):146 x = rec["feats"][enc]["raw"][np.asarray(rows, dtype=np.int64)]147 if col_idx is not None:148 x = x[:, np.asarray(col_idx, dtype=np.int64)]149 with _timed("to_fp32"):150 x = x.to(torch.float32).numpy()151 return x.reshape(x.shape[0], -1)152153def _linear_grams(z_tr, evals):154 """fp32 Gram products promoted to fp64 for the ridge solve."""155 with _timed("gram_gemm"):156 g = np.asarray(z_tr @ z_tr.T, dtype=np.float64) # fp32 GEMM → fp64 Gram157 with _timed("eval_kernels"):158 kern = {name: np.asarray(z @ z_tr.T, dtype=np.float64)159 for name, (z, _) in evals.items()}160 return g, kern161162def _lam_grid(z_tr, y_tr, evals):163 """Score the fixed lambda grid; use the primal only when feature width < train rows."""164 if len(y_tr) < 2:165 return {name: {m: float("nan") for m in LAM_MULTS} for name in evals}166 if z_tr.shape[1] < z_tr.shape[0]:167 return _lam_grid_primal(z_tr, y_tr, evals)168 g, kern = _linear_grams(z_tr, evals)169 n = g.shape[0]170 with _timed("eigh"):171 w, V = np.linalg.eigh(g) # G symmetric PSD ⇒ w >= 0172 c = V.T @ np.asarray(y_tr, dtype=np.float64)173 base = float(np.sum(w) / max(n, 1)) # trace(G)/n — the λ scale174 out: dict = {name: {} for name in evals}175 with _timed("lam_sweep"):176 for m in LAM_MULTS:177 alpha = V @ (c / (w + m * base))178 for name, (_, y) in evals.items():179 s = kern[name] @ alpha180 out[name][m] = (auroc(s, y) if len(y) >= 2 else float("nan"))181 return out182183def _lam_grid_primal(z_tr, y_tr, evals):184 """Equivalent ridge in feature space, retaining trace(Z.T @ Z) / n as the lambda scale."""185 with _timed("gram_gemm"):186 a_mat = np.asarray(z_tr.T @ z_tr, dtype=np.float64) # (d, d)187 n = z_tr.shape[0]188 with _timed("eigh"):189 w, V = np.linalg.eigh(a_mat)190 c = V.T @ (z_tr.T @ np.asarray(y_tr, dtype=np.float64))191 base = float(np.trace(a_mat) / max(n, 1)) # == sum(eig(G))/n, the dual's scale192 with _timed("lam_sweep"):193 lam = np.asarray(LAM_MULTS, dtype=np.float64) * base194 beta = V @ (c[:, None] / (w[:, None] + lam[None, :])) # (d, |LAM_MULTS|)195 out: dict = {}196 for name, (z, y) in evals.items():197 if len(y) < 2:198 out[name] = {m: float("nan") for m in LAM_MULTS}199 continue200 s = np.asarray(z, dtype=np.float64) @ beta # (n_e, |LAM_MULTS|)201 out[name] = {m: auroc(s[:, i], y) for i, m in enumerate(LAM_MULTS)}202 return out203204def _select_lam(d) -> dict:205 """Select maximum validation AUROC, keeping the first (smallest) tied lambda.206207 All-NaN validation returns NaN. Boundary flags describe grid position,208 not evidence that the optimum lies outside the grid."""209 finite = [(m, va) for m, va in d["val"].items() if not np.isnan(va)]210 if not finite:211 return {"val": float("nan"), "test": float("nan"), "lam_mult": float("nan"),212 "lam_pin": "", "lam_pinned": False, "n_tied": 0}213 best_val = max(va for _, va in finite)214 # Insertion order is LAM_MULTS ascending, so tied[0] is the smallest tied λ — exactly what the215 # strict-`>` loop this replaced selected. That equivalence is what keeps "argmax" byte-faithful.216 tied = [m for m, va in finite if va == best_val]217 m = tied[0]218 pin = "lo" if m == LAM_MULTS[0] else ("hi" if m == LAM_MULTS[-1] else "")219 out = {"val": best_val, "test": d["test"][m], "lam_mult": float(m),220 "lam_pin": pin, "lam_pinned": pin == "lo", "n_tied": len(tied)}221 return out222223def _cell_key(tap, norm) -> str:224 return f"{tap}|{norm}"225226def _grid_cells(grid) -> dict:227 """Keep a separately validation-selected result for every tap."""228 return {_cell_key(t, nm): _select_lam(d) for (t, nm), d in grid.items()}229230def _region_cols(anchor_rec, test_rec):231 """Anchor∩test region columns, aligned BY ATLAS ID (not by position).232233 The reserved 'unknown' id is NOT an anatomical location: two electrodes carrying it are234 not in the same place, so aligning subjects on it would be a free unearned column. It is235 a no-op on the Lite board (anchor S2T4 has no unmapped electrodes) but fires on any236 corpus that does, so it is excluded here rather than assumed away.237 """238 a_p = np.asarray(anchor_rec["present_parcels"], dtype=np.int64)239 t_p = np.asarray(test_rec["present_parcels"], dtype=np.int64)240 common = np.intersect1d(a_p, t_p)241 if UNKNOWN_REGION_ID in common:242 print(f"[check] dropping unknown region {UNKNOWN_REGION_ID} from intersection", flush=True)243 common = common[common != UNKNOWN_REGION_ID]244 if common.size == 0:245 return None, None, common246 a_idx = [int(np.where(a_p == c)[0][0]) for c in common]247 t_idx = [int(np.where(t_p == c)[0][0]) for c in common]248 return a_idx, t_idx, common249250def _run_norms(grid, enc, z_tr, z_va, z_te, y_tr, y_va, y_te):251 """Standardize in place using training statistics, then score linear ridge."""252 a, (b, c) = _standardize_inplace(z_tr, [z_va, z_te])253 grid[(enc, "std")] = _lam_grid(a, y_tr, {"val": (b, y_va), "test": (c, y_te)})254255def _ws_cell(rec, task, taps) -> dict:256 """Within-session: board KFold(2). Per fold fit train, λ-select on the val half, report the257 test half; average the two folds' test AUROCs, per (tap, norm)."""258 y = np.asarray(rec["labels"][task], dtype=np.float64)259 folds = []260 fold_ids = []261 for _fold, sp in sorted(rec["ws_split"][task].items()):262 tr, va, te = (_finite(y, sp["train"]), _finite(y, sp["val"]), _finite(y, sp["test"]))263 if len(tr) < 2 or len(te) < 2:264 continue265 grid = {}266 for enc in taps:267 if not _have(rec, enc):268 continue269 z_tr = _feat(rec, enc, tr)270 z_va, z_te = _feat(rec, enc, va), _feat(rec, enc, te)271 _run_norms(grid, enc, z_tr, z_va, z_te, y[tr], y[va], y[te])272 if grid:273 folds.append(_grid_cells(grid))274 fold_ids.append(int(_fold))275 if not folds:276 return {"cells": {}}277 keys = sorted({k for f in folds for k in f})278 out = {}279 for k in keys:280 vals = [f[k]["test"] for f in folds if k in f]281 out[k] = {"test": float(np.nanmean(vals)) if vals else float("nan"),282 "folds": [{"fold_idx": idx, "test_roc_auc": float(f[k]["test"])}283 for idx, f in zip(fold_ids, folds) if k in f],284 "lam_pinned": bool(any(f[k]["lam_pinned"] for f in folds if k in f)),285 "lam_sat": bool(any(f[k].get("lam_pin") == "hi" for f in folds if k in f)),286 "lam_mult": [f[k]["lam_mult"] for f in folds if k in f]}287 return {"cells": out}288289def _cs_cell(anchor_rec, test_rec, task, taps) -> dict:290 """Cross-subject: fit the anchor's finite rows, λ-select on the test cell's val half, report291 its test half. Features are the anchor∩test region intersection (atlas-id aligned)."""292 y_a = np.asarray(anchor_rec["labels"][task], dtype=np.float64)293 y_t = np.asarray(test_rec["labels"][task], dtype=np.float64)294 tr = _finite(y_a, np.arange(len(y_a)))295 va = _finite(y_t, test_rec["cs_split"][task]["val"])296 te = _finite(y_t, test_rec["cs_split"][task]["test"])297 if len(tr) < 2 or len(te) < 2:298 return {"cells": {}}299 a_idx, t_idx, common = _region_cols(anchor_rec, test_rec)300 if common.size == 0:301 return {"cells": {}}302 grid: dict = {}303 for enc in taps:304 if not (_have(anchor_rec, enc) and _have(test_rec, enc)):305 continue306 z_tr = _feat(anchor_rec, enc, tr, a_idx)307 z_va, z_te = _feat(test_rec, enc, va, t_idx), _feat(test_rec, enc, te, t_idx)308 _run_norms(grid, enc, z_tr, z_va, z_te, y_a[tr], y_t[va], y_t[te])309 if not grid:310 return {"cells": {}}311 return {"cells": _grid_cells(grid), "n_parcels": int(common.size)}312313def _elec_cols(train_rec, test_rec):314 """Shared electrodes between two SAME-SUBJECT sessions, aligned BY LABEL (elec_labels).315316 Sibling Lite trials drop DIFFERENT bad channels (measured: 4/6 subjects differ, e.g. subj-3317 100 vs 102), so the per-electrode axis is NOT positionally aligned across sessions —318 intersecting by identity is the only correct alignment. Returns (train_idx, test_idx,319 n_shared); (None, None, 0) if either cache lacks elec_labels (pre-edit cache) or no overlap."""320 a = train_rec.get("elec_labels")321 t = test_rec.get("elec_labels")322 if a is None or t is None:323 return None, None, 0324 a = np.asarray(a)325 t = np.asarray(t)326 common = np.intersect1d(a, t)327 if common.size == 0:328 return None, None, 0329 a_idx = [int(np.where(a == c)[0][0]) for c in common]330 t_idx = [int(np.where(t == c)[0][0]) for c in common]331 return a_idx, t_idx, int(common.size)332333def _csession_cell(train_rec, test_rec, task, taps) -> dict:334 """Cross-session: train on the SIBLING trial of the SAME subject, λ-select on this cell's val335 half, report its test half.336337 Upstream ``generate_splits_cross_session`` halves the test session IDENTICALLY to338 ``generate_splits_cross_subject`` (train_test_splits.py:153-156 == 66-69: val=range(size//2),339 test=range(size//2,size)), so ``cs_split``'s val/test are reused VERBATIM — only the train340 anchor differs (the sibling trial, not S2T4). Region taps align by atlas id (``_region_cols``);341 per-electrode taps align by electrode IDENTITY (``_elec_cols``), on the shared-electrode subset.342 """343 y_a = np.asarray(train_rec["labels"][task], dtype=np.float64)344 y_t = np.asarray(test_rec["labels"][task], dtype=np.float64)345 tr = _finite(y_a, np.arange(len(y_a)))346 va = _finite(y_t, test_rec["cs_split"][task]["val"])347 te = _finite(y_t, test_rec["cs_split"][task]["test"])348 if len(tr) < 2 or len(te) < 2:349 return {"cells": {}}350 p_a, p_t, p_common = _region_cols(train_rec, test_rec)351 e_a, e_t, n_elec = _elec_cols(train_rec, test_rec)352 grid: dict = {}353 for enc in taps:354 if not (_have(train_rec, enc) and _have(test_rec, enc)):355 continue356 if _is_elec(enc):357 if e_a is None:358 continue359 col_a, col_t = e_a, e_t360 else:361 if p_a is None:362 continue363 col_a, col_t = p_a, p_t364 z_tr = _feat(train_rec, enc, tr, col_a)365 z_va, z_te = _feat(test_rec, enc, va, col_t), _feat(test_rec, enc, te, col_t)366 _run_norms(grid, enc, z_tr, z_va, z_te, y_a[tr], y_t[va], y_t[te])367 if not grid:368 return {"cells": {}}369 return {"cells": _grid_cells(grid),370 "n_parcels": int(p_common.size) if p_a is not None else 0,371 "n_elec": n_elec}372373MMAP_DEFAULT = {"ws": False, "cs": False, "csession": False}374375def _load(cache_dir, session, tag, mmap=False):376 """Load a session cache. ``mmap=True`` defers the read into the gathers (see MMAP_DEFAULT).377378 Pages arrive only where a tensor is actually indexed, so selectivity is free: a CS shard379 never gathers enc12_elec and therefore never reads those 34 GB. No tap-filter argument is380 needed — not touching IS not loading. But laziness is not a speedup, and for WS it is a381 slowdown; pick with MMAP_DEFAULT and A/B with --mmap/--no-mmap before trusting a change.382 """383 s, t = session384 rec = torch.load(f"{cache_dir}/enc_s{s}_t{t}_{tag}.pt", map_location="cpu",385 weights_only=False, mmap=mmap)386 if _ELEC_LABELS_SIDECAR is not None and rec.get("elec_labels") is None:387 lab = _ELEC_LABELS_SIDECAR.get(f"s{s}_t{t}")388 if lab is not None:389 lab = np.asarray(lab)390 if lab.ndim != 1 or len(np.unique(lab)) != len(lab):391 raise ValueError("sidecar contact labels must be a unique one-dimensional sequence")392 if lab.shape[0] != rec["feats"]["enc12_elec"]["raw"].shape[1]:393 raise ValueError(394 f"sidecar labels ({lab.shape[0]}) != enc12_elec electrodes "395 f"({rec['feats']['enc12_elec']['raw'].shape[1]}) for s{s}_t{t}")396 rec["elec_labels"] = lab397 return rec398399_SHARED: dict = {}400401def _task_worker(task):402 """Runs in a forked child. Reads the cache from _SHARED — NEVER take it as an argument:403 Pool pickles arguments, which would serialize a multi-GB cache per task."""404 fn, taps = _SHARED["fn"], _SHARED["taps"]405 return task, fn(task, taps)406407def _map_tasks(fn, taps, workers) -> dict:408 """{task: cell} over BOARD_TASKS, optionally across forked workers."""409 if workers <= 1:410 return {task: fn(task, taps) for task in BOARD_TASKS}411 _SHARED["fn"], _SHARED["taps"] = fn, taps # set BEFORE fork ⇒ inherited, not pickled412 with mp.get_context("fork").Pool(workers) as pool:413 return dict(pool.map(_task_worker, BOARD_TASKS))414415def _ws_shard(cache_dir, tag, session, taps=WS_TAPS, workers=1,416 mmap=MMAP_DEFAULT["ws"]) -> dict:417 rec = _load(cache_dir, session, tag, mmap=mmap)418 out = _map_tasks(lambda task, tp: _ws_cell(rec, task, tp), taps, workers)419 return {"kind": "ws", "name": f"S{session[0]}T{session[1]}",420 "cells": {f"{tag}|{k}": v for k, v in out.items()}}421422def _cs_shard(cache_dir, tag, cell, taps=CS_TAPS, workers=1,423 mmap=MMAP_DEFAULT["cs"]) -> dict:424 taps = tuple(t for t in taps if t not in ELEC_TAPS) # CS is region-bridged by necessity425 anchor_rec = _load(cache_dir, CS_TRAIN_ANCHOR, tag, mmap=mmap)426 test_rec = _load(cache_dir, cell, tag, mmap=mmap)427 out = _map_tasks(lambda task, tp: _cs_cell(anchor_rec, test_rec, task, tp), taps, workers)428 return {"kind": "cs", "name": f"S{cell[0]}T{cell[1]}",429 "cells": {f"{tag}|{k}": v for k, v in out.items()}}430431def _csession_shard(cache_dir, tag, cell, taps=CSESSION_TAPS, workers=1,432 mmap=MMAP_DEFAULT["csession"]) -> dict:433 """Cross-session cell: train on the sibling trial (same subject), test on this session's434 held-out half. Keeps the per-electrode taps (electrode identity IS shared within subject)."""435 train_rec = _load(cache_dir, _sibling(cell), tag, mmap=mmap)436 test_rec = _load(cache_dir, cell, tag, mmap=mmap)437 out = _map_tasks(lambda task, tp: _csession_cell(train_rec, test_rec, task, tp), taps, workers)438 return {"kind": "csession", "name": f"S{cell[0]}T{cell[1]}",439 "cells": {f"{tag}|{k}": v for k, v in out.items()}}440441def _blank(tags) -> dict:442 return {f"{tag}|{t}": {"ws": {}, "cs": {}, "csession": {}, "pinned": {}, "sat": {},443 "n_parcels": {}, "n_elec": {}}444 for tag in tags for t in BOARD_TASKS}445446def _absorb(res, sh) -> None:447 """Fold one shard in. res[task][kind]["tap|norm"][cell_name] = test AUROC — the full grid,448 every entry populated over every cell. No axis is collapsed at merge time either."""449 kind = sh["kind"]450 for k, val in sh["cells"].items():451 for gk, s in (val.get("cells") or {}).items():452 res[k][kind].setdefault(gk, {})[sh["name"]] = s["test"]453 if s.get("lam_pinned"):454 res[k]["pinned"].setdefault(f"{kind}:{gk}", []).append(sh["name"])455 if s.get("lam_sat"):456 res[k]["sat"].setdefault(f"{kind}:{gk}", []).append(sh["name"])457 if val.get("n_parcels") is not None:458 res[k]["n_parcels"][sh["name"]] = val["n_parcels"]459 if val.get("n_elec") is not None:460 res[k]["n_elec"][sh["name"]] = val["n_elec"]461462def _merge(tags, shard_dir) -> dict:463 res = _blank(tags)464 for kind in ("ws", "cs", "csession"):465 for path in sorted(glob.glob(f"{shard_dir}/{kind}_*.json")):466 with open(path) as f:467 _absorb(res, json.load(f))468 return _finalize(res)469470def _finalize(res: dict) -> dict:471 """Cohort-mean each grid entry over its cells (12 WS sessions / 10 CS cells / 12 cross-session)."""472 for c in res.values():473 for kind in ("ws", "cs", "csession"):474 c[f"{kind}_mean"] = {gk: float(np.nanmean(list(d.values())))475 for gk, d in c[kind].items() if d}476 return res477478def _compute_all(cache_dir, tags, ws_taps=WS_TAPS, cs_taps=CS_TAPS) -> dict:479 res = _blank(tags)480 for tag in tags:481 for session in LITE_SESSIONS:482 sh = _ws_shard(cache_dir, tag, session, ws_taps)483 _absorb(res, sh)484 print(f"[{tag}] WS done {sh['name']}", flush=True)485 for cell in CS_TEST_CELLS:486 sh = _cs_shard(cache_dir, tag, cell, cs_taps)487 _absorb(res, sh)488 print(f"[{tag}] CS done {sh['name']}", flush=True)489 for cell in CSESSION_CELLS:490 sh = _csession_shard(cache_dir, tag, cell, CSESSION_TAPS)491 _absorb(res, sh)492 print(f"[{tag}] CSession done {sh['name']}", flush=True)493 return _finalize(res)494495def _macro(res, tag, kind, gk) -> float:496 """Macro over the 15 board tasks of one grid entry's cohort mean."""497 v = [res[f"{tag}|{t}"].get(f"{kind}_mean", {}).get(gk, np.nan) for t in BOARD_TASKS]498 return float(np.nanmean(v)) if not all(np.isnan(x) for x in v) else float("nan")499500def _report(tags, res) -> None:501 """Print each tap’s per-task and macro test AUROC without selecting across taps."""502 for tag in tags:503 for kind, label in (("cs", "CS (anchor S2T4 → 10 cells)"),504 ("csession", "CSession (12 cells, sibling-trained)"),505 ("ws", "WS (12 sessions)")):506 gks = sorted({g for t in BOARD_TASKS for g in res[f"{tag}|{t}"].get(kind, {})},507 key=lambda g: (g.split("|")[1], g.split("|")[0]))508 if not gks:509 continue510 print(f"\n=== {label} test-half AUROC — tag={tag} ===", flush=True)511 print(f" {'task':18s}" + "".join(f"{g:>18s}" for g in gks), flush=True)512 for t in BOARD_TASKS:513 m = res[f"{tag}|{t}"].get(f"{kind}_mean", {})514 print(f" {t:18s}" + "".join(f"{m.get(g, float('nan')):18.4f}" for g in gks),515 flush=True)516 print(f" {'MACRO(15)':18s}"517 + "".join(f"{_macro(res, tag, kind, g):18.4f}" for g in gks), flush=True)518519def main() -> None:520 p = argparse.ArgumentParser()521 p.add_argument("--cache-dir", required=True)522 p.add_argument("--tags", required=True)523 p.add_argument("--out", required=True)524 p.add_argument("--mode", choices=("all", "ws", "cs", "csession", "merge"), default="all")525 p.add_argument("--index", type=int, help="shard index (mode=ws|cs|csession)")526 p.add_argument("--shard-dir")527 p.add_argument("--workers", type=int, default=1,528 help="fork this many task-workers per shard (memory-bound: WS ~3, CS ~8). "529 "Set OMP_NUM_THREADS = cpus-per-task / workers.")530 p.add_argument("--mmap", dest="mmap", default=None, action="store_true",531 help=f"force lazy paging (default per mode: {MMAP_DEFAULT}). Measured: it "532 f"DEFERS the read into the gathers at ~1/4 sequential bandwidth — a "533 f"win for CS (skips the 34 GB elec tap) and a loss for WS.")534 p.add_argument("--no-mmap", dest="mmap", action="store_false",535 help="force one eager sequential read of the whole cache.")536 p.add_argument("--taps", default="",537 help=f"comma-separated subset of {ALL_TAPS} (default: per-regime; CS drops "538 f"{ELEC_TAPS} automatically)")539 p.add_argument("--elec-labels-sidecar",540 help="pickle {'s{S}_t{T}': labels} to attach to records that lack elec_labels "541 "(caches encoded before the field was stored, e.g. arm0/r4b). Validated "542 "same-set upstream; _load re-asserts count-match before attaching.")543 args = p.parse_args()544545 global _ELEC_LABELS_SIDECAR546 if args.elec_labels_sidecar:547 with open(args.elec_labels_sidecar, "rb") as fh:548 _ELEC_LABELS_SIDECAR = pickle.load(fh)549 print(f"[sidecar] attaching elec_labels for {len(_ELEC_LABELS_SIDECAR)} sessions",550 flush=True)551552 tags = tuple(t.strip() for t in args.tags.split(","))553 # Default taps are per-REGIME (WS/CSession electrode-only, CS region-only); --taps overrides.554 _mode_taps = {"ws": WS_TAPS, "cs": CS_TAPS, "csession": CSESSION_TAPS}555 taps = (tuple(t.strip() for t in args.taps.split(",") if t.strip())556 or _mode_taps.get(args.mode, ALL_TAPS))557 _validate_taps(taps)558559 if args.mode in ("ws", "cs", "csession"):560 cells = {"ws": LITE_SESSIONS, "cs": CS_TEST_CELLS, "csession": CSESSION_CELLS}[args.mode]561 cell = cells[args.index]562 fn = {"ws": _ws_shard, "cs": _cs_shard, "csession": _csession_shard}[args.mode]563 t0 = time.perf_counter()564 use_mmap = MMAP_DEFAULT[args.mode] if args.mmap is None else args.mmap565 sh = fn(args.cache_dir, tags[0], cell, taps, workers=args.workers, mmap=use_mmap)566 os.makedirs(args.shard_dir, exist_ok=True)567 out = f"{args.shard_dir}/{args.mode}_{sh['name']}.json"568 with open(out, "w") as f:569 json.dump(sh, f, indent=2)570 print(f"wrote {out}", flush=True)571 # Phase totals are per-PROCESS: with --workers>1 the children's timers die with them,572 # so this table is the parent's view (load + merge) only. Profile with --workers 1.573 _phase_report(f"{args.mode} {sh['name']} workers={args.workers} "574 f"mmap={use_mmap} wall={(time.perf_counter() - t0) / 60:.1f} min")575 return576577 res = _merge(tags, args.shard_dir) if args.mode == "merge" else _compute_all(578 args.cache_dir, tags,579 tuple(t for t in taps if t in ELEC_TAPS),580 tuple(t for t in taps if t not in ELEC_TAPS))581 _report(tags, res)582 with open(args.out, "w") as f:583 json.dump(res, f, indent=2)584 print(f"\nwrote {args.out}\nMERGE_DONE", flush=True)585586if __name__ == "__main__":587 main()