alignment_example.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/iljung0810/anaconda3/envs/DiffAlign/lib/python3.9/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import copy\n",
    "import torch\n",
    "import py3Dmol\n",
    "\n",
    "from rdkit import Chem\n",
    "from torch_geometric.data import Batch\n",
    "from diffalign.utils.chem import set_rdmol_positions, mol_to_graph_data_obj\n",
    "from diffalign.models.epsnet.diffusion import DiffAlign"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "device = 'cuda:0'\n",
    "model = DiffAlign().to(device)\n",
    "model.load_state_dict(torch.load(\"./param/cross_align.pt\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "mol = Chem.SDMolSupplier(f\"./mol/ligand_init.sdf\", sanitize=False)[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "sample: 0it [00:00, ?it/s]/home/iljung0810/Lab/DiffAlign_workspace/DiffAlign/diffalign/models/epsnet/diffusion.py:18: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\n",
      "  merge_batch.graph_idx = torch.tensor(merge_batch.batch)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/iljung0810/Lab/DiffAlign_workspace/DiffAlign/diffalign/models/common.py:123: UserWarning: torch.sparse.SparseTensor(indices, values, shape, *, device=) is deprecated.  Please use torch.sparse_coo_tensor(indices, values, shape, dtype=, device=). (Triggered internally at /opt/conda/conda-bld/pytorch_1708025847094/work/torch/csrc/utils/tensor_new.cpp:618.)\n",
      "  bgraph_adj = torch.sparse.LongTensor(\n",
      "sample: 199it [00:02, 71.26it/s]\n"
     ]
    }
   ],
   "source": [
    "batch1 = mol_to_graph_data_obj(mol)\n",
    "batch2 = copy.deepcopy(batch1)\n",
    "\n",
    "mean_pos = batch2['pos'].mean(0)\n",
    "batch2['pos'] = batch2['pos'] - mean_pos\n",
    "query_batch = Batch.from_data_list([batch1]).cuda()\n",
    "reference_batch = Batch.from_data_list([batch2]).cuda()\n",
    "query_batch.pos.normal_()\n",
    "\n",
    "model.eval()\n",
    "\n",
    "pos_gen, pos_gen_traj = model.DDIM_CFG_Sampling(\n",
    "    query_batch=query_batch,\n",
    "    reference_batch=reference_batch,\n",
    "    n_steps=200\n",
    ")\n",
    "reference_batch['pos'] = reference_batch['pos'] + mean_pos.cuda()\n",
    "pos_gen = pos_gen.cpu() + mean_pos.cpu()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "init_mol = set_rdmol_positions(mol, reference_batch[0]['pos'])\n",
    "gen_mol = set_rdmol_positions(mol, pos_gen)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/3dmoljs_load.v0": "<div id=\"3dmolviewer_17333795616564157\"  style=\"position: relative; width: 500px; height: 500px;\">\n        <p id=\"3dmolwarning_17333795616564157\" style=\"background-color:#ffcccc;color:black\">3Dmol.js failed to load for some reason.  Please check your browser console for error messages.<br></p>\n        </div>\n<script>\n\nvar loadScriptAsync = function(uri){\n  return new Promise((resolve, reject) => {\n    //this is to ignore the existence of requirejs amd\n    var savedexports, savedmodule;\n    if (typeof exports !== 'undefined') savedexports = exports;\n    else exports = {}\n    if (typeof module !== 'undefined') savedmodule = module;\n    else module = {}\n\n    var tag = document.createElement('script');\n    tag.src = uri;\n    tag.async = true;\n    tag.onload = () => {\n        exports = savedexports;\n        module = savedmodule;\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://cdnjs.cloudflare.com/ajax/libs/3Dmol/2.3.0/3Dmol-min.js');\n}\n\nvar viewer_17333795616564157 = null;\nvar warn = document.getElementById(\"3dmolwarning_17333795616564157\");\nif(warn) {\n    warn.parentNode.removeChild(warn);\n}\n$3Dmolpromise.then(function() {\nviewer_17333795616564157 = $3Dmol.createViewer(document.getElementById(\"3dmolviewer_17333795616564157\"),{backgroundColor:\"white\"});\nviewer_17333795616564157.zoomTo();\n\tviewer_17333795616564157.addModel(\"1a69_ligand\\n     RDKit          3D\\n\\n 31 33  0  0  1  0  0  0  0  0999 V2000\\n  -18.6290   72.4960   64.7810 N   0  0  0  0  0  0  0  0  0  0  0  0\\n  -17.6590   72.8430   65.6920 C   0  0  0  0  0  0  0  0  0  0  0  0\\n  -17.4340   72.1880   66.8280 N   0  0  0  0  0  0  0  0  0  0  0  0\\n  -18.2650   71.1350   66.9710 C   0  0  0  0  0  0  0  0  0  0  0  0\\n  -19.2640   70.7080   66.1180 C   0  0  0  0  0  0  0  0  0  0  0  0\\n  -19.4960   71.4120   64.9200 C   0  0  0  0  0  0  0  0  0  0  0  0\\n  -20.3350   71.1770   64.0460 O   0  0  0  0  0  0  0  0  0  0  0  0\\n  -19.9000   69.5840   66.6230 N   0  0  0  0  0  2  0  0  0  0  0  0\\n  -19.2890   69.3560   67.7600 N   0  0  0  0  0  4  0  0  0  0  0  0\\n  -18.2820   70.2630   68.0340 C   0  0  0  0  0  0  0  0  0  0  0  0\\n  -17.5510   69.7970   69.2390 C   0  0  1  0  0  0  0  0  0  0  0  0\\n  -16.4120   70.7300   69.6340 C   0  0  2  0  0  0  0  0  0  0  0  0\\n  -15.2040   69.9950   69.6560 O   0  0  0  0  0  0  0  0  0  0  0  0\\n  -16.8460   71.3300   70.9950 C   0  0  2  0  0  0  0  0  0  0  0  0\\n  -16.1160   70.8000   72.1220 O   0  0  0  0  0  0  0  0  0  0  0  0\\n  -18.3190   70.9540   71.1130 C   0  0  2  0  0  0  0  0  0  0  0  0\\n  -18.4590   69.7460   70.3290 O   0  0  0  0  0  0  0  0  0  0  0  0\\n  -19.2870   72.0220   70.6090 C   0  0  1  0  0  0  0  0  0  0  0  0\\n  -20.6560   71.5800   70.5560 O   0  0  0  0  0  0  0  0  0  0  0  0\\n  -18.7201   73.0664   63.9525 H   0  0  0  0  0  0  0  0  0  0  0  0\\n  -17.0412   73.7009   65.4715 H   0  0  0  0  0  0  0  0  0  0  0  0\\n  -17.1523   68.7993   69.0552 H   0  0  0  0  0  0  0  0  0  0  0  0\\n  -16.3319   71.5302   68.8982 H   0  0  0  0  0  0  0  0  0  0  0  0\\n  -16.7442   72.4149   70.9680 H   0  0  0  0  0  0  0  0  0  0  0  0\\n  -18.5469   70.7304   72.1552 H   0  0  0  0  0  0  0  0  0  0  0  0\\n  -19.2283   72.8823   71.2757 H   0  0  0  0  0  0  0  0  0  0  0  0\\n  -18.9793   72.3309   69.6100 H   0  0  0  0  0  0  0  0  0  0  0  0\\n  -19.5302   68.5925   68.3757 H   0  0  0  0  0  0  0  0  0  0  0  0\\n  -21.2109   72.2936   70.2328 H   0  0  0  0  0  0  0  0  0  0  0  0\\n  -15.1894   71.0386   72.0444 H   0  0  0  0  0  0  0  0  0  0  0  0\\n  -14.4804   70.5751   69.9040 H   0  0  0  0  0  0  0  0  0  0  0  0\\n  1  2  1  0\\n  1  6  1  0\\n  2  3  2  0\\n  3  4  1  0\\n  4  5  2  0\\n  4 10  1  0\\n  5  6  1  0\\n  5  8  1  0\\n  6  7  2  0\\n  8  9  1  0\\n  9 10  2  0\\n 10 11  1  0\\n 11 12  1  0\\n 11 17  1  0\\n 11 22  1  6\\n 12 13  1  0\\n 12 14  1  0\\n 12 23  1  6\\n 14 15  1  0\\n 14 16  1  0\\n 14 24  1  6\\n 16 17  1  0\\n 16 18  1  0\\n 16 25  1  1\\n 18 19  1  0\\n 20  1  1  0\\n 21  2  1  0\\n 18 26  1  1\\n 27 18  1  0\\n 28  9  1  0\\n 29 19  1  0\\n 30 15  1  0\\n 31 13  1  0\\nM  END\\n\",\"mol\");\n\tviewer_17333795616564157.addModel(\"1a69_ligand\\n     RDKit          3D\\n\\n 31 33  0  0  1  0  0  0  0  0999 V2000\\n  -18.5428   72.4761   64.7699 N   0  0  0  0  0  0  0  0  0  0  0  0\\n  -17.5279   72.8025   65.7890 C   0  0  0  0  0  0  0  0  0  0  0  0\\n  -17.4808   72.0211   66.9008 N   0  0  0  0  0  0  0  0  0  0  0  0\\n  -18.2611   71.0120   67.0622 C   0  0  0  0  0  0  0  0  0  0  0  0\\n  -19.2738   70.6054   66.1560 C   0  0  0  0  0  0  0  0  0  0  0  0\\n  -19.4461   71.4386   64.9924 C   0  0  0  0  0  0  0  0  0  0  0  0\\n  -20.3034   71.2558   64.0788 O   0  0  0  0  0  0  0  0  0  0  0  0\\n  -19.8618   69.5291   66.6005 N   0  0  0  0  0  2  0  0  0  0  0  0\\n  -19.2045   69.2367   67.8442 N   0  0  0  0  0  4  0  0  0  0  0  0\\n  -18.2542   70.1145   68.1358 C   0  0  0  0  0  0  0  0  0  0  0  0\\n  -17.4280   69.7394   69.3457 C   0  0  1  0  0  0  0  0  0  0  0  0\\n  -16.4033   70.6036   69.8435 C   0  0  2  0  0  0  0  0  0  0  0  0\\n  -15.2576   69.8840   70.4399 O   0  0  0  0  0  0  0  0  0  0  0  0\\n  -16.9491   71.3994   70.9893 C   0  0  1  0  0  0  0  0  0  0  0  0\\n  -17.1736   72.7079   70.8005 O   0  0  0  0  0  0  0  0  0  0  0  0\\n  -18.2152   70.8805   71.2271 C   0  0  2  0  0  0  0  0  0  0  0  0\\n  -18.2766   69.4845   70.3956 O   0  0  0  0  0  0  0  0  0  0  0  0\\n  -19.3365   71.8231   70.5433 C   0  0  1  0  0  0  0  0  0  0  0  0\\n  -20.6425   71.1930   70.4165 O   0  0  0  0  0  0  0  0  0  0  0  0\\n  -18.5609   73.0909   63.9111 H   0  0  0  0  0  0  0  0  0  0  0  0\\n  -16.9075   73.6236   65.5429 H   0  0  0  0  0  0  0  0  0  0  0  0\\n  -16.7253   68.7607   69.4245 H   0  0  0  0  0  0  0  0  0  0  0  0\\n  -16.0003   71.3191   69.0734 H   0  0  0  0  0  0  0  0  0  0  0  0\\n  -16.5039   71.2800   71.9821 H   0  0  0  0  0  0  0  0  0  0  0  0\\n  -18.5177   70.7443   72.3112 H   0  0  0  0  0  0  0  0  0  0  0  0\\n  -19.4675   72.7215   71.2376 H   0  0  0  0  0  0  0  0  0  0  0  0\\n  -18.9996   72.2268   69.6087 H   0  0  0  0  0  0  0  0  0  0  0  0\\n  -19.5330   68.4172   68.4454 H   0  0  0  0  0  0  0  0  0  0  0  0\\n  -21.4495   71.7778   70.4659 H   0  0  0  0  0  0  0  0  0  0  0  0\\n  -16.9761   72.8353   69.9911 H   0  0  0  0  0  0  0  0  0  0  0  0\\n  -15.1423   69.1269   69.9059 H   0  0  0  0  0  0  0  0  0  0  0  0\\n  1  2  1  0\\n  1  6  1  0\\n  2  3  2  0\\n  3  4  1  0\\n  4  5  2  0\\n  4 10  1  0\\n  5  6  1  0\\n  5  8  1  0\\n  6  7  2  0\\n  8  9  1  0\\n  9 10  2  0\\n 10 11  1  0\\n 11 12  1  0\\n 11 17  1  0\\n 11 22  1  6\\n 12 13  1  0\\n 12 14  1  0\\n 12 23  1  6\\n 14 15  1  0\\n 14 16  1  0\\n 14 24  1  6\\n 16 17  1  0\\n 16 18  1  0\\n 16 25  1  1\\n 18 19  1  0\\n 20  1  1  0\\n 21  2  1  0\\n 18 26  1  1\\n 27 18  1  0\\n 28  9  1  0\\n 29 19  1  0\\n 30 15  1  0\\n 31 13  1  0\\nM  END\\n\",\"mol\");\n\tviewer_17333795616564157.setStyle({\"model\": 0},{\"stick\": {\"color\": \"orange\"}});\n\tviewer_17333795616564157.setStyle({\"model\": 1},{\"stick\": {\"color\": \"green\"}});\n\tviewer_17333795616564157.zoomTo();\nviewer_17333795616564157.render();\n});\n</script>",
      "text/html": [
       "<div id=\"3dmolviewer_17333795616564157\"  style=\"position: relative; width: 500px; height: 500px;\">\n",
       "        <p id=\"3dmolwarning_17333795616564157\" style=\"background-color:#ffcccc;color:black\">3Dmol.js failed to load for some reason.  Please check your browser console for error messages.<br></p>\n",
       "        </div>\n",
       "<script>\n",
       "\n",
       "var loadScriptAsync = function(uri){\n",
       "  return new Promise((resolve, reject) => {\n",
       "    //this is to ignore the existence of requirejs amd\n",
       "    var savedexports, savedmodule;\n",
       "    if (typeof exports !== 'undefined') savedexports = exports;\n",
       "    else exports = {}\n",
       "    if (typeof module !== 'undefined') savedmodule = module;\n",
       "    else module = {}\n",
       "\n",
       "    var tag = document.createElement('script');\n",
       "    tag.src = uri;\n",
       "    tag.async = true;\n",
       "    tag.onload = () => {\n",
       "        exports = savedexports;\n",
       "        module = savedmodule;\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://cdnjs.cloudflare.com/ajax/libs/3Dmol/2.3.0/3Dmol-min.js');\n",
       "}\n",
       "\n",
       "var viewer_17333795616564157 = null;\n",
       "var warn = document.getElementById(\"3dmolwarning_17333795616564157\");\n",
       "if(warn) {\n",
       "    warn.parentNode.removeChild(warn);\n",
       "}\n",
       "$3Dmolpromise.then(function() {\n",
       "viewer_17333795616564157 = $3Dmol.createViewer(document.getElementById(\"3dmolviewer_17333795616564157\"),{backgroundColor:\"white\"});\n",
       "viewer_17333795616564157.zoomTo();\n",
       "\tviewer_17333795616564157.addModel(\"1a69_ligand\\n     RDKit          3D\\n\\n 31 33  0  0  1  0  0  0  0  0999 V2000\\n  -18.6290   72.4960   64.7810 N   0  0  0  0  0  0  0  0  0  0  0  0\\n  -17.6590   72.8430   65.6920 C   0  0  0  0  0  0  0  0  0  0  0  0\\n  -17.4340   72.1880   66.8280 N   0  0  0  0  0  0  0  0  0  0  0  0\\n  -18.2650   71.1350   66.9710 C   0  0  0  0  0  0  0  0  0  0  0  0\\n  -19.2640   70.7080   66.1180 C   0  0  0  0  0  0  0  0  0  0  0  0\\n  -19.4960   71.4120   64.9200 C   0  0  0  0  0  0  0  0  0  0  0  0\\n  -20.3350   71.1770   64.0460 O   0  0  0  0  0  0  0  0  0  0  0  0\\n  -19.9000   69.5840   66.6230 N   0  0  0  0  0  2  0  0  0  0  0  0\\n  -19.2890   69.3560   67.7600 N   0  0  0  0  0  4  0  0  0  0  0  0\\n  -18.2820   70.2630   68.0340 C   0  0  0  0  0  0  0  0  0  0  0  0\\n  -17.5510   69.7970   69.2390 C   0  0  1  0  0  0  0  0  0  0  0  0\\n  -16.4120   70.7300   69.6340 C   0  0  2  0  0  0  0  0  0  0  0  0\\n  -15.2040   69.9950   69.6560 O   0  0  0  0  0  0  0  0  0  0  0  0\\n  -16.8460   71.3300   70.9950 C   0  0  2  0  0  0  0  0  0  0  0  0\\n  -16.1160   70.8000   72.1220 O   0  0  0  0  0  0  0  0  0  0  0  0\\n  -18.3190   70.9540   71.1130 C   0  0  2  0  0  0  0  0  0  0  0  0\\n  -18.4590   69.7460   70.3290 O   0  0  0  0  0  0  0  0  0  0  0  0\\n  -19.2870   72.0220   70.6090 C   0  0  1  0  0  0  0  0  0  0  0  0\\n  -20.6560   71.5800   70.5560 O   0  0  0  0  0  0  0  0  0  0  0  0\\n  -18.7201   73.0664   63.9525 H   0  0  0  0  0  0  0  0  0  0  0  0\\n  -17.0412   73.7009   65.4715 H   0  0  0  0  0  0  0  0  0  0  0  0\\n  -17.1523   68.7993   69.0552 H   0  0  0  0  0  0  0  0  0  0  0  0\\n  -16.3319   71.5302   68.8982 H   0  0  0  0  0  0  0  0  0  0  0  0\\n  -16.7442   72.4149   70.9680 H   0  0  0  0  0  0  0  0  0  0  0  0\\n  -18.5469   70.7304   72.1552 H   0  0  0  0  0  0  0  0  0  0  0  0\\n  -19.2283   72.8823   71.2757 H   0  0  0  0  0  0  0  0  0  0  0  0\\n  -18.9793   72.3309   69.6100 H   0  0  0  0  0  0  0  0  0  0  0  0\\n  -19.5302   68.5925   68.3757 H   0  0  0  0  0  0  0  0  0  0  0  0\\n  -21.2109   72.2936   70.2328 H   0  0  0  0  0  0  0  0  0  0  0  0\\n  -15.1894   71.0386   72.0444 H   0  0  0  0  0  0  0  0  0  0  0  0\\n  -14.4804   70.5751   69.9040 H   0  0  0  0  0  0  0  0  0  0  0  0\\n  1  2  1  0\\n  1  6  1  0\\n  2  3  2  0\\n  3  4  1  0\\n  4  5  2  0\\n  4 10  1  0\\n  5  6  1  0\\n  5  8  1  0\\n  6  7  2  0\\n  8  9  1  0\\n  9 10  2  0\\n 10 11  1  0\\n 11 12  1  0\\n 11 17  1  0\\n 11 22  1  6\\n 12 13  1  0\\n 12 14  1  0\\n 12 23  1  6\\n 14 15  1  0\\n 14 16  1  0\\n 14 24  1  6\\n 16 17  1  0\\n 16 18  1  0\\n 16 25  1  1\\n 18 19  1  0\\n 20  1  1  0\\n 21  2  1  0\\n 18 26  1  1\\n 27 18  1  0\\n 28  9  1  0\\n 29 19  1  0\\n 30 15  1  0\\n 31 13  1  0\\nM  END\\n\",\"mol\");\n",
       "\tviewer_17333795616564157.addModel(\"1a69_ligand\\n     RDKit          3D\\n\\n 31 33  0  0  1  0  0  0  0  0999 V2000\\n  -18.5428   72.4761   64.7699 N   0  0  0  0  0  0  0  0  0  0  0  0\\n  -17.5279   72.8025   65.7890 C   0  0  0  0  0  0  0  0  0  0  0  0\\n  -17.4808   72.0211   66.9008 N   0  0  0  0  0  0  0  0  0  0  0  0\\n  -18.2611   71.0120   67.0622 C   0  0  0  0  0  0  0  0  0  0  0  0\\n  -19.2738   70.6054   66.1560 C   0  0  0  0  0  0  0  0  0  0  0  0\\n  -19.4461   71.4386   64.9924 C   0  0  0  0  0  0  0  0  0  0  0  0\\n  -20.3034   71.2558   64.0788 O   0  0  0  0  0  0  0  0  0  0  0  0\\n  -19.8618   69.5291   66.6005 N   0  0  0  0  0  2  0  0  0  0  0  0\\n  -19.2045   69.2367   67.8442 N   0  0  0  0  0  4  0  0  0  0  0  0\\n  -18.2542   70.1145   68.1358 C   0  0  0  0  0  0  0  0  0  0  0  0\\n  -17.4280   69.7394   69.3457 C   0  0  1  0  0  0  0  0  0  0  0  0\\n  -16.4033   70.6036   69.8435 C   0  0  2  0  0  0  0  0  0  0  0  0\\n  -15.2576   69.8840   70.4399 O   0  0  0  0  0  0  0  0  0  0  0  0\\n  -16.9491   71.3994   70.9893 C   0  0  1  0  0  0  0  0  0  0  0  0\\n  -17.1736   72.7079   70.8005 O   0  0  0  0  0  0  0  0  0  0  0  0\\n  -18.2152   70.8805   71.2271 C   0  0  2  0  0  0  0  0  0  0  0  0\\n  -18.2766   69.4845   70.3956 O   0  0  0  0  0  0  0  0  0  0  0  0\\n  -19.3365   71.8231   70.5433 C   0  0  1  0  0  0  0  0  0  0  0  0\\n  -20.6425   71.1930   70.4165 O   0  0  0  0  0  0  0  0  0  0  0  0\\n  -18.5609   73.0909   63.9111 H   0  0  0  0  0  0  0  0  0  0  0  0\\n  -16.9075   73.6236   65.5429 H   0  0  0  0  0  0  0  0  0  0  0  0\\n  -16.7253   68.7607   69.4245 H   0  0  0  0  0  0  0  0  0  0  0  0\\n  -16.0003   71.3191   69.0734 H   0  0  0  0  0  0  0  0  0  0  0  0\\n  -16.5039   71.2800   71.9821 H   0  0  0  0  0  0  0  0  0  0  0  0\\n  -18.5177   70.7443   72.3112 H   0  0  0  0  0  0  0  0  0  0  0  0\\n  -19.4675   72.7215   71.2376 H   0  0  0  0  0  0  0  0  0  0  0  0\\n  -18.9996   72.2268   69.6087 H   0  0  0  0  0  0  0  0  0  0  0  0\\n  -19.5330   68.4172   68.4454 H   0  0  0  0  0  0  0  0  0  0  0  0\\n  -21.4495   71.7778   70.4659 H   0  0  0  0  0  0  0  0  0  0  0  0\\n  -16.9761   72.8353   69.9911 H   0  0  0  0  0  0  0  0  0  0  0  0\\n  -15.1423   69.1269   69.9059 H   0  0  0  0  0  0  0  0  0  0  0  0\\n  1  2  1  0\\n  1  6  1  0\\n  2  3  2  0\\n  3  4  1  0\\n  4  5  2  0\\n  4 10  1  0\\n  5  6  1  0\\n  5  8  1  0\\n  6  7  2  0\\n  8  9  1  0\\n  9 10  2  0\\n 10 11  1  0\\n 11 12  1  0\\n 11 17  1  0\\n 11 22  1  6\\n 12 13  1  0\\n 12 14  1  0\\n 12 23  1  6\\n 14 15  1  0\\n 14 16  1  0\\n 14 24  1  6\\n 16 17  1  0\\n 16 18  1  0\\n 16 25  1  1\\n 18 19  1  0\\n 20  1  1  0\\n 21  2  1  0\\n 18 26  1  1\\n 27 18  1  0\\n 28  9  1  0\\n 29 19  1  0\\n 30 15  1  0\\n 31 13  1  0\\nM  END\\n\",\"mol\");\n",
       "\tviewer_17333795616564157.setStyle({\"model\": 0},{\"stick\": {\"color\": \"orange\"}});\n",
       "\tviewer_17333795616564157.setStyle({\"model\": 1},{\"stick\": {\"color\": \"green\"}});\n",
       "\tviewer_17333795616564157.zoomTo();\n",
       "viewer_17333795616564157.render();\n",
       "});\n",
       "</script>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "mblock = Chem.MolToMolBlock(init_mol)\n",
    "view = py3Dmol.view(width=500, height=500)\n",
    "view.addModel(mblock, 'mol')\n",
    "mblock = Chem.MolToMolBlock(gen_mol)\n",
    "view.addModel(mblock, 'mol')\n",
    "view.setStyle({'model':0},{'stick':{'color': 'orange'}})\n",
    "view.setStyle({'model':1},{'stick':{'color': 'green'}})\n",
    "view.zoomTo()\n",
    "view.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "DiffAlign",
   "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.9.19"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}