data.py

import torch
from torch_geometric.data import Data, Batch, InMemoryDataset, Dataset
from torch_geometric.utils import from_networkx, from_smiles, to_networkx
from torch_geometric.datasets import TUDataset
import networkx as nx
import pickle
import random
import os
import numpy as np

from rdkit import Chem
from ACG import unknown_product
from ogb.graphproppred import PygGraphPropPredDataset

from multiprocessing import Pool, Manager

class Single_dataset(InMemoryDataset):
    def __init__(self, root, name, transform=None, pre_transform=None, pre_filter=None):
        self.name = name
        super(Single_dataset, self).__init__(root, transform, pre_transform, pre_filter)
        self.root = root
        self.data, self.slices = torch.load(self.processed_paths[0])

    @property
    def raw_file_names(self):
        return ["graphs.pkl"]

    @property
    def processed_file_names(self):
        if self.name == 'A':
            return ['data_A.pt']
        else:
            return ['data_B.pt']
        
    def download(self):
        pass  

    def process(self):
        # Read data into huge `Data` list.
        path = self.root + '/raw/'
        if self.name == 'A':
            data_list, _ = pickle.load(open(path+'graphs.pkl', "rb"))
        else:
            _, data_list = pickle.load(open(path+'graphs.pkl', "rb"))

        self.data, self.slices = self.collate(data_list)
        torch.save((self.data, self.slices), self.processed_paths[0])
        
class PairDataset(Dataset):
    def __init__(self, root, transform=None, pre_transform=None, pre_filter=None):
        super().__init__(root, transform, pre_transform, pre_filter)
        self.data_a = Single_dataset(root, 'A')
        self.data_b = Single_dataset(root, 'B')

    def len(self):
        return len(self.data_a)

    def get(self, idx):
        return (self.data_a[idx], self.data_b[idx])
    
class PairDatasetProduct(Dataset):
    def __init__(self, root, transform=None, pre_transform=None, pre_filter=None):
        super().__init__(root, transform, pre_transform, pre_filter)
        self.data_a = Single_dataset(root, 'A')
        self.data_b = Single_dataset(root, 'B')

    def len(self):
        return len(self.data_a)

    def get(self, idx):
        data_s = self.data_a[idx]
        data_t = self.data_b[idx]
        data_s.edge_label = data_s.edge_attr
        data_t.edge_label = data_t.edge_attr
        G_s, G_t = to_networkx(data_s, to_undirected=True, node_attrs=['x'], edge_attrs=['edge_label']), to_networkx(data_t, to_undirected=True, node_attrs=['x'], edge_attrs=['edge_label'])
        G_st, G_complement_st = unknown_product(G_s, G_t)
        for node in G_st.nodes():
            G_st.nodes[node]['label'] = node
        data_st = from_networkx(G_st)
        return (data_s, data_t, data_st)


x_map = {
    'atomic_num':
    list(range(0, 119)),
    'chirality': [
        'CHI_UNSPECIFIED',
        'CHI_TETRAHEDRAL_CW',
        'CHI_TETRAHEDRAL_CCW',
        'CHI_OTHER',
        'CHI_TETRAHEDRAL',
        'CHI_ALLENE',
        'CHI_SQUAREPLANAR',
        'CHI_TRIGONALBIPYRAMIDAL',
        'CHI_OCTAHEDRAL',
    ],
    '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': [
        'UNSPECIFIED',
        'SINGLE',
        'DOUBLE',
        'TRIPLE',
        'QUADRUPLE',
        'QUINTUPLE',
        'HEXTUPLE',
        'ONEANDAHALF',
        'TWOANDAHALF',
        'THREEANDAHALF',
        'FOURANDAHALF',
        'FIVEANDAHALF',
        'AROMATIC',
        'IONIC',
        'HYDROGEN',
        'THREECENTER',
        'DATIVEONE',
        'DATIVE',
        'DATIVEL',
        'DATIVER',
        'OTHER',
        'ZERO',
    ],
    'stereo': [
        'STEREONONE',
        'STEREOANY',
        'STEREOZ',
        'STEREOE',
        'STEREOCIS',
        'STEREOTRANS',
    ],
    'is_conjugated': [False, True],
}


