D4CMPP2 0.4.0__py3-none-any.whl

D4CMPP2/src/DataManager/Dataset/ISAGraphDataset.py

@@ -0,0 +1,384 @@
+import torch
+import warnings
+from torch.utils.data import Dataset
+from torch_geometric.data import Batch
+import numpy as np
+from ..contracts import ISA_BATCH_CONTRACT, LEGACY_ISA_BATCH_CONTRACT
+
+class ISAGraphDataset_legacy(Dataset):
+    batch_contract = LEGACY_ISA_BATCH_CONTRACT
+    def __init__(self, graphs=None, target=None, smiles=None):
+        if graphs is None: return
+        self.graphs = graphs
+        self.r_node = [g['r_nd'].x for g in graphs]
+        self.r2r_edge = [g['r_nd', 'r2r', 'r_nd'].edge_attr for g in graphs]
+        self.i_node = [g['i_nd'].x for g in graphs]
+        self.d_node = [g['d_nd'].x for g in graphs]
+        self.d2d_edge = [g['d_nd', 'd2d', 'd_nd'].edge_attr for g in graphs]
+        self.target = torch.tensor(target).float()
+        self.smiles = smiles
+
+
+    def __len__(self):
+        return len(self.graphs)
+
+    def __getitem__(self, idx):
+        return self.graphs[idx], self.r_node[idx], self.r2r_edge[idx], self.i_node[idx], self.d_node[idx], self.d2d_edge[idx], self.target[idx], self.smiles[idx]
+    
+    def reload(self, data):
+        self.graphs, self.r_node, self.r2r_edge, self.i_node, self.d_node, self.d2d_edge, self.target, self.smiles = data
+        self.target = torch.stack(self.target)
+        if self.target.dim() == 1:
+            self.target = self.target.unsqueeze(-1)
+
+    # def subDataset(self, idx):
+    #     self.graphs = [self.graphs[i] for i in idx]
+    #     self.r_node = [self.r_node[i] for i in idx]
+    #     self.r2r_edge = [self.r2r_edge[i] for i in idx]
+    #     self.i_node = [self.i_node[i] for i in idx]
+    #     self.d_node = [self.d_node[i] for i in idx]
+    #     self.d2d_edge = [self.d2d_edge[i] for i in idx]
+    #     self.target = self.target[np.array(idx, dtype=int)]
+    #     self.smiles = [self.smiles[i] for i in idx]
+
+    def get_subDataset(self, idx):
+        graphs = [self.graphs[i] for i in idx]
+        r_node = [self.r_node[i] for i in idx]
+        r2r_edge = [self.r2r_edge[i] for i in idx]
+        i_node = [self.i_node[i] for i in idx]
+        d_node = [self.d_node[i] for i in idx]
+        d2d_edge = [self.d2d_edge[i] for i in idx]
+        target = [self.target[i] for i in idx]
+        smiles = [self.smiles[i] for i in idx]
+        
+        dataset = ISAGraphDataset_legacy()
+        dataset.reload((graphs, r_node, r2r_edge, i_node, d_node, d2d_edge, target, smiles))
+        return dataset
+
+    @staticmethod
+    def collate(samples):
+        graphs, r_node, r2r_edge, i_node, d_node, d2d_edge, target, smiles = map(list, zip(*samples))
+        batched_graph = Batch.from_data_list(graphs)
+        return batched_graph, torch.concat(r_node,dim=0), torch.concat(r2r_edge,dim=0), torch.concat(i_node,dim=0), torch.concat(d_node,dim=0), torch.concat(d2d_edge,dim=0), torch.stack(target,dim=0), smiles
+    
+    @staticmethod
+    def unwrapper(graph, r_node, r2r_edge, i_node, d_node, d2d_edge, target, smiles, device='cpu'):
+        graph = graph.to(device=device)
+        r_node = r_node.float().to(device=device)
+        r2r_edge = r2r_edge.float().to(device=device)
+        i_node = i_node.float().to(device=device)
+        d_node = d_node.float().to(device=device)
+        d2d_edge = d2d_edge.float().to(device=device)
+        target = target.float().to(device=device)
+        return {'graph':graph, 'r_node':r_node, 'r_edge':r2r_edge, 'i_node':i_node, 'd_node':d_node, 'd_edge':d2d_edge, 'target':target, 'smiles':smiles}
+    
+    
+class ISAGraphDataset(ISAGraphDataset_legacy):
+    batch_contract = ISA_BATCH_CONTRACT
+    def __init__(self, graphs=None, target=None, smiles=None, numeric_inputs=None, row_indices=None):
+        if graphs is None: 
+            self.graphs = {}
+            self.target = None
+            self.smiles = {}
+            self.data_keys = []
+            return
+
+
+        for key in graphs:
+            setattr(self, key + '_graphs', graphs[key])
+            try:
+                setattr(self, key + '_r_node', [g['r_nd'].x for g in graphs[key]])
+            except KeyError:
+                warnings.warn(
+                    f"Graph key {key!r} has no 'r_nd' node type; using zero features.",
+                    RuntimeWarning,
+                    stacklevel=2,
+                )
+                setattr(self, key + '_r_node', [torch.zeros((g['r_nd'].num_nodes, 1)) for g in graphs[key]])
+
+            try:
+                setattr(self, key + '_r2r_edge', [g['r_nd', 'r2r', 'r_nd'].edge_attr for g in graphs[key]])
+            except KeyError:
+                warnings.warn(
+                    f"Graph key {key!r} has no 'r2r' edge type; using zero features.",
+                    RuntimeWarning,
+                    stacklevel=2,
+                )
+                setattr(self, key + '_r2r_edge', [torch.zeros((g['r_nd', 'r2r', 'r_nd'].num_edges, 1)) for g in graphs[key]])
+            # if 'i_nd' in graphs[key][0].nodes:
+            #     setattr(self, key + '_i_node', [g.nodes['i_nd'].data['f'] for g in graphs[key]])
+            # else:
+            #     setattr(self, key + '_i_node', [torch.zeros((g.num_nodes(), 0)) for g in graphs[key]])
+            try:
+                setattr(self, key + '_i2i_edge', [g['i_nd', 'i2i', 'i_nd'].edge_attr for g in graphs[key]])
+            except KeyError:
+                warnings.warn(
+                    f"Graph key {key!r} has no 'i2i' edge type; using zero features.",
+                    RuntimeWarning,
+                    stacklevel=2,
+                )
+                setattr(self, key + '_i2i_edge', [torch.zeros((g['i_nd', 'i2i', 'i_nd'].num_edges, 1)) for g in graphs[key]])
+
+
+
+            # if 'i2i' in graphs[key][0].nodes:
+            #     setattr(self, key + '_i_node', [g.nodes['i_nd'].data['f'] for g in graphs[key]])
+            # else:
+            #     setattr(self, key + '_i_node', [torch.zeros((g.num_nodes(), 0)) for g in graphs[key]])
+            try:
+                setattr(self, key + '_i_node', [g['i_nd'].x for g in graphs[key]])
+            except KeyError:
+                warnings.warn(
+                    f"Graph key {key!r} has no 'i_nd' node type; using zero features.",
+                    RuntimeWarning,
+                    stacklevel=2,
+                )
+                setattr(self, key + '_i_node', [torch.zeros((g['i_nd'].num_nodes, 1)) for g in graphs[key]])
+
+            try:
+                setattr(self, key + '_d_node', [g['d_nd'].x for g in graphs[key]])
+            except KeyError:
+                warnings.warn(
+                    f"Graph key {key!r} has no 'd_nd' node type; using zero features.",
+                    RuntimeWarning,
+                    stacklevel=2,
+                )
+                setattr(self, key + '_d_node', [torch.zeros((g['d_nd'].num_nodes, 1)) for g in graphs[key]])
+
+            try:
+                setattr(self, key + '_d2d_edge', [g['d_nd', 'd2d', 'd_nd'].edge_attr for g in graphs[key]])
+            except KeyError:
+                warnings.warn(
+                    f"Graph key {key!r} has no 'd2d' edge type; using zero features.",
+                    RuntimeWarning,
+                    stacklevel=2,
+                )
+                setattr(self, key + '_d2d_edge', [torch.zeros((g['d_nd', 'd2d', 'd_nd'].num_edges, 1)) for g in graphs[key]])
+
+            # if 'd_nd' in graphs[key][0].nodes:
+            #     setattr(self, key + '_d_node', [g.nodes['d_nd'].data['f'] for g in graphs[key]])
+            # else:
+            #     setattr(self, key + '_d_node', [torch.zeros((g.num_nodes(), 0)) for g in graphs[key]])
+
+            # if 'd2d' in graphs[key][0].edges:
+            #     setattr(self, key + '_d2d_edge', [g.edges['d2d'].data['dist'] for g in graphs[key]])
+            # else:
+            #     setattr(self, key + '_d2d_edge', [torch.zeros((g.num_edges(), 0)) for g in graphs[key]])
+
+        if numeric_inputs is not None:
+            for key in numeric_inputs:
+                setattr(self, key + '_var', numeric_inputs[key])
+        else:
+            numeric_inputs = {}
+
+        for key in smiles:
+            if key not in graphs:
+                raise ValueError(f"Key '{key}' in smiles is not found in graphs.")
+            setattr(self, key + '_smiles', smiles[key])
+
+        self.target = torch.tensor(target).float() if target is not None else None
+        self.original_row_index = (
+            torch.as_tensor(row_indices, dtype=torch.long) if row_indices is not None else None
+        )
+        self.data_keys = list(graphs.keys()) + list(numeric_inputs.keys()) + ['target']
+
+    def __len__(self):
+        if getattr(self, 'target', None) is not None:
+            return len(self.target)
+        else:
+            if len(self.data_keys) == 0:
+                return 0
+            key = self.data_keys[0]
+            if hasattr(self, key + '_graphs'):
+                return len(getattr(self, key + '_graphs'))
+            elif hasattr(self, key + '_var'):
+                return len(getattr(self, key + '_var'))
+            elif hasattr(self, key + '_smiles'):
+                return len(getattr(self, key + '_smiles'))
+            else:
+                return 0
+
+    def __getitem__(self, idx):
+        item = {}
+        for key in self.data_keys:
+            if hasattr(self, key + '_graphs'):
+                item[key + '_graphs'] = getattr(self, key + '_graphs')[idx]
+                if hasattr(self, key + '_smiles'):
+                    item[key + '_smiles'] = getattr(self, key + '_smiles')[idx]
+                # if hasattr(self, key + '_node_feature'):
+                #     nf = getattr(self, key + '_node_feature')[idx]
+                #     if type(nf) is not torch.Tensor:
+                #         nf = torch.tensor(nf, dtype=torch.float32)
+                #     item[key + '_node_feature'] = nf
+                # if hasattr(self, key + '_edge_feature'):
+                #     ef = getattr(self, key + '_edge_feature')[idx]
+                #     if type(ef) is not torch.Tensor:
+                #         ef = torch.tensor(ef, dtype=torch.float32)
+                #     item[key + '_edge_feature'] = ef
+                if hasattr(self, key + '_r_node'):
+                    f = getattr(self, key + '_r_node')[idx]
+                    if type(f) is not torch.Tensor:
+                        f = torch.tensor(f, dtype=torch.float32)
+                    item[key + '_r_node'] = f
+                if hasattr(self, key + '_r2r_edge'):
+                    f = getattr(self, key + '_r2r_edge')[idx]
+                    if type(f) is not torch.Tensor:
+                        f = torch.tensor(f, dtype=torch.float32)
+                    item[key + '_r2r_edge'] = f
+                if hasattr(self, key + '_i_node'):
+                    f = getattr(self, key + '_i_node')[idx]
+                    if type(f) is not torch.Tensor:
+                        f = torch.tensor(f, dtype=torch.float32)
+                    item[key + '_i_node'] = f
+                if hasattr(self, key + '_i2i_edge'):
+                    f = getattr(self, key + '_i2i_edge')[idx]
+                    if type(f) is not torch.Tensor:
+                        f = torch.tensor(f, dtype=torch.float32)
+                    item[key + '_i2i_edge'] = f
+                if hasattr(self, key + '_d_node'):
+                    f = getattr(self, key + '_d_node')[idx]
+                    if type(f) is not torch.Tensor:
+                        f = torch.tensor(f, dtype=torch.float32)
+                    item[key + '_d_node'] = f
+                if hasattr(self, key + '_d2d_edge'):
+                    f = getattr(self, key + '_d2d_edge')[idx]
+                    if type(f) is not torch.Tensor:
+                        f = torch.tensor(f, dtype=torch.float32)
+                    item[key + '_d2d_edge'] = f
+
+            elif hasattr(self, key + '_var'):
+                var = getattr(self, key + '_var')[idx]
+                if torch.is_tensor(var):
+                    var = var.float()
+                else:
+                    var = torch.tensor(var, dtype=torch.float32)
+                item[key + '_var'] = var
+            elif hasattr(self, key + '_smiles'):
+                item[key + '_smiles'] = getattr(self, key + '_smiles')[idx]
+        if hasattr(self, 'target') and self.target is not None:
+            item['target'] = self.target[idx]
+        if self.original_row_index is not None:
+            item['original_row_index'] = self.original_row_index[idx]
+        return item
+
+    def reload(self, data):
+        if len(data) == 0:
+            self.target = None
+            self.data_keys = []
+            return
+
+        new_data_keys = []
+        for key in data[0]:
+            values = [d[key] for d in data]
+            if key == 'target':
+                if len(values) == 0:
+                    setattr(self, key, None)
+                elif torch.is_tensor(values[0]):
+                    target = torch.stack(values, dim=0).float()
+                    if target.dim() == 1:
+                        target = target.unsqueeze(-1)
+                    setattr(self, key, target)
+                else:
+                    target = torch.tensor(values, dtype=torch.float32)
+                    if target.dim() == 1:
+                        target = target.unsqueeze(-1)
+                    setattr(self, key, target)
+            elif key == 'original_row_index':
+                setattr(self, key, torch.as_tensor(values, dtype=torch.long))
+            elif key.endswith('_var'):
+                if len(values) == 0:
+                    setattr(self, key, torch.empty((0,), dtype=torch.float32))
+                elif torch.is_tensor(values[0]):
+                    setattr(self, key, torch.stack(values, dim=0).float())
+                else:
+                    setattr(self, key, torch.tensor(values, dtype=torch.float32))
+            else:
+                setattr(self, key, values)
+            if key.endswith('_graphs'):
+                new_data_keys.append(key[:-7])  # Remove '_graphs'
+            elif key.endswith('_var'):
+                new_data_keys.append(key[:-4])  # Remove '_var'
+            elif key == 'target':
+                new_data_keys.append(key)
+
+        self.data_keys = new_data_keys
+
+
+    def subDataset(self, idx):
+        for key in self.data_keys:
+            if hasattr(self, key + '_graphs'):
+                setattr(self, key + '_graphs', [getattr(self, key + '_graphs')[i] for i in idx])
+                # setattr(self, key + '_node_feature', [getattr(self, key + '_node_feature')[i] for i in idx])
+                # setattr(self, key + '_edge_feature', [getattr(self, key + '_edge_feature')[i] for i in idx])
+            if hasattr(self, key + '_r_node'):
+                setattr(self, key + '_r_node', [getattr(self, key + '_r_node')[i] for i in idx])
+            if hasattr(self, key + '_r2r_edge'):
+                setattr(self, key + '_r2r_edge', [getattr(self, key + '_r2r_edge')[i] for i in idx])
+            if hasattr(self, key + '_i_node'):
+                setattr(self, key + '_i_node', [getattr(self, key + '_i_node')[i] for i in idx])
+            if hasattr(self, key + '_i2i_edge'):
+                setattr(self, key + '_i2i_edge', [getattr(self, key + '_i2i_edge')[i] for i in idx])
+            if hasattr(self, key + '_d_node'):
+                setattr(self, key + '_d_node', [getattr(self, key + '_d_node')[i] for i in idx])
+            if hasattr(self, key + '_d2d_edge'):
+                setattr(self, key + '_d2d_edge', [getattr(self, key + '_d2d_edge')[i] for i in idx])
+
+            if hasattr(self, key + '_var'):
+                var = getattr(self, key + '_var')
+                if torch.is_tensor(var):
+                    setattr(self, key + '_var', var[torch.tensor(idx, dtype=torch.long)])
+                else:
+                    setattr(self, key + '_var', var[np.array(idx, dtype=int)])
+            if hasattr(self, key + '_smiles'):
+                setattr(self, key + '_smiles', [getattr(self, key + '_smiles')[i] for i in idx])
+        if self.target is not None:
+            self.target = self.target[np.array(idx, dtype=int)]
+        if self.original_row_index is not None:
+            self.original_row_index = self.original_row_index[torch.as_tensor(idx, dtype=torch.long)]
+
+    def get_subDataset(self, idx):
+        dataset = ISAGraphDataset()
+        dataset.reload([self[i] for i in idx])
+        return dataset
+
+
+    @staticmethod
+    def collate(samples):
+        batched_data = {}
+        for key in samples[0].keys():
+            if key.endswith('_graphs'):
+                batched_data[key] = Batch.from_data_list([s[key] for s in samples])
+            # elif key.endswith('_node_feature'):
+            #     batched_data[key] = torch.concat([s[key] for s in samples], dim=0)
+            # elif key.endswith('_edge_feature'):
+            #     batched_data[key] = torch.concat([s[key] for s in samples], dim=0)
+            elif key.endswith('_r_node') or key.endswith('_i_node') or key.endswith('_d_node'):
+                batched_data[key] = torch.concat([s[key] for s in samples], dim=0)
+            elif key.endswith('_r2r_edge') or key.endswith('_i2i_edge') or key.endswith('_d2d_edge'):
+                batched_data[key] = torch.concat([s[key] for s in samples], dim=0)
+            elif key.endswith('_var'):
+                batched_data[key] = torch.stack([s[key].reshape(-1) for s in samples], dim=0)
+            elif key.endswith('_smiles'):
+                batched_data[key] = [s[key] for s in samples]
+            elif key == 'target':
+                batched_data[key] = torch.stack([s[key].reshape(-1) for s in samples], dim=0)
+            elif key == 'original_row_index':
+                batched_data[key] = torch.stack([s[key].reshape(()) for s in samples], dim=0)
+        return batched_data
+
+    @staticmethod
+    def unwrapper(device='cpu', **batched_data):
+        for key in batched_data.keys():
+            if key.endswith('_graphs'):
+                batched_data[key] = batched_data[key].to(device=device)
+            elif key.endswith('_r_node') or key.endswith('_i_node') or key.endswith('_d_node'):
+                batched_data[key] = batched_data[key].float().to(device=device)
+            elif key.endswith('_r2r_edge') or key.endswith('_i2i_edge') or key.endswith('_d2d_edge'):
+                batched_data[key] = batched_data[key].float().to(device=device)
+            elif key.endswith('_var'):
+                batched_data[key] = batched_data[key].float().to(device=device)
+            elif key == 'target':
+                batched_data[key] = batched_data[key].float().to(device=device)
+            elif key == 'original_row_index':
+                batched_data[key] = batched_data[key].long().to(device=device)
+        return batched_data

