Buckets:
| """cwm3/probes/gather_features.py — materialise every FROZEN representation for the probe rows (read-only w.r.t. all caches). | |
| Per (regime, split) -> results/probes/raw_{r}_{s}.npz compact target/contract arrays of the probe rows (data/cwm3_dataset.load_split | |
| + results/targets_pressure, uid/t witness asserted) | |
| -> results/probes/feats_{r}_{s}.npz frozen visual representations, row-aligned with raw_*: | |
| raw1408 f16[n,1408] frontfeat clip-mean token of the LAST bank slot (window [tau-2, tau], tau = floor(2t)/2 <= t); NaN if no slot | |
| region_file -> results/probes/region_{r}_{s}.npy f16[n_reg,16,1408] (np.load mmap_mode="r"): frontfeat 4x4 region tokens of the same | |
| window for the region-subset rows only (`region_rows` bool[n], `region_filled` bool[n_reg]) | |
| tok16 f16[n,16,4,64] token bank (PCA-64, 2x2 pooled) for the 16 causal slots (zero where missing), tok_ok bool[n,16] | |
| pooled16 f16[n,16,256] bank pooled256 (PCA-256 of the clip-mean token) for the 16 slots | |
| bank_t16 f32[n,16] bank clip END times (NaN where missing) | |
| Vf4224 / Vc4224 / crop4224 f16 s12 pooled streams [mean, last, trend] over the <=5 s5 clips ending in [t-2, t] (NaN if row_ok False) | |
| Phases are independent and resumable: --phase raw | bank | frontfeat | vis | all. OMP threads <= 4 (shared host). | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import os | |
| import sys | |
| import time | |
| from multiprocessing import Pool | |
| 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 | |
| LOG = C.OUT / "gather_features.log" | |
| SETS = [(r, s) for r in C.REGIMES for s in C.SPLITS] | |
| def log(m): | |
| C.log(m, LOG) | |
| # ----------------------------------------------------------------------------- phase raw | |
| def phase_raw(): | |
| import cwm3_dataset as ds | |
| for regime, split in SETS: | |
| out = C.raw_path(regime, split) | |
| if out.exists(): | |
| log(f"[raw] {out.name} exists — skipped") | |
| continue | |
| t0 = time.time() | |
| pr = C.probe_rows(regime, split, log) | |
| rows = pr["rows"] | |
| full = ds.load_split(regime, split, ds.ALL_STREAMS, verbose=log) | |
| assert full["meta"]["N_file"] == pr["N_file"] | |
| raw = ds.subset(full, rows) | |
| del full | |
| with np.load(C.PRESS / f"pressure_{regime}_{split}.npz", allow_pickle=True) as z: | |
| assert np.array_equal(z["uid"][rows], raw["uid"]) and np.allclose(z["t"][rows], raw["t"]), "pressure uid/t witness" | |
| for k, k2 in (("P", "P"), ("Pvalid", "Pvalid"), ("P0", "P0"), ("P0valid", "P0valid"), ("P_first", "P_first"), ("S0", "PS0"), | |
| ("M0", "M0"), ("M0valid", "M0valid"), ("Moving", "PMoving"), ("Moving0", "PMoving0"), ("InRange", "PInRange")): | |
| raw[k2] = z[k][rows] | |
| raw["P_families"] = z["families"] | |
| raw["M0_names"] = z["margin_names"] | |
| raw["region_rows"] = pr["region"] | |
| raw["rows"] = rows | |
| meta = raw.pop("meta") | |
| np.savez(out, **{k: v for k, v in raw.items() if isinstance(v, np.ndarray)}, meta=np.array(str(meta))) | |
| log(f"[raw] {out.name}: n={len(rows):,} keys={len(raw)} {time.time() - t0:.0f}s") | |
| # ----------------------------------------------------------------------------- phase bank | |
| def phase_bank(): | |
| tok = C.npz_memmap(C.BANKS / "front_tok.npz", "tok") | |
| pooled = C.npz_memmap(C.BANKS / "front_tok.npz", "pooled256") | |
| bank_t = np.asarray(C.npz_memmap(C.BANKS / "front_tok.npz", "t")) | |
| for regime, split in SETS: | |
| out = C.OUT / f"bank_{regime}_{split}.npz" | |
| if out.exists(): | |
| log(f"[bank] {out.name} exists — skipped") | |
| continue | |
| t0 = time.time() | |
| raw = np.load(C.raw_path(regime, split), allow_pickle=True) | |
| ti = raw["tok_idx"].astype(np.int64) | |
| ok = ti >= 0 | |
| safe = np.where(ok, ti, 0) | |
| flat = safe.reshape(-1) | |
| order = np.argsort(flat, kind="stable") # sorted reads from the memmap | |
| tok16 = np.zeros((len(flat), 4, 64), np.float16) | |
| pooled16 = np.zeros((len(flat), 256), np.float16) | |
| tok16[order] = tok[flat[order]] | |
| pooled16[order] = pooled[flat[order]] | |
| tok16 = tok16.reshape(ti.shape + (4, 64)) * ok[:, :, None, None] | |
| pooled16 = pooled16.reshape(ti.shape + (256,)) * ok[:, :, None] | |
| bt = np.where(ok, bank_t[safe], np.nan).astype(np.float32) | |
| t_anchor = raw["t"].astype(np.float64) | |
| lag = t_anchor[:, None] - bt | |
| assert np.all(lag[ok] >= -1e-6), "non-causal gather" | |
| np.savez(out, tok16=tok16, pooled16=pooled16, bank_t16=bt, tok_ok=ok) | |
| log(f"[bank] {out.name}: n={len(ti):,} all16={ok.all(1).mean():.4f} last={ok[:, -1].mean():.4f} max_lag_last={np.nanmax(lag[:, -1]):.2f}s {time.time() - t0:.0f}s") | |
| # ----------------------------------------------------------------------------- phase frontfeat | |
| def _ff_job(job): | |
| uid, items = job # items: list of (set_i, row_i, local, need_region) | |
| t0 = time.time() | |
| p = C.FF / f"{uid}.npz" | |
| loc = np.array([it[2] for it in items], np.int64) | |
| order = np.argsort(loc) | |
| want = loc[order] | |
| feats = C.gather_member_rows(p, "feats", want) | |
| out_f = np.zeros((len(items), 1408), np.float16) | |
| out_f[order] = feats | |
| need = np.array([it[3] for it in items], bool) | |
| out_r = None | |
| if need.any(): | |
| loc_r = loc[need] | |
| order_r = np.argsort(loc_r) | |
| reg = C.gather_member_rows(p, "region", loc_r[order_r]) | |
| out_r = np.zeros((len(loc_r), 16, 1408), np.float16) | |
| out_r[order_r] = reg | |
| with np.load(p) as z: | |
| t_route = z["t"][want[[0, -1]]] if len(want) else None | |
| return uid, out_f, out_r, need, round(time.time() - t0, 1), t_route | |
| def phase_frontfeat(workers=4): | |
| todo = [(r, s) for r, s in SETS if not (C.OUT / f"frontfeat_{r}_{s}.npz").exists()] | |
| if not todo: | |
| log("[frontfeat] all exist — skipped") | |
| return | |
| with np.load(C.BANKS / "front_tok.npz", allow_pickle=True) as z: | |
| names, off = [str(u) for u in z["uid_index"]], z["route_off"] | |
| per_route = {} | |
| sets = [] | |
| for si, (r, s) in enumerate(todo): | |
| raw = np.load(C.raw_path(r, s), allow_pickle=True) | |
| ti = raw["tok_idx"][:, -1].astype(np.int64) | |
| reg = raw["region_rows"].astype(bool) | |
| sets.append({"n": len(ti), "ok": ti >= 0, "region": reg, "uid": raw["uid"], "t": raw["t"]}) | |
| for i in np.flatnonzero(ti >= 0): | |
| rt = int(np.searchsorted(off, ti[i], side="right") - 1) | |
| local = int(ti[i] - off[rt]) | |
| per_route.setdefault(names[rt], []).append((si, int(i), local, bool(reg[i]))) | |
| jobs = sorted(per_route.items(), key=lambda kv: -len(kv[1])) | |
| n_reg = sum(it[3] for kv in jobs for it in kv[1]) | |
| log(f"[frontfeat] {len(jobs)} routes, {sum(len(v) for _, v in jobs):,} rows ({n_reg:,} with region) for {todo}") | |
| res = [{"raw1408": np.full((st["n"], 1408), np.nan, np.float16), | |
| "region": np.lib.format.open_memmap(C.OUT / f"region_{r}_{s}.npy", mode="w+", dtype=np.float16, shape=(int(st["region"].sum()), 16, 1408)), | |
| "region_filled": np.zeros(int(st["region"].sum()), bool)} for st, (r, s) in zip(sets, todo)] | |
| reg_pos = [np.cumsum(st["region"]) - 1 for st in sets] | |
| t0 = time.time() | |
| done = 0 | |
| with Pool(workers, initializer=lambda: os.environ.update(OMP_NUM_THREADS="1")) as pool: | |
| for uid, out_f, out_r, need, dt, t_route in pool.imap_unordered(_ff_job, jobs, chunksize=1): | |
| items = per_route[uid] | |
| k_r = 0 | |
| for k, (si, i, local, nr) in enumerate(items): | |
| res[si]["raw1408"][i] = out_f[k] | |
| if nr: | |
| j = reg_pos[si][i] | |
| res[si]["region"][j] = out_r[k_r] | |
| res[si]["region_filled"][j] = True | |
| k_r += 1 | |
| done += 1 | |
| if done % 40 == 0 or done == len(jobs): | |
| log(f"[frontfeat] {done}/{len(jobs)} routes ({time.time() - t0:.0f}s)") | |
| for si, (r, s) in enumerate(todo): | |
| out = C.OUT / f"frontfeat_{r}_{s}.npz" | |
| res[si]["region"].flush() | |
| np.savez(out, raw1408=res[si]["raw1408"], region_filled=res[si]["region_filled"], region_rows=sets[si]["region"], | |
| region_file=np.array(str(C.OUT / f"region_{r}_{s}.npy"))) | |
| log(f"[frontfeat] {out.name}: raw ok {np.isfinite(res[si]['raw1408'][:, 0]).mean():.4f}, region filled {res[si]['region_filled'].mean():.4f}") | |
| # ----------------------------------------------------------------------------- phase vis (s5-derived pooled streams) | |
| def _vis_job(job): | |
| regime, split, fam, key, tag = job | |
| t0 = time.time() | |
| src = C.S12R / f"{fam}_{regime}_{split}.npz" | |
| raw = np.load(C.raw_path(regime, split), allow_pickle=True) | |
| rows = raw["rows"].astype(np.int64) | |
| with np.load(src, allow_pickle=False) as z: | |
| row_ok = z["row_ok"].astype(bool) | |
| pos = np.flatnonzero(row_ok) | |
| ok = row_ok[rows] | |
| want = np.searchsorted(pos, rows[ok]) | |
| assert np.array_equal(pos[want], rows[ok]) | |
| got = C.gather_member_rows(src, key, want) | |
| out = np.full((len(rows), 4224), np.nan, np.float16) | |
| if got is not None: | |
| out[ok] = got | |
| np.save(C.OUT / f"vis_{tag}_{regime}_{split}.npy", out) | |
| return f"[vis] {tag} {regime}/{split}: ok={ok.mean():.4f} {time.time() - t0:.0f}s" | |
| def phase_vis(workers=3): | |
| jobs = [] | |
| for r, s in SETS: | |
| for fam, key, tag in (("vis", "Vf", "Vf4224"), ("vis", "Vc", "Vc4224"), ("crop", "V", "crop4224")): | |
| if not (C.OUT / f"vis_{tag}_{r}_{s}.npy").exists(): | |
| jobs.append((r, s, fam, key, tag)) | |
| if not jobs: | |
| log("[vis] all exist — skipped") | |
| return | |
| with Pool(workers, initializer=lambda: os.environ.update(OMP_NUM_THREADS="1")) as pool: | |
| for msg in pool.imap_unordered(_vis_job, jobs): | |
| log(msg) | |
| # ----------------------------------------------------------------------------- assemble | |
| def phase_assemble(): | |
| for r, s in SETS: | |
| out = C.feats_path(r, s) | |
| if out.exists(): | |
| log(f"[assemble] {out.name} exists — skipped") | |
| continue | |
| parts = {} | |
| ff = np.load(C.OUT / f"frontfeat_{r}_{s}.npz") | |
| parts.update(raw1408=ff["raw1408"], region_filled=ff["region_filled"], region_rows=ff["region_rows"], region_file=ff["region_file"]) | |
| bk = np.load(C.OUT / f"bank_{r}_{s}.npz") | |
| parts.update(tok16=bk["tok16"], pooled16=bk["pooled16"], bank_t16=bk["bank_t16"], tok_ok=bk["tok_ok"]) | |
| for tag in ("Vf4224", "Vc4224", "crop4224"): | |
| parts[tag] = np.load(C.OUT / f"vis_{tag}_{r}_{s}.npy") | |
| np.savez(out, **parts) | |
| log(f"[assemble] {out.name}: " + ", ".join(f"{k}{v.shape}" for k, v in parts.items())) | |
| def main(): | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--phase", default="all", choices=["raw", "bank", "frontfeat", "vis", "assemble", "all"]) | |
| ap.add_argument("--workers", type=int, default=4) | |
| a = ap.parse_args() | |
| C.OUT.mkdir(parents=True, exist_ok=True) | |
| t0 = time.time() | |
| if a.phase in ("raw", "all"): | |
| phase_raw() | |
| if a.phase in ("bank", "all"): | |
| phase_bank() | |
| if a.phase in ("frontfeat", "all"): | |
| phase_frontfeat(a.workers) | |
| if a.phase in ("vis", "all"): | |
| phase_vis(min(a.workers, 3)) | |
| if a.phase in ("assemble", "all"): | |
| phase_assemble() | |
| log(f"DONE phase={a.phase} {time.time() - t0:.0f}s") | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 11.7 kB
- Xet hash:
- d0fa22a7fb6e8c5351648e28ee9bbf29fc418910d3e2510327f2f69da8269790
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.