{"cells":[{"metadata":{},"cell_type":"markdown","source":"This kernel shows how to use RDKit and PyTorch Geometric (PyG) to prepare dataset for graph model. \n\nI am new for Kaggle and graph model, let me know if you see any typo or have any suggestion. \n\nSome useful links:\n\n1 PyG official tutorial: https://pytorch-geometric.readthedocs.io/en/latest/index.html\n\n2 PyG tutorial by Huang: https://towardsdatascience.com/hands-on-graph-neural-networks-with-pytorch-pytorch-geometric-359487e221a8\n\n3 GCN tutorial by Tobias: https://towardsdatascience.com/how-to-do-deep-learning-on-graphs-with-graph-convolutional-networks-7d2250723780\n\n4 MPNN kernel by fnand: https://www.kaggle.com/fnands/1-mpnn\n\nIn this kernel, I didn't include the target value (scalar coupling), I will update it in the future work. "},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"_kg_hide-input":true,"_kg_hide-output":true},"cell_type":"code","source":"import numpy as np \nimport pandas as pd\n\nimport os\nimport sys\n\nimport time\nfrom tqdm import tqdm","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Install RDKit and XYZ2MOL"},{"metadata":{},"cell_type":"markdown","source":"Taken from Jo's kernel with a little bit change. \n\nhttps://www.kaggle.com/sunhwan/using-rdkit-for-atomic-feature-and-visualization"},{"metadata":{"_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","trusted":true,"_kg_hide-output":true,"collapsed":true},"cell_type":"code","source":"!conda install --yes --prefix {sys.prefix} -c rdkit rdkit","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"collapsed":true},"cell_type":"code","source":"%%bash -e\nif ! [[ -f ./xyz2mol.py ]]; then\n  wget https://raw.githubusercontent.com/jensengroup/xyz2mol/master/xyz2mol.py\nfi","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# rdkit & xyz2mol\nfrom rdkit import Chem\nfrom rdkit.Chem import AllChem\nfrom rdkit.Chem.Draw import IPythonConsole #Needed to show molecules\nfrom rdkit.Chem import Draw\nfrom rdkit.Chem.Draw.MolDrawing import MolDrawing, DrawingOptions #Only needed if modifying defaults\nDrawingOptions.bondLineWidth=1.8\nfrom rdkit.Chem.rdmolops import SanitizeFlags\n\n# https://github.com/jensengroup/xyz2mol\nfrom xyz2mol import xyz2mol, xyz2AC, AC2mol, read_xyz_file\nfrom pathlib import Path\nimport pickle\n\nCACHEDIR = Path('./')\n\ndef chiral_stereo_check(mol):\n    # avoid sanitization error e.g., dsgdb9nsd_037900.xyz\n    Chem.SanitizeMol(mol, SanitizeFlags.SANITIZE_ALL - SanitizeFlags.SANITIZE_PROPERTIES)\n    Chem.DetectBondStereochemistry(mol,-1)\n    # ignore stereochemistry for now\n    #Chem.AssignStereochemistry(mol, flagPossibleStereoCenters=True, force=True)\n    #Chem.AssignAtomChiralTagsFromStructure(mol,-1)\n    return mol\n\ndef xyz2mol(atomicNumList,charge,xyz_coordinates,charged_fragments,quick):\n    AC,mol = xyz2AC(atomicNumList,xyz_coordinates)\n    new_mol = AC2mol(mol,AC,atomicNumList,charge,charged_fragments,quick)\n    new_mol = chiral_stereo_check(new_mol)\n    return new_mol\n\ndef MolFromXYZ(filename):\n    charged_fragments = True\n    quick = True\n    cache_filename = CACHEDIR/f'{filename.stem}.pkl'\n    if cache_filename.exists():\n        return pickle.load(open(cache_filename, 'rb'))\n    else:\n        try:\n            atomicNumList, charge, xyz_coordinates = read_xyz_file(filename)\n            mol = xyz2mol(atomicNumList, charge, xyz_coordinates, charged_fragments, quick)\n        except:\n            print(filename)\n    return mol","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Creating Dataset with PyG"},{"metadata":{},"cell_type":"markdown","source":"This part is inspired by:\n\n1 AlchemyDataset (https://github.com/tencent-alchemy/Alchemy/blob/master/pyg/Alchemy_dataset.py)\n\n2 Heng's discussion and starter kit (https://www.kaggle.com/c/champs-scalar-coupling/discussion/93972)\n"},{"metadata":{"trusted":true},"cell_type":"code","source":"! pip install --verbose --no-cache-dir torch-scatter\n! pip install --verbose --no-cache-dir torch-sparse\n! pip install --verbose --no-cache-dir torch-cluster\n! pip install --verbose --no-cache-dir torch-spline-conv\n! pip install torch-geometric","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import os.path as osp\nimport torch\nfrom torch_geometric.data import Data\nfrom torch_geometric.data import InMemoryDataset\nfrom rdkit import Chem\nfrom rdkit.Chem import ChemicalFeatures\nfrom rdkit import RDConfig\nimport networkx as nx\nimport pathlib\nimport pandas as pd\n\n\nDATA_DIR = '../input'\n\nclass ChampsDataset(InMemoryDataset):\n    fdef_name = osp.join(RDConfig.RDDataDir, 'BaseFeatures.fdef')\n    chem_feature_factory = ChemicalFeatures.BuildFeatureFactory(fdef_name)\n\n    def __init__(self, root, mode='train', transform=None, pre_transform=None, pre_filter=None):\n        self.mode = mode\n        self.df = pd.read_csv(DATA_DIR + '/%s.csv'%self.mode)\n        self.id = self.df.molecule_name.unique()\n        \n        super(ChampsDataset, self).__init__(root, transform, pre_transform, pre_filter)\n        self.data, self.slices = torch.load(self.processed_paths[0])\n\n    @property\n    def raw_file_names(self):\n        return []\n\n    @property\n    def processed_file_names(self):\n        return 'ChampsDataset_%s.pt'%self.mode\n\n    def download(self):\n        pass\n\n    def node_features(self, g):\n        feat = []\n        for n, d in g.nodes(data=True):\n            h_t = []\n            # Atom type (One-hot H, C, N, O F)\n            h_t += [int(d['a_type'] == x) for x in ['H', 'C', 'N', 'O', 'F', 'S', 'Cl']]\n            # Atomic number\n            h_t.append(d['a_num'])\n            # Acceptor\n            h_t.append(d['acceptor'])\n            # Donor\n            h_t.append(d['donor'])\n            # Aromatic\n            h_t.append(int(d['aromatic']))\n            # Hybradization\n            h_t += [int(d['hybridization'] == x) \\\n                    for x in (Chem.rdchem.HybridizationType.SP, \\\n                        Chem.rdchem.HybridizationType.SP2,\n                        Chem.rdchem.HybridizationType.SP3)]\n            h_t.append(d['num_h'])\n            feat.append((n, h_t))\n        feat.sort(key=lambda item: item[0])\n        node_attr = torch.FloatTensor([item[1] for item in feat])\n        return node_attr\n    \n\n    def edge_features(self, g):\n        e={}\n        for n1, n2, d in g.edges(data=True):\n            e_t = [int(d['b_type'] == x)\n                    for x in (Chem.rdchem.BondType.SINGLE, \\\n                            Chem.rdchem.BondType.DOUBLE, \\\n                            Chem.rdchem.BondType.TRIPLE, \\\n                            Chem.rdchem.BondType.AROMATIC)]\n            e[(n1, n2)] = e_t\n\n        edge_index = torch.LongTensor(list(e.keys())).transpose(0, 1)\n        edge_attr = torch.FloatTensor(list(e.values()))\n        return edge_index, edge_attr\n\n    # xyz file reader for dataset\n    def xyz_graph_reader(self, xyzfiles):\n        \n        mol = MolFromXYZ(xyzfiles)\n        feats = self.chem_feature_factory.GetFeaturesForMol(mol)\n\n        g = nx.DiGraph()\n\n        # Create nodes\n        assert len(mol.GetConformers()) == 1\n        geom = mol.GetConformers()[0].GetPositions()\n        for i in range(mol.GetNumAtoms()):\n            atom_i = mol.GetAtomWithIdx(i)\n            g.add_node(i, a_type=atom_i.GetSymbol(), a_num=atom_i.GetAtomicNum(), acceptor=0, donor=0,\n                    aromatic=atom_i.GetIsAromatic(), hybridization=atom_i.GetHybridization(),\n                    num_h=atom_i.GetTotalNumHs())\n\n        for i in range(len(feats)):\n            if feats[i].GetFamily() == 'Donor':\n                node_list = feats[i].GetAtomIds()\n                for i in node_list:\n                    g.node[i]['donor'] = 1\n            elif feats[i].GetFamily() == 'Acceptor':\n                node_list = feats[i].GetAtomIds()\n                for i in node_list:\n                    g.node[i]['acceptor'] = 1\n        # Read Edges\n        for i in range(mol.GetNumAtoms()):\n            for j in range(mol.GetNumAtoms()):\n                e_ij = mol.GetBondBetweenAtoms(i, j)\n                if e_ij is not None:\n                    g.add_edge(i, j, b_type=e_ij.GetBondType())\n\n        node_attr = self.node_features(g)\n        edge_index, edge_attr = self.edge_features(g)\n        data = Data(\n                x=node_attr,\n                pos=torch.FloatTensor(geom),\n                edge_index=edge_index,\n                edge_attr=edge_attr,\n                )\n        return data\n\n    def process(self):\n        data_list = []\n        molecule_names = self.df.molecule_name.unique() \n        xyzfiles = [Path(f'../input/structures/{f}.xyz') for f in molecule_names]\n        \n        for i in tqdm(range(len(molecule_names))):\n            champs_data = self.xyz_graph_reader(xyzfiles[i])\n            if champs_data is not None:\n                data_list.append(champs_data)\n\n        if self.pre_filter is not None:\n            data_list = [data for data in data_list if self.pre_filter(data)]\n\n        if self.pre_transform is not None:\n            data_list = [self.pre_transform(data) for data in data_list]\n\n        data, slices = self.collate(data_list)\n        torch.save((data, slices), self.processed_paths[0])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_data = ChampsDataset(root ='../', mode = 'train')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_data = ChampsDataset(root ='../', mode = 'test')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Take a look of a random graph data."},{"metadata":{"trusted":true},"cell_type":"code","source":"train_data[50000]","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Save Dataset"},{"metadata":{"trusted":true},"cell_type":"code","source":"torch.save(train_data, 'train.pt')\ntorch.save(test_data, 'test.pt')","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.6.4","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"}},"nbformat":4,"nbformat_minor":1}