hjxdl 0.0.1__py3-none-any.whl

hdl/data/dataset/fp/fp_dataset.py

@@ -0,0 +1,122 @@
+import typing as t
+
+import torch
+from rdkit import Chem
+# import torch.utils.data as tud
+
+from hdl.data.dataset.base_dataset import CSVDataset, CSVRDataset
+from hdl.features.fp.features_generators import (
+    get_features_generator,
+    get_available_features_generators,
+    FP_BITS_DICT 
+)
+# from hdl.features.fp.rxn import get_rxnrep_fingerprint
+
+
+class FPDataset(CSVDataset):
+    def __init__(
+        self,
+        csv_file: str,
+        splitter: str,
+        smiles_cols: t.List,
+        target_cols: t.List = [],
+        missing_labels: t.List = [],
+        num_classes: t.List = [],
+        target_transform: t.Union[str, t.List[str]] = None,
+        fp_type: str = 'morgan_count',
+        **kwargs
+    ) -> None:
+        super().__init__(
+            csv_file,
+            splitter=splitter,
+            smiles_col=smiles_cols,
+            target_cols=target_cols,
+            num_classes=num_classes,
+            target_transform=target_transform,
+            **kwargs
+        )
+        self.smiles_cols = smiles_cols 
+        assert fp_type in get_available_features_generators()
+        self.fp_type = fp_type
+        self.fp_generator = get_features_generator(self.fp_type)
+        self.fp_numbits = FP_BITS_DICT[self.fp_type]
+        self.missing_labels = missing_labels
+    
+    def __getitem__(self, index):
+        smiles_list = self.df.loc[index][self.smiles_cols].tolist()
+
+        fingerprint_list = list(
+            map(
+                lambda x: torch.LongTensor(self.fp_generator(Chem.MolFromSmiles(x))),
+                smiles_list
+            )
+        )
+        if any(self.target_cols):
+            target_list = self.df.loc[index][self.target_cols].tolist()
+            
+            # process with missing label
+            final_targets = []
+            for target, missing_label in zip(target_list, self.missing_labels):
+                if missing_label is not None and target == missing_label:
+                    final_targets.append(float('nan'))
+                else:
+                    final_targets.append(target)
+ 
+            if self.target_transform is None:
+                return fingerprint_list, final_targets 
+            else:
+                # print(final_targets)
+                target_tensors = [
+                    trans(target, num_class, missing_label=float('nan'))
+                    for trans, target, num_class in zip(
+                        self.target_transform,
+                        final_targets,
+                        self.num_classes
+                    )
+                ]
+                # print(target_tensors)
+                return fingerprint_list, target_tensors, final_targets 
+        else:
+            return fingerprint_list
+
+
+class FPRDataset(CSVRDataset):
+    def __init__(
+        self,
+        csv_file: str,
+        splitter: str,
+        smiles_col: str,
+        target_col: str = None,
+        missing_label: str = None,
+        target_transform: t.Union[str, t.List[str]] = None,
+        fp_type: str = 'morgan_count',
+        **kwargs
+    ) -> None:
+        super().__init__(
+            csv_file,
+            splitter=splitter,
+            smiles_col=smiles_col,
+            target_col=target_col,
+            target_transform=target_transform,
+            missing_label=missing_label,
+            **kwargs
+        )
+        assert fp_type in get_available_features_generators()
+        self.fp_type = fp_type
+        self.fp_generator = get_features_generator(self.fp_type)
+        self.fp_numbits = FP_BITS_DICT[self.fp_type]
+        self.missing_label = missing_label
+    
+    def __getitem__(self, index):
+        smiles = self.df.loc[index][self.smiles_col]
+        try:
+            fp = torch.LongTensor(self.fp_generator(Chem.MolFromSmiles(smiles)))
+        except Exception as _:
+            fp = torch.zeros(self.fp_numbits).long()
+
+        if self.target_col is not None: 
+            target = self.df.loc[index][self.target_col]
+            target = (target, )
+            return fp, target
+        else:
+            return fp 

hdl/data/dataset/graph/chiral.py

