hjxdl 0.0.1__py3-none-any.whl

hdl/data/dataset/graph/molnet.py

@@ -0,0 +1,362 @@
+# import os
+import copy
+import os.path as osp
+import re
+from itertools import product, repeat
+from typing import List, Tuple, Dict, Optional, Union
+from collections import Sequence
+
+import torch
+from torch import Tensor
+from torch_geometric.data import (
+    InMemoryDataset, Data 
+)
+from torch_geometric.data.dataset import Dataset, IndexType
+import numpy as np
+
+try:
+    from rdkit import Chem
+except ImportError:
+    Chem = None
+
+from jupyfuncs.pbar import tqdm
+
+
+x_map = {
+    'atomic_num':
+    list(range(0, 119)),
+    'chirality': [
+        'CHI_UNSPECIFIED',
+        'CHI_TETRAHEDRAL_CW',
+        'CHI_TETRAHEDRAL_CCW',
+        'CHI_OTHER',
+    ],
+    'degree':
+    list(range(0, 11)),
+    'formal_charge':
+    list(range(-5, 7)),
+    'num_hs':
+    list(range(0, 9)),
+    'num_radical_electrons':
+    list(range(0, 5)),
+    'hybridization': [
+        'UNSPECIFIED',
+        'S',
+        'SP',
+        'SP2',
+        'SP3',
+        'SP3D',
+        'SP3D2',
+        'OTHER',
+    ],
+    'is_aromatic': [False, True],
+    'is_in_ring': [False, True],
+}
+
+e_map = {
+    'bond_type': [
+        'misc',
+        'SINGLE',
+        'DOUBLE',
+        'TRIPLE',
+        'AROMATIC',
+    ],
+    'stereo': [
+        'STEREONONE',
+        'STEREOZ',
+        'STEREOE',
+        'STEREOCIS',
+        'STEREOTRANS',
+        'STEREOANY',
+    ],
+    'is_conjugated': [False, True],
+}
+
+
+class MoleculeNet(torch.utils.data.Dataset):
+    r"""The `MoleculeNet <http://moleculenet.ai/datasets-1>`_ benchmark
+    collection  from the `"MoleculeNet: A Benchmark for Molecular Machine
+    Learning" <https://arxiv.org/abs/1703.00564>`_ paper, containing datasets
+    from physical chemistry, biophysics and physiology.
+    All datasets come with the additional node and edge features introduced by
+    the `Open Graph Benchmark <https://ogb.stanford.edu/docs/graphprop/>`_.
+    Args:
+        root (string): Root directory where the dataset should be saved.
+        name (string): The name of the dataset (:obj:`"ESOL"`,
+            :obj:`"FreeSolv"`, :obj:`"Lipo"`, :obj:`"PCBA"`, :obj:`"MUV"`,
+            :obj:`"HIV"`, :obj:`"BACE"`, :obj:`"BBPB"`, :obj:`"Tox21"`,
+            :obj:`"ToxCast"`, :obj:`"SIDER"`, :obj:`"ClinTox"`).
+        transform (callable, optional): A function/transform that takes in an
+            :obj:`torch_geometric.data.Data` object and returns a transformed
+            version. The data object will be transformed before every access.
+            (default: :obj:`None`)
+        pre_transform (callable, optional): A function/transform that takes in
+            an :obj:`torch_geometric.data.Data` object and returns a
+            transformed version. The data object will be transformed before
+            being saved to disk. (default: :obj:`None`)
+        pre_filter (callable, optional): A function that takes in an
+            :obj:`torch_geometric.data.Data` object and returns a boolean
+            value, indicating whether the data object should be included in the
+            final dataset. (default: :obj:`None`)
+    """
+
+    def __init__(self, root, file_type='smi_in_csv', 
+                 transform=None, pre_transform=None,
+                 pre_filter=None):
+
+        self.file_type = 'smi_in_csv'
+
+        if Chem is None:
+            raise ImportError('`MoleculeNet` requires `rdkit`.')
+
+        self.transform = transform
+        self.pre_transform = pre_transform
+        self.pre_filter = pre_filter
+
+        # self.name = name.lower()
+        # assert self.name in self.names.keys()
+        # super(MoleculeNet, self).__init__(root, transform, pre_transform,
+        #                                   pre_filter)
+        self.root_dir = ''
+        self.processed_file = osp.join(
+            self.root_dir, 'processed.pt'
+        )
+        self.process()
+        self.data, self.slices = torch.load(self.processed_file)
+
+    def process(self):
+        with open(self.raw_paths[0], 'r') as f:
+            dataset = f.read().split('\n')[1:-1]
+            dataset = [x for x in dataset if len(x) > 0]  # Filter empty lines.
+
+        data_list = []
+        for line in tqdm(dataset):
+            line = re.sub(r'\".*\"', '', line)  # Replace ".*" strings.
+            line = line.split(',')
+
+            smiles = line[self.names[self.name][3]]
+            ys = line[self.names[self.name][4]]
+            ys = ys if isinstance(ys, list) else [ys]
+
+            ys = [float(y) if len(y) > 0 else float('NaN') for y in ys]
+            y = torch.tensor(ys, dtype=torch.float).view(1, -1)
+
+            mol = Chem.MolFromSmiles(smiles)
+            if mol is None:
+                continue
+
+            xs = []
+            for atom in mol.GetAtoms():
+                x = []
+                x.append(x_map['atomic_num'].index(atom.GetAtomicNum()))
+                x.append(x_map['chirality'].index(str(atom.GetChiralTag())))
+                x.append(x_map['degree'].index(atom.GetTotalDegree()))
+                x.append(x_map['formal_charge'].index(atom.GetFormalCharge()))
+                x.append(x_map['num_hs'].index(atom.GetTotalNumHs()))
+                x.append(x_map['num_radical_electrons'].index(
+                    atom.GetNumRadicalElectrons()))
+                x.append(x_map['hybridization'].index(
+                    str(atom.GetHybridization())))
+                x.append(x_map['is_aromatic'].index(atom.GetIsAromatic()))
+                x.append(x_map['is_in_ring'].index(atom.IsInRing()))
+                xs.append(x)
+
+            x = torch.tensor(xs, dtype=torch.long).view(-1, 9)
+
+            edge_indices, edge_attrs = [], []
+            for bond in mol.GetBonds():
+                i = bond.GetBeginAtomIdx()
+                j = bond.GetEndAtomIdx()
+
+                e = []
+                e.append(e_map['bond_type'].index(str(bond.GetBondType())))
+                e.append(e_map['stereo'].index(str(bond.GetStereo())))
+                e.append(e_map['is_conjugated'].index(bond.GetIsConjugated()))
+
+                edge_indices += [[i, j], [j, i]]
+                edge_attrs += [e, e]
+
+            edge_index = torch.tensor(edge_indices)
+            edge_index = edge_index.t().to(torch.long).view(2, -1)
+            edge_attr = torch.tensor(edge_attrs, dtype=torch.long).view(-1, 3)
+
+            # Sort indices.
+            if edge_index.numel() > 0:
+                perm = (edge_index[0] * x.size(0) + edge_index[1]).argsort()
+                edge_index, edge_attr = edge_index[:, perm], edge_attr[perm]
+
+            data = Data(x=x, edge_index=edge_index, edge_attr=edge_attr, y=y,
+                        smiles=smiles)
+
+            if self.pre_filter is not None and not self.pre_filter(data):
+                continue
+
+            if self.pre_transform is not None:
+                data = self.pre_transform(data)
+
+            data_list.append(data)
+
+        torch.save(
+            self.collate(data_list),
+            self.processed_file
+        )
+
+    def collate(data_list: List[Data]) -> Tuple[Data, Dict[str, Tensor]]:
+        r"""Collates a python list of data objects to the internal storage
+        format of :class:`torch_geometric.data.InMemoryDataset`."""
+        keys = data_list[0].keys
+        data = data_list[0].__class__()
+
+        for key in keys:
+            data[key] = []
+        slices = {key: [0] for key in keys}
+
+        for item, key in product(data_list, keys):
+            data[key].append(item[key])
+            if isinstance(item[key], Tensor) and item[key].dim() > 0:
+                cat_dim = item.__cat_dim__(key, item[key])
+                cat_dim = 0 if cat_dim is None else cat_dim
+                s = slices[key][-1] + item[key].size(cat_dim)
+            else:
+                s = slices[key][-1] + 1
+            slices[key].append(s)
+
+    def __repr__(self):
+        return 'MoleculeNet ({})'.format(len(self))
+    
+    def copy(self, idx: Optional[IndexType] = None):
+        if idx is None:
+            data_list = [self.get(i) for i in range(len(self))]
+        else:
+            data_list = [self.get(i) for i in self.index_select(idx).indices()]
+
+        dataset = copy.copy(self)
+        dataset._indices = None
+        dataset._data_list = data_list
+        dataset.data, dataset.slices = self.collate(data_list)
+        return dataset
+
+    @property
+    def num_classes(self) -> int:
+        r"""The number of classes in the dataset."""
+        y = self.data.y
+        if y is None:
+            return 0
+        elif y.numel() == y.size(0) and not torch.is_floating_point(y):
+            return int(self.data.y.max()) + 1
+        elif y.numel() == y.size(0) and torch.is_floating_point(y):
+            return torch.unique(y).numel()
+        else:
+            return self.data.y.size(-1)
+
+    def len(self) -> int:
+        for item in self.slices.values():
+            return len(item) - 1
+        return 0
+
+    def get(self, idx: int) -> Data:
+        if hasattr(self, '_data_list'):
+            if self._data_list is None:
+                self._data_list = self.len() * [None]
+            else:
+                data = self._data_list[idx]
+                if data is not None:
+                    return copy.copy(data)
+
+        data = self.data.__class__()
+        if hasattr(self.data, '__num_nodes__'):
+            data.num_nodes = self.data.__num_nodes__[idx]
+
+        for key in self.data.keys:
+            item, slices = self.data[key], self.slices[key]
+            start, end = slices[idx].item(), slices[idx + 1].item()
+            if torch.is_tensor(item):
+                s = list(repeat(slice(None), item.dim()))
+                cat_dim = self.data.__cat_dim__(key, item)
+                if cat_dim is None:
+                    cat_dim = 0
+                s[cat_dim] = slice(start, end)
+            elif start + 1 == end:
+                s = slices[start]
+            else:
+                s = slice(start, end)
+            data[key] = item[s]
+
+        if hasattr(self, '_data_list'):
+            self._data_list[idx] = copy.copy(data)
+
+        return data
+
+    def __len__(self) -> int:
+        r"""The number of examples in the dataset."""
+        return len(self.indices())
+
+    def __getitem__(
+        self,
+        idx: Union[int, np.integer, IndexType],
+    ) -> Union['Dataset', Data]:
+        r"""In case :obj:`idx` is of type integer, will return the data object
+        at index :obj:`idx` (and transforms it in case :obj:`transform` is
+        present).
+        In case :obj:`idx` is a slicing object, *e.g.*, :obj:`[2:5]`, a list, a
+        tuple, a PyTorch :obj:`LongTensor` or a :obj:`BoolTensor`, or a numpy
+        :obj:`np.array`, will return a subset of the dataset at the specified
+        indices."""
+        if (isinstance(idx, (int, np.integer))
+                or (isinstance(idx, Tensor) and idx.dim() == 0)
+                or (isinstance(idx, np.ndarray) and np.isscalar(idx))):
+
+            data = self.get(self.indices()[idx])
+            data = data if self.transform is None else self.transform(data)
+            return data
+
+        else:
+            return self.index_select(idx)
+
+    def index_select(self, idx: IndexType) -> 'Dataset':
+        indices = self.indices()
+
+        if isinstance(idx, slice):
+            indices = indices[idx]
+
+        elif isinstance(idx, Tensor) and idx.dtype == torch.long:
+            return self.index_select(idx.flatten().tolist())
+
+        elif isinstance(idx, Tensor) and idx.dtype == torch.bool:
+            idx = idx.flatten().nonzero(as_tuple=False)
+            return self.index_select(idx.flatten().tolist())
+
+        elif isinstance(idx, np.ndarray) and idx.dtype == np.int64:
+            return self.index_select(idx.flatten().tolist())
+
+        elif isinstance(idx, np.ndarray) and idx.dtype == np.bool:
+            idx = idx.flatten().nonzero()[0]
+            return self.index_select(idx.flatten().tolist())
+
+        elif isinstance(idx, Sequence) and not isinstance(idx, str):
+            indices = [indices[i] for i in idx]
+
+        else:
+            raise IndexError(
+                f"Only integers, slices (':'), list, tuples, torch.tensor and "
+                f"np.ndarray of dtype long or bool are valid indices (got "
+                f"'{type(idx).__name__}')")
+
+        dataset = copy.copy(self)
+        dataset._indices = indices
+        return dataset
+
+    def shuffle(
+        self,
+        return_perm: bool = False,
+    ) -> Union['Dataset', Tuple['Dataset', Tensor]]:
+        r"""Randomly shuffles the examples in the dataset.
+
+        Args:
+            return_perm (bool, optional): If set to :obj:`True`, will return
+                the random permutation used to shuffle the dataset in addition.
+                (default: :obj:`False`)
+        """
+        perm = torch.randperm(len(self))
+        dataset = self.index_select(perm)
+        return (dataset, perm) if return_perm is True else dataset

