| from rdkit import Chem | |
| import codecs | |
| import numpy as np | |
| import pandas as pd | |
| from rdkit.Chem import Descriptors | |
| from subword_nmt.apply_bpe import BPE | |
| import deepchem as dc | |
| import deepchem.molnet as dcm | |
| from rdkit import RDLogger | |
| from torch.utils.data import DataLoader, Dataset, Subset | |
| from pytorch_lightning import LightningDataModule | |
| lg = RDLogger.logger() | |
| lg.setLevel(RDLogger.CRITICAL) | |
| def smiles2index(s1, words2idx, bpe): | |
| t1 = bpe.process_line(s1).split() | |
| i1 = [words2idx[i] for i in t1] | |
| return i1 | |
| def index2multi_hot(i1, idx2word): | |
| v1 = np.zeros(len(idx2word)) | |
| v1[i1] = 1 | |
| return v1 | |
| def index2multi_hot_fg(molecule, fgroups_list): | |
| v1 = np.zeros(len(fgroups_list)) | |
| for idx in range(len(fgroups_list)): | |
| if molecule.HasSubstructMatch(fgroups_list[idx]): | |
| v1[idx] = 1 | |
| return v1 | |
| def smiles2vector_fgr(s1, words2idx, bpe, idx2word, fgroups_list): | |
| i1 = smiles2index(s1, words2idx, bpe) | |
| mfg = index2multi_hot(i1, idx2word) | |
| molecule = Chem.MolFromSmiles(s1) | |
| fg = index2multi_hot_fg(molecule, fgroups_list) | |
| return fg, mfg | |
| class FsrFgDataset(Dataset): | |
| def __init__(self, data, idx2word, fgroups_list, words2idx, bpe, descriptor_funcs): | |
| self.mols = data.X | |
| self.y = data.y | |
| self.smiles = data.ids | |
| self.idx2word = idx2word | |
| self.fgroups_list = fgroups_list | |
| self.words2idx = words2idx | |
| self.bpe = bpe | |
| self.descriptor_funcs = descriptor_funcs | |
| def __len__(self): | |
| return len(self.y) | |
| def __getitem__(self, idx): | |
| mol = self.mols[idx] | |
| smile = self.smiles[idx] | |
| target = self.y[idx] | |
| fg, mfg = smiles2vector_fgr(smile, self.words2idx, self.bpe, self.idx2word, self.fgroups_list) | |
| num_features = np.asarray([self.descriptor_funcs[key](mol) for key in self.descriptor_funcs.keys()]) | |
| return np.float32(fg), np.float32(mfg), np.float32(num_features), int(target) | |
| class FsrFgDataModule(LightningDataModule): | |
| def __init__(self, root: str = 'Data', task_name: str = 'ecoli', batch_size: int = 8, num_workers: int = 8, | |
| pin_memory: bool = True, split_type: str = 'scaffold', num_folds: int = 5, fold_index: int = 0): | |
| super().__init__() | |
| self.root = root | |
| self.task_name = task_name | |
| self.num_workers = num_workers | |
| self.pin_memory = pin_memory | |
| self.batch_size = batch_size | |
| self.split_type = split_type | |
| self.num_folds = num_folds | |
| self.fold_index = fold_index | |
| self.test_fold = None | |
| self.val_fold = None | |
| self.train_fold = None | |
| self.fgroups_list = None | |
| self.descriptor_funcs = None | |
| self.bpe = None | |
| self.idx2word = None | |
| self.words2idx = None | |
| self.val_data = None | |
| self.test_data = None | |
| self.train_data = None | |
| self.splits = None | |
| def prepare_data(self) -> None: | |
| df = pd.read_csv('Data/ecoli_inhibition.csv') | |
| ids = df['smiles'] | |
| labels = df['Encoded_Activity'] | |
| X = [Chem.MolFromSmiles(smiles) for smiles in ids] | |
| data = dc.data.DiskDataset.from_numpy(X=X, y=labels, ids=ids) | |
| fgroups = pd.read_csv(self.root + '/Functional_groups_filtered.csv') | |
| fgroups_list = list(map(lambda x: Chem.MolFromSmarts(x), fgroups['SMARTS'].tolist())) | |
| self.fgroups_list = [i for i in fgroups_list if i] | |
| self.descriptor_funcs = {name: func for name, func in Descriptors.descList} | |
| vocab_path = self.root + '/codes_drug_chembl_1500.txt' | |
| bpe_codes_fin = codecs.open(vocab_path) | |
| self.bpe = BPE(bpe_codes_fin, merges=-1, separator='') | |
| vocab_map = pd.read_csv(self.root + '/subword_units_map_drug_chembl_1500.csv') | |
| self.idx2word = vocab_map['index'].values | |
| self.words2idx = dict(zip(self.idx2word, range(0, len(self.idx2word)))) | |
| self.weights = data.w | |
| if self.split_type == 'scaffold': | |
| splitter = dc.splits.ScaffoldSplitter() | |
| self.splits = [] | |
| for i in range(self.num_folds): | |
| new_data = data.complete_shuffle() | |
| self.splits.append((new_data, splitter.split(dataset=new_data, seed=i))) | |
| self.dataset = FsrFgDataset(self.splits[self.fold_index][0], self.idx2word, self.fgroups_list, self.words2idx, self.bpe, | |
| self.descriptor_funcs) | |
| self.train_ind, self.val_ind, self.test_ind = self.splits[self.fold_index][1] | |
| else: | |
| splitter = dc.splits.RandomStratifiedSplitter() | |
| self.splits = [splitter.split(dataset=data, seed=fold_num) for fold_num in range(self.num_folds)] | |
| self.dataset = FsrFgDataset(data, self.idx2word, self.fgroups_list, self.words2idx, self.bpe, | |
| self.descriptor_funcs) | |
| self.train_ind, self.val_ind, self.test_ind = self.splits[self.fold_index] | |
| def setup(self, stage=None): | |
| """define train, test and validation datasets """ | |
| self.train_fold = Subset(self.dataset, self.train_ind) | |
| self.val_fold = Subset(self.dataset, self.val_ind) | |
| self.test_fold = Subset(self.dataset, self.test_ind) | |
| def train_dataloader(self): | |
| """returns train dataloader""" | |
| loader = DataLoader(self.train_fold, batch_size=self.batch_size, shuffle=False, | |
| num_workers=self.num_workers, pin_memory=self.pin_memory, drop_last=True) | |
| return loader | |
| def val_dataloader(self): | |
| """returns val dataloader""" | |
| loader = DataLoader(self.val_fold, batch_size=1, shuffle=False, | |
| num_workers=self.num_workers, pin_memory=self.pin_memory) | |
| return loader | |
| def test_dataloader(self): | |
| """returns test dataloader""" | |
| loader = DataLoader(self.test_fold, batch_size=1, shuffle=False, | |
| num_workers=self.num_workers, pin_memory=self.pin_memory) | |
| return loader | |