Fig3c+Fig6+tSNE+attention_calc.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "30958abf",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using backend: tensorflow\n"
     ]
    }
   ],
   "source": [
    "import pandas as pd\n",
    "import tensorflow as tf\n",
    "import rdkit\n",
    "from main import create_input, create_model, predict\n",
    "from argparse import ArgumentParser\n",
    "from molgraph import single_mol_from_supergraph\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "93588a95",
   "metadata": {},
   "source": [
    "# Run predictive model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "f2c00078",
   "metadata": {},
   "outputs": [],
   "source": [
    "parser = ArgumentParser()\n",
    "parser.add_argument('-predict', action=\"store_true\", default=True)\n",
    "parser.add_argument('-watsoneq', action='store_true', default=False)\n",
    "parser.add_argument('-K_fold', action='store_true', default=False)\n",
    "parser.add_argument('-maxatoms', type=int, default=64)\n",
    "parser.add_argument('-lr', type=float, default=5.0e-4)\n",
    "parser.add_argument('-epoch', type=int, default=100)\n",
    "parser.add_argument('-batchsize', type=int, default=256)\n",
    "parser.add_argument('-layers', type=int, default=5)\n",
    "parser.add_argument('-heads', type=int, default=5)\n",
    "parser.add_argument('-residcon', action=\"store_true\", default=True)\n",
    "parser.add_argument('-explicitH', action=\"store_true\", default=False)\n",
    "parser.add_argument('-dropout', type=float, default=0.0)\n",
    "parser.add_argument('-modelname', type=str, default='best_211007')\n",
    "parser.add_argument('-num_hidden', type=int, default=32)\n",
    "parser.add_argument('-train_only', action=\"store_true\", default=False)\n",
    "parser.add_argument('-sw_thr', type=float, default=0.0)\n",
    "parser.add_argument('-sw_decay', type=int, default=1)\n",
    "parser.add_argument('-loss', type=str, default='kl_div_normal')\n",
    "args = parser.parse_args('')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "fd76e4dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.read_csv('data/Data_211005.csv')\n",
    "tmp = []\n",
    "#for atom feature t-SNE and atom attention weights (no temperature dependence)\n",
    "for smi in df.smiles.unique():\n",
    "    tmp.append([smi,298.15])\n",
    "pd.DataFrame(tmp).to_csv('molecules_to_predict.csv', header=['smiles','temperature'], index=False)\n",
    "\n",
    "\n",
    "#for temperature dependence\n",
    "#pd.DataFrame(df[['smiles','Train/Valid/Test','temperature','HoV (kJ/mol)']]).to_csv('molecules_to_predict.csv', index=False)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "f172106a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<gnn.GAT_unc object at 0x7fa26e475e90>\n",
      "0\n",
      "1\n",
      "2\n",
      "3\n",
      "4\n"
     ]
    }
   ],
   "source": [
    "device ='/cpu:0'\n",
    "data = pd.read_csv('molecules_to_predict.csv')\n",
    "data['total_atoms'] = [ rdkit.Chem.MolFromSmiles(smi).GetNumHeavyAtoms() for smi in data.smiles]\n",
    "INPUT = create_input(data, device, args)\n",
    "data, num_mols, T, Graphs, seg = INPUT\n",
    "atom_feat_dim = Graphs.ndata['feat'].shape[-1]\n",
    "model, model_name = create_model(args, atom_feat_dim, '')\n",
    "################\n",
    "Predicted_HoV, atom_features_each_layer, Attention_each_layer, T_part, updated_T = predict(model, Graphs.ndata['feat'], Graphs, \\\n",
    "                                                    seg, args.maxatoms, T, '', num_mols, mu_s_NLR=[])\n",
    "pred_mean, pred_stddev = Predicted_HoV[:, 0], Predicted_HoV[:, 1]\n",
    "data['Predicted'] = pred_mean.numpy()\n",
    "data['ML_unc'] = pred_stddev.numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "3812a2fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "#If full DB with temperature dependence is tested\n",
    "#valid = data[data['Train/Valid/Test'] == 'Valid']\n",
    "#np.abs(valid.Predicted - valid['HoV (kJ/mol)']).mean()\n",
    "#train = data[data['Train/Valid/Test'] == 'Train']\n",
    "#np.abs(train.Predicted - train['HoV (kJ/mol)']).mean()\n",
    "#test = data[data['Train/Valid/Test'] == 'Test']\n",
    "#np.abs(test.Predicted - test['HoV (kJ/mol)']).mean()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74c9a66e",
   "metadata": {},
   "source": [
    "# Run t-SNE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "4513dd33",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "94200"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "indices_nonzeros = []\n",
    "\n",
    "start = 0\n",
    "for tot_atom in data.total_atoms:\n",
    "    indices_nonzeros += list(range(start, start+tot_atom))\n",
    "    start += 64\n",
    "\n",
    "#initial_onehot = Graphs.ndata['feat'].numpy()[indices_nonzeros]\n",
    "atom_features_layer1 = atom_features_each_layer[0].numpy()[indices_nonzeros]\n",
    "atom_features_layer3 = atom_features_each_layer[2].numpy()[indices_nonzeros]\n",
    "atom_features_layer5 = atom_features_each_layer[4].numpy()[indices_nonzeros]\n",
    "\n",
    "len(atom_features_layer1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "b0b00693",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.decomposition import PCA\n",
    "from sklearn.pipeline import Pipeline\n",
    "from sklearn.manifold import TSNE\n",
    "pipe = Pipeline(steps=(\n",
    "    ('PCA', PCA(n_components=10)),\n",
    "    ('TSNE', TSNE(n_components=2))\n",
    "))\n",
    "\n",
    "#xt_onehot = pipe.fit_transform(initial_onehot)\n",
    "xt_1 = pipe.fit_transform(atom_features_layer1)\n",
    "xt_3 = pipe.fit_transform(atom_features_layer3)\n",
    "xt_5 = pipe.fit_transform(atom_features_layer5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "7103ccde",
   "metadata": {},
   "outputs": [],
   "source": [
    "feat_vectors_2d = np.concatenate([xt_1,xt_3,xt_5], axis=1)\n",
    "df = pd.DataFrame(feat_vectors_2d, columns = ['xt_1_x', 'xt_1_y',\n",
    "                                         'xt_3_x', 'xt_3_y',\n",
    "                                        'xt_5_x', 'xt_5_y'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53b570d4",
   "metadata": {},
   "source": [
    "# Attention score analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "d06b92ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "Att_each_layer_results = []\n",
    "\n",
    "for l in range(len(Attention_each_layer)):\n",
    "    A_per_each_mol, src_per_each_mol, dst_per_each_mol = single_mol_from_supergraph(Graphs, \n",
    "                                                                                    atom_features_each_layer[l], \n",
    "                                                                                    Attention_each_layer[l], \n",
    "                                                                                    Max_atoms = 64)\n",
    "    Att_each_layer_results.append([A_per_each_mol, src_per_each_mol, dst_per_each_mol])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "8a945550",
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_attention_matrix(A, src, dst):\n",
    "    A_matrix = np.zeros((max(src) + 1, max(src) +1 ))\n",
    "    for i in range(len(A)):\n",
    "        A_matrix[src[i], dst[i]] = A[i]\n",
    "    return A_matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "02ae7768",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "94200 94200\n"
     ]
    }
   ],
   "source": [
    "unique_smiles = list(data.smiles.unique())\n",
    "unique_smiles_numatoms = [rdkit.Chem.MolFromSmiles(smi).GetNumHeavyAtoms() for smi in unique_smiles]\n",
    "\n",
    "atom_indices_each_molecule = []\n",
    "each_atom_belongs_to_which_smiles = []\n",
    "\n",
    "for i in range(len(unique_smiles)):\n",
    "    atom_indices_each_molecule += list(range(unique_smiles_numatoms[i]))\n",
    "    each_atom_belongs_to_which_smiles += [unique_smiles[i]] * unique_smiles_numatoms[i]\n",
    "\n",
    "print(len(atom_indices_each_molecule), len(each_atom_belongs_to_which_smiles))\n",
    "    \n",
    "Aij_rowsum_all_layers = []\n",
    "for i in range(len(Att_each_layer_results)):\n",
    "    Aij_rowsum = []\n",
    "    \n",
    "    Att_ith = Att_each_layer_results[i]\n",
    "    A_per_each_mol, src_per_each_mol, dst_per_each_mol = Att_ith\n",
    "    \n",
    "    for molecule_index in range(len(A_per_each_mol)):\n",
    "        A, src, dst = A_per_each_mol[molecule_index], \\\n",
    "                      src_per_each_mol[molecule_index], \\\n",
    "                      dst_per_each_mol[molecule_index]\n",
    "        \n",
    "        A_matrix = make_attention_matrix(A, src, dst)\n",
    "        \n",
    "        if unique_smiles_numatoms[molecule_index] != len(A_matrix):\n",
    "            print('???????????????')\n",
    "        \n",
    "        Aij_rowsum_one_molecule = []\n",
    "        for atom_index in range(len(A_matrix)):\n",
    "            #Aij_rowsum.append(np.sum(A_matrix[atom_index]) / np.count_nonzero(A_matrix[atom_index]))\n",
    "            Aij_rowsum_one_molecule.append(np.sum(A_matrix[atom_index]) / len(A_matrix))\n",
    "        Aij_rowsum_one_molecule = [ x / max(Aij_rowsum_one_molecule) for x in Aij_rowsum_one_molecule]\n",
    "        \n",
    "        Aij_rowsum += Aij_rowsum_one_molecule\n",
    "    Aij_rowsum_all_layers.append(Aij_rowsum)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "9de8726e",
   "metadata": {},
   "outputs": [],
   "source": [
    "df['smiles'] = each_atom_belongs_to_which_smiles\n",
    "df['atom_index'] = atom_indices_each_molecule\n",
    "\n",
    "for i in range(len(Aij_rowsum_all_layers)):\n",
    "    df['Att_'+str(i+1)] = Aij_rowsum_all_layers[i]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eca0fa74",
   "metadata": {},
   "source": [
    "# Atom type/substructure analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "dc534421",
   "metadata": {},
   "outputs": [],
   "source": [
    "from rdkit import Chem\n",
    "\n",
    "def substruc_search(smiles, atom_index, radius=1):\n",
    "    m = Chem.MolFromSmiles(smiles)\n",
    "    \n",
    "    env = Chem.FindAtomEnvironmentOfRadiusN(m,radius,int(atom_index))\n",
    "    amap={}\n",
    "    submol=Chem.PathToSubmol(m,env,atomMap=amap)\n",
    "    \n",
    "    \n",
    "    mol_smarts=Chem.MolToSmarts(submol)\n",
    "    \n",
    "    return mol_smarts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "07ac84c3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "94200\n"
     ]
    }
   ],
   "source": [
    "smarts_radius_1_each_atom = []\n",
    "for _, row in df.iterrows():\n",
    "    a_ind, smi = row['atom_index'], row['smiles']\n",
    "    smarts = substruc_search(smi, a_ind)\n",
    "    smarts_radius_1_each_atom.append(smarts)\n",
    "    \n",
    "print(len(smarts_radius_1_each_atom))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "8f15d52a",
   "metadata": {},
   "outputs": [],
   "source": [
    "df['smarts'] = smarts_radius_1_each_atom"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "505f7fdc",
   "metadata": {},
   "outputs": [],
   "source": [
    "def atom_type(smi, a_ind):\n",
    "    mol = rdkit.Chem.MolFromSmiles(smi)\n",
    "        \n",
    "    atom = mol.GetAtomWithIdx(a_ind)    \n",
    "    \n",
    "    atom_type = ''\n",
    "    atom_type += atom.GetSymbol()\n",
    "\n",
    "    num_bonds = atom.GetDegree() + atom.GetTotalNumHs()\n",
    "\n",
    "    if atom.GetSymbol() == 'C':\n",
    "        if num_bonds == 4:\n",
    "            atom_type += 'sp3'\n",
    "        elif num_bonds == 3:\n",
    "            atom_type += 'sp2'\n",
    "        elif num_bonds <= 2:\n",
    "            atom_type += 'sp'\n",
    "    elif atom.GetSymbol() == 'O':\n",
    "        if num_bonds == 2:\n",
    "            atom_type += 'sp3'\n",
    "        elif num_bonds == 1:\n",
    "            atom_type += 'sp2'\n",
    "        else:\n",
    "            print(\"?!?!?!\")\n",
    "\n",
    "    if atom.GetIsAromatic():\n",
    "        atom_type += '(aro)'\n",
    "    elif atom.IsInRing():\n",
    "        atom_type += '(ring)'\n",
    "\n",
    "    if atom_type == 'Csp(aro)':\n",
    "        atom_type = 'C(ring)'\n",
    "    elif atom_type == 'Csp2(aro)':\n",
    "        atom_type = 'C(aro)'\n",
    "    elif atom_type == 'Osp3(aro)':\n",
    "        atom_type = 'O(aro)'\n",
    "\n",
    "    if atom_type == 'Csp(ring)' or atom_type == 'Csp2(ring)' or atom_type == 'Csp3(ring)':\n",
    "        atom_type = 'C(ring)'\n",
    "\n",
    "    if atom_type == 'Osp3':\n",
    "        atom_type = '-O-'\n",
    "    if atom_type == 'Osp2':\n",
    "        atom_type = 'O='\n",
    "    if atom_type == 'Osp3(ring)':\n",
    "        atom_type = 'O(ring)'\n",
    "        \n",
    "    return atom_type"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "d2c9d8d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "df['atom_type'] = [atom_type( row['smiles'], row['atom_index']) for _, row in df.iterrows() ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "1e4a2795",
   "metadata": {},
   "outputs": [],
   "source": [
    "df.to_csv('feat_vectors_2d+Att.csv', index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "687e3c8f",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}