{"cells":[{"metadata":{"trusted":true},"cell_type":"code","source":"!conda install -y -c rdkit rdkit","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"pip install --pre dgl","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nimport matplotlib.pyplot as plt\nimport gc\nimport numpy as np\nimport pandas as pd\nfrom itertools import combinations\nfrom glob import glob\n\nfrom rdkit.Chem import rdmolops, ChemicalFeatures\n\n#from utils import print_progress\n#import constants as C\n\nimport os\n%cd /kaggle/input/champs-scalar-coupling\nprint(os.listdir(\".\"))\n\n# Any results you write to the current directory are saved as output.","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"##\n# Copied from: https://github.com/jensengroup/xyz2mol\n# \n#\n# Written by Jan H. Jensen based on this paper Yeonjoon Kim and Woo Youn Kim\n# \"Universal Structure Conversion Method for Organic Molecules: From Atomic Connectivity\n# to Three-Dimensional Geometry\" Bull. Korean Chem. Soc. 2015, Vol. 36, 1769-1777 DOI: 10.1002/bkcs.10334\n#\nfrom rdkit import Chem\nfrom rdkit.Chem import AllChem\nimport itertools\nfrom rdkit.Chem import rdmolops\nfrom collections import defaultdict\nimport copy\nimport networkx as nx #uncomment if you don't want to use \"quick\"/install networkx\n\n\nglobal __ATOM_LIST__\n__ATOM_LIST__ = [ x.strip() for x in ['h ','he', \\\n      'li','be','b ','c ','n ','o ','f ','ne', \\\n      'na','mg','al','si','p ','s ','cl','ar', \\\n      'k ','ca','sc','ti','v ','cr','mn','fe','co','ni','cu', \\\n      'zn','ga','ge','as','se','br','kr', \\\n      'rb','sr','y ','zr','nb','mo','tc','ru','rh','pd','ag', \\\n      'cd','in','sn','sb','te','i ','xe', \\\n      'cs','ba','la','ce','pr','nd','pm','sm','eu','gd','tb','dy', \\\n      'ho','er','tm','yb','lu','hf','ta','w ','re','os','ir','pt', \\\n      'au','hg','tl','pb','bi','po','at','rn', \\\n      'fr','ra','ac','th','pa','u ','np','pu'] ]\n\n\ndef get_atom(atom):\n    global __ATOM_LIST__\n    atom = atom.lower()\n    return __ATOM_LIST__.index(atom) + 1\n\n\ndef getUA(maxValence_list, valence_list):\n    UA = []\n    DU = []\n    for i, (maxValence,valence) in enumerate(zip(maxValence_list, valence_list)):\n        if maxValence - valence > 0:\n            UA.append(i)\n            DU.append(maxValence - valence)\n    return UA,DU\n\n\ndef get_BO(AC,UA,DU,valences,UA_pairs,quick):\n    BO = AC.copy()\n    DU_save = []\n\n    while DU_save != DU:\n        for i,j in UA_pairs:\n            BO[i,j] += 1\n            BO[j,i] += 1\n\n        BO_valence = list(BO.sum(axis=1))\n        DU_save = copy.copy(DU)\n        UA, DU = getUA(valences, BO_valence)\n        UA_pairs = get_UA_pairs(UA,AC,quick)[0]\n\n    return BO\n\n\ndef valences_not_too_large(BO,valences):\n    number_of_bonds_list = BO.sum(axis=1)\n    for valence, number_of_bonds in zip(valences,number_of_bonds_list):\n        if number_of_bonds > valence:\n            return False\n\n    return True\n\n\ndef BO_is_OK(BO,AC,charge,DU,atomic_valence_electrons,atomicNumList,charged_fragments):\n    Q = 0 # total charge\n    q_list = []\n    if charged_fragments:\n        BO_valences = list(BO.sum(axis=1))\n        for i,atom in enumerate(atomicNumList):\n            q = get_atomic_charge(atom,atomic_valence_electrons[atom],BO_valences[i])\n            Q += q\n            if atom == 6:\n                number_of_single_bonds_to_C = list(BO[i,:]).count(1)\n                if number_of_single_bonds_to_C == 2 and BO_valences[i] == 2:\n                    Q += 1\n                    q = 2\n                if number_of_single_bonds_to_C == 3 and Q + 1 < charge:\n                    Q += 2\n                    q = 1\n\n            if q != 0:\n                q_list.append(q)\n\n    if (BO-AC).sum() == sum(DU) and charge == Q and len(q_list) <= abs(charge):\n        return True\n    else:\n        return False\n\n\ndef get_atomic_charge(atom,atomic_valence_electrons,BO_valence):\n    if atom == 1:\n        charge = 1 - BO_valence\n    elif atom == 5:\n        charge = 3 - BO_valence\n    elif atom == 15 and BO_valence == 5:\n        charge = 0\n    elif atom == 16 and BO_valence == 6:\n        charge = 0\n    else:\n        charge = atomic_valence_electrons - 8 + BO_valence\n\n    return charge\n\ndef clean_charges(mol):\n    # this hack should not be needed any more but is kept just in case\n\n    rxn_smarts = ['[N+:1]=[*:2]-[C-:3]>>[N+0:1]-[*:2]=[C-0:3]',\n                  '[N+:1]=[*:2]-[O-:3]>>[N+0:1]-[*:2]=[O-0:3]',\n                  '[N+:1]=[*:2]-[*:3]=[*:4]-[O-:5]>>[N+0:1]-[*:2]=[*:3]-[*:4]=[O-0:5]',\n                  '[#8:1]=[#6:2]([!-:6])[*:3]=[*:4][#6-:5]>>[*-:1][*:2]([*:6])=[*:3][*:4]=[*+0:5]',\n                  '[O:1]=[c:2][c-:3]>>[*-:1][*:2][*+0:3]',\n                  '[O:1]=[C:2][C-:3]>>[*-:1][*:2]=[*+0:3]']\n\n    fragments = Chem.GetMolFrags(mol,asMols=True,sanitizeFrags=False)\n\n    for i,fragment in enumerate(fragments):\n        for smarts in rxn_smarts:\n            patt = Chem.MolFromSmarts(smarts.split(\">>\")[0])\n            while fragment.HasSubstructMatch(patt):\n                rxn = AllChem.ReactionFromSmarts(smarts)\n                ps = rxn.RunReactants((fragment,))\n                fragment = ps[0][0]\n        if i == 0:\n            mol = fragment\n        else:\n            mol = Chem.CombineMols(mol,fragment)\n\n    return mol\n\n\ndef BO2mol(mol,BO_matrix, atomicNumList,atomic_valence_electrons,mol_charge,charged_fragments):\n    # based on code written by Paolo Toscani\n\n    l = len(BO_matrix)\n    l2 = len(atomicNumList)\n    BO_valences = list(BO_matrix.sum(axis=1))\n\n    if (l != l2):\n        raise RuntimeError('sizes of adjMat ({0:d}) and atomicNumList '\n            '{1:d} differ'.format(l, l2))\n\n    rwMol = Chem.RWMol(mol)\n\n    bondTypeDict = {\n        1: Chem.BondType.SINGLE,\n        2: Chem.BondType.DOUBLE,\n        3: Chem.BondType.TRIPLE\n    }\n\n    for i in range(l):\n        for j in range(i + 1, l):\n            bo = int(round(BO_matrix[i, j]))\n            if (bo == 0):\n                continue\n            bt = bondTypeDict.get(bo, Chem.BondType.SINGLE)\n            rwMol.AddBond(i, j, bt)\n    mol = rwMol.GetMol()\n\n    if charged_fragments:\n        mol = set_atomic_charges(mol,atomicNumList,atomic_valence_electrons,BO_valences,BO_matrix,mol_charge)\n    else:\n        mol = set_atomic_radicals(mol,atomicNumList,atomic_valence_electrons,BO_valences)\n\n    return mol\n\ndef set_atomic_charges(mol,atomicNumList,atomic_valence_electrons,BO_valences,BO_matrix,mol_charge):\n    q = 0\n    for i,atom in enumerate(atomicNumList):\n        a = mol.GetAtomWithIdx(i)\n        charge = get_atomic_charge(atom,atomic_valence_electrons[atom],BO_valences[i])\n        q += charge\n        if atom == 6:\n            number_of_single_bonds_to_C = list(BO_matrix[i,:]).count(1)\n            if number_of_single_bonds_to_C == 2 and BO_valences[i] == 2:\n                    q += 1\n                    charge = 0\n            if number_of_single_bonds_to_C == 3 and q + 1 < mol_charge:\n                    q += 2\n                    charge = 1\n\n        if (abs(charge) > 0):\n            a.SetFormalCharge(int(charge))\n\n    # shouldn't be needed anymore bit is kept just in case\n    #mol = clean_charges(mol)\n\n    return mol\n\n\ndef set_atomic_radicals(mol,atomicNumList,atomic_valence_electrons,BO_valences):\n    # The number of radical electrons = absolute atomic charge\n    for i,atom in enumerate(atomicNumList):\n        a = mol.GetAtomWithIdx(i)\n        charge = get_atomic_charge(atom,atomic_valence_electrons[atom],BO_valences[i])\n\n        if (abs(charge) > 0):\n            a.SetNumRadicalElectrons(abs(int(charge)))\n\n    return mol\n\ndef get_bonds(UA,AC):\n    bonds = []\n\n    for k,i in enumerate(UA):\n        for j in UA[k+1:]:\n            if AC[i,j] == 1:\n                bonds.append(tuple(sorted([i,j])))\n\n    return bonds\n\ndef get_UA_pairs(UA,AC,quick):\n    bonds = get_bonds(UA,AC)\n    if len(bonds) == 0:\n        return [()]\n\n    if quick:\n        G=nx.Graph()\n        G.add_edges_from(bonds)\n        UA_pairs = [list(nx.max_weight_matching(G))]\n        return UA_pairs\n\n    max_atoms_in_combo = 0\n    UA_pairs = [()]\n    for combo in list(itertools.combinations(bonds, int(len(UA)/2))):\n        flat_list = [item for sublist in combo for item in sublist]\n        atoms_in_combo = len(set(flat_list))\n        if atoms_in_combo > max_atoms_in_combo:\n            max_atoms_in_combo = atoms_in_combo\n            UA_pairs = [combo]\n #           if quick and max_atoms_in_combo == 2*int(len(UA)/2):\n #               return UA_pairs\n        elif atoms_in_combo == max_atoms_in_combo:\n            UA_pairs.append(combo)\n\n    return UA_pairs\n\ndef AC2BO(AC,atomicNumList,charge,charged_fragments,quick):\n    # TODO\n    atomic_valence = defaultdict(list)\n    atomic_valence[1] = [1]\n    atomic_valence[6] = [4]\n    atomic_valence[7] = [4,3]\n    atomic_valence[8] = [2,1]\n    atomic_valence[9] = [1]\n    atomic_valence[14] = [4]\n    atomic_valence[15] = [5,4,3]\n    atomic_valence[16] = [6,4,2]\n    atomic_valence[17] = [1]\n    atomic_valence[32] = [4]\n    atomic_valence[35] = [1]\n    atomic_valence[53] = [1]\n\n\n    atomic_valence_electrons = {}\n    atomic_valence_electrons[1] = 1\n    atomic_valence_electrons[6] = 4\n    atomic_valence_electrons[7] = 5\n    atomic_valence_electrons[8] = 6\n    atomic_valence_electrons[9] = 7\n    atomic_valence_electrons[14] = 4\n    atomic_valence_electrons[15] = 5\n    atomic_valence_electrons[16] = 6\n    atomic_valence_electrons[17] = 7\n    atomic_valence_electrons[32] = 4\n    atomic_valence_electrons[35] = 7\n    atomic_valence_electrons[53] = 7\n\n    # make a list of valences, e.g. for CO: [[4],[2,1]]\n    valences_list_of_lists = []\n    for atomicNum in atomicNumList:\n        valences_list_of_lists.append(atomic_valence[atomicNum])\n\n    # convert [[4],[2,1]] to [[4,2],[4,1]]\n    valences_list = list(itertools.product(*valences_list_of_lists))\n\n    best_BO = AC.copy()\n\n    # implemenation of algorithm shown in Figure 2\n    # UA: unsaturated atoms\n    # DU: degree of unsaturation (u matrix in Figure)\n    # best_BO: Bcurr in Figure\n    #\n\n    for valences in valences_list:\n        AC_valence = list(AC.sum(axis=1))\n        UA,DU_from_AC = getUA(valences, AC_valence)\n\n        if len(UA) == 0 and BO_is_OK(AC,AC,charge,DU_from_AC,atomic_valence_electrons,atomicNumList,charged_fragments):\n            return AC,atomic_valence_electrons\n\n        UA_pairs_list = get_UA_pairs(UA,AC,quick)\n        for UA_pairs in UA_pairs_list:\n            BO = get_BO(AC,UA,DU_from_AC,valences,UA_pairs,quick)\n            if BO_is_OK(BO,AC,charge,DU_from_AC,atomic_valence_electrons,atomicNumList,charged_fragments):\n                return BO,atomic_valence_electrons\n\n            elif BO.sum() >= best_BO.sum() and valences_not_too_large(BO,valences):\n                best_BO = BO.copy()\n\n    return best_BO,atomic_valence_electrons\n\n\ndef AC2mol(mol,AC,atomicNumList,charge,charged_fragments,quick):\n    # convert AC matrix to bond order (BO) matrix\n    BO,atomic_valence_electrons = AC2BO(AC,atomicNumList,charge,charged_fragments,quick)\n\n    # add BO connectivity and charge info to mol object\n    mol = BO2mol(mol,BO, atomicNumList,atomic_valence_electrons,charge,charged_fragments)\n\n    return mol\n\n\ndef get_proto_mol(atomicNumList):\n    mol = Chem.MolFromSmarts(\"[#\"+str(atomicNumList[0])+\"]\")\n    rwMol = Chem.RWMol(mol)\n    for i in range(1,len(atomicNumList)):\n        a = Chem.Atom(atomicNumList[i])\n        rwMol.AddAtom(a)\n\n    mol = rwMol.GetMol()\n\n    return mol\n\n\ndef get_atomicNumList(atomic_symbols):\n    atomicNumList = []\n    for symbol in atomic_symbols:\n        atomicNumList.append(get_atom(symbol))\n    return atomicNumList\n\n\ndef read_xyz_file(filename):\n\n    atomic_symbols = []\n    xyz_coordinates = []\n\n    with open(filename, \"r\") as file:\n        for line_number,line in enumerate(file):\n            if line_number == 0:\n                num_atoms = int(line)\n            elif line_number == 1:\n                if \"charge=\" in line:\n                    charge = int(line.split(\"=\")[1])\n                else:\n                    charge = 0\n            else:\n                atomic_symbol, x, y, z = line.split()\n                atomic_symbols.append(atomic_symbol)\n                xyz_coordinates.append([float(x),float(y),float(z)])\n\n    atomicNumList = get_atomicNumList(atomic_symbols)\n\n    return atomicNumList,charge,xyz_coordinates\n\ndef xyz2AC(atomicNumList,xyz):\n    import numpy as np\n    mol = get_proto_mol(atomicNumList)\n\n    conf = Chem.Conformer(mol.GetNumAtoms())\n    for i in range(mol.GetNumAtoms()):\n        conf.SetAtomPosition(i,(xyz[i][0],xyz[i][1],xyz[i][2]))\n    mol.AddConformer(conf)\n\n    dMat = Chem.Get3DDistanceMatrix(mol)\n    pt = Chem.GetPeriodicTable()\n\n    num_atoms = len(atomicNumList)\n    AC = np.zeros((num_atoms,num_atoms)).astype(int)\n\n    for i in range(num_atoms):\n        a_i = mol.GetAtomWithIdx(i)\n        Rcov_i = pt.GetRcovalent(a_i.GetAtomicNum())*1.30\n        for j in range(i+1,num_atoms):\n            a_j = mol.GetAtomWithIdx(j)\n            Rcov_j = pt.GetRcovalent(a_j.GetAtomicNum())*1.30\n            if dMat[i,j] <= Rcov_i + Rcov_j:\n                AC[i,j] = 1\n                AC[j,i] = 1\n\n    return AC,mol,dMat\n\ndef chiral_stereo_check(mol):\n    Chem.DetectBondStereochemistry(mol,-1)\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            check_chiral_stereo=True):\n\n    # Get atom connectivity (AC) matrix, list of atomic numbers, molecular charge,\n    # and mol object with no connectivity information\n    AC,mol,dMat = xyz2AC(atomicNumList, xyz_coordinates)\n\n    # Convert AC to bond order matrix and add connectivity and charge info to mol object\n    new_mol = AC2mol(mol, AC, atomicNumList, charge, charged_fragments, quick)\n\n    # sanitize\n    try: Chem.SanitizeMol(new_mol)\n    except ValueError as e: print(e)\n\n    # Check for stereocenters and chiral centers\n    if check_chiral_stereo:\n        try: new_mol = chiral_stereo_check(new_mol)\n        except ValueError as e: print(e)\n\n    return new_mol,dMat","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import os\n%cd /kaggle/input/champs-scalar-coupling\nprint(os.listdir(\".\"))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import gc\nimport numpy as np\nimport pandas as pd\nfrom itertools import combinations\nfrom glob import glob\n\nfrom rdkit.Chem import rdmolops, ChemicalFeatures\n\n\n\nmol_feat_columns = ['ave_bond_length', 'std_bond_length', 'ave_atom_weight']\nxyz_filepath_list = list(glob('structures/*.xyz'))\nxyz_filepath_list.sort()\n\n\n## Functions to create the RDKit mol objects\ndef mol_from_xyz(filepath, add_hs=True, compute_dist_centre=False):\n    \"\"\"Wrapper function for calling xyz2mol function.\"\"\"\n    charged_fragments = True  # alternatively radicals are made\n\n    # quick is faster for large systems but requires networkx\n    # if you don't want to install networkx set quick=False and\n    # uncomment 'import networkx as nx' at the top of the file\n    quick = True\n\n    atomicNumList, charge, xyz_coordinates = read_xyz_file(filepath)\n    mol, dMat = xyz2mol(atomicNumList, charge, xyz_coordinates,\n                        charged_fragments, quick, check_chiral_stereo=False)\n\n    return mol, np.array(xyz_coordinates), dMat\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"mol, xyz_coordinates, dMat = mol_from_xyz(xyz_filepath_list[0])","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"mol: rdkit.Chem.rdchem.Mol\ndMat: Distance Matrix"},{"metadata":{"trusted":true},"cell_type":"code","source":"from dgl.data.chem import mol_to_bigraph","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"g = mol_to_bigraph(mol)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from dgl.data.chem import mol_to_bigraph\ndef get_graph(filepath='./', add_hs=True, compute_dist_centre=False):\n    graph_list = []\n    xyz_filepath_list = list(glob(filepath+'structures/*.xyz'))\n    xyz_filepath_list.sort()\n    for i in range(len(xyz_filepath_list)):\n        mol, xyz_coordinates, _ = mol_from_xyz(xyz_filepath_list[i])\n        try:\n            g = mol_to_bigraph(mol)\n            g.ndata['xyz'] = xyz_coordinates\n            graph_list.append(g)\n        except RuntimeError:\n            pass\n            \n    return graph_list","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_molecules():\n    \"\"\"\n    Constructs rdkit mol objects derrived from the .xyz files. Also returns:\n        - mol ids (unique numerical ids)\n        - set of molecule level features\n        - arrays of xyz coordinates\n        - euclidean distance matrices\n        - graph distance matrices.\n    All objects are returned in dictionaries with 'mol_name' as keys.\n    \"\"\"\n    N_MOLS =  130775\n    MAX_N_ATOMS = 29\n    mols, mol_ids, mol_feats = {}, {}, {}\n    xyzs, dist_matrices, graph_dist_matrices = {}, {}, {}\n    print('Create molecules and distance matrices.')\n    for i in range(N_MOLS):\n        #print_progress(i, N_MOLS)\n        filepath = xyz_filepath_list[i]\n        mol_name = filepath.split('/')[-1][:-4]\n        mol, xyz, dist_matrix = mol_from_xyz(filepath)\n        mols[mol_name] = mol\n        xyzs[mol_name] = xyz\n        dist_matrices[mol_name] = dist_matrix\n        mol_ids[mol_name] = i\n\n        # make padded graph distance matrix dataframes\n        n_atoms = len(xyz)\n        graph_dist_matrix = pd.DataFrame(np.pad(\n            rdmolops.GetDistanceMatrix(mol),\n            [(0, 0), (0, MAX_N_ATOMS - n_atoms)], 'constant'\n        ))\n        graph_dist_matrix['molecule_id'] = n_atoms * [i]\n        graph_dist_matrices[mol_name] = graph_dist_matrix\n\n        # compute molecule level features\n        adj_matrix = rdmolops.GetAdjacencyMatrix(mol)\n        atomic_num_list, _, _ = read_xyz_file(filepath)\n        dists = dist_matrix.ravel()[np.tril(adj_matrix).ravel()==1]\n        mol_feats[mol_name] = pd.Series(\n            [np.mean(dists), np.std(dists), np.mean(atomic_num_list)],\n            index=mol_feat_columns\n        )\n\n    return mols, mol_ids, mol_feats, xyzs, dist_matrices, graph_dist_matrices\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from dgl.data.chem import mol_to_bigraph\nimport torch as th\n\ndef get_molecules(filepath='./', add_hs=True, compute_dist_centre=False):\n    N_MOLS =  130775\n    MAX_N_ATOMS = 29\n    graph_list = []\n    mol_feat_columns = ['ave_bond_length', 'std_bond_length', 'ave_atom_weight']\n\n    xyz_filepath_list = list(glob(filepath+'structures/*.xyz'))\n    xyz_filepath_list.sort()\n    \n    for i in range(100):\n        mol, xyz, dist_matrix = mol_from_xyz(xyz_filepath_list[i])\n        try:\n            g = mol_to_bigraph(mol)\n            \n            g.ndata['xyz'] = xyz\n            g.gdata = {}\n            \n            mol_name = xyz_filepath_list[i].split('/')[-1][:-4]\n            g.gdata['name'] = mol_name \n            \n            g.gdata['dist_matrices'] = th.tensor(dist_matrix)\n            g.gdata['mol_id'] = i\n            \n            g.ndata['atom_ids'] = th.tensor((len(xyz) * [i]))\n            graph_list.append(g)\n            \n            graph_dist_matrix = pd.DataFrame(np.pad(\n            rdmolops.GetDistanceMatrix(mol),\n            [(0, 0), (0, MAX_N_ATOMS - len(xyz))], 'constant'))\n            g.gdata['graph_dist_matrix'] = graph_dist_matrix\n            \n            # compute molecule level features\n            adj_matrix = rdmolops.GetAdjacencyMatrix(mol)\n            atomic_num_list, _, _ = read_xyz_file(xyz_filepath_list[i])\n            dists = dist_matrix.ravel()[np.tril(adj_matrix).ravel()==1]\n            g.gdata['mol_feat'] = pd.Series(\n            [np.mean(dists), np.std(dists), np.mean(atomic_num_list)],\n            index=mol_feat_columns\n        )\n        except RuntimeError:\n            pass\n            \n    return graph_list","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"graph_list = get_molecules()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train = pd.read_csv('train.csv')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"len(train['molecule_name'].unique())","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"'dsgdb9nsd_000001' in train['molecule_name'].unique()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train.numpy()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"graph_list[0]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"graph_list[0].gdata","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"graph_list[0].ndata['xyz']","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from dgl.data.utils import save_graphs\nsave_graphs('graph_list.bin', graph_list)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"len(graph_list)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"graph_list[0].gdata","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def map_atom_info(df, atom_idx, struct_df):\n    \"\"\"Adds xyz-coordinates of atom_{atom_idx} to 'df'.\"\"\"\n    df = pd.merge(df, struct_df, how = 'left',\n                  left_on  = ['molecule_name', f'atom_index_{atom_idx}'],\n                  right_on = ['molecule_name', 'atom_index'])\n    df = df.drop('atom_index', axis=1)\n    df = df.rename(columns={'atom': f'atom_{atom_idx}',\n                            'x': f'x_{atom_idx}',\n                            'y': f'y_{atom_idx}',\n                            'z': f'z_{atom_idx}'})\n    return df","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def add_dist(df, struct_df):\n    \"\"\"Adds euclidean distance between scalar coupling atoms to 'df'.\"\"\"\n    df = map_atom_info(df, 0, struct_df)\n    df = map_atom_info(df, 1, struct_df)\n    p_0 = df[['x_0', 'y_0', 'z_0']].values\n    p_1 = df[['x_1', 'y_1', 'z_1']].values\n    df['dist'] = np.linalg.norm(p_0 - p_1, axis=1)\n    df.drop(columns=['x_0', 'y_0', 'z_0', 'x_1', 'y_1', 'z_1'], inplace=True)\n    return df","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\ndef transform_per_atom_group(df, a_idx, col='dist', trans='mean'):\n    \"\"\"Apply transformation 'trans' on feature in 'col' to scalar coupling\n    constants grouped at the atom level.\"\"\"\n    return df.groupby(\n        ['molecule_name', f'atom_index_{a_idx}'])[col].transform(trans)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport pickle\n\nfrom dgl.data.chem.utils import mol_to_bigraph,\\\n                                CanonicalAtomFeaturizer\nimport torch\nfrom rdkit import Chem\nfrom rdkit.Chem import ChemicalFeatures, rdmolops\nfrom rdkit import RDConfig\n\nfrom glob import glob\n\n\n\nxyz_filepath_list = list(glob('structures/*.xyz'))\nxyz_filepath_list.sort()\n\n\n## Functions to create the RDKit mol objects\ndef mol_from_xyz(filepath, add_hs=True, compute_dist_centre=False):\n    \"\"\"Wrapper function for calling xyz2mol function.\"\"\"\n    charged_fragments = True  # alternatively radicals are made\n\n    # quick is faster for large systems but requires networkx\n    # if you don't want to install networkx set quick=False and\n    # uncomment 'import networkx as nx' at the top of the file\n    quick = True\n\n    atomicNumList, charge, xyz_coordinates = read_xyz_file(filepath)\n    mol, dMat = xyz2mol(atomicNumList, charge, xyz_coordinates,\n                        charged_fragments, quick, check_chiral_stereo=False)\n\n    return mol, np.array(xyz_coordinates), dMat\n\n\ndef bond_featurizer(mol, self_loop=True):\n    \"\"\"Featurization for all bonds in a molecule.\n    The bond indices will be preserved.\n    Parameters\n    ----------\n    mol : rdkit.Chem.rdchem.Mol\n        RDKit molecule object\n    self_loop : bool\n        Whether to add self loops. Default to be False.\n    Returns\n    -------\n    bond_feats_dict : dict\n        Dictionary for bond features\n    \"\"\"\n    bond_feats_dict = defaultdict(list)\n\n    mol_conformers = mol.GetConformers()\n    assert len(mol_conformers) == 1\n    geom = mol_conformers[0].GetPositions()\n\n    num_atoms = mol.GetNumAtoms()\n    for u in range(num_atoms):\n        for v in range(num_atoms):\n            if u == v and not self_loop:\n                continue\n\n            e_uv = mol.GetBondBetweenAtoms(u, v)\n            if e_uv is None:\n                continue\n            else:\n                bond_type = e_uv.GetBondType()\n            bond_feats_dict['e_feat'].append([\n                float(bond_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            ])\n            bond_feats_dict['distance'].append(\n                np.linalg.norm(geom[u] - geom[v]))\n\n    bond_feats_dict['e_feat'] = torch.tensor(\n        np.array(bond_feats_dict['e_feat']).astype(np.float32))\n    bond_feats_dict['distance'] = torch.tensor(\n        np.array(bond_feats_dict['distance']).astype(np.float32)).reshape(-1 , 1)\n\n    return bond_feats_dict\n\n\nclass KaggleMolDataset(object):\n    def __init__(self, \n                 file_list = xyz_filepath_list,\n                 label_filepath = './',\n                 store_path = './' ,\n                 mode='train', \n                 from_raw=True,\n                 mol_to_graph = mol_to_bigraph,\n                 atom_featurizer=CanonicalAtomFeaturizer,\n                 bond_featurizer=bond_featurizer):\n\n        assert mode in ['train', 'test'], \\\n           'Expect mode to be train or test, got {}.'.format(mode)\n\n        self.mode = mode\n        \n        self.from_raw = from_raw\n        \"\"\"\n        if not from_raw:\n            file_name = \"%s_processed\" % (mode)\n        else:\n            file_name = \"structures\"\n        self.file_dir = pathlib.Path(file_dir, file_name)\n        \"\"\"\n        self.file_list = file_list\n        self.store_path = store_path\n        self.label_filepath = label_filepath\n        self.graphs, self.labels = [],[]\n        self._load(mol_to_graph, atom_featurizer, bond_featurizer)\n\n    def _load(self, mol_to_graph, atom_featurizer, bond_featurizer):\n        if not self.from_raw:\n            pass\n        #    with open(osp.join(self.file_dir, \"%s_graphs.pkl\" % self.mode), \"rb\") as f:\n        #        self.graphs = pickle.load(f)\n        #    with open(osp.join(self.file_dir, \"%s_labels.pkl\" % self.mode), \"rb\") as f:\n        #        self.labels = pickle.load(f)\n        else:\n            print('Start preprocessing dataset...')\n            labels  = pd.read_csv(self.label_filepath +self.mode + '.csv')\n            cnt = 0\n            dataset_size = len(labels['molecule_name'].unique())\n            mol_names = labels['molecule_name'].unique()\n            \n            for i in range(len(self.file_list)):\n                mol_name = self.file_list[i].split('/')[-1][:-4] \n                if mol_name in mol_names:\n                    cnt += 1\n                    print('Processing molecule {:d}/{:d}'.format(cnt, dataset_size))\n                    mol, xyz, dist_matrix = mol_from_xyz(self.file_list[i])\n                    \n                    graph = mol_to_graph(mol, bond_featurizer=bond_featurizer)  \n                    graph.gdata = {}    \n                    smiles = Chem.MolToSmiles(mol)\n                    graph.gdata['smiles'] = smiles    \n                    graph.gdata['mol_name'] = mol_name \n                    graph.ndata['h'] = torch.cat([graph.ndata['h'], torch.tensor(xyz).float()],\n                                                  dim = 1)\n                    self.graphs.append(graph)\n                    label = labels[labels['molecule_name'] ==mol_name ].drop([\n                                                                        'molecule_name', \n                                                                        'type',\n                                                                        'id'\n                                                                      ],\n                                                                         axis = 1\n                                                                    )\n                    self.labels.append(label)\n\n            with open(osp.join(self.store_path, \"%s_grapgs.pkl\" % self.mode), \"wb\") as f:\n                pickle.dump(self.graphs, f)\n            with open(osp.join(self.store_path, \"%s_labels.pkl\" % self.mode), \"wb\") as f:\n                pickle.dump(self.labels, f)\n\n        print(len(self.graphs), \"loaded!\")\n\n    def __getitem__(self, item):\n        \"\"\"Get datapoint with index\n        Parameters\n        ----------\n        item : int\n            Datapoint index\n        Returns\n        -------\n        str\n            SMILES for the ith datapoint\n        DGLGraph\n            DGLGraph for the ith datapoint\n        Tensor of dtype float32\n            Labels of the datapoint for all tasks\n        \"\"\"\n        g, l = self.graphs[item], self.labels[item]\n        return g.smile, g, l\n\n    def __len__(self):\n        \"\"\"Length of the dataset\n        Returns\n        -------\n        int\n            Length of Dataset\n        \"\"\"\n        return len(self.graphs)\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"dataset = KaggleMolDataset()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from dgl.data.chem.utils import smile_to_bigraph\nsmile_to_bigraph('[H]C([H])([H])[H]')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"mol","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"mol = Chem.MolFromSmiles('[H]C([H])([H])[H]')\nmol.GetNumAtoms()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"f = CanonicalAtomFeaturizer()\nf(mol)['h'].shape","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}