Buckets:

illusion-brain
/

brain-domains

Files

xet

illusion-brain/brain-domains / scripts /exp68_train_encoding.py

redmingwei

4 days ago

download

raw

18 kB

#!/usr/bin/env python3
"""exp68 — Tier 0 编码的 train split 版：1654 个概念。
## 为什么做这个
exp67 在 **test split（200 个概念）** 上得到：ViT-H-14 特征能预测 EEG 脑响应，
R² 中位 +0.0196，10/10 被试显著，且约等于跨被试脑对脑上限（比值中位 1.03）。
但 200 张图配 1024 维特征是**严重欠定**的（选出的 α 全是 30，强正则），
R² 绝对值只有 1%–4%。train split 有 **1654 个概念**，样本量 ×8，
编码强度应当显著提高 —— 这是判断「模型特征究竟能解释多少脑活动」的关键数字。
## 三处与 exp67 不同的设计
1. **primal 形式**：n=1654 > d=1024，Gram 矩阵不再是最小的那一边。
   XᵀX 是 1024×1024，解一次 1e9 flops。
2. **置换零改为打乱训练集的 Y 行**：X 完全不动 ⇒ 同一折的 XᵀX 复用，
   每次置换省掉一次 1.4e9 的矩阵乘（总省 ~40%）。
   在零假设「X 不含 Y 的信息」下，这与打乱「特征↔图像对应关系」是等价的。
3. **试次平均响应缓存到桶**：每个被试的 (1654, 630) float32 只有 4.2 MB。
   原始 npz 是 4.17 GB/被试，BBB 又需要所有被试 —— 先各自算好试次平均并落盘，
   再从缓存里做 BBB，内存峰值从 42 GB 降到 ~5 GB。
## 必须先验证顺序对齐
train split 的 ViT 特征（`ViT-H-14_features_train.pt`）与
`omnibrain4d/eeg_train/` 的概念顺序**必须对齐**，否则整个实验作废。
本脚本第一步就用 exp66 的办法独立验证：拿 train .pt 里的概念文本构造字符串 RDM，
与 ViT text_features 的 RDM 求相关 + 置换零。**这一步不通过，后面都不用跑。**
## 判据
    顺序对齐   r(字符串RDM, ViT text RDM) 显著高于置换零
    编码       R² img 显著高于置换零，且 img/BBB 比值可解读
    跨被试     10 个被试中显著者占比
"""

from __future__ import annotations

import json
import os
import re
import sys
import time
from collections import Counter
from pathlib import Path

import numpy as np

for _v in ("OMP_NUM_THREADS", "OPENBLAS_NUM_THREADS", "MKL_NUM_THREADS"):
    os.environ.setdefault(_v, "8")

FS = 250.0
WIN = 25
BASELINE = 12
N_FOLDS = 5
N_SUBJECTS = 10
N_NULL = 20
ALPHAS = [0.01, 0.1, 1.0, 3.0, 10.0, 30.0, 100.0, 300.0, 1000.0, 3000.0]
SEED = 0

TRAIN_PT = "/training/eeg/things-eeg/Preprocessed_data_250Hz_whiten/sub-01/train.pt"
FEAT_TRAIN_PT = "/training/eeg/eeg-image-decode/ViT-H-14_features_train.pt"
EegTRAIN = Path("/domains/omnibrain4d/eeg_train")
OUT_DIR = Path("/domains/runs/tier0-encoding")
TA_CACHE = Path("/domains/runs/tier0-encoding/cache")


def log(*a):
    print(*a, flush=True)


# ================================================================ 工具
def make_folds(n, k, rng):
    idx = rng.permutation(n)
    return [(np.setdiff1d(idx, te), te) for te in np.array_split(idx, k)]


def _rank(a):
    from scipy.stats import rankdata
    return rankdata(a)


def spearman(a, b):
    ra = _rank(a).astype(np.float64)
    rb = _rank(b).astype(np.float64)
    ra -= ra.mean()
    rb -= rb.mean()
    den = np.sqrt((ra * ra).sum() * (rb * rb).sum())
    return float((ra * rb).sum() / den) if den > 0 else 0.0


def rdm_full(mat):
    m = np.asarray(mat, dtype=np.float64)
    m = m - m.mean(axis=1, keepdims=True)
    m = m / (np.linalg.norm(m, axis=1, keepdims=True) + 1e-12)
    r = m @ m.T
    np.fill_diagonal(r, 1.0)
    return 1.0 - r


def tri(R):
    return R[np.triu_indices_from(R, k=1)]


def tokenize(s):
    return re.findall(r"[a-z0-9]+", s.lower())


def tfidf_matrix(docs):
    toks = [tokenize(d) for d in docs]
    df = Counter()
    for t in toks:
        df.update(set(t))
    vocab = sorted(df.keys())
    idx = {w: i for i, w in enumerate(vocab)}
    n = len(docs)
    M = np.zeros((n, len(vocab)), dtype=np.float64)
    for i, t in enumerate(toks):
        for w in t:
            M[i, idx[w]] += 1.0
    dfv = np.array([df[w] for w in vocab], dtype=np.float64)
    idf = np.log((1.0 + n) / (1.0 + dfv)) + 1.0
    M *= idf
    M /= (np.linalg.norm(M, axis=1, keepdims=True) + 1e-12)
    return M


