from typing import Tuple, Union, Any, Dict, Optional, List from typing_extensions import Self from enum import Enum import json import os import sys from tqdm import tqdm import numpy as np import pandas as pd from rdkit import Chem from torch.utils.data import Dataset from rdkit.Chem import AllChem, Descriptors from open_biomed.data import Molecule, Text from open_biomed.datasets.base_dataset import BaseDataset, assign_split, featurize from open_biomed.utils.config import Config from open_biomed.utils.featurizer import Featurizer, Featurized from open_biomed.utils.split_utils import random_split, scaffold_split class MoleculePropertyPredictionDataset(BaseDataset): def __init__(self, cfg: Config, featurizer: Featurizer) -> None: self.molecules, self.texts, self.labels = [], [], [] super(MoleculePropertyPredictionDataset, self).__init__(cfg, featurizer) def __len__(self) -> int: return len(self.molecules) @featurize def __getitem__(self, index) -> Dict[str, Featurized[Any]]: return { "molecule": self.molecules[index], "label": self.labels[index], } class MoleculePropertyPredictionEvalDataset(Dataset): def __init__(self) -> None: super(MoleculePropertyPredictionEvalDataset, self).__init__() self.molecules, self.texts, self.labels = [], [], [] # @classmethod # def from_train_set(cls, dataset: MoleculePropertyPredictionDataset) -> Self: # # TODO: when encoutering multiple results for the same original molecule and text, load will be crashed # # so we ignore same smiles multi labels # # just be multiple mol-labels pairs for multi-labels # # NOTE: # # Given the same original molecule and text, multiple results are acceptable # # We combine these results for evaluation # mol2label = dict() # for i in range(len(dataset)): # molecule = dataset.molecules[i] # label = dataset.labels[i] # dict_key = str(molecule) # if dict_key not in mol2label: # mol2label[dict_key] = [] # mol2label[dict_key].append((molecule, label)) # new_dataset = cls() # for k, v in mol2label.items(): # new_dataset.molecules.append(v[0][0]) # new_dataset.labels.append([x[1] for x in v]) # new_dataset.featurizer = dataset.featurizer # return new_dataset @classmethod def from_train_set(cls, dataset: MoleculePropertyPredictionDataset) -> Self: # NOTE: # WARNING: ignore same smiles different labels new_dataset = cls() new_dataset.molecules = dataset.molecules new_dataset.labels = [] for label in dataset.labels: new_dataset.labels.append([label]) new_dataset.featurizer = dataset.featurizer return new_dataset def __len__(self) -> int: return len(self.molecules) @featurize def __getitem__(self, index) -> Dict[str, Featurized[Any]]: return { "molecule": self.molecules[index], "label": self.labels[index], } def get_labels(self) -> List[List[Text]]: return self.labels def _load_bbbp_dataset(input_path): input_df = pd.read_csv(input_path, sep=',') smiles_list = input_df['smiles'] rdkit_mol_objs_list = [AllChem.MolFromSmiles(s) for s in smiles_list] preprocessed_rdkit_mol_objs_list = [m if m is not None else None for m in rdkit_mol_objs_list] preprocessed_smiles_list = [AllChem.MolToSmiles(m) if m is not None else None for m in preprocessed_rdkit_mol_objs_list] labels = input_df['p_np'] # convert 0 to -1 labels = labels.replace(0, -1) # there are no nans assert len(smiles_list) == len(preprocessed_rdkit_mol_objs_list) assert len(smiles_list) == len(preprocessed_smiles_list) assert len(smiles_list) == len(labels) return preprocessed_smiles_list, \ preprocessed_rdkit_mol_objs_list, labels.values def _load_clintox_dataset(input_path): input_df = pd.read_csv(input_path, sep=',') smiles_list = input_df['smiles'] rdkit_mol_objs_list = [AllChem.MolFromSmiles(s) for s in smiles_list] preprocessed_rdkit_mol_objs_list = [m if m is not None else None for m in rdkit_mol_objs_list] preprocessed_smiles_list = [AllChem.MolToSmiles(m) if m is not None else None for m in preprocessed_rdkit_mol_objs_list] tasks = ['FDA_APPROVED', 'CT_TOX'] labels = input_df[tasks] # convert 0 to -1 labels = labels.replace(0, -1) # there are no nans assert len(smiles_list) == len(preprocessed_rdkit_mol_objs_list) assert len(smiles_list) == len(preprocessed_smiles_list) assert len(smiles_list) == len(labels) return preprocessed_smiles_list, \ preprocessed_rdkit_mol_objs_list, labels.values # input_path = 'dataset/clintox/raw/clintox.csv' # smiles_list, rdkit_mol_objs_list, labels = _load_clintox_dataset(input_path) def _load_esol_dataset(input_path): # NB: some examples have multiple species input_df = pd.read_csv(input_path, sep=',') smiles_list = input_df['smiles'] rdkit_mol_objs_list = [AllChem.MolFromSmiles(s) for s in smiles_list] labels = input_df['measured log solubility in mols per litre'] assert len(smiles_list) == len(rdkit_mol_objs_list) assert len(smiles_list) == len(labels) return smiles_list, rdkit_mol_objs_list, labels.values # input_path = 'dataset/esol/raw/delaney-processed.csv' # smiles_list, rdkit_mol_objs_list, labels = _load_esol_dataset(input_path) def _load_tdc_classification_dataset(input_path): input_df_trainval = pd.read_csv(os.path.join(input_path, 'train_val.csv'), sep=',') input_df_test = pd.read_csv(os.path.join(input_path, 'test.csv'), sep=',') smiles_list_trainval = input_df_trainval['Drug'].tolist() smiles_list_test = input_df_test['Drug'].tolist() smiles_list_trainval_copy = smiles_list_trainval.copy() rdkit_mol_objs_list_trainval = [AllChem.MolFromSmiles(s) for s in smiles_list_trainval] rdkit_mol_objs_list_test = [AllChem.MolFromSmiles(s) for s in smiles_list_test] rdkit_mol_objs_list_trainval_copy = rdkit_mol_objs_list_trainval.copy() labels_trainval = input_df_trainval['Y'].replace(0, -1).tolist() labels_test = input_df_test['Y'].replace(0, -1).tolist() labels_trainval_copy = labels_trainval.copy() trainval_indices = list(range(len(smiles_list_trainval))) test_indices = list(range(len(smiles_list_trainval), len(smiles_list_trainval) + len(smiles_list_test))) assert len(smiles_list_trainval) == len(rdkit_mol_objs_list_trainval) assert len(smiles_list_trainval) == len(labels_trainval) assert len(smiles_list_test) == len(rdkit_mol_objs_list_test) assert len(smiles_list_test) == len(labels_test) assert len(rdkit_mol_objs_list_trainval) + len(rdkit_mol_objs_list_test) == len(smiles_list_trainval) + len(smiles_list_test) smiles_list_trainval.extend(smiles_list_test) rdkit_mol_objs_list_trainval.extend(rdkit_mol_objs_list_test) labels_trainval.extend(labels_test) smiles_list = np.array(smiles_list_trainval) rdkit_mol_objs_list = np.array(rdkit_mol_objs_list_trainval) labels = np.array(labels_trainval) # TODO now just concat return return smiles_list, \ rdkit_mol_objs_list, \ labels, \ smiles_list_trainval_copy, \ rdkit_mol_objs_list_trainval_copy, \ labels_trainval_copy, \ test_indices def _load_tdc_regression_dataset(input_path): input_df_trainval = pd.read_csv(os.path.join(input_path, 'train_val.csv'), sep=',') input_df_test = pd.read_csv(os.path.join(input_path, 'test.csv'), sep=',') smiles_list_trainval = input_df_trainval['Drug'].tolist() smiles_list_test = input_df_test['Drug'].tolist() smiles_list_trainval_copy = smiles_list_trainval.copy() rdkit_mol_objs_list_trainval = [AllChem.MolFromSmiles(s) for s in smiles_list_trainval] rdkit_mol_objs_list_test = [AllChem.MolFromSmiles(s) for s in smiles_list_test] rdkit_mol_objs_list_trainval_copy = rdkit_mol_objs_list_trainval.copy() labels_trainval = input_df_trainval['Y'].tolist() labels_test = input_df_test['Y'].tolist() labels_trainval_copy = labels_trainval.copy() trainval_indices = list(range(len(smiles_list_trainval))) test_indices = list(range(len(smiles_list_trainval), len(smiles_list_trainval) + len(smiles_list_test))) assert len(smiles_list_trainval) == len(rdkit_mol_objs_list_trainval) assert len(smiles_list_trainval) == len(labels_trainval) assert len(smiles_list_test) == len(rdkit_mol_objs_list_test) assert len(smiles_list_test) == len(labels_test) assert len(rdkit_mol_objs_list_trainval) + len(rdkit_mol_objs_list_test) == len(smiles_list_trainval) + len(smiles_list_test) smiles_list_trainval.extend(smiles_list_test) rdkit_mol_objs_list_trainval.extend(rdkit_mol_objs_list_test) labels_trainval.extend(labels_test) smiles_list = np.array(smiles_list_trainval) rdkit_mol_objs_list = np.array(rdkit_mol_objs_list_trainval) labels = np.array(labels_trainval) # TODO now just concat return return smiles_list, \ rdkit_mol_objs_list, \ labels, \ smiles_list_trainval_copy, \ rdkit_mol_objs_list_trainval_copy, \ labels_trainval_copy, \ test_indices def _load_caco2_dataset(input_path): return _load_tdc_regression_dataset(input_path) def _load_bbb_dataset(input_path): return _load_tdc_classification_dataset(input_path) def _load_cyp2c9_inhibition_dataset(input_path): return _load_tdc_classification_dataset(input_path) def _load_half_life_dataset(input_path): return _load_tdc_regression_dataset(input_path) def _load_ld50_dataset(input_path): return _load_tdc_regression_dataset(input_path) # input_path = 'dataset/admet_group/caco2_wang/' def _load_freesolv_dataset(input_path): input_df = pd.read_csv(input_path, sep=',') smiles_list = input_df['smiles'] rdkit_mol_objs_list = [AllChem.MolFromSmiles(s) for s in smiles_list] labels = input_df['expt'] assert len(smiles_list) == len(rdkit_mol_objs_list) assert len(smiles_list) == len(labels) return smiles_list, rdkit_mol_objs_list, labels.values def _load_lipophilicity_dataset(input_path): input_df = pd.read_csv(input_path, sep=',') smiles_list = input_df['smiles'] rdkit_mol_objs_list = [AllChem.MolFromSmiles(s) for s in smiles_list] labels = input_df['exp'] assert len(smiles_list) == len(rdkit_mol_objs_list) assert len(smiles_list) == len(labels) return smiles_list, rdkit_mol_objs_list, labels.values def _load_malaria_dataset(input_path): input_df = pd.read_csv(input_path, sep=',') smiles_list = input_df['smiles'] rdkit_mol_objs_list = [AllChem.MolFromSmiles(s) for s in smiles_list] labels = input_df['activity'] assert len(smiles_list) == len(rdkit_mol_objs_list) assert len(smiles_list) == len(labels) return smiles_list, rdkit_mol_objs_list, labels.values def _load_cep_dataset(input_path): input_df = pd.read_csv(input_path, sep=',') smiles_list = input_df['smiles'] rdkit_mol_objs_list = [AllChem.MolFromSmiles(s) for s in smiles_list] labels = input_df['PCE'] assert len(smiles_list) == len(rdkit_mol_objs_list) assert len(smiles_list) == len(labels) return smiles_list, rdkit_mol_objs_list, labels.values def _load_muv_dataset(input_path): input_df = pd.read_csv(input_path, sep=',') smiles_list = input_df['smiles'] rdkit_mol_objs_list = [AllChem.MolFromSmiles(s) for s in smiles_list] tasks = ['MUV-466', 'MUV-548', 'MUV-600', 'MUV-644', 'MUV-652', 'MUV-689', 'MUV-692', 'MUV-712', 'MUV-713', 'MUV-733', 'MUV-737', 'MUV-810', 'MUV-832', 'MUV-846', 'MUV-852', 'MUV-858', 'MUV-859'] labels = input_df[tasks] # convert 0 to -1 labels = labels.replace(0, -1) # convert nan to 0 labels = labels.fillna(0) assert len(smiles_list) == len(rdkit_mol_objs_list) assert len(smiles_list) == len(labels) return smiles_list, rdkit_mol_objs_list, labels.values def check_columns(df, tasks, N): bad_tasks = [] total_missing_count = 0 for task in tasks: value_list = df[task] pos_count = sum(value_list == 1) neg_count = sum(value_list == -1) missing_count = sum(value_list == 0) total_missing_count += missing_count pos_ratio = 100. * pos_count / (pos_count + neg_count) missing_ratio = 100. * missing_count / N assert pos_count + neg_count + missing_count == N if missing_ratio >= 50: bad_tasks.append(task) print('task {}\t\tpos_ratio: {:.5f}\tmissing ratio: {:.5f}'.format(task, pos_ratio, missing_ratio)) print('total missing ratio: {:.5f}'.format(100. * total_missing_count / len(tasks) / N)) return bad_tasks def check_rows(labels, N): from collections import defaultdict p, n, m = defaultdict(int), defaultdict(int), defaultdict(int) bad_count = 0 for i in range(N): value_list = labels[i] pos_count = sum(value_list == 1) neg_count = sum(value_list == -1) missing_count = sum(value_list == 0) p[pos_count] += 1 n[neg_count] += 1 m[missing_count] += 1 if pos_count + neg_count == 0: bad_count += 1 print('bad_count\t', bad_count) print('pos\t', p) print('neg\t', n) print('missing\t', m) return def _load_pcba_dataset(input_path): input_df = pd.read_csv(input_path, sep=',') tasks = list(input_df.columns)[:-2] N = input_df.shape[0] temp_df = input_df[tasks] temp_df = temp_df.replace(0, -1) temp_df = temp_df.fillna(0) bad_tasks = check_columns(temp_df, tasks, N) for task in bad_tasks: tasks.remove(task) print('good tasks\t', len(tasks)) labels = input_df[tasks] labels = labels.replace(0, -1) labels = labels.fillna(0) labels = labels.values print(labels.shape) # 439863, 92 check_rows(labels, N) input_df.dropna(subset=tasks, how='all', inplace=True) # convert 0 to -1 # input_df = input_df.replace(0, -1) # convert nan to 0 input_df = input_df.fillna(0) labels = input_df[tasks].values print(input_df.shape) # 435685, 92 N = input_df.shape[0] check_rows(labels, N) smiles_list = input_df['smiles'].tolist() rdkit_mol_objs_list = [AllChem.MolFromSmiles(s) for s in smiles_list] assert len(smiles_list) == len(rdkit_mol_objs_list) assert len(smiles_list) == len(labels) return smiles_list, rdkit_mol_objs_list, labels def _load_sider_dataset(input_path): input_df = pd.read_csv(input_path, sep=',') smiles_list = input_df['smiles'] rdkit_mol_objs_list = [AllChem.MolFromSmiles(s) for s in smiles_list] tasks = ['Hepatobiliary disorders', 'Metabolism and nutrition disorders', 'Product issues', 'Eye disorders', 'Investigations', 'Musculoskeletal and connective tissue disorders', 'Gastrointestinal disorders', 'Social circumstances', 'Immune system disorders', 'Reproductive system and breast disorders', 'Neoplasms benign, malignant and unspecified (incl cysts and polyps)', 'General disorders and administration site conditions', 'Endocrine disorders', 'Surgical and medical procedures', 'Vascular disorders', 'Blood and lymphatic system disorders', 'Skin and subcutaneous tissue disorders', 'Congenital, familial and genetic disorders', 'Infections and infestations', 'Respiratory, thoracic and mediastinal disorders', 'Psychiatric disorders', 'Renal and urinary disorders', 'Pregnancy, puerperium and perinatal conditions', 'Ear and labyrinth disorders', 'Cardiac disorders', 'Nervous system disorders', 'Injury, poisoning and procedural complications'] labels = input_df[tasks] # convert 0 to -1 labels = labels.replace(0, -1) assert len(smiles_list) == len(rdkit_mol_objs_list) assert len(smiles_list) == len(labels) return smiles_list, rdkit_mol_objs_list, labels.values def _load_toxcast_dataset(input_path): # NB: some examples have multiple species, some example smiles are invalid input_df = pd.read_csv(input_path, sep=',') smiles_list = input_df['smiles'] rdkit_mol_objs_list = [AllChem.MolFromSmiles(s) for s in smiles_list] # Some smiles could not be successfully converted # to rdkit mol object so them to None preprocessed_rdkit_mol_objs_list = [m if m is not None else None for m in rdkit_mol_objs_list] preprocessed_smiles_list = [AllChem.MolToSmiles(m) if m is not None else None for m in preprocessed_rdkit_mol_objs_list] tasks = list(input_df.columns)[1:] labels = input_df[tasks] # convert 0 to -1 labels = labels.replace(0, -1) # convert nan to 0 labels = labels.fillna(0) assert len(smiles_list) == len(preprocessed_rdkit_mol_objs_list) assert len(smiles_list) == len(preprocessed_smiles_list) assert len(smiles_list) == len(labels) return preprocessed_smiles_list, \ preprocessed_rdkit_mol_objs_list, labels.values # root_path = 'dataset/chembl_with_labels' def check_smiles_validity(smiles): try: m = Chem.MolFromSmiles(smiles) if m: return True else: return False except: return False def split_rdkit_mol_obj(mol): """ Split rdkit mol object containing multiple species or one species into a list of mol objects or a list containing a single object respectively """ smiles = AllChem.MolToSmiles(mol, isomericSmiles=True) smiles_list = smiles.split('.') mol_species_list = [] for s in smiles_list: if check_smiles_validity(s): mol_species_list.append(AllChem.MolFromSmiles(s)) return mol_species_list def get_largest_mol(mol_list): """ Given a list of rdkit mol objects, returns mol object containing the largest num of atoms. If multiple containing largest num of atoms, picks the first one """ num_atoms_list = [len(m.GetAtoms()) for m in mol_list] largest_mol_idx = num_atoms_list.index(max(num_atoms_list)) return mol_list[largest_mol_idx] def _load_tox21_dataset(input_path): input_df = pd.read_csv(input_path, sep=',') smiles_list = input_df['smiles'] rdkit_mol_objs_list = [AllChem.MolFromSmiles(s) for s in smiles_list] tasks = ['NR-AR', 'NR-AR-LBD', 'NR-AhR', 'NR-Aromatase', 'NR-ER', 'NR-ER-LBD', 'NR-PPAR-gamma', 'SR-ARE', 'SR-ATAD5', 'SR-HSE', 'SR-MMP', 'SR-p53'] labels = input_df[tasks] # convert 0 to -1 labels = labels.replace(0, -1) # convert nan to 0 labels = labels.fillna(0) assert len(smiles_list) == len(rdkit_mol_objs_list) assert len(smiles_list) == len(labels) return smiles_list, rdkit_mol_objs_list, labels.values def _load_hiv_dataset(input_path): input_df = pd.read_csv(input_path, sep=',') smiles_list = input_df['smiles'] rdkit_mol_objs_list = [AllChem.MolFromSmiles(s) for s in smiles_list] labels = input_df['HIV_active'] # convert 0 to -1 labels = labels.replace(0, -1) # there are no nans assert len(smiles_list) == len(rdkit_mol_objs_list) assert len(smiles_list) == len(labels) return smiles_list, rdkit_mol_objs_list, labels.values def _load_bace_dataset(input_path): input_df = pd.read_csv(input_path, sep=',') smiles_list = input_df['smiles'] rdkit_mol_objs_list = [AllChem.MolFromSmiles(s) for s in smiles_list] labels = input_df['Class'] # convert 0 to -1 labels = labels.replace(0, -1) # there are no nans folds = input_df['Model'] folds = folds.replace('Train', 0) # 0 -> train folds = folds.replace('Valid', 1) # 1 -> valid folds = folds.replace('Test', 2) # 2 -> test assert len(smiles_list) == len(rdkit_mol_objs_list) assert len(smiles_list) == len(labels) assert len(smiles_list) == len(folds) # return smiles_list, rdkit_mol_objs_list, folds.values, labels.values return smiles_list, rdkit_mol_objs_list, labels.values datasetname2function = { "bbbp": _load_bbbp_dataset, "clintox": _load_clintox_dataset, "tox21": _load_tox21_dataset, "toxcast": _load_toxcast_dataset, "sider": _load_sider_dataset, "hiv": _load_hiv_dataset, "bace": _load_bace_dataset, "muv": _load_muv_dataset, "freesolv": _load_freesolv_dataset, "esol": _load_esol_dataset, "lipophilicity": _load_lipophilicity_dataset, "caco2_wang": _load_caco2_dataset, "bbb_martins": _load_bbb_dataset, "cyp2c9_veith": _load_cyp2c9_inhibition_dataset, "half_life_obach": _load_half_life_dataset, "ld50_zhu": _load_ld50_dataset, } class Task(Enum): CLASSFICATION = 0 REGRESSION = 1 class MoleculeNet(MoleculePropertyPredictionDataset): name2target = { "BBBP": ["p_np"], "Tox21": ["NR-AR", "NR-AR-LBD", "NR-AhR", "NR-Aromatase", "NR-ER", "NR-ER-LBD", "NR-PPAR-gamma", "SR-ARE", "SR-ATAD5", "SR-HSE", "SR-MMP", "SR-p53"], "ClinTox": ["CT_TOX", "FDA_APPROVED"], "HIV": ["HIV_active"], "Bace": ["class"], "SIDER": ["Hepatobiliary disorders", "Metabolism and nutrition disorders", "Product issues", "Eye disorders", "Investigations", "Musculoskeletal and connective tissue disorders", "Gastrointestinal disorders", "Social circumstances", "Immune system disorders", "Reproductive system and breast disorders", "Neoplasms benign, malignant and unspecified (incl cysts and polyps)", "General disorders and administration site conditions", "Endocrine disorders", "Surgical and medical procedures", "Vascular disorders", "Blood and lymphatic system disorders", "Skin and subcutaneous tissue disorders", "Congenital, familial and genetic disorders", "Infections and infestations", "Respiratory, thoracic and mediastinal disorders", "Psychiatric disorders", "Renal and urinary disorders", "Pregnancy, puerperium and perinatal conditions", "Ear and labyrinth disorders", "Cardiac disorders", "Nervous system disorders", "Injury, poisoning and procedural complications"], "MUV": ['MUV-692', 'MUV-689', 'MUV-846', 'MUV-859', 'MUV-644', 'MUV-548', 'MUV-852', 'MUV-600', 'MUV-810', 'MUV-712', 'MUV-737', 'MUV-858', 'MUV-713', 'MUV-733', 'MUV-652', 'MUV-466', 'MUV-832'], "Toxcast": [""], # 617 "FreeSolv": ["expt"], "ESOL": ["measured log solubility in mols per litre"], "Lipo": ["exp"], "qm7": ["u0_atom"], "qm8": ["E1-CC2", "E2-CC2", "f1-CC2", "f2-CC2", "E1-PBE0", "E2-PBE0", "f1-PBE0", "f2-PBE0", "E1-CAM", "E2-CAM", "f1-CAM","f2-CAM"], "qm9": ['mu', 'alpha', 'homo', 'lumo', 'gap', 'r2', 'zpve', 'cv'], "caco2_wang": [""], "bbb_martins": [""], "cyp2c9_veith": [""], "half_life_obach":[""], "ld50_zhu": [""], } name2task = { "BBBP": Task.CLASSFICATION, "Tox21": Task.CLASSFICATION, "ClinTox": Task.CLASSFICATION, "HIV": Task.CLASSFICATION, "Bace": Task.CLASSFICATION, "SIDER": Task.CLASSFICATION, "MUV": Task.CLASSFICATION, "Toxcast": Task.CLASSFICATION, "bbb_martins": Task.CLASSFICATION, "cyp2c9_veith": Task.CLASSFICATION, "FreeSolv": Task.REGRESSION, "ESOL": Task.REGRESSION, "Lipo": Task.REGRESSION, "qm7": Task.REGRESSION, "qm8": Task.REGRESSION, "qm9": Task.REGRESSION, "caco2_wang": Task.REGRESSION, "half_life_obach": Task.REGRESSION, "ld50_zhu": Task.REGRESSION, } name2text = { "bbbp": "Binary labels of blood-brain barrier penetration(permeability)", "clintox": "Qualitative data of drugs that failed clinical trials for toxicity reasons", "tox21": "Qualitative toxicity measurements including nuclear receptors and stress response pathways", "toxcast": "Compounds based on in vitro high-throughput screening", "sider": "marketed drugs and adverse drug reactions (ADR), grouped into 27 system organ classes", "hiv": "Experimentally measured abilities to inhibit HIV replication", "bace": "Quantitative (IC50) and qualitative (binary label) binding results for human β-secretase 1(BACE-1)", "muv": "Subset of PubChem BioAssay designed for validation of virtual screening techniques", } tdc_names = ["caco2_wang","bbb_martins","cyp2c9_veith","half_life_obach","ld50_zhu"] def __init__(self, cfg: Config, featurizer: Featurizer) -> None: self.name = cfg.name self.targets = self.name2target[cfg.name] # TODO: 看一下graphmvp这里是干什么用的，后续上regression任务的时候要考虑 if self.name not in self.tdc_names: self.task = self.name2task[cfg.name] file_name = os.listdir(os.path.join(cfg.path, self.name.lower(), "raw"))[0] assert file_name[-4:] == ".csv" self.path = os.path.join(cfg.path, self.name.lower(), "raw", file_name) else: self.task = self.name2task[cfg.name] # file_name = os.listdir(cfg.path)[0] original_path = cfg.path parent_dir = os.path.dirname(original_path) dataset_dir = os.path.basename(original_path) new_path = os.path.join(parent_dir, 'admet_group', dataset_dir) file_name = os.listdir(new_path)[0] assert file_name[-4:] == ".csv" self.path = new_path cfg.path = new_path super(MoleculeNet, self).__init__(cfg, featurizer) def _train_test_split(self, strategy="scaffold"): if strategy == "random": self.train_index, self.validation_index, self.test_index = random_split(len(self), 0.1, 0.1) elif strategy == "scaffold": self.train_index, self.validation_index, self.test_index = scaffold_split(self, 0.1, 0.1, is_standard=True) def _normalize(self): if self.name in ["qm7", "qm9"]: self.normalizer = [] for i in range(len(self.targets)): self.normalizer.append(Normalizer(self.labels[:, i])) self.labels[:, i] = self.normalizer[i].norm(self.labels[:, i]) else: # TODO: self.normalizer = [None] * len(self.targets) def _load_data(self) -> None: if self.name not in self.tdc_names: smiles_list, rdkit_mol_objs, labels = datasetname2function[self.name.lower()](self.path) else: smiles_list, rdkit_mol_objs, labels, smiles_list_trainval, rdkit_mol_objs_trainval, labels_trainval, test_ids = datasetname2function[self.name.lower()](self.path) if labels.ndim == 1: labels = np.expand_dims(labels, axis=1) def load_smiles_and_labels(smiles_list, rdkit_mol_objs, labels): self.molecules = [] self.labels = [] for i in range(len(smiles_list)): rdkit_mol = rdkit_mol_objs[i] if rdkit_mol is None: continue # TODO: drugs and smiles are all get from AllChem.MolFromSmiles() #self.smiles.append(smiles_list[i]) self.molecules.append(Molecule.from_smiles(smiles_list[i])) self.labels.append(labels[i]) load_smiles_and_labels(smiles_list, rdkit_mol_objs, labels) if self.name not in self.tdc_names: self._train_test_split() else: load_smiles_and_labels(smiles_list_trainval, rdkit_mol_objs_trainval, labels_trainval) self.train_index, self.validation_index, self.test_index = scaffold_split(self, 0.05, 0.05, is_standard=True) self.validation_index.extend(self.test_index) self.test_index = test_ids load_smiles_and_labels(smiles_list, rdkit_mol_objs, labels) self._normalize() self.split_indexes = {} for split in ["train", "validation", "test"]: self.split_indexes[split] = getattr(self, f"{split}_index") @assign_split def split(self, split_cfg: Optional[Config] = None) -> Tuple[Any, Any, Any]: attrs = ["molecules", "labels"] ret = ( self.get_subset(self.split_indexes["train"], attrs), MoleculePropertyPredictionEvalDataset.from_train_set(self.get_subset(self.split_indexes["validation"], attrs)), MoleculePropertyPredictionEvalDataset.from_train_set(self.get_subset(self.split_indexes["test"], attrs)), ) del self return ret