Buckets:

AsianPlayer
/

TriDrive

Files
xet
AsianPlayer/TriDrive / code /baselines_dms /dataset_adapter.py
AsianPlayer
16 days ago
download
raw
25.9 kB
"""cwm3/baselines_dms/dataset_adapter.py — layer 1 of the DMS-baseline evaluation: feed a baseline's per-frame channels to the UNCHANGED U0 trainer in
place of the 11 driver-kinematics channels (plan §2 第一层 / §3 fit_u_base, 2026-09-14).
HOW THE UNCHANGED U0 TRAINER CONSUMES A SUBSTITUTE KINEMATICS PACK (verified by reading the code; file:line as of 2026-09-14)
  1. train/train_cwm3.py main() :606 parses `--kin-source {cwm2,causal_routes,oldin_ctgt}` :622 and `--kin-version {v1,v2}` :624 and calls
     ds.set_kin_source(args.kin_source, args.kin_version) :632.  There is NO flag that names a directory.
  2. data/cwm3_dataset.py: set_kin_source :65 stores the pair; kin_dir() :82-84 returns the MODULE CONSTANT KINC_V2 :61 (= results/targets_kin_causal_routes_v2)
     when kin_version == "v2", else KINC :60; files() :120-121 builds
         "kinhist": kin_dir()/f"kinhist_causal_{regime}_{split}.npz"   (when kin_source() == "causal_routes")      <- the INPUT history Kh
         "kin":     kin_dir()/f"kin_causal_{regime}_{split}.npz"       (causal_routes | oldin_ctgt)                <- the TARGETS K / S / S0 / EM
     load_split :170-178 (stream "Kh") reads exactly ("uid", "t", "Kh_c", "Khvalid", "row_ok") :174, asserts the uid / t witness against the s12 targets,
     sets Kh[~Khvalid] = NaN :177 and Kh[~row_ok] = NaN :178, stores f16 + kin_row_ok.  The kinhist file's `channels` key is NOT asserted (only the
     targets file is: :189 `== KIN_CHANNELS`), so the 11 input channels may be anything as long as the SHAPE is kept.
  3. train/train_cwm3.py need_streams :100 ("U": X, Kh, kin, Y) -> both files are read; stage_input_files :116 hashes them into run.json; run.json
     records kin_dir = str(ds.kin_dir()) :734.  The GPU child is created with os.fork() :557, so module state set in the parent is inherited.
  4. train/batch.py: DK = 11 :25; compute_stats :66 (Kh branch :82) = nan-aware float64 mean / sd over Kh.reshape(-1, DK) on TRAIN rows (both regimes);
     prep_U :207 -> std_kh :195: valid = isfinite(Kh[..., 10]) & isfinite(Kh[..., 0]) :197, Z = nan_to_num(clamp((Kh - kmu) / ksd, -clip, clip)) with
     clip = cfg["kclip"] (:202; 8 for kin_can_noK_clip8_ln, train/stages.yaml :52), Xs = clamp((X - xmu) / xsd, -8, 8), dt = -(39 - k) * 0.1;
     Kx = cat[Z (11), valid (1), Xs (5), dt (1)] = [B, 40, 18] :216.  models/branches.py DriverBranch: N_KIN_IN = 12 :89, cin = 18 :91 (input LayerNorm
     over the first 12 channels when u_in_ln).
  => CONTRACT of the substitute pack  <sub_dir>/kinhist_causal_{regime}_{split}.npz  (this module writes it):
         uid, t                 == s12_wm/results/targets_{regime}_{split}.npz (byte-equal)
         Kh_c    f16[N, 40, 11] the baseline channels at t-3.9..t @ 10 Hz, EXACTLY 11 columns: C <= 11 baseline channels first, then `pad_i` columns
                                (0.0 where the step is valid, NaN where not) — or, for C > 11 (VideoMAEv2 384-d, OMZ action logits), the 11 leading PCA
                                components fitted on TRAIN valid frames (PCAReducer; recorded in `adapter`); columns 0 and 10 finite <=> the step is valid
                                (std_kh's validity rule), NaN elsewhere
         Khvalid bool[N, 40]    the step has a valid baseline output;  row_ok bool[N]  the route has a cache and a frame at t
         channels str[11], units, source "baselines_dms:<b>", baseline, adapter json (channel selection / pads / reducer / cache sha256s)
     and  <sub_dir>/kin_causal_{regime}_{split}.npz = SYMLINKS to the ORIGINAL causal targets of results/targets_kin_causal_routes_v2 (link_targets):
     the driver-onset targets (hands_on / head_turned chains, K residual) stay exactly those of the U0 reference run.
  => MINIMAL HOOK (trainer, data layer, batch builder and model all UNCHANGED): point the module constant at the substitute directory in-process, then
     run train_cwm3.main() with the survivor recipe flags:
         import cwm3_dataset as ds; ds.KINC_V2 = Path(sub_dir)                       # (KINC for --kin-version v1)
         sys.argv = trainer_argv(seed, out_dir)   # train_cwm3.py --stage U0 --variant kin_can_noK_clip8_ln --seed k --kin-source causal_routes
                                                  #   --kin-version v2 --stats-version 2 --out results/runs/U0_<b>_s<k> [--minutes 18]
         import train_cwm3; train_cwm3.main()
     (launch_snippet() prints the exact one-liner; fit_u_base.py wraps it).  train/train_u0_from_init.py forwards only --kin-source / --kin-version (no
     directory) and warning_features_ext.dwell_block(kin_dir=) is the readout, not the trainer — neither is a substitute hook.
  Standardisation: the trainer computes kmu / ksd itself from the TRAIN rows of the substitute pack (compute_stats v2), so this module stores RAW
  channel values; build_kx_numpy() is the numpy mirror of prep_U (tests / documentation) and compute_k_stats / compute_x_stats mirror compute_stats.
Nothing here reads a video, a model or any openpilot field; the pack directory lives under results/baselines_dms/ (never inside results/targets_*).
  python baselines_dms/dataset_adapter.py --self-test
  python baselines_dms/dataset_adapter.py --baseline l2cs --regime manual --split val --out-dir results/baselines_dms/kin_sub/l2cs [--cache-dir ...] [--reducer ...]
"""
from __future__ import annotations