@@ -0,0 +1,62 @@
+import typing as t
+
+import numpy as np
+import torch
+from torch._C import dtype
+import torch_geometric as tg
+from torch_geometric.data import Dataset 
+
+from hdl.features.graph.featurization import MolGraph
+
+
+class MolDataset(Dataset):
+
+    def __init__(
+        self,
+        smiles: t.List,
+        labels: t.List,
+        chiral_features: bool = False,
+        global_chiral_features: bool = False,
+    ):
+        super(MolDataset, self).__init__()
+
+        # self.split = list(range(len(smiles)))  # fix this
+        # self.smiles = [smiles[i] for i in self.split]
+        # self.labels = [labels[i] for i in self.split]
+        self.smiles = smiles
+        self.labels = labels
+        # self.data_map = {k: v for k, v in zip(range(len(self.smiles)), self.split)}
+        # self.args = args
+        self.chiral_features = chiral_features
+        self.global_chiral_features = global_chiral_features
+
+        self.mean = np.mean(self.labels)
+        self.std = np.std(self.labels)
+
+    def process_key(self, key):
+        smi = self.smiles[key]
+        molgraph = MolGraph(
+            smi,
+            self.chiral_features,
+            self.global_chiral_features
+        )
+        mol = self.molgraph2data(molgraph, key)
+        return mol
+
+    def molgraph2data(self, molgraph, key):
+        data = tg.data.Data()
+        data.x = torch.tensor(molgraph.f_atoms, dtype=torch.float)
+        data.edge_index = torch.tensor(molgraph.edge_index, dtype=torch.long).t().contiguous()
+        data.edge_attr = torch.tensor(molgraph.f_bonds, dtype=torch.float)
+        data.y = torch.tensor([self.labels[key]], dtype=torch.float)
+        data.parity_atoms = torch.tensor(molgraph.parity_atoms, dtype=torch.long)
+        data.parity_bond_index = torch.tensor(molgraph.parity_bond_index, dtype=torch.long)
+        data.smiles = self.smiles[key]
+
+        return data
+
+    def __len__(self):
+        return len(self.smiles)
+
+    def __getitem__(self, key):
+        return self.process_key(key)

hdl/data/dataset/graph/gin.py

