| |
| """Run reproducible CupOFMoCA PDC experiments. |
| |
| This wrapper keeps the submitted generation programs intact while providing one |
| portable interface, bundled checkpoint defaults, named targets, and collision- |
| safe outputs. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import shlex |
| import subprocess |
| import sys |
| from pathlib import Path |
|
|
|
|
| ROOT = Path(__file__).resolve().parents[1] |
| CONFIG = json.loads((ROOT / "configs" / "paper.json").read_text()) |
|
|
| CHECKPOINTS = { |
| "smiles": ROOT / "checkpoints" / "SMILESReDi_base.ckpt", |
| "aa": ROOT / "checkpoints" / "PepReDi_base.pt", |
| "uaa": ROOT / "checkpoints" / "UAAReDi_base.pt", |
| } |
|
|
|
|
| def read_fasta(path: Path) -> dict[str, str]: |
| records: dict[str, str] = {} |
| name = None |
| chunks: list[str] = [] |
| for raw in path.read_text().splitlines(): |
| line = raw.strip() |
| if not line: |
| continue |
| if line.startswith(">"): |
| if name is not None: |
| records[name] = "".join(chunks) |
| name, chunks = line[1:].split()[0], [] |
| elif name is None: |
| raise ValueError(f"Sequence before FASTA header in {path}") |
| else: |
| chunks.append(line) |
| if name is not None: |
| records[name] = "".join(chunks) |
| return records |
|
|
|
|
| def parse_args() -> argparse.Namespace: |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--representation", choices=("smiles", "aa", "uaa"), required=True) |
| parser.add_argument("--variant", choices=("baseline", "cupofmoca"), default="cupofmoca") |
| parser.add_argument("--method", choices=("guided", "staged"), default="guided") |
| parser.add_argument("--target", choices=tuple(CONFIG["targets"]), required=True) |
| parser.add_argument( |
| "--seeds", type=int, nargs="+", default=CONFIG["default_seeds"], |
| help="Random seeds for generation.", |
| ) |
| parser.add_argument("--num-candidates", type=int, default=CONFIG["num_candidates_per_seed"]) |
| parser.add_argument("--gen-len", type=int, default=None) |
| parser.add_argument("--checkpoint", type=Path, default=None) |
| parser.add_argument("--output-root", type=Path, default=ROOT / "outputs") |
| parser.add_argument("--python", default=sys.executable) |
| parser.add_argument("--dry-run", action="store_true") |
| parser.add_argument("--force", action="store_true", help="Allow an existing output to be replaced.") |
| parser.add_argument( |
| "extra_args", nargs=argparse.REMAINDER, |
| help="Arguments after '--' are appended to the underlying generator command.", |
| ) |
| args = parser.parse_args() |
| if args.num_candidates < 1: |
| parser.error("--num-candidates must be positive") |
| return args |
|
|
|
|
| def build_command(args: argparse.Namespace, seed: int, target_seq: str, output: Path) -> list[str]: |
| suffix = "" if args.representation == "smiles" else f"_{args.representation}" |
| prefix = "cupofmoca_" if args.variant == "cupofmoca" else "" |
| program = f"smiles.generation.{prefix}{args.method}{suffix}" |
| checkpoint = (args.checkpoint or CHECKPOINTS[args.representation]).resolve() |
| target_cfg = CONFIG["targets"][args.target] |
| gen_len = args.gen_len or ( |
| target_cfg["smiles_length"] |
| if args.representation == "smiles" |
| else CONFIG["representation_lengths"][args.representation] |
| ) |
|
|
| cmd = [ |
| args.python, "-u", "-m", program, |
| "--checkpoint", str(checkpoint), |
| "--target", target_seq, |
| "--num_samples", str(CONFIG["num_samples_per_batch"]), |
| "--num_batches", str(args.num_candidates), |
| "--gen_len", str(gen_len), |
| "--optimization_steps", str(CONFIG["optimization_steps"]), |
| "--weights", *map(str, CONFIG["weights"]), |
| "--top_p", str(CONFIG["top_p"]), |
| "--seed", str(seed), |
| "--output_file", str(output), |
| ] |
|
|
| if args.representation in {"aa", "uaa"}: |
| cmd += ["--tokenizer_type", "amino"] |
| if args.representation == "aa": |
| |
| |
| cmd += ["--affinity_type", "smiles"] |
|
|
| if args.method == "guided": |
| cmd += [ |
| "--builder_type", target_cfg["builder"], |
| "--require_valid_x0", |
| "--max_x0_retries", str(CONFIG["max_x0_retries"]), |
| "--output_dir", str(output.parent / "logs"), |
| ] |
|
|
| if args.variant == "cupofmoca": |
| cmd += [ |
| "--use_acceptance", |
| "--beta_min", str(CONFIG["beta_min"]), |
| "--beta_max", str(CONFIG["beta_max"]), |
| ] |
| if args.method == "guided": |
| cmd += ["--guidance_method", "exp_tilt"] |
| if args.representation in {"aa", "uaa"}: |
| cmd += ["--cascade_k", str(CONFIG["cascade_k"])] |
| else: |
| cmd += ["--output_dir", str(output.parent / "logs")] |
|
|
| extra = list(args.extra_args) |
| if extra and extra[0] == "--": |
| extra = extra[1:] |
| return cmd + extra |
|
|
|
|
| def main() -> int: |
| args = parse_args() |
| targets = read_fasta(ROOT / "configs" / "targets.fasta") |
| target_seq = targets[args.target] |
| failures = 0 |
|
|
| for seed in args.seeds: |
| run_dir = ( |
| args.output_root.resolve() |
| / args.target |
| / args.representation |
| / f"{args.variant}_{args.method}" |
| / f"seed_{seed}" |
| ) |
| output = run_dir / "molecules.txt" |
| if output.exists() and not args.force and not args.dry_run: |
| print(f"Refusing to overwrite {output}; pass --force to replace it.", file=sys.stderr) |
| failures += 1 |
| continue |
| if output.exists() and args.force: |
| output.unlink() |
| command = build_command(args, seed, target_seq, output) |
| display = " ".join( |
| f"<{args.target} sequence, {len(target_seq)} aa>" if item == target_seq else shlex.quote(item) |
| for item in command |
| ) |
| print(display) |
| if args.dry_run: |
| continue |
| run_dir.mkdir(parents=True, exist_ok=True) |
| completed = subprocess.run(command, cwd=ROOT, check=False) |
| if completed.returncode: |
| failures += 1 |
| print(f"Seed {seed} failed with exit code {completed.returncode}.", file=sys.stderr) |
|
|
| return 1 if failures else 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
|
|