import json
import os
import sys
import time
from pathlib import Path

import numpy as np

PKG = Path(__file__).resolve().parent
CWM3 = PKG.parent
EXPLORE = CWM3.parent
S12R = EXPLORE / "s12_wm" / "results"
for p in (str(PKG), str(CWM3 / "data")):
    if p not in sys.path:
        sys.path.insert(0, p)
import cache_schema as cs                                   # noqa: E402
import summarise as SM                                      # noqa: E402

DK, HIST, HIST8, DX = 11, 40, 80, 5
X_CLIP = 8.0
TRAINER = CWM3 / "train" / "train_cwm3.py"
VARIANT = "kin_can_noK_clip8_ln"                            # the U0 survivor recipe (train/stages.yaml :52): noK, kclip 8, input LayerNorm(12)
KIN_V2 = CWM3 / "results" / "targets_kin_causal_routes_v2"
KIN_V1 = CWM3 / "results" / "targets_kin_causal_routes"
SUB_ROOT = cs.RESULTS_ROOT / "kin_sub"
PAD_PREFIX = "pad_"


# ----------------------------------------------------------------------------- numpy mirrors of train/batch.py (documentation + tests)
def _nan_stats(A, D):
    sm, sq, cnt = np.zeros(D, np.float64), np.zeros(D, np.float64), np.zeros(D, np.float64)
    A = A.reshape(-1, D)
    for i in range(0, len(A), 1_000_000):
        kb = A[i:i + 1_000_000].astype(np.float64)
        ok = np.isfinite(kb)
        sm += np.where(ok, kb, 0.0).sum(0)
        sq += np.where(ok, kb ** 2, 0.0).sum(0)
        cnt += ok.sum(0)
    mu = sm / np.maximum(cnt, 1)
    var = sq / np.maximum(cnt, 1) - mu ** 2
    return mu.astype(np.float32), (np.sqrt(np.maximum(var, 0)) + 1e-6).astype(np.float32)


def compute_k_stats(Kh):
    """train/batch.compute_stats v2 rule for Kh (nan-aware float64 over all steps of the TRAIN rows) -> (kmu[11], ksd[11])."""
    return _nan_stats(np.asarray(Kh), DK)