def to_rdkit(data, kekulize: bool = False, use_original_index=False, dataset_name=None):
    from rdkit import Chem
    if dataset_name == 'AIDS':
        index = torch.tensor([ 6,  8,  7, 17,  9, 16, 34, 15, 11, 53, 27, 35,  3, 14, 12, 29, 33,  5,
            78, 44, 19, 46, 79, 52, 74, 45, 30, 83, 82, 32, 51, 50, 31, 80, 67, 81,
            28, 65])
    elif dataset_name == 'MCF-7':
        index = torch.tensor([8, 7, 6, 35, 16, 17, 9, 11, 78, 30, 28, 25, 15, 53, 34, 50, 26, 82, 14, 24, 80, 33, 5, 31, 22, 83, 19, 29, 40, 77, 3, 46, 79, 74, 51, 27, 12, 47, 45, 44, 48, 68, 23, 89, 81, 32]
                    )
    mol = Chem.RWMol()

    for i in range(data.num_nodes):
        if use_original_index:
            atom = Chem.Atom(int(data.x[i]))
        else:
            atom = Chem.Atom(index[data.x[i]].item())

        mol.AddAtom(atom)

    edges = [tuple(i) for i in data.edge_index.t().tolist()]
    visited = set()

    for i in range(len(edges)):
        src, dst = edges[i]
        if tuple(sorted(edges[i])) in visited:
            continue

        bond_type = Chem.BondType.values[data.edge_attr[i].item()]
        mol.AddBond(src, dst, bond_type)

        visited.add(tuple(sorted(edges[i])))

    mol = mol.GetMol()

    if kekulize:
        Chem.Kekulize(mol)

    aa = Chem.SanitizeMol(mol, catchErrors=True)
    if aa == Chem.rdmolops.SanitizeFlags.SANITIZE_PROPERTIES:
        return None
    
    return mol


def to_rdmol(
    data,
    kekulize: bool = False,
):
    """Converts a :class:`torch_geometric.data.Data` instance to a
    :class:`rdkit.Chem.Mol` instance.

    Args:
        data (torch_geometric.data.Data): The molecular graph data.
        kekulize (bool, optional): If set to :obj:`True`, converts aromatic
            bonds to single/double bonds. (default: :obj:`False`)
    """
    from rdkit import Chem

    mol = Chem.RWMol()

    assert data.x is not None
    assert data.num_nodes is not None
    assert data.edge_index is not None
    assert data.edge_attr is not None
    for i in range(data.num_nodes):
        atom = Chem.Atom(int(data.x[i]))
        mol.AddAtom(atom)

    edges = [tuple(i) for i in data.edge_index.t().tolist()]
    visited = set()

    for i in range(len(edges)):
        src, dst = edges[i]
        if tuple(sorted(edges[i])) in visited:
            continue

        bond_type = Chem.BondType.values[int(data.edge_attr[i])]
        mol.AddBond(src, dst, bond_type)

        visited.add(tuple(sorted(edges[i])))

    mol = mol.GetMol()

    if kekulize:
        Chem.Kekulize(mol)

    Chem.SanitizeMol(mol)
    Chem.AssignStereochemistry(mol)

    return mol