# ================================================================ PART 0
def part0_probe():
    """探查 train split 的数据 + **独立验证顺序对齐**。"""
    import torch
    log("=" * 80)
    log("PART 0 — train split 探查 + 顺序对齐独立验证")
    log("=" * 80)
    info = {}

    f = torch.load(FEAT_TRAIN_PT, map_location="cpu", weights_only=False)
    log(f"\n[ViT train 特征] keys={list(f.keys())}")
    for k in f.keys():
        sh = tuple(f[k].shape)
        info[f"feat_{k}_shape"] = list(sh)
        log(f"  {k}: {sh} {f[k].dtype}")
    IMG = f["img_features"].numpy().astype(np.float64)
    TXT = f["text_features"].numpy().astype(np.float64)
    n_concept = IMG.shape[0]

    log(f"\n[原始 train.pt] {TRAIN_PT}")
    d = torch.load(TRAIN_PT, map_location="cpu", weights_only=False)
    log(f"  keys = {list(d.keys())}")
    for k in d.keys():
        v = d[k]
        if hasattr(v, "shape"):
            log(f"  {k}: {tuple(v.shape)} {getattr(v,'dtype','')}")
    eeg_shape = tuple(d["eeg"].shape)
    info["raw_train_eeg_shape"] = list(eeg_shape)
    lab = np.asarray(d["label"])
    txt = np.asarray(d["text"]) if "text" in d else None
    imgp = np.asarray(d["img"]) if "img" in d else None
    log(f"  eeg {eeg_shape} | label {lab.shape} | text {None if txt is None else txt.shape}"
        f" | img {None if imgp is None else imgp.shape}")
    assert eeg_shape[0] == n_concept, f"概念数不一致: eeg {eeg_shape[0]} vs 特征 {n_concept}"

    # ---- 顺序对齐独立检验（不碰 EEG）----
    align = {}
    if txt is not None:
        texts = [str(txt[c, 0]) for c in range(n_concept)]
        R_str = rdm_full(tfidf_matrix(texts))
        R_vit_txt = rdm_full(TXT)
        R_vit_img = rdm_full(IMG)
        rng = np.random.default_rng(SEED)
        r_t, null_t = _rsa_null(R_str, R_vit_txt, 200, rng)
        r_i, null_i = _rsa_null(R_str, R_vit_img, 200, rng)
        p_t = float((np.sum(np.abs(null_t) >= abs(r_t)) + 1) / (len(null_t) + 1))
        p_i = float((np.sum(np.abs(null_i) >= abs(r_i)) + 1) / (len(null_i) + 1))
        log(f"\n  [对齐检验] r(字符串RDM, ViT text RDM) = {r_t:+.4f}  p={p_t:.4f} "
            f"(零分布 sd={null_t.std():.4f})")
        log(f"  [对齐检验] r(字符串RDM, ViT img  RDM) = {r_i:+.4f}  p={p_i:.4f}")
        align = {"r_text": r_t, "p_text": p_t, "null_sd_text": float(null_t.std()),
                 "r_img": r_i, "p_img": p_i, "ok": bool(p_t < 0.05 and r_t > 0.1)}
        log(f"  → {'✅ train split 顺序对齐确认' if align['ok'] else '❌ 对齐未通过，实验作废'}")
        for c in [0, 1, n_concept - 1]:
            log(f"    [{c}] label={lab[c,0]} text={texts[c][:40]!r} img={str(imgp[c,0])[:70]}")
    info["alignment"] = align
    info["label_is_row_index"] = bool(
        np.array_equal(lab[:, 0], np.arange(n_concept)))
    log(f"  label == 行号? {info['label_is_row_index']}")

    del d            # 释放 4.17 GB
    return IMG, TXT, info


def _rsa_null(R_a, R_b, n_perm, rng):
    a = tri(R_a)
    r_obs = spearman(a, tri(R_b))
    null = np.empty(n_perm)
    n = R_b.shape[0]
    for i in range(n_perm):
        p = rng.permutation(n)
        null[i] = spearman(a, tri(R_b[np.ix_(p, p)]))
    return r_obs, null


# ================================================================ PART 1
def build_targets_chunked(eeg, chunk=64):
    """(n_img, n_rep, n_ch, T) → (n_img, n_ch*n_bin) 试次平均。
    分块处理，避免一次性造出 8.3 GB 的 float32 副本（1654×80×63×250 float32）。
    """
    n_img, n_rep, n_ch, T = eeg.shape
    bins = [(a, min(a + WIN, T)) for a in range(0, T, WIN)]
    n_bin = len(bins)
    out = np.empty((n_img, n_bin * n_ch), dtype=np.float32)
    for s in range(0, n_img, chunk):
        e = eeg[s:s + chunk].astype(np.float32)
        e -= e[:, :, :, :BASELINE].mean(axis=3, keepdims=True)
        acc = np.empty((e.shape[0], n_rep, n_bin * n_ch), dtype=np.float32)
        for bi, (a, b) in enumerate(bins):
            acc[:, :, bi * n_ch:(bi + 1) * n_ch] = e[:, :, :, a:b].mean(axis=3)
        out[s:s + chunk] = acc.mean(axis=1)
        del e, acc
    return out


