| """Run IgFold structure prediction from antibody sequences or a FASTA file.""" |
|
|
| import argparse |
| import json |
| from pathlib import Path |
|
|
| import torch |
|
|
| from igfold import IgFoldRunner |
|
|
|
|
| DEFAULT_HEAVY = ( |
| "EVQLVQSGPEVKKPGTSVKVSCKASGFTFMSSAVQWVRQARGQRLEWIGWIVIGSGNTNYAQKF" |
| "QERVTITRDMSTSTAYMELSSLRSEDTAVYYCAAPYCSSISCNDGFDIWGQGTMVTVS" |
| ) |
| DEFAULT_LIGHT = ( |
| "DVVMTQTPFSLPVSLGDQASISCRSSQSLVHSNGNTYLHWYLQKPGQSPKLLIYKVSNRFSGV" |
| "PDRFSGSGSGTDFTLKISRVEAEDLGVYFCSQSTHVPYTFGGGTKLEIK" |
| ) |
|
|
|
|
| def parse_args(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--fasta", type=Path, help="Input FASTA containing H/L chains.") |
| parser.add_argument("--heavy", help="Heavy-chain amino-acid sequence.") |
| parser.add_argument("--light", help="Light-chain amino-acid sequence.") |
| parser.add_argument("--output", type=Path, default=Path("output/inference/antibody.pdb")) |
| parser.add_argument("--num-models", type=int, choices=range(1, 5), default=4) |
| parser.add_argument("--refine", action="store_true", help="Enable structure refinement.") |
| parser.add_argument("--openmm", action="store_true", help="Use OpenMM for refinement.") |
| parser.add_argument("--renum", action="store_true", help="Apply Chothia numbering.") |
| parser.add_argument("--cpu", action="store_true", help="Force CPU inference.") |
| return parser.parse_args() |
|
|
|
|
| def main(): |
| args = parse_args() |
| root = Path(__file__).resolve().parents[1] |
| checkpoints = sorted((root / "weight" / "IgFold").glob("*.ckpt"))[: args.num_models] |
| if len(checkpoints) != args.num_models: |
| raise FileNotFoundError( |
| f"Expected {args.num_models} checkpoints in {root / 'weight' / 'IgFold'}, " |
| f"found {len(checkpoints)}." |
| ) |
|
|
| if args.fasta is not None: |
| fasta_file = str(args.fasta.resolve()) |
| sequences = None |
| else: |
| fasta_file = None |
| sequences = {"H": args.heavy or DEFAULT_HEAVY} |
| light = args.light if args.light is not None else DEFAULT_LIGHT |
| if light: |
| sequences["L"] = light |
|
|
| if args.refine and not args.openmm: |
| from igfold.refine.pyrosetta_ref import init_pyrosetta |
|
|
| init_pyrosetta() |
|
|
| args.output.parent.mkdir(parents=True, exist_ok=True) |
| runner = IgFoldRunner( |
| model_ckpts=[str(path) for path in checkpoints], |
| try_gpu=not args.cpu, |
| ) |
| output = runner.fold( |
| str(args.output), |
| fasta_file=fasta_file, |
| sequences=sequences, |
| do_refine=args.refine, |
| use_openmm=args.openmm, |
| do_renum=args.renum, |
| ) |
|
|
| result = { |
| "status": "PASS", |
| "output_pdb": str(args.output.resolve()), |
| "models": len(runner.models), |
| "device": str(next(runner.models[0].parameters()).device), |
| "torch": torch.__version__, |
| "coords_shape": list(output.coords.shape), |
| "prmsd_shape": list(output.prmsd.shape), |
| } |
| print("IGFOLD_INFERENCE_RESULT=" + json.dumps(result, sort_keys=True)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