hdl/data/dataset/loaders/chiral_graph.py

@@ -0,0 +1,71 @@
+import typing as t
+
+import pandas as pd
+import numpy as np
+from torch_geometric.loader import DataLoader
+
+from hdl.data.dataset.graph.chiral import MolDataset
+from hdl.data.dataset.samplers.chiral import StereoSampler 
+from hdl.data.dataset.loaders.spliter import split_data
+
+
+def get_chiralgraph_loader(
+    data_path: str = None,
+    smiles_list: t.List = [],
+    label_list: t.List = [], 
+    batch_size: int = 1,
+    shuffle: bool = False,
+    smiles_col: str = 'SMILES',
+    label_col: str = 'label',
+    num_workers: int = 10,
+    shuffle_pairs: bool = False,
+    chiral_features: bool = True,
+    global_chiral_features: bool = True 
+):
+
+    if data_path is not None:
+        data_df = pd.read_csv(data_path)
+
+        # smiles = data_df.iloc[:, 0].values
+        # labels = data_df.iloc[:, 1].values.astype(np.float32)
+        smiles = data_df[smiles_col].tolist()
+        labels = data_df[label_col].to_numpy()
+    else:
+        smiles = smiles_list
+        labels = np.array(label_list)
+    
+    dataset = MolDataset(
+        smiles=smiles,
+        labels=labels,
+        chiral_features=chiral_features,
+        global_chiral_features=global_chiral_features
+    )
+    loader = DataLoader(
+        dataset=dataset,
+        batch_size=batch_size,
+        shuffle=shuffle,
+        num_workers=num_workers,
+        pin_memory=True,
+        sampler=StereoSampler(dataset) if shuffle_pairs else None)
+    return loader, dataset
+    
+    split_loader_list = []
+    split_data_list = split_data(smiles, labels, split_type="random")
+    for split_smiles, split_labels in split_data_list:
+        dataset = MolDataset(
+            smiles=split_smiles,
+            labels=split_labels,
+            chiral_features=chiral_features,
+            global_chiral_features=global_chiral_features,
+        )
+    
+    # train_dataset = dataset
+        loader = DataLoader(dataset=dataset,
+                            batch_size=batch_size,
+                            shuffle=shuffle,
+                            num_workers=num_workers,
+                            pin_memory=True,
+                            sampler=StereoSampler(dataset) if shuffle_pairs else None)
+        split_loader_list.append(loader)
+
+    return split_loader_list, dataset

