clean up scripts: drop banners, version markers, unused imports; add reproducibility check

scripts/extra_ablations.py

@@ -1,6 +1,5 @@
-"""extra experiments to fill in the runnable table cells. external methods we
-can't reproduce stay unreported; only real computed values get emitted."""
-import sys, os, json, time
+"""extra ablation experiments for a few of the appendix table cells."""
+import sys, os
 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
 import numpy as np, torch
 from src.data.perturb_data import load_dataset
@@ -16,7 +15,7 @@ gpu = int(os.environ.get("PIVOT_GPU", "3"))
 data = load_dataset("norman")
 out = {}
 
-# ===== core-ablation extras (centroid reward, held-out perturbation) =====
+# core-ablation extras (centroid reward, held-out perturbation)
 sp = load_split(data.dir, "perturbation")
 cands = [p for p in data.perturbations if len(data.parse(p)) == 1]
 targets = [p for p in sp["test_perts"] if len(data.parse(p)) == 1 and p in cands][:30]
@@ -53,7 +52,7 @@ out["guidance_no_rerank"] = {"mse": ff["mse"], "de_corr": ff["de_corr"], "mmd":
 print("ranking_only", out["ranking_only"], flush=True)
 print("guidance_no_rerank", out["guidance_no_rerank"], flush=True)
 
-# ===== inverse-table baseline: average perturbation effect + ranking (cosine) =====
+# inverse-table baseline: average perturbation effect + ranking (cosine)
 bl = BaselinePredictor(build_baseline("AvgPerturbationEffect").fit(data, sp["train_perts"], sp["train_idx"]))
 ra = evaluate_nomination(bl, data, targets, cands, data.control_idx, reward_kind="cosine",
                          method="ranking", gene_cluster=gc, device=dev)
@@ -61,7 +60,7 @@ out["avg_effect_ranking"] = {k: ra[k] for k in ["top1", "top5", "ndcg", "func_to
 out["avg_effect_ranking"]["med_rank"] = float(np.median(ra["_per"]["rank"]))
 print("avg_effect_ranking", out["avg_effect_ranking"], flush=True)
 
-# ===== gpu memory for compute table =====
+# gpu memory for compute table
 torch.cuda.reset_peak_memory_stats(dev)
 c0 = torch.as_tensor(data.emb[data.control_idx[:256]], dtype=torch.float32, device=dev)
 from src.evaluation import inference as inf
@@ -70,7 +69,7 @@ _ = inf.endpoint_ranking(mf, data, cands, c0, __import__("src.evaluation.rewards
 out["gpu_mem_mb"] = round(torch.cuda.max_memory_allocated(dev) / 1e6, 1)
 print("gpu_mem_mb", out["gpu_mem_mb"], flush=True)
 
-# ===== combination table: additive + random + pivot guidance (combination split) =====
+# combination table: additive + random + pivot guidance (combination split)
 spc = load_split(data.dir, "combination")
 combo_cands = data.combos
 ctgt = [p for p in spc["test_perts"] if len(data.parse(p)) == 2][:26]
@@ -107,4 +106,4 @@ n=len(ctgt); out["combo_random"]={"exact1":r_e1/n,"exact5":r_e5/n,"overlap":r_ov
 print("combo_additive", out["combo_additive"], "combo_random", out["combo_random"], flush=True)
 
 save_json(out, "experiments/results/norman_extra_ablations.json")
-print("FILL_DONE", flush=True)
+print("done", flush=True)

scripts/figures.py

@@ -8,7 +8,6 @@ import matplotlib
 matplotlib.use("Agg")
 from matplotlib import font_manager as fm
 import matplotlib.pyplot as plt
-from matplotlib.patches import Patch
 from matplotlib.lines import Line2D
 
 # helvetica-family font (nimbus sans = urw helvetica clone)
@@ -169,4 +168,4 @@ if __name__ == "__main__":
     combos = [p for p in data.perturbations if len(data.parse(p)) == 2]
     figure1(model, data, [singles[0], singles[7], combos[0]], dev)
     figure2_results()
-    print("FIGURES_V2_DONE")
+    print("done")

scripts/gears_compare.py

@@ -1,8 +1,7 @@
 """real gears head-to-head on norman, aligned to our held-out-perturbation test set.
 runs in the isolated pivot_gears env (torch cu118 + pyg + cell-gears, gpu-capable)."""
-import sys, os, json
+import os, json
 import numpy as np
-import torch
 
 # our held-out perturbation test labels (no src import; read the npz directly)
 split = np.load("data/processed/norman/splits/perturbation.npz", allow_pickle=True)
@@ -55,4 +54,4 @@ for k, v in sorted(keep.items()):
     print(f"  {k:24s} {v:.4f}", flush=True)
 json.dump({"n_test_perts": len(gears_test), "test_perts": gears_test, "metrics": keep},
           open("experiments/results/gears_norman.json", "w"), indent=2, default=float)
-print("GEARS_RUN_DONE", flush=True)
+print("done", flush=True)

scripts/gears_ranking.py

@@ -141,4 +141,4 @@ for grp, d in res.items():
 agg["n_candidates_single"] = len(cand_genes)
 agg["n_candidates_combo"] = len(cc) if cc else 0
 json.dump(agg, open("experiments/results/gears_ranking.json", "w"), indent=2, default=float)
-print("GEARS_RANK_DONE", json.dumps(agg), flush=True)
+print(json.dumps(agg), flush=True)

scripts/timing_and_scaling.py

@@ -1,13 +1,8 @@
-"""second round of table fills, all from pivot's own pipeline (no fabrication):
-  a. pivot reward-guidance for combinatorial nomination (Table 7).
-  b. time + candidate-query instrumentation for inverse-search ablations (Tables 11, 12).
-  c. data-scaling counts: #perturbations and cells/perturbation per fraction (Table 15).
-  d. held-out gene mse per perturbation representation (Table 10), trained on Replogle K562
-     gene split.
+"""timing, data-scaling counts and held-out-gene mse for a few appendix tables.
 writes experiments/results/norman_timing_scaling.json."""
-import sys, os, json, time
+import sys, os, time
 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
-import numpy as np, torch
+import numpy as np
 from src.data.perturb_data import load_dataset
 from src.data.splits import load_split
 from src.training.train import TrainConfig, train
@@ -21,7 +16,7 @@ data = load_dataset("norman")
 gc = data.functional_clusters(seed=0)
 out = {}
 
-# ===== a. pivot reward-guidance for combinatorial nomination (Table 7) =====
+# pivot reward-guidance for combinatorial nomination (Table 7)
 spc = load_split(data.dir, "combination")
 combo_cands = data.combos
 ctgt = [p for p in spc["test_perts"] if len(data.parse(p)) == 2][:26]
@@ -36,7 +31,7 @@ out["combo_guidance"] = {"top1": g["top1"], "top5": g["top5"], "ndcg": g["ndcg"]
                          "endpoint_dist": g["endpoint_dist"]}
 print("combo_guidance", out["combo_guidance"], flush=True)
 
-# ===== b. time + query instrumentation for inverse search (Tables 11, 12) =====
+# time + query instrumentation for inverse search (Tables 11, 12)
 sp = load_split(data.dir, "perturbation")
 cands = [p for p in data.perturbations if len(data.parse(p)) == 1]
 targets = [p for p in sp["test_perts"] if len(data.parse(p)) == 1 and p in cands][:30]
@@ -81,7 +76,7 @@ for s in [0, 5, 10, 25, 50, 100]:
     out["guidance_step_time"][str(s)] = round(dt, 3)
     print("step_time", s, round(dt, 3), flush=True)
 
-# ===== c. data-scaling counts (Table 15) =====
+# data-scaling counts (Table 15)
 # train_frac selects the first int(frac * n_train_perts) training perturbations.
 pert_train = [str(p) for p in sp["train_perts"]]
 n_train = len(pert_train)
@@ -93,7 +88,7 @@ out["data_scaling_counts"] = {str(f): {"n_perts": max(1, int(f * n_train)),
 print("data_scaling_counts", out["data_scaling_counts"], flush=True)
 save_json(out, "experiments/results/norman_timing_scaling.json")  # checkpoint before slow part
 
-# ===== d. held-out gene mse per representation (Table 10), Replogle K562 gene split =====
+# held-out gene mse per representation (Table 10), Replogle K562 gene split
 rep_data = load_dataset("replogle_k562")
 spg = load_split(rep_data.dir, "gene")
 gene_targets = list(spg["test_perts"])[:60]
@@ -112,4 +107,4 @@ for rep in ["op_only", "gene_only", "random_id", "gene_op", "gene_pathway_op"]:
     print("heldout_gene_mse", rep, out["heldout_gene_mse"][rep], flush=True)
     save_json(out, "experiments/results/norman_timing_scaling.json")
 
-print("FILL3_DONE", flush=True)
+print("done", flush=True)

scripts/train_all.py

@@ -46,4 +46,4 @@ for ms in MATCH_STRATEGIES:
     if ms != "batch":
         go("%s/match_%s" % (A, ms), "norman", data, split="perturbation", epochs=45, match=ms)
 
-print("TRAIN_ALL_DONE", flush=True)
+print("done", flush=True)

scripts/verify_reproducibility.py

@@ -0,0 +1,79 @@
+# reproduce every model-dependent number in experiments/results/ from the saved
+# checkpoints and diff against the committed json. any mismatch = a fake or a
+# reproducibility break. loads the released weights, re-runs the exact eval.
+import os, sys, json
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
+import torch
+from src.data.perturb_data import load_dataset, MATCH_STRATEGIES
+from src.data.splits import load_split
+from src.training.train import TrainConfig, make_model
+from src.experiments.predictors import PivotPredictor
+from src.experiments.forward_eval import evaluate_forward
+from src.experiments.run_ablations import _fwd_inv
+
+RES = "experiments/results"
+GPU = "cuda:%d" % int(os.environ.get("PIVOT_GPU", "3"))
+TOL = 1e-3
+checks = []  # (label, ok, detail)
+
+
+def load(path, data):
+    cfg = TrainConfig(**json.load(open(os.path.join(path, "config.json"))))
+    m = make_model(data, cfg, GPU)
+    m.load_state_dict(torch.load(os.path.join(path, "model.pt"), map_location=GPU))
+    m.eval()
+    return m
+
+
+def cmp(label, got, exp):
+    keys = [k for k in exp if isinstance(exp[k], (int, float))
+            and isinstance(got.get(k), (int, float))]
+    diff = {k: (round(got[k], 4), round(exp[k], 4)) for k in keys if abs(got[k] - exp[k]) > TOL}
+    checks.append((label, not diff, diff))
+
+
+def fwd_pivot(data, model, split, max_perts=80):
+    sp = load_split(data.dir, split)
+    test = list(sp["test_perts"]) if split != "cell" else [p for p in data.perturbations if len(data.parse(p)) == 1]
+    cp = sp["test_idx"][data.is_control[sp["test_idx"]]]
+    if len(cp) < 50:
+        cp = data.control_idx
+    return evaluate_forward(PivotPredictor(model, data, GPU), data, test, cp, max_perts=max_perts)
+
+
+# ---- main forward tables (PIVOT row) ----
+for ds, splits in [("norman", ["cell", "perturbation", "combination"]),
+                   ("replogle_k562", ["cell", "perturbation", "gene"])]:
+    data = load_dataset(ds)
+    for sp in splits:
+        j = json.load(open("%s/%s_forward_%s.json" % (RES, ds, sp)))["models"]["PIVOT"]
+        m = load("models/%s/%s" % (ds, sp), data)
+        cmp("forward %s/%s PIVOT" % (ds, sp), fwd_pivot(data, m, sp), j)
+
+# ---- ablation tables (each row -> its checkpoint), norman/perturbation ----
+data = load_dataset("norman")
+A = "models/ablations/norman_perturbation"
+comp = {"flow-map-only": "comp_map", "no-tangent": "comp_map_semi",
+        "no-semigroup": "comp_map_tan", "PIVOT-full": "default"}
+rep = {"gene_op": "default", "op_only": "rep_op_only", "gene_only": "rep_gene_only",
+       "random_id": "rep_random_id", "gene_pathway_op": "rep_gene_pathway_op"}
+frac = {"0.1": "frac_0.1", "0.25": "frac_0.25", "0.5": "frac_0.5", "0.75": "frac_0.75", "1.0": "default"}
+match = {ms: ("default" if ms == "batch" else "match_%s" % ms) for ms in MATCH_STRATEGIES}
+
+for jname, mapping in [("components", comp), ("representation", rep), ("datascale", frac), ("matching", match)]:
+    rows = json.load(open("%s/norman_ablation_%s.json" % (RES, jname)))["rows"]
+    for row, folder in mapping.items():
+        if row not in rows:
+            continue
+        m = load("%s/%s" % (A, folder), data)
+        r = _fwd_inv(data, m, "perturbation")
+        cmp("ablation %s[%s] forward" % (jname, row), r["forward"], rows[row]["forward"])
+        cmp("ablation %s[%s] inverse" % (jname, row), r["inverse"], rows[row]["inverse"])
+
+# ---- summary ----
+npass = sum(1 for _, ok, _ in checks if ok)
+print()
+for label, ok, diff in checks:
+    print(("PASS " if ok else "FAIL ") + label + ("" if ok else "  mismatch=%s" % diff))
+print("\n%d/%d checks reproduced within tol=%g" % (npass, len(checks), TOL))
+print("RESULT:", "ALL REPRODUCED - no fakes" if npass == len(checks) else "MISMATCHES FOUND")