Buckets:
| """cwm3/data/pca_streams.py — PCA-256 of the pooled V-JEPA2 streams (CPU only). | |
| Streams (s12_wm/results): Vf = vis_{r}_{s}.npz["Vf"] front pooled, causal ([t-2,t] window) | |
| Vc = vis_{r}_{s}.npz["Vc"] cabin full-frame pooled, causal | |
| crop= crop_{r}_{s}.npz["V"] cabin route-median crop pooled, NON-causal ROI (diagnostic only) | |
| Each is [n_ok, 4224] f16 with a row_ok[N] mask against the s12 anchor grid (targets_{r}_{s}.npz). | |
| --fit : a seeded uniform sample of <= --n-fit rows of MANUAL TRAIN (row_ok rows) per stream -> exact covariance PCA | |
| (float64 eigendecomposition) -> results/pca/pca_{stream}_256.npz {mu, W[4224,256], evr, n_fit, seed, source_sha256} | |
| + results/pca/pca_manifest.jsonl (append-only: file, sha256, explained variance). | |
| --pack : every (regime, split, stream) -> results/pack/{vf256,vc256,crop256}_{regime}_{split}.npz | |
| {Z f16[N,256] (zero where row_ok is False), row_ok bool[N], pca_sha256, source, n} | |
| Both phases stream the zip member chunk by chunk (grid_dataset.scatter_by_row_ok idea) — never the whole 5-6 GB array. | |
| No openpilot DMS field is read (inputs: V-JEPA2 pooled latents only). | |
| python pca_streams.py --fit --pack [--workers 4] [--streams Vf,Vc,crop] | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import hashlib | |
| import json | |
| import sys | |
| import time | |
| import zipfile | |
| from multiprocessing import Pool | |
| from pathlib import Path | |
| import numpy as np | |
| from numpy.lib import format as npf | |
| CWM3 = Path(__file__).resolve().parents[1] | |
| EXPLORE = CWM3.parent | |
| S12R = EXPLORE / "s12_wm" / "results" | |
| PCA_DIR = CWM3 / "results" / "pca" | |
| PACK_DIR = CWM3 / "results" / "pack" | |
| sys.path.insert(0, str(CWM3 / "common")) | |
| from runmeta import write_run_json, sha256_file # noqa: E402 | |
| D_IN, K = 4224, 256 | |
| STREAMS = {"Vf": ("vis", "Vf", "vf256"), "Vc": ("vis", "Vc", "vc256"), "crop": ("crop", "V", "crop256")} | |
| REGIMES, SPLITS = ("manual", "engaged"), ("train", "val", "test") | |
| CHUNK = 16384 | |
| def src_path(stream, regime, split): | |
| fam = STREAMS[stream][0] | |
| return S12R / f"{fam}_{regime}_{split}.npz" | |
| def iter_member(path, key, chunk=CHUNK): | |
| """Yield (row0, block[k, ...]) of a .npy member of a (compressed) npz, streamed through the zip decompressor.""" | |
| with zipfile.ZipFile(path) as zf, zf.open(key + ".npy") as fh: | |
| version = npf.read_magic(fh) | |
| shape, fortran, dtype = npf._read_array_header(fh, version) | |
| assert not fortran and dtype != object, f"{path}:{key}" | |
| rowsz = int(np.prod(shape[1:], dtype=np.int64)) * dtype.itemsize | |
| n = shape[0] | |
| for i in range(0, n, chunk): | |
| k = min(chunk, n - i) | |
| buf = fh.read(k * rowsz) | |
| assert len(buf) == k * rowsz, f"{path}:{key} short read at {i}" | |
| yield i, np.frombuffer(buf, dtype).reshape((k,) + tuple(shape[1:])) | |
| def read_small(path, key): | |
| with np.load(path, allow_pickle=False) as z: | |
| return z[key] | |
| # ----------------------------------------------------------------------------- fit | |
| def fit_stream(stream, n_fit, seed, log): | |
| t0 = time.time() | |
| fam, key, _ = STREAMS[stream] | |
| p = src_path(stream, "manual", "train") | |
| row_ok = read_small(p, "row_ok").astype(bool) | |
| n_ok = int(row_ok.sum()) | |
| rng = np.random.default_rng(seed) | |
| pick = np.zeros(n_ok, bool) | |
| pick[rng.choice(n_ok, min(n_fit, n_ok), replace=False)] = True | |
| S = np.zeros(D_IN, np.float64) | |
| C = np.zeros((D_IN, D_IN), np.float64) | |
| n = 0 | |
| for i, blk in iter_member(p, key): | |
| sel = pick[i:i + len(blk)] | |
| if not sel.any(): | |
| continue | |
| x = blk[sel].astype(np.float64) | |
| S += x.sum(0) | |
| C += x.T @ x | |
| n += len(x) | |
| mu = S / n | |
| cov = C / n - np.outer(mu, mu) | |
| evals, evecs = np.linalg.eigh(cov) | |
| order = np.argsort(evals)[::-1] | |
| evals, evecs = evals[order], evecs[:, order] | |
| W = evecs[:, :K].astype(np.float32) | |
| evr = (evals[:K] / evals.sum()).astype(np.float32) | |
| out = PCA_DIR / f"pca_{stream}_{K}.npz" | |
| assert not out.exists(), f"{out} exists (append-only results; delete manually if a refit is intended)" | |
| np.savez(out, mu=mu.astype(np.float32), W=W, evr=evr, n_fit=np.array(n), seed=np.array(seed), | |
| source=np.array(str(p)), source_sha256=np.array(sha256_file(p)), stream=np.array(stream), | |
| created_utc=np.array(time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()))) | |
| sha = sha256_file(out) | |
| with open(PCA_DIR / "pca_manifest.jsonl", "a") as f: | |
| f.write(json.dumps({"file": out.name, "sha256": sha, "stream": stream, "n_fit": n, "k": K, | |
| "evr_sum": float(evr.sum()), "evr_first8": [float(v) for v in evr[:8]], | |
| "source": str(p), "seed": seed, "wall_s": round(time.time() - t0, 1)}) + "\n") | |
| log(f"[fit] {stream}: n={n} evr(256)={evr.sum():.4f} top8={np.round(evr[:8], 3).tolist()} -> {out.name} sha={sha[:12]} {time.time() - t0:.0f}s") | |
| return out | |
| # ----------------------------------------------------------------------------- pack | |
| def pack_one(job): | |
| stream, regime, split = job | |
| t0 = time.time() | |
| fam, key, tag = STREAMS[stream] | |
| out = PACK_DIR / f"{tag}_{regime}_{split}.npz" | |
| if out.exists(): | |
| return f"[pack] {out.name} exists — skipped" | |
| pca_p = PCA_DIR / f"pca_{stream}_{K}.npz" | |
| pz = np.load(pca_p) | |
| mu, W = pz["mu"].astype(np.float32), pz["W"].astype(np.float32) | |
| p = src_path(stream, regime, split) | |
| row_ok = read_small(p, "row_ok").astype(bool) | |
| N = len(row_ok) | |
| pos = np.flatnonzero(row_ok) | |
| Z = np.zeros((N, K), np.float16) | |
| n_seen = 0 | |
| for i, blk in iter_member(p, key): | |
| z = (blk.astype(np.float32) - mu) @ W | |
| Z[pos[i:i + len(blk)]] = z.astype(np.float16) | |
| n_seen += len(blk) | |
| assert n_seen == len(pos), f"{p}:{key} has {n_seen} rows but row_ok.sum()={len(pos)}" | |
| tg = S12R / f"targets_{regime}_{split}.npz" | |
| with np.load(tg, allow_pickle=True) as z: | |
| assert len(z["t"]) == N, f"{p}: row_ok has {N} rows, targets {len(z['t'])}" | |
| np.savez(out, Z=Z, row_ok=row_ok, pca_file=np.array(pca_p.name), pca_sha256=np.array(sha256_file(pca_p)), | |
| source=np.array(str(p)), targets=np.array(str(tg)), n=np.array(N), stream=np.array(stream), | |
| roi=np.array("route-median (non-causal)" if stream == "crop" else "fullframe_causal")) | |
| return f"[pack] {out.name}: N={N} ok={row_ok.mean():.4f} {time.time() - t0:.0f}s" | |
| def main(): | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--fit", action="store_true") | |
| ap.add_argument("--pack", action="store_true") | |
| ap.add_argument("--streams", default="Vf,Vc,crop") | |
| ap.add_argument("--n-fit", type=int, default=200_000) | |
| ap.add_argument("--seed", type=int, default=0) | |
| ap.add_argument("--workers", type=int, default=4) | |
| a = ap.parse_args() | |
| t_start = time.time() | |
| PCA_DIR.mkdir(parents=True, exist_ok=True) | |
| PACK_DIR.mkdir(parents=True, exist_ok=True) | |
| logp = PCA_DIR / "pca_streams.log" | |
| def log(m): | |
| line = f"[{time.strftime('%H:%M:%S')}] {m}" | |
| print(line, flush=True) | |
| with open(logp, "a") as f: | |
| f.write(line + "\n") | |
| streams = [s for s in a.streams.split(",") if s] | |
| if a.fit: | |
| for s in streams: | |
| if (PCA_DIR / f"pca_{s}_{K}.npz").exists(): | |
| log(f"[fit] pca_{s}_{K}.npz exists — skipped") | |
| continue | |
| fit_stream(s, a.n_fit, a.seed, log) | |
| if a.pack: | |
| jobs = [(s, r, sp) for s in streams for r in REGIMES for sp in SPLITS] | |
| with Pool(a.workers) as pool: | |
| for msg in pool.imap_unordered(pack_one, jobs): | |
| log(msg) | |
| write_run_json(PCA_DIR, "pca_streams", [Path(__file__)], | |
| [src_path(s, r, sp) for s in streams for r in REGIMES for sp in SPLITS if src_path(s, r, sp).exists()], | |
| seed=a.seed, t_start=t_start, input_listing=False, | |
| extra={"fit": a.fit, "pack": a.pack, "streams": streams, "n_fit": a.n_fit, "k": K, "dms_fields_used": [], | |
| "outputs": sorted(p.name for p in PACK_DIR.glob("*.npz")) + sorted(p.name for p in PCA_DIR.glob("pca_*_256.npz"))}) | |
| log("DONE") | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 8.28 kB
- Xet hash:
- d9f682c7af5e40f7da9348939c1cec241335b56b8987067e34280ce69c8aeb6c
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.