Buckets:
| """cwm3/data/cwm3_dataset.py — the row-aligned join for the cwm3 branches (both regimes concatenated, `mode` column). | |
| Per (regime, split), row-aligned on the s12 anchor grid (uid/t witness asserted against targets_{r}_{s}.npz): | |
| s12_wm/results/targets_{r}_{s}.npz uid dongle t X f32[N,40,5] (t-3.9..t @10 Hz) Y f32[N,25,5] M bool[N,25] slice_tag u8 | |
| cwm3/results/pack/xpre_{r}_{s}.npz Xpre f32[N,40,5] (t-7.9..t-4.0 @10 Hz) ok -> X8 = concat(Xpre, X)[:, 1::2] (8 s @5 Hz) | |
| cwm2/results/targets_kin/kinhist_* Kh f16[N,40,11] (t-3.9..t @10 Hz, NaN = invalid; route-median ROI -> NON-causal, disclosed) | |
| cwm2/results/targets_kin/kin_* K f16[N,25,11] Kvalid[N,25] S f32[N,25,2] (NaN censored) S0[N,2] EM bool[N,25,3] | |
| cwm2/results/targets_demand/demand_* D Dvalid bool[N,25,6] D0 D0valid[N,6] Resp RespValid[N,25] Resp0 Resp0valid[N] NoResp NoRespValid Moving0 | |
| s12_wm/results/future_* F f16[N,5,256] Fok bool[N,5] | |
| cwm3/results/pack/{vf256,vc256,crop256}_{r}_{s}.npz Z f16[N,256] + row_ok (PCA-256 of Vf / Vc full-frame causal / crop non-causal) | |
| pack version v2 (set_pack_version("v2"), 2026-09-10 latent audit): vf256 / vc256 and future_* from cwm3/results/pack_v2/ (frame/20 time axis; | |
| the v1 s5-library files lag by up to 18 s on ~20 % of routes); crop256, xpre, tokens and every non-visual stream are unchanged. | |
| cwm3/results/banks/anchor_index_{r}_{s}.npy int32[N,16] rows of results/banks/front_tok.npz (-1 missing); bank via load_bank() | |
| cwm3/results/pack_v2/regionpca256_{r}_{s}.npz vector f16[N,512] (PCA-256 of the flattened 16x1408 region grid: last window (+) 8 s mean), | |
| ok bool[N]; pack v2 ONLY (2026-09-11 road stream, data/build_regionpca256.py); row order = vf256 | |
| cwm3/results/targets_pressure/pressure_{r}_{s}.npz P bool[N,25,6] Pvalid P0[N,6] P0valid (road-pressure chains ttc, ttc_drop, cut_in, lead_braking, | |
| lane_excursion, badas; data/build_pressure_targets.py) — stream `pressure`, E0 `*_P` variants (2026-09-11) | |
| kin_source causal_routes / oldin_ctgt, kin_version v1 | v2 (set_kin_source(v, version)): v2 = results/targets_kin_causal_routes_v2 (74 routes | |
| re-extracted with the population driver-side prior; same keys / contract; docs/k4_handoff.md) | |
| Derived here (CPU, once): phase0/phase (steer-phase chain, s12_wm/episodes hysteresis), dt channels are NOT materialised (constant | |
| grids: X/Kh step k -> -(39-k)*0.1 s; X8 step k -> -(39-k)*0.2 s; token slot j -> t_bank - t, real) — train/batch adds them on the GPU. | |
| `need`: the set of streams a stage loads (memory discipline); everything is kept in its compact storage dtype. | |
| Both regimes: load_regimes() concatenates the per-regime dicts and adds mode (0 manual / 1 engaged), regime_row (row in its own | |
| targets file) and integer codes for dongle / uid. --max-anchors: seeded uniform subset (rows chosen before any stream is read). | |
| Never reads grid10 / driver_state.csv / any openpilot DMS field (the kin files are pose-derived; --check asserts the key sets). | |
| python cwm3_dataset.py --check [--regimes manual,engaged] [--split train] # shapes, coverage, RSS footprint | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import os | |
| import sys | |
| import time | |
| from pathlib import Path | |
| import numpy as np | |
| CWM3 = Path(__file__).resolve().parents[1] | |
| EXPLORE = CWM3.parent | |
| CWM2 = EXPLORE / "cwm2" | |
| S12R = EXPLORE / "s12_wm" / "results" | |
| KIN = CWM2 / "results" / "targets_kin" | |
| DEM = CWM2 / "results" / "targets_demand" | |
| PACK = CWM3 / "results" / "pack" | |
| PACK_V2 = CWM3 / "results" / "pack_v2" # 2026-09-10 latent audit: vf256 / vc256 / F rebuilt on the frame/20 time axis | |
| _PACK_VERSION = {"v": "v1"} | |
| def set_pack_version(v): | |
| """v1 (default): results/pack + s12 future files; v2: results/pack_v2 for vf256 / vc256 and the aligned F (crop256 has no v2).""" | |
| assert v in ("v1", "v2"), v | |
| _PACK_VERSION["v"] = v | |
| def pack_version(): | |
| return _PACK_VERSION["v"] | |
| KINC = CWM3 / "results" / "targets_kin_causal_routes" # K4 (2026-09-11): causal route-level driver kinematics (docs/k4_handoff.md) | |
| KINC_V2 = CWM3 / "results" / "targets_kin_causal_routes_v2" # K4 v2 (2026-09-11): 74 co-occupant routes re-extracted with a population driver-side prior | |
| KINC_V3 = CWM3 / "results" / "targets_kin_causal_routes_v3" # v3 (2026-09-16): manual-derived wheel reference + hysteresis + 1 s persistence for hands_on / hand events; INPUTS identical to v2 (kinhist hard-linked, K_c asserted) | |
| KINC_V2P = CWM3 / "results" / "targets_kin_causal_routes_v2p" # v2p (2026-09-18, plan A arm; pre-reg docs/decisions.md 2026-09-18T05:34:40Z): v2 + 3 keypoint proxies as INPUT columns 11-13 of Kh_c / Kh8_c / K0_c | |
| # (head_pitch_kp, wrist_face_kp, hands_lap_kp; data/build_kin_causal_v2p.py); K_c / S_c / E_c / validity BIT-IDENTICAL to v2 (targets unchanged) | |
| _KIN_SOURCE = {"v": "cwm2", "version": "v1"} | |
| def set_kin_source(v, version="v1"): | |
| """cwm2 (default): cwm2/results/targets_kin (route-median ROI, NON-causal); causal_routes: results/targets_kin_causal_routes[_v2] (causal box per | |
| frame; keys *_c renamed to the old contract at load time, invalid history steps -> NaN, rows without a causal box -> kin_row_ok False). | |
| version: v1 (default) | v2 = the population-prior re-extraction (same contract); applies to causal_routes AND oldin_ctgt (its CAUSAL targets).""" | |
| assert v in ("cwm2", "causal_routes", "oldin_ctgt"), v # oldin_ctgt (K4 input-isolation control): OLD route-median Kh inputs + CAUSAL targets | |
| assert version in ("v1", "v2", "v3", "v2p"), version # v3 (2026-09-16): same contract as v2, only S_c / S0_c / E_c / EM_c / E0_c differ; v2p (2026-09-18): v2 + 3 input proxies (Kh 14 wide) | |
| _KIN_SOURCE["v"], _KIN_SOURCE["version"] = v, version | |
| def kin_source(): | |
| return _KIN_SOURCE["v"] | |
| def kin_version(): | |
| return _KIN_SOURCE["version"] | |
| def kin_dir(): | |
| """the causal kin pack directory of the selected version (v1 dir when kin_source is cwm2, for reference only).""" | |
| return {"v2": KINC_V2, "v3": KINC_V3, "v2p": KINC_V2P}.get(kin_version(), KINC) | |
| def n_kin(): | |
| """width of the kinematics INPUT history Kh of the selected kin version: 14 for v2p (11 v2 channels + 3 keypoint proxies), 11 otherwise.""" | |
| return len(KIN_CHANNELS_V2P) if kin_version() == "v2p" else len(KIN_CHANNELS) | |
| BANKS = CWM3 / "results" / "banks" | |
| PRESS = CWM3 / "results" / "targets_pressure" # road-pressure chains (P), 2026-09-11 E0 P-head variants | |
| ROAD_DYN = CWM3 / "results" / "targets_road_dyn" # 2026-09-16 WP4 F5b: E side streams vfcan256 (vf256 PCA 0..215 + 40 CAN anchor features) / can256 (CAN only); data/build_road_dynamics_targets.py | |
| GAZE = CWM3 / "results" / "targets_gaze_causal_routes_v2" # 2026-09-12 gaze chain: causal gaze_off state chain on the s12 grid (data/build_gaze_targets.py); U `u_gaze_chain` | |
| KPF_DIR = None # LANE F1 (2026-09-17): set_kpf_dir(<KPF run dir>) -> stream "kpf" = kpf_{regime}_{split}.npz feat [N,10] f16 | |
| FACE = CWM3 / "results" / "targets_face_causal_routes_v2" # 2026-09-12: our own causal face-landmark channels (data/build_face_causal.py, F = 9 `lm_*`); U `u_face` inputs | |
| sys.path.insert(0, str(EXPLORE / "s12_wm")) | |
| sys.path.insert(0, str(CWM3 / "common")) | |
| import episodes as ep # noqa: E402 | |
| from pagecache import drop_all as _drop_cache # noqa: E402 (2026-09-11: evict loaded npz files from the page cache; free-memory guard) | |
| HIST, H, DX, DK, NSLOT = 40, 25, 5, 11, 16 | |
| MODE = {"manual": 0, "engaged": 1} | |
| REGIMES, SPLITS = ("manual", "engaged"), ("train", "val", "test") | |
| FAMILIES = ["steer_rate", "decel", "ttc", "ttc_drop", "lane", "badas"] | |
| P_FAMILIES = ["ttc", "ttc_drop", "cut_in", "lead_braking", "lane_excursion", "badas"] # road-pressure chain order (build_pressure_targets.FAMILIES) | |
| KIN_CHANNELS = ["head_cx", "head_cy", "Lwri_x", "Lwri_y", "Rwri_x", "Rwri_y", "head_turn", "head_turn_dev", | |
| "Lhand_rest_dist", "Rhand_rest_dist", "min_hand_rest_dist"] | |
| IDX_MIN_DIST, IDX_HEAD_CX = 10, 0 | |
| KIN_PROXY_CHANNELS = ["head_pitch_kp", "wrist_face_kp", "hands_lap_kp"] # v2p (2026-09-18): our own HALPE-136 keypoint geometry (data/build_kin_causal_v2p.py); never a DMS quantity | |
| KIN_CHANNELS_V2P = KIN_CHANNELS + KIN_PROXY_CHANNELS | |
| FACE_CHANNELS = ["lm_yaw", "lm_yaw_dev", "lm_pitch_dev", "lm_roll_dev", "ear_conf_asym", "eye_open", "blink_rate_4s", "gaze_off", "gaze_off_dwell_s"] # build_face_causal.CHANNELS | |
| N_FACE = len(FACE_CHANNELS) | |
| GAZE_STATES, GAZE_EVENTS = ["gaze_off"], ["gaze_off_on", "gaze_off_off"] | |
| ALL_STREAMS = frozenset({"X", "X8", "Kh", "kin", "demand", "F", "vf256", "vc256", "crop256", "tokens", "Y", "regionpca256", "pressure", "gaze", "face", "vfcan256", "can256", "vfcan256s"}) | |
| NEED = {"V": frozenset({"X", "X8", "Y", "demand"}), # demand only for Resp (response chain) | |
| "U": frozenset({"X", "Kh", "kin", "Y", "vc256", "crop256", "gaze", "face"}), # gaze / face: opt-in U variants (2026-09-12) | |
| "E": frozenset({"Y", "demand", "F", "vf256", "tokens", "regionpca256", "pressure", "vfcan256", "can256", "vfcan256s"})} # regionpca256: pack v2 only; pressure: E0 `*_P` variants; vfcan256 / can256: F5b side streams (2026-09-16) | |
| FORBIDDEN_KEYS = {"faceProb", "isDistracted", "face_yaw", "face_pitch", "driver_state", "dstate", "grid10"} | |
| ROI = {"Kh": "route-median (non-causal)", "K": "route-median (non-causal)", "S": "route-median (non-causal)", | |
| "crop256": "route-median (non-causal)", "vc256": "fullframe_causal", "vf256": "fullframe_causal", "tokens": "fullframe_causal", | |
| "regionpca256": "fullframe_causal", "Fh": "causal_clip_roi", "Sg": "causal_clip_roi", | |
| "vfcan256": "fullframe_causal (216 PCA) + CAN anchor features", "can256": "CAN anchor features only (no vision)", "vfcan256s": "fullframe_causal (216 PCA) + CAN anchor features x4 (post-hoc scale arm)"} # Fh / Sg: causal per-frame driver box (route-level pose), 2026-09-12 | |
| def set_kpf_dir(d): | |
| """LANE F1: the keypoint-forecaster run dir whose kpf_{regime}_{split}.npz feed the joint model's `kpf` stream (None = stream unavailable).""" | |
| global KPF_DIR | |
| KPF_DIR = Path(d) if d else None | |
| def files(regime, split): | |
| v2 = pack_version() == "v2" | |
| fut = (PACK_V2 if v2 else S12R) / f"future_{regime}_{split}.npz" | |
| vis = PACK_V2 if v2 else PACK | |
| return {"targets": S12R / f"targets_{regime}_{split}.npz", "xpre": PACK / f"xpre_{regime}_{split}.npz", | |
| "kinhist": (kin_dir() / f"kinhist_causal_{regime}_{split}.npz" if kin_source() == "causal_routes" else KIN / f"kinhist_{regime}_{split}.npz"), | |
| "kin": (kin_dir() / f"kin_causal_{regime}_{split}.npz" if kin_source() in ("causal_routes", "oldin_ctgt") else KIN / f"kin_{regime}_{split}.npz"), | |
| "regionpca256": PACK_V2 / f"regionpca256_{regime}_{split}.npz", # pack v2 only (asserted in load_split) | |
| "vfcan256": ROAD_DYN / f"vfcan256_{regime}_{split}.npz", "can256": ROAD_DYN / f"can256_{regime}_{split}.npz", # F5b side streams (Z / row_ok / n contract of vf256) | |
| "vfcan256s": ROAD_DYN / f"vfcan256s_{regime}_{split}.npz", # post-hoc scale-sensitivity side stream (CAN block x4) | |
| "pressure": PRESS / f"pressure_{regime}_{split}.npz", | |
| "gaze": GAZE / f"gaze_{regime}_{split}.npz", "face": FACE / f"facehist_causal_{regime}_{split}.npz", # opt-in U streams (2026-09-12) | |
| "kpf": (KPF_DIR / f"kpf_{regime}_{split}.npz" if KPF_DIR is not None else None), # LANE F1 (2026-09-17) opt-in U stream | |
| "demand": DEM / f"demand_{regime}_{split}.npz", "future": fut, | |
| "vf256": vis / f"vf256_{regime}_{split}.npz", "vc256": vis / f"vc256_{regime}_{split}.npz", | |
| "crop256": PACK / f"crop256_{regime}_{split}.npz", "anchor_index": BANKS / f"anchor_index_{regime}_{split}.npy", | |
| "bank": BANKS / "front_tok.npz", "target_stats": S12R / "target_stats.npz"} | |
| def _rd(path, keys, sub, allow_pickle=False): | |
| out = {} | |
| with np.load(path, allow_pickle=allow_pickle) as z: | |
| for k in keys: | |
| out[k] = sub(z[k]) | |
| return out | |
| def load_split(regime, split, need=ALL_STREAMS, max_anchors=None, seed=0, verbose=print): | |
| """-> dict of compact numpy arrays for one (regime, split); keys depend on `need`. Always: uid dongle t M slice_tag mode idx.""" | |
| assert regime in REGIMES and split in SPLITS | |
| need = frozenset(need) | |
| t0 = time.time() | |
| f = files(regime, split) | |
| with np.load(f["targets"], allow_pickle=True) as z: | |
| N_file = len(z["t"]) | |
| sel = None | |
| if max_anchors and max_anchors < N_file: | |
| sel = np.sort(np.random.default_rng(seed).choice(N_file, max_anchors, replace=False)) | |
| sub = (lambda a: a[sel]) if sel is not None else (lambda a: a) | |
| out = {k: sub(z[k]) for k in ("uid", "dongle", "t", "M", "slice_tag")} | |
| X = sub(z["X"]).astype(np.float32) | |
| if "Y" in need or "X" in need: | |
| out["Y"] = sub(z["Y"]).astype(np.float32) | |
| N = len(out["t"]) | |
| out["M"] = out["M"].astype(bool) | |
| out["slice_tag"] = out["slice_tag"].astype(np.uint8) | |
| out["idx"] = sel.astype(np.int64) if sel is not None else np.arange(N, dtype=np.int64) | |
| out["mode"] = np.full(N, MODE[regime], np.int8) | |
| if "X" in need or "X8" in need: | |
| out["X"] = X | |
| if "Y" in need or "X" in need: | |
| # steer-phase chain (s12 hysteresis on |dθ/dt|; s0 from the context end) — the only CPU-derived target | |
| s_prev, s = ep.phase_chain_targets(out["Y"], X) | |
| out["phase0"] = s_prev[:, 0].astype(np.uint8) | |
| out["phase"] = s.astype(np.uint8) | |
| if "X8" in need: | |
| d = _rd(f["xpre"], ("Xpre", "ok"), sub) | |
| out["Xpre"], out["x8_ok"] = d["Xpre"].astype(np.float32), d["ok"].astype(bool) | |
| causal = kin_source() == "causal_routes" | |
| causal_tgt = kin_source() in ("causal_routes", "oldin_ctgt") | |
| if "Kh" in need: | |
| if causal: | |
| d = _rd(f["kinhist"], ("uid", "t", "Kh_c", "Khvalid", "row_ok"), sub) | |
| assert np.array_equal(d["uid"], out["uid"]) and np.allclose(d["t"], out["t"]), f"{f['kinhist']}: uid/t witness" | |
| Kh = d["Kh_c"].astype(np.float32) | |
| Kh[~d["Khvalid"].astype(bool)] = np.nan # the old contract: NaN = invalid step | |
| Kh[~d["row_ok"].astype(bool)] = np.nan | |
| out["Kh"] = Kh.astype(np.float16) | |
| out["kin_row_ok"] = d["row_ok"].astype(bool) | |
| else: | |
| d = _rd(f["kinhist"], ("uid", "t", "Kh"), sub) | |
| assert np.array_equal(d["uid"], out["uid"]) and np.allclose(d["t"], out["t"]), f"{f['kinhist']}: uid/t witness" | |
| out["Kh"] = d["Kh"].astype(np.float16) | |
| if "kin" in need: | |
| if causal_tgt: | |
| d = _rd(f["kin"], ("uid", "t", "K_c", "Kvalid", "S_c", "S0_c", "EM_c", "channels", "row_ok"), (lambda a: a), allow_pickle=True) | |
| assert np.array_equal(sub(d["uid"]), out["uid"]) and np.allclose(sub(d["t"]), out["t"]), f"{f['kin']}: uid/t witness" | |
| assert [str(c) for c in d["channels"]] == KIN_CHANNELS | |
| ok = sub(d["row_ok"]).astype(bool) | |
| out["K"] = sub(d["K_c"]).astype(np.float16) | |
| out["Kvalid"] = sub(d["Kvalid"]).astype(bool) & ok[:, None] | |
| S = sub(d["S_c"]).astype(np.float32) | |
| S[~ok] = np.nan # no causal box at the anchor -> censored states | |
| out["S"] = S.astype(np.float16) | |
| S0 = sub(d["S0_c"]).astype(np.float32) | |
| S0[~ok] = np.nan | |
| out["S0"] = S0 | |
| out["EM"] = sub(d["EM_c"]).astype(bool) & ok[:, None, None] | |
| out["kin_row_ok"] = ok | |
| else: | |
| d = _rd(f["kin"], ("uid", "t", "K", "Kvalid", "S", "S0", "EM", "channels"), sub if False else (lambda a: a), allow_pickle=True) | |
| assert np.array_equal(sub(d["uid"]), out["uid"]) and np.allclose(sub(d["t"]), out["t"]), f"{f['kin']}: uid/t witness" | |
| assert [str(c) for c in d["channels"]] == KIN_CHANNELS | |
| out["K"] = sub(d["K"]).astype(np.float16) | |
| out["Kvalid"] = sub(d["Kvalid"]).astype(bool) | |
| out["S"] = sub(d["S"]).astype(np.float16) # NaN kept (censored) | |
| out["S0"] = sub(d["S0"]).astype(np.float32) | |
| out["EM"] = sub(d["EM"]).astype(bool) | |
| if "demand" in need: | |
| d = _rd(f["demand"], ("uid", "t", "D", "Dvalid", "D0", "D0valid", "Resp", "RespValid", "Resp0", "Resp0valid", | |
| "NoResp", "NoRespValid", "Moving0", "families"), (lambda a: a), allow_pickle=True) | |
| assert np.array_equal(sub(d["uid"]), out["uid"]) and np.allclose(sub(d["t"]), out["t"]), f"{f['demand']}: uid/t witness" | |
| assert [str(c) for c in d["families"]] == FAMILIES | |
| for k in ("D", "Dvalid", "D0", "D0valid", "Resp", "RespValid", "Resp0", "Resp0valid", "NoResp", "NoRespValid", "Moving0"): | |
| out[k] = sub(d[k]).astype(bool) | |
| if "F" in need: | |
| d = _rd(f["future"], ("F", "ok"), sub) | |
| out["F"], out["Fok"] = d["F"].astype(np.float16), d["ok"].astype(bool) | |
| for s in ("vf256", "vc256", "crop256", "vfcan256", "can256", "vfcan256s"): # vfcan256 / can256 (+ vfcan256s scale arm) (2026-09-16 F5b): the same Z / row_ok / n contract, results/targets_road_dyn | |
| if s in need: | |
| d = _rd(f[s], ("Z", "row_ok", "n"), (lambda a: a)) | |
| assert int(d["n"]) == N_file, f"{f[s]}: n={int(d['n'])} != targets N={N_file}" | |
| out[s], out[s + "_ok"] = sub(d["Z"]).astype(np.float16), sub(d["row_ok"]).astype(bool) | |
| if "pressure" in need: | |
| d = _rd(f["pressure"], ("uid", "t", "P", "Pvalid", "P0", "P0valid", "families"), (lambda a: a), allow_pickle=True) | |
| assert np.array_equal(sub(d["uid"]), out["uid"]) and np.allclose(sub(d["t"]), out["t"]), f"{f['pressure']}: uid/t witness" | |
| assert [str(c) for c in d["families"]] == P_FAMILIES, list(d["families"]) | |
| for k in ("P", "Pvalid", "P0", "P0valid"): | |
| out[k] = sub(d[k]).astype(bool) | |
| if "gaze" in need: | |
| # 2026-09-12 gaze chain (data/build_gaze_targets.py): the causal gaze_off state at t+0.2..t+5.0 (S_g, NaN censored) / at t (S0_g) and the onset masks | |
| # EM_g [N,25,2] (gaze_off_on, gaze_off_off) — the kin convention (S / S0 / EM) for a third driver state; rows without a decoded frame censored. | |
| d = _rd(f["gaze"], ("uid", "t", "row_ok", "S_g", "S0_g", "EM_g", "state_names"), (lambda a: a), allow_pickle=True) | |
| assert np.array_equal(sub(d["uid"]), out["uid"]) and np.allclose(sub(d["t"]), out["t"]), f"{f['gaze']}: uid/t witness" | |
| assert [str(c) for c in d["state_names"]] == GAZE_STATES | |
| ok = sub(d["row_ok"]).astype(bool) | |
| Sg = sub(d["S_g"]).astype(np.float32) | |
| Sg[~ok] = np.nan | |
| out["Sg"] = Sg.astype(np.float16) | |
| Sg0 = sub(d["S0_g"]).astype(np.float32) | |
| Sg0[~ok] = np.nan | |
| out["Sg0"] = Sg0 | |
| out["EMg"] = sub(d["EM_g"]).astype(bool) & ok[:, None, None] | |
| if "face" in need: | |
| # 2026-09-12 `u_face`: the last 4 s (40 steps @10 Hz, t-3.9..t) of the 9 causal face-landmark channels (NaN = invalid) + per-step validity | |
| d = _rd(f["face"], ("uid", "t", "Fh8_c", "Fh8valid", "channels"), (lambda a: a), allow_pickle=True) | |
| assert np.array_equal(sub(d["uid"]), out["uid"]) and np.allclose(sub(d["t"]), out["t"]), f"{f['face']}: uid/t witness" | |
| assert [str(c) for c in d["channels"]] == FACE_CHANNELS | |
| out["Fh"] = np.ascontiguousarray(sub(d["Fh8_c"])[:, -HIST:, :]).astype(np.float16) | |
| out["Fhvalid"] = np.ascontiguousarray(sub(d["Fh8valid"])[:, -HIST:]).astype(bool) | |
| if "kpf" in need: | |
| # LANE F1 (2026-09-17): the DETACHED keypoint-forecaster channel of one KPF run (seed-matched by the caller), row-aligned with the grid (uid / t witness) | |
| assert f["kpf"] is not None, "stream kpf needs cwm3_dataset.set_kpf_dir(<KPF run dir>)" | |
| d = _rd(f["kpf"], ("route", "route_names", "t", "feat", "valid"), (lambda a: a), allow_pickle=True) | |
| kpf_uid = d["route_names"].astype(str)[d["route"]] # the KPF files carry the route index (uid == route name of the s12 grid) | |
| assert np.array_equal(sub(kpf_uid), out["uid"]) and np.allclose(sub(d["t"]), out["t"]), f"{f['kpf']}: uid/t witness" | |
| out["kpf"] = np.ascontiguousarray(sub(d["feat"])).astype(np.float16) | |
| out["kpf_valid"] = sub(d["valid"]).astype(bool) | |
| if "regionpca256" in need: | |
| assert pack_version() == "v2", "the regionpca256 road stream exists in pack v2 only (set_pack_version('v2'))" | |
| d = _rd(f["regionpca256"], ("vector", "ok"), (lambda a: a)) | |
| assert d["vector"].shape == (N_file, 512), f"{f['regionpca256']}: vector {d['vector'].shape} != ({N_file}, 512)" | |
| out["regionpca256"], out["regionpca256_ok"] = sub(d["vector"]).astype(np.float16), sub(d["ok"]).astype(bool) | |
| if "tokens" in need: | |
| ai = np.load(f["anchor_index"]) | |
| assert ai.shape == (N_file, NSLOT), f"{f['anchor_index']}: shape {ai.shape}" | |
| out["tok_idx"] = sub(ai).astype(np.int32) | |
| for k in out: | |
| assert k not in FORBIDDEN_KEYS | |
| _drop_cache(v for k, v in f.items() if k not in ("bank", "target_stats")) # arrays are in RSS now; the cached file pages are redundant | |
| out["meta"] = {"regime": regime, "split": split, "N_file": N_file, "N": N, "need": sorted(need), "pack_version": pack_version(), "kin_source": kin_source(), | |
| "kin_version": kin_version(), | |
| "files": {k: str(v) for k, v in f.items() if k != "bank"}, "roi": ROI, "dms_fields_used": [], | |
| "subsample": None if sel is None else {"max_anchors": max_anchors, "seed": seed, "rule": "uniform without replacement"}} | |
| verbose(f"[cwm3_dataset] {regime}/{split} N={N:,}/{N_file:,} need={sorted(need)} {time.time() - t0:.1f}s") | |
| return out | |
| def load_bank(need_tok=True, need_pooled=True): | |
| """results/banks/front_tok.npz -> dict {tok f16[T,4,64], pooled256 f16[T,256], t f32[T], uid_index, route_off, causality_rule}.""" | |
| with np.load(BANKS / "front_tok.npz", allow_pickle=True) as z: | |
| out = {"t": z["t"], "uid_index": z["uid_index"], "route_off": z["route_off"], "causality_rule": str(z["causality_rule"])} | |
| if need_tok: | |
| out["tok"] = z["tok"] | |
| if need_pooled: | |
| out["pooled256"] = z["pooled256"] | |
| _drop_cache([BANKS / "front_tok.npz"]) | |
| return out | |
| def concat_regimes(parts): | |
| """list of per-regime dicts (same split, same need) -> one dict with mode / regime_row / dongle_code / uid_code.""" | |
| keys = [k for k in parts[0] if k != "meta"] | |
| out = {} | |
| for k in keys: | |
| out[k] = np.concatenate([p[k] for p in parts]) | |
| out["regime_row"] = np.concatenate([p["idx"] for p in parts]).astype(np.int64) | |
| out["regime_id"] = np.concatenate([p["mode"].astype(np.int8) for p in parts]) | |
| dong, out["dongle_code"] = np.unique(out["dongle"], return_inverse=True) | |
| uids, out["uid_code"] = np.unique(out["uid"], return_inverse=True) | |
| out["dongle_code"] = out["dongle_code"].astype(np.int32) | |
| out["uid_code"] = out["uid_code"].astype(np.int32) | |
| out["meta"] = {"parts": [p["meta"] for p in parts], "N": int(len(out["t"])), "n_dongles": int(len(dong)), "n_uids": int(len(uids)), | |
| "dms_fields_used": [], "roi": ROI} | |
| return out | |
| def load_regimes(regimes, split, need=ALL_STREAMS, max_anchors=None, seed=0, verbose=print): | |
| parts = [load_split(r, split, need, max_anchors, seed + i, verbose) for i, r in enumerate(regimes)] | |
| return concat_regimes(parts) if len(parts) > 1 else dict(parts[0], regime_row=parts[0]["idx"], | |
| regime_id=np.full(parts[0]["meta"]["N"], MODE[regimes[0]], np.int8), | |
| dongle_code=np.unique(parts[0]["dongle"], return_inverse=True)[1].astype(np.int32), | |
| uid_code=np.unique(parts[0]["uid"], return_inverse=True)[1].astype(np.int32)) | |
| def subset(raw, idx): | |
| N = len(raw["t"]) | |
| return {k: (v[idx] if isinstance(v, np.ndarray) and v.ndim >= 1 and len(v) == N else v) for k, v in raw.items()} | |
| def nbytes(raw): | |
| return int(sum(v.nbytes for v in raw.values() if isinstance(v, np.ndarray))) | |
| def rss_gb(): | |
| try: | |
| import psutil | |
| return psutil.Process().memory_info().rss / 2**30 | |
| except ImportError: | |
| return float(open("/proc/self/statm").read().split()[1]) * os.sysconf("SC_PAGE_SIZE") / 2**30 | |
| # ----------------------------------------------------------------------------- synthetic (tests / smoke only) | |
| def synthetic_raw(N, seed=0, mode=0, need=ALL_STREAMS): | |
| """Compact-contract arrays with plausible marginals; NOT data, never a result.""" | |
| rng = np.random.default_rng(seed) | |
| X = np.zeros((N, HIST, DX), np.float32) | |
| X[:, :, 0] = np.cumsum(rng.normal(0, 3, (N, HIST)), 1) | |
| X[:, :, 1] = rng.random((N, HIST)) < 0.4 | |
| X[:, :, 2] = rng.random((N, HIST)) < 0.1 | |
| X[:, :, 3] = np.abs(rng.normal(12, 6, (N, 1)) + rng.normal(0, 0.3, (N, HIST))) | |
| X[:, :, 4] = rng.normal(0, 0.7, (N, HIST)) | |
| Y = np.zeros((N, H, DX), np.float32) | |
| Y[:, :, 0] = X[:, -1:, 0] + np.cumsum(rng.normal(0, 4, (N, H)), 1) | |
| Y[:, :, 1] = rng.random((N, H)) < 0.4 | |
| Y[:, :, 2] = rng.random((N, H)) < 0.1 | |
| Y[:, :, 3] = np.abs(X[:, -1:, 3] + np.cumsum(rng.normal(0, 0.2, (N, H)), 1)) | |
| Y[:, :, 4] = rng.normal(0, 0.7, (N, H)) | |
| M = np.ones((N, H), bool) | |
| M[rng.random(N) < 0.05, 15:] = False | |
| out = {"uid": np.array([f"route_{i % 5}" for i in range(N)]), "dongle": np.array([f"d{i % 3}" for i in range(N)]), | |
| "t": np.arange(N) * 0.2 + 10.0, "M": M, "slice_tag": rng.integers(0, 4, N).astype(np.uint8), "idx": np.arange(N, dtype=np.int64), | |
| "mode": np.full(N, mode, np.int8), "X": X, "Y": Y} | |
| s_prev, s = ep.phase_chain_targets(Y, X) | |
| out["phase0"], out["phase"] = s_prev[:, 0].astype(np.uint8), s.astype(np.uint8) | |
| out["Xpre"] = np.roll(X, 1, 0) | |
| out["x8_ok"] = rng.random(N) < 0.95 | |
| Kh = rng.normal(0, 1, (N, HIST, DK)).astype(np.float16) | |
| Kh[rng.random(N) < 0.08] = np.nan | |
| out["Kh"] = Kh | |
| out["K"] = rng.normal(0, 1, (N, H, DK)).astype(np.float16) | |
| out["Kvalid"] = rng.random((N, H)) < 0.9 | |
| S = (rng.random((N, H, 2)) < np.array([0.9, 0.15])).astype(np.float32) | |
| S[rng.random((N, H, 2)) < 0.05] = np.nan | |
| out["S"] = S.astype(np.float16) | |
| S0 = (rng.random((N, 2)) < np.array([0.9, 0.15])).astype(np.float32) | |
| S0[rng.random((N, 2)) < 0.05] = np.nan | |
| out["S0"] = S0 | |
| out["EM"] = np.isfinite(S[:, :, [0, 0, 1]]) | |
| out["D"] = rng.random((N, H, 6)) < 0.03 | |
| out["Dvalid"] = rng.random((N, H, 6)) < 0.9 | |
| out["D0"] = rng.random((N, 6)) < 0.03 | |
| out["D0valid"] = rng.random((N, 6)) < 0.9 | |
| out["Resp"] = rng.random((N, H)) < 0.3 | |
| out["RespValid"] = rng.random((N, H)) < 0.95 | |
| out["Resp0"] = rng.random(N) < 0.3 | |
| out["Resp0valid"] = rng.random(N) < 0.95 | |
| out["NoResp"] = rng.random((N, H)) < 0.5 | |
| out["NoRespValid"] = rng.random((N, H)) < 0.9 | |
| out["Moving0"] = rng.random(N) < 0.9 | |
| out["F"] = rng.normal(0, 1, (N, 5, 256)).astype(np.float16) | |
| out["Fok"] = rng.random((N, 5)) < 0.95 | |
| for s_ in ("vf256", "vc256", "crop256", "vfcan256", "can256", "vfcan256s"): | |
| ok = rng.random(N) < 0.9 | |
| out[s_] = (rng.normal(0, 1, (N, 256)) * ok[:, None]).astype(np.float16) | |
| out[s_ + "_ok"] = ok | |
| ok = rng.random(N) < 0.9 | |
| out["regionpca256"] = (rng.normal(0, 1, (N, 512)) * ok[:, None]).astype(np.float16) | |
| out["regionpca256_ok"] = ok | |
| out["P"], out["Pvalid"] = rng.random((N, H, 6)) < 0.03, rng.random((N, H, 6)) < 0.9 | |
| out["P0"], out["P0valid"] = rng.random((N, 6)) < 0.03, rng.random((N, 6)) < 0.9 | |
| Sg = (rng.random((N, H)) < 0.15).astype(np.float32) | |
| Sg[rng.random((N, H)) < 0.05] = np.nan | |
| out["Sg"] = Sg.astype(np.float16) | |
| Sg0 = (rng.random(N) < 0.15).astype(np.float32) | |
| Sg0[rng.random(N) < 0.05] = np.nan | |
| out["Sg0"] = Sg0 | |
| emg = np.isfinite(np.concatenate([Sg0[:, None], Sg[:, :-1]], 1)) & np.isfinite(Sg) | |
| out["EMg"] = np.stack([emg, emg], -1) | |
| Fh = rng.normal(0, 1, (N, HIST, N_FACE)).astype(np.float16) | |
| Fh[:, :, 7] = rng.random((N, HIST)) < 0.15 | |
| Fhv = rng.random((N, HIST)) < 0.95 | |
| Fh[~Fhv] = np.nan | |
| out["Fh"], out["Fhvalid"] = Fh, Fhv | |
| ti = np.tile(np.arange(NSLOT), (N, 1)) + (np.arange(N) % 64)[:, None] * NSLOT | |
| ti[rng.random((N, NSLOT)) < 0.05] = -1 | |
| out["tok_idx"] = ti.astype(np.int32) | |
| out["meta"] = {"regime": "synthetic", "split": "synthetic", "N_file": N, "N": N, "need": sorted(need), "dms_fields_used": [], "roi": ROI} | |
| keep = {"uid", "dongle", "t", "M", "slice_tag", "idx", "mode", "meta", "phase0", "phase", "Y"} | |
| for stream, keys in STREAM_KEYS.items(): | |
| if stream in need: | |
| keep |= set(keys) | |
| return {k: v for k, v in out.items() if k in keep} | |
| STREAM_KEYS = {"X": ("X",), "X8": ("X", "Xpre", "x8_ok"), "Kh": ("Kh", "kin_row_ok"), "kin": ("K", "Kvalid", "S", "S0", "EM", "kin_row_ok"), | |
| "demand": ("D", "Dvalid", "D0", "D0valid", "Resp", "RespValid", "Resp0", "Resp0valid", "NoResp", "NoRespValid", "Moving0"), | |
| "F": ("F", "Fok"), "vf256": ("vf256", "vf256_ok"), "vc256": ("vc256", "vc256_ok"), "crop256": ("crop256", "crop256_ok"), | |
| "tokens": ("tok_idx",), "Y": ("Y",), "regionpca256": ("regionpca256", "regionpca256_ok"), "pressure": ("P", "Pvalid", "P0", "P0valid"), | |
| "vfcan256": ("vfcan256", "vfcan256_ok"), "can256": ("can256", "can256_ok"), "vfcan256s": ("vfcan256s", "vfcan256s_ok"), | |
| "gaze": ("Sg", "Sg0", "EMg"), "face": ("Fh", "Fhvalid"), "kpf": ("kpf", "kpf_valid")} | |
| def synthetic_bank(n_rows=64 * NSLOT, seed=0): | |
| rng = np.random.default_rng(seed) | |
| return {"tok": rng.normal(0, 1, (n_rows, 4, 64)).astype(np.float16), "pooled256": rng.normal(0, 1, (n_rows, 256)).astype(np.float16), | |
| "t": (np.arange(n_rows) * 0.5).astype(np.float32), "uid_index": np.array(["synthetic"]), "route_off": np.array([0, n_rows]), | |
| "causality_rule": "synthetic"} | |
| # ----------------------------------------------------------------------------- CLI check | |
| def main(): | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--check", action="store_true") | |
| ap.add_argument("--regimes", default="manual,engaged") | |
| ap.add_argument("--split", default="train") | |
| ap.add_argument("--need", default="all") | |
| ap.add_argument("--max-anchors", type=int, default=0) | |
| a = ap.parse_args() | |
| need = ALL_STREAMS if a.need == "all" else frozenset(a.need.split(",")) | |
| r0 = rss_gb() | |
| t0 = time.time() | |
| raw = load_regimes(a.regimes.split(","), a.split, need, a.max_anchors or None) | |
| r1 = rss_gb() | |
| print(f"N={raw['meta']['N']:,} arrays={nbytes(raw) / 2**30:.2f} GiB RSS {r0:.2f} -> {r1:.2f} GiB ({time.time() - t0:.0f}s)") | |
| for k, v in sorted(raw.items()): | |
| if isinstance(v, np.ndarray): | |
| extra = "" | |
| if v.dtype.kind == "b": | |
| extra = f" mean={v.mean():.4f}" | |
| elif k in ("Kh", "K", "S"): | |
| extra = f" finite={np.isfinite(v.astype(np.float32)).mean():.4f}" | |
| elif k == "tok_idx": | |
| extra = f" all16={np.mean((v >= 0).all(1)):.4f}" | |
| print(f" {k:12s} {str(v.shape):22s} {str(v.dtype):8s} {v.nbytes / 2**20:9.1f} MiB{extra}") | |
| if a.check: | |
| for k in raw: | |
| assert k not in FORBIDDEN_KEYS | |
| print("no-DMS key check: OK") | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 31.9 kB
- Xet hash:
- 33b491f3c4d3c31c8cec568a571a932c030823ea005dc8f03f7127727bf88675
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.