def part1_cache():
    """逐被试算试次平均并落盘（BBB 需要所有被试，原始 npz 又太大）。"""
    log("\n" + "=" * 80)
    log("PART 1 — 试次平均响应缓存（10 个被试）")
    log("=" * 80)
    TA_CACHE.mkdir(parents=True, exist_ok=True)
    for sid in range(1, N_SUBJECTS + 1):
        fp = EegTRAIN / f"sub-{sid:02d}.npz"
        if not fp.exists():
            log(f"  [skip] sub-{sid:02d} 无文件")
            continue
        cache = TA_CACHE / f"train_ta_sub-{sid:02d}.npy"
        if cache.exists():
            log(f"  sub-{sid:02d}: 已有缓存，跳过")
            continue
        t0 = time.time()
        eeg = np.load(fp)["eeg"]
        log(f"  sub-{sid:02d}: 载入 {eeg.shape} {eeg.dtype}  [{time.time()-t0:.0f}s]")
        Y = build_targets_chunked(eeg)
        del eeg
        np.save(cache, Y)
        log(f"  sub-{sid:02d}: 试次平均 {Y.shape} → {cache.name}  [{time.time()-t0:.0f}s]")


# ================================================================ PART 2
def ridge_cv_primal(X, Y, folds, alphas, n_null, rng):
    """primal ridge（n > d）。返回 (best_alpha, 逐窗R², 逐窗零分布数组)。
    置换零：**打乱训练折的 Y 行**（X 不动 ⇒ XᵀX 复用，每次置换省一次
    1323×1024×1024 的矩阵乘）。零假设下与打乱「特征↔图像对应」等价。
    """
    n, d = X.shape
    n_ch = 63
    n_bin = Y.shape[1] // n_ch
    assert Y.shape[1] == n_ch * n_bin, f"目标数 {Y.shape[1]} 不是 {n_ch} 的整数倍"
    real_w = {a: None for a in alphas}
    # 每个置换一个槽，跨折收集（与真实值同样要跨折平均）
    null_w = [[] for _ in range(n_null)]

    for tr, te in folds:
        Xtr, Xte = X[tr], X[te]
        mu = Xtr.mean(axis=0, keepdims=True)
        sd = Xtr.std(axis=0, keepdims=True) + 1e-6
        Xtr = (Xtr - mu) / sd
        Xte = (Xte - mu) / sd
        XtX = (Xtr.T @ Xtr).astype(np.float64)
        ybar = Y[tr].mean(axis=0, keepdims=True)
        Ytr = Y[tr] - ybar
        Yte = Y[te]
        ss_tot = ((Yte - ybar) ** 2).sum(axis=0)
        scale = float(np.trace(XtX)) / d
        XtY = Xtr.T @ Ytr
        for a in alphas:
            A = XtX.copy()
            A.flat[:: d + 1] += a * scale
            W = np.linalg.solve(A, XtY)
            pred = Xte @ W + ybar
            r2 = 1.0 - ((Yte - pred) ** 2).sum(axis=0) / np.maximum(ss_tot, 1e-12)
            real_w[a] = r2 if real_w[a] is None else real_w[a] + r2
        # 置换零：每个置换在每个折里独立打乱训练 Y 的行
        for pi in range(n_null):
            Yp = Ytr[rng.permutation(len(tr))]
            XtYp = Xtr.T @ Yp
            vals = []
            for a in alphas:
                A = XtX.copy()
                A.flat[:: d + 1] += a * scale
                W = np.linalg.solve(A, XtYp)
                pred = Xte @ W + ybar
                r2 = 1.0 - ((Yte - pred) ** 2).sum(axis=0) / np.maximum(ss_tot, 1e-12)
                vals.append(r2)
            null_w[pi].append(vals)
    nf = len(folds)
    best = None
    for a in alphas:
        w = real_w[a] / nf
        M = w.reshape(n_bin, n_ch)
        pw = np.median(M, axis=1)
        if best is None or pw.max() > best[1].max():
            best = (a, pw)
    # 零分布：每次置换按同一规则（先跨折平均，再取逐窗中位的最大者）选 alpha
    nulls = np.zeros((n_null, n_bin))
    for pi in range(n_null):
        acc = np.zeros((len(alphas), Y.shape[1]))
        for vals in null_w[pi]:
            acc += np.asarray(vals)      # (n_alphas, n_targets) 逐折累加
        acc /= nf
        bestp = None
        for ai in range(len(alphas)):
            pw = np.median(acc[ai].reshape(n_bin, n_ch), axis=1)
            if bestp is None or pw.max() > bestp.max():
                bestp = pw
        nulls[pi] = bestp
    return best[0], best[1], nulls