hdl/data/dataset/loaders/collate_funcs/fp.py

@@ -0,0 +1,56 @@
+r""""Contains definitions of the methods used by the _BaseDataLoaderIter workers to
+collate samples fetched from dataset into Tensor(s).
+
+These **needs** to be in global scope since Py2 doesn't support serializing
+static methods.
+"""
+import typing as t
+
+import numpy as np
+import pandas as pd
+import torch
+
+
+int_types = (
+    int,
+    np.int32,
+    np.int64,
+    pd.Int16Dtype,
+    pd.Int32Dtype,
+    pd.Int64Dtype,
+    # torch.int32
+)
+
+
+def fp_collate(batch):
+    transposed = list(zip(*batch))
+
+    # fps
+    fps = list(zip(*transposed[0]))
+    fps = [torch.vstack(fp).float() for fp in fps]
+    if len(transposed) == 1:
+        return fps
+
+    # target_list
+    targets = list(zip(*transposed[-1]))
+    targets_list = []
+    for target_labels in targets:
+        if not isinstance(target_labels[0], t.Iterable):
+            target_labels = torch.Tensor(target_labels)
+        else:
+            target_labels = list(target_labels)
+        # if isinstance(target_labels[0], int_types):
+        #     target_labels = torch.LongTensor(target_labels)
+        targets_list.append(target_labels) 
+ 
+    # target_tensors
+    if len(transposed) == 3:
+        target_tensors = list(zip(*transposed[1])) 
+        target_tensors = [
+            torch.vstack(target_tensor).float()
+            for target_tensor in target_tensors
+        ]
+    
+        return fps, target_tensors, targets_list
+    else:
+        return fps, targets, targets_list 

