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