File size: 3,643 Bytes
6fa9282 | 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 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 | """Execute corrected response, nomination, baseline, and loss-ablation comparisons."""
import argparse, subprocess, sys
from pathlib import Path
ROOT = Path(__file__).resolve().parents[1]
def run(*args):
subprocess.run([sys.executable, "-m", "pivot.cli", *map(str, args)], check=True)
if __name__ == "__main__":
p = argparse.ArgumentParser()
p.add_argument("--raw", required=True)
p.add_argument("--dataset", choices=["norman", "replogle_k562"], required=True)
p.add_argument(
"--split",
choices=["cell", "perturbation", "combination", "gene"],
default="combination",
)
p.add_argument("--output", required=True)
p.add_argument("--seeds", type=int, nargs="+", default=[0, 1, 2])
p.add_argument("--device", default="cuda")
p.add_argument("--ablations", action="store_true")
p.add_argument("--input-scale", choices=["counts", "log1p"], default="counts")
a = p.parse_args()
variants = ["full", "distribution_2", "distribution_10"]
if a.ablations:
variants += [
"map_only",
"map_tangent",
"map_semigroup",
"gene_only",
"random_pairing",
"nearest_pairing",
]
catalog = "combination" if a.split == "combination" else "single"
for seed in a.seeds:
out = Path(a.output) / f"seed_{seed}"
cache = out / "cache"
if not (cache / "meta.json").exists():
run(
"prepare",
"--raw",
a.raw,
"--dataset",
a.dataset,
"--split",
a.split,
"--seed",
seed,
"--output",
cache,
"--batch-col",
"batch" if a.dataset == "replogle_k562" else "gemgroup",
"--celltype-col",
"cell_line" if a.dataset == "replogle_k562" else "celltype",
"--input-scale",
a.input_scale,
)
for variant in variants:
model = out / variant
run(
"train",
"--cache",
cache,
"--config",
ROOT / "configs" / f"{variant}.json",
"--device",
a.device,
"--seed",
seed,
"--output",
model,
)
for initialization in ["random", "best"]:
run(
"evaluate",
"--cache",
cache,
"--checkpoint",
model / "best.pt",
"--catalog",
catalog,
"--initialization",
initialization,
"--device",
a.device,
"--seed",
seed,
"--output",
out / f"{variant}_{initialization}.json",
)
for baseline in [
"mean_control",
"average_effect",
"additive",
"ridge",
"endpoint_mlp",
"conditional_mlp",
]:
run(
"evaluate",
"--cache",
cache,
"--baseline",
baseline,
"--catalog",
catalog,
"--device",
a.device,
"--seed",
seed,
"--output",
out / f"{baseline}.json",
)
|