# Copyright (c) DP Technology. # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. from functools import lru_cache import numpy as np from unicore.data import BaseWrapperDataset import pickle from . import data_utils class AffinityDataset(BaseWrapperDataset): def __init__( self, dataset, seed, atoms, coordinates, pocket_atoms, pocket_coordinates, affinity, is_train=False, pocket="pocket" ): self.dataset = dataset self.seed = seed self.atoms = atoms self.coordinates = coordinates self.pocket_atoms = pocket_atoms self.pocket_coordinates = pocket_coordinates self.affinity = affinity self.is_train = is_train self.pocket=pocket self.set_epoch(None) def set_epoch(self, epoch, **unused): super().set_epoch(epoch) self.epoch = epoch def pocket_atom(self, atom): if atom[0] in ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9']: return atom[1] else: return atom[0] @lru_cache(maxsize=16) def __cached_item__(self, index: int, epoch: int): atoms = np.array(self.dataset[index][self.atoms]) ori_mol_length = len(atoms) #coordinates = self.dataset[index][self.coordinates] size = len(self.dataset[index][self.coordinates]) if self.is_train: with data_utils.numpy_seed(self.seed, epoch, index): sample_idx = np.random.randint(size) else: with data_utils.numpy_seed(self.seed, 1, index): sample_idx = np.random.randint(size) #print(len(self.dataset[index][self.coordinates][sample_idx])) coordinates = self.dataset[index][self.coordinates][sample_idx] #print(coordinates.shape) pocket_atoms = np.array( [self.pocket_atom(item) for item in self.dataset[index][self.pocket_atoms]] ) ori_pocket_length = len(pocket_atoms) pocket_coordinates = np.stack(self.dataset[index][self.pocket_coordinates]) smi = self.dataset[index]["smi"] pocket = self.dataset[index][self.pocket] if self.affinity in self.dataset[index]: affinity = float(self.dataset[index][self.affinity]) else: affinity = 1 return { "atoms": atoms, "coordinates": coordinates.astype(np.float32), "holo_coordinates": coordinates.astype(np.float32),#placeholder "pocket_atoms": pocket_atoms, "pocket_coordinates": pocket_coordinates.astype(np.float32), "holo_pocket_coordinates": pocket_coordinates.astype(np.float32),#placeholder "smi": smi, "pocket": pocket, "affinity": affinity, "ori_mol_length": ori_mol_length, "ori_pocket_length": ori_pocket_length } def __getitem__(self, index: int): return self.__cached_item__(index, self.epoch) class AffinityAugDataset(BaseWrapperDataset): def __init__( self, dataset, seed, atoms, coordinates, pocket_atoms, pocket_coordinates, affinity, is_train=False, pocket="pocket_id" ): self.dataset = dataset self.seed = seed self.atoms = atoms self.coordinates = coordinates self.pocket_atoms = pocket_atoms self.pocket_coordinates = pocket_coordinates self.affinity = affinity self.is_train = is_train self.pocket=pocket self.set_epoch(None) def set_epoch(self, epoch, **unused): super().set_epoch(epoch) self.epoch = epoch def pocket_atom(self, atom): if atom[0] in ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9']: return atom[1] else: return atom[0] @lru_cache(maxsize=16) def __cached_item__(self, index: int, epoch: int): #mol_atoms_list = self.dataset[index][self.atoms] with data_utils.numpy_seed(self.seed, epoch, index): mol_idx = np.random.randint(len(self.dataset[index][self.atoms])) atoms = np.array(self.dataset[index][self.atoms][mol_idx]) ori_mol_length = len(atoms) #coordinates = self.dataset[index][self.coordinates] size = len(self.dataset[index][self.coordinates][mol_idx]) if self.is_train: with data_utils.numpy_seed(self.seed, epoch, index): sample_idx = np.random.randint(size) else: with data_utils.numpy_seed(self.seed, 1, index): sample_idx = np.random.randint(size) #print(len(self.dataset[index][self.coordinates][sample_idx])) coordinates = self.dataset[index][self.coordinates][mol_idx][sample_idx] #pocket_list = self.dataset[index][self.pocket_atoms] with data_utils.numpy_seed(self.seed, epoch, index): pocket_idx = np.random.randint(len(self.dataset[index][self.pocket_atoms])) pocket_atoms = np.array( [self.pocket_atom(item) for item in self.dataset[index][self.pocket_atoms][pocket_idx]] ) ori_pocket_length = len(pocket_atoms) pocket_coordinates = np.stack(self.dataset[index][self.pocket_coordinates][pocket_idx]) smi = self.dataset[index]["smiles"][mol_idx] pocket = self.dataset[index][self.pocket][0] if self.affinity in self.dataset[index]: affinity = float(self.dataset[index][self.affinity]) else: affinity = 1 return { "atoms": atoms, "coordinates": coordinates.astype(np.float32), "holo_coordinates": coordinates.astype(np.float32),#placeholder "pocket_atoms": pocket_atoms, "pocket_coordinates": pocket_coordinates.astype(np.float32), "holo_pocket_coordinates": pocket_coordinates.astype(np.float32),#placeholder "smi": smi, "pocket": pocket, "affinity": affinity, "ori_mol_length": ori_mol_length, "ori_pocket_length": ori_pocket_length } def __getitem__(self, index: int): return self.__cached_item__(index, self.epoch) class AffinityHNSDataset(BaseWrapperDataset): def __init__( self, dataset, seed, atoms, coordinates, atoms_hns, coordinates_hns, pocket_atoms, pocket_coordinates, affinity, is_train=False, pocket="pocket" ): self.dataset = dataset self.seed = seed self.atoms = atoms self.coordinates = coordinates self.atoms_hns = atoms_hns self.coordinates_hns = coordinates_hns self.pocket_atoms = pocket_atoms self.pocket_coordinates = pocket_coordinates self.affinity = affinity self.is_train = is_train self.pocket=pocket self.set_epoch(None) def set_epoch(self, epoch, **unused): super().set_epoch(epoch) self.epoch = epoch def pocket_atom(self, atom): if atom[0] in ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9']: return atom[1] else: return atom[0] @lru_cache(maxsize=16) def __cached_item__(self, index: int, epoch: int): atoms = np.array(self.dataset[index][self.atoms]) ori_mol_length = len(atoms) #coordinates = self.dataset[index][self.coordinates] size = len(self.dataset[index][self.coordinates]) if self.is_train: with data_utils.numpy_seed(self.seed, epoch, index): sample_idx = np.random.randint(size) else: with data_utils.numpy_seed(self.seed, 1, index): sample_idx = np.random.randint(size) #print(len(self.dataset[index][self.coordinates][sample_idx])) coordinates = self.dataset[index][self.coordinates][sample_idx] atoms_hns = np.array(self.dataset[index][self.atoms_hns]) coordinates_hns = self.dataset[index][self.coordinates_hns][0] pocket_atoms = np.array( [self.pocket_atom(item) for item in self.dataset[index][self.pocket_atoms]] ) ori_pocket_length = len(pocket_atoms) pocket_coordinates = np.stack(self.dataset[index][self.pocket_coordinates]) smi = self.dataset[index]["smi"] pocket = self.dataset[index][self.pocket] if self.affinity in self.dataset[index]: affinity = float(self.dataset[index][self.affinity]) else: affinity = 1 return { "atoms": atoms, "coordinates": coordinates.astype(np.float32), "atoms_hns": atoms_hns, "coordinates_hns": coordinates_hns.astype(np.float32), "holo_coordinates": coordinates.astype(np.float32),#placeholder "pocket_atoms": pocket_atoms, "pocket_coordinates": pocket_coordinates.astype(np.float32), "holo_pocket_coordinates": pocket_coordinates.astype(np.float32),#placeholder "smi": smi, "pocket": pocket, "affinity": affinity, "ori_mol_length": ori_mol_length, "ori_pocket_length": ori_pocket_length } def __getitem__(self, index: int): return self.__cached_item__(index, self.epoch) class AffinityTestDataset(BaseWrapperDataset): def __init__( self, dataset, seed, atoms, coordinates, pocket_atoms, pocket_coordinates, affinity=None, is_train=False, pocket="pocket" ): self.dataset = dataset self.seed = seed self.atoms = atoms self.coordinates = coordinates self.pocket_atoms = pocket_atoms self.pocket_coordinates = pocket_coordinates self.affinity = affinity self.is_train = is_train self.pocket=pocket self.set_epoch(None) def set_epoch(self, epoch, **unused): super().set_epoch(epoch) self.epoch = epoch def pocket_atom(self, atom): if atom[0] in ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9']: return atom[1] else: return atom[0] @lru_cache(maxsize=16) def __cached_item__(self, index: int, epoch: int): atoms = np.array(self.dataset[index][self.atoms]) ori_length = len(atoms) #coordinates = self.dataset[index][self.coordinates] size = len(self.dataset[index][self.coordinates]) with data_utils.numpy_seed(self.seed, epoch, index): sample_idx = np.random.randint(size) coordinates = self.dataset[index][self.coordinates][sample_idx] pocket_atoms = np.array( [self.pocket_atom(item) for item in self.dataset[index][self.pocket_atoms]] ) #print(len(self.dataset[index][self.pocket_coordinates])) pocket_coordinates = np.stack(self.dataset[index][self.pocket_coordinates]) smi = self.dataset[index]["smi"] pocket = self.dataset[index][self.pocket] affinity = self.dataset[index][self.affinity] return { "atoms": atoms, "coordinates": coordinates.astype(np.float32), "holo_coordinates": coordinates.astype(np.float32),#placeholder "pocket_atoms": pocket_atoms, "pocket_coordinates": pocket_coordinates.astype(np.float32), "holo_pocket_coordinates": pocket_coordinates.astype(np.float32),#placeholder "smi": smi, "pocket": pocket, "affinity": affinity.astype(np.float32), "ori_length": ori_length } def __getitem__(self, index: int): return self.__cached_item__(index, self.epoch) class AffinityMolDataset(BaseWrapperDataset): def __init__( self, dataset, seed, atoms, coordinates, is_train=False, ): self.dataset = dataset self.seed = seed self.atoms = atoms self.coordinates = coordinates self.is_train = is_train self.set_epoch(None) def set_epoch(self, epoch, **unused): super().set_epoch(epoch) self.epoch = epoch def pocket_atom(self, atom): if atom[0] in ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9']: return atom[1] else: return atom[0] @lru_cache(maxsize=16) def __cached_item__(self, index: int, epoch: int): #print(self.dataset[index]) atoms = np.array(self.dataset[index][self.atoms]) ori_length = len(atoms) #coordinates = self.dataset[index][self.coordinates] size = len(self.dataset[index][self.coordinates]) #print(size) # TODO: FB: introduce enough random when training using pairwise data # with data_utils.numpy_seed(self.seed, epoch, index): if self.is_train: sample_idx = np.random.randint(size) else: with data_utils.numpy_seed(self.seed, index): sample_idx = np.random.randint(size) # check coordinates is 2 dimension or not if len(self.dataset[index][self.coordinates][sample_idx].shape) == 2: coordinates = self.dataset[index][self.coordinates][sample_idx] else: coordinates = self.dataset[index][self.coordinates] #coordinates = self.dataset[index][self.coordinates][sample_idx] #coordinates = self.dataset[index][self.coordinates] smi = self.dataset[index]["smi"] name = self.dataset[index].get("name", None) mol = pickle.dumps(self.dataset[index].get("mol", None)) return { "atoms": atoms, "coordinates": coordinates.astype(np.float32), "holo_coordinates": coordinates.astype(np.float32),#placeholder "smi": smi, "ori_length": ori_length, "name": name, "mol": mol } def __getitem__(self, index: int): return self.__cached_item__(index, self.epoch) class AffinityPocketDataset(BaseWrapperDataset): def __init__( self, dataset, seed, pocket_atoms, pocket_coordinates, is_train=False, pocket="pocket" ): self.dataset = dataset self.seed = seed self.pocket_atoms = pocket_atoms self.pocket_coordinates = pocket_coordinates self.is_train = is_train self.pocket=pocket self.set_epoch(None) def set_epoch(self, epoch, **unused): super().set_epoch(epoch) self.epoch = epoch def pocket_atom(self, atom): if atom[0] in ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9']: return atom[1] else: return atom[0] @lru_cache(maxsize=16) def __cached_item__(self, index: int, epoch: int): # print(self.dataset[index].keys()) pocket_atoms = np.array( [self.pocket_atom(item) for item in self.dataset[index][self.pocket_atoms]] ) ori_length = len(pocket_atoms) pocket_coordinates = np.stack(self.dataset[index][self.pocket_coordinates]) if self.pocket in self.dataset[index]: pocket = self.dataset[index][self.pocket] else: pocket = "" if "pocket_residue_name" in self.dataset[index]: pocket_residue_name = self.dataset[index]["pocket_residue_name"] pocket_residue_name_noH = [] for res, atom in zip(pocket_residue_name, pocket_atoms): if atom == "H": continue pocket_residue_name_noH.append(res) else: pocket_residue_name_noH = [""] return { "pocket_atoms": pocket_atoms, "pocket_coordinates": pocket_coordinates.astype(np.float32), "holo_pocket_coordinates": pocket_coordinates.astype(np.float32),#placeholder "pocket": pocket, "pocket_residue_name": pocket_residue_name_noH, "ori_length": ori_length } def __getitem__(self, index: int): return self.__cached_item__(index, self.epoch) class AffinityValidDataset(BaseWrapperDataset): def __init__( self, dataset, seed, atoms, coordinates, pocket_atoms, pocket_coordinates, pocket="pocket" ): self.dataset = dataset self.seed = seed self.atoms = atoms self.coordinates = coordinates self.pocket_atoms = pocket_atoms self.pocket_coordinates = pocket_coordinates self.pocket=pocket self.set_epoch(None) def set_epoch(self, epoch, **unused): super().set_epoch(epoch) self.epoch = epoch def pocket_atom(self, atom): if atom[0] in ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9']: return atom[1] else: return atom[0] @lru_cache(maxsize=16) def __cached_item__(self, index: int, epoch: int): atoms = np.array(self.dataset[index][self.atoms]) ori_mol_length = len(atoms) #coordinates = self.dataset[index][self.coordinates] size = len(self.dataset[index][self.coordinates]) with data_utils.numpy_seed(self.seed, epoch, index): sample_idx = np.random.randint(size) coordinates = self.dataset[index][self.coordinates][sample_idx] pocket_atoms = np.array( [self.pocket_atom(item) for item in self.dataset[index][self.pocket_atoms]] ) ori_pocket_length = len(pocket_atoms) pocket_coordinates = np.stack(self.dataset[index][self.pocket_coordinates]) smi = self.dataset[index]["smi"] pocket = self.dataset[index][self.pocket] return { "atoms": atoms, "coordinates": coordinates.astype(np.float32), "holo_coordinates": coordinates.astype(np.float32),#placeholder "pocket_atoms": pocket_atoms, "pocket_coordinates": pocket_coordinates.astype(np.float32), "holo_pocket_coordinates": pocket_coordinates.astype(np.float32),#placeholder "smi": smi, "pocket": pocket, "ori_mol_length": ori_mol_length, "ori_pocket_length": ori_pocket_length } def __getitem__(self, index: int): return self.__cached_item__(index, self.epoch)