hdl/data/dataset/loaders/collate_funcs/rxn.py

@@ -0,0 +1,40 @@
+from rxnfp.tokenization import (
+    SmilesTokenizer,
+    convert_reaction_to_valid_features_batch,
+)
+import torch
+import numpy as np
+import pkg_resources
+
+
+__all__ = [
+    'collate_rxn',
+]
+
+
+def collate_rxn(
+    rxn_list,
+    labels,
+    vocab_path: str = None,
+    max_len: int = 512
+):
+    if vocab_path is None:
+        vocab_path = pkg_resources.resource_filename(
+            "rxnfp",
+            "models/transformers/bert_ft/vocab.txt"
+        )
+    tokenizer = SmilesTokenizer(
+        vocab_path, max_len=max_len
+    )
+
+    feats = convert_reaction_to_valid_features_batch(
+        rxn_list,
+        tokenizer
+    )
+    X = [
+        torch.tensor(feats.input_ids.astype(np.int64)),
+        torch.tensor(feats.input_mask.astype(np.int64)),
+        torch.tensor(feats.segment_ids.astype(np.int64))
+    ]
+    y = torch.LongTensor(labels)
+    return X, y

hdl/data/dataset/loaders/general.py

@@ -0,0 +1,23 @@
+import typing as t
+
+import torch.utils.data as tud
+
+from hdl.data.dataset.loaders.collate_funcs.fp import fp_collate
+
+
+class Loader(tud.DataLoader):
+    def __init__(
+        self,
+        dataset,
+        batch_size: int = 128,
+        shuffle: bool = True,
+        num_workers: int = 12,
+        collate_fn: t.Callable = fp_collate
+    ):
+        super().__init__(
+            dataset,
+            batch_size=batch_size,
+            shuffle=shuffle,
+            num_workers=num_workers,
+            collate_fn=collate_fn
+        )