def part2_encode(rng):
    log("\n" + "=" * 80)
    log("PART 2 — Ridge 编码（1654 概念，primal，5-fold 图像级 CV）")
    log("=" * 80)
    import torch
    f = torch.load(FEAT_TRAIN_PT, map_location="cpu", weights_only=False)
    IMG = f["img_features"].numpy().astype(np.float64)
    TXT = f["text_features"].numpy().astype(np.float64)

    tas = {}
    for sid in range(1, N_SUBJECTS + 1):
        c = TA_CACHE / f"train_ta_sub-{sid:02d}.npy"
        if c.exists():
            tas[sid] = np.load(c)
    log(f"  载入试次平均缓存: {sorted(tas)}")
    if not tas:
        return {}

    rows = []
    for sid, Y in tas.items():
        t0 = time.time()
        n_img = Y.shape[0]
        folds = make_folds(n_img, N_FOLDS, rng)
        others = [v for k, v in tas.items() if k != sid]
        Y_bbb = np.mean(np.stack(others, axis=0), axis=0) if others else None

        a_i, w_i, null_i = ridge_cv_primal(IMG, Y, folds, ALPHAS, N_NULL, rng)
        a_t, w_t, null_t = ridge_cv_primal(TXT, Y, folds, ALPHAS, N_NULL, rng)
        a_b, w_b, _ = ridge_cv_primal(Y_bbb, Y, folds, ALPHAS, 0, rng) if Y_bbb is not None \
            else (None, None, None)

        bi = int(np.argmax(w_i))
        bt = int(np.argmax(w_t))
        p_i = float((np.sum(null_i[:, bi] >= w_i[bi]) + 1) / (N_NULL + 1))
        p_t = float((np.sum(null_t[:, bt] >= w_t[bt]) + 1) / (N_NULL + 1))
        rec = {
            "subject": f"sub-{sid:02d}", "n_img": n_img,
            "alpha_img": a_i, "alpha_txt": a_t, "alpha_bbb": a_b,
            "best_win_img_ms": [bi * 100, (bi + 1) * 100],
            "best_win_txt_ms": [bt * 100, (bt + 1) * 100],
            "r2_img_best_window": float(w_i[bi]),
            "r2_txt_best_window": float(w_t[bt]),
            "r2_bbb_best_window": float(w_b[bi]) if w_b is not None else None,
            "null_img_at_best_median": float(np.median(null_i[:, bi])),
            "null_img_at_best_p95": float(np.percentile(null_i[:, bi], 95)),
            "p_img": p_i, "p_txt": p_t,
            "ratio_img_over_bbb": (float(w_i[bi]) / float(w_b[bi]))
            if w_b is not None and abs(w_b[bi]) > 1e-9 else None,
            "per_window_r2_img": [float(x) for x in w_i],
            "per_window_r2_bbb": [float(x) for x in w_b] if w_b is not None else None,
            "elapsed_s": round(time.time() - t0, 1),
        }
        rows.append(rec)
        log(f"  sub-{sid:02d}: 最佳窗 {bi*100}-{bi*100+100}ms  "
            f"R² img {w_i[bi]:+.4f} (零中位 {rec['null_img_at_best_median']:+.4f}, p={p_i:.3f}) | "
            f"txt {w_t[bt]:+.4f} (p={p_t:.3f}) | BBB {rec['r2_bbb_best_window']:+.4f} | "
            f"img/BBB {rec['ratio_img_over_bbb']:.2f}  [{rec['elapsed_s']}s]")

        OUT_DIR.mkdir(parents=True, exist_ok=True)
        (OUT_DIR / "exp68_train_encoding.json").write_text(
            json.dumps({"per_subject": rows}, indent=2, ensure_ascii=False),
            encoding="utf-8")
    return {"per_subject": rows}


