| |
| """ |
| EGFR-VHL ablation driver. Deduplicated layout: |
| |
| 3 baseline runs (weights = 1,1,1,1,1, seeds 42-44) |
| 5 scorers x 2 non-baseline weights {0,2} x 3 seeds = 30 runs |
| 3 acceptance=on runs (baseline weights, seeds 42-44) |
| total = 36 runs |
| |
| The baseline runs cover every (scorer, w=1) cell at once. The acceptance=off |
| case is the baseline (same runs reused), so only acceptance=on is run here. |
| |
| Each run = 100 samples on EGFR-VHL. Outputs go to |
| generation_logs/ablation_egfr/<tag>/. Resumable (skips done .smi files). |
| """ |
| import argparse |
| import itertools |
| import os |
| import subprocess |
| import time |
| from pathlib import Path |
|
|
| |
| |
| |
| SCRIPT_DIR = Path(__file__).resolve().parent.parent |
| REPO_ROOT = SCRIPT_DIR.parent |
| CHECKPOINT = REPO_ROOT / "checkpoints" / "protac_linker_finetune" / "epoch=99-v1.ckpt" |
| CLF_PATH = REPO_ROOT / "src" / "classifier" / "protac_linker_clf.pkl" |
| GENERATE = SCRIPT_DIR / "generate.py" |
| OUT_BASE = REPO_ROOT / "generation_logs" / "ablation_egfr" |
|
|
| PROTAC_MODULE = "ErlotinibVHL" |
| N_SAMPLES = 100 |
| N_STEPS = 1500 |
| WARMUP_STEPS = 100 |
| SCORER_NAMES = ["classifier", "fast", "sa", "chem", "conn"] |
| WEIGHTS = [0.0, 1.0, 2.0, 4.0] |
| NON_BASELINE = [w for w in WEIGHTS if w != 1.0] |
| SEEDS = [42, 43, 44] |
| BASELINE_W = [1.0, 1.0, 1.0, 1.0, 1.0] |
|
|
|
|
| def weight_vec(scorer_idx: int, w: float) -> list[float]: |
| """Baseline (1,1,1,1,1) with one slot replaced.""" |
| v = list(BASELINE_W) |
| v[scorer_idx] = w |
| return v |
|
|
|
|
| def build_runs() -> list[dict]: |
| """Return list of run configs. Each config has a unique tag.""" |
| runs = [] |
| |
| |
| for seed in SEEDS: |
| runs.append({ |
| "tag": f"egfr_baseline_seed{seed}", |
| "weights": list(BASELINE_W), |
| "seed": seed, |
| "accept": False, |
| }) |
| |
| for idx, name in enumerate(SCORER_NAMES): |
| for w in NON_BASELINE: |
| for seed in SEEDS: |
| runs.append({ |
| "tag": f"egfr_{name}_w{w:g}_seed{seed}", |
| "weights": weight_vec(idx, w), |
| "seed": seed, |
| "accept": False, |
| }) |
| |
| for seed in SEEDS: |
| runs.append({ |
| "tag": f"egfr_accept_on_seed{seed}", |
| "weights": list(BASELINE_W), |
| "seed": seed, |
| "accept": True, |
| }) |
| return runs |
|
|
|
|
| def run_command(run: dict, gpu: int) -> list[str]: |
| out_dir = OUT_BASE / run["tag"] |
| out_dir.mkdir(parents=True, exist_ok=True) |
| cmd = [ |
| "python", str(GENERATE), |
| "--checkpoint", str(CHECKPOINT), |
| "--protac_module", PROTAC_MODULE, |
| "--fix_termini", "cooh", |
| "--fast_scoring", |
| "--clf_path", str(CLF_PATH), |
| "--num_samples", str(N_SAMPLES), |
| "--num_batches", "1", |
| "--num_steps", str(N_STEPS), |
| "--warmup_steps", str(WARMUP_STEPS), |
| "--weights", *[str(w) for w in run["weights"]], |
| "--seed", str(run["seed"]), |
| "--output_file", str(out_dir / f"{run['tag']}.smi"), |
| "--output_dir", str(out_dir), |
| ] |
| if run["accept"]: |
| cmd.append("--use_acceptance") |
| return cmd |
|
|
|
|
| def is_done(run: dict) -> bool: |
| out_dir = OUT_BASE / run["tag"] |
| smi = out_dir / f"{run['tag']}.smi" |
| if not smi.exists(): |
| return False |
| |
| return smi.stat().st_size > 0 |
|
|
|
|
| def launch(cmd: list[str], gpu: int, log_path: Path): |
| env = os.environ.copy() |
| env["CUDA_VISIBLE_DEVICES"] = str(gpu) |
| with open(log_path, "w") as logf: |
| return subprocess.Popen(cmd, env=env, stdout=logf, stderr=subprocess.STDOUT) |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--gpus", default="2,3,4,5,6", |
| help="comma-separated GPU IDs to use (default: 2,3,4,5,6)") |
| ap.add_argument("--stagger", type=float, default=2.0, |
| help="seconds between successive run starts (avoid CSV collisions)") |
| ap.add_argument("--dry-run", action="store_true", |
| help="print commands without launching") |
| args = ap.parse_args() |
|
|
| gpu_list = [int(g) for g in args.gpus.split(",")] |
| runs = build_runs() |
| todo = [r for r in runs if not is_done(r)] |
| done = [r for r in runs if is_done(r)] |
|
|
| print(f"Total runs: {len(runs)}") |
| print(f"Already done: {len(done)}") |
| print(f"To launch: {len(todo)}") |
| print(f"GPUs available: {gpu_list}") |
|
|
| if args.dry_run: |
| for r in todo[:5]: |
| cmd = run_command(r, gpu_list[0]) |
| print("[DRY]", " ".join(cmd)) |
| if len(todo) > 5: |
| print(f"[DRY] ... and {len(todo)-5} more") |
| return |
|
|
| OUT_BASE.mkdir(parents=True, exist_ok=True) |
|
|
| active: dict[int, dict] = {} |
| queue = list(todo) |
|
|
| def dispatch(gpu: int): |
| if not queue: |
| return False |
| run = queue.pop(0) |
| cmd = run_command(run, gpu) |
| log_path = OUT_BASE / run["tag"] / "run.log" |
| proc = launch(cmd, gpu, log_path) |
| active[gpu] = {"proc": proc, "run": run, "log": log_path, |
| "start": time.time()} |
| print(f"[GPU {gpu}] start {run['tag']} (pid {proc.pid})") |
| time.sleep(args.stagger) |
| return True |
|
|
| |
| for gpu in gpu_list: |
| dispatch(gpu) |
|
|
| |
| while active or queue: |
| for gpu, info in list(active.items()): |
| ret = info["proc"].poll() |
| if ret is not None: |
| elapsed = time.time() - info["start"] |
| status = "OK" if ret == 0 else f"FAIL({ret})" |
| print(f"[GPU {gpu}] done {info['run']['tag']} " |
| f"({status}, {elapsed/60:.1f} min)") |
| del active[gpu] |
| if queue: |
| dispatch(gpu) |
| time.sleep(10) |
|
|
| print(f"\nAll runs finished. Outputs under {OUT_BASE}/") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