D4CMPP2/src/DataManager/GraphGenerator/ISAGraphGenerator.py

@@ -0,0 +1,223 @@
+from rdkit import Chem
+import torch
+import numpy as np
+import traceback
+import warnings
+from collections import deque
+from torch_geometric.data import HeteroData
+
+from .MolGraphGenerator import MolGraphGenerator
+from D4CMPP2.src.utils.featureizer import InvalidAtomError
+from D4CMPP2.src.utils.sculptor import SubgroupSplitter
+
+class ISAGraphGenerator(MolGraphGenerator):
+    def __init__(self, frag_ref=None, sculptor_index=(6,2,0)):
+        self.sculptor = SubgroupSplitter(frag_ref,
+                                        get_index=True,
+                                        split_order=sculptor_index[0],
+                                        combine_rest_order=sculptor_index[1],
+                                        absorb_neighbor_order=sculptor_index[2],
+                                        overlapped_ring_combine=True
+                                )
+        self.r_node_dim = None
+        self.i_node_dim = None
+        self.d_node_dim = None
+        self.r_edge_dim = None
+        self.i_edge_dim = None
+        self.d_edge_dim = None
+        self.node_dim= None
+        self.edge_dim = None
+        super().__init__()
+
+
+        self.set_feature_dim()
+        self.node_dim = self.r_node_dim 
+        self.edge_dim = self.r_edge_dim
+        self.verbose = True
+
+    def set_feature_dim(self):
+        try:
+            get_graph = self.get_graph('FC1CCCCC1CCCOCCC')
+        except InvalidAtomError as e:
+            if self.verbose:
+                warnings.warn(
+                    f"ISA feature-dimension probe encountered an invalid atom: {e}",
+                    RuntimeWarning,
+                    stacklevel=2,
+                )
+
+    def get_graph(self, smi, **kwargs):
+        mol = Chem.MolFromSmiles(smi)
+        if mol is None: 
+            raise Exception("Invalid SMILES: failed to generate mol object")
+        if kwargs.get("explicit_h",False):
+            mol = Chem.AddHs(mol)
+        g = self.generate_graph(mol, **kwargs)
+        atom_feature = self.af(mol)
+        bond_feature = self.bf(mol)
+
+        g['r_nd'].x = torch.tensor(atom_feature).float()
+        if self.r_node_dim is None:
+            self.r_node_dim = atom_feature.shape[1]
+
+        relation = ('r_nd', 'r2r', 'r_nd')
+        if mol.GetNumBonds() == 0:
+            edge_dim = self.r_edge_dim
+            if edge_dim is None:
+                edge_dim = self.bf(Chem.MolFromSmiles("C-C")).shape[1]
+            edge_count = g[relation].edge_index.shape[1]
+            edata = torch.zeros((edge_count, edge_dim), dtype=torch.float32)
+        else:
+            edata = torch.tensor(bond_feature).float()
+            edata = torch.cat([edata,edata],dim=0)
+        g[relation].edge_attr = edata
+        if self.r_edge_dim is None:
+            self.r_edge_dim = edata.shape[1]
+
+        g['i_nd'].x = torch.zeros((g['i_nd'].num_nodes, 1)).float()
+        if self.i_node_dim is None:
+            self.i_node_dim = 1
+        if self.i_edge_dim is None:
+            self.i_edge_dim = 0
+
+        g['d_nd'].x = torch.zeros((g['d_nd'].num_nodes, 1)).float()
+        if self.d_node_dim is None:
+            self.d_node_dim = 1
+        if self.d_edge_dim is None:
+            self.d_edge_dim = 0
+
+
+        return g
+    
+    def generate_sub_graph(self, mol, frags):
+        src, dst = [], []
+        for bond in mol.GetBonds():
+            start, end = bond.GetBeginAtomIdx(), bond.GetEndAtomIdx()
+            for frag in frags:
+                if start in frag:
+                    if end in frag:
+                        src.append(start)
+                        dst.append(end)
+                        break
+                    else:
+                        break
+                elif end in frag:
+                    if start in frag:
+                        src.append(start)
+                        dst.append(end)
+                        break
+                    else:
+                        break
+                    
+        return (src+dst, dst+src)
+        
+
+    def generate_dot_graph(self, mol, frags, max_dist = 4):
+        src, dst = [], []
+        for bond in mol.GetBonds():
+            start, end = bond.GetBeginAtomIdx(), bond.GetEndAtomIdx()
+            for i,frag in enumerate(frags):
+                if start in frag:
+                    start_frag = i
+                if end in frag:
+                    end_frag = i
+            if start_frag != end_frag:
+                src.append(start_frag)
+                dst.append(end_frag)
+        if len(src) == 0:
+            return ([],[]), []
+        adjacency = [[] for _ in frags]
+        for start, end in zip(src, dst):
+            adjacency[start].append(end)
+            adjacency[end].append(start)
+
+        out_src, out_dst, distances = [], [], []
+        for root in range(len(frags)):
+            shortest = [-1] * len(frags)
+            shortest[root] = 0
+            queue = deque([root])
+            while queue:
+                current = queue.popleft()
+                if shortest[current] >= max_dist:
+                    continue
+                for neighbor in adjacency[current]:
+                    if shortest[neighbor] == -1:
+                        shortest[neighbor] = shortest[current] + 1
+                        queue.append(neighbor)
+            for target in range(root + 1, len(frags)):
+                distance = shortest[target]
+                if 0 < distance <= max_dist:
+                    out_src.append(root)
+                    out_dst.append(target)
+                    distances.append(distance)
+
+        directed_distances = distances + distances
+        features = np.zeros((len(directed_distances), max_dist), dtype=np.float32)
+        for row, distance in enumerate(directed_distances):
+            features[row, distance - 1] = 1.0
+        return (out_src + out_dst, out_dst + out_src), features
+
+    def generate_graph(self,mol, **kwargs):
+        max_dist = kwargs.get('max_dist',4)
+        num_atoms = mol.GetNumAtoms()
+        mol_data = self.generate_mol_graph(mol)
+        
+        frag = self.sculptor.fragmentation_with_condition(mol)
+        frag_data = self.generate_sub_graph(mol, frag) 
+
+        i2d_src = []
+        i2d_dst = []
+        for i, f in enumerate(frag):
+            for a in f:
+                i2d_src.append(a)
+                i2d_dst.append(i)
+
+        if len(frag) ==1:
+            dot_data = ([0],[0])
+            dist = np.zeros((1,max_dist))
+        else:
+            dot_data, dist = self.generate_dot_graph(mol, frag, max_dist = max_dist)
+
+        graph_data = {
+            ('r_nd', 'r2r', 'r_nd'): mol_data,
+            ('r_nd', 'r2i', 'i_nd'): (list(range(num_atoms)), list(range(num_atoms))),
+            ('i_nd', 'i2i', 'i_nd'): frag_data,
+            ('i_nd', 'i2d', 'd_nd'): (i2d_src, i2d_dst),
+            ('d_nd', 'd2d', 'd_nd'): dot_data,
+            ('d_nd', 'd2r', 'r_nd'): (i2d_dst, i2d_src),
+
+        }
+
+
+        g = HeteroData()
+        g['r_nd'].num_nodes = num_atoms
+        g['i_nd'].num_nodes = num_atoms
+        g['d_nd'].num_nodes = len(frag)
+        for edge_type, (src, dst) in graph_data.items():
+            g[edge_type].edge_index = torch.tensor([src, dst], dtype=torch.long)
+        g['r_nd', 'r2r', 'r_nd'].edge_attr = torch.empty((len(mol_data[0]), 0))
+        g['i_nd', 'i2i', 'i_nd'].edge_attr = torch.empty((len(frag_data[0]), 0))
+        g['d_nd', 'd2d', 'd_nd'].edge_attr = torch.tensor(dist).float()
+        if self.d_edge_dim is None:
+            self.d_edge_dim = dist.shape[1]
+
+        return g
+
+    def get_empty_graph(self):
+        graph = HeteroData()
+        for node_type, dim in (("r_nd", self.r_node_dim), ("i_nd", 1), ("d_nd", 1)):
+            graph[node_type].x = torch.zeros((0, dim or 1), dtype=torch.float32)
+        relations = (
+            ("r_nd", "r2r", "r_nd"),
+            ("r_nd", "r2i", "i_nd"),
+            ("i_nd", "i2i", "i_nd"),
+            ("i_nd", "i2d", "d_nd"),
+            ("d_nd", "d2d", "d_nd"),
+            ("d_nd", "d2r", "r_nd"),
+        )
+        for relation in relations:
+            graph[relation].edge_index = torch.empty((2, 0), dtype=torch.long)
+        graph["r_nd", "r2r", "r_nd"].edge_attr = torch.zeros((0, self.r_edge_dim or 1))
+        graph["i_nd", "i2i", "i_nd"].edge_attr = torch.zeros((0, 0))
+        graph["d_nd", "d2d", "d_nd"].edge_attr = torch.zeros((0, self.d_edge_dim or 1))
+        return graph