hdl/data/dataset/loaders/spliter.py

@@ -0,0 +1,86 @@
+from typing import DefaultDict, Tuple
+from random import Random
+from collections import defaultdict
+from rdkit import Chem
+from rdkit.Chem.Scaffolds import MurckoScaffold
+from hdl.data.dataset.graph.chiral import MolDataset
+
+
+def split_data(
+    smis: Tuple[str],
+    labels: Tuple,
+    split_type: str = "random",
+    sizes: Tuple[float, float, float] = (0.8, 0.2, 0.0),
+    seed: int = 999,
+    num_folds: int = 1,
+    balanced: bool = True,
+    args=None,
+) -> Tuple[Tuple[str], Tuple[str], Tuple[str]]:
+    random = Random(seed)
+
+    if split_type == "random":
+        indices = list(range(len(smis)))
+        random.shuffle(indices)
+
+        train_size = int(sizes[0] * len(smis))
+        train_val_size = int((sizes[0] + sizes[1]) * len(smis))
+        train = [
+            [smis[i] for i in indices[:train_size]],
+            [labels[i] for i in indices[:train_size]],
+        ]
+        val = [
+            [smis[i] for i in indices[train_size:train_val_size]],
+            [labels[i] for i in indices[train_size:train_val_size]],
+        ]
+        test = [
+            [smis[i] for i in indices[train_val_size:]],
+            [labels[i] for i in indices[train_val_size:]],
+        ]
+    elif split_type == "scaffold_balanced":
+        train_size, val_size, test_size = (
+            sizes[0] * len(data),
+            sizes[1] * len(data),
+            sizes[2] * len(data),
+        )
+        train, val, test = [], [], []
+        train_scaffold_count, val_scaffold_count, test_scaffold_count = 0, 0, 0
+        scaffold_to_indices = defaultdict(set)
+        rdmols = [Chem.MolFromSmiles(s) for s in smis]
+        for i, rdmol in enumerate(rdmols):
+            scaffold = MurckoScaffold.MurckoScaffoldSmiles(
+                mol=rdmol, includeChirality=False
+            )
+            scaffold_to_indices[scaffold].add(i)
+        if balanced:
+            index_sets = list(scaffold_to_indices.values())
+            big_index_sets = []
+            small_index_sets = []
+            for index_set in index_sets:
+                if len(index_set) > val_size / 2 or len(index_set) > test_size / 2:
+                    big_index_sets.append(index_set)
+                else:
+                    small_index_sets.append(index_set)
+            random.seed(seed)
+            random.shuffle(big_index_sets)
+            random.shuffle(small_index_sets)
+            index_sets = big_index_sets + small_index_sets
+        else:
+            index_sets = sorted(
+                list(scaffold_to_indices.values()),
+                key=lambda index_set: len(index_set),
+                reverse=True,
+            )
+        for index_set in index_sets:
+            if len(train) + len(index_set) <= train_size:
+                train += index_set
+                train_scaffold_count += 1
+            elif len(val) + len(index_set) <= val_size:
+                val += index_set
+                val_scaffold_count += 1
+            else:
+                test += index_set
+                test_scaffold_count += 1
+        train = [smis[i] for i in train]
+        val = [smis[i] for i in val]
+        test = [smis[i] for i in test]
+    return train, val, test

hdl/data/dataset/samplers/chiral.py

@@ -0,0 +1,19 @@
+from itertools import chain
+
+import numpy as np
+from torch.utils.data.sampler import Sampler
+
+
+class StereoSampler(Sampler):
+
+    def __init__(self, data_source):
+        self.data_source = data_source
+
+    def __iter__(self):
+        groups = [[i, i + 1] for i in range(0, len(self.data_source), 2)]
+        np.random.shuffle(groups)
+        indices = list(chain(*groups))
+        return iter(indices)
+
+    def __len__(self):
+        return len(self.data_source)