# ================================================================ 主
def main():
    t0 = time.time()
    OUT_DIR.mkdir(parents=True, exist_ok=True)
    rng = np.random.default_rng(SEED)
    report = {"schema_version": "1.0", "script": "exp68_train_encoding",
              "started": time.strftime("%Y-%m-%dT%H:%M:%S"),
              "design": {"n_concept": 1654, "ridge": "primal (n>d)",
                         "cv": "5-fold over IMAGES",
                         "null": "permuting training-fold Y rows (X fixed -> Gram reused)",
                         "baselines": ["置换零", "BBB 跨被试脑对脑上限"]}}

    IMG, TXT, info = part0_probe()
    report["files"] = info
    if not info.get("alignment", {}).get("ok", False):
        log("\n❌ 顺序对齐未通过 —— 终止，不做后续编码")
        report["verdicts"] = {"alignment": {"ok": False}}
        (OUT_DIR / "exp68_train_encoding.json").write_text(
            json.dumps(report, indent=2, ensure_ascii=False), encoding="utf-8")
        return

    part1_cache()
    report["encoding"] = part2_encode(rng)

    rows = report["encoding"].get("per_subject", [])
    if rows:
        med_img = float(np.median([r["r2_img_best_window"] for r in rows]))
        med_txt = float(np.median([r["r2_txt_best_window"] for r in rows]))
        med_bbb = float(np.median([r["r2_bbb_best_window"] for r in rows]))
        med_null = float(np.median([r["null_img_at_best_median"] for r in rows]))
        n_sig = sum(1 for r in rows if r["p_img"] < 0.05)
        ratio = float(np.median([r["ratio_img_over_bbb"] for r in rows]))
        log("\n" + "=" * 80)
        log("判定")
        log("=" * 80)
        log(f"  编码显著   {'✅' if n_sig >= max(3, len(rows)//2) else '❌'}  "
            f"最佳窗 R² 中位 img {med_img:+.4f} vs 零 {med_null:+.4f}，显著 {n_sig}/{len(rows)}")
        log(f"  vs BBB     BBB {med_bbb:+.4f}，img/BBB = {ratio:.3f}")
        log(f"  txt 中位 {med_txt:+.4f}")
        report["verdicts"] = {
            "alignment": {"ok": True},
            "encoding_significant": {"ok": bool(n_sig >= max(3, len(rows) // 2)),
                                     "median_r2_img": med_img, "n_sig": n_sig,
                                     "n_subjects": len(rows)},
            "vs_bbb": {"ok": bool(ratio > 0), "median_r2_bbb": med_bbb, "ratio": ratio},
        }
    report["elapsed_s"] = round(time.time() - t0, 1)
    (OUT_DIR / "exp68_train_encoding_full.json").write_text(
        json.dumps(report, indent=2, ensure_ascii=False), encoding="utf-8")
    print("=== JSON BEGIN ===")
    print(json.dumps(report, indent=2, ensure_ascii=False))
    print("=== JSON END ===")
    log(f"\n报告 → {OUT_DIR}/exp68_train_encoding_full.json  耗时 {time.time()-t0:.0f}s")


if __name__ == "__main__":
    main()

Xet Storage Details

Size:: 18 kB
Xet hash:: d2e2c0a68815fdf384ef01a70223d1c7a78afc1a15f260aacdb2dc114c5531c4

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.

	#!/usr/bin/env python3
	"""exp68 — Tier 0 编码的 train split 版：1654 个概念。

	## 为什么做这个

	exp67 在 test split（200 个概念）上得到：ViT-H-14 特征能预测 EEG 脑响应，
	R² 中位 +0.0196，10/10 被试显著，且约等于跨被试脑对脑上限（比值中位 1.03）。

	但 200 张图配 1024 维特征是严重欠定的（选出的 α 全是 30，强正则），
	R² 绝对值只有 1%–4%。train split 有 1654 个概念，样本量 ×8，
	编码强度应当显著提高 —— 这是判断「模型特征究竟能解释多少脑活动」的关键数字。

	## 三处与 exp67 不同的设计

	1. primal 形式：n=1654 > d=1024，Gram 矩阵不再是最小的那一边。
	XᵀX 是 1024×1024，解一次 1e9 flops。
	2. 置换零改为打乱训练集的 Y 行：X 完全不动 ⇒ 同一折的 XᵀX 复用，
	每次置换省掉一次 1.4e9 的矩阵乘（总省 ~40%）。
	在零假设「X 不含 Y 的信息」下，这与打乱「特征↔图像对应关系」是等价的。
	3. 试次平均响应缓存到桶：每个被试的 (1654, 630) float32 只有 4.2 MB。
	原始 npz 是 4.17 GB/被试，BBB 又需要所有被试 —— 先各自算好试次平均并落盘，
	再从缓存里做 BBB，内存峰值从 42 GB 降到 ~5 GB。

	## 必须先验证顺序对齐

	train split 的 ViT 特征（`ViT-H-14_features_train.pt`）与
	`omnibrain4d/eeg_train/` 的概念顺序必须对齐，否则整个实验作废。
	本脚本第一步就用 exp66 的办法独立验证：拿 train .pt 里的概念文本构造字符串 RDM，
	与 ViT text_features 的 RDM 求相关 + 置换零。这一步不通过，后面都不用跑。

	## 判据

	顺序对齐 r(字符串RDM, ViT text RDM) 显著高于置换零
	编码 R² img 显著高于置换零，且 img/BBB 比值可解读
	跨被试 10 个被试中显著者占比
	"""

	from __future__ import annotations

	import json
	import os
	import re
	import sys
	import time
	from collections import Counter
	from pathlib import Path

	import numpy as np

	for _v in ("OMP_NUM_THREADS", "OPENBLAS_NUM_THREADS", "MKL_NUM_THREADS"):
	os.environ.setdefault(_v, "8")

	FS = 250.0
	WIN = 25
	BASELINE = 12
	N_FOLDS = 5
	N_SUBJECTS = 10
	N_NULL = 20
	ALPHAS = [0.01, 0.1, 1.0, 3.0, 10.0, 30.0, 100.0, 300.0, 1000.0, 3000.0]
	SEED = 0

	TRAIN_PT = "/training/eeg/things-eeg/Preprocessed_data_250Hz_whiten/sub-01/train.pt"
	FEAT_TRAIN_PT = "/training/eeg/eeg-image-decode/ViT-H-14_features_train.pt"
	EegTRAIN = Path("/domains/omnibrain4d/eeg_train")
	OUT_DIR = Path("/domains/runs/tier0-encoding")
	TA_CACHE = Path("/domains/runs/tier0-encoding/cache")


	def log(*a):
	print(*a, flush=True)


	# ================================================================ 工具
	def make_folds(n, k, rng):
	idx = rng.permutation(n)
	return [(np.setdiff1d(idx, te), te) for te in np.array_split(idx, k)]


	def _rank(a):
	from scipy.stats import rankdata
	return rankdata(a)


	def spearman(a, b):
	ra = _rank(a).astype(np.float64)
	rb = _rank(b).astype(np.float64)
	ra -= ra.mean()
	rb -= rb.mean()
	den = np.sqrt((ra * ra).sum() * (rb * rb).sum())
	return float((ra * rb).sum() / den) if den > 0 else 0.0


	def rdm_full(mat):
	m = np.asarray(mat, dtype=np.float64)
	m = m - m.mean(axis=1, keepdims=True)
	m = m / (np.linalg.norm(m, axis=1, keepdims=True) + 1e-12)
	r = m @ m.T
	np.fill_diagonal(r, 1.0)
	return 1.0 - r


	def tri(R):
	return R[np.triu_indices_from(R, k=1)]


	def tokenize(s):
	return re.findall(r"[a-z0-9]+", s.lower())


	def tfidf_matrix(docs):
	toks = [tokenize(d) for d in docs]
	df = Counter()
	for t in toks:
	df.update(set(t))
	vocab = sorted(df.keys())
	idx = {w: i for i, w in enumerate(vocab)}
	n = len(docs)
	M = np.zeros((n, len(vocab)), dtype=np.float64)
	for i, t in enumerate(toks):
	for w in t:
	M[i, idx[w]] += 1.0
	dfv = np.array([df[w] for w in vocab], dtype=np.float64)
	idf = np.log((1.0 + n) / (1.0 + dfv)) + 1.0
	M *= idf
	M /= (np.linalg.norm(M, axis=1, keepdims=True) + 1e-12)
	return M


	# ================================================================ PART 0
	def part0_probe():
	"""探查 train split 的数据 + 独立验证顺序对齐。"""
	import torch
	log("=" * 80)
	log("PART 0 — train split 探查 + 顺序对齐独立验证")
	log("=" * 80)
	info = {}

	f = torch.load(FEAT_TRAIN_PT, map_location="cpu", weights_only=False)
	log(f"\n[ViT train 特征] keys={list(f.keys())}")
	for k in f.keys():
	sh = tuple(f[k].shape)
	info[f"feat_{k}_shape"] = list(sh)
	log(f" {k}: {sh} {f[k].dtype}")
	IMG = f["img_features"].numpy().astype(np.float64)
	TXT = f["text_features"].numpy().astype(np.float64)
	n_concept = IMG.shape[0]

	log(f"\n[原始 train.pt] {TRAIN_PT}")
	d = torch.load(TRAIN_PT, map_location="cpu", weights_only=False)
	log(f" keys = {list(d.keys())}")
	for k in d.keys():
	v = d[k]
	if hasattr(v, "shape"):
	log(f" {k}: {tuple(v.shape)} {getattr(v,'dtype','')}")
	eeg_shape = tuple(d["eeg"].shape)
	info["raw_train_eeg_shape"] = list(eeg_shape)
	lab = np.asarray(d["label"])
	txt = np.asarray(d["text"]) if "text" in d else None
	imgp = np.asarray(d["img"]) if "img" in d else None
	log(f" eeg {eeg_shape} \| label {lab.shape} \| text {None if txt is None else txt.shape}"
	f" \| img {None if imgp is None else imgp.shape}")
	assert eeg_shape[0] == n_concept, f"概念数不一致: eeg {eeg_shape[0]} vs 特征 {n_concept}"

	# ---- 顺序对齐独立检验（不碰 EEG）----
	align = {}
	if txt is not None:
	texts = [str(txt[c, 0]) for c in range(n_concept)]
	R_str = rdm_full(tfidf_matrix(texts))
	R_vit_txt = rdm_full(TXT)
	R_vit_img = rdm_full(IMG)
	rng = np.random.default_rng(SEED)
	r_t, null_t = _rsa_null(R_str, R_vit_txt, 200, rng)
	r_i, null_i = _rsa_null(R_str, R_vit_img, 200, rng)
	p_t = float((np.sum(np.abs(null_t) >= abs(r_t)) + 1) / (len(null_t) + 1))
	p_i = float((np.sum(np.abs(null_i) >= abs(r_i)) + 1) / (len(null_i) + 1))
	log(f"\n [对齐检验] r(字符串RDM, ViT text RDM) = {r_t:+.4f} p={p_t:.4f} "
	f"(零分布 sd={null_t.std():.4f})")
	log(f" [对齐检验] r(字符串RDM, ViT img RDM) = {r_i:+.4f} p={p_i:.4f}")
	align = {"r_text": r_t, "p_text": p_t, "null_sd_text": float(null_t.std()),
	"r_img": r_i, "p_img": p_i, "ok": bool(p_t < 0.05 and r_t > 0.1)}
	log(f" → {'✅ train split 顺序对齐确认' if align['ok'] else '❌ 对齐未通过，实验作废'}")
	for c in [0, 1, n_concept - 1]:
	log(f" [{c}] label={lab[c,0]} text={texts[c][:40]!r} img={str(imgp[c,0])[:70]}")
	info["alignment"] = align
	info["label_is_row_index"] = bool(
	np.array_equal(lab[:, 0], np.arange(n_concept)))
	log(f" label == 行号? {info['label_is_row_index']}")

	del d # 释放 4.17 GB
	return IMG, TXT, info


	def _rsa_null(R_a, R_b, n_perm, rng):
	a = tri(R_a)
	r_obs = spearman(a, tri(R_b))
	null = np.empty(n_perm)
	n = R_b.shape[0]
	for i in range(n_perm):
	p = rng.permutation(n)
	null[i] = spearman(a, tri(R_b[np.ix_(p, p)]))
	return r_obs, null


	# ================================================================ PART 1
	def build_targets_chunked(eeg, chunk=64):
	"""(n_img, n_rep, n_ch, T) → (n_img, n_ch*n_bin) 试次平均。

	分块处理，避免一次性造出 8.3 GB 的 float32 副本（1654×80×63×250 float32）。
	"""
	n_img, n_rep, n_ch, T = eeg.shape
	bins = [(a, min(a + WIN, T)) for a in range(0, T, WIN)]
	n_bin = len(bins)
	out = np.empty((n_img, n_bin * n_ch), dtype=np.float32)
	for s in range(0, n_img, chunk):
	e = eeg[s:s + chunk].astype(np.float32)
	e -= e[:, :, :, :BASELINE].mean(axis=3, keepdims=True)
	acc = np.empty((e.shape[0], n_rep, n_bin * n_ch), dtype=np.float32)
	for bi, (a, b) in enumerate(bins):
	acc[:, :, bi * n_ch:(bi + 1) * n_ch] = e[:, :, :, a:b].mean(axis=3)
	out[s:s + chunk] = acc.mean(axis=1)
	del e, acc
	return out


	def part1_cache():
	"""逐被试算试次平均并落盘（BBB 需要所有被试，原始 npz 又太大）。"""
	log("\n" + "=" * 80)
	log("PART 1 — 试次平均响应缓存（10 个被试）")
	log("=" * 80)
	TA_CACHE.mkdir(parents=True, exist_ok=True)
	for sid in range(1, N_SUBJECTS + 1):
	fp = EegTRAIN / f"sub-{sid:02d}.npz"
	if not fp.exists():
	log(f" [skip] sub-{sid:02d} 无文件")
	continue
	cache = TA_CACHE / f"train_ta_sub-{sid:02d}.npy"
	if cache.exists():
	log(f" sub-{sid:02d}: 已有缓存，跳过")
	continue
	t0 = time.time()
	eeg = np.load(fp)["eeg"]
	log(f" sub-{sid:02d}: 载入 {eeg.shape} {eeg.dtype} [{time.time()-t0:.0f}s]")
	Y = build_targets_chunked(eeg)
	del eeg
	np.save(cache, Y)
	log(f" sub-{sid:02d}: 试次平均 {Y.shape} → {cache.name} [{time.time()-t0:.0f}s]")


	# ================================================================ PART 2
	def ridge_cv_primal(X, Y, folds, alphas, n_null, rng):
	"""primal ridge（n > d）。返回 (best_alpha, 逐窗R², 逐窗零分布数组)。

	置换零：打乱训练折的 Y 行（X 不动 ⇒ XᵀX 复用，每次置换省一次
	1323×1024×1024 的矩阵乘）。零假设下与打乱「特征↔图像对应」等价。
	"""
	n, d = X.shape
	n_ch = 63
	n_bin = Y.shape[1] // n_ch
	assert Y.shape[1] == n_ch * n_bin, f"目标数 {Y.shape[1]} 不是 {n_ch} 的整数倍"
	real_w = {a: None for a in alphas}
	# 每个置换一个槽，跨折收集（与真实值同样要跨折平均）
	null_w = [[] for _ in range(n_null)]

	for tr, te in folds:
	Xtr, Xte = X[tr], X[te]
	mu = Xtr.mean(axis=0, keepdims=True)
	sd = Xtr.std(axis=0, keepdims=True) + 1e-6
	Xtr = (Xtr - mu) / sd
	Xte = (Xte - mu) / sd
	XtX = (Xtr.T @ Xtr).astype(np.float64)
	ybar = Y[tr].mean(axis=0, keepdims=True)
	Ytr = Y[tr] - ybar
	Yte = Y[te]
	ss_tot = ((Yte - ybar) ** 2).sum(axis=0)
	scale = float(np.trace(XtX)) / d
	XtY = Xtr.T @ Ytr
	for a in alphas:
	A = XtX.copy()
	A.flat[:: d + 1] += a * scale
	W = np.linalg.solve(A, XtY)
	pred = Xte @ W + ybar
	r2 = 1.0 - ((Yte - pred) ** 2).sum(axis=0) / np.maximum(ss_tot, 1e-12)
	real_w[a] = r2 if real_w[a] is None else real_w[a] + r2
	# 置换零：每个置换在每个折里独立打乱训练 Y 的行
	for pi in range(n_null):
	Yp = Ytr[rng.permutation(len(tr))]
	XtYp = Xtr.T @ Yp
	vals = []
	for a in alphas:
	A = XtX.copy()
	A.flat[:: d + 1] += a * scale
	W = np.linalg.solve(A, XtYp)
	pred = Xte @ W + ybar
	r2 = 1.0 - ((Yte - pred) ** 2).sum(axis=0) / np.maximum(ss_tot, 1e-12)
	vals.append(r2)
	null_w[pi].append(vals)
	nf = len(folds)
	best = None
	for a in alphas:
	w = real_w[a] / nf
	M = w.reshape(n_bin, n_ch)
	pw = np.median(M, axis=1)
	if best is None or pw.max() > best[1].max():
	best = (a, pw)
	# 零分布：每次置换按同一规则（先跨折平均，再取逐窗中位的最大者）选 alpha
	nulls = np.zeros((n_null, n_bin))
	for pi in range(n_null):
	acc = np.zeros((len(alphas), Y.shape[1]))
	for vals in null_w[pi]:
	acc += np.asarray(vals) # (n_alphas, n_targets) 逐折累加
	acc /= nf
	bestp = None
	for ai in range(len(alphas)):
	pw = np.median(acc[ai].reshape(n_bin, n_ch), axis=1)
	if bestp is None or pw.max() > bestp.max():
	bestp = pw
	nulls[pi] = bestp
	return best[0], best[1], nulls


	def part2_encode(rng):
	log("\n" + "=" * 80)
	log("PART 2 — Ridge 编码（1654 概念，primal，5-fold 图像级 CV）")
	log("=" * 80)
	import torch
	f = torch.load(FEAT_TRAIN_PT, map_location="cpu", weights_only=False)
	IMG = f["img_features"].numpy().astype(np.float64)
	TXT = f["text_features"].numpy().astype(np.float64)

	tas = {}
	for sid in range(1, N_SUBJECTS + 1):
	c = TA_CACHE / f"train_ta_sub-{sid:02d}.npy"
	if c.exists():
	tas[sid] = np.load(c)
	log(f" 载入试次平均缓存: {sorted(tas)}")
	if not tas:
	return {}

	rows = []
	for sid, Y in tas.items():
	t0 = time.time()
	n_img = Y.shape[0]
	folds = make_folds(n_img, N_FOLDS, rng)
	others = [v for k, v in tas.items() if k != sid]
	Y_bbb = np.mean(np.stack(others, axis=0), axis=0) if others else None

	a_i, w_i, null_i = ridge_cv_primal(IMG, Y, folds, ALPHAS, N_NULL, rng)
	a_t, w_t, null_t = ridge_cv_primal(TXT, Y, folds, ALPHAS, N_NULL, rng)
	a_b, w_b, _ = ridge_cv_primal(Y_bbb, Y, folds, ALPHAS, 0, rng) if Y_bbb is not None \
	else (None, None, None)

	bi = int(np.argmax(w_i))
	bt = int(np.argmax(w_t))
	p_i = float((np.sum(null_i[:, bi] >= w_i[bi]) + 1) / (N_NULL + 1))
	p_t = float((np.sum(null_t[:, bt] >= w_t[bt]) + 1) / (N_NULL + 1))
	rec = {
	"subject": f"sub-{sid:02d}", "n_img": n_img,
	"alpha_img": a_i, "alpha_txt": a_t, "alpha_bbb": a_b,
	"best_win_img_ms": [bi * 100, (bi + 1) * 100],
	"best_win_txt_ms": [bt * 100, (bt + 1) * 100],
	"r2_img_best_window": float(w_i[bi]),
	"r2_txt_best_window": float(w_t[bt]),
	"r2_bbb_best_window": float(w_b[bi]) if w_b is not None else None,
	"null_img_at_best_median": float(np.median(null_i[:, bi])),
	"null_img_at_best_p95": float(np.percentile(null_i[:, bi], 95)),
	"p_img": p_i, "p_txt": p_t,
	"ratio_img_over_bbb": (float(w_i[bi]) / float(w_b[bi]))
	if w_b is not None and abs(w_b[bi]) > 1e-9 else None,
	"per_window_r2_img": [float(x) for x in w_i],
	"per_window_r2_bbb": [float(x) for x in w_b] if w_b is not None else None,
	"elapsed_s": round(time.time() - t0, 1),
	}
	rows.append(rec)
	log(f" sub-{sid:02d}: 最佳窗 {bi100}-{bi100+100}ms "
	f"R² img {w_i[bi]:+.4f} (零中位 {rec['null_img_at_best_median']:+.4f}, p={p_i:.3f}) \| "
	f"txt {w_t[bt]:+.4f} (p={p_t:.3f}) \| BBB {rec['r2_bbb_best_window']:+.4f} \| "
	f"img/BBB {rec['ratio_img_over_bbb']:.2f} [{rec['elapsed_s']}s]")

	OUT_DIR.mkdir(parents=True, exist_ok=True)
	(OUT_DIR / "exp68_train_encoding.json").write_text(
	json.dumps({"per_subject": rows}, indent=2, ensure_ascii=False),
	encoding="utf-8")
	return {"per_subject": rows}


	# ================================================================ 主
	def main():
	t0 = time.time()
	OUT_DIR.mkdir(parents=True, exist_ok=True)
	rng = np.random.default_rng(SEED)
	report = {"schema_version": "1.0", "script": "exp68_train_encoding",
	"started": time.strftime("%Y-%m-%dT%H:%M:%S"),
	"design": {"n_concept": 1654, "ridge": "primal (n>d)",
	"cv": "5-fold over IMAGES",
	"null": "permuting training-fold Y rows (X fixed -> Gram reused)",
	"baselines": ["置换零", "BBB 跨被试脑对脑上限"]}}

	IMG, TXT, info = part0_probe()
	report["files"] = info
	if not info.get("alignment", {}).get("ok", False):
	log("\n❌ 顺序对齐未通过 —— 终止，不做后续编码")
	report["verdicts"] = {"alignment": {"ok": False}}
	(OUT_DIR / "exp68_train_encoding.json").write_text(
	json.dumps(report, indent=2, ensure_ascii=False), encoding="utf-8")
	return

	part1_cache()
	report["encoding"] = part2_encode(rng)

	rows = report["encoding"].get("per_subject", [])
	if rows:
	med_img = float(np.median([r["r2_img_best_window"] for r in rows]))
	med_txt = float(np.median([r["r2_txt_best_window"] for r in rows]))
	med_bbb = float(np.median([r["r2_bbb_best_window"] for r in rows]))
	med_null = float(np.median([r["null_img_at_best_median"] for r in rows]))
	n_sig = sum(1 for r in rows if r["p_img"] < 0.05)
	ratio = float(np.median([r["ratio_img_over_bbb"] for r in rows]))
	log("\n" + "=" * 80)
	log("判定")
	log("=" * 80)
	log(f" 编码显著 {'✅' if n_sig >= max(3, len(rows)//2) else '❌'} "
	f"最佳窗 R² 中位 img {med_img:+.4f} vs 零 {med_null:+.4f}，显著 {n_sig}/{len(rows)}")
	log(f" vs BBB BBB {med_bbb:+.4f}，img/BBB = {ratio:.3f}")
	log(f" txt 中位 {med_txt:+.4f}")
	report["verdicts"] = {
	"alignment": {"ok": True},
	"encoding_significant": {"ok": bool(n_sig >= max(3, len(rows) // 2)),
	"median_r2_img": med_img, "n_sig": n_sig,
	"n_subjects": len(rows)},
	"vs_bbb": {"ok": bool(ratio > 0), "median_r2_bbb": med_bbb, "ratio": ratio},
	}
	report["elapsed_s"] = round(time.time() - t0, 1)
	(OUT_DIR / "exp68_train_encoding_full.json").write_text(
	json.dumps(report, indent=2, ensure_ascii=False), encoding="utf-8")
	print("=== JSON BEGIN ===")
	print(json.dumps(report, indent=2, ensure_ascii=False))
	print("=== JSON END ===")
	log(f"\n报告 → {OUT_DIR}/exp68_train_encoding_full.json 耗时 {time.time()-t0:.0f}s")


	if __name__ == "__main__":
	main()