mirror of
https://github.com/gcorso/DiffDock.git
synced 2026-06-04 18:04:23 +08:00
548 lines
29 KiB
Python
548 lines
29 KiB
Python
import os
|
|
import pickle
|
|
from multiprocessing import Pool
|
|
import random
|
|
import copy
|
|
from torch_geometric.data import Batch
|
|
|
|
import numpy as np
|
|
import torch
|
|
from prody import confProDy
|
|
from rdkit import Chem
|
|
from rdkit.Chem import RemoveHs
|
|
from torch_geometric.data import Dataset, HeteroData
|
|
from torch_geometric.utils import subgraph
|
|
from tqdm import tqdm
|
|
confProDy(verbosity='none')
|
|
from datasets.process_mols import get_lig_graph_with_matching, moad_extract_receptor_structure
|
|
from utils.utils import read_strings_from_txt
|
|
|
|
class MOAD(Dataset):
|
|
def __init__(self, root, transform=None, cache_path='data/cache', split='train', limit_complexes=0, chain_cutoff=None,
|
|
receptor_radius=30, num_workers=1, c_alpha_max_neighbors=None, popsize=15, maxiter=15,
|
|
matching=True, keep_original=False, max_lig_size=None, remove_hs=False, num_conformers=1, all_atoms=False,
|
|
atom_radius=5, atom_max_neighbors=None, esm_embeddings_path=None, esm_embeddings_sequences_path=None, require_ligand=False,
|
|
include_miscellaneous_atoms=False, keep_local_structures=False,
|
|
min_ligand_size=0, knn_only_graph=False, matching_tries=1, multiplicity=1,
|
|
max_receptor_size=None, remove_promiscuous_targets=None, unroll_clusters=False, remove_pdbbind=False,
|
|
enforce_timesplit=False, no_randomness=False, single_cluster_name=None, total_dataset_size=None, skip_matching=False):
|
|
|
|
super(MOAD, self).__init__(root, transform)
|
|
self.moad_dir = root
|
|
self.include_miscellaneous_atoms = include_miscellaneous_atoms
|
|
self.max_lig_size = max_lig_size
|
|
self.split = split
|
|
self.limit_complexes = limit_complexes
|
|
self.receptor_radius = receptor_radius
|
|
self.num_workers = num_workers
|
|
self.c_alpha_max_neighbors = c_alpha_max_neighbors
|
|
self.remove_hs = remove_hs
|
|
self.require_ligand = require_ligand
|
|
self.esm_embeddings_path = esm_embeddings_path
|
|
self.esm_embeddings_sequences_path = esm_embeddings_sequences_path
|
|
self.keep_local_structures = keep_local_structures
|
|
self.knn_only_graph = knn_only_graph
|
|
self.matching_tries = matching_tries
|
|
self.all_atoms = all_atoms
|
|
self.multiplicity = multiplicity
|
|
self.chain_cutoff = chain_cutoff
|
|
self.no_randomness = no_randomness
|
|
self.total_dataset_size = total_dataset_size
|
|
self.skip_matching = skip_matching
|
|
|
|
self.prot_cache_path = os.path.join(cache_path, f'MOAD12_limit{self.limit_complexes}_INDEX{self.split}'
|
|
f'_recRad{self.receptor_radius}_recMax{self.c_alpha_max_neighbors}'
|
|
+ (''if not all_atoms else f'_atomRad{atom_radius}_atomMax{atom_max_neighbors}')
|
|
+ ('' if self.esm_embeddings_path is None else f'_esmEmbeddings')
|
|
+ ('' if not self.include_miscellaneous_atoms else '_miscAtoms')
|
|
+ ('' if not self.knn_only_graph else '_knnOnly'))
|
|
|
|
self.lig_cache_path = os.path.join(cache_path, f'MOAD12_limit{self.limit_complexes}_INDEX{self.split}'
|
|
f'_maxLigSize{self.max_lig_size}_H{int(not self.remove_hs)}'
|
|
+ ('' if not matching else f'_matching')
|
|
+ ('' if not skip_matching else f'skip')
|
|
+ (''if not matching or num_conformers == 1 else f'_confs{num_conformers}')
|
|
+ ('' if not keep_local_structures else f'_keptLocalStruct')
|
|
+ ('' if self.matching_tries == 1 else f'_tries{matching_tries}'))
|
|
|
|
self.popsize, self.maxiter = popsize, maxiter
|
|
self.matching, self.keep_original = matching, keep_original
|
|
self.num_conformers = num_conformers
|
|
self.single_cluster_name = single_cluster_name
|
|
if split == 'train':
|
|
split = 'PDBBind'
|
|
|
|
with open("./data/splits/MOAD_generalisation_splits.pkl", "rb") as f:
|
|
self.split_clusters = pickle.load(f)[split]
|
|
|
|
clustes_path = os.path.join(self.moad_dir, "new_cluster_to_ligands.pkl")
|
|
with open(clustes_path, "rb") as f:
|
|
self.cluster_to_ligands = pickle.load(f)
|
|
#self.cluster_to_ligands = {k: [s.split('.')[0] for s in v] for k, v in self.cluster_to_ligands.items()}
|
|
|
|
self.atom_radius, self.atom_max_neighbors = atom_radius, atom_max_neighbors
|
|
if not self.check_all_receptors():
|
|
os.makedirs(self.prot_cache_path, exist_ok=True)
|
|
self.preprocessing_receptors()
|
|
|
|
self.atom_radius, self.atom_max_neighbors = atom_radius, atom_max_neighbors
|
|
if not os.path.exists(os.path.join(self.lig_cache_path, "ligands.pkl")):
|
|
os.makedirs(self.lig_cache_path, exist_ok=True)
|
|
self.preprocessing_ligands()
|
|
|
|
print('loading ligands from memory: ', os.path.join(self.lig_cache_path, "ligands.pkl"))
|
|
with open(os.path.join(self.lig_cache_path, "ligands.pkl"), 'rb') as f:
|
|
self.ligands = pickle.load(f)
|
|
|
|
if require_ligand:
|
|
with open(os.path.join(self.lig_cache_path, "rdkit_ligands.pkl"), 'rb') as f:
|
|
self.rdkit_ligands = pickle.load(f)
|
|
self.rdkit_ligands = {lig.name:mol for mol, lig in zip(self.rdkit_ligands, self.ligands)}
|
|
|
|
len_before = len(self.ligands)
|
|
if not self.single_cluster_name is None:
|
|
self.ligands = [lig for lig in self.ligands if lig.name in self.cluster_to_ligands[self.single_cluster_name]]
|
|
print('Kept', len(self.ligands), f'ligands in {self.single_cluster_name} out of', len_before)
|
|
|
|
len_before = len(self.ligands)
|
|
self.ligands = {lig.name: lig for lig in self.ligands if min_ligand_size == 0 or lig['ligand'].x.shape[0] >= min_ligand_size}
|
|
print('removed', len_before - len(self.ligands), 'ligands below minimum size out of', len_before)
|
|
|
|
receptors_names = set([lig.name[:6] for lig in self.ligands.values()])
|
|
self.collect_receptors(receptors_names, max_receptor_size, remove_promiscuous_targets)
|
|
|
|
# filter ligands for which the receptor failed
|
|
tot_before = len(self.ligands)
|
|
self.ligands = {k:v for k, v in self.ligands.items() if k[:6] in self.receptors}
|
|
print('removed', tot_before - len(self.ligands), 'ligands with no receptor out of', tot_before)
|
|
|
|
if remove_pdbbind:
|
|
complexes_pdbbind = read_strings_from_txt('data/splits/timesplit_no_lig_overlap_train') + read_strings_from_txt('data/splits/timesplit_no_lig_overlap_val')
|
|
with open('data/BindingMOAD_2020_ab_processed_biounit/ecod_t_group_binding_site_assignment_dict_major_domain.pkl', 'rb') as f:
|
|
pdbbind_to_cluster = pickle.load(f)
|
|
clusters_pdbbind = set([pdbbind_to_cluster[c] for c in complexes_pdbbind])
|
|
self.split_clusters = [c for c in self.split_clusters if c not in clusters_pdbbind]
|
|
self.cluster_to_ligands = {k: v for k, v in self.cluster_to_ligands.items() if k not in clusters_pdbbind}
|
|
ligand_accepted = []
|
|
for c, ligands in self.cluster_to_ligands.items():
|
|
ligand_accepted += ligands
|
|
ligand_accepted = set(ligand_accepted)
|
|
tot_before = len(self.ligands)
|
|
self.ligands = {k: v for k, v in self.ligands.items() if k in ligand_accepted}
|
|
print('removed', tot_before - len(self.ligands), 'ligands in overlap with PDBBind out of', tot_before)
|
|
|
|
if enforce_timesplit:
|
|
with open("data/splits/pdbids_2019", "r") as f:
|
|
lines = f.readlines()
|
|
pdbids_from2019 = []
|
|
for i in range(6, len(lines), 4):
|
|
pdbids_from2019.append(lines[i][18:22])
|
|
|
|
pdbids_from2019 = set(pdbids_from2019)
|
|
len_before = len(self.ligands)
|
|
self.ligands = {k: v for k, v in self.ligands.items() if k[:4].upper() not in pdbids_from2019}
|
|
print('removed', len_before - len(self.ligands), 'ligands from 2019 out of', len_before)
|
|
|
|
if unroll_clusters:
|
|
rec_keys = set([k[:6] for k in self.ligands.keys()])
|
|
self.cluster_to_ligands = {k:[k2 for k2 in self.ligands.keys() if k2[:6] == k] for k in rec_keys}
|
|
self.split_clusters = list(rec_keys)
|
|
else:
|
|
for c in self.cluster_to_ligands.keys():
|
|
self.cluster_to_ligands[c] = [v for v in self.cluster_to_ligands[c] if v in self.ligands]
|
|
self.split_clusters = [c for c in self.split_clusters if len(self.cluster_to_ligands[c])>0]
|
|
|
|
print_statistics(self)
|
|
list_names = [name for cluster in self.split_clusters for name in self.cluster_to_ligands[cluster]]
|
|
with open(os.path.join(self.prot_cache_path, f'moad_{self.split}_names.txt'), 'w') as f:
|
|
f.write('\n'.join(list_names))
|
|
|
|
def len(self):
|
|
return len(self.split_clusters) * self.multiplicity if self.total_dataset_size is None else self.total_dataset_size
|
|
|
|
def get_by_name(self, ligand_name, cluster):
|
|
ligand_graph = copy.deepcopy(self.ligands[ligand_name])
|
|
complex_graph = copy.deepcopy(self.receptors[ligand_name[:6]])
|
|
|
|
if False and self.keep_original and hasattr(ligand_graph['ligand'], 'orig_pos'):
|
|
lig_path = os.path.join(self.moad_dir, 'pdb_superligand', ligand_name + '.pdb')
|
|
lig = Chem.MolFromPDBFile(lig_path)
|
|
formula = np.asarray([atom.GetSymbol() for atom in lig.GetAtoms()])
|
|
|
|
# check for same receptor/ligand pair with a different binding position
|
|
for ligand_comp in self.cluster_to_ligands[cluster]:
|
|
if ligand_comp == ligand_name or ligand_comp[:6] != ligand_name[:6]:
|
|
continue
|
|
|
|
lig_path_comp = os.path.join(self.moad_dir, 'pdb_superligand', ligand_comp + '.pdb')
|
|
if not os.path.exists(lig_path_comp):
|
|
continue
|
|
|
|
lig_comp = Chem.MolFromPDBFile(lig_path_comp)
|
|
formula_comp = np.asarray([atom.GetSymbol() for atom in lig_comp.GetAtoms()])
|
|
|
|
if formula.shape == formula_comp.shape and np.all(formula == formula_comp) and hasattr(
|
|
self.ligands[ligand_comp], 'orig_pos'):
|
|
print(f'Found complex {ligand_comp} to have the same complex/ligand pair, adding it into orig_pos')
|
|
# add the orig_pos of the binding position
|
|
if not isinstance(ligand_graph['ligand'].orig_pos, list):
|
|
ligand_graph['ligand'].orig_pos = [ligand_graph['ligand'].orig_pos]
|
|
ligand_graph['ligand'].orig_pos.append(self.ligands[ligand_comp].orig_pos)
|
|
|
|
for type in ligand_graph.node_types + ligand_graph.edge_types:
|
|
for key, value in ligand_graph[type].items():
|
|
complex_graph[type][key] = value
|
|
complex_graph.name = ligand_graph.name
|
|
if isinstance(complex_graph['ligand'].pos, list):
|
|
for i in range(len(complex_graph['ligand'].pos)):
|
|
complex_graph['ligand'].pos[i] -= complex_graph.original_center
|
|
else:
|
|
complex_graph['ligand'].pos -= complex_graph.original_center
|
|
if self.require_ligand:
|
|
complex_graph.mol = copy.deepcopy(self.rdkit_ligands[ligand_name])
|
|
|
|
if self.chain_cutoff:
|
|
distances = torch.norm(
|
|
(torch.from_numpy(complex_graph['ligand'].orig_pos[0]) - complex_graph.original_center).unsqueeze(1) - complex_graph['receptor'].pos.unsqueeze(0), dim=2)
|
|
distances = distances.min(dim=0)[0]
|
|
if torch.min(distances) >= self.chain_cutoff:
|
|
print('minimum distance', torch.min(distances), 'too large', ligand_name,
|
|
'skipping and returning random. Number of chains',
|
|
torch.max(complex_graph['receptor'].chain_ids) + 1)
|
|
return self.get(random.randint(0, self.len()))
|
|
|
|
within_cutoff = distances < self.chain_cutoff
|
|
chains_within_cutoff = torch.zeros(torch.max(complex_graph['receptor'].chain_ids) + 1)
|
|
chains_within_cutoff.index_add_(0, complex_graph['receptor'].chain_ids, within_cutoff.float())
|
|
chains_within_cutoff_bool = chains_within_cutoff > 0
|
|
residues_to_keep = chains_within_cutoff_bool[complex_graph['receptor'].chain_ids]
|
|
|
|
if self.all_atoms:
|
|
atom_to_res_mapping = complex_graph['atom', 'atom_rec_contact', 'receptor'].edge_index[1]
|
|
atoms_to_keep = residues_to_keep[atom_to_res_mapping]
|
|
rec_remapper = (torch.cumsum(residues_to_keep.long(), dim=0) - 1)
|
|
atom_to_res_new_mapping = rec_remapper[atom_to_res_mapping][atoms_to_keep]
|
|
atom_res_edge_index = torch.stack([torch.arange(len(atom_to_res_new_mapping)), atom_to_res_new_mapping])
|
|
|
|
complex_graph['atom'].x = complex_graph['atom'].x[atoms_to_keep]
|
|
complex_graph['atom'].pos = complex_graph['atom'].pos[atoms_to_keep]
|
|
complex_graph['atom', 'atom_contact', 'atom'].edge_index = \
|
|
subgraph(atoms_to_keep, complex_graph['atom', 'atom_contact', 'atom'].edge_index,
|
|
relabel_nodes=True)[0]
|
|
complex_graph['atom', 'atom_rec_contact', 'receptor'].edge_index = atom_res_edge_index
|
|
|
|
complex_graph['receptor'].pos = complex_graph['receptor'].pos[residues_to_keep]
|
|
complex_graph['receptor'].x = complex_graph['receptor'].x[residues_to_keep]
|
|
complex_graph['receptor'].side_chain_vecs = complex_graph['receptor'].side_chain_vecs[residues_to_keep]
|
|
complex_graph['receptor', 'rec_contact', 'receptor'].edge_index = \
|
|
subgraph(residues_to_keep, complex_graph['receptor', 'rec_contact', 'receptor'].edge_index,
|
|
relabel_nodes=True)[0]
|
|
|
|
extra_center = torch.mean(complex_graph['receptor'].pos, dim=0, keepdim=True)
|
|
complex_graph['receptor'].pos -= extra_center
|
|
if isinstance(complex_graph['ligand'].pos, list):
|
|
for i in range(len(complex_graph['ligand'].pos)):
|
|
complex_graph['ligand'].pos[i] -= extra_center
|
|
else:
|
|
complex_graph['ligand'].pos -= extra_center
|
|
complex_graph.original_center += extra_center
|
|
|
|
complex_graph['receptor'].pop('chain_ids')
|
|
|
|
for a in ['random_coords', 'coords', 'seq', 'sequence', 'mask', 'rmsd_matching', 'cluster', 'orig_seq',
|
|
'to_keep', 'chain_ids']:
|
|
if hasattr(complex_graph, a):
|
|
delattr(complex_graph, a)
|
|
if hasattr(complex_graph['receptor'], a):
|
|
delattr(complex_graph['receptor'], a)
|
|
|
|
return complex_graph
|
|
|
|
def get(self, idx):
|
|
if self.total_dataset_size is not None:
|
|
idx = random.randint(0, len(self.split_clusters) - 1)
|
|
|
|
idx = idx % len(self.split_clusters)
|
|
cluster = self.split_clusters[idx]
|
|
|
|
if self.no_randomness:
|
|
ligand_name = sorted(self.cluster_to_ligands[cluster])[0]
|
|
else:
|
|
ligand_name = random.choice(self.cluster_to_ligands[cluster])
|
|
|
|
complex_graph = self.get_by_name(ligand_name, cluster)
|
|
|
|
if self.total_dataset_size is not None:
|
|
complex_graph = Batch.from_data_list([complex_graph])
|
|
|
|
return complex_graph
|
|
|
|
def get_all_complexes(self):
|
|
complexes = {}
|
|
for cluster in self.split_clusters:
|
|
for ligand_name in self.cluster_to_ligands[cluster]:
|
|
complexes[ligand_name] = self.get_by_name(ligand_name, cluster)
|
|
return complexes
|
|
|
|
def preprocessing_receptors(self):
|
|
print(f'Processing receptors from [{self.split}] and saving it to [{self.prot_cache_path}]')
|
|
|
|
complex_names_all = sorted([l for c in self.split_clusters for l in self.cluster_to_ligands[c]])
|
|
if self.limit_complexes is not None and self.limit_complexes != 0:
|
|
complex_names_all = complex_names_all[:self.limit_complexes]
|
|
|
|
receptor_names_all = [l[:6] for l in complex_names_all]
|
|
receptor_names_all = sorted(list(dict.fromkeys(receptor_names_all)))
|
|
print(f'Loading {len(receptor_names_all)} receptors.')
|
|
|
|
if self.esm_embeddings_path is not None:
|
|
id_to_embeddings = torch.load(self.esm_embeddings_path)
|
|
sequences_list = read_strings_from_txt(self.esm_embeddings_sequences_path)
|
|
sequences_to_embeddings = {}
|
|
for i, seq in enumerate(sequences_list):
|
|
sequences_to_embeddings[seq] = id_to_embeddings[str(i)]
|
|
|
|
else:
|
|
sequences_to_embeddings = None
|
|
|
|
# running preprocessing in parallel on multiple workers and saving the progress every 1000 complexes
|
|
list_indices = list(range(len(receptor_names_all)//1000+1))
|
|
random.shuffle(list_indices)
|
|
for i in list_indices:
|
|
if os.path.exists(os.path.join(self.prot_cache_path, f"receptors{i}.pkl")):
|
|
continue
|
|
receptor_names = receptor_names_all[1000*i:1000*(i+1)]
|
|
receptor_graphs = []
|
|
if self.num_workers > 1:
|
|
p = Pool(self.num_workers, maxtasksperchild=1)
|
|
p.__enter__()
|
|
with tqdm(total=len(receptor_names), desc=f'loading receptors {i}/{len(receptor_names_all)//1000+1}') as pbar:
|
|
map_fn = p.imap_unordered if self.num_workers > 1 else map
|
|
for t in map_fn(self.get_receptor, zip(receptor_names, [sequences_to_embeddings]*len(receptor_names))):
|
|
if t is not None:
|
|
print(len(receptor_graphs))
|
|
receptor_graphs.append(t)
|
|
pbar.update()
|
|
if self.num_workers > 1: p.__exit__(None, None, None)
|
|
|
|
print('Number of receptors: ', len(receptor_graphs))
|
|
with open(os.path.join(self.prot_cache_path, f"receptors{i}.pkl"), 'wb') as f:
|
|
pickle.dump((receptor_graphs), f)
|
|
return receptor_names_all
|
|
|
|
def check_all_receptors(self):
|
|
complex_names_all = sorted([l for c in self.split_clusters for l in self.cluster_to_ligands[c]])
|
|
if self.limit_complexes is not None and self.limit_complexes != 0:
|
|
complex_names_all = complex_names_all[:self.limit_complexes]
|
|
receptor_names_all = [l[:6] for l in complex_names_all]
|
|
receptor_names_all = list(dict.fromkeys(receptor_names_all))
|
|
for i in range(len(receptor_names_all)//1000+1):
|
|
if not os.path.exists(os.path.join(self.prot_cache_path, f"receptors{i}.pkl")):
|
|
return False
|
|
return True
|
|
|
|
def collect_receptors(self, receptors_to_keep=None, max_receptor_size=None, remove_promiscuous_targets=None):
|
|
complex_names_all = sorted([l for c in self.split_clusters for l in self.cluster_to_ligands[c]])
|
|
if self.limit_complexes is not None and self.limit_complexes != 0:
|
|
complex_names_all = complex_names_all[:self.limit_complexes]
|
|
receptor_names_all = [l[:6] for l in complex_names_all]
|
|
receptor_names_all = sorted(list(dict.fromkeys(receptor_names_all)))
|
|
|
|
receptor_graphs_all = []
|
|
total_recovered = 0
|
|
print(f'Loading {len(receptor_names_all)} receptors to keep {len(receptors_to_keep)}.')
|
|
for i in range(len(receptor_names_all)//1000+1):
|
|
print(f'prot path: {os.path.join(self.prot_cache_path, f"receptors{i}.pkl")}')
|
|
with open(os.path.join(self.prot_cache_path, f"receptors{i}.pkl"), 'rb') as f:
|
|
l = pickle.load(f)
|
|
total_recovered += len(l)
|
|
if receptors_to_keep is not None:
|
|
l = [t for t in l if t['receptor_name'] in receptors_to_keep]
|
|
receptor_graphs_all.extend(l)
|
|
|
|
cur_len = len(receptor_graphs_all)
|
|
print(f"Kept {len(receptor_graphs_all)} receptors out of {len(receptor_names_all)} total and recovered {total_recovered}")
|
|
|
|
if max_receptor_size is not None:
|
|
receptor_graphs_all = [rec for rec in receptor_graphs_all if rec["receptor"].pos.shape[0] <= max_receptor_size]
|
|
print(f"Kept {len(receptor_graphs_all)} receptors out of {cur_len} after filtering by size")
|
|
cur_len = len(receptor_graphs_all)
|
|
|
|
if remove_promiscuous_targets is not None:
|
|
promiscuous_targets = set()
|
|
for name in complex_names_all:
|
|
l = name.split('_')
|
|
if int(l[3]) > remove_promiscuous_targets:
|
|
promiscuous_targets.add(name[:6])
|
|
receptor_graphs_all = [rec for rec in receptor_graphs_all if rec["receptor_name"] not in promiscuous_targets]
|
|
print(f"Kept {len(receptor_graphs_all)} receptors out of {cur_len} after removing promiscuous targets")
|
|
|
|
self.receptors = {}
|
|
for r in receptor_graphs_all:
|
|
self.receptors[r['receptor_name']] = r
|
|
return
|
|
|
|
def get_receptor(self, par):
|
|
name, sequences_to_embeddings = par
|
|
rec_path = os.path.join(self.moad_dir, 'pdb_protein', name + '_protein.pdb')
|
|
if not os.path.exists(rec_path):
|
|
print("Receptor not found", name, rec_path)
|
|
return None
|
|
|
|
complex_graph = HeteroData()
|
|
complex_graph['receptor_name'] = name
|
|
try:
|
|
moad_extract_receptor_structure(path=rec_path, complex_graph=complex_graph, neighbor_cutoff=self.receptor_radius,
|
|
max_neighbors=self.c_alpha_max_neighbors, sequences_to_embeddings=sequences_to_embeddings,
|
|
knn_only_graph=self.knn_only_graph, all_atoms=self.all_atoms, atom_cutoff=self.atom_radius,
|
|
atom_max_neighbors=self.atom_max_neighbors)
|
|
|
|
except Exception as e:
|
|
print(f'Skipping {name} because of the error:')
|
|
print(e)
|
|
return None
|
|
|
|
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
|
|
complex_graph.original_center = protein_center
|
|
return complex_graph
|
|
|
|
|
|
def preprocessing_ligands(self):
|
|
print(f'Processing complexes from [{self.split}] and saving it to [{self.lig_cache_path}]')
|
|
|
|
complex_names_all = sorted([l for c in self.split_clusters for l in self.cluster_to_ligands[c]])
|
|
if self.limit_complexes is not None and self.limit_complexes != 0:
|
|
complex_names_all = complex_names_all[:self.limit_complexes]
|
|
print(f'Loading {len(complex_names_all)} ligands.')
|
|
|
|
# running preprocessing in parallel on multiple workers and saving the progress every 1000 complexes
|
|
list_indices = list(range(len(complex_names_all)//1000+1))
|
|
random.shuffle(list_indices)
|
|
for i in list_indices:
|
|
if os.path.exists(os.path.join(self.lig_cache_path, f"ligands{i}.pkl")):
|
|
continue
|
|
complex_names = complex_names_all[1000*i:1000*(i+1)]
|
|
ligand_graphs, rdkit_ligands = [], []
|
|
if self.num_workers > 1:
|
|
p = Pool(self.num_workers, maxtasksperchild=1)
|
|
p.__enter__()
|
|
with tqdm(total=len(complex_names), desc=f'loading complexes {i}/{len(complex_names_all)//1000+1}') as pbar:
|
|
map_fn = p.imap_unordered if self.num_workers > 1 else map
|
|
for t in map_fn(self.get_ligand, complex_names):
|
|
if t is not None:
|
|
ligand_graphs.append(t[0])
|
|
rdkit_ligands.append(t[1])
|
|
pbar.update()
|
|
if self.num_workers > 1: p.__exit__(None, None, None)
|
|
|
|
with open(os.path.join(self.lig_cache_path, f"ligands{i}.pkl"), 'wb') as f:
|
|
pickle.dump((ligand_graphs), f)
|
|
with open(os.path.join(self.lig_cache_path, f"rdkit_ligands{i}.pkl"), 'wb') as f:
|
|
pickle.dump((rdkit_ligands), f)
|
|
|
|
ligand_graphs_all = []
|
|
for i in range(len(complex_names_all)//1000+1):
|
|
with open(os.path.join(self.lig_cache_path, f"ligands{i}.pkl"), 'rb') as f:
|
|
l = pickle.load(f)
|
|
ligand_graphs_all.extend(l)
|
|
with open(os.path.join(self.lig_cache_path, f"ligands.pkl"), 'wb') as f:
|
|
pickle.dump((ligand_graphs_all), f)
|
|
|
|
rdkit_ligands_all = []
|
|
for i in range(len(complex_names_all) // 1000 + 1):
|
|
with open(os.path.join(self.lig_cache_path, f"rdkit_ligands{i}.pkl"), 'rb') as f:
|
|
l = pickle.load(f)
|
|
rdkit_ligands_all.extend(l)
|
|
with open(os.path.join(self.lig_cache_path, f"rdkit_ligands.pkl"), 'wb') as f:
|
|
pickle.dump((rdkit_ligands_all), f)
|
|
|
|
def get_ligand(self, name):
|
|
if self.split == 'train':
|
|
lig_path = os.path.join(self.moad_dir, 'pdb_superligand', name + '.pdb')
|
|
else:
|
|
lig_path = os.path.join(self.moad_dir, 'pdb_ligand', name + '.pdb')
|
|
|
|
if not os.path.exists(lig_path):
|
|
print("Ligand not found", name, lig_path)
|
|
return None
|
|
|
|
# read pickle
|
|
lig = Chem.MolFromPDBFile(lig_path)
|
|
|
|
if self.max_lig_size is not None and lig.GetNumHeavyAtoms() > self.max_lig_size:
|
|
print(f'Ligand with {lig.GetNumHeavyAtoms()} heavy atoms is larger than max_lig_size {self.max_lig_size}. Not including {name} in preprocessed data.')
|
|
return None
|
|
|
|
try:
|
|
if self.matching:
|
|
smile = Chem.MolToSmiles(lig)
|
|
if '.' in smile:
|
|
print(f'Ligand {name} has multiple fragments and we are doing matching. Not including {name} in preprocessed data.')
|
|
return None
|
|
|
|
complex_graph = HeteroData()
|
|
complex_graph['name'] = name
|
|
|
|
Chem.SanitizeMol(lig)
|
|
get_lig_graph_with_matching(lig, complex_graph, self.popsize, self.maxiter, self.matching, self.keep_original,
|
|
self.num_conformers, remove_hs=self.remove_hs, tries=self.matching_tries, skip_matching=self.skip_matching)
|
|
except Exception as e:
|
|
print(f'Skipping {name} because of the error:')
|
|
print(e)
|
|
return None
|
|
|
|
if self.split != 'train':
|
|
other_positions = [complex_graph['ligand'].orig_pos]
|
|
nsplit = name.split('_')
|
|
for i in range(100):
|
|
new_file = os.path.join(self.moad_dir, 'pdb_ligand', f'{nsplit[0]}_{nsplit[1]}_{nsplit[2]}_{i}.pdb')
|
|
if os.path.exists(new_file):
|
|
if i != int(nsplit[3]):
|
|
lig = Chem.MolFromPDBFile(new_file)
|
|
lig = RemoveHs(lig, sanitize=True)
|
|
other_positions.append(lig.GetConformer().GetPositions())
|
|
else:
|
|
break
|
|
complex_graph['ligand'].orig_pos = np.asarray(other_positions)
|
|
|
|
return complex_graph, lig
|
|
|
|
|
|
def print_statistics(dataset):
|
|
statistics = ([], [], [], [], [], [])
|
|
receptor_sizes = []
|
|
|
|
for i in range(len(dataset)):
|
|
complex_graph = dataset[i]
|
|
lig_pos = complex_graph['ligand'].pos if torch.is_tensor(complex_graph['ligand'].pos) else complex_graph['ligand'].pos[0]
|
|
receptor_sizes.append(complex_graph['receptor'].pos.shape[0])
|
|
radius_protein = torch.max(torch.linalg.vector_norm(complex_graph['receptor'].pos, dim=1))
|
|
molecule_center = torch.mean(lig_pos, dim=0)
|
|
radius_molecule = torch.max(
|
|
torch.linalg.vector_norm(lig_pos - molecule_center.unsqueeze(0), dim=1))
|
|
distance_center = torch.linalg.vector_norm(molecule_center)
|
|
statistics[0].append(radius_protein)
|
|
statistics[1].append(radius_molecule)
|
|
statistics[2].append(distance_center)
|
|
if "rmsd_matching" in complex_graph:
|
|
statistics[3].append(complex_graph.rmsd_matching)
|
|
else:
|
|
statistics[3].append(0)
|
|
statistics[4].append(int(complex_graph.random_coords) if "random_coords" in complex_graph else -1)
|
|
if "random_coords" in complex_graph and complex_graph.random_coords and "rmsd_matching" in complex_graph:
|
|
statistics[5].append(complex_graph.rmsd_matching)
|
|
|
|
if len(statistics[5]) == 0:
|
|
statistics[5].append(-1)
|
|
name = ['radius protein', 'radius molecule', 'distance protein-mol', 'rmsd matching', 'random coordinates', 'random rmsd matching']
|
|
print('Number of complexes: ', len(dataset))
|
|
for i in range(len(name)):
|
|
array = np.asarray(statistics[i])
|
|
print(f"{name[i]}: mean {np.mean(array)}, std {np.std(array)}, max {np.max(array)}")
|
|
|
|
return
|
|
|