{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!--NOTEBOOK_HEADER-->\n",
    "*This notebook contains material from [PyRosetta](https://RosettaCommons.github.io/PyRosetta.notebooks);\n",
    "content is available [on Github](https://github.com/RosettaCommons/PyRosetta.notebooks.git).*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!--NAVIGATION-->\n",
    "< [PyRosettaCluster Tutorial 3. Multiple decoys](http://nbviewer.jupyter.org/github/RosettaCommons/PyRosetta.notebooks/blob/master/notebooks/16.09-PyRosettaCluster-Multiple-decoys.ipynb) | [Contents](toc.ipynb) | [Index](index.ipynb) | [Command Reference](http://nbviewer.jupyter.org/github/RosettaCommons/PyRosetta.notebooks/blob/master/notebooks/A.00-Appendix-A.ipynb) ><p><a href=\"https://colab.research.google.com/github/RosettaCommons/PyRosetta.notebooks/blob/master/notebooks/16.10-PyRosettaCluster-Ligand-params.ipynb\"><img align=\"left\" src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open in Colab\" title=\"Open in Google Colaboratory\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# PyRosettaCluster Tutorial 4. Ligand params\n",
    "\n",
    "PyRosettaCluster Tutorial 4 is an example of how to use a non-canonical residue or ligand `.params` file with `PyRosettaCluster`. If a structure contains a ligand that requires a `.params` file, then PyRosetta must be initialized prior to job distribution with `PyRosettaCluster`. For reproducibility outside of `PyRosettaCluster`, PyRosetta should always be initialized with a constant seed."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Warning*: This notebook uses `pyrosetta.distributed.viewer` code, which runs in `jupyter notebook` and might not run if you're using `jupyterlab`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Note:* This Jupyter notebook uses parallelization and is **not** meant to be executed within a Google Colab environment."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Note:* This Jupyter notebook requires the PyRosetta distributed layer which is obtained by building PyRosetta with the `--serialization` flag or installing PyRosetta from the RosettaCommons conda channel \n",
    "\n",
    "**Please see Chapter 16.00 for setup instructions**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Note:* This Jupyter notebook is intended to be run within **Jupyter Lab**, but may still be run as a standalone Jupyter notebook."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1. Import packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<style>.container { width:100% !important; }</style>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import bz2\n",
    "import glob\n",
    "import logging\n",
    "import os\n",
    "import pyrosetta\n",
    "import pyrosetta.distributed.io as io\n",
    "import pyrosetta.distributed.viewer as viewer\n",
    "\n",
    "from pyrosetta.distributed.cluster import PyRosettaCluster\n",
    "\n",
    "logging.basicConfig(level=logging.INFO)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. Initialize a compute cluster using `dask`:\n",
    "\n",
    "See Tutorial 1A for review:\n",
    "1. Click the \"Dask\" tab in Jupyter Lab <i>(arrow, left)</i>\n",
    "2. Click the \"+ NEW\" button to launch a new compute cluster <i>(arrow, lower)</i>\n",
    "3. Once the cluster has started, click the brackets to \"inject client code\" for the cluster into your notebook\n",
    "\n",
    "Inject client code here, then run the cell:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table style=\"border: 2px solid white;\">\n",
       "<tr>\n",
       "<td style=\"vertical-align: top; border: 0px solid white\">\n",
       "<h3 style=\"text-align: left;\">Client</h3>\n",
       "<ul style=\"text-align: left; list-style: none; margin: 0; padding: 0;\">\n",
       "  <li><b>Scheduler: </b>tcp://127.0.0.1:40329</li>\n",
       "  <li><b>Dashboard: </b><a href='http://127.0.0.1:8787/status' target='_blank'>http://127.0.0.1:8787/status</a>\n",
       "</ul>\n",
       "</td>\n",
       "<td style=\"vertical-align: top; border: 0px solid white\">\n",
       "<h3 style=\"text-align: left;\">Cluster</h3>\n",
       "<ul style=\"text-align: left; list-style:none; margin: 0; padding: 0;\">\n",
       "  <li><b>Workers: </b>4</li>\n",
       "  <li><b>Cores: </b>4</li>\n",
       "  <li><b>Memory: </b>16.63 GB</li>\n",
       "</ul>\n",
       "</td>\n",
       "</tr>\n",
       "</table>"
      ],
      "text/plain": [
       "<Client: 'tcp://127.0.0.1:40329' processes=4 threads=4, memory=16.63 GB>"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "if not os.getenv(\"DEBUG\"):\n",
    "    from dask.distributed import Client\n",
    "\n",
    "    client = Client(\"tcp://127.0.0.1:40329\")\n",
    "else:\n",
    "    client = None\n",
    "client"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3. Define ligand `.params` file(s) and initialize PyRosetta with a constant seed:\n",
    "\n",
    "The `-run:constant_seed 1` flag defines a default constant seed of `1111111` and is necessary for reproducibility of your simulation! Initialization is necessary prior to distributing jobs that return a `Pose` or `PackedPose` with ligand or non-canonical residues. If you do not propery initialize PyRosetta within the Jupyter Notebook, then your Jupyter Notebook kernel may die and the job distribution may fail."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.getenv(\"DEBUG\"):\n",
    "    params = os.path.join(os.getcwd(), \"inputs\", \"TPA.am1-bcc.fa.params\")\n",
    "    pyrosetta.distributed.init(f\"-extra_res_fa {params} -run:constant_seed 1 -multithreading:total_threads 1\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4. Define the user-provided PyRosetta protocol:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def protocol1(packed_pose_in, **kwargs):\n",
    "    \"\"\"\n",
    "    Relax residue 1X (i.e. the ligand).\n",
    "    \n",
    "    Args:\n",
    "        packed_pose_in: A `PackedPose` object. Optional.\n",
    "        **kwargs: PyRosettaCluster keyword arguments.\n",
    "\n",
    "    Returns:\n",
    "        A `PackedPose` object.\n",
    "    \"\"\"\n",
    "    import pyrosetta\n",
    "    import pyrosetta.distributed.io as io\n",
    "    import pyrosetta.distributed.tasks.rosetta_scripts as rosetta_scripts\n",
    "\n",
    "    xml = \"\"\"\n",
    "        <ROSETTASCRIPTS>\n",
    "          <RESIDUE_SELECTORS>\n",
    "            <Index name=\"ligand_selector\" resnums=\"1X\"/>\n",
    "            <Not name=\"not_ligand_selector\" selector=\"ligand_selector\"/>\n",
    "          </RESIDUE_SELECTORS>\n",
    "          <TASKOPERATIONS>\n",
    "            <ResfileCommandOperation name=\"repack_ligand\" command=\"NATAA\" residue_selector=\"ligand_selector\"/>\n",
    "            <OperateOnResidueSubset name=\"prevent_repacking\" selector=\"not_ligand_selector\">\n",
    "              <PreventRepackingRLT/>\n",
    "            </OperateOnResidueSubset>\n",
    "          </TASKOPERATIONS>\n",
    "          <MOVERS>\n",
    "            <FastRelax name=\"relax\" task_operations=\"repack_ligand,prevent_repacking\">\n",
    "              <MoveMap bb=\"0\" chi=\"0\" jump=\"1\">\n",
    "                <ResidueSelector selector=\"ligand_selector\" chi=\"1\" bb=\"1\" bondangle=\"0\" bondlength=\"0\"/>\n",
    "              </MoveMap>\n",
    "            </FastRelax>\n",
    "          </MOVERS>\n",
    "          <PROTOCOLS>\n",
    "            <Add mover=\"relax\"/>\n",
    "          </PROTOCOLS>\n",
    "        </ROSETTASCRIPTS>\n",
    "        \"\"\"\n",
    "    \n",
    "    return rosetta_scripts.SingleoutputRosettaScriptsTask(xml)(io.pose_from_file(kwargs[\"s\"]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 5. Launch the original simulation using `distribute()`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.getenv(\"DEBUG\"):\n",
    "    my_task = {\n",
    "        \"options\": \"-ex1\",\n",
    "        \"extra_options\": f\"-out:level 300 -multithreading:total_threads 1 -extra_res_fa {params}\",\n",
    "        \"s\": os.path.join(os.getcwd(), \"inputs\", \"test_lig.pdb\"), \n",
    "    }\n",
    "\n",
    "    output_path = os.path.join(os.getcwd(), \"outputs_4\")\n",
    "\n",
    "    PyRosettaCluster(\n",
    "        tasks=my_task,\n",
    "        client=client,\n",
    "        scratch_dir=output_path,\n",
    "        output_path=output_path,\n",
    "    ).distribute(protocols=[protocol1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "While jobs are running, you may monitor their progress using the dask dashboard diagnostics within Jupyter Lab!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 6. Visualize the resultant decoy:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Gather the input and output decoys from disk into memory:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.getenv(\"DEBUG\"):\n",
    "    input_file = os.path.join(os.getcwd(), \"inputs\", \"test_lig.pdb\")\n",
    "    output_files = glob.glob(os.path.join(output_path, \"decoys\", \"*\", \"*.pdb.bz2\"))\n",
    "\n",
    "    packed_poses = []\n",
    "    for pdbfile in [input_file] + output_files:\n",
    "        if pdbfile.endswith(\".bz2\"):\n",
    "            with open(pdbfile, \"rb\") as f:\n",
    "                packed_poses.append(io.pose_from_pdbstring(bz2.decompress(f.read()).decode()))\n",
    "        elif pdbfile.endswith(\".pdb\"):\n",
    "            with open(pdbfile, \"r\") as f:\n",
    "                packed_poses.append(io.pose_from_pdbstring(f.read()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "View the poses in memory:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<script>$('.output_scroll').removeClass('output_scroll')</script>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<script>$('.output_scroll').removeClass('output_scroll')</script>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<script>$('.output_scroll').removeClass('output_scroll')</script>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<script>$('.output_scroll').removeClass('output_scroll')</script>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<script>$('.output_scroll').removeClass('output_scroll')</script>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9b9033859c4e40689ceb12aa078006cd",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "interactive(children=(IntSlider(value=0, continuous_update=False, description='Decoys', max=1), Output()), _do…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<function pyrosetta.distributed.viewer.core.Viewer.show.<locals>.view(i=0)>"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "if not os.getenv(\"DEBUG\"):\n",
    "    chX = pyrosetta.rosetta.core.select.residue_selector.ChainSelector(\"X\")\n",
    "\n",
    "    view = viewer.init(packed_poses, window_size=(800, 600))\n",
    "    view.add(viewer.setStyle())\n",
    "    view.add(viewer.setStyle(residue_selector=chX, colorscheme=\"magentaCarbon\", radius=0.35))\n",
    "    view.add(viewer.setHydrogenBonds())\n",
    "    view.add(viewer.setHydrogens(polar_only=True))\n",
    "    view()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Congrats! \n",
    "You have successfully executed a PyRosetta simulation that modifies a ligand residue with `PyRosettaCluster`!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!--NAVIGATION-->\n",
    "< [PyRosettaCluster Tutorial 3. Multiple decoys](http://nbviewer.jupyter.org/github/RosettaCommons/PyRosetta.notebooks/blob/master/notebooks/16.09-PyRosettaCluster-Multiple-decoys.ipynb) | [Contents](toc.ipynb) | [Index](index.ipynb) | [Command Reference](http://nbviewer.jupyter.org/github/RosettaCommons/PyRosetta.notebooks/blob/master/notebooks/A.00-Appendix-A.ipynb) ><p><a href=\"https://colab.research.google.com/github/RosettaCommons/PyRosetta.notebooks/blob/master/notebooks/16.10-PyRosettaCluster-Ligand-params.ipynb\"><img align=\"left\" src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open in Colab\" title=\"Open in Google Colaboratory\"></a>"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:PyRosetta.notebooks]",
   "language": "python",
   "name": "pyrosetta.notebooks"
  },
  "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}