def compute_x_stats(X):
    """train/batch.compute_stats v2 rule for X (float64 over all 40 steps) -> (xmu[5], xsd[5])."""
    X = np.asarray(X).reshape(-1, DX)
    mu, var = np.zeros(DX, np.float64), np.zeros(DX, np.float64)
    for i in range(0, len(X), 1_000_000):
        xb = X[i:i + 1_000_000].astype(np.float64)
        mu += xb.sum(0)
        var += (xb ** 2).sum(0)
    mu /= len(X)
    var = var / len(X) - mu ** 2
    return mu.astype(np.float32), (np.sqrt(np.maximum(var, 0)) + 1e-6).astype(np.float32)


def build_kx_numpy(Kh, X, st, clip=8.0):
    """numpy mirror of train/batch.prep_U (:207-216) WITHOUT the opt-in face block: Kh f32[N,40,11] (NaN invalid), X f32[N,40,5] -> Kx f32[N,40,18]
    = [standardised + clamped Kh (NaN -> 0), valid bit (Kh[...,10] & Kh[...,0] finite), standardised CAN X (clamp 8), dt = -(39-k)*0.1]."""
    Kh = np.asarray(Kh, np.float32)
    X = np.asarray(X, np.float32)
    N = len(Kh)
    valid = np.isfinite(Kh[..., 10]) & np.isfinite(Kh[..., 0])
    with np.errstate(invalid="ignore"):
        Z = np.nan_to_num(np.clip((Kh - st["kmu"]) / st["ksd"], -clip, clip))
        Xs = np.clip((X - st["xmu"]) / st["xsd"], -X_CLIP, X_CLIP)
    dt = np.broadcast_to((-(39 - np.arange(HIST)) * 0.1).astype(np.float32)[None, :, None], (N, HIST, 1))
    return np.concatenate([Z, valid[..., None].astype(np.float32), Xs, dt], -1).astype(np.float32)


# ----------------------------------------------------------------------------- channel reduction (C > 11)
class PCAReducer:
    """C -> n_out linear map fitted on TRAIN valid frames (mean-centred PCA, np.linalg.eigh); saved as npz (mu, W, evr, channels, n_fit, fit_uids)."""

    def __init__(self, mu, W, evr, channels, n_fit=0, fit_uids=(), baseline=""):
        self.mu, self.W, self.evr, self.channels, self.n_fit, self.fit_uids, self.baseline = np.asarray(mu, np.float32), np.asarray(W, np.float32), np.asarray(evr, np.float32), list(channels), int(n_fit), list(fit_uids), baseline

    @classmethod
    def fit(cls, F, channels, n_out=DK, baseline="", fit_uids=()):
        F = np.asarray(F, np.float64)
        F = F[np.isfinite(F).all(1)]
        mu = F.mean(0)
        cov = (F - mu).T @ (F - mu) / max(len(F), 1)
        evals, evecs = np.linalg.eigh(cov)
        order = np.argsort(evals)[::-1]
        evals, evecs = evals[order], evecs[:, order]
        return cls(mu, evecs[:, :n_out], evals[:n_out] / max(evals.sum(), 1e-12), channels, len(F), fit_uids, baseline)

    def transform(self, F):
        F = np.asarray(F, np.float32)
        out = (np.nan_to_num(F) - self.mu) @ self.W
        out[~np.isfinite(F).all(-1)] = np.nan
        return out.astype(np.float32)

    @property
    def out_names(self):
        return [f"pca{j}" for j in range(self.W.shape[1])]

    def save(self, path):
        path = Path(path)
        path.parent.mkdir(parents=True, exist_ok=True)
        np.savez(path, mu=self.mu, W=self.W, evr=self.evr, channels=np.array(self.channels), n_fit=np.int64(self.n_fit), fit_uids=np.array(self.fit_uids), baseline=np.array(self.baseline),
                 rule=np.array("PCA fitted on TRAIN-split valid frames of the baseline caches only"))
        return path

    @classmethod
    def load(cls, path):
        with np.load(path, allow_pickle=False) as z:
            return cls(z["mu"], z["W"], z["evr"], [str(c) for c in z["channels"]], int(z["n_fit"]), [str(u) for u in z["fit_uids"]], str(z["baseline"]))

    def describe(self):
        return {"kind": "pca", "n_in": len(self.channels), "n_out": int(self.W.shape[1]), "evr_sum": float(self.evr.sum()), "n_fit": self.n_fit, "n_fit_uids": len(self.fit_uids), "baseline": self.baseline}


def train_uids():
    out = set()
    for r in ("manual", "engaged"):
        with np.load(S12R / f"targets_{r}_train.npz", allow_pickle=True) as z:
            out |= set(np.unique(z["uid"]).tolist())
    return out