@@ -0,0 +1,255 @@
+# import os
+import csv
+import math
+# import time
+import random
+# import networkx as nx
+import numpy as np
+from copy import deepcopy
+import typing as t
+
+import torch
+# import torch.nn.functional as F
+# from torch.utils.data import Dataset, DataLoader
+from torch.utils.data.sampler import SubsetRandomSampler
+# import torchvision.transforms as transforms
+
+# from torch_scatter import scatter
+from torch_geometric.data import Data, Dataset
+from torch_geometric.loader import DataLoader
+
+# import rdkit
+from rdkit import Chem
+# from rdkit.Chem.rdchem import HybridizationType
+from rdkit.Chem.rdchem import BondType as BT
+# from rdkit.Chem import AllChem
+
+from hdl.data.dataset.utils import read_smiles
+
+
+__all__ = [
+    "MoleculeDataset",
+    "MoleculeDatasetWrapper"
+]
+
+
+ATOM_LIST = list(range(1, 119))
+CHIRALITY_LIST = [
+    Chem.rdchem.ChiralType.CHI_UNSPECIFIED,
+    Chem.rdchem.ChiralType.CHI_TETRAHEDRAL_CW,
+    Chem.rdchem.ChiralType.CHI_TETRAHEDRAL_CCW,
+    Chem.rdchem.ChiralType.CHI_OTHER
+]
+BOND_LIST = [
+    BT.SINGLE, 
+    BT.DOUBLE, 
+    BT.TRIPLE, 
+    BT.AROMATIC
+]
+BONDDIR_LIST = [
+    Chem.rdchem.BondDir.NONE,
+    Chem.rdchem.BondDir.ENDUPRIGHT,
+    Chem.rdchem.BondDir.ENDDOWNRIGHT
+]
+
+
+class MoleculeDataset(Dataset):
+    def __init__(
+        self,
+        data_path,
+        file_type: str = 'smi',
+        smi_col_names: t.List = [],
+        y_col_name: str = None, 
+    ):
+        super(Dataset, self).__init__()
+        self.smiles_data = read_smiles(
+            data_path=data_path,
+            file_type=file_type,
+            smi_col_names=smi_col_names,
+            y_col_name=y_col_name
+        )
+        self.smi_col_names = smi_col_names
+        self.y_col_name = y_col_name
+    
+    def __getitem__(
+        self,
+        idx: int
+    ):
+        if any(self.smi_col_names):
+            item = [
+                self.getitem(smiles)
+                for smiles in self.smiles_data[idx][: len(self.smi_col_names)]
+            ]
+            if self.y_col_name is not None:
+                item.append(float(self.smiles_data[idx][-1]))
+            return item
+        else:
+            return self.getitem(self.smiles_data[idx])
+
+    def getitem(self, smiles):
+        mol = Chem.MolFromSmiles(smiles)
+        # mol = Chem.AddHs(mol)
+
+        N = mol.GetNumAtoms()
+        M = mol.GetNumBonds()
+
+        type_idx = []
+        chirality_idx = []
+        atomic_number = []
+        # aromatic = []
+        # sp, sp2, sp3, sp3d = [], [], [], []
+        # num_hs = []
+        for atom in mol.GetAtoms():
+            type_idx.append(ATOM_LIST.index(atom.GetAtomicNum()))
+            chirality_idx.append(CHIRALITY_LIST.index(atom.GetChiralTag()))
+            atomic_number.append(atom.GetAtomicNum())
+            # aromatic.append(1 if atom.GetIsAromatic() else 0)
+            # hybridization = atom.GetHybridization()
+            # sp.append(1 if hybridization == HybridizationType.SP else 0)
+            # sp2.append(1 if hybridization == HybridizationType.SP2 else 0)
+            # sp3.append(1 if hybridization == HybridizationType.SP3 else 0)
+            # sp3d.append(1 if hybridization == HybridizationType.SP3D else 0)
+
+        # z = torch.tensor(atomic_number, dtype=torch.long)
+        x1 = torch.tensor(type_idx, dtype=torch.long).view(-1,1)
+        x2 = torch.tensor(chirality_idx, dtype=torch.long).view(-1,1)
+        x = torch.cat([x1, x2], dim=-1)
+        # x2 = torch.tensor([atomic_number, aromatic, sp, sp2, sp3, sp3d, num_hs],
+        #                     dtype=torch.float).t().contiguous()
+        # x = torch.cat([x1.to(torch.float), x2], dim=-1)
+
+        row, col, edge_feat = [], [], []
+        for bond in mol.GetBonds():
+            start, end = bond.GetBeginAtomIdx(), bond.GetEndAtomIdx()
+            row += [start, end]
+            col += [end, start]
+            # edge_type += 2 * [MOL_BONDS[bond.GetBondType()]]
+            edge_feat.append([
+                BOND_LIST.index(bond.GetBondType()),
+                BONDDIR_LIST.index(bond.GetBondDir())
+            ])
+            edge_feat.append([
+                BOND_LIST.index(bond.GetBondType()),
+                BONDDIR_LIST.index(bond.GetBondDir())
+            ])
+
+        edge_index = torch.tensor([row, col], dtype=torch.long)
+        edge_attr = torch.tensor(np.array(edge_feat), dtype=torch.long)
+
+        # random mask a subgraph of the molecule
+        num_mask_nodes = max([1, math.floor(0.25*N)])
+        num_mask_edges = max([0, math.floor(0.25*M)])
+        mask_nodes_i = random.sample(list(range(N)), num_mask_nodes)
+        mask_nodes_j = random.sample(list(range(N)), num_mask_nodes)
+
+        mask_edges_i_single = random.sample(list(range(M)), num_mask_edges)
+        mask_edges_j_single = random.sample(list(range(M)), num_mask_edges)
+        mask_edges_i = [2*i for i in mask_edges_i_single] + [2*i+1 for i in mask_edges_i_single]
+        mask_edges_j = [2*i for i in mask_edges_j_single] + [2*i+1 for i in mask_edges_j_single]
+
+        x_i = deepcopy(x)
+        for atom_idx in mask_nodes_i:
+            x_i[atom_idx,:] = torch.tensor([len(ATOM_LIST), 0])
+        edge_index_i = torch.zeros((2, 2*(M-num_mask_edges)), dtype=torch.long)
+        edge_attr_i = torch.zeros((2*(M-num_mask_edges), 2), dtype=torch.long)
+        count = 0
+        for bond_idx in range(2*M):
+            if bond_idx not in mask_edges_i:
+                edge_index_i[:,count] = edge_index[:,bond_idx]
+                edge_attr_i[count,:] = edge_attr[bond_idx,:]
+                count += 1
+        data_i = Data(x=x_i, edge_index=edge_index_i, edge_attr=edge_attr_i)
+
+        x_j = deepcopy(x)
+        for atom_idx in mask_nodes_j:
+            x_j[atom_idx,:] = torch.tensor([len(ATOM_LIST), 0])
+        edge_index_j = torch.zeros((2, 2*(M-num_mask_edges)), dtype=torch.long)
+        edge_attr_j = torch.zeros((2*(M-num_mask_edges), 2), dtype=torch.long)
+        count = 0
+        for bond_idx in range(2*M):
+            if bond_idx not in mask_edges_j:
+                edge_index_j[:,count] = edge_index[:,bond_idx]
+                edge_attr_j[count,:] = edge_attr[bond_idx,:]
+                count += 1
+        data_j = Data(x=x_j, edge_index=edge_index_j, edge_attr=edge_attr_j)
+        
+        return data_i, data_j
+
+    def __len__(self):
+        return len(self.smiles_data)
+
+
+class MoleculeDatasetWrapper(object):
+    def __init__(
+        self,
+        batch_size,
+        num_workers,
+        valid_size,
+        data_path,
+        file_type: str = 'smi',
+        smi_col_names: t.List = [], 
+        y_col_name: str = None,
+    ):
+        super(object, self).__init__()
+        self.data_path = data_path
+        self.batch_size = batch_size
+        self.num_workers = num_workers
+        self.valid_size = valid_size
+        self.file_type = file_type
+        self.smi_col_names = smi_col_names
+        self.y_col_name = y_col_name
+
+    def get_data_loaders(self):
+        train_dataset = MoleculeDataset(data_path=self.data_path)
+        train_loader, valid_loader = self.get_train_validation_data_loaders(train_dataset)
+        return train_loader, valid_loader
+    
+    def get_test_loader(self, shuffle=False):
+        test_dataset = MoleculeDataset(
+            data_path=self.data_path,
+            file_type=self.file_type,
+            smi_col_names=self.smi_col_names,
+            y_col_name=self.y_col_name
+        )
+        test_loader = self.get_test_data_loader(
+            test_dataset,
+            shuffle=shuffle
+        )
+        return test_loader
+
+    def get_train_validation_data_loaders(self, train_dataset):
+        # obtain training indices that will be used for validation
+        num_train = len(train_dataset)
+        indices = list(range(num_train))
+        np.random.shuffle(indices)
+
+        split = int(np.floor(self.valid_size * num_train))
+        train_idx, valid_idx = indices[split:], indices[:split]
+
+        # define samplers for obtaining training and validation batches
+        train_sampler = SubsetRandomSampler(train_idx)
+        valid_sampler = SubsetRandomSampler(valid_idx)
+
+        train_loader = DataLoader(train_dataset, batch_size=self.batch_size, sampler=train_sampler,
+                                  num_workers=self.num_workers, drop_last=True)
+
+        valid_loader = DataLoader(train_dataset, batch_size=self.batch_size, sampler=valid_sampler,
+                                  num_workers=self.num_workers, drop_last=True)
+
+        return train_loader, valid_loader
+
+    def get_test_data_loader(
+        self,
+        test_dataset,
+        shuffle=False
+    ):
+        # num_test = len(test_dataset)
+        # indices = list(range(num_test))
+        test_loader = DataLoader(
+            test_dataset,
+            batch_size=self.batch_size,
+            num_workers=self.num_workers,
+            drop_last=False,
+            shuffle=shuffle
+        )
+        return test_loader