def data_gen(dataset, dataset_name, num_samples, N_min=50, N_max=15, path=None):
    graph_list_A = []
    graph_list_B = []
    
    from rdkit.Chem import rdRascalMCES
    opts = rdRascalMCES.RascalOptions()
    opts.similarityThreshold = 0.0
    opts.completeAromaticRings = False
    opts.maxBondMatchPairs = 100000
    opts.timeout = 20000
    
    while len(graph_list_A) < num_samples:
        # filter graph size
        random_sample = random.choice(dataset)
        while not (N_min <= random_sample.num_nodes <= N_max):
            random_sample = random.choice(dataset)
        data_A = random_sample.clone()
        random_sample = random.choice(dataset)
        while not (N_min <= random_sample.num_nodes <= N_max):
            random_sample = random.choice(dataset)
        data_B = random_sample.clone()
    
        if dataset_name == 'AIDS' or dataset_name == 'MCF-7':
        # from one hot encoding to decimal encoding
            A_x = data_A.x.argmax(dim=1)
            A_edge_attr = data_A.edge_attr.argmax(dim=1)
            data_A.x = A_x
            data_A.edge_attr = A_edge_attr + 1
            B_x = data_B.x.argmax(dim=1)
            B_edge_attr = data_B.edge_attr.argmax(dim=1)
            data_B.x = B_x
            data_B.edge_attr = B_edge_attr + 1
            use_original_index = False
        elif dataset_name == 'MOLHIV':
            data_A.x = data_A.x[:, 0] + 1
            data_A.edge_attr = data_A.edge_attr[:, 0] + 1
            data_A.edge_attr[data_A.edge_attr == 4] = 12
            data_B.x = data_B.x[:, 0] + 1
            data_B.edge_attr = data_B.edge_attr[:, 0] + 1
            data_B.edge_attr[data_B.edge_attr == 4] = 12
            use_original_index = True
        
        # G_A = to_networkx(data_A)
        # G_B = to_networkx(data_B)
        
        # rascal
        mol_A = to_rdkit(data_A, dataset_name=dataset_name, kekulize=False, use_original_index=use_original_index)
        mol_B = to_rdkit(data_B, dataset_name=dataset_name, kekulize=False, use_original_index=use_original_index)
        if mol_A is None or mol_B is None:
            continue
        
        data_A = from_smiles(Chem.MolToSmiles(mol_A), kekulize=False)
        data_B = from_smiles(Chem.MolToSmiles(mol_B), kekulize=False)
        
        if data_A.num_nodes == 0 or data_B.num_nodes == 0:
            continue
        
        # print(Chem.MolToSmiles(mol_A))
        # print(Chem.MolToSmiles(mol_B))
        allowed_values = torch.tensor([1, 2, 3, 12])
        is_outside1 = ~torch.isin(data_A.edge_attr[:, 0], allowed_values)
        is_outside2 = ~torch.isin(data_B.edge_attr[:, 0], allowed_values)
        if is_outside1.any() or is_outside2.any():
            continue
        
        results = rdRascalMCES.FindMCES(mol_A, mol_B, opts)
        
        if results == []:
            continue
        res = results[0]
        bond_matches = res.bondMatches()
        molA_bonds = []
        molB_bonds = []

        for bond1_idx, bond2_idx in bond_matches:
            bond1 = mol_A.GetBondWithIdx(bond1_idx)
            atomA_1 = bond1.GetBeginAtomIdx()
            atomA_2 = bond1.GetEndAtomIdx()

            bond2 = mol_B.GetBondWithIdx(bond2_idx)
            atomB_1 = bond2.GetBeginAtomIdx()
            atomB_2 = bond2.GetEndAtomIdx()
            
            molA_bonds.append([atomA_1, atomA_2])
            molB_bonds.append([atomB_1, atomB_2])
        
        atom_matches = res.atomMatches()   
        r_indices = [match[0] for match in atom_matches]
        r_indices = np.array(r_indices)
        t_indices = [match[1] for match in atom_matches]
        t_indices = np.array(t_indices)
        
        mces_size_edge = torch.tensor(len(res.bondMatches()))
        mces_size_node = torch.tensor(len(res.atomMatches()))
        
        sim = res.similarity
        
        data_A.y = torch.stack([mces_size_edge, mces_size_node, torch.tensor(sim)], dim=-1).unsqueeze(0)
        data_B.y = torch.stack([mces_size_edge, mces_size_node, torch.tensor(sim)], dim=-1).unsqueeze(0)
        
        data_A.x = data_A.x[:, 0]
        data_A.edge_attr = data_A.edge_attr[:, 0]
        data_B.x = data_B.x[:, 0]
        data_B.edge_attr = data_B.edge_attr[:, 0]
        
        data_A.match_indices = torch.tensor(r_indices)
        data_B.match_indices = torch.tensor(t_indices)
        
        graph_list_A.append(data_A)
        graph_list_B.append(data_B)

    os.makedirs(path+'raw/', exist_ok=True)
        
    if args.process_id >= 0:
        path = path + f"raw/graphs_{args.process_id}.pkl"
    else:
        path = path + "raw/graphs.pkl"
    with open(path, "wb+") as fp:
        pickle.dump([graph_list_A, graph_list_B], fp)
        
    return graph_list_A, graph_list_B


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--num_samples', type=int, default=100)
    parser.add_argument('--min_num_nodes', type=int, default=30)
    parser.add_argument('--num_nodes', type=int, default=60)
    parser.add_argument('--num_workers', type=int, default=4)
    parser.add_argument('--process_id', type=int, default=-1)
    parser.add_argument('--dataset', type=str, default='MOLHIV')
    args, unknown = parser.parse_known_args()
    
    if args.dataset == 'AIDS':
        path = './data/MCES/AIDS-train/'
        ori_dataset = TUDataset(root='./Data/'+args.dataset, name=args.dataset)
    elif args.dataset == 'MCF-7':
        path = './data/MCES/MCF-7-train/'
        ori_dataset = TUDataset(root='./Data/'+args.dataset, name=args.dataset)
    elif args.dataset == 'MOLHIV':
        path = './data/MCES/MOLHIV-train/'
        ori_dataset = PygGraphPropPredDataset(root='./Data/'+args.dataset, name='ogbg-molhiv')

    graph_list_A, graph_list_B = data_gen(dataset=ori_dataset, dataset_name=args.dataset, num_samples=args.num_samples, N_max=args.num_nodes, N_min=args.min_num_nodes, path=path)