| """`casmd-predict` CLI: prep / ingest / select.""" |
| from __future__ import annotations |
| import argparse |
| import shutil |
| import sys |
| from pathlib import Path |
|
|
| from casmd.predict.prep import prep_inputs, SERVER_URLS |
| from casmd.predict.boltz2 import parse_boltz2 |
| from casmd.predict.protenix import parse_protenix |
| from casmd.predict.af3 import parse_af3 |
| from casmd.predict.selector import select_best, rank |
|
|
|
|
| _PARSERS = { |
| "boltz": parse_boltz2, |
| "boltz-2": parse_boltz2, |
| "protenix": parse_protenix, |
| "af3": parse_af3, |
| } |
|
|
|
|
| def main(argv: list[str] | None = None) -> int: |
| parser = argparse.ArgumentParser(prog="casmd-predict") |
| sub = parser.add_subparsers(dest="cmd", required=True) |
|
|
| p_prep = sub.add_parser("prep", help="Write submission-ready inputs for all 3 servers.") |
| p_prep.add_argument("name", help="Job name (used as filename / title)") |
| p_prep.add_argument("--protein", type=Path, required=True, help="Protein FASTA file") |
| p_prep.add_argument("--rna", help="RNA sequence (single line)") |
| p_prep.add_argument("--dna", help="DNA sequence (single line)") |
| p_prep.add_argument("--target", help="DNA target-complement sequence") |
| p_prep.add_argument("-o", "--output", type=Path, required=True, help="Output directory") |
|
|
| p_ing = sub.add_parser("ingest", help="Parse a downloaded prediction bundle.") |
| p_ing.add_argument("bundle_dir", type=Path) |
| p_ing.add_argument("--backend", required=True, choices=list(_PARSERS.keys())) |
|
|
| p_sel = sub.add_parser("select", help="Pick best model across one or more bundles.") |
| p_sel.add_argument("bundles", type=str, nargs="+", |
| help="Paths in the form '<backend>:<dir>', e.g. 'af3:./af3_run/'") |
| p_sel.add_argument("-o", "--output", type=Path, required=True, help="Where to copy the best PDB") |
| p_sel.add_argument("--by", default="iptm", choices=["iptm", "ptm", "plddt_mean"]) |
|
|
| args = parser.parse_args(argv) |
|
|
| if args.cmd == "prep": |
| result = prep_inputs( |
| protein_fasta=args.protein.read_text(), |
| rna_sequence=args.rna, |
| dna_sequence=args.dna, |
| target_complement=args.target, |
| output_dir=args.output, |
| job_name=args.name, |
| ) |
| print(f"Wrote:\n {result.boltz_input}\n {result.protenix_input}\n {result.af3_input}") |
| print("\nSubmit at:") |
| for k, url in SERVER_URLS.items(): |
| print(f" {k:9s} -> {url}") |
| return 0 |
|
|
| if args.cmd == "ingest": |
| bundle = _PARSERS[args.backend](args.bundle_dir) |
| print(f"Parsed {len(bundle.models)} models from {bundle.backend}") |
| for m in bundle.models: |
| iptm = f"{m.iptm:.3f}" if m.iptm is not None else " -- " |
| ptm = f"{m.ptm:.3f}" if m.ptm is not None else " -- " |
| print(f" model {m.model_id} iPTM={iptm} pTM={ptm} -> {m.pdb_path}") |
| return 0 |
|
|
| if args.cmd == "select": |
| bundles = [] |
| for spec in args.bundles: |
| parts = spec.split(":", 1) |
| if len(parts) != 2: |
| print(f"error: bundle spec must be '<backend>:<dir>', got {spec!r}", file=sys.stderr) |
| return 2 |
| backend, dirpath = parts |
| if backend not in _PARSERS: |
| print(f"error: unknown backend {backend!r}", file=sys.stderr) |
| return 2 |
| bundles.append(_PARSERS[backend](Path(dirpath))) |
| best = select_best(bundles, by=args.by) |
| args.output.parent.mkdir(parents=True, exist_ok=True) |
| shutil.copy(best.pdb_path, args.output) |
| score = getattr(best, args.by) |
| print(f"Best: {best.backend} model {best.model_id} {args.by}={score:.3f}") |
| print(f"Copied to: {args.output}") |
| return 0 |
|
|
| return 1 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
