File size: 3,051 Bytes
1fdc49a | 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 | """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()
|