scripts/run_matrix.py

"""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",
            )