File size: 10,836 Bytes
96f168d | 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 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 | import os
import torch
from Bio.PDB import PDBParser
from esm import FastaBatchedDataset, pretrained
from rdkit.Chem import AddHs, MolFromSmiles
from torch_geometric.data import Dataset, HeteroData
import esm
from datasets.process_mols import parse_pdb_from_path, generate_conformer, read_molecule, get_lig_graph_with_matching, \
extract_receptor_structure, get_rec_graph
three_to_one = {'ALA': 'A',
'ARG': 'R',
'ASN': 'N',
'ASP': 'D',
'CYS': 'C',
'GLN': 'Q',
'GLU': 'E',
'GLY': 'G',
'HIS': 'H',
'ILE': 'I',
'LEU': 'L',
'LYS': 'K',
'MET': 'M',
'MSE': 'M', # MSE this is almost the same AA as MET. The sulfur is just replaced by Selen
'PHE': 'F',
'PRO': 'P',
'PYL': 'O',
'SER': 'S',
'SEC': 'U',
'THR': 'T',
'TRP': 'W',
'TYR': 'Y',
'VAL': 'V',
'ASX': 'B',
'GLX': 'Z',
'XAA': 'X',
'XLE': 'J'}
def get_sequences_from_pdbfile(file_path):
biopython_parser = PDBParser()
structure = biopython_parser.get_structure('random_id', file_path)
structure = structure[0]
sequence = None
for i, chain in enumerate(structure):
seq = ''
for res_idx, residue in enumerate(chain):
if residue.get_resname() == 'HOH':
continue
residue_coords = []
c_alpha, n, c = None, None, None
for atom in residue:
if atom.name == 'CA':
c_alpha = list(atom.get_vector())
if atom.name == 'N':
n = list(atom.get_vector())
if atom.name == 'C':
c = list(atom.get_vector())
if c_alpha != None and n != None and c != None: # only append residue if it is an amino acid
try:
seq += three_to_one[residue.get_resname()]
except Exception as e:
seq += '-'
print("encountered unknown AA: ", residue.get_resname(), ' in the complex. Replacing it with a dash - .')
if sequence is None:
sequence = seq
else:
sequence += (":" + seq)
return sequence
def set_nones(l):
return [s if str(s) != 'nan' else None for s in l]
def get_sequences(protein_files, protein_sequences):
new_sequences = []
for i in range(len(protein_files)):
if protein_files[i] is not None:
new_sequences.append(get_sequences_from_pdbfile(protein_files[i]))
else:
new_sequences.append(protein_sequences[i])
return new_sequences
def compute_ESM_embeddings(model, alphabet, labels, sequences):
# settings used
toks_per_batch = 4096
repr_layers = [33]
include = "per_tok"
truncation_seq_length = 1022
dataset = FastaBatchedDataset(labels, sequences)
batches = dataset.get_batch_indices(toks_per_batch, extra_toks_per_seq=1)
data_loader = torch.utils.data.DataLoader(
dataset, collate_fn=alphabet.get_batch_converter(truncation_seq_length), batch_sampler=batches
)
assert all(-(model.num_layers + 1) <= i <= model.num_layers for i in repr_layers)
repr_layers = [(i + model.num_layers + 1) % (model.num_layers + 1) for i in repr_layers]
embeddings = {}
with torch.no_grad():
for batch_idx, (labels, strs, toks) in enumerate(data_loader):
print(f"Processing {batch_idx + 1} of {len(batches)} batches ({toks.size(0)} sequences)")
if torch.cuda.is_available():
toks = toks.to(device="cuda", non_blocking=True)
out = model(toks, repr_layers=repr_layers, return_contacts=False)
representations = {layer: t.to(device="cpu") for layer, t in out["representations"].items()}
for i, label in enumerate(labels):
truncate_len = min(truncation_seq_length, len(strs[i]))
embeddings[label] = representations[33][i, 1: truncate_len + 1].clone()
return embeddings
def generate_ESM_structure(model, filename, sequence):
model.set_chunk_size(256)
chunk_size = 256
output = None
while output is None:
try:
with torch.no_grad():
output = model.infer_pdb(sequence)
with open(filename, "w") as f:
f.write(output)
print("saved", filename)
except RuntimeError as e:
if 'out of memory' in str(e):
print('| WARNING: ran out of memory on chunk_size', chunk_size)
for p in model.parameters():
if p.grad is not None:
del p.grad # free some memory
torch.cuda.empty_cache()
chunk_size = chunk_size // 2
if chunk_size > 2:
model.set_chunk_size(chunk_size)
else:
print("Not enough memory for ESMFold")
break
else:
raise e
return output is not None
class InferenceDataset(Dataset):
def __init__(self, out_dir, complex_names, protein_files, ligand_descriptions, protein_sequences, lm_embeddings,
receptor_radius=30, c_alpha_max_neighbors=None, precomputed_lm_embeddings=None,
remove_hs=False, all_atoms=False, atom_radius=5, atom_max_neighbors=None):
super(InferenceDataset, self).__init__()
self.receptor_radius = receptor_radius
self.c_alpha_max_neighbors = c_alpha_max_neighbors
self.remove_hs = remove_hs
self.all_atoms = all_atoms
self.atom_radius, self.atom_max_neighbors = atom_radius, atom_max_neighbors
self.complex_names = complex_names
self.protein_files = protein_files
self.ligand_descriptions = ligand_descriptions
self.protein_sequences = protein_sequences
# generate LM embeddings
if lm_embeddings and (precomputed_lm_embeddings is None or precomputed_lm_embeddings[0] is None):
print("Generating ESM language model embeddings")
model_location = "esm2_t33_650M_UR50D"
model, alphabet = pretrained.load_model_and_alphabet(model_location)
model.eval()
if torch.cuda.is_available():
model = model.cuda()
protein_sequences = get_sequences(protein_files, protein_sequences)
labels, sequences = [], []
for i in range(len(protein_sequences)):
s = protein_sequences[i].split(':')
sequences.extend(s)
labels.extend([complex_names[i] + '_chain_' + str(j) for j in range(len(s))])
lm_embeddings = compute_ESM_embeddings(model, alphabet, labels, sequences)
self.lm_embeddings = []
for i in range(len(protein_sequences)):
s = protein_sequences[i].split(':')
self.lm_embeddings.append([lm_embeddings[f'{complex_names[i]}chain{j}'] for j in range(len(s))])
elif not lm_embeddings:
self.lm_embeddings = [None] * len(self.complex_names)
else:
self.lm_embeddings = precomputed_lm_embeddings
# generate structures with ESMFold
if None in protein_files:
print("generating missing structures with ESMFold")
model = esm.pretrained.esmfold_v1()
model = model.eval().cuda()
for i in range(len(protein_files)):
if protein_files[i] is None:
self.protein_files[i] = f"{out_dir}/{complex_names[i]}/{complex_names[i]}_esmfold.pdb"
if not os.path.exists(self.protein_files[i]):
print("generating", self.protein_files[i])
generate_ESM_structure(model, self.protein_files[i], protein_sequences[i])
def len(self):
return len(self.complex_names)
def get(self, idx):
name, protein_file, ligand_description, lm_embedding = \
self.complex_names[idx], self.protein_files[idx], self.ligand_descriptions[idx], self.lm_embeddings[idx]
# build the pytorch geometric heterogeneous graph
complex_graph = HeteroData()
complex_graph['name'] = name
# parse the ligand, either from file or smile
try:
mol = MolFromSmiles(ligand_description) # check if it is a smiles or a path
if mol is not None:
mol = AddHs(mol)
generate_conformer(mol)
else:
mol = read_molecule(ligand_description, remove_hs=False, sanitize=True)
if mol is None:
raise Exception('RDKit could not read the molecule ', ligand_description)
mol.RemoveAllConformers()
mol = AddHs(mol)
generate_conformer(mol)
except Exception as e:
print('Failed to read molecule ', ligand_description, ' We are skipping it. The reason is the exception: ', e)
complex_graph['success'] = False
return complex_graph
try:
# parse the receptor from the pdb file
rec_model = parse_pdb_from_path(protein_file)
get_lig_graph_with_matching(mol, complex_graph, popsize=None, maxiter=None, matching=False, keep_original=False,
num_conformers=1, remove_hs=self.remove_hs)
rec, rec_coords, c_alpha_coords, n_coords, c_coords, lm_embeddings = extract_receptor_structure(rec_model, mol, lm_embedding_chains=lm_embedding)
if lm_embeddings is not None and len(c_alpha_coords) != len(lm_embeddings):
print(f'LM embeddings for complex {name} did not have the right length for the protein. Skipping {name}.')
complex_graph['success'] = False
return complex_graph
get_rec_graph(rec, rec_coords, c_alpha_coords, n_coords, c_coords, complex_graph, rec_radius=self.receptor_radius,
c_alpha_max_neighbors=self.c_alpha_max_neighbors, all_atoms=self.all_atoms,
atom_radius=self.atom_radius, atom_max_neighbors=self.atom_max_neighbors, remove_hs=self.remove_hs, lm_embeddings=lm_embeddings)
except Exception as e:
print(f'Skipping {name} because of the error:')
print(e)
complex_graph['success'] = False
return complex_graph
protein_center = torch.mean(complex_graph['receptor'].pos, dim=0, keepdim=True)
complex_graph['receptor'].pos -= protein_center
if self.all_atoms:
complex_graph['atom'].pos -= protein_center
ligand_center = torch.mean(complex_graph['ligand'].pos, dim=0, keepdim=True)
complex_graph['ligand'].pos -= ligand_center
complex_graph.original_center = protein_center
complex_graph.mol = mol
complex_graph['success'] = True
return complex_graph
|