scripts/inference.py

"""Run CombFold combinatorial assembly on AlphaFold-Multimer PDB files."""

import argparse
import json
from pathlib import Path


ROOT = Path(__file__).resolve().parents[1]
DEFAULT_CONFIG = ROOT / "conf" / "config.json"


def load_config(path: Path) -> dict:
    with path.open("r", encoding="utf-8") as handle:
        return json.load(handle)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--subunits", type=Path, required=True, help="CombFold subunits JSON file.")
    parser.add_argument("--pdbs", type=Path, required=True, help="Folder containing AFM PDB predictions.")
    parser.add_argument("--output", type=Path, required=True, help="Empty output folder.")
    parser.add_argument("--crosslinks", type=Path, help="Optional crosslinks file.")
    parser.add_argument("--config", type=Path, default=DEFAULT_CONFIG)
    parser.add_argument("--max-results", type=int, help="Maximum assembled structures to write.")
    parser.add_argument("--output-cif", action="store_true", help="Write CIF instead of PDB output.")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    from run_on_pdbs import run_on_pdbs_folder

    config_path = args.config.resolve()
    subunits_path = args.subunits.resolve()
    pdbs_path = args.pdbs.resolve()
    output_path = args.output.resolve()
    crosslinks_path = args.crosslinks.resolve() if args.crosslinks else None

    config = load_config(config_path)
    max_results = args.max_results if args.max_results is not None else int(config.get("max_results", 5))
    output_cif = args.output_cif or bool(config.get("output_cif", False))

    run_on_pdbs_folder(
        str(subunits_path),
        str(pdbs_path),
        str(output_path),
        crosslinks_path=str(crosslinks_path) if crosslinks_path else None,
        output_cif=output_cif,
        max_results_number=max_results,
    )

    assembled_dir = output_path / "assembled_results"
    structures = sorted(assembled_dir.glob("*.cif" if output_cif else "*.pdb"))
    result = {
        "status": "PASS" if structures else "FAILED",
        "output": str(output_path),
        "assembled_structures": len(structures),
        "format": "cif" if output_cif else "pdb",
    }
    print("COMBFOLD_INFERENCE_RESULT=" + json.dumps(result, sort_keys=True))
    if not structures:
        raise RuntimeError("CombFold did not produce any assembled structures.")


if __name__ == "__main__":
    main()