File size: 6,480 Bytes
a4cdcc4 | 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 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 | #!/usr/bin/env python3
"""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":
# This is the scorer used by the PDC paper runs. The optional AA-input
# scorer is not needed and its multi-GB checkpoint is intentionally omitted.
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())
|