EquiBind.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "EquiBind.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "authorship_tag": "ABX9TyPQcJReQQd2ht7qYg0xhb/j",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "gpuClass": "standard"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/pablo-arantes/Cloud-Bind/blob/main/EquiBind.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# **Hi there!**\n",
        "\n",
        "This is a Jupyter notebook to perform direct-shot prediction of both the receptor binding location (blind docking) and the ligand’s bound pose and orientation using [EquiBind](https://github.com/HannesStark/EquiBind) (geometric deep learning model).\n",
        "\n",
        "The main goal of this notebook is to demonstrate how to harness the power of cloud-computing to perform drug binding structure prediction in a cheap and yet feasible fashion.\n",
        "\n",
        "--- \n",
        "**Bugs**\n",
        "- If you encounter any bugs, please report the issue to https://github.com/pablo-arantes/Cloud-Bind/issues\n",
        "\n",
        "**Acknowledgments**\n",
        "- We would like to thank the [EquiBind](https://github.com/HannesStark/EquiBind) and [GNINA](https://github.com/gnina/gnina) team for doing an excellent job open sourcing the software.\n",
        "- We would like to thank the [Roitberg](https://roitberg.chem.ufl.edu/) team for developing the fantastic [TorchANI](https://github.com/aiqm/torchani).\n",
        "- We would like to thank [@ruiz_moreno_aj](https://twitter.com/ruiz_moreno_aj) for his work on [Jupyter Dock](https://github.com/AngelRuizMoreno/Jupyter_Dock) \n",
        "- We would like to thank the ChemosimLab ([@ChemosimLab](https://twitter.com/ChemosimLab)) team for their incredible [ProLIF](https://prolif.readthedocs.io/en/latest/index.html#) (Protein-Ligand Interaction Fingerprints) tool.\n",
        "- Also, credit to [David Koes](https://github.com/dkoes) for his awesome [py3Dmol](https://3dmol.csb.pitt.edu/) plugin.\n",
        "- Finally, we would like to thank [Making it rain](https://github.com/pablo-arantes/making-it-rain) team, **Pablo R. Arantes** ([@pablitoarantes](https://twitter.com/pablitoarantes)), **Marcelo D. Polêto** ([@mdpoleto](https://twitter.com/mdpoleto)), **Conrado Pedebos** ([@ConradoPedebos](https://twitter.com/ConradoPedebos)) and **Rodrigo Ligabue-Braun** ([@ligabue_braun](https://twitter.com/ligabue_braun)), for their amazing work.\n",
        "- A Cloud-Bind by **Pablo R. Arantes** ([@pablitoarantes](https://twitter.com/pablitoarantes))\n",
        "\n",
        "- For related notebooks see: [Cloud-Bind](https://github.com/pablo-arantes/Cloud-Bind)"
      ],
      "metadata": {
        "id": "n-Xw3hPplMsi"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "HBqsph65FmvO"
      },
      "outputs": [],
      "source": [
        "#@title **Install dependencies**\n",
        "#@markdown It will take a few minutes, please, drink a coffee and wait. ;-)\n",
        "# install dependencies\n",
        "%%capture\n",
        "import sys\n",
        "!pip -q install py3Dmol 2>&1 1>/dev/null\n",
        "!pip install Cython\n",
        "\n",
        "# install conda\n",
        "!wget -qnc https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh \n",
        "!bash Miniconda3-latest-Linux-x86_64.sh -bfp /usr/local 2>&1 1>/dev/null\n",
        "!rm -r Miniconda3-latest-Linux-x86_64.sh\n",
        "!conda install -y -q -c conda-forge python=3.7 pdbfixer 2>&1 1>/dev/null\n",
        "!conda install -c conda-forge parmed  --yes 2>&1 1>/dev/null\n",
        "!conda install -c bioconda pybel --yes\n",
        "!conda install -c openbabel openbabel --yes\n",
        "!git clone https://github.com/HannesStark/EquiBind\n",
        "!pip install torch\n",
        "!pip install torchvision\n",
        "!pip install dgl-cu102\n",
        "!pip install bio\n",
        "!pip install biopandas\n",
        "!pip install joblib\n",
        "!pip install tensorboard\n",
        "!pip install POT\n",
        "!pip install pyaml\n",
        "!pip install rdkit-pypi\n",
        "!pip install torchani\n",
        "!pip install ase\n",
        "!pip install pyfiglet\n",
        "\n",
        "#load dependencies\n",
        "sys.path.append('/usr/local/lib/python3.7/site-packages/')\n",
        "import parmed as pmd\n",
        "from biopandas.pdb import PandasPdb\n",
        "import os\n",
        "import urllib.request  \n",
        "import numpy as np\n",
        "import py3Dmol\n",
        "from pdbfixer import PDBFixer\n",
        "import platform\n",
        "import scipy.cluster.hierarchy\n",
        "from scipy.spatial.distance import squareform\n",
        "import scipy.stats as stats\n",
        "import matplotlib.pyplot as plt\n",
        "import pandas as pd\n",
        "from scipy.interpolate import griddata\n",
        "import seaborn as sb\n",
        "from statistics import mean, stdev\n",
        "from matplotlib import colors\n",
        "from IPython.display import set_matplotlib_formats\n",
        "import pyfiglet"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@title **Check if you correctly allocated GPU nodes**\n",
        "\n",
        "gpu_info = !nvidia-smi\n",
        "gpu_info = '\\n'.join(gpu_info)\n",
        "if gpu_info.find('failed') >= 0:\n",
        "  print('Select the Runtime > \"Change runtime type\" menu to enable a GPU accelerator, ')\n",
        "  print('and then re-execute this cell.')\n",
        "else:\n",
        "  print(gpu_info)"
      ],
      "metadata": {
        "cellView": "form",
        "id": "RA1YuWNGGb70"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "---\n",
        "# **Loading the necessary input files**\n",
        "\n",
        "At this point, we should have all libraries and dependencies installed.\n",
        "\n",
        "**Important**: Make sure the PDB ID points to the correct structure. \n",
        "\n",
        "Below, you should provide the PDB ID of your protein.\n",
        "\n",
        "**Please, don't use spaces in the PDB ID names, i.e., 3HTB.**"
      ],
      "metadata": {
        "id": "qntfzVRElUbb"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from openmm.app.pdbfile import PDBFile\n",
        "#@title **Please, provide the necessary input files below for receptor**: \n",
        "#@markdown **Important:** Run the cell to prepare your receptor.\n",
        "\n",
        "import warnings\n",
        "warnings.filterwarnings('ignore')\n",
        "import os\n",
        "from Bio.PDB import PDBParser, PDBIO, Select\n",
        "from Bio.PDB import is_aa\n",
        "import pandas as pd\n",
        "from pdbfixer import PDBFixer\n",
        "\n",
        "\n",
        "# Google_Drive_Path = '/content/' #@param {type:\"string\"}\n",
        "workDir = '/content/'\n",
        "\n",
        "if os.path.exists(os.path.join(workDir, \"name_residues_receptor.txt\")):\n",
        "  os.remove(os.path.join(workDir,\"name_residues_receptor.txt\"))\n",
        "else:\n",
        "  pass\n",
        "\n",
        "receptor = os.path.join(workDir, \"receptor.pdb\")\n",
        "ligand = os.path.join(workDir, \"ligand.sdf\")\n",
        "\n",
        "\n",
        "Query_PDB_ID = '3HTB' #@param {type:\"string\"}\n",
        "\n",
        "pdbfn = Query_PDB_ID + \".pdb\"\n",
        "url = 'https://files.rcsb.org/download/' + pdbfn\n",
        "outfnm = os.path.join(workDir, pdbfn)\n",
        "urllib.request.urlretrieve(url, outfnm)\n",
        "\n",
        "#prepare receptor\n",
        "ppdb = PandasPdb().read_pdb(outfnm)\n",
        "ppdb.df['ATOM'] = ppdb.df['ATOM']\n",
        "ppdb.df['HETATM'] = ppdb.df['HETATM'][ppdb.df['HETATM']['residue_name'] != 'HOH']\n",
        "ppdb.df['ATOM'] = ppdb.df['ATOM'][ppdb.df['ATOM']['atom_name'] != 'OXT']\n",
        "ppdb.df['ATOM']= ppdb.df['ATOM'][ppdb.df['ATOM']['element_symbol'] != 'H']\n",
        "ppdb.to_pdb(path=receptor, records=['ATOM', 'HETATM'], gz=False, append_newline=True)\n",
        "\n",
        "fixer = PDBFixer(filename=receptor)\n",
        "fixer.removeHeterogens()\n",
        "fixer.findMissingResidues()\n",
        "fixer.findMissingAtoms()\n",
        "fixer.addMissingAtoms()\n",
        "fixer.addMissingHydrogens(pH=7.4)\n",
        "PDBFile.writeFile(fixer.topology, fixer.positions, open(receptor, 'w'))\n",
        "\n",
        "path = '/content/'\n",
        "\n",
        "def is_het(residue):\n",
        "    res = residue.id[0]\n",
        "    return res != \" \" and res != \"W\"\n",
        "\n",
        "def aa(residue):\n",
        "    res = residue.id[0]\n",
        "    return res != \"W\" \n",
        "\n",
        "\n",
        "class ResidueSelect(Select):\n",
        "    def __init__(self, chain, residue):\n",
        "        self.chain = chain\n",
        "        self.residue = residue\n",
        "\n",
        "    def accept_chain(self, chain):\n",
        "        return chain.id == self.chain.id\n",
        "\n",
        "    def accept_residue(self, residue):\n",
        "        return residue == self.residue and aa(residue)\n",
        "\n",
        "def extract_ligands(path):\n",
        "    pdb = PDBParser().get_structure(receptor, receptor)\n",
        "    io = PDBIO()\n",
        "    io.set_structure(pdb)\n",
        "    i2 = 1\n",
        "    name_residues2 = []\n",
        "    for model in pdb:\n",
        "      for chain in model:\n",
        "        for residue in chain:\n",
        "          if not aa(residue):\n",
        "            continue\n",
        "          # print(f\"{chain[i].resname} {i}\")\n",
        "          name_residues2.append(residue)\n",
        "          print((f\"saving {residue}\"), file=open(os.path.join(workDir, \"name_residues_receptor.txt\"), \"a\",))\n",
        "          i2 += 1\n",
        "\n",
        "extract_ligands(path)\n",
        "\n",
        "\n",
        "dataset2 = pd.read_csv(os.path.join(workDir, 'name_residues_receptor.txt'), delimiter = \" \", header=None)\n",
        "df2 = pd.DataFrame(dataset2)\n",
        "df2 = df2.iloc[:, [2]]\n",
        "new2 = df2.to_numpy()\n",
        "\n",
        "if os.path.exists(receptor):\n",
        "  print(\"Successfully generated the files! :-)\")\n",
        "else:\n",
        "  print(\"ERROR: Check your inputs! \")"
      ],
      "metadata": {
        "cellView": "form",
        "id": "CM2kzZ-3Gfz_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title **Receptor Visualization**: \n",
        "#@markdown Now the protein has been sanitized, it would be recomended to visualize and check the protein.\n",
        "import py3Dmol\n",
        "\n",
        "view = py3Dmol.view()\n",
        "view.removeAllModels()\n",
        "view.setViewStyle({'style':'outline','color':'black','width':0.1})\n",
        "\n",
        "view.addModel(open(receptor,'r').read(),format='pdb')\n",
        "Prot=view.getModel()\n",
        "Prot.setStyle({'cartoon':{'arrows':True, 'tubes':True, 'style':'oval', 'color':'white'}})\n",
        "view.addSurface(py3Dmol.VDW,{'opacity':0.6,'color':'white'})\n",
        "\n",
        "view.zoomTo()\n",
        "view.show()"
      ],
      "metadata": {
        "cellView": "form",
        "id": "YMRen7pTGnFD"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import requests\n",
        "import rdkit\n",
        "from rdkit import Chem\n",
        "from rdkit.Chem import AllChem,Draw\n",
        "import IPython\n",
        "from IPython.display import Image\n",
        "import pybel\n",
        "import matplotlib.pyplot as plt\n",
        "import matplotlib.image as mpimg\n",
        "\n",
        "\n",
        "# from rdkit.Chem.Draw import IPythonConsole\n",
        "\n",
        "#@title **Please, provide the necessary input files for the ligand**: \n",
        "\n",
        "#@markdown Type the name of your ligand. If you don't know the exactly name, please, check at https://pubchem.ncbi.nlm.nih.gov/\n",
        "Name = '2-PROPYLPHENOL' #@param {type:\"string\"}\n",
        "\n",
        "\n",
        "def get_cids(text):\n",
        "    \"\"\"\n",
        "    Search pubchem and return best matching record\n",
        "    \"\"\"  \n",
        "    url = 'https://pubchem.ncbi.nlm.nih.gov/rest/pug/compound/name/{}/cids/TXT'.format(text)\n",
        "    response = requests.get(url)\n",
        "    cids = response.text.split()\n",
        "    if len(cids) == 0:\n",
        "        return None\n",
        "    else:\n",
        "        return cids\n",
        "def get_record(cid):\n",
        "    \"\"\"\n",
        "    Get pubchem record for a given cid and returns molecule as rdkit\n",
        "    \"\"\"\n",
        "    url = 'https://pubchem.ncbi.nlm.nih.gov/rest/pug/compound/cid/{}/record/SDF'.format(cid)\n",
        "    response = requests.get(url)\n",
        "    mol = Chem.MolFromMolBlock(response.text)\n",
        "    smi = Draw.MolToFile(mol, size=(600, 600), filename=str(Name) + '.png')\n",
        "    img = mpimg.imread(str(Name) + '.png')\n",
        "    plt.figure(figsize = (8,8))\n",
        "    imgplot = plt.imshow(img)\n",
        "    plt.axis('off')\n",
        "    plt.show()\n",
        "    hmol = Chem.AddHs(mol)\n",
        "    AllChem.EmbedMolecule(hmol)\n",
        "    hmol.GetConformer(0)\n",
        "    mp = AllChem.MMFFGetMoleculeProperties(hmol)\n",
        "    ff = AllChem.MMFFGetMoleculeForceField(hmol, mp)\n",
        "    # Optimize\n",
        "    AllChem.OptimizeMoleculeConfs(hmol, ff, numThreads=1, maxIters=1000)\n",
        "    AllChem.MolToMolFile(hmol, (os.path.join(workDir, \"ligand.mol\")))\n",
        "    AllChem.MolToPDBFile(hmol, (os.path.join(workDir, \"ligand.pdb\")))\n",
        "    # mol= [m for m in pybel.readfile(filename=os.path.join(workDir, \"ligand.mol\"), format='mol')][0]\n",
        "    # out=pybel.Outputfile(filename=ligand,format='sdf',overwrite=True)\n",
        "    # out.write(mol)\n",
        "    # out.close()\n",
        "  # return mol\n",
        "\n",
        "def get_molecule(text, n_results=1):\n",
        "    \"\"\"\n",
        "    Search pubchem and return best matching record\n",
        "    \"\"\"  \n",
        "    cids = get_cids(text)\n",
        "    if cids is None:\n",
        "        return None\n",
        "    else:\n",
        "        if n_results == 1:\n",
        "            return get_record(cids[0])\n",
        "        else:\n",
        "            return [get_record(cid) for cid in cids[:n_results]]\n",
        "\n",
        "get_molecule(Name, n_results=1)"
      ],
      "metadata": {
        "id": "y-a05s1XGzqI",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from typing import List\n",
        "from ase import Atoms\n",
        "from ase.lattice.cubic import Diamond\n",
        "from ase.md.langevin import Langevin\n",
        "from ase.optimize import BFGS\n",
        "from ase import io\n",
        "from ase.io import read, write\n",
        "from ase import units\n",
        "from ase.constraints import ExternalForce, FixInternals\n",
        "import torch\n",
        "import torchani\n",
        "import pandas as pd\n",
        "import numpy as np\n",
        "import pybel\n",
        "from torchani.units import HARTREE_TO_KCALMOL\n",
        "\n",
        "\n",
        "#@title **Ligand geometry optimization using TorchANI**: \n",
        "\n",
        "#@markdown Geometry optimization for the ligand 3D structure, using ANI-1x, ANI-1ccx or ANI-2x as the optimizing engine.\n",
        "\n",
        "#@markdown If you want to know more about **TorchANI**, please, check at https://aiqm.github.io/torchani/\n",
        "model_name = \"ANI-2x\" #@param [\"ANI-1x\", \"ANI-1ccx\", \"ANI-2x\"]\n",
        "\n",
        "#@markdown Convergence threshold for geometry optimization:\n",
        "\n",
        "opt_tol = 0.0001 #@param {type:\"slider\", min:0.0001, max:0.01, step:0.0001}\n",
        "\n",
        "\n",
        "\n",
        "# device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
        "device = torch.device('cpu')\n",
        "\n",
        "if model_name == \"ANI-2x\":\n",
        "  model = torchani.models.ANI2x(periodic_table_index=True).to(device)\n",
        "  calculator = torchani.models.ANI2x().ase()\n",
        "  print(\"Model = ANI2x\")\n",
        "elif model_name == \"ANI-1ccx\":\n",
        "  model = torchani.models.ANI1ccx(periodic_table_index=True).to(device)\n",
        "  calculator = torchani.models.ANI1ccx().ase()\n",
        "  print(\"Model = ANI1ccx\")\n",
        "elif model_name == \"ANI-1x\":\n",
        "  model = torchani.models.ANI1x(periodic_table_index=True).to(device)\n",
        "  calculator = torchani.models.ANI1x().ase()\n",
        "  print(\"Model = ANI1x\")\n",
        "else:\n",
        "  pass\n",
        "\n",
        "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\n",
        "\n",
        "mol_deg = AllChem.MolFromMolFile ((os.path.join(workDir, \"ligand.mol\")), removeHs=False)\n",
        "mol = io.read(os.path.join(workDir, \"ligand.mol\"))\n",
        "coordinates, species = mol2arr([mol_deg], device)\n",
        "tensor1 = coordinates.detach().numpy()\n",
        "atoms = Atoms(mol, positions=tensor1[0])\n",
        "atoms.center(vacuum=3.0)\n",
        "atoms.set_calculator(calculator)\n",
        "print(\"Begin Geometry Optimization \")\n",
        "opt = BFGS(atoms)\n",
        "opt.run(fmax=opt_tol)\n",
        "# print()\n",
        "write((os.path.join(workDir, \"ligand_min.xyz\")), format=\"xyz\", images=atoms)\n",
        "\n",
        "atomic_symbols = []\n",
        "xyz_coordinates = []\n",
        "\n",
        "with open((os.path.join(workDir, \"ligand_min.xyz\")), \"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",
        "          comment = line # might have useful information\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",
        "from rdkit.Geometry import Point3D\n",
        "conf = mol_deg.GetConformer()\n",
        "\n",
        "for i in range(mol_deg.GetNumAtoms()):\n",
        "  x,y,z = xyz_coordinates[i]\n",
        "  conf.SetAtomPosition(i,Point3D(x,y,z))\n",
        "AllChem.MolToMolFile(mol_deg, (os.path.join(workDir, \"ligand_min.mol\")))\n",
        "AllChem.MolToPDBFile(mol_deg, (os.path.join(workDir, \"ligand_min.pdb\")))\n",
        "#convert to sdf format\n",
        "mol= [m for m in pybel.readfile(filename=os.path.join(workDir, \"ligand_min.mol\"), format='mol')][0]\n",
        "out=pybel.Outputfile(filename=ligand,format='sdf',overwrite=True)\n",
        "out.write(mol)\n",
        "out.close()\n",
        "\n",
        "\n",
        "#TorchANI Energies\n",
        "mol_deg = AllChem.MolFromMolFile ((os.path.join(workDir, \"ligand_min.mol\")), removeHs=False)\n",
        "coordinates, species = mol2arr([mol_deg], device)\n",
        "energy = model((species, coordinates)).energies\n",
        "derivative = torch.autograd.grad(energy.sum(), coordinates)[0]\n",
        "force = -derivative\n",
        "print('Energy:', energy.item())\n",
        "# print('Force:', force.squeeze())"
      ],
      "metadata": {
        "cellView": "form",
        "id": "bEYsHGvq7gmw"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title **Ligand Visualization**: \n",
        "#@markdown Now the ligand has been sanitized, it would be recomended to visualize and check the ligand.\n",
        "\n",
        "view = py3Dmol.view()\n",
        "view.removeAllModels()\n",
        "view.setViewStyle({'style':'outline','color':'black','width':0.05})\n",
        "\n",
        "# view.addModel(open(receptor,'r').read(),format='pdb')\n",
        "# Prot=view.getModel()\n",
        "# Prot.setStyle({'cartoon':{'arrows':True, 'tubes':True, 'style':'oval', 'color':'white'}})\n",
        "# view.addSurface(py3Dmol.VDW,{'opacity':0.6,'color':'white'})\n",
        "\n",
        "view.addModel(open(ligand,'r').read(),format='mol2')\n",
        "ref_m = view.getModel()\n",
        "ref_m.setStyle({},{'stick':{'colorscheme':'greenCarbon','radius':0.2}})\n",
        "\n",
        "view.zoomTo()\n",
        "view.show()"
      ],
      "metadata": {
        "cellView": "form",
        "id": "cnY-Ouuxt6pL"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title **Parameters for the EquiBind calculation:**\n",
        "Output_folder_name = 'final' #@param {type:\"string\"}\n",
        "\n",
        "#@markdown Explicit random seed:\n",
        "seed = \"0\" #@param {type:\"string\"}\n",
        "\n",
        "font = pyfiglet.figlet_format(\"EquiBind\")\n",
        "print(font)\n",
        "\n",
        "docking_output = os.path.join(workDir, Output_folder_name)\n",
        "\n",
        "# if os.path.exists(docking_output):\n",
        "#   os.remove(docking_output)\n",
        "# else:\n",
        "#   pass\n",
        "\n",
        "equibind = \"python3 /content/EquiBind/multiligand_inference.py -o \" + str(docking_output) + \" -r /content/receptor.pdb -l /content/ligand.sdf --seed \" + str(int(seed))\n",
        "\n",
        "original_stdout = sys.stdout # Save a reference to the original standard output\n",
        "with open('equibind.sh', 'w') as f:\n",
        "    sys.stdout = f # Change the standard output to the file we created.\n",
        "    print(equibind)\n",
        "    sys.stdout = original_stdout # Reset the standard output to its original value\n",
        "\n",
        "!chmod 700 equibind.sh 2>&1 1>/dev/null\n",
        "!bash equibind.sh\n",
        "\n",
        "import gzip\n",
        "v = py3Dmol.view()\n",
        "v.setViewStyle({'style':'outline','color':'black','width':0.05})\n",
        "v.addModel(open(receptor).read())\n",
        "v.setStyle({'cartoon':{},'stick':{'colorscheme':'white','radius':.1}})\n",
        "# v.addModel(open(res_box).read())\n",
        "# v.setStyle({'model':1},{'stick':{'colorscheme':'dimgrayCarbon','radius':.175}})\n",
        "v.addModelsAsFrames(open(os.path.join(workDir, Output_folder_name + \"/output.sdf\")).read())\n",
        "v.setStyle({'model':1},{'stick':{'colorscheme':'greenCarbon'}})\n",
        "v.zoomTo({'model':1})\n",
        "v.rotate(90)"
      ],
      "metadata": {
        "cellView": "form",
        "id": "Xwdv_ZltqfUn"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title **Download your results**\n",
        "from google.colab import files\n",
        "import shutil\n",
        "\n",
        "if os.path.exists(os.path.join(workDir, \"sample_data\")):\n",
        "  shutil.rmtree(os.path.join(workDir, \"sample_data\"))\n",
        "else:\n",
        "  pass\n",
        "\n",
        "if os.path.exists(os.path.join(workDir, \"results_EquiBind.zip\")):\n",
        "  os.remove(os.path.join(workDir, \"results_EquiBind.zip\"))\n",
        "else:\n",
        "  pass\n",
        "\n",
        "!zip -FSr \"results_EquiBind.zip\" \"/content/\" -x \"/content/EquiBind/*\" 2>&1 1>/dev/null\n",
        "files.download('results_EquiBind.zip')"
      ],
      "metadata": {
        "cellView": "form",
        "id": "f_L76NfaMUqi"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}