File size: 3,764 Bytes
fe43f4f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 | # 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")
|