File size: 4,907 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 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 | import os
from time import time
from glob import glob
import torch
from antiberty import AntiBERTyRunner
import igfold
from igfold.model.IgFold import IgFold
from igfold.utils.folding import fold
from igfold.utils.embed import embed
from igfold.utils.general import exists
def display_license():
license_url = "https://github.com/Graylab/IgFold/blob/main/LICENSE.md"
license_message = f"""
The code, data, and weights for this work are made available for non-commercial use
(including at commercial entities) under the terms of the JHU Academic Software License
Agreement. For commercial inquiries, please contact awichma2[at]jhu.edu.
License: {license_url}
"""
print(license_message)
class IgFoldRunner():
"""
Wrapper for IgFold model predictions.
"""
def __init__(self, num_models=4, model_ckpts=None, try_gpu=True):
"""
Initialize IgFoldRunner.
:param num_models: Number of pre-trained IgFold models to use for prediction.
:param model_ckpts: List of model checkpoints to use (instead of pre-trained).
"""
display_license()
if exists(model_ckpts):
num_models = len(model_ckpts)
else:
if num_models < 1 or num_models > 4:
raise ValueError("num_models must be between 1 and 4.")
if not exists(model_ckpts):
project_path = os.path.dirname(
os.path.realpath(igfold.__file__))
ckpt_path = os.path.join(
project_path,
"trained_models/IgFold/*.ckpt",
)
model_ckpts = list(glob(ckpt_path))
model_ckpts = list(sorted(model_ckpts))[:num_models]
print(f"Loading {num_models} IgFold models...")
device = torch.device(
"cuda:0" if torch.cuda.is_available() and try_gpu else "cpu")
print(f"Using device: {device}")
self.models = []
for ckpt_file in model_ckpts:
print(f"Loading {ckpt_file}...")
self.models.append(
IgFold.load_from_checkpoint(ckpt_file).eval().to(device))
print(f"Successfully loaded {num_models} IgFold models.")
self.antiberty = AntiBERTyRunner()
self.antiberty.model.eval()
self.antiberty.model.to(device)
print("Loaded AntiBERTy model.")
def fold(
self,
pdb_file,
fasta_file=None,
sequences=None,
template_pdb=None,
ignore_cdrs=None,
ignore_chain=None,
skip_pdb=False,
do_refine=True,
use_openmm=False,
do_renum=True,
truncate_sequences=False,
):
"""
Predict antibody structure with IgFold.
:param pdb_file: PDB file to predict.
:param fasta_file: FASTA file containing sequences.
:param sequences: Dictionary of sequences.
:param template_pdb: PDB file containing template structure.
:param ignore_cdrs: List of CDRs to ignore.
:param ignore_chain: Chain to ignore.
:param skip_pdb: Skip PDB processing.
:param do_refine: Perform PyRosetta refinement.
:param do_renum: Renumber PDB to Chothia with AbNum.
:param truncate_sequences: Truncate sequences with AbNumber.
"""
start_time = time()
model_out = fold(
self.antiberty,
self.models,
pdb_file=pdb_file,
fasta_file=fasta_file,
sequences=sequences,
template_pdb=template_pdb,
ignore_cdrs=ignore_cdrs,
ignore_chain=ignore_chain,
skip_pdb=skip_pdb,
do_refine=do_refine,
use_openmm=use_openmm,
do_renum=do_renum,
truncate_sequences=truncate_sequences,
)
print(f"Completed folding in {time() - start_time:.2f} seconds.")
return model_out
def embed(
self,
model_idx=0,
fasta_file=None,
sequences=None,
template_pdb=None,
ignore_cdrs=None,
ignore_chain=None,
):
"""
Embed antibody sequences with IgFold.
:param fasta_file: FASTA file containing sequences.
:param sequences: Dictionary of sequences.
:param template_pdb: PDB file containing template structure.
:param ignore_cdrs: List of CDRs to ignore.
:param ignore_chain: Chain to ignore.
"""
start_time = time()
model_out = embed(
self.antiberty,
self.models[model_idx],
fasta_file=fasta_file,
sequences=sequences,
template_pdb=template_pdb,
ignore_cdrs=ignore_cdrs,
ignore_chain=ignore_chain,
)
print(f"Completed embedding in {time() - start_time:.2f} seconds.")
return model_out
|