torchani_test.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from rdkit import Chem\n",
    "from rdkit.Chem import AllChem\n",
    "from rdkit.Chem.Draw import IPythonConsole\n",
    "import py3Dmol"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torchani\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from psikit import Psikit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "pk = Psikit()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "smiles = 'c1ccccc1-c2c(C)cncc2'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "pk.read_from_smiles(smiles)\n",
    "pk.rdkit_optimize()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1min 10s, sys: 1.79 s, total: 1min 12s\n",
      "Wall time: 18.5 s\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "-518.5098949584047"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%time pk.energy(basis_sets='WB97X/6-31g*')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/3dmoljs_load.v0": "<div id=\"3dmolviewer_1599310546698088\"  style=\"position: relative; width: 400px; height: 400px\">\n        <p id=\"3dmolwarning_1599310546698088\" style=\"background-color:#ffcccc;color:black\">You appear to be running in JupyterLab (or JavaScript failed to load for some other reason).  You need to install the 3dmol extension: <br>\n        <tt>jupyter labextension install jupyterlab_3dmol</tt></p>\n        </div>\n<script>\n\nvar loadScriptAsync = function(uri){\n  return new Promise((resolve, reject) => {\n    var tag = document.createElement('script');\n    tag.src = uri;\n    tag.async = true;\n    tag.onload = () => {\n      resolve();\n    };\n  var firstScriptTag = document.getElementsByTagName('script')[0];\n  firstScriptTag.parentNode.insertBefore(tag, firstScriptTag);\n});\n};\n\nif(typeof $3Dmolpromise === 'undefined') {\n$3Dmolpromise = null;\n  $3Dmolpromise = loadScriptAsync('https://3dmol.csb.pitt.edu/build/3Dmol.js');\n}\n\nvar viewer_1599310546698088 = null;\nvar warn = document.getElementById(\"3dmolwarning_1599310546698088\");\nif(warn) {\n    warn.parentNode.removeChild(warn);\n}\n$3Dmolpromise.then(function() {\nviewer_1599310546698088 = $3Dmol.createViewer($(\"#3dmolviewer_1599310546698088\"),{backgroundColor:\"white\"});\n\tviewer_1599310546698088.removeAllModels();\n\tviewer_1599310546698088.addModel(\"\\n     RDKit          3D\\n\\n 24 25  0  0  0  0  0  0  0  0999 V2000\\n    1.6294    0.1673   -1.1889 C   0  0  0  0  0  0  0  0  0  0  0  0\\n    2.9678    0.5618   -1.0972 C   0  0  0  0  0  0  0  0  0  0  0  0\\n    3.5260    0.8617    0.1483 C   0  0  0  0  0  0  0  0  0  0  0  0\\n    2.7480    0.7632    1.3051 C   0  0  0  0  0  0  0  0  0  0  0  0\\n    1.4092    0.3692    1.2180 C   0  0  0  0  0  0  0  0  0  0  0  0\\n    0.8368    0.0784   -0.0321 C   0  0  0  0  0  0  0  0  0  0  0  0\\n   -0.5767   -0.3514   -0.1311 C   0  0  0  0  0  0  0  0  0  0  0  0\\n   -1.6452    0.4632    0.3162 C   0  0  0  0  0  0  0  0  0  0  0  0\\n   -1.4205    1.8350    0.8956 C   0  0  0  0  0  0  0  0  0  0  0  0\\n   -2.9549   -0.0360    0.2283 C   0  0  0  0  0  0  0  0  0  0  0  0\\n   -3.1912   -1.2610   -0.3093 N   0  0  0  0  0  0  0  0  0  0  0  0\\n   -2.1839   -2.0446   -0.7712 C   0  0  0  0  0  0  0  0  0  0  0  0\\n   -0.8629   -1.6054   -0.6917 C   0  0  0  0  0  0  0  0  0  0  0  0\\n    1.2082   -0.0600   -2.1606 H   0  0  0  0  0  0  0  0  0  0  0  0\\n    3.5723    0.6347   -1.9920 H   0  0  0  0  0  0  0  0  0  0  0  0\\n    4.5624    1.1661    0.2177 H   0  0  0  0  0  0  0  0  0  0  0  0\\n    3.1837    0.9876    2.2701 H   0  0  0  0  0  0  0  0  0  0  0  0\\n    0.8238    0.2791    2.1243 H   0  0  0  0  0  0  0  0  0  0  0  0\\n   -1.1603    1.7485    1.9708 H   0  0  0  0  0  0  0  0  0  0  0  0\\n   -0.6039    2.3605    0.3574 H   0  0  0  0  0  0  0  0  0  0  0  0\\n   -2.3341    2.4602    0.8029 H   0  0  0  0  0  0  0  0  0  0  0  0\\n   -3.7922    0.5535    0.5797 H   0  0  0  0  0  0  0  0  0  0  0  0\\n   -2.4084   -3.0143   -1.1949 H   0  0  0  0  0  0  0  0  0  0  0  0\\n   -0.0662   -2.2463   -1.0500 H   0  0  0  0  0  0  0  0  0  0  0  0\\n  1  2  2  0\\n  2  3  1  0\\n  3  4  2  0\\n  4  5  1  0\\n  5  6  2  0\\n  6  7  1  0\\n  7  8  2  0\\n  8  9  1  0\\n  8 10  1  0\\n 10 11  2  0\\n 11 12  1  0\\n 12 13  2  0\\n  6  1  1  0\\n 13  7  1  0\\n  1 14  1  0\\n  2 15  1  0\\n  3 16  1  0\\n  4 17  1  0\\n  5 18  1  0\\n  9 19  1  0\\n  9 20  1  0\\n  9 21  1  0\\n 10 22  1  0\\n 12 23  1  0\\n 13 24  1  0\\nM  END\\n\",\"sdf\");\n\tviewer_1599310546698088.setStyle({\"stick\": {}});\n\tviewer_1599310546698088.setBackgroundColor(\"0xeeeeee\");\n\tviewer_1599310546698088.zoomTo();\nviewer_1599310546698088.render();\n});\n</script>",
      "text/html": [
       "<div id=\"3dmolviewer_1599310546698088\"  style=\"position: relative; width: 400px; height: 400px\">\n",
       "        <p id=\"3dmolwarning_1599310546698088\" style=\"background-color:#ffcccc;color:black\">You appear to be running in JupyterLab (or JavaScript failed to load for some other reason).  You need to install the 3dmol extension: <br>\n",
       "        <tt>jupyter labextension install jupyterlab_3dmol</tt></p>\n",
       "        </div>\n",
       "<script>\n",
       "\n",
       "var loadScriptAsync = function(uri){\n",
       "  return new Promise((resolve, reject) => {\n",
       "    var tag = document.createElement('script');\n",
       "    tag.src = uri;\n",
       "    tag.async = true;\n",
       "    tag.onload = () => {\n",
       "      resolve();\n",
       "    };\n",
       "  var firstScriptTag = document.getElementsByTagName('script')[0];\n",
       "  firstScriptTag.parentNode.insertBefore(tag, firstScriptTag);\n",
       "});\n",
       "};\n",
       "\n",
       "if(typeof $3Dmolpromise === 'undefined') {\n",
       "$3Dmolpromise = null;\n",
       "  $3Dmolpromise = loadScriptAsync('https://3dmol.csb.pitt.edu/build/3Dmol.js');\n",
       "}\n",
       "\n",
       "var viewer_1599310546698088 = null;\n",
       "var warn = document.getElementById(\"3dmolwarning_1599310546698088\");\n",
       "if(warn) {\n",
       "    warn.parentNode.removeChild(warn);\n",
       "}\n",
       "$3Dmolpromise.then(function() {\n",
       "viewer_1599310546698088 = $3Dmol.createViewer($(\"#3dmolviewer_1599310546698088\"),{backgroundColor:\"white\"});\n",
       "\tviewer_1599310546698088.removeAllModels();\n",
       "\tviewer_1599310546698088.addModel(\"\\n     RDKit          3D\\n\\n 24 25  0  0  0  0  0  0  0  0999 V2000\\n    1.6294    0.1673   -1.1889 C   0  0  0  0  0  0  0  0  0  0  0  0\\n    2.9678    0.5618   -1.0972 C   0  0  0  0  0  0  0  0  0  0  0  0\\n    3.5260    0.8617    0.1483 C   0  0  0  0  0  0  0  0  0  0  0  0\\n    2.7480    0.7632    1.3051 C   0  0  0  0  0  0  0  0  0  0  0  0\\n    1.4092    0.3692    1.2180 C   0  0  0  0  0  0  0  0  0  0  0  0\\n    0.8368    0.0784   -0.0321 C   0  0  0  0  0  0  0  0  0  0  0  0\\n   -0.5767   -0.3514   -0.1311 C   0  0  0  0  0  0  0  0  0  0  0  0\\n   -1.6452    0.4632    0.3162 C   0  0  0  0  0  0  0  0  0  0  0  0\\n   -1.4205    1.8350    0.8956 C   0  0  0  0  0  0  0  0  0  0  0  0\\n   -2.9549   -0.0360    0.2283 C   0  0  0  0  0  0  0  0  0  0  0  0\\n   -3.1912   -1.2610   -0.3093 N   0  0  0  0  0  0  0  0  0  0  0  0\\n   -2.1839   -2.0446   -0.7712 C   0  0  0  0  0  0  0  0  0  0  0  0\\n   -0.8629   -1.6054   -0.6917 C   0  0  0  0  0  0  0  0  0  0  0  0\\n    1.2082   -0.0600   -2.1606 H   0  0  0  0  0  0  0  0  0  0  0  0\\n    3.5723    0.6347   -1.9920 H   0  0  0  0  0  0  0  0  0  0  0  0\\n    4.5624    1.1661    0.2177 H   0  0  0  0  0  0  0  0  0  0  0  0\\n    3.1837    0.9876    2.2701 H   0  0  0  0  0  0  0  0  0  0  0  0\\n    0.8238    0.2791    2.1243 H   0  0  0  0  0  0  0  0  0  0  0  0\\n   -1.1603    1.7485    1.9708 H   0  0  0  0  0  0  0  0  0  0  0  0\\n   -0.6039    2.3605    0.3574 H   0  0  0  0  0  0  0  0  0  0  0  0\\n   -2.3341    2.4602    0.8029 H   0  0  0  0  0  0  0  0  0  0  0  0\\n   -3.7922    0.5535    0.5797 H   0  0  0  0  0  0  0  0  0  0  0  0\\n   -2.4084   -3.0143   -1.1949 H   0  0  0  0  0  0  0  0  0  0  0  0\\n   -0.0662   -2.2463   -1.0500 H   0  0  0  0  0  0  0  0  0  0  0  0\\n  1  2  2  0\\n  2  3  1  0\\n  3  4  2  0\\n  4  5  1  0\\n  5  6  2  0\\n  6  7  1  0\\n  7  8  2  0\\n  8  9  1  0\\n  8 10  1  0\\n 10 11  2  0\\n 11 12  1  0\\n 12 13  2  0\\n  6  1  1  0\\n 13  7  1  0\\n  1 14  1  0\\n  2 15  1  0\\n  3 16  1  0\\n  4 17  1  0\\n  5 18  1  0\\n  9 19  1  0\\n  9 20  1  0\\n  9 21  1  0\\n 10 22  1  0\\n 12 23  1  0\\n 13 24  1  0\\nM  END\\n\",\"sdf\");\n",
       "\tviewer_1599310546698088.setStyle({\"stick\": {}});\n",
       "\tviewer_1599310546698088.setBackgroundColor(\"0xeeeeee\");\n",
       "\tviewer_1599310546698088.zoomTo();\n",
       "viewer_1599310546698088.render();\n",
       "});\n",
       "</script>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "IPythonConsole.drawMol3D(pk.mol)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = torchani.models.ANI2x(periodic_table_index=True).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mol2arr(mols, device=device):\n",
    "    coordinates = []\n",
    "    spices = []\n",
    "    for mol in mols:\n",
    "        pos = mol.GetConformer().GetPositions().tolist()\n",
    "        atomnums = [a.GetAtomicNum() for a in mol.GetAtoms()]\n",
    "        coordinates.append(pos)\n",
    "        spices.append(atomnums)\n",
    "    coordinates = torch.tensor(coordinates,\n",
    "                               requires_grad=True,\n",
    "                               device=device)\n",
    "    species = torch.tensor(spices, device=device)\n",
    "    return coordinates, species"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "coordinates, species = mol2arr([pk.mol], device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 16.5 ms, sys: 15 µs, total: 16.5 ms\n",
      "Wall time: 16.2 ms\n"
     ]
    }
   ],
   "source": [
    "%time energy = model((species, coordinates)).energies\n",
    "derivative = torch.autograd.grad(energy.sum(), coordinates)[0]\n",
    "force = -derivative"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Energy: -518.5092850781618\n",
      "Force: tensor([[ 1.3592e-02,  3.5025e-03,  2.0906e-02],\n",
      "        [-1.4324e-02, -1.6894e-04,  1.4033e-02],\n",
      "        [-1.3329e-02, -3.7380e-03,  2.3463e-04],\n",
      "        [-9.2622e-03, -5.9881e-03, -1.5998e-02],\n",
      "        [ 2.1978e-02,  4.3879e-03, -1.6717e-02],\n",
      "        [ 1.4468e-02,  7.4550e-03,  3.5641e-03],\n",
      "        [-2.8785e-02, -9.5195e-03, -5.3672e-03],\n",
      "        [-2.6472e-03, -3.8259e-03,  8.0281e-04],\n",
      "        [-9.3185e-04,  1.7754e-02, -1.3840e-03],\n",
      "        [ 6.6303e-02, -7.4541e-02, -4.3813e-02],\n",
      "        [-1.2796e-02, -1.2576e-03,  1.3008e-03],\n",
      "        [-2.4707e-03,  7.7882e-02,  3.8845e-02],\n",
      "        [-1.5166e-02,  6.1120e-03,  1.5764e-03],\n",
      "        [-5.9885e-03, -2.7664e-03, -4.3181e-03],\n",
      "        [ 2.6414e-03, -5.7588e-06, -5.2600e-03],\n",
      "        [ 5.9540e-03,  1.8318e-03,  4.1510e-04],\n",
      "        [ 1.0910e-03,  1.0532e-03,  4.8398e-03],\n",
      "        [-7.2534e-03, -1.1427e-03,  3.0659e-03],\n",
      "        [-1.5687e-03,  1.0146e-02, -1.6901e-02],\n",
      "        [-1.1689e-02, -1.0652e-02,  1.5877e-02],\n",
      "        [ 1.5274e-02, -1.3123e-02,  6.3923e-03],\n",
      "        [-9.5655e-03, -1.6000e-03, -5.9007e-04],\n",
      "        [-1.3352e-02, -2.3642e-03, -2.1185e-04],\n",
      "        [ 7.8282e-03,  5.6739e-04, -1.2904e-03]], device='cuda:0')\n"
     ]
    }
   ],
   "source": [
    "print('Energy:', energy.item())\n",
    "print('Force:', force.squeeze())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[ 1.3592e-02,  3.5025e-03,  2.0906e-02],\n",
       "         [-1.4324e-02, -1.6894e-04,  1.4033e-02],\n",
       "         [-1.3329e-02, -3.7380e-03,  2.3463e-04],\n",
       "         [-9.2622e-03, -5.9881e-03, -1.5998e-02],\n",
       "         [ 2.1978e-02,  4.3879e-03, -1.6717e-02],\n",
       "         [ 1.4468e-02,  7.4550e-03,  3.5641e-03],\n",
       "         [-2.8785e-02, -9.5195e-03, -5.3672e-03],\n",
       "         [-2.6472e-03, -3.8259e-03,  8.0281e-04],\n",
       "         [-9.3185e-04,  1.7754e-02, -1.3840e-03],\n",
       "         [ 6.6303e-02, -7.4541e-02, -4.3813e-02],\n",
       "         [-1.2796e-02, -1.2576e-03,  1.3008e-03],\n",
       "         [-2.4707e-03,  7.7882e-02,  3.8845e-02],\n",
       "         [-1.5166e-02,  6.1120e-03,  1.5764e-03],\n",
       "         [-5.9885e-03, -2.7664e-03, -4.3181e-03],\n",
       "         [ 2.6414e-03, -5.7588e-06, -5.2600e-03],\n",
       "         [ 5.9540e-03,  1.8318e-03,  4.1510e-04],\n",
       "         [ 1.0910e-03,  1.0532e-03,  4.8398e-03],\n",
       "         [-7.2534e-03, -1.1427e-03,  3.0659e-03],\n",
       "         [-1.5687e-03,  1.0146e-02, -1.6901e-02],\n",
       "         [-1.1689e-02, -1.0652e-02,  1.5877e-02],\n",
       "         [ 1.5274e-02, -1.3123e-02,  6.3923e-03],\n",
       "         [-9.5655e-03, -1.6000e-03, -5.9007e-04],\n",
       "         [-1.3352e-02, -2.3642e-03, -2.1185e-04],\n",
       "         [ 7.8282e-03,  5.6739e-04, -1.2904e-03]]], device='cuda:0')"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "force"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}