def fit_reducer(baseline, cache_dir=None, n_out=DK, max_frames=200_000, seed=0, out_path=None, log=print):
    """PCA-n_out over the valid frames of the TRAIN-split route caches of `baseline` (never val / test) -> PCAReducer (saved to out_path when given)."""
    cache_dir = Path(cache_dir) if cache_dir else cs.cache_dir(baseline)
    tr = train_uids()
    rng = np.random.default_rng(seed)
    parts, used, chans = [], [], None
    files = sorted(p for p in cache_dir.glob("*.npz") if p.stem in tr)
    per = max(max_frames // max(len(files), 1), 200)
    for p in files:
        rc = cs.RouteCache.load(p)
        chans = chans or rc.channels
        assert rc.channels == chans, p
        F = rc.feat[rc.valid].astype(np.float32)
        if len(F):
            parts.append(F[np.sort(rng.choice(len(F), min(per, len(F)), replace=False))])
            used.append(rc.uid)
    assert parts, f"no TRAIN caches with valid frames in {cache_dir}"
    red = PCAReducer.fit(np.concatenate(parts), chans, n_out, baseline, used)
    log(f"reducer {baseline}: {red.describe()}")
    if out_path:
        red.save(out_path)
    return red


# ----------------------------------------------------------------------------- pack writer
def route_hist(cache, anchors_t, hist=HIST):
    """cache + anchor times -> (Kh f32[n,hist,C] NaN invalid, Kv bool[n,hist], row_ok bool[n]) on the 10 Hz grid t-(hist-1)/10..t (causal; summarise.window_index asserts)."""
    D = SM.dense_series(cache)
    idx = SM.window_index(D["k0"], D["M"], anchors_t, hist)
    Gh, Vh = SM.gather(D["G"], D["V"], idx)
    ka = SM.anchor_steps(D["k0"], anchors_t)
    inside = (ka >= 0) & (ka < D["M"])
    row_ok = inside.copy()
    row_ok[inside] = D["F"][ka[inside]]
    return Gh, Vh, row_ok


def select_channels(channels, want=None, reducer=None):
    """-> (indices into the cache channels, output channel names[11]).  C <= 11: all (or `want`) + pads; C > 11: a reducer is required."""
    ch = list(channels)
    if reducer is not None:
        assert reducer.channels == ch, "reducer was fitted on different channels"
        names = reducer.out_names
        assert len(names) == DK, len(names)
        return None, names
    sel = [ch.index(c) for c in want] if want else list(range(len(ch)))
    assert len(sel) <= DK, f"{len(sel)} channels > {DK}: pass a PCAReducer (fit_reducer) or a `channels` subset"
    names = [ch[i] for i in sel] + [f"{PAD_PREFIX}{j}" for j in range(len(sel), DK)]
    return sel, names


def to_dk(Gh, Vh, sel, reducer):
    """[n,T,C] -> [n,T,11]: selected channels + pads (0 where valid), or reducer output; NaN wherever the step is invalid or any used channel is non-finite."""
    n, T, C = Gh.shape
    if reducer is not None:
        Z = reducer.transform(Gh.reshape(-1, C)).reshape(n, T, DK)
    else:
        Z = np.full((n, T, DK), np.nan, np.float32)
        Z[:, :, :len(sel)] = Gh[:, :, sel]
        Z[:, :, len(sel):] = 0.0
    ok = Vh & np.isfinite(Z).all(-1)
    Z[~ok] = np.nan
    return Z.astype(np.float32), ok


def write_substitute_kin_pack(baseline, regime, split, out_dir=None, cache_dir=None, channels=None, reducer=None, with_hist8=False, log=print):
    """-> path of <out_dir>/kinhist_causal_{regime}_{split}.npz (see the module docstring for the contract).  reducer: PCAReducer | path | None."""
    t0 = time.time()
    cache_dir = Path(cache_dir) if cache_dir else cs.cache_dir(baseline)
    out_dir = Path(out_dir) if out_dir else SUB_ROOT / baseline
    if isinstance(reducer, (str, Path)):
        reducer = PCAReducer.load(reducer)
    with np.load(S12R / f"targets_{regime}_{split}.npz", allow_pickle=True) as z:
        uid, t = z["uid"], z["t"]
    N = len(uid)
    Kh = np.full((N, HIST, DK), np.nan, np.float16)
    Kv = np.zeros((N, HIST), bool)
    Kh8 = np.full((N, HIST8, DK), np.nan, np.float16) if with_hist8 else None
    Kv8 = np.zeros((N, HIST8), bool) if with_hist8 else None
    row_ok = np.zeros(N, bool)
    route_dec = np.zeros(N, bool)
    names = None
    shas = {}
    n_routes = n_dec = 0
    for u, rows in SM._route_blocks(uid):
        n_routes += 1
        p = cache_dir / f"{u}.npz"
        if not p.exists():
            continue
        rc = cs.RouteCache.load(p)
        sel, nm = select_channels(rc.channels, channels, reducer)
        if names is None:
            names = nm
        assert nm == names, f"{u}: channel set differs"
        shas[u] = cs.sha256_file(p)
        ts = t[rows].astype(np.float64)
        Gh, Vh, ok = route_hist(rc, ts, HIST)
        Z, V = to_dk(Gh, Vh, sel, reducer)
        Kh[rows], Kv[rows], row_ok[rows] = Z.astype(np.float16), V, ok
        if with_hist8:
            Gh8, Vh8, _ = route_hist(rc, ts, HIST8)
            Z8, V8 = to_dk(Gh8, Vh8, sel, reducer)
            Kh8[rows], Kv8[rows] = Z8.astype(np.float16), V8
        route_dec[rows] = True
        n_dec += 1
    if names is None:
        names = [f"{PAD_PREFIX}{j}" for j in range(DK)]
    Kh[~row_ok] = np.nan
    Kv[~row_ok] = False
    adapter = {"baseline": baseline, "cache_dir": str(cache_dir), "channels_used": [c for c in names if not c.startswith(PAD_PREFIX) and not c.startswith("pca")],
               "pads": [c for c in names if c.startswith(PAD_PREFIX)], "pad_rule": "0.0 where the step is valid, NaN where not (std_kh validity = channels 0 and 10 finite)",
               "reducer": (reducer.describe() if reducer is not None else None), "hist_steps": HIST, "hz": SM.HZ, "source_cache_sha256": shas,
               "trainer_contract": "data/cwm3_dataset.load_split reads uid, t, Kh_c, Khvalid, row_ok (kin_source causal_routes; kin_dir() -> this directory via ds.KINC_V2)",
               "written_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())}
    out_dir.mkdir(parents=True, exist_ok=True)
    path = out_dir / f"kinhist_causal_{regime}_{split}.npz"
    extra = {} if not with_hist8 else {"Kh8_c": Kh8, "Kh8valid": Kv8}
    np.savez_compressed(path, uid=uid, t=t, row_ok=row_ok, route_decoded=route_dec, Kh_c=Kh, Khvalid=Kv, channels=np.array(names),
                        units=np.array(f"raw {baseline} channels (standardised by the trainer's TRAIN statistics); pads 0"), source=np.array(f"baselines_dms:{baseline}"),
                        baseline=np.array(baseline), adapter=np.array(json.dumps(adapter)), **extra)
    log(f"{baseline} {regime}/{split}: N={N:,} routes {n_routes} decoded {n_dec} row_ok {row_ok.mean():.4f} Khvalid|row_ok {Kv[row_ok].mean() if row_ok.any() else float('nan'):.4f} "
        f"channels {names} {time.time() - t0:.1f}s -> {path}")
    return path


def link_targets(out_dir, kin_version="v2", regimes=("manual", "engaged"), splits=("train", "val", "test")):
    """symlink the ORIGINAL causal target files kin_causal_{r}_{s}.npz of results/targets_kin_causal_routes[_v2] into out_dir (the trainer reads them from kin_dir())."""
    src_dir = KIN_V2 if kin_version == "v2" else KIN_V1
    out_dir = Path(out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    links = []
    for r in regimes:
        for s in splits:
            src = src_dir / f"kin_causal_{r}_{s}.npz"
            dst = out_dir / src.name
            if not dst.exists():
                if dst.is_symlink():
                    dst.unlink()
                os.symlink(src, dst)
            links.append(dst)
    return links


def trainer_argv(seed, out_dir, variant=VARIANT, kin_version="v2", minutes=18, extra=()):
    """the exact train_cwm3.py argv of a substitute-kinematics U0 run (survivor recipe; stats v2; pack default)."""
    return [str(TRAINER), "--stage", "U0", "--variant", variant, "--seed", str(int(seed)), "--kin-source", "causal_routes", "--kin-version", kin_version,
            "--stats-version", "2", "--out", str(out_dir), "--minutes", str(minutes), *extra]


def launch_snippet(sub_dir, seed, out_dir, variant=VARIANT, kin_version="v2", minutes=18):
    """the in-process hook (trainer unchanged): repoint cwm3_dataset.KINC_V2 (or KINC) at sub_dir, then train_cwm3.main() with trainer_argv."""
    const = "KINC_V2" if kin_version == "v2" else "KINC"
    argv = trainer_argv(seed, out_dir, variant, kin_version, minutes)
    return ("import sys; from pathlib import Path; sys.path[:0] = [%r, %r, %r]; import cwm3_dataset as ds; ds.%s = Path(%r); "
            "import train_cwm3; sys.argv = %r; train_cwm3.main()" % (str(CWM3 / "data"), str(CWM3 / "train"), str(CWM3 / "common"), const, str(sub_dir), argv))


def load_via_dataset(sub_dir, regime, split, kin_version="v2", need=("X", "Kh"), max_anchors=None):
    """load one (regime, split) through the UNCHANGED data layer with kin_dir() repointed at sub_dir (constant restored afterwards) -> raw dict (Kh f16[N,40,11], kin_row_ok, X, ...)."""
    import cwm3_dataset as ds
    sub_dir = Path(sub_dir)
    link_targets(sub_dir, kin_version)
    const = "KINC_V2" if kin_version == "v2" else "KINC"
    saved = (getattr(ds, const), dict(ds._KIN_SOURCE))
    try:
        setattr(ds, const, sub_dir)
        ds.set_kin_source("causal_routes", kin_version)
        assert ds.kin_dir() == sub_dir and ds.files(regime, split)["kinhist"].parent == sub_dir
        return ds.load_split(regime, split, frozenset(need), max_anchors=max_anchors, verbose=lambda m: None)
    finally:
        setattr(ds, const, saved[0])
        ds._KIN_SOURCE.update(saved[1])


# ----------------------------------------------------------------------------- self-test
def _self_test():
    import tempfile
    with np.load(S12R / "targets_manual_val.npz", allow_pickle=True) as z:
        uid, t = z["uid"], z["t"]
    routes = [str(u) for u in np.unique(uid)[:2]]
    rep = {"self_test": "dataset_adapter", "routes": routes}
    with tempfile.TemporaryDirectory() as d:
        cdir = Path(d) / "cache" / "synthetic"
        for i, u in enumerate(routes):
            ts = t[uid == u].astype(np.float64)
            cs.synthetic_route_cache(uid=u, t0=max(0.0, ts.min() - 8.2), t1=ts.max() + 0.2, channels=("yaw_deg", "pitch_deg", "closed_p", "conf"), seed=i, meta={"baseline": "synthetic"}).save(cdir / f"{u}.npz")
        sub = Path(d) / "kin_sub"
        p = write_substitute_kin_pack("synthetic", "manual", "val", sub, cache_dir=cdir, log=lambda m: rep.__setitem__("writer_log", m))
        with np.load(p, allow_pickle=True) as z:
            assert np.array_equal(z["uid"], uid) and np.array_equal(z["t"], t) and z["Kh_c"].shape == (len(uid), HIST, DK) and z["Kh_c"].dtype == np.float16
            ok = z["row_ok"]
            Kh, V = z["Kh_c"].astype(np.float32), z["Khvalid"]
            assert ok.any() and not ok[~np.isin(uid, routes)].any() and np.isfinite(Kh[V][:, [0, DK - 1]]).all() and np.isnan(Kh[~V]).all()
            names = [str(c) for c in z["channels"]]
            rep.update(n_rows=int(len(uid)), rows_ok=int(ok.sum()), channels=names, adapter=json.loads(str(z["adapter"]))["pad_rule"])
        raw = load_via_dataset(sub, "manual", "val")
        assert raw["Kh"].shape == (len(uid), HIST, DK) and raw["kin_row_ok"].sum() == ok.sum() and raw["meta"]["kin_source"] == "causal_routes"
        rep["data_layer"] = {"Kh": list(raw["Kh"].shape), "kin_row_ok": int(raw["kin_row_ok"].sum()), "kinhist_file": raw["meta"]["files"]["kinhist"], "kin_file": raw["meta"]["files"]["kin"]}
        # numpy mirror vs the real prep_U on CPU
        try:
            import torch
            sys.path.insert(0, str(CWM3 / "train"))
            import batch as bt
            rows = np.flatnonzero(raw["kin_row_ok"])[:512]
            st = dict(zip(("kmu", "ksd"), compute_k_stats(raw["Kh"][raw["kin_row_ok"]])))
            st.update(zip(("xmu", "xsd"), compute_x_stats(raw["X"])))
            Kx_np = build_kx_numpy(raw["Kh"][rows].astype(np.float32), raw["X"][rows], st, clip=8.0)
            store = bt.GpuStore({"t": raw["t"][rows], "Kh": raw["Kh"][rows], "X": raw["X"][rows], "mode": raw["mode"][rows]}, ["Kh", "X", "mode"], torch.device("cpu"))
            b, _ = bt.prep_U(store, torch.arange(len(rows)), st, {"kclip": 8.0}, need_targets=False)
            assert b["Kx"].shape == (len(rows), HIST, DK + 1 + DX + 1) and np.allclose(b["Kx"].numpy(), Kx_np, atol=1e-5)
            rep["prep_U_mirror"] = {"rows": int(len(rows)), "Kx_shape": list(b["Kx"].shape), "max_abs_diff": float(np.abs(b["Kx"].numpy() - Kx_np).max())}
        except ImportError as e:
            rep["prep_U_mirror"] = f"skipped ({e})"
        # reducer path (C > 11): 16 synthetic channels -> 11 PCA components
        cdir2 = Path(d) / "cache" / "wide"
        u = routes[0]
        ts = t[uid == u].astype(np.float64)
        cs.synthetic_route_cache(uid=u, t0=max(0.0, ts.min() - 8.2), t1=ts.max() + 0.2, channels=tuple(f"e{j}_p" for j in range(16)), seed=5, meta={"baseline": "wide"}).save(cdir2 / f"{u}.npz")
        rc = cs.RouteCache.load(cdir2 / f"{u}.npz")
        red = PCAReducer.fit(rc.feat[rc.valid].astype(np.float32), rc.channels, DK, "wide", [u])
        p2 = write_substitute_kin_pack("wide", "manual", "val", Path(d) / "kin_sub_wide", cache_dir=cdir2, reducer=red, log=lambda m: None)
        with np.load(p2, allow_pickle=True) as z:
            assert z["Kh_c"].shape == (len(uid), HIST, DK) and [str(c) for c in z["channels"]] == red.out_names
        rep["reducer"] = red.describe()
        rep["launch_snippet"] = launch_snippet("results/baselines_dms/kin_sub/<b>", 0, "results/runs/U0_<b>_s0")
    print(json.dumps(rep, indent=1, default=str))


if __name__ == "__main__":
    import argparse
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("--self-test", action="store_true")
    ap.add_argument("--baseline", default="")
    ap.add_argument("--regime", default="manual")
    ap.add_argument("--split", default="val")
    ap.add_argument("--out-dir", default="")
    ap.add_argument("--cache-dir", default="")
    ap.add_argument("--channels", default="", help="comma list of cache channels to use (<= 11); default all")
    ap.add_argument("--reducer", default="", help="PCAReducer npz for C > 11 (fit with --fit-reducer)")
    ap.add_argument("--fit-reducer", action="store_true", help="fit the PCA-11 reducer on TRAIN caches and save it to <out-dir>/reducer.npz")
    ap.add_argument("--with-hist8", action="store_true")
    a = ap.parse_args()
    if a.baseline:
        out_dir = Path(a.out_dir) if a.out_dir else SUB_ROOT / a.baseline
        red = None
        if a.fit_reducer:
            red = fit_reducer(a.baseline, a.cache_dir or None, out_path=out_dir / "reducer.npz")
        elif a.reducer:
            red = a.reducer
        write_substitute_kin_pack(a.baseline, a.regime, a.split, out_dir, a.cache_dir or None, a.channels.split(",") if a.channels else None, red, a.with_hist8)
        link_targets(out_dir)
        print("launch:", launch_snippet(out_dir, 0, CWM3 / "results" / "runs" / f"U0_{a.baseline}_s0"))
    if a.self_test or not a.baseline:
        _self_test()
Xet Storage Details

Size:: 25.9 kB
Xet hash:: e2fb405b896a74c3fdec357aea677ec19978d3373b09365154926c59237123a0
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.