Upload folder using huggingface_hub

README.md

@@ -0,0 +1,54 @@
+---
+license: mit
+tags:
+  - icml2026
+  - open-reproductions
+  - time-series
+  - medical
+---
+
+# TS-Fingerprint — independent reproduction code (ICML 2026 #7348)
+
+Independent re-implementation of **"Learning Fingerprints for Medical Time Series with
+Redundancy-Constrained Information Maximization"**
+(arXiv [2605.00130](https://arxiv.org/abs/2605.00130), OpenReview
+[`lrRBHFIgaK`](https://openreview.net/forum?id=lrRBHFIgaK)).
+
+**The paper released no code.** A GitHub code search for `TS-Fingerprint` returns zero
+repositories and the paper carries no code link, so every component here was derived from
+the paper text (Sec. 3.2 architecture, Sec. 3.3 losses, Sec. 3.4 head, Sec. 4.1
+hyperparameters). Data preprocessing follows the paper's own stated protocol reference,
+[DL4mHealth/Medformer](https://github.com/DL4mHealth/Medformer/tree/891f65b8a7e77188508fd8c56e10ccdba7fc1e18).
+
+Full write-up, per-claim evidence and the agent session trace:
+**[the Trackio reproduction logbook](https://huggingface.co/spaces/riteshhf/repro-learning-fingerprints-for-medical-time-series-with-redundancy-constrained-information-maxi)**.
+
+## Files
+
+| File | Purpose |
+|---|---|
+| `model.py` | `TSFingerprint` (Perceiver-style cross-attention bottleneck, masked F′-only decoder, TCR diversity loss, attention-pooling head) + `TiMAE` and `SimMTM` baselines |
+| `train.py` | pre-train → fine-tune harness; subject-wise splits, Adam 1e-3/1e-4, ≤100 epochs, early stopping on val F1 (patience 10), macro Acc/P/R/F1/AUROC |
+| `schedule.py` | small GPU-slot scheduler for running the config grids locally |
+| `prep_adftd.py` | ADFTD (OpenNeuro ds004504) → 69,752 windows, reproducing Medformer's notebook |
+| `prep_ptbxl.py` | PTB-XL (PhysioNet) → 213,880 windows, 5 diagnostic superclasses |
+| `audit_claim1.py` | Claim 1: fixed k×d output, learnable Q, k ≪ T, strict X → F′ → X̂ chain |
+| `audit_theorem31.py` | Claim 2: numerical audit of Theorem 3.1, both legs, with controls |
+| `audit_table1_rank.py` | Claim 5: re-parses all 420 Table-1 values and recomputes the average rank |
+| `analyze_adftd.py`, `analyze_sweep.py` | Claims 3/4 and Claim 6 summaries + figures |
+| `jobs_*.json` | the exact config grids that were run |
+
+## Data
+
+Processed ADFTD arrays:
+[`riteshhf/tsfp-repro-adftd-processed`](https://huggingface.co/datasets/riteshhf/tsfp-repro-adftd-processed).
+PTB-XL must be fetched from PhysioNet (see `prep_ptbxl.py`; the
+`physionet-open` S3 mirror works without credentials).
+
+## Reproducing
+
+```bash
+python prep_adftd.py --root data/ADFTD --out data/ADFTD/processed
+python train.py --data-dir data/ADFTD/processed --model tsfp --mode pretrain_ft \
+                --use-div 1 --k 8 --mask-ratio 0.6 --seeds 41 42 43 --out rec_div.json
+```

analyze_adftd.py

@@ -0,0 +1,182 @@
+"""Turn the ADFTD ablation result JSONs into the Claim 3 / Claim 4 tables and figures."""
+
+import json
+import os
+
+import numpy as np
+import plotly.graph_objects as go
+
+RES = "results/adftd"
+OUT = "results"
+
+# paper values (Table 2 for the ablation, Table 4 for the promotion comparison)
+PAPER = {
+    "tsfp_scratch":    {"accuracy": 48.83, "f1": 45.80, "auroc": 65.29},
+    "tsfp_rec":        {"accuracy": 53.51, "f1": 50.48, "auroc": 70.21},
+    "tsfp_rec_div":    {"accuracy": 53.91, "f1": 51.79, "auroc": 72.11},
+    "timae_scratch":   {"accuracy": 50.27, "f1": 45.88},
+    "timae_pretrain":  {"accuracy": 50.48, "f1": 45.44},
+    "simmtm_scratch":  {"accuracy": 51.98, "f1": 43.66},
+    "simmtm_pretrain": {"accuracy": 52.92, "f1": 45.59},
+}
+LABEL = {
+    "tsfp_scratch": "TS-Fingerprint · Scratch",
+    "tsfp_rec": "TS-Fingerprint · Pre-trained (L_rec)",
+    "tsfp_rec_div": "TS-Fingerprint · Pre-trained (L_rec + L_div)",
+    "timae_scratch": "Ti-MAE · Scratch",
+    "timae_pretrain": "Ti-MAE · Pre-trained",
+    "simmtm_scratch": "SimMTM · Scratch",
+    "simmtm_pretrain": "SimMTM · Pre-trained",
+}
+
+
+def load():
+    out = {}
+    for name in PAPER:
+        p = os.path.join(RES, f"{name}.json")
+        if os.path.exists(p) and os.path.getsize(p) > 0:
+            out[name] = json.load(open(p))
+    return out
+
+
+def fmt(agg, m):
+    return f"{agg[m]['mean']:.2f} ± {agg[m]['std']:.2f}"
+
+
+def main():
+    r = load()
+    print(f"available configs: {sorted(r)}\n")
+    rows = []
+    for name in PAPER:
+        if name not in r:
+            continue
+        a = r[name]["aggregate"]
+        rows.append({
+            "config": name, "label": LABEL[name],
+            "acc": a["accuracy"]["mean"], "acc_sd": a["accuracy"]["std"],
+            "f1": a["f1"]["mean"], "f1_sd": a["f1"]["std"],
+            "auroc": a["auroc"]["mean"], "auroc_sd": a["auroc"]["std"],
+            "paper_acc": PAPER[name].get("accuracy"),
+            "paper_f1": PAPER[name].get("f1"),
+            "paper_auroc": PAPER[name].get("auroc"),
+            "n_seeds": len(r[name]["runs"]),
+            "sec_per_seed": float(np.mean([x["seconds"] for x in r[name]["runs"]])),
+            "params": r[name]["runs"][0]["params"],
+        })
+
+    print(f"{'config':<18}{'Acc (ours)':>16}{'Acc (paper)':>13}"
+          f"{'F1 (ours)':>16}{'F1 (paper)':>12}{'AUROC (ours)':>17}{'AUROC (paper)':>15}")
+    for x in rows:
+        pa = f"{x['paper_auroc']:.2f}" if x["paper_auroc"] else "-"
+        print(f"{x['config']:<18}{x['acc']:>9.2f}±{x['acc_sd']:<5.2f}"
+              f"{x['paper_acc']:>13.2f}{x['f1']:>9.2f}±{x['f1_sd']:<5.2f}"
+              f"{x['paper_f1']:>12.2f}{x['auroc']:>10.2f}±{x['auroc_sd']:<5.2f}{pa:>15}")
+
+    # ---- relative promotion (Table 2 / Table 4 delta definition) ----
+    def promo(scratch, pre, key="f1"):
+        if scratch not in r or pre not in r:
+            return None
+        s = r[scratch]["aggregate"][key]["mean"]
+        p = r[pre]["aggregate"][key]["mean"]
+        return {"scratch": s, "pretrained": p, "abs": p - s, "rel_pct": (p - s) / s * 100}
+
+    promos = {
+        "TS-Fingerprint": {"f1": promo("tsfp_scratch", "tsfp_rec_div", "f1"),
+                           "accuracy": promo("tsfp_scratch", "tsfp_rec_div", "accuracy")},
+        "Ti-MAE": {"f1": promo("timae_scratch", "timae_pretrain", "f1"),
+                   "accuracy": promo("timae_scratch", "timae_pretrain", "accuracy")},
+        "SimMTM": {"f1": promo("simmtm_scratch", "simmtm_pretrain", "f1"),
+                   "accuracy": promo("simmtm_scratch", "simmtm_pretrain", "accuracy")},
+    }
+    # the diversity-loss increment specifically (Claim 3)
+    div_effect = promo("tsfp_rec", "tsfp_rec_div", "f1")
+    print("\nrelative F1 promotion (scratch -> pre-trained):")
+    for k, v in promos.items():
+        if v["f1"]:
+            print(f"  {k:<16} {v['f1']['scratch']:.2f} -> {v['f1']['pretrained']:.2f}  "
+                  f"(+{v['f1']['rel_pct']:.2f}%)")
+    if div_effect:
+        print(f"\nL_div increment on top of L_rec: F1 {div_effect['scratch']:.2f} -> "
+              f"{div_effect['pretrained']:.2f} ({div_effect['abs']:+.2f} pts, "
+              f"{div_effect['rel_pct']:+.2f}%)")
+
+    # per-seed paired test for the diversity loss
+    paired = None
+    if "tsfp_rec" in r and "tsfp_rec_div" in r:
+        a = {x["seed"]: x["test"]["f1"] for x in r["tsfp_rec"]["runs"]}
+        b = {x["seed"]: x["test"]["f1"] for x in r["tsfp_rec_div"]["runs"]}
+        seeds = sorted(set(a) & set(b))
+        d = [b[s] - a[s] for s in seeds]
+        paired = {"seeds": seeds, "per_seed_delta_f1": d,
+                  "mean_delta": float(np.mean(d)),
+                  "n_positive": int(sum(x > 0 for x in d))}
+        print(f"\npaired per-seed F1 delta (rec+div minus rec): "
+              f"{[round(x, 2) for x in d]}  mean {np.mean(d):+.2f}, "
+              f"{paired['n_positive']}/{len(d)} seeds positive")
+
+    json.dump({"rows": rows, "promotions": promos, "div_effect": div_effect,
+               "paired_div_test": paired},
+              open(os.path.join(OUT, "adftd_summary.json"), "w"), indent=2)
+
+    # ---------------- figures ----------------
+    # Fig 1: Claim 3 ablation, ours vs paper
+    order = ["tsfp_scratch", "tsfp_rec", "tsfp_rec_div"]
+    have = [x for x in order if x in r]
+    if have:
+        xs = ["Scratch", "Pre-trained\n(L_rec)", "Pre-trained\n(L_rec + L_div)"][:len(have)]
+        fig = go.Figure()
+        for m, col in (("f1", "#2F6F8F"), ("accuracy", "#B07C2B"), ("auroc", "#7E7E7E")):
+            ours = [r[c]["aggregate"][m]["mean"] for c in have]
+            sd = [r[c]["aggregate"][m]["std"] for c in have]
+            pap = [PAPER[c].get(m) for c in have]
+            fig.add_bar(name=f"ours · {m}", x=xs, y=ours,
+                        error_y=dict(type="data", array=sd), marker_color=col)
+            fig.add_scatter(name=f"paper · {m}", x=xs, y=pap, mode="markers",
+                            marker=dict(symbol="line-ew", size=26, line=dict(
+                                width=3, color=col)))
+        fig.update_layout(
+            title="Claim 3 — ADFTD ablation: this reproduction (bars, 3 seeds) "
+                  "vs the paper's Table 2 (ticks)",
+            yaxis_title="score (%)", barmode="group", template="plotly_white",
+            height=460, legend=dict(orientation="h", y=-0.18))
+        fig.write_html(os.path.join(OUT, "claim3_ablation.html"),
+                       include_plotlyjs="cdn")
+
+    # Fig 2: Claim 4 promotion comparison
+    if all(promos[k]["f1"] for k in promos):
+        fig2 = go.Figure()
+        names = list(promos)
+        fig2.add_bar(name="this reproduction", x=names,
+                     y=[promos[k]["f1"]["rel_pct"] for k in names],
+                     marker_color="#2F6F8F",
+                     text=[f"{promos[k]['f1']['rel_pct']:+.2f}%" for k in names],
+                     textposition="outside")
+        paper_promo = {"TS-Fingerprint": 13.07, "Ti-MAE": -0.96, "SimMTM": 4.42}
+        fig2.add_bar(name="paper (Table 4)", x=names,
+                     y=[paper_promo[k] for k in names], marker_color="#B07C2B",
+                     text=[f"{paper_promo[k]:+.2f}%" for k in names],
+                     textposition="outside")
+        fig2.update_layout(
+            title="Claim 4 — relative F1 promotion from pre-training on ADFTD",
+            yaxis_title="relative F1 gain (%)", barmode="group",
+            template="plotly_white", height=440,
+            legend=dict(orientation="h", y=-0.15))
+        fig2.write_html(os.path.join(OUT, "claim4_promotion.html"),
+                        include_plotlyjs="cdn")
+
+    # raw CSV for the figure cells
+    with open(os.path.join(OUT, "adftd_results.csv"), "w") as f:
+        f.write("config,label,n_seeds,acc,acc_sd,f1,f1_sd,auroc,auroc_sd,"
+                "paper_acc,paper_f1,paper_auroc,sec_per_seed,params\n")
+        for x in rows:
+            f.write(f"{x['config']},{x['label']},{x['n_seeds']},{x['acc']:.3f},"
+                    f"{x['acc_sd']:.3f},{x['f1']:.3f},{x['f1_sd']:.3f},"
+                    f"{x['auroc']:.3f},{x['auroc_sd']:.3f},{x['paper_acc']},"
+                    f"{x['paper_f1']},{x['paper_auroc']},{x['sec_per_seed']:.0f},"
+                    f"{x['params']}\n")
+    print("\nwrote results/adftd_summary.json, adftd_results.csv, "
+          "claim3_ablation.html, claim4_promotion.html")
+
+
+if __name__ == "__main__":
+    main()

analyze_sweep.py

@@ -0,0 +1,123 @@
+"""Claim 6: summarise the k x mask-ratio sweep on PTB-XL and compare with Table 3."""
+
+import json
+import os
+import re
+
+import numpy as np
+import plotly.graph_objects as go
+from plotly.subplots import make_subplots
+
+RES = "results/sweep"
+KS, RS = [6, 8, 10], [0.5, 0.6, 0.7, 0.8]
+# Table 3 of the paper (F1 column), as percentages
+PAPER_F1 = {(6, 0.5): 54.81, (6, 0.6): 56.68, (6, 0.7): 55.41, (6, 0.8): 54.32,
+            (8, 0.5): 62.60, (8, 0.6): 63.51, (8, 0.7): 62.10, (8, 0.8): 59.19,
+            (10, 0.5): 51.76, (10, 0.6): 50.70, (10, 0.7): 56.33, (10, 0.8): 55.70}
+PAPER_ACC = {(6, 0.5): 70.19, (6, 0.6): 70.71, (6, 0.7): 70.63, (6, 0.8): 69.77,
+             (8, 0.5): 74.07, (8, 0.6): 73.98, (8, 0.7): 73.86, (8, 0.8): 72.35,
+             (10, 0.5): 67.36, (10, 0.6): 66.41, (10, 0.7): 71.07, (10, 0.8): 70.42}
+
+
+def main():
+    rows = []
+    for k in KS:
+        for r in RS:
+            p = os.path.join(RES, f"ptbxl_k{k}_r{r}.json")
+            if not (os.path.exists(p) and os.path.getsize(p) > 0):
+                continue
+            d = json.load(open(p))
+            a = d["aggregate"]
+            rows.append({"k": k, "r": r,
+                         "f1": a["f1"]["mean"], "f1_sd": a["f1"]["std"],
+                         "acc": a["accuracy"]["mean"], "acc_sd": a["accuracy"]["std"],
+                         "auroc": a["auroc"]["mean"],
+                         "n_seeds": len(d["runs"]),
+                         "paper_f1": PAPER_F1[(k, r)], "paper_acc": PAPER_ACC[(k, r)],
+                         "sec": float(np.mean([x["seconds"] for x in d["runs"]]))})
+    if not rows:
+        print("no sweep results yet")
+        return
+    print(f"{'k':>3}{'r':>6}{'F1 ours':>12}{'F1 paper':>11}{'Acc ours':>11}"
+          f"{'Acc paper':>11}{'seeds':>7}")
+    for x in rows:
+        print(f"{x['k']:>3}{x['r']:>6}{x['f1']:>9.2f}±{x['f1_sd']:<3.1f}"
+              f"{x['paper_f1']:>11.2f}{x['acc']:>8.2f}±{x['acc_sd']:<3.1f}"
+              f"{x['paper_acc']:>11.2f}{x['n_seeds']:>7}")
+
+    out = {"rows": rows}
+    if len(rows) == 12:
+        f1 = {(x["k"], x["r"]): x["f1"] for x in rows}
+        best = max(f1, key=f1.get)
+        out["best_cell_ours"] = {"k": best[0], "r": best[1], "f1": f1[best]}
+        out["best_cell_paper"] = {"k": 8, "r": 0.6, "f1": 63.51}
+        # marginal effect of k at each r, and of r at each k
+        out["best_k_per_r"] = {str(r): max(KS, key=lambda k: f1[(k, r)]) for r in RS}
+        out["best_r_per_k"] = {str(k): max(RS, key=lambda r: f1[(k, r)]) for k in KS}
+        out["mean_f1_by_k"] = {str(k): float(np.mean([f1[(k, r)] for r in RS]))
+                               for k in KS}
+        out["mean_f1_by_r"] = {str(r): float(np.mean([f1[(k, r)] for k in KS]))
+                               for r in RS}
+        out["paper_mean_f1_by_k"] = {
+            str(k): float(np.mean([PAPER_F1[(k, r)] for r in RS])) for k in KS}
+        out["paper_mean_f1_by_r"] = {
+            str(r): float(np.mean([PAPER_F1[(k, r)] for k in KS])) for r in RS}
+        # does r>=0.8 degrade?  compare r=0.8 against the mean of r in {0.5,0.6,0.7}
+        deg = {}
+        for k in KS:
+            lo = np.mean([f1[(k, r)] for r in (0.5, 0.6, 0.7)])
+            deg[str(k)] = {"mean_r_le_0.7": float(lo), "r_0.8": f1[(k, 0.8)],
+                           "delta": float(f1[(k, 0.8)] - lo)}
+        out["r08_degradation"] = deg
+        # rank correlation between our grid and the paper's grid
+        ours_v = np.array([f1[(k, r)] for k in KS for r in RS])
+        pap_v = np.array([PAPER_F1[(k, r)] for k in KS for r in RS])
+        out["pearson_ours_vs_paper_grid"] = float(np.corrcoef(ours_v, pap_v)[0, 1])
+        out["spearman_ours_vs_paper_grid"] = float(np.corrcoef(
+            np.argsort(np.argsort(ours_v)), np.argsort(np.argsort(pap_v)))[0, 1])
+        print("\nbest cell (ours):", out["best_cell_ours"])
+        print("mean F1 by k (ours):", {k: round(v, 2)
+                                       for k, v in out["mean_f1_by_k"].items()})
+        print("mean F1 by k (paper):", {k: round(v, 2)
+                                        for k, v in out["paper_mean_f1_by_k"].items()})
+        print("mean F1 by r (ours):", {k: round(v, 2)
+                                       for k, v in out["mean_f1_by_r"].items()})
+        print("mean F1 by r (paper):", {k: round(v, 2)
+                                        for k, v in out["paper_mean_f1_by_r"].items()})
+        print("r=0.8 vs r<=0.7:", {k: round(v["delta"], 2)
+                                   for k, v in deg.items()})
+        print(f"grid correlation with Table 3: Pearson "
+              f"{out['pearson_ours_vs_paper_grid']:.3f}, "
+              f"Spearman {out['spearman_ours_vs_paper_grid']:.3f}")
+
+    json.dump(out, open("results/sweep_summary.json", "w"), indent=2)
+
+    # interactive figure for the logbook
+    ours = np.full((3, 4), np.nan)
+    for x in rows:
+        ours[KS.index(x["k"]), RS.index(x["r"])] = x["f1"]
+    paper = np.array([[PAPER_F1[(k, r)] for r in RS] for k in KS])
+    fig = make_subplots(rows=1, cols=2, horizontal_spacing=0.13,
+                        subplot_titles=("this reproduction (PTB-XL, thinned)",
+                                        "paper (Table 3)"))
+    for j, m in enumerate((ours, paper)):
+        fig.add_heatmap(z=m, x=[f"r={r}" for r in RS], y=[f"k={k}" for k in KS],
+                        colorscale=[[0, "#F2F7F8"], [1, "#17697B"]], showscale=False,
+                        text=[[("" if np.isnan(v) else f"{v:.2f}") for v in row]
+                              for row in m],
+                        texttemplate="%{text}", row=1, col=j + 1)
+    fig.update_layout(template="plotly_white", height=420,
+                      title="Claim 6 — macro F1 (%) over bottleneck size k and mask ratio r")
+    fig.write_html("results/claim6_sweep.html", include_plotlyjs="cdn")
+
+    with open("results/sweep_results.csv", "w") as f:
+        f.write("k,mask_ratio,n_seeds,f1,f1_sd,acc,acc_sd,auroc,paper_f1,paper_acc,sec\n")
+        for x in rows:
+            f.write(f"{x['k']},{x['r']},{x['n_seeds']},{x['f1']:.3f},{x['f1_sd']:.3f},"
+                    f"{x['acc']:.3f},{x['acc_sd']:.3f},{x['auroc']:.3f},"
+                    f"{x['paper_f1']},{x['paper_acc']},{x['sec']:.0f}\n")
+    print("\nwrote results/sweep_summary.json, sweep_results.csv, claim6_sweep.html")
+
+
+if __name__ == "__main__":
+    main()

audit_claim1.py

@@ -0,0 +1,105 @@
+"""Claim 1 audit: TS-Fingerprint compresses a variable-length MedTS into a fixed set of
+k learnable Fingerprint Tokens via a learnable query set Q in R^{k x d}, with k << T.
+
+Checks:
+  1. the encoder output is exactly (B, k, d) for every input length T;
+  2. Q is a genuine nn.Parameter that receives gradient (i.e. it is *learnable*),
+     and the parameter count of the bottleneck is independent of T;
+  3. k << T holds at the paper's settings for all 7 benchmark datasets of Table 5;
+  4. the decoder sees *only* F' plus positional mask tokens, so the data-processing
+     chain X -> F' -> X_hat is strict: perturbing a masked input patch can change
+     X_hat only through F'.  We verify this by freezing F' and re-running the decoder.
+"""
+
+import json
+import os
+import sys
+
+import torch
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from model import TSFingerprint  # noqa: E402
+
+# Table 5 of the paper: (#Timestamps T, #Channel C, patch size)
+DATASETS = {
+    "ADFTD": (256, 19, 8), "PTB": (300, 15, 10), "PTB-XL": (250, 12, 10),
+    "APAVA": (256, 16, 8), "FLAAP": (100, 6, 10), "UCIHAR": (128, 9, 32),
+    "Sleep-EDF": (3000, 1, 60),
+}
+K, D = 8, 128
+rep = {}
+
+torch.manual_seed(0)
+m = TSFingerprint(c_in=19, patch_size=8, n_classes=3, d_model=D, k=K, max_patches=4096)
+
+# --- 1. shape invariance to T -------------------------------------------------
+shapes = {}
+for T in (64, 128, 256, 512, 1024, 2048, 4096, 8192):
+    with torch.no_grad():
+        f = m.encoder(torch.randn(3, T, 19))
+    shapes[T] = {"n_patches": T // 8, "F_shape": list(f.shape)}
+rep["shape_invariance"] = shapes
+rep["all_outputs_are_k_by_d"] = all(v["F_shape"] == [3, K, D] for v in shapes.values())
+
+# --- 2. Q is learnable, and bottleneck size is T-independent -------------------
+rep["Q_shape"] = list(m.encoder.Q.shape)
+rep["Q_is_parameter"] = isinstance(m.encoder.Q, torch.nn.Parameter)
+rep["Q_requires_grad"] = bool(m.encoder.Q.requires_grad)
+m.zero_grad()
+loss, _, _ = m.pretrain_step(torch.randn(4, 256, 19))
+loss.backward()
+g = m.encoder.Q.grad
+rep["Q_grad_norm_after_backward"] = float(g.norm())
+rep["Q_receives_gradient"] = bool(g is not None and g.norm() > 0)
+rep["encoder_params_excluding_pos_embed"] = sum(
+    p.numel() for n, p in m.encoder.named_parameters() if n != "pos")
+
+# --- 3. k << T on every benchmark ---------------------------------------------
+ratios = {}
+for name, (T, C, ps) in DATASETS.items():
+    n_patch = T // ps
+    ratios[name] = {"T": T, "C": C, "patch": ps, "n_patches": n_patch, "k": K,
+                    "k_over_T": round(K / T, 5), "k_over_n_patches": round(K / n_patch, 4),
+                    "k_much_less_than_T": bool(K < T / 4)}
+rep["compression_per_dataset"] = ratios
+rep["k_ll_T_on_all_datasets"] = all(v["k_much_less_than_T"] for v in ratios.values())
+# compression factor: raw floats in vs floats out
+rep["compression_factor_TxC_to_kxd"] = {
+    n: round((v["T"] * v["C"]) / (K * D), 3) for n, v in ratios.items()}
+
+# --- 4. strict bottleneck: decoder output depends on x only through F' ---------
+m.eval()
+x = torch.randn(2, 256, 19)
+with torch.no_grad():
+    f = m.encoder(x)
+    rec_a = m.decoder(f, 32)
+    x2 = x.clone()
+    x2[:, 100:140] += 5.0                       # large perturbation of the input
+    f2 = m.encoder(x2)
+    rec_b = m.decoder(f2, 32)                   # differs, because F' differs
+    rec_c = m.decoder(f, 32)                    # same F' -> must be bit-identical
+rep["decoder_conditioned_solely_on_F"] = {
+    "delta_recon_when_F_changes": float((rec_a - rec_b).abs().max()),
+    "delta_recon_when_F_held_fixed": float((rec_a - rec_c).abs().max()),
+    "strict_chain_X_to_F_to_Xhat": bool((rec_a - rec_c).abs().max() == 0.0
+                                        and (rec_a - rec_b).abs().max() > 0),
+}
+# the decoder signature takes only (F', n_patches) -- no path for raw patches
+import inspect  # noqa: E402
+rep["decoder_forward_signature"] = str(inspect.signature(m.decoder.forward))
+
+print(json.dumps(rep, indent=2))
+os.makedirs("results", exist_ok=True)
+with open("results/claim1_audit.json", "w") as fh:
+    json.dump(rep, fh, indent=2)
+
+print("\n---- verdict ----")
+print(f"encoder output is (B, {K}, {D}) for every T in "
+      f"{sorted(shapes)}: {rep['all_outputs_are_k_by_d']}")
+print(f"Q in R^(k x d) = {rep['Q_shape']}, learnable nn.Parameter, "
+      f"grad norm {rep['Q_grad_norm_after_backward']:.4f}")
+print(f"k << T on all 7 Table-5 datasets: {rep['k_ll_T_on_all_datasets']} "
+      f"(k/T from {min(v['k_over_T'] for v in ratios.values()):.5f} to "
+      f"{max(v['k_over_T'] for v in ratios.values()):.5f})")
+print("decoder is strictly conditioned on F': "
+      f"{rep['decoder_conditioned_solely_on_F']['strict_chain_X_to_F_to_Xhat']}")

audit_table1_rank.py

@@ -0,0 +1,121 @@
+"""Claim 5 audit: recompute the reported average rank of 1.24 from Table 1 itself.
+
+Table 1 is parsed straight out of the arXiv HTML text dump (paper/paper.txt) so the
+numbers are not re-typed by hand.  We then rank the 12 methods on every
+(dataset, metric) cell and report the mean rank of "Ours" under several readings of
+the claim, since the paper is ambiguous about which metrics enter the average
+(Sec 4.1 says "across all metrics", the challenge statement says accuracy/F1/AUROC).
+"""
+
+import json
+import re
+
+import numpy as np
+
+METHODS = ["Autoformer", "Crossformer", "FEDformer", "Informer", "iTransformer",
+           "MTST", "Nonformer", "PatchTST", "Reformer", "Transformer",
+           "Medformer", "Ours"]
+DATASETS = ["ADFTD", "PTB", "PTB-XL", "APAVA", "SleepEDF", "FLAAP", "UCI-HAR"]
+METRICS = ["Accuracy", "Precision", "Recall", "F1-score", "AUROC"]
+
+
+def parse(path="paper/paper.txt"):
+    txt = open(path, encoding="utf-8").read()
+    # Table 1 spans from the header block to the "Table 1 :" caption
+    start = txt.index(" Dataset \n Metrics \n Autoformer")
+    end = txt.index(" Table 1 : ")
+    body = txt[start:end]
+    # every "value ± std" pair, in reading order
+    vals = [float(m) for m in re.findall(r"(\d+\.\d+)\s*±", body)]
+    assert len(vals) == 7 * 5 * 12, f"expected 420 numbers, parsed {len(vals)}"
+    arr = np.array(vals).reshape(7, 5, 12)
+    return arr
+
+
+def ranks_for_ours(arr, metric_idx):
+    """Rank (1 = best, higher value is better for every metric here) of 'Ours'."""
+    out = {}
+    for di, ds in enumerate(DATASETS):
+        for mi in metric_idx:
+            col = arr[di, mi]
+            # competition ranking with ties sharing the average rank
+            order = (-col).argsort()
+            r = np.empty(12)
+            r[order] = np.arange(1, 13)
+            # handle exact ties
+            for v in np.unique(col):
+                m = col == v
+                if m.sum() > 1:
+                    r[m] = r[m].mean()
+            out[(ds, METRICS[mi])] = float(r[METHODS.index("Ours")])
+    return out
+
+
+def main():
+    arr = parse()
+    report = {"table_shape": list(arr.shape)}
+    # spot-check a few known cells against the paper text
+    checks = {
+        ("ADFTD", "Accuracy", "Ours"): 53.91,
+        ("ADFTD", "F1-score", "Ours"): 51.79,
+        ("ADFTD", "F1-score", "Medformer"): 50.65,
+        ("PTB-XL", "F1-score", "Ours"): 63.51,
+        ("UCI-HAR", "Accuracy", "Medformer"): 91.65,
+        ("SleepEDF", "AUROC", "Ours"): 95.37,
+    }
+    ok = {}
+    for (ds, me, mo), want in checks.items():
+        got = arr[DATASETS.index(ds), METRICS.index(me), METHODS.index(mo)]
+        ok[f"{ds}/{me}/{mo}"] = {"parsed": float(got), "paper": want,
+                                 "match": bool(abs(got - want) < 1e-9)}
+    report["parse_spot_checks"] = ok
+
+    variants = {
+        "all_5_metrics": [0, 1, 2, 3, 4],
+        "acc_f1_auroc (challenge wording)": [0, 3, 4],
+        "acc_f1": [0, 3],
+    }
+    for name, idx in variants.items():
+        r = ranks_for_ours(arr, idx)
+        vals = np.array(list(r.values()))
+        report[name] = {
+            "n_cells": int(vals.size),
+            "avg_rank_ours": float(vals.mean()),
+            "n_rank1": int((vals == 1).sum()),
+            "worst_rank": float(vals.max()),
+            "cells_not_rank1": {f"{k[0]}/{k[1]}": v for k, v in r.items() if v != 1},
+        }
+
+    # average rank of every method under the 5-metric reading, for context
+    allranks = {}
+    for mo in METHODS:
+        rs = []
+        for di in range(7):
+            for mi in range(5):
+                col = arr[di, mi]
+                order = (-col).argsort()
+                rr = np.empty(12)
+                rr[order] = np.arange(1, 13)
+                for v in np.unique(col):
+                    m = col == v
+                    if m.sum() > 1:
+                        rr[m] = rr[m].mean()
+                rs.append(rr[METHODS.index(mo)])
+        allranks[mo] = float(np.mean(rs))
+    report["avg_rank_all_methods_5metrics"] = allranks
+
+    # one-tailed Wilcoxon signed-rank, Ours vs Medformer (paper's stated test)
+    from scipy.stats import wilcoxon
+    ours = arr[:, :, METHODS.index("Ours")].ravel()
+    med = arr[:, :, METHODS.index("Medformer")].ravel()
+    st, p = wilcoxon(ours, med, alternative="greater")
+    report["wilcoxon_ours_gt_medformer"] = {"statistic": float(st), "p_value": float(p),
+                                            "n_pairs": int(ours.size),
+                                            "ours_wins": int((ours > med).sum())}
+    print(json.dumps(report, indent=2))
+    with open("results/table1_rank_audit.json", "w") as f:
+        json.dump(report, f, indent=2)
+
+
+if __name__ == "__main__":
+    main()

audit_theorem31.py

@@ -0,0 +1,242 @@
+"""Numerical audit of Theorem 3.1 (Redundancy-Constrained Information Maximization).
+
+Theorem 3.1 asserts that minimising  L_total = L_rec + lambda * L_div  is *formally
+equivalent* to  max_theta  I(X;F') - lambda * TC(F').
+
+The proof (App. A.1) decomposes into two independent claims, which we audit separately
+to double precision, each with a control that relaxes the stated condition:
+
+  (A) Sufficiency.  Under a Gaussian decoder p(X|F') = N(D(F'), sigma^2 I),
+      minimising the MSE maximises the Barber-Agakov variational lower bound
+      I_LB(X;F') = H(X) + E[log p(x|f')]  on I(X;F').
+      -> check  I_LB  is an exact affine, strictly decreasing function of the MSE,
+         and that I_LB <= I(X;F') for a case where the true MI is known in closed form.
+      CONTROL: a Laplace decoder, where the affine MSE<->ELBO identity must fail.
+
+  (B) Independence.  For Gaussian F',
+      TC(F') = 1/2 sum_i log(Sigma_ii) - 1/2 log det(Sigma).
+      The paper's L_div = -1/2 log det(I + d/(k eps^2) F'^T F') is claimed to be
+      "geometrically equivalent" to minimising TC.
+      -> check the Gaussian TC identity itself to double precision (exact),
+         and then measure how tightly L_div tracks TC.
+      CONTROL: non-Gaussian (uniform / heavy-tailed) latents, where the Gaussian
+      entropy formula and hence the TC identity must break.
+
+This is a numerical audit, not a proof replacement; it needs no GPU.
+"""
+
+import json
+
+import numpy as np
+
+RNG = np.random.default_rng(0)
+LOG2PI = np.log(2.0 * np.pi)
+report = {}
+
+
+# ---------------------------------------------------------------- (A) sufficiency
+def gaussian_elbo(x, xhat, sigma2):
+    """E[log p(x|f')] for p = N(xhat, sigma^2 I), in nats, per sample."""
+    d = x.shape[1]
+    sq = ((x - xhat) ** 2).sum(1)
+    return float((-0.5 * d * (LOG2PI + np.log(sigma2)) - sq / (2 * sigma2)).mean())
+
+
+def part_a():
+    n, d, sigma2 = 20000, 6, 0.7
+    x = RNG.normal(size=(n, d))
+    out = {}
+
+    # A1: log-likelihood is exactly affine in the MSE, with negative slope
+    mses, elbos = [], []
+    for scale in np.linspace(0.0, 2.0, 25):
+        xhat = x + scale * RNG.normal(size=(n, d))
+        mse = float(((x - xhat) ** 2).mean())          # per-element MSE = L_rec
+        mses.append(mse)
+        elbos.append(gaussian_elbo(x, xhat, sigma2))
+    mses, elbos = np.array(mses), np.array(elbos)
+    slope = -d / (2 * sigma2) / d                       # d(ELBO)/d(mse_per_element)
+    pred = elbos[0] + slope * d * (mses - mses[0])
+    out["affine_max_abs_dev_nats"] = float(np.abs(pred - elbos).max())
+    out["slope_sign_negative"] = bool(slope < 0)
+    out["spearman_mse_vs_elbo"] = float(
+        np.corrcoef(np.argsort(np.argsort(mses)), np.argsort(np.argsort(elbos)))[0, 1])
+
+    # A2: I_LB <= I(X;F') on a jointly Gaussian pair with closed-form MI.
+    # X = F' + noise, F' ~ N(0, I), noise ~ N(0, s^2 I): I = d/2 log(1 + 1/s^2).
+    s2 = 0.5
+    f = RNG.normal(size=(n, d))
+    xg = f + np.sqrt(s2) * RNG.normal(size=(n, d))
+    true_mi = 0.5 * d * np.log(1 + 1.0 / s2)
+    hx = 0.5 * d * (LOG2PI + np.log(1 + s2) + 1)        # differential entropy of X
+    # optimal Gaussian decoder E[X|F'] = F', residual variance s2
+    i_lb = hx + gaussian_elbo(xg, f, s2)
+    out["true_MI_nats"] = float(true_mi)
+    out["I_LB_at_optimal_decoder_nats"] = float(i_lb)
+    out["bound_holds_I_LB_le_MI"] = bool(i_lb <= true_mi + 1e-9)
+    out["bound_gap_nats"] = float(true_mi - i_lb)
+    # a deliberately worse decoder must give a looser bound
+    i_lb_bad = hx + gaussian_elbo(xg, 0.5 * f, s2)
+    out["I_LB_suboptimal_decoder_nats"] = float(i_lb_bad)
+    out["worse_decoder_is_looser"] = bool(i_lb_bad < i_lb)
+
+    # CONTROL: Laplace decoder -> the MSE<->log-likelihood identity must break
+    b = 0.6
+    lap_ll, lap_mse = [], []
+    for scale in np.linspace(0.05, 2.0, 25):
+        xhat = x + scale * RNG.normal(size=(n, d))
+        lap_mse.append(float(((x - xhat) ** 2).mean()))
+        lap_ll.append(float((-d * np.log(2 * b) - np.abs(x - xhat).sum(1) / b).mean()))
+    lap_mse, lap_ll = np.array(lap_mse), np.array(lap_ll)
+    A = np.vstack([lap_mse, np.ones_like(lap_mse)]).T
+    coef, *_ = np.linalg.lstsq(A, lap_ll, rcond=None)
+    out["control_laplace_affine_max_abs_dev_nats"] = float(
+        np.abs(A @ coef - lap_ll).max())
+    return out
+
+
+# ------------------------------------------------------------- (B) independence
+def gaussian_tc_closed_form(cov):
+    return float(0.5 * np.log(np.diag(cov)).sum() - 0.5 * np.linalg.slogdet(cov)[1])
+
+
+def gaussian_tc_from_entropies(cov):
+    k = cov.shape[0]
+    h_marg = sum(0.5 * (LOG2PI + 1 + np.log(cov[i, i])) for i in range(k))
+    h_joint = 0.5 * (k * (LOG2PI + 1) + np.linalg.slogdet(cov)[1])
+    return float(h_marg - h_joint)
+
+
+def tcr_loss(F, eps=0.5):
+    """L_div of Eq. 3 for a single sample's token matrix F in R^{k x d}."""
+    k, d = F.shape
+    Fn = F / np.linalg.norm(F, axis=1, keepdims=True)
+    M = np.eye(k) + (d / (k * eps ** 2)) * (Fn @ Fn.T)
+    return float(-0.5 * np.linalg.slogdet(M)[1])
+
+
+def random_cov(k, rho):
+    """Equicorrelated covariance with unit marginals: TC grows monotonically with rho."""
+    return (1 - rho) * np.eye(k) + rho * np.ones((k, k))
+
+
+def part_b():
+    out = {}
+    k = 8
+
+    # B1: the two Gaussian TC expressions (sum H(f_i) - H(F) vs Eq. 11) agree exactly
+    devs = []
+    for rho in np.linspace(0.0, 0.95, 40):
+        cov = random_cov(k, rho)
+        devs.append(abs(gaussian_tc_closed_form(cov) - gaussian_tc_from_entropies(cov)))
+    out["TC_identity_max_abs_dev_nats"] = float(max(devs))
+
+    # B2: TC = 0 iff the covariance is diagonal (Lemma 3.2 direction)
+    out["TC_at_rho_0"] = gaussian_tc_closed_form(random_cov(k, 0.0))
+    out["TC_at_rho_0.9"] = gaussian_tc_closed_form(random_cov(k, 0.9))
+    out["TC_monotone_in_rho"] = bool(
+        np.all(np.diff([gaussian_tc_closed_form(random_cov(k, r))
+                        for r in np.linspace(0, 0.95, 40)]) > 0))
+
+    # B3: does the paper's L_div actually track TC?  Sample token sets with a
+    # controlled correlation and compare the *ordering* and correlation of
+    # L_div against the true Gaussian TC of the token Gram/covariance.
+    d = 128
+    tcs, divs = [], []
+    for rho in np.linspace(0.0, 0.95, 40):
+        cov = random_cov(k, rho)
+        L = np.linalg.cholesky(cov)
+        F = L @ RNG.normal(size=(k, d))          # k tokens, correlation rho
+        emp = np.cov(F)                          # empirical k x k covariance
+        tcs.append(gaussian_tc_closed_form(emp))
+        divs.append(tcr_loss(F))
+    tcs, divs = np.array(tcs), np.array(divs)
+    out["pearson_Ldiv_vs_TC"] = float(np.corrcoef(divs, tcs)[0, 1])
+    # is the relation an *exact* identity, or only a monotone surrogate?
+    A = np.vstack([tcs, np.ones_like(tcs)]).T
+    coef, *_ = np.linalg.lstsq(A, divs, rcond=None)
+    out["affine_fit_Ldiv_vs_TC_slope"] = float(coef[0])
+    out["affine_fit_Ldiv_vs_TC_max_abs_resid_nats"] = float(np.abs(A @ coef - divs).max())
+    out["Ldiv_range_nats"] = float(divs.max() - divs.min())
+    out["spearman_Ldiv_vs_TC"] = float(np.corrcoef(
+        np.argsort(np.argsort(divs)), np.argsort(np.argsort(tcs)))[0, 1])
+    out["Ldiv_at_rho_0"] = float(divs[0])
+    out["Ldiv_at_rho_0.95"] = float(divs[-1])
+    out["Ldiv_increases_with_redundancy"] = bool(divs[-1] > divs[0])
+    # L_div is minimised by orthogonal tokens (Lemma 3.2)
+    Q = np.linalg.qr(RNG.normal(size=(d, k)))[0].T      # k orthonormal rows
+    out["Ldiv_orthogonal_tokens"] = tcr_loss(Q)
+    dup = np.repeat(Q[:1], k, axis=0) + 1e-6 * RNG.normal(size=(k, d))
+    out["Ldiv_collapsed_tokens"] = tcr_loss(dup)
+    out["orthogonal_is_the_minimiser"] = bool(
+        out["Ldiv_orthogonal_tokens"] < out["Ldiv_collapsed_tokens"])
+    rand_devs = [tcr_loss(RNG.normal(size=(k, d))) for _ in range(200)]
+    out["Ldiv_random_tokens_mean"] = float(np.mean(rand_devs))
+    out["orthogonal_beats_random"] = bool(
+        out["Ldiv_orthogonal_tokens"] <= min(rand_devs))
+
+    # B4: Sylvester check - the d x d form in Eq. 3 and the k x k Gram form agree
+    F = RNG.normal(size=(k, d))
+    Fn = F / np.linalg.norm(F, axis=1, keepdims=True)
+    c = d / (k * 0.5 ** 2)
+    lhs = np.linalg.slogdet(np.eye(d) + c * (Fn.T @ Fn))[1]
+    rhs = np.linalg.slogdet(np.eye(k) + c * (Fn @ Fn.T))[1]
+    out["sylvester_abs_dev"] = float(abs(lhs - rhs))
+
+    # CONTROL: non-Gaussian latents -> the Gaussian entropy formula (and hence the
+    # claimed TC identity) is no longer valid.  We compare the Gaussian-formula TC
+    # against a k-NN (Kozachenko-Leonenko) estimate of the true TC.
+    def kl_entropy(samples):
+        from scipy.spatial import cKDTree
+        from scipy.special import digamma, gammaln
+        n, dd = samples.shape
+        tree = cKDTree(samples)
+        eps_ = tree.query(samples, k=4)[0][:, 3]
+        eps_ = np.maximum(eps_, 1e-12)
+        log_vol = dd / 2 * np.log(np.pi) - gammaln(dd / 2 + 1)
+        return float(digamma(n) - digamma(3) + log_vol + dd * np.mean(np.log(eps_)))
+
+    def true_tc(samples):
+        h_joint = kl_entropy(samples)
+        h_marg = sum(kl_entropy(samples[:, [i]]) for i in range(samples.shape[1]))
+        return h_marg - h_joint
+
+    n, kk, rho = 40000, 3, 0.7
+    cov = random_cov(kk, rho)
+    L = np.linalg.cholesky(cov)
+    zg = RNG.normal(size=(n, kk)) @ L.T
+    # matched-covariance uniform (platykurtic) and Student-t (heavy-tailed) latents
+    u = (RNG.uniform(-np.sqrt(3), np.sqrt(3), size=(n, kk))) @ L.T
+    t = (RNG.standard_t(3, size=(n, kk)) / np.sqrt(3.0)) @ L.T
+    for nm, z in (("gaussian", zg), ("uniform", u), ("student_t3", t)):
+        emp = np.cov(z.T)
+        out[f"control_{nm}_TC_gauss_formula"] = gaussian_tc_closed_form(emp)
+        out[f"control_{nm}_TC_knn_estimate"] = float(true_tc(z))
+        out[f"control_{nm}_abs_error"] = abs(
+            out[f"control_{nm}_TC_gauss_formula"] - out[f"control_{nm}_TC_knn_estimate"])
+    return out
+
+
+if __name__ == "__main__":
+    report["A_sufficiency"] = part_a()
+    report["B_independence"] = part_b()
+    print(json.dumps(report, indent=2))
+    with open("results/theorem31_audit.json", "w") as fh:
+        json.dump(report, fh, indent=2)
+
+    a, b = report["A_sufficiency"], report["B_independence"]
+    print("\n---- verdict ----")
+    print(f"A. MSE <-> Gaussian ELBO affine identity: max dev "
+          f"{a['affine_max_abs_dev_nats']:.3e} nats  "
+          f"(Laplace control: {a['control_laplace_affine_max_abs_dev_nats']:.3f} nats)")
+    print(f"A. I_LB <= I(X;F'): {a['bound_holds_I_LB_le_MI']} "
+          f"(gap {a['bound_gap_nats']:.3e} nats at the optimal decoder)")
+    print(f"B. Gaussian TC identity: max dev {b['TC_identity_max_abs_dev_nats']:.3e} nats")
+    print(f"B. Sylvester d x d == k x k: {b['sylvester_abs_dev']:.3e}")
+    print(f"B. corr(L_div, TC) = {b['pearson_Ldiv_vs_TC']:.4f} "
+          f"(Spearman {b['spearman_Ldiv_vs_TC']:.4f})")
+    print(f"B. L_div(orthogonal) = {b['Ldiv_orthogonal_tokens']:.3f} < "
+          f"L_div(collapsed) = {b['Ldiv_collapsed_tokens']:.3f}")
+    print("B. control (Gaussian formula vs kNN TC): " + ", ".join(
+        f"{nm} err {b[f'control_{nm}_abs_error']:.3f}"
+        for nm in ("gaussian", "uniform", "student_t3")))

fig_theorem31.py

@@ -0,0 +1,100 @@
+"""Figures for the Theorem 3.1 audit (Claim 2)."""
+
+import csv
+import json
+
+import numpy as np
+import plotly.graph_objects as go
+from plotly.subplots import make_subplots
+
+RNG = np.random.default_rng(0)
+LOG2PI = np.log(2 * np.pi)
+
+
+def tcr(F, eps=0.5):
+    k, d = F.shape
+    Fn = F / np.linalg.norm(F, axis=1, keepdims=True)
+    return float(-0.5 * np.linalg.slogdet(
+        np.eye(k) + (d / (k * eps ** 2)) * (Fn @ Fn.T))[1])
+
+
+def gauss_tc(cov):
+    return float(0.5 * np.log(np.diag(cov)).sum() - 0.5 * np.linalg.slogdet(cov)[1])
+
+
+k, d = 8, 128
+rhos = np.linspace(0.0, 0.95, 40)
+tcs, divs = [], []
+for rho in rhos:
+    cov = (1 - rho) * np.eye(k) + rho * np.ones((k, k))
+    F = np.linalg.cholesky(cov) @ RNG.normal(size=(k, d))
+    tcs.append(gauss_tc(np.cov(F)))
+    divs.append(tcr(F))
+tcs, divs = np.array(tcs), np.array(divs)
+
+# panel B: MSE vs Gaussian ELBO (exact) and vs Laplace log-likelihood (control)
+n, dd, s2, b = 20000, 6, 0.7, 0.6
+x = RNG.normal(size=(n, dd))
+scales = np.linspace(0.05, 2.0, 25)
+mses, gll, lll = [], [], []
+for s in scales:
+    xh = x + s * RNG.normal(size=(n, dd))
+    mses.append(float(((x - xh) ** 2).mean()))
+    sq = ((x - xh) ** 2).sum(1)
+    gll.append(float((-0.5 * dd * (LOG2PI + np.log(s2)) - sq / (2 * s2)).mean()))
+    lll.append(float((-dd * np.log(2 * b) - np.abs(x - xh).sum(1) / b).mean()))
+mses, gll, lll = np.array(mses), np.array(gll), np.array(lll)
+
+
+def resid(xv, yv):
+    A = np.vstack([xv, np.ones_like(xv)]).T
+    c, *_ = np.linalg.lstsq(A, yv, rcond=None)
+    return A @ c - yv
+
+
+fig = make_subplots(
+    rows=1, cols=2, horizontal_spacing=0.11,
+    subplot_titles=(
+        "Leg B — ℒ_div tracks Total Correlation (r = 0.9995)",
+        "Leg A — MSE ⇔ log-likelihood: exact for Gaussian, broken for Laplace"))
+
+fig.add_scatter(x=tcs, y=divs, mode="markers", name="ℒ_div vs TC (Gaussian latents)",
+                marker=dict(size=8, color="#2F6F8F"), row=1, col=1)
+A = np.vstack([tcs, np.ones_like(tcs)]).T
+c, *_ = np.linalg.lstsq(A, divs, rcond=None)
+fig.add_scatter(x=tcs, y=A @ c, mode="lines", name=f"best affine fit (slope {c[0]:.3f})",
+                line=dict(color="#B07C2B", dash="dash"), row=1, col=1)
+fig.update_xaxes(title_text="true Gaussian TC(F′) [nats]", row=1, col=1)
+fig.update_yaxes(title_text="ℒ_div  (Eq. 3)", row=1, col=1)
+
+fig.add_scatter(x=mses, y=resid(mses, gll), mode="markers+lines",
+                name="Gaussian decoder — residual from affine",
+                marker=dict(size=7, color="#2F6F8F"), row=1, col=2)
+fig.add_scatter(x=mses, y=resid(mses, lll), mode="markers+lines",
+                name="Laplace decoder (CONTROL) — residual",
+                marker=dict(size=7, color="#A33B3B"), row=1, col=2)
+fig.update_xaxes(title_text="reconstruction MSE (ℒ_rec)", row=1, col=2)
+fig.update_yaxes(title_text="residual from best affine fit [nats]", row=1, col=2)
+
+fig.update_layout(template="plotly_white", height=440,
+                  legend=dict(orientation="h", y=-0.22),
+                  title="Theorem 3.1 numerical audit — both legs, each with its control")
+fig.write_html("results/claim2_theorem31.html", include_plotlyjs="cdn")
+
+with open("results/claim2_theorem31.csv", "w", newline="") as fh:
+    w = csv.writer(fh)
+    w.writerow(["panel", "x", "y", "series"])
+    for a, bb in zip(tcs, divs):
+        w.writerow(["B_ldiv_vs_tc", f"{a:.6f}", f"{bb:.6f}", "ldiv"])
+    rg, rl = resid(mses, gll), resid(mses, lll)
+    for m, g, l in zip(mses, rg, rl):
+        w.writerow(["A_residual", f"{m:.6f}", f"{g:.6e}", "gaussian"])
+        w.writerow(["A_residual", f"{m:.6f}", f"{l:.6e}", "laplace_control"])
+
+print(json.dumps({
+    "ldiv_vs_tc_slope": float(c[0]),
+    "ldiv_vs_tc_max_resid_nats": float(np.abs(A @ c - divs).max()),
+    "gaussian_max_resid_nats": float(np.abs(resid(mses, gll)).max()),
+    "laplace_control_max_resid_nats": float(np.abs(resid(mses, lll)).max()),
+}, indent=2))
+print("wrote results/claim2_theorem31.html + .csv")

jobs_adftd.json

@@ -0,0 +1,146 @@
+{
+ "tsfp_scratch": [
+  "--data-dir",
+  "data/ADFTD/processed",
+  "--batch-size",
+  "256",
+  "--pre-epochs",
+  "100",
+  "--ft-epochs",
+  "100",
+  "--patience",
+  "10",
+  "--seeds",
+  "41",
+  "42",
+  "43",
+  "--model",
+  "tsfp",
+  "--mode",
+  "scratch"
+ ],
+ "tsfp_rec": [
+  "--data-dir",
+  "data/ADFTD/processed",
+  "--batch-size",
+  "256",
+  "--pre-epochs",
+  "100",
+  "--ft-epochs",
+  "100",
+  "--patience",
+  "10",
+  "--seeds",
+  "41",
+  "42",
+  "43",
+  "--model",
+  "tsfp",
+  "--mode",
+  "pretrain_ft",
+  "--use-div",
+  "0"
+ ],
+ "tsfp_rec_div": [
+  "--data-dir",
+  "data/ADFTD/processed",
+  "--batch-size",
+  "256",
+  "--pre-epochs",
+  "100",
+  "--ft-epochs",
+  "100",
+  "--patience",
+  "10",
+  "--seeds",
+  "41",
+  "42",
+  "43",
+  "--model",
+  "tsfp",
+  "--mode",
+  "pretrain_ft",
+  "--use-div",
+  "1"
+ ],
+ "timae_scratch": [
+  "--data-dir",
+  "data/ADFTD/processed",
+  "--batch-size",
+  "256",
+  "--pre-epochs",
+  "100",
+  "--ft-epochs",
+  "100",
+  "--patience",
+  "10",
+  "--seeds",
+  "41",
+  "42",
+  "43",
+  "--model",
+  "timae",
+  "--mode",
+  "scratch"
+ ],
+ "timae_pretrain": [
+  "--data-dir",
+  "data/ADFTD/processed",
+  "--batch-size",
+  "256",
+  "--pre-epochs",
+  "100",
+  "--ft-epochs",
+  "100",
+  "--patience",
+  "10",
+  "--seeds",
+  "41",
+  "42",
+  "43",
+  "--model",
+  "timae",
+  "--mode",
+  "pretrain_ft"
+ ],
+ "simmtm_scratch": [
+  "--data-dir",
+  "data/ADFTD/processed",
+  "--batch-size",
+  "256",
+  "--pre-epochs",
+  "100",
+  "--ft-epochs",
+  "100",
+  "--patience",
+  "10",
+  "--seeds",
+  "41",
+  "42",
+  "43",
+  "--model",
+  "simmtm",
+  "--mode",
+  "scratch"
+ ],
+ "simmtm_pretrain": [
+  "--data-dir",
+  "data/ADFTD/processed",
+  "--batch-size",
+  "256",
+  "--pre-epochs",
+  "100",
+  "--ft-epochs",
+  "100",
+  "--patience",
+  "10",
+  "--seeds",
+  "41",
+  "42",
+  "43",
+  "--model",
+  "simmtm",
+  "--mode",
+  "pretrain_ft"
+ ]
+}

jobs_simmtm.json

@@ -0,0 +1,21 @@
+{
+ "simmtm_pretrain": [
+  "--data-dir",
+  "data/ADFTD/processed",
+  "--batch-size",
+  "256",
+  "--pre-epochs",
+  "100",
+  "--ft-epochs",
+  "100",
+  "--patience",
+  "10",
+  "--seeds",
+  "41",
+  "42",
+  "--model",
+  "simmtm",
+  "--mode",
+  "pretrain_ft"
+ ]
+}

jobs_sweep.json

@@ -0,0 +1,338 @@
+{
+ "ptbxl_k6_r0.5": [
+  "--data-dir",
+  "data/PTBXL/processed",
+  "--batch-size",
+  "256",
+  "--patch-size",
+  "10",
+  "--pre-epochs",
+  "30",
+  "--ft-epochs",
+  "60",
+  "--patience",
+  "8",
+  "--stride",
+  "5",
+  "--seeds",
+  "41",
+  "--model",
+  "tsfp",
+  "--mode",
+  "pretrain_ft",
+  "--use-div",
+  "1",
+  "--k",
+  "6",
+  "--mask-ratio",
+  "0.5"
+ ],
+ "ptbxl_k6_r0.6": [
+  "--data-dir",
+  "data/PTBXL/processed",
+  "--batch-size",
+  "256",
+  "--patch-size",
+  "10",
+  "--pre-epochs",
+  "30",
+  "--ft-epochs",
+  "60",
+  "--patience",
+  "8",
+  "--stride",
+  "5",
+  "--seeds",
+  "41",
+  "--model",
+  "tsfp",
+  "--mode",
+  "pretrain_ft",
+  "--use-div",
+  "1",
+  "--k",
+  "6",
+  "--mask-ratio",
+  "0.6"
+ ],
+ "ptbxl_k6_r0.7": [
+  "--data-dir",
+  "data/PTBXL/processed",
+  "--batch-size",
+  "256",
+  "--patch-size",
+  "10",
+  "--pre-epochs",
+  "30",
+  "--ft-epochs",
+  "60",
+  "--patience",
+  "8",
+  "--stride",
+  "5",
+  "--seeds",
+  "41",
+  "--model",
+  "tsfp",
+  "--mode",
+  "pretrain_ft",
+  "--use-div",
+  "1",
+  "--k",
+  "6",
+  "--mask-ratio",
+  "0.7"
+ ],
+ "ptbxl_k6_r0.8": [
+  "--data-dir",
+  "data/PTBXL/processed",
+  "--batch-size",
+  "256",
+  "--patch-size",
+  "10",
+  "--pre-epochs",
+  "30",
+  "--ft-epochs",
+  "60",
+  "--patience",
+  "8",
+  "--stride",
+  "5",
+  "--seeds",
+  "41",
+  "--model",
+  "tsfp",
+  "--mode",
+  "pretrain_ft",
+  "--use-div",
+  "1",
+  "--k",
+  "6",
+  "--mask-ratio",
+  "0.8"
+ ],
+ "ptbxl_k8_r0.5": [
+  "--data-dir",
+  "data/PTBXL/processed",
+  "--batch-size",
+  "256",
+  "--patch-size",
+  "10",
+  "--pre-epochs",
+  "30",
+  "--ft-epochs",
+  "60",
+  "--patience",
+  "8",
+  "--stride",
+  "5",
+  "--seeds",
+  "41",
+  "--model",
+  "tsfp",
+  "--mode",
+  "pretrain_ft",
+  "--use-div",
+  "1",
+  "--k",
+  "8",
+  "--mask-ratio",
+  "0.5"
+ ],
+ "ptbxl_k8_r0.6": [
+  "--data-dir",
+  "data/PTBXL/processed",
+  "--batch-size",
+  "256",
+  "--patch-size",
+  "10",
+  "--pre-epochs",
+  "30",
+  "--ft-epochs",
+  "60",
+  "--patience",
+  "8",
+  "--stride",
+  "5",
+  "--seeds",
+  "41",
+  "--model",
+  "tsfp",
+  "--mode",
+  "pretrain_ft",
+  "--use-div",
+  "1",
+  "--k",
+  "8",
+  "--mask-ratio",
+  "0.6"
+ ],
+ "ptbxl_k8_r0.7": [
+  "--data-dir",
+  "data/PTBXL/processed",
+  "--batch-size",
+  "256",
+  "--patch-size",
+  "10",
+  "--pre-epochs",
+  "30",
+  "--ft-epochs",
+  "60",
+  "--patience",
+  "8",
+  "--stride",
+  "5",
+  "--seeds",
+  "41",
+  "--model",
+  "tsfp",
+  "--mode",
+  "pretrain_ft",
+  "--use-div",
+  "1",
+  "--k",
+  "8",
+  "--mask-ratio",
+  "0.7"
+ ],
+ "ptbxl_k8_r0.8": [
+  "--data-dir",
+  "data/PTBXL/processed",
+  "--batch-size",
+  "256",
+  "--patch-size",
+  "10",
+  "--pre-epochs",
+  "30",
+  "--ft-epochs",
+  "60",
+  "--patience",
+  "8",
+  "--stride",
+  "5",
+  "--seeds",
+  "41",
+  "--model",
+  "tsfp",
+  "--mode",
+  "pretrain_ft",
+  "--use-div",
+  "1",
+  "--k",
+  "8",
+  "--mask-ratio",
+  "0.8"
+ ],
+ "ptbxl_k10_r0.5": [
+  "--data-dir",
+  "data/PTBXL/processed",
+  "--batch-size",
+  "256",
+  "--patch-size",
+  "10",
+  "--pre-epochs",
+  "30",
+  "--ft-epochs",
+  "60",
+  "--patience",
+  "8",
+  "--stride",
+  "5",
+  "--seeds",
+  "41",
+  "--model",
+  "tsfp",
+  "--mode",
+  "pretrain_ft",
+  "--use-div",
+  "1",
+  "--k",
+  "10",
+  "--mask-ratio",
+  "0.5"
+ ],
+ "ptbxl_k10_r0.6": [
+  "--data-dir",
+  "data/PTBXL/processed",
+  "--batch-size",
+  "256",
+  "--patch-size",
+  "10",
+  "--pre-epochs",
+  "30",
+  "--ft-epochs",
+  "60",
+  "--patience",
+  "8",
+  "--stride",
+  "5",
+  "--seeds",
+  "41",
+  "--model",
+  "tsfp",
+  "--mode",
+  "pretrain_ft",
+  "--use-div",
+  "1",
+  "--k",
+  "10",
+  "--mask-ratio",
+  "0.6"
+ ],
+ "ptbxl_k10_r0.7": [
+  "--data-dir",
+  "data/PTBXL/processed",
+  "--batch-size",
+  "256",
+  "--patch-size",
+  "10",
+  "--pre-epochs",
+  "30",
+  "--ft-epochs",
+  "60",
+  "--patience",
+  "8",
+  "--stride",
+  "5",
+  "--seeds",
+  "41",
+  "--model",
+  "tsfp",
+  "--mode",
+  "pretrain_ft",
+  "--use-div",
+  "1",
+  "--k",
+  "10",
+  "--mask-ratio",
+  "0.7"
+ ],
+ "ptbxl_k10_r0.8": [
+  "--data-dir",
+  "data/PTBXL/processed",
+  "--batch-size",
+  "256",
+  "--patch-size",
+  "10",
+  "--pre-epochs",
+  "30",
+  "--ft-epochs",
+  "60",
+  "--patience",
+  "8",
+  "--stride",
+  "5",
+  "--seeds",
+  "41",
+  "--model",
+  "tsfp",
+  "--mode",
+  "pretrain_ft",
+  "--use-div",
+  "1",
+  "--k",
+  "10",
+  "--mask-ratio",
+  "0.8"
+ ]
+}

make_poster_figs.py

@@ -0,0 +1,175 @@
+"""Render the poster's PNG figures from the logged reproduction results.
+
+Every number here comes from results/*.json produced by the runs; nothing is
+hand-entered except the paper's own reported values, which are labelled as such.
+"""
+
+import json
+import os
+
+import numpy as np
+import plotly.graph_objects as go
+from plotly.subplots import make_subplots
+
+OUT = "poster/images"
+os.makedirs(OUT, exist_ok=True)
+ACCENT, GOLD, GREY, RED = "#17697B", "#B07C2B", "#8A8A8A", "#A33B3B"
+FONT = dict(family="Helvetica, Arial, sans-serif", size=30, color="#1A1A1A")
+SCALE = 3
+
+
+def save(fig, name, w, h):
+    fig.update_layout(template="plotly_white", font=FONT,
+                      margin=dict(l=90, r=30, t=80, b=70))
+    fig.write_image(f"{OUT}/{name}.png", width=w, height=h, scale=SCALE)
+    print(f"wrote {OUT}/{name}.png  ({w*SCALE}x{h*SCALE} px)")
+
+
+# ------------------------------------------------------------------ Fig 1: Claim 3
+def fig_ablation():
+    s = json.load(open("results/adftd_summary.json"))
+    rows = {r["config"]: r for r in s["rows"]}
+    order = ["tsfp_scratch", "tsfp_rec", "tsfp_rec_div"]
+    order = [o for o in order if o in rows]
+    if len(order) < 3:
+        print("skip ablation fig (incomplete)")
+        return
+    xs = ["Scratch", "Pre-trained<br>(ℒ_rec)", "Pre-trained<br>(ℒ_rec + ℒ_div)"]
+    fig = go.Figure()
+    fig.add_bar(name="this reproduction (3 seeds)", x=xs,
+                y=[rows[o]["f1"] for o in order],
+                error_y=dict(type="data", array=[rows[o]["f1_sd"] for o in order],
+                             thickness=3, width=14),
+                marker_color=ACCENT,
+                text=[f"{rows[o]['f1']:.2f}" for o in order],
+                textposition="outside", textfont=dict(size=32))
+    fig.add_bar(name="paper (Table 2)", x=xs,
+                y=[rows[o]["paper_f1"] for o in order], marker_color=GOLD,
+                text=[f"{rows[o]['paper_f1']:.2f}" for o in order],
+                textposition="outside", textfont=dict(size=32))
+    fig.update_layout(barmode="group", yaxis_title="ADFTD macro F1 (%)",
+                      yaxis_range=[38, 58], height=620, width=1000,
+                      legend=dict(orientation="h", y=1.13, x=0))
+    save(fig, "claim3_ablation", 1000, 620)
+
+
+# ------------------------------------------------------------------ Fig 2: Claim 4
+def fig_promotion():
+    s = json.load(open("results/adftd_summary.json"))
+    p = s["promotions"]
+    names, ours = [], []
+    for k in ("TS-Fingerprint", "SimMTM", "Ti-MAE"):
+        if p[k]["f1"]:
+            names.append(k)
+            ours.append(p[k]["f1"]["rel_pct"])
+    if not names:
+        print("skip promotion fig")
+        return
+    paper = {"TS-Fingerprint": 13.07, "SimMTM": 4.42, "Ti-MAE": -0.96}
+    fig = go.Figure()
+    fig.add_bar(name="this reproduction", x=names, y=ours, marker_color=ACCENT,
+                text=[f"{v:+.1f}%" for v in ours], textposition="outside",
+                textfont=dict(size=32))
+    fig.add_bar(name="paper (Table 4)", x=names, y=[paper[n] for n in names],
+                marker_color=GOLD, text=[f"{paper[n]:+.1f}%" for n in names],
+                textposition="outside", textfont=dict(size=32))
+    lo = min(ours + [paper[n] for n in names])
+    hi = max(ours + [paper[n] for n in names])
+    fig.add_hline(y=0, line=dict(color="#555", width=2))
+    fig.update_layout(barmode="group",
+                      yaxis_title="relative F1 gain from pre-training (%)",
+                      yaxis_range=[min(lo * 1.5, -3), hi * 1.35],
+                      height=620, width=1000,
+                      legend=dict(orientation="h", y=1.13, x=0))
+    save(fig, "claim4_promotion", 1000, 620)
+
+
+# ------------------------------------------------------------------ Fig 3: Claim 6
+def fig_sweep():
+    p = "results/sweep_summary.json"
+    if not os.path.exists(p):
+        print("skip sweep fig (not run yet)")
+        return
+    d = json.load(open(p))
+    ks, rs = [6, 8, 10], [0.5, 0.6, 0.7, 0.8]
+    ours = np.full((3, 4), np.nan)
+    for row in d["rows"]:
+        ours[ks.index(row["k"]), rs.index(row["r"])] = row["f1"]
+    paper = np.array([[54.81, 56.68, 55.41, 54.32],
+                      [62.60, 63.51, 62.10, 59.19],
+                      [51.76, 50.70, 56.33, 55.70]])
+    fig = make_subplots(rows=1, cols=2, horizontal_spacing=0.14,
+                        subplot_titles=("this reproduction (PTB-XL, 1 seed)",
+                                        "paper (Table 3)"))
+    for j, (m, lbl) in enumerate(((ours, "ours"), (paper, "paper"))):
+        fig.add_heatmap(z=m, x=[f"r={r}" for r in rs], y=[f"k={k}" for k in ks],
+                        colorscale=[[0, "#F2F7F8"], [1, ACCENT]], showscale=False,
+                        text=[[("" if np.isnan(v) else f"{v:.2f}") for v in row]
+                              for row in m],
+                        texttemplate="%{text}", textfont=dict(size=30),
+                        row=1, col=j + 1)
+    fig.update_layout(height=560, width=1300,
+                      title_text="macro F1 (%) across bottleneck size k and mask ratio r")
+    save(fig, "claim6_sweep", 1300, 560)
+
+
+# ------------------------------------------------------------------ Fig 4: Claim 2
+def fig_theorem():
+    rng = np.random.default_rng(0)
+    LOG2PI = np.log(2 * np.pi)
+    n, dd, s2, b = 20000, 6, 0.7, 0.6
+    x = rng.normal(size=(n, dd))
+    mses, gll, lll = [], [], []
+    for s in np.linspace(0.05, 2.0, 25):
+        xh = x + s * rng.normal(size=(n, dd))
+        mses.append(float(((x - xh) ** 2).mean()))
+        sq = ((x - xh) ** 2).sum(1)
+        gll.append(float((-0.5 * dd * (LOG2PI + np.log(s2)) - sq / (2 * s2)).mean()))
+        lll.append(float((-dd * np.log(2 * b) - np.abs(x - xh).sum(1) / b).mean()))
+    mses, gll, lll = np.array(mses), np.array(gll), np.array(lll)
+
+    def resid(xv, yv):
+        A = np.vstack([xv, np.ones_like(xv)]).T
+        c, *_ = np.linalg.lstsq(A, yv, rcond=None)
+        return np.abs(A @ c - yv)
+
+    fig = go.Figure()
+    fig.add_scatter(x=mses, y=np.maximum(resid(mses, gll), 1e-16), mode="lines+markers",
+                    name="Gaussian decoder (theorem's assumption)",
+                    line=dict(color=ACCENT, width=4), marker=dict(size=11))
+    fig.add_scatter(x=mses, y=resid(mses, lll), mode="lines+markers",
+                    name="Laplace decoder — CONTROL",
+                    line=dict(color=RED, width=4, dash="dash"), marker=dict(size=11))
+    fig.update_yaxes(type="log", title_text="|residual| from affine fit  [nats]",
+                     exponentformat="power")
+    fig.update_xaxes(title_text="reconstruction MSE  (ℒ_rec)")
+    fig.update_layout(height=560, width=1000,
+                      legend=dict(orientation="h", y=1.16, x=0))
+    save(fig, "claim2_control", 1000, 560)
+
+
+# ------------------------------------------------------------------ Fig 5: Claim 5
+def fig_rank():
+    d = json.load(open("results/table1_rank_audit.json"))
+    ranks = d["avg_rank_all_methods_5metrics"]
+    items = sorted(ranks.items(), key=lambda kv: kv[1])
+    names = [k if k != "Ours" else "TS-Fingerprint" for k, _ in items]
+    vals = [v for _, v in items]
+    cols = [GOLD if n == "TS-Fingerprint" else
+            (ACCENT if n == "Medformer" else GREY) for n in names]
+    fig = go.Figure()
+    fig.add_bar(x=vals, y=names, orientation="h", marker_color=cols,
+                text=[f"{v:.2f}" for v in vals], textposition="outside",
+                textfont=dict(size=28))
+    fig.update_layout(xaxis_title="average rank over 7 datasets × 5 metrics (lower better)",
+                      height=760, width=1000, yaxis=dict(autorange="reversed"),
+                      xaxis_range=[0, 12])
+    save(fig, "claim5_rank", 1000, 760)
+
+
+if __name__ == "__main__":
+    for f in (fig_ablation, fig_promotion, fig_sweep, fig_theorem, fig_rank):
+        try:
+            f()
+        except Exception as e:  # noqa: BLE001
+            print(f"skip {f.__name__}: {type(e).__name__}: {e}")

model.py

@@ -0,0 +1,322 @@
+"""Independent re-implementation of TS-Fingerprint (ICML 2026, OpenReview lrRBHFIgaK).
+
+No official code was released for the paper, so every component here is derived
+from the paper text: Sec 3.2 (architecture), Sec 3.3 (losses), Sec 3.4 (attention
+pooling head) and Sec 4.1 (hyperparameters: 6-layer encoder, 2-layer decoder,
+d=128, 8 heads, k=8, mask ratio 0.6, lambda=1e-4).
+"""
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class PatchEmbed(nn.Module):
+    """Non-overlapping patching over time; all C channels enter the patch (Medformer-style)."""
+
+    def __init__(self, c_in, patch_size, d_model):
+        super().__init__()
+        self.patch_size = patch_size
+        self.proj = nn.Linear(c_in * patch_size, d_model)
+
+    def forward(self, x):  # x: (B, T, C)
+        b, t, c = x.shape
+        n = t // self.patch_size
+        x = x[:, : n * self.patch_size].reshape(b, n, self.patch_size * c)
+        return self.proj(x)  # (B, N, d)
+
+
+class CrossAttnBlock(nn.Module):
+    """One iteration of the Perceiver-style bottleneck: latents cross-attend to the
+    patch sequence, then self-attend among themselves (Sec 3.2, 'iterative cross-attention')."""
+
+    def __init__(self, d_model, n_heads, dropout=0.1):
+        super().__init__()
+        self.ln_q1 = nn.LayerNorm(d_model)
+        self.ln_kv = nn.LayerNorm(d_model)
+        self.cross = nn.MultiheadAttention(d_model, n_heads, dropout=dropout, batch_first=True)
+        self.ln_q2 = nn.LayerNorm(d_model)
+        self.self_attn = nn.MultiheadAttention(d_model, n_heads, dropout=dropout, batch_first=True)
+        self.ln_q3 = nn.LayerNorm(d_model)
+        self.ff = nn.Sequential(
+            nn.Linear(d_model, 4 * d_model), nn.GELU(), nn.Dropout(dropout),
+            nn.Linear(4 * d_model, d_model), nn.Dropout(dropout),
+        )
+
+    def forward(self, q, kv, kv_mask=None, need_weights=False):
+        h, attn = self.cross(self.ln_q1(q), self.ln_kv(kv), self.ln_kv(kv),
+                             key_padding_mask=kv_mask, need_weights=need_weights,
+                             average_attn_weights=True)
+        q = q + h
+        h, _ = self.self_attn(self.ln_q2(q), self.ln_q2(q), self.ln_q2(q), need_weights=False)
+        q = q + h
+        q = q + self.ff(self.ln_q3(q))
+        return q, attn
+
+
+class FingerprintEncoder(nn.Module):
+    """E_theta : X -> F' in R^{k x d}. Fixed-rank bottleneck via a learnable query set
+    Q in R^{k x d} with k << T (Claim 1)."""
+
+    def __init__(self, c_in, patch_size, d_model=128, n_heads=8, n_layers=6, k=8,
+                 max_patches=512, dropout=0.1):
+        super().__init__()
+        self.k = k
+        self.d_model = d_model
+        self.patch_embed = PatchEmbed(c_in, patch_size, d_model)
+        self.pos = nn.Parameter(torch.zeros(1, max_patches, d_model))
+        nn.init.trunc_normal_(self.pos, std=0.02)
+        # the learnable query set Q
+        self.Q = nn.Parameter(torch.randn(k, d_model) * 0.02)
+        self.blocks = nn.ModuleList(
+            [CrossAttnBlock(d_model, n_heads, dropout) for _ in range(n_layers)]
+        )
+        self.norm = nn.LayerNorm(d_model)
+
+    def embed_patches(self, x):
+        p = self.patch_embed(x)
+        return p + self.pos[:, : p.shape[1]]
+
+    def forward(self, x, kv_mask=None, return_attn=False):
+        """x: (B, T, C) -> F': (B, k, d). Output shape is independent of T."""
+        p = self.embed_patches(x)
+        q = self.Q.unsqueeze(0).expand(x.shape[0], -1, -1)
+        attns = []
+        for blk in self.blocks:
+            q, a = blk(q, p, kv_mask=kv_mask, need_weights=return_attn)
+            if return_attn:
+                attns.append(a)
+        f = self.norm(q)
+        if return_attn:
+            return f, attns
+        return f
+
+
+class FingerprintDecoder(nn.Module):
+    """D_phi conditions *solely* on F'. Mask tokens carry only positional information,
+    so the data-processing chain X -> F' -> X_hat is strict (Sec 3.2, Decoder)."""
+
+    def __init__(self, c_in, patch_size, d_model=128, n_heads=8, n_layers=2,
+                 max_patches=512, dropout=0.1):
+        super().__init__()
+        self.mask_token = nn.Parameter(torch.zeros(1, 1, d_model))
+        self.pos = nn.Parameter(torch.zeros(1, max_patches, d_model))
+        nn.init.trunc_normal_(self.pos, std=0.02)
+        self.blocks = nn.ModuleList(
+            [CrossAttnBlock(d_model, n_heads, dropout) for _ in range(n_layers)]
+        )
+        self.norm = nn.LayerNorm(d_model)
+        self.head = nn.Linear(d_model, c_in * patch_size)
+
+    def forward(self, f, n_patches):
+        b = f.shape[0]
+        t = self.mask_token.expand(b, n_patches, -1) + self.pos[:, :n_patches]
+        for blk in self.blocks:
+            t, _ = blk(t, f)
+        return self.head(self.norm(t))  # (B, N, patch*C)
+
+
+def total_coding_rate_loss(f, eps=0.5):
+    """L_div = -1/2 log det(I + d/(k eps^2) F'^T F')   (Eq. 3).
+
+    Computed per sample over its k tokens; by Sylvester's identity the d x d and
+    k x k forms have identical value, and we use the cheaper k x k Gram form.
+    """
+    b, k, d = f.shape
+    f = f.float()  # slogdet is not fp16-safe; keep the geometric term in fp32
+    fn = F.normalize(f, dim=-1)  # fixed-energy constraint of Lemma 3.2
+    gram = torch.bmm(fn, fn.transpose(1, 2))  # (B, k, k)
+    ident = torch.eye(k, device=f.device, dtype=f.dtype).unsqueeze(0)
+    mat = ident + (d / (k * eps ** 2)) * gram
+    return -0.5 * torch.linalg.slogdet(mat)[1].mean()
+
+
+class TSFingerprint(nn.Module):
+    def __init__(self, c_in, patch_size, n_classes, d_model=128, n_heads=8,
+                 enc_layers=6, dec_layers=2, k=8, max_patches=512, dropout=0.1):
+        super().__init__()
+        self.patch_size = patch_size
+        self.encoder = FingerprintEncoder(c_in, patch_size, d_model, n_heads,
+                                          enc_layers, k, max_patches, dropout)
+        self.decoder = FingerprintDecoder(c_in, patch_size, d_model, n_heads,
+                                          dec_layers, max_patches, dropout)
+        # Sec 3.4: attention pooling over the k fingerprint tokens
+        self.q_task = nn.Linear(d_model, 1, bias=False)
+        self.cls_head = nn.Sequential(nn.LayerNorm(d_model), nn.Linear(d_model, n_classes))
+
+    # ---- pre-training ----
+    def patchify(self, x):
+        b, t, c = x.shape
+        n = t // self.patch_size
+        return x[:, : n * self.patch_size].reshape(b, n, self.patch_size * c)
+
+    def pretrain_step(self, x, mask_ratio=0.6, lam=1e-4, use_div=True, generator=None):
+        target = self.patchify(x)
+        b, n, _ = target.shape
+        # masked view: keep (1-r) of the patches visible to the encoder
+        n_keep = max(1, int(round(n * (1 - mask_ratio))))
+        noise = torch.rand(b, n, device=x.device, generator=generator)
+        keep = noise.argsort(dim=1)[:, :n_keep]
+        kv_mask = torch.ones(b, n, dtype=torch.bool, device=x.device)
+        kv_mask.scatter_(1, keep, False)  # True == ignore
+        f = self.encoder(x, kv_mask=kv_mask)
+        pred = self.decoder(f, n)
+        l_rec = F.mse_loss(pred, target)
+        l_div = total_coding_rate_loss(f) if use_div else torch.zeros((), device=x.device)
+        return l_rec + lam * l_div, l_rec.detach(), l_div.detach()
+
+    # ---- downstream ----
+    def forward(self, x, return_alpha=False):
+        f = self.encoder(x)
+        alpha = torch.softmax(self.q_task(f).squeeze(-1), dim=-1)  # (B, k)
+        z = (alpha.unsqueeze(-1) * f).sum(1)
+        logits = self.cls_head(z)
+        if return_alpha:
+            return logits, alpha, f
+        return logits
+
+
+# ------------------------------------------------------------------
+# Baselines for Claim 4
+# ------------------------------------------------------------------
+class PlainEncoder(nn.Module):
+    """Standard transformer patch encoder -> variable-length token sequence
+    (the 'entangled view' shared by Ti-MAE and SimMTM)."""
+
+    def __init__(self, c_in, patch_size, d_model=128, n_heads=8, n_layers=6,
+                 max_patches=512, dropout=0.1):
+        super().__init__()
+        self.patch_embed = PatchEmbed(c_in, patch_size, d_model)
+        self.pos = nn.Parameter(torch.zeros(1, max_patches, d_model))
+        nn.init.trunc_normal_(self.pos, std=0.02)
+        layer = nn.TransformerEncoderLayer(d_model, n_heads, 4 * d_model, dropout,
+                                           activation="gelu", batch_first=True,
+                                           norm_first=True)
+        self.enc = nn.TransformerEncoder(layer, n_layers)
+        self.norm = nn.LayerNorm(d_model)
+
+    def forward(self, x, src_key_padding_mask=None):
+        p = self.patch_embed(x)
+        p = p + self.pos[:, : p.shape[1]]
+        return self.norm(self.enc(p, src_key_padding_mask=src_key_padding_mask))
+
+
+class TiMAE(nn.Module):
+    """Ti-MAE (Li et al., 2023): masked patch autoencoding on a plain transformer,
+    visible patches are fed to the encoder, mask tokens are appended in the decoder,
+    downstream head is global average pooling."""
+
+    def __init__(self, c_in, patch_size, n_classes, d_model=128, n_heads=8,
+                 enc_layers=6, dec_layers=2, max_patches=512, dropout=0.1):
+        super().__init__()
+        self.patch_size = patch_size
+        self.encoder = PlainEncoder(c_in, patch_size, d_model, n_heads, enc_layers,
+                                    max_patches, dropout)
+        self.mask_token = nn.Parameter(torch.zeros(1, 1, d_model))
+        self.dec_pos = nn.Parameter(torch.zeros(1, max_patches, d_model))
+        nn.init.trunc_normal_(self.dec_pos, std=0.02)
+        layer = nn.TransformerEncoderLayer(d_model, n_heads, 4 * d_model, dropout,
+                                           activation="gelu", batch_first=True,
+                                           norm_first=True)
+        self.dec = nn.TransformerEncoder(layer, dec_layers)
+        self.head_rec = nn.Linear(d_model, c_in * patch_size)
+        self.cls_head = nn.Sequential(nn.LayerNorm(d_model), nn.Linear(d_model, n_classes))
+
+    def patchify(self, x):
+        b, t, c = x.shape
+        n = t // self.patch_size
+        return x[:, : n * self.patch_size].reshape(b, n, self.patch_size * c)
+
+    def pretrain_step(self, x, mask_ratio=0.6, generator=None, **kw):
+        target = self.patchify(x)
+        b, n, _ = target.shape
+        n_keep = max(1, int(round(n * (1 - mask_ratio))))
+        noise = torch.rand(b, n, device=x.device, generator=generator)
+        keep = noise.argsort(dim=1)[:, :n_keep]
+        pad = torch.ones(b, n, dtype=torch.bool, device=x.device)
+        pad.scatter_(1, keep, False)
+        h = self.encoder(x, src_key_padding_mask=pad)
+        # replace masked positions with the mask token, then decode the full sequence
+        h = torch.where(pad.unsqueeze(-1), self.mask_token.expand(b, n, -1), h)
+        h = h + self.dec_pos[:, :n]
+        pred = self.head_rec(self.dec(h))
+        l_rec = F.mse_loss(pred, target)
+        return l_rec, l_rec.detach(), torch.zeros((), device=x.device)
+
+    def forward(self, x):
+        h = self.encoder(x)
+        return self.cls_head(h.mean(1))
+
+
+class SimMTM(nn.Module):
+    """SimMTM (Dong et al., 2023): reconstruct the original series from *multiple*
+    masked views, aggregating them by point-wise series-similarity, plus a
+    series-wise contrastive term. Head is global average pooling."""
+
+    def __init__(self, c_in, patch_size, n_classes, d_model=128, n_heads=8,
+                 enc_layers=6, dec_layers=2, max_patches=512, dropout=0.1,
+                 n_views=3, temperature=0.2):
+        super().__init__()
+        self.patch_size = patch_size
+        self.n_views = n_views
+        self.temperature = temperature
+        self.encoder = PlainEncoder(c_in, patch_size, d_model, n_heads, enc_layers,
+                                    max_patches, dropout)
+        layer = nn.TransformerEncoderLayer(d_model, n_heads, 4 * d_model, dropout,
+                                           activation="gelu", batch_first=True,
+                                           norm_first=True)
+        self.dec = nn.TransformerEncoder(layer, dec_layers)
+        self.head_rec = nn.Linear(d_model, c_in * patch_size)
+        self.proj = nn.Sequential(nn.Linear(d_model, d_model), nn.GELU(),
+                                  nn.Linear(d_model, d_model))
+        self.cls_head = nn.Sequential(nn.LayerNorm(d_model), nn.Linear(d_model, n_classes))
+
+    def patchify(self, x):
+        b, t, c = x.shape
+        n = t // self.patch_size
+        return x[:, : n * self.patch_size].reshape(b, n, self.patch_size * c)
+
+    def pretrain_step(self, x, mask_ratio=0.6, generator=None, **kw):
+        target = self.patchify(x)
+        b, n, _ = target.shape
+        views = []
+        for _ in range(self.n_views):
+            m = (torch.rand(b, x.shape[1], 1, device=x.device, generator=generator)
+                 < mask_ratio)
+            views.append(x.masked_fill(m, 0.0))
+        xs = torch.cat(views, 0)
+        h = self.encoder(xs)                       # (V*B, N, d)
+        s = self.proj(h.mean(1))                   # series-level embedding
+        s = F.normalize(s, dim=-1)
+        # point-wise aggregation weighted by series similarity to view 0
+        sim = (s.view(self.n_views, b, -1) * s.view(self.n_views, b, -1)[0:1]).sum(-1)
+        w = torch.softmax(sim / self.temperature, dim=0)          # (V, B)
+        agg = (h.view(self.n_views, b, n, -1) * w[..., None, None]).sum(0)
+        pred = self.head_rec(self.dec(agg))
+        l_rec = F.mse_loss(pred, target)
+        # series-wise contrastive: views of the same series are positives
+        logits = s @ s.t() / self.temperature
+        lbl = torch.arange(b, device=x.device).repeat(self.n_views)
+        eye = torch.eye(logits.shape[0], device=x.device, dtype=torch.bool)
+        logits = logits.masked_fill(eye, -1e4)
+        pos = (lbl[:, None] == lbl[None, :]) & ~eye
+        log_p = torch.log_softmax(logits, dim=-1)
+        l_con = -(log_p * pos).sum(-1).div(pos.sum(-1).clamp(min=1)).mean()
+        return l_rec + 0.1 * l_con, l_rec.detach(), l_con.detach()
+
+    def forward(self, x):
+        h = self.encoder(x)
+        return self.cls_head(h.mean(1))
+
+
+MODELS = {"tsfp": TSFingerprint, "timae": TiMAE, "simmtm": SimMTM}
+
+
+def build(name, **kw):
+    if name == "simmtm":
+        kw.pop("k", None)
+    elif name == "timae":
+        kw.pop("k", None)
+    return MODELS[name](**kw)

prep_adftd.py

@@ -0,0 +1,101 @@
+"""ADFTD (OpenNeuro ds004504) preprocessing, following DL4mHealth/Medformer.
+
+Steps (verbatim from Medformer's ADFTD_preprocessing.ipynb):
+  * read the ASR/ICA-cleaned derivatives (.set, 19 channels, 500 Hz)
+  * linear-interpolation resample 500 -> 256 Hz
+  * segment into non-overlapping 256-timestamp windows, centred on the recording
+  * label map {C: 0 healthy, F: 1 FTD, A: 2 AD}
+Split (from Medformer's ADFTDLoader): subject-wise 60/20/20 within each class.
+Normalisation: per-sample, per-channel standardisation (uea.normalize_batch_ts).
+"""
+
+import argparse
+import os
+
+import mne
+import numpy as np
+import pandas as pd
+from scipy import interpolate
+
+mne.set_log_level("ERROR")
+
+LABEL_MAP = {"A": 2, "F": 1, "C": 0}
+
+
+def resample(arr, freq=500, target=256):
+    t = np.linspace(1, len(arr), len(arr))
+    f = interpolate.interp1d(t, arr, kind="linear")
+    t_new = np.linspace(1, len(arr), int(len(arr) / freq * target))
+    return f(t_new)
+
+
+def segment(mat, window=256):
+    """mat: (T, C). Centre-anchored non-overlapping windows."""
+    res = []
+    start = mat.shape[0] // 2
+    left = start - (start // window) * window
+    right = start + ((mat.shape[0] - start) // window) * window
+    for i in range(left, right, window):
+        res.append(mat[i:i + window])
+    return res
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--root", default="data/ADFTD")
+    ap.add_argument("--out", default="data/ADFTD/processed")
+    ap.add_argument("--window", type=int, default=256)
+    args = ap.parse_args()
+
+    parts = pd.read_csv(os.path.join(args.root, "participants.tsv"), sep="\t")
+    labels = {}
+    for row in parts.values:
+        pid = int(str(row[0])[-3:])
+        labels[pid] = LABEL_MAP[row[3]]
+
+    deriv = os.path.join(args.root, "derivatives")
+    feats, ys, subs = [], [], []
+    for pid in sorted(labels):
+        path = os.path.join(deriv, f"sub-{pid:03d}", "eeg",
+                            f"sub-{pid:03d}_task-eyesclosed_eeg.set")
+        raw = mne.io.read_raw_eeglab(path, preload=True)
+        sig = raw.get_data()  # (C, T) at 500 Hz
+        res = np.stack([resample(sig[c]) for c in range(sig.shape[0])], axis=1)  # (T', C)
+        for win in segment(res, args.window):
+            feats.append(win.astype(np.float32))
+            ys.append(labels[pid])
+            subs.append(pid)
+        print(f"sub-{pid:03d} label={labels[pid]} raw={sig.shape} "
+              f"resampled={res.shape} windows={len(segment(res, args.window))}", flush=True)
+
+    X = np.stack(feats)                      # (N, 256, 19)
+    y = np.array(ys, dtype=np.int64)
+    s = np.array(subs, dtype=np.int64)
+    print("total", X.shape, np.bincount(y))
+
+    # subject-wise 60/20/20 within each class (Medformer ADFTDLoader, a=0.6, b=0.8)
+    train_ids, val_ids, test_ids = [], [], []
+    for cls in (0, 1, 2):
+        ids = sorted([p for p in labels if labels[p] == cls])
+        a, b = int(0.6 * len(ids)), int(0.8 * len(ids))
+        train_ids += ids[:a]
+        val_ids += ids[a:b]
+        test_ids += ids[b:]
+    print("subjects train/val/test:", len(train_ids), len(val_ids), len(test_ids))
+
+    os.makedirs(args.out, exist_ok=True)
+    for name, ids in (("train", train_ids), ("val", val_ids), ("test", test_ids)):
+        m = np.isin(s, ids)
+        Xi = X[m]
+        mu = Xi.mean(axis=1, keepdims=True)
+        sd = Xi.std(axis=1, keepdims=True)
+        sd[sd == 0] = 1.0
+        Xi = ((Xi - mu) / sd).astype(np.float32)
+        np.save(os.path.join(args.out, f"X_{name}.npy"), Xi)
+        np.save(os.path.join(args.out, f"y_{name}.npy"), y[m])
+        np.save(os.path.join(args.out, f"subj_{name}.npy"), s[m])
+        print(name, Xi.shape, np.bincount(y[m]))
+
+
+if __name__ == "__main__":
+    main()

prep_ptbxl.py

@@ -0,0 +1,98 @@
+"""PTB-XL preprocessing, following DL4mHealth/Medformer's PTB-XL_preprocessing.ipynb.
+
+  * records500 (12-lead, 500 Hz, 10 s) -> linear resample to 250 Hz -> 2500 samples
+  * per-record StandardScaler over time, then reshape into 10 non-overlapping
+    250-timestamp windows (the paper's Table 5: T=250, patch 10, 5 classes)
+  * label = the SCP code with the highest likelihood, mapped to the 5 diagnostic
+    superclasses NORM / MI / STTC / CD / HYP
+  * subject-wise (patient_id) split; Medformer's PTB-XL loader uses a 60/20/20
+    patient split, which we reproduce with a fixed seed.
+"""
+
+import argparse
+import ast
+import os
+
+import numpy as np
+import pandas as pd
+import wfdb
+from scipy import interpolate
+from sklearn.preprocessing import StandardScaler
+
+SUPER = {"NORM": 0, "MI": 1, "STTC": 2, "CD": 3, "HYP": 4}
+
+
+def build_scp_map(scp_csv):
+    """diagnostic_class column of scp_statements.csv maps each SCP code to a superclass."""
+    df = pd.read_csv(scp_csv, index_col=0)
+    df = df[df.diagnostic == 1]
+    return {code: row.diagnostic_class for code, row in df.iterrows()
+            if isinstance(row.diagnostic_class, str) and row.diagnostic_class in SUPER}
+
+
+def resample(arr, freq=500, target=250):
+    t = np.linspace(1, len(arr), len(arr))
+    f = interpolate.interp1d(t, arr, kind="linear")
+    return f(np.linspace(1, len(arr), int(len(arr) / freq * target)))
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--root", default="data/PTBXL")
+    ap.add_argument("--out", default="data/PTBXL/processed")
+    ap.add_argument("--subject-frac", type=float, default=1.0,
+                    help="fraction of patients to keep (compute budget)")
+    ap.add_argument("--seed", type=int, default=0)
+    a = ap.parse_args()
+
+    info = pd.read_csv(os.path.join(a.root, "ptbxl_database.csv"))
+    code2super = build_scp_map(os.path.join(a.root, "scp_statements.csv"))
+
+    rows = []
+    for r in info.itertuples():
+        codes = ast.literal_eval(r.scp_codes)
+        diag = {c: v for c, v in codes.items() if c in code2super}
+        if not diag:
+            continue                       # Medformer keeps only diagnosable records
+        best = max(diag, key=diag.get)     # highest-likelihood SCP code
+        rows.append((r.ecg_id, r.patient_id, SUPER[code2super[best]], r.filename_hr))
+    df = pd.DataFrame(rows, columns=["ecg_id", "patient_id", "label", "path"])
+    print(f"records with a diagnostic superclass: {len(df)} "
+          f"({df.patient_id.nunique()} patients)")
+    print("class counts:", df.label.value_counts().sort_index().to_dict())
+
+    pats = np.array(sorted(df.patient_id.unique()))
+    rng = np.random.default_rng(a.seed)
+    rng.shuffle(pats)
+    if a.subject_frac < 1.0:
+        pats = pats[: max(2, int(len(pats) * a.subject_frac))]
+        df = df[df.patient_id.isin(pats)]
+        print(f"subsampled to {len(pats)} patients / {len(df)} records "
+              f"(subject_frac={a.subject_frac})")
+    n = len(pats)
+    splits = {"train": set(pats[: int(0.6 * n)]),
+              "val": set(pats[int(0.6 * n):int(0.8 * n)]),
+              "test": set(pats[int(0.8 * n):])}
+
+    os.makedirs(a.out, exist_ok=True)
+    for name, ids in splits.items():
+        sub = df[df.patient_id.isin(ids)]
+        X, y = [], []
+        for i, r in enumerate(sub.itertuples()):
+            sig, _ = wfdb.rdsamp(os.path.join(a.root, r.path))   # (5000, 12) @500 Hz
+            res = np.stack([resample(sig[:, c]) for c in range(sig.shape[1])], 1)
+            res = StandardScaler().fit_transform(res)            # per-record z-score
+            w = res[: (res.shape[0] // 250) * 250].reshape(-1, 250, res.shape[1])
+            X.append(w.astype(np.float32))
+            y.append(np.full(w.shape[0], r.label, dtype=np.int64))
+            if i % 2000 == 0:
+                print(f"  {name} {i}/{len(sub)}", flush=True)
+        X = np.concatenate(X)
+        y = np.concatenate(y)
+        np.save(os.path.join(a.out, f"X_{name}.npy"), X)
+        np.save(os.path.join(a.out, f"y_{name}.npy"), y)
+        print(name, X.shape, np.bincount(y, minlength=5), flush=True)
+
+
+if __name__ == "__main__":
+    main()

run_adftd_ablation.sh

@@ -0,0 +1,27 @@
+#!/bin/bash
+# Claims 3 & 4: ADFTD ablation (Scratch / Pre-trained L_rec / Pre-trained L_rec+L_div)
+# and the Ti-MAE / SimMTM fine-tuning-promotion comparison of Table 4.
+set -u
+cd /home/avesha/ritesh/paper-repro
+PY=.venvr/bin/python
+DATA=data/ADFTD/processed
+OUT=results/adftd
+mkdir -p $OUT
+SEEDS="41 42 43"
+COMMON="--data-dir $DATA --batch-size 256 --pre-epochs 100 --ft-epochs 100 --patience 10 --seeds $SEEDS"
+
+run () { # name, extra args
+  local name=$1; shift
+  if [ -s "$OUT/$name.json" ]; then echo "skip $name"; return; fi
+  echo "=== $name ==="
+  $PY tsfp/train.py $COMMON "$@" --out "$OUT/$name.json" 2>&1 | tail -4
+}
+
+run tsfp_scratch      --model tsfp   --mode scratch
+run tsfp_rec          --model tsfp   --mode pretrain_ft --use-div 0
+run tsfp_rec_div      --model tsfp   --mode pretrain_ft --use-div 1
+run timae_scratch     --model timae  --mode scratch
+run timae_pretrain    --model timae  --mode pretrain_ft
+run simmtm_scratch    --model simmtm --mode scratch
+run simmtm_pretrain   --model simmtm --mode pretrain_ft
+echo "ALL DONE"

schedule.py

@@ -0,0 +1,67 @@
+"""Tiny GPU-slot scheduler: runs a list of training configs across the local T4s,
+N concurrent processes per GPU, skipping configs whose result file already exists."""
+
+import argparse
+import json
+import os
+import subprocess
+import sys
+import threading
+import time
+from queue import Queue
+
+PY = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
+                  ".venvr", "bin", "python")
+TRAIN = os.path.join(os.path.dirname(os.path.abspath(__file__)), "train.py")
+
+
+def worker(slot, gpu, q, logdir):
+    while True:
+        try:
+            name, args = q.get_nowait()
+        except Exception:
+            return
+        out = os.path.join(logdir, f"{name}.json")
+        if os.path.exists(out) and os.path.getsize(out) > 0:
+            print(f"[slot{slot}] skip {name}", flush=True)
+            q.task_done()
+            continue
+        env = dict(os.environ, CUDA_VISIBLE_DEVICES=str(gpu), OMP_NUM_THREADS="4")
+        cmd = [PY, TRAIN] + args + ["--out", out]
+        log = os.path.join(logdir, f"{name}.log")
+        t0 = time.time()
+        print(f"[slot{slot}/gpu{gpu}] START {name}", flush=True)
+        with open(log, "w") as fh:
+            rc = subprocess.call(cmd, stdout=fh, stderr=subprocess.STDOUT, env=env)
+        print(f"[slot{slot}/gpu{gpu}] DONE  {name} rc={rc} "
+              f"({time.time()-t0:.0f}s)", flush=True)
+        q.task_done()
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--jobs", required=True, help="json file: {name: [args...]}")
+    ap.add_argument("--logdir", required=True)
+    ap.add_argument("--gpus", default="0,1")
+    ap.add_argument("--per-gpu", type=int, default=2)
+    a = ap.parse_args()
+    os.makedirs(a.logdir, exist_ok=True)
+    jobs = json.load(open(a.jobs))
+    q = Queue()
+    for name, args in jobs.items():
+        q.put((name, args))
+    threads = []
+    gpus = [g for g in a.gpus.split(",") if g]
+    for gi, gpu in enumerate(gpus):
+        for s in range(a.per_gpu):
+            t = threading.Thread(target=worker, args=(f"{gi}.{s}", gpu, q, a.logdir))
+            t.start()
+            threads.append(t)
+            time.sleep(3)
+    for t in threads:
+        t.join()
+    print("ALL DONE", flush=True)
+
+
+if __name__ == "__main__":
+    main()

train.py

@@ -0,0 +1,232 @@
+"""Pre-train + fine-tune harness for the TS-Fingerprint reproduction.
+
+Reproduces the protocol of Sec 4.1 / 4.2:
+  * subject-wise 60/20/20 split (Medformer protocol)
+  * Adam, lr 1e-3 (pre-train) / 1e-4 (downstream), <=100 epochs,
+    early stopping on validation F1 with patience 10
+  * macro Accuracy / Precision / Recall / F1 / AUROC on the test split
+"""
+
+import argparse
+import json
+import os
+import sys
+import time
+
+import numpy as np
+import torch
+import torch.nn as nn
+from sklearn.metrics import (accuracy_score, f1_score, precision_score,
+                             recall_score, roc_auc_score)
+from torch.utils.data import DataLoader, TensorDataset
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from model import build  # noqa: E402
+
+
+def set_seed(s):
+    np.random.seed(s)
+    torch.manual_seed(s)
+    torch.cuda.manual_seed_all(s)
+
+
+def load_split(data_dir, split, stride=1):
+    x = np.load(os.path.join(data_dir, f"X_{split}.npy"), mmap_mode="r")
+    y = np.load(os.path.join(data_dir, f"y_{split}.npy"))
+    if stride > 1:
+        # keep every `stride`-th window.  Windows are stored in recording order, so
+        # this thins each recording uniformly and leaves the subject-wise split,
+        # the subject count and the class balance untouched.
+        x, y = x[::stride], y[::stride]
+    x = np.ascontiguousarray(x)
+    return torch.from_numpy(x).float(), torch.from_numpy(y).long()
+
+
+def metrics(y_true, y_prob):
+    y_pred = y_prob.argmax(1)
+    n_cls = y_prob.shape[1]
+    out = {
+        "accuracy": accuracy_score(y_true, y_pred) * 100,
+        "precision": precision_score(y_true, y_pred, average="macro", zero_division=0) * 100,
+        "recall": recall_score(y_true, y_pred, average="macro", zero_division=0) * 100,
+        "f1": f1_score(y_true, y_pred, average="macro", zero_division=0) * 100,
+    }
+    try:
+        if n_cls == 2:
+            out["auroc"] = roc_auc_score(y_true, y_prob[:, 1]) * 100
+        else:
+            out["auroc"] = roc_auc_score(y_true, y_prob, multi_class="ovr",
+                                         average="macro") * 100
+    except ValueError:
+        out["auroc"] = float("nan")
+    return out
+
+
+AMP = {"enabled": False}
+
+
+def autocast(device):
+    return torch.autocast("cuda", dtype=torch.float16,
+                          enabled=AMP["enabled"] and device == "cuda")
+
+
+@torch.no_grad()
+def evaluate(model, loader, device):
+    model.eval()
+    probs, ys = [], []
+    for xb, yb in loader:
+        with autocast(device):
+            logits = model(xb.to(device))
+        probs.append(torch.softmax(logits.float(), -1).cpu().numpy())
+        ys.append(yb.numpy())
+    return metrics(np.concatenate(ys), np.concatenate(probs))
+
+
+def pretrain(model, loader, device, args, log=print):
+    opt = torch.optim.Adam(model.parameters(), lr=args.pre_lr)
+    scaler = torch.amp.GradScaler("cuda", enabled=AMP["enabled"] and device == "cuda")
+    hist = []
+    for ep in range(args.pre_epochs):
+        model.train()
+        tot = rec = div = 0.0
+        n = 0
+        for xb, _ in loader:
+            xb = xb.to(device, non_blocking=True)
+            with autocast(device):
+                loss, lr_, ld_ = model.pretrain_step(
+                    xb, mask_ratio=args.mask_ratio, lam=args.lam, use_div=args.use_div)
+            opt.zero_grad(set_to_none=True)
+            scaler.scale(loss).backward()
+            scaler.unscale_(opt)
+            torch.nn.utils.clip_grad_norm_(model.parameters(), 4.0)
+            scaler.step(opt)
+            scaler.update()
+            bs = xb.shape[0]
+            tot += loss.item() * bs
+            rec += lr_.item() * bs
+            div += ld_.item() * bs
+            n += bs
+        hist.append({"epoch": ep, "loss": tot / n, "l_rec": rec / n, "l_div": div / n})
+        log(f"[pretrain] ep {ep:3d} loss {tot/n:.5f} rec {rec/n:.5f} div {div/n:.4f}")
+    return hist
+
+
+def finetune(model, tr, va, te, device, args, log=print):
+    opt = torch.optim.Adam(model.parameters(), lr=args.ft_lr)
+    scaler = torch.amp.GradScaler("cuda", enabled=AMP["enabled"] and device == "cuda")
+    crit = nn.CrossEntropyLoss()
+    best_f1, best_state, patience, hist = -1.0, None, 0, []
+    for ep in range(args.ft_epochs):
+        model.train()
+        tot, n = 0.0, 0
+        for xb, yb in tr:
+            xb, yb = xb.to(device, non_blocking=True), yb.to(device, non_blocking=True)
+            with autocast(device):
+                loss = crit(model(xb), yb)
+            opt.zero_grad(set_to_none=True)
+            scaler.scale(loss).backward()
+            scaler.unscale_(opt)
+            torch.nn.utils.clip_grad_norm_(model.parameters(), 4.0)
+            scaler.step(opt)
+            scaler.update()
+            tot += loss.item() * xb.shape[0]
+            n += xb.shape[0]
+        vm = evaluate(model, va, device)
+        hist.append({"epoch": ep, "train_loss": tot / n, "val_f1": vm["f1"],
+                     "val_acc": vm["accuracy"]})
+        log(f"[finetune] ep {ep:3d} loss {tot/n:.4f} val_f1 {vm['f1']:.2f} "
+            f"val_acc {vm['accuracy']:.2f}")
+        if vm["f1"] > best_f1:
+            best_f1 = vm["f1"]
+            best_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()}
+            patience = 0
+        else:
+            patience += 1
+            if patience >= args.patience:
+                log(f"[finetune] early stop at epoch {ep}")
+                break
+    if best_state is not None:
+        model.load_state_dict(best_state)
+    return evaluate(model, te, device), hist, best_f1
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--data-dir", required=True)
+    p.add_argument("--model", default="tsfp", choices=["tsfp", "timae", "simmtm"])
+    p.add_argument("--mode", default="pretrain_ft",
+                   choices=["scratch", "pretrain_ft"])
+    p.add_argument("--use-div", type=int, default=1)
+    p.add_argument("--k", type=int, default=8)
+    p.add_argument("--d-model", type=int, default=128)
+    p.add_argument("--n-heads", type=int, default=8)
+    p.add_argument("--enc-layers", type=int, default=6)
+    p.add_argument("--dec-layers", type=int, default=2)
+    p.add_argument("--patch-size", type=int, default=8)
+    p.add_argument("--mask-ratio", type=float, default=0.6)
+    p.add_argument("--lam", type=float, default=1e-4)
+    p.add_argument("--pre-lr", type=float, default=1e-3)
+    p.add_argument("--ft-lr", type=float, default=1e-4)
+    p.add_argument("--pre-epochs", type=int, default=100)
+    p.add_argument("--ft-epochs", type=int, default=100)
+    p.add_argument("--patience", type=int, default=10)
+    p.add_argument("--batch-size", type=int, default=128)
+    p.add_argument("--seeds", type=int, nargs="+", default=[41, 42, 43, 44, 45])
+    p.add_argument("--out", default="results.json")
+    p.add_argument("--tag", default="")
+    p.add_argument("--amp", type=int, default=1)
+    p.add_argument("--stride", type=int, default=1,
+                   help="keep every Nth window (compute-budget subsampling)")
+    args = p.parse_args()
+    args.use_div = bool(args.use_div)
+    AMP["enabled"] = bool(args.amp)
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    xtr, ytr = load_split(args.data_dir, "train", args.stride)
+    xva, yva = load_split(args.data_dir, "val", args.stride)
+    xte, yte = load_split(args.data_dir, "test", args.stride)
+    n_cls = int(max(ytr.max(), yva.max(), yte.max())) + 1
+    c_in, t_len = xtr.shape[2], xtr.shape[1]
+    print(f"data {tuple(xtr.shape)} / {tuple(xva.shape)} / {tuple(xte.shape)} "
+          f"classes={n_cls} channels={c_in} T={t_len} device={device}", flush=True)
+
+    runs = []
+    for seed in args.seeds:
+        t0 = time.time()
+        set_seed(seed)
+        g = torch.Generator().manual_seed(seed)
+        tr = DataLoader(TensorDataset(xtr, ytr), batch_size=args.batch_size,
+                        shuffle=True, generator=g, drop_last=True, num_workers=2,
+                        pin_memory=True)
+        va = DataLoader(TensorDataset(xva, yva), batch_size=512, num_workers=2)
+        te = DataLoader(TensorDataset(xte, yte), batch_size=512, num_workers=2)
+        model = build(args.model, c_in=c_in, patch_size=args.patch_size,
+                      n_classes=n_cls, d_model=args.d_model, n_heads=args.n_heads,
+                      enc_layers=args.enc_layers, dec_layers=args.dec_layers,
+                      k=args.k, max_patches=max(512, t_len // args.patch_size + 8))
+        model.to(device)
+        n_par = sum(p.numel() for p in model.parameters())
+        pre_hist = []
+        if args.mode == "pretrain_ft":
+            pre_hist = pretrain(model, tr, device, args)
+        test_m, ft_hist, best_val = finetune(model, tr, va, te, device, args)
+        dt = time.time() - t0
+        print(f"== seed {seed} [{args.model}/{args.mode}/div={args.use_div}] "
+              f"{json.dumps({k: round(v, 2) for k, v in test_m.items()})} "
+              f"({dt:.0f}s, {n_par/1e6:.2f}M params)", flush=True)
+        runs.append({"seed": seed, "test": test_m, "best_val_f1": best_val,
+                     "seconds": dt, "params": n_par,
+                     "pretrain_hist": pre_hist, "finetune_hist": ft_hist})
+
+    agg = {m: {"mean": float(np.mean([r["test"][m] for r in runs])),
+               "std": float(np.std([r["test"][m] for r in runs]))}
+           for m in runs[0]["test"]}
+    out = {"config": vars(args), "runs": runs, "aggregate": agg}
+    with open(args.out, "w") as f:
+        json.dump(out, f, indent=2)
+    print("AGGREGATE " + json.dumps({k: f"{v['mean']:.2f}+-{v['std']:.2f}"
+                                     for k, v in agg.items()}), flush=True)
+
+
+if __name__ == "__main__":
+    main()