D4CMPP2/src/DataManager/GraphGenerator/MolGraphGenerator.py

@@ -0,0 +1,73 @@
+import torch
+import rdkit.Chem as Chem
+import numpy as np
+import traceback
+from torch_geometric.data import Data
+from D4CMPP2.src.utils.featureizer import get_atom_features, get_bond_features, InvalidAtomError
+
+class MolGraphGenerator:
+    def __init__(self):
+        self.af = get_atom_features
+        self.bf = get_bond_features
+        self.set_feature_dim()
+
+    # Set the feature dimensions by generating a dummy graph
+    def set_feature_dim(self):
+        self.node_dim = self.af(Chem.MolFromSmiles('C')).shape[1]
+        self.edge_dim = self.bf(Chem.MolFromSmiles('C-C')).shape[1]
+
+    # Get the graph from the SMILES
+    def get_graph(self,smi,**kwargs):
+        mol = Chem.MolFromSmiles(smi)
+            
+        if kwargs.get("explicit_h",False) or mol.GetNumAtoms() == 1:
+            mol = Chem.AddHs(mol)
+        if mol is None: 
+            raise Exception("Invalid SMILES: failed to generate mol object")
+        g = self.generate_graph(mol)
+        g = self.add_feature(g,mol)
+        
+        return g
+    
+    def get_empty_graph(self):
+        return Data(
+            x=torch.zeros((0, self.node_dim), dtype=torch.float32),
+            edge_index=torch.empty((2, 0), dtype=torch.long),
+            edge_attr=torch.zeros((0, self.edge_dim), dtype=torch.float32),
+            num_nodes=0,
+        )
+
+    # Add the features to the graph
+    def add_feature(self, g, mol):
+        atom_feature = self.af(mol)
+        g.x = torch.tensor(atom_feature).float()
+        
+        if mol.GetNumBonds() == 0:
+            g.edge_attr = torch.zeros((0,self.edge_dim)).float()
+        else:
+            bond_feature = self.bf(mol)
+            edata = torch.tensor(bond_feature).float()
+            edata = torch.cat([edata,edata],dim=0)
+            g.edge_attr = edata
+        return g
+    
+    # Generate the graph from the molecule object
+    def generate_mol_graph(self,mol):
+        num_atoms = mol.GetNumAtoms()
+        if num_atoms == 1:
+            return ([0],[0])
+        src, dst = [], []
+        for bond in mol.GetBonds():
+            start, end = bond.GetBeginAtomIdx(), bond.GetEndAtomIdx()
+            src.append(start)
+            dst.append(end)
+
+        return (src+dst, dst+src)
+    
+    def generate_graph(self,mol):
+        mol_data = self.generate_mol_graph(mol)            
+        g = Data(
+            edge_index=torch.tensor(mol_data, dtype=torch.long),
+            num_nodes=mol.GetNumAtoms(),
+        )
+        return g