docs/tutorials/07-tutorial-cb6-but-namd.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# pAPRika tutorial 7 - APR/NAMD with Colvars restraints\n",
    "\n",
    "In this tutorial, we will perform APR calculations for the butane (BUT)--cucurbit[6]uril (CB6) host-guest system with the [Generalized-Born](https://en.wikipedia.org/wiki/Implicit_solvation#Generalized_Born) implicit solvent. This tutorial will use `NAMD` for the simulation and `Colvars` for the restraints. Although it is possible to run `NAMD` with `Plumed` we will not use `Plumed` in this tutorial. Users should refer to the previous tutorial to see how to use `Plumed` restraints.\n",
    "\n",
    "🔵 Since we have a prepared the host-guest-dummy setup from the first tutorial, we will skip the initial tleap steps and go right into initializing the restraints. If you haven't completed [tutorial 1](01-tutorial-cb6-but.ipynb) please go back and run the notebook to generate the starting structure."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initialize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define names\n",
    "We will store the files created in this tutorial in a folder called `namd` so we don't mix files with the previous tutorial."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "base_name = \"cb6-but-dum\"\n",
    "work_dir = \"namd\"\n",
    "complex_dir = \"complex\"\n",
    "data = \"../../../paprika/data/cb6-but\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Configure APR Restraints"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define anchor atoms\n",
    "See [tutorial 1](01-tutorial-cb6-but.ipynb) for the choice of selection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Guest atoms\n",
    "G1 = \":BUT@C\"\n",
    "G2 = \":BUT@C3\"\n",
    "\n",
    "# Host atoms\n",
    "H1 = \":CB6@C\"\n",
    "H2 = \":CB6@C31\"\n",
    "H3 = \":CB6@C18\"\n",
    "\n",
    "# Dummy atoms\n",
    "D1 = \":DM1\"\n",
    "D2 = \":DM2\"\n",
    "D3 = \":DM3\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Determine the number of windows\n",
    "Before we add the restraints, it is helpful to set the $\\lambda$ fractions that control the strength of the force constants during attach and release, and to define the distances for the pulling phase.\n",
    "\n",
    "The attach fractions go from 0 to 1 and we place more points at the bottom of the range to sample the curvature of $dU/d \\lambda$. Next, we generally apply a distance restraint until the guest is ~18 Angstroms away from the host, in increments of 0.4 Angstroms. This distance should be at least twice the Lennard-Jones cutoff in the system. These values have worked well for us, but this is one aspect that should be carefully checked for new systems."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "There are [15, 18, 0] windows in this attach-pull-release calculation.\n"
     ]
    }
   ],
   "source": [
    "attach_string = \"0.00 0.40 0.80 1.60 2.40 4.00 5.50 8.65 11.80 18.10 24.40 37.00 49.60 74.80 100.00\"\n",
    "attach_fractions = [float(i) / 100 for i in attach_string.split()]\n",
    "\n",
    "initial_distance = 6.0\n",
    "pull_distances = np.arange(0.0 + initial_distance, 18.0 + initial_distance, 1.0)\n",
    "\n",
    "release_fractions = []\n",
    "\n",
    "windows = [len(attach_fractions), len(pull_distances), len(release_fractions)]\n",
    "print(f\"There are {windows} windows in this attach-pull-release calculation.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Load structure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "import parmed as pmd"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* Load complex structure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "structure = pmd.load_file(\n",
    "    os.path.join(complex_dir, f\"{base_name}.prmtop\"),\n",
    "    os.path.join(complex_dir, f\"{base_name}.rst7\"),\n",
    "    structure = True,\n",
    ") "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Host Static Restraints\n",
    "See [tutorial 1](01-tutorial-cb6-but.ipynb#host_static) for an explanation of the static restraints."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "import paprika.restraints as restraints"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "static_restraints = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "r = restraints.static_DAT_restraint(restraint_mask_list = [D1, H1],\n",
    "                                    num_window_list = windows,\n",
    "                                    ref_structure = structure,\n",
    "                                    force_constant = 5.0,\n",
    "                                    amber_index=True)\n",
    "\n",
    "static_restraints.append(r)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "r = restraints.static_DAT_restraint(restraint_mask_list = [D2, D1, H1],\n",
    "                                    num_window_list = windows,\n",
    "                                    ref_structure = structure,\n",
    "                                    force_constant = 100.0,\n",
    "                                    amber_index=True)\n",
    "\n",
    "static_restraints.append(r)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "r = restraints.static_DAT_restraint(restraint_mask_list = [D3, D2, D1, H1],\n",
    "                                    num_window_list = windows,\n",
    "                                    ref_structure = structure,\n",
    "                                    force_constant = 100.0,\n",
    "                                    amber_index=True)\n",
    "\n",
    "static_restraints.append(r)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "r = restraints.static_DAT_restraint(restraint_mask_list = [D1, H1, H2],\n",
    "                                    num_window_list = windows,\n",
    "                                    ref_structure = structure,\n",
    "                                    force_constant = 100.0,\n",
    "                                    amber_index=True)\n",
    "\n",
    "static_restraints.append(r)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "r = restraints.static_DAT_restraint(restraint_mask_list = [D2, D1, H1, H2],\n",
    "                                    num_window_list = windows,\n",
    "                                    ref_structure = structure,\n",
    "                                    force_constant = 100.0,\n",
    "                                    amber_index=True)\n",
    "\n",
    "static_restraints.append(r)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "r = restraints.static_DAT_restraint(restraint_mask_list = [D1, H1, H2, H3],\n",
    "                                    num_window_list = windows,\n",
    "                                    ref_structure = structure,\n",
    "                                    force_constant = 100.0,\n",
    "                                    amber_index=True)\n",
    "\n",
    "static_restraints.append(r)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Guest translational and rotational restraints\n",
    "See [tutorial 1](01-tutorial-cb6-but.ipynb#guest) for an explanation of the guest restraints"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "guest_restraints = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "r = restraints.DAT_restraint()\n",
    "r.mask1 = D1\n",
    "r.mask2 = G1\n",
    "r.topology = structure\n",
    "r.auto_apr = True\n",
    "r.continuous_apr = True\n",
    "r.amber_index = True\n",
    "\n",
    "r.attach[\"target\"] = pull_distances[0]          # Angstroms\n",
    "r.attach[\"fraction_list\"] = attach_fractions\n",
    "r.attach[\"fc_final\"] = 5.0                      # kcal/mol/Angstroms**2\n",
    "\n",
    "r.pull[\"target_final\"] = 24.0                   # Angstroms\n",
    "r.pull[\"num_windows\"] = windows[1]\n",
    "\n",
    "r.initialize()\n",
    "guest_restraints.append(r)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "r = restraints.DAT_restraint()\n",
    "r.mask1 = D2\n",
    "r.mask2 = D1\n",
    "r.mask3 = G1\n",
    "r.topology = structure\n",
    "r.auto_apr = True\n",
    "r.continuous_apr = True\n",
    "r.amber_index = True\n",
    "\n",
    "r.attach[\"target\"] = 180.0                      # Degrees\n",
    "r.attach[\"fraction_list\"] = attach_fractions\n",
    "r.attach[\"fc_final\"] = 100.0                    # kcal/mol/radian**2\n",
    "\n",
    "r.pull[\"target_final\"] = 180.0                  # Degrees\n",
    "r.pull[\"num_windows\"] = windows[1]\n",
    "\n",
    "r.initialize()\n",
    "guest_restraints.append(r)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "r = restraints.DAT_restraint()\n",
    "r.mask1 = D1\n",
    "r.mask2 = G1\n",
    "r.mask3 = G2\n",
    "r.topology = structure\n",
    "r.auto_apr = True\n",
    "r.continuous_apr = True\n",
    "r.amber_index = True\n",
    "\n",
    "r.attach[\"target\"] = 180.0                      # Degrees\n",
    "r.attach[\"fraction_list\"] = attach_fractions\n",
    "r.attach[\"fc_final\"] = 100.0                    # kcal/mol/radian**2\n",
    "\n",
    "r.pull[\"target_final\"] = 180.0                  # Degrees\n",
    "r.pull[\"num_windows\"] = windows[1]\n",
    "\n",
    "r.initialize()\n",
    "guest_restraints.append(r)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Create APR windows\n",
    "We use the guest restraints to create a list of windows with the appropriate names and then create the directories."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "from paprika.restraints.restraints import create_window_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "window_list = create_window_list(guest_restraints)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.isdir(work_dir):\n",
    "    os.makedirs(work_dir)\n",
    "    \n",
    "for window in window_list:\n",
    "    folder = os.path.join(work_dir, window)\n",
    "    if not os.path.isdir(folder):\n",
    "        os.makedirs(os.path.join(work_dir, window))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Write APR restraints to a Colvars module file\n",
    "In this section we create an instance of `Colvars()` from `paprika.restraints.colvars`, which is a class to generate the `Colvars` restraint files. This class is analogous to the `Plumed` class. We need to specify the list of restraints used throughout the APR calculations and the corresponding windows list. In this tutorial we will print the **host static** restraints and the **guest** restraints. The `Colvars` class includes a method to add restraints to dummy atoms (`add_dummy_atoms_to_file`). We will add the positional restraints on the dummy atoms after generating the structures for each window.\n",
    "\n",
    "<font size='4'><ins>**Note**</ins></font>: be careful when specifying the force constants in `DAT_restraints`. We follow the AMBER (and CHARMM) convention where the force constant is already multiplied by a factor of 1/2 but `Colvars` (and `Plumed`) requires the user to specify the force constant without this factor, i.e.\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "U_{amber} &= K_{amber} (r-r_{0})^2 \\\\\n",
    "U_{colvars} &= \\frac{1}{2} k_{colvars} (r - r_{0})^2 \n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "thus $k_{colvars} = 2 \\times  K_{amber}$. If AMBER force constants was used in generating the `DAT_restraints` (the case in this tutorial) we need to set the variable `uses_legacy_k` to `True` (this is on by default)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "from paprika.restraints.colvars import Colvars"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "restraints_list = (static_restraints + guest_restraints)\n",
    "\n",
    "colvars = Colvars()\n",
    "colvars.file_name = \"colvars.tcl\"\n",
    "colvars.path = work_dir\n",
    "colvars.window_list = window_list\n",
    "colvars.restraint_list = restraints_list\n",
    "colvars_uses_legacy_k = True\n",
    "\n",
    "colvars.dump_to_file()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Prepare host-guest system"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Translate guest molecule\n",
    "For the attach windows, we will use the initial, bound coordinates for the host-guest complex. Only the force constants change during this phase, so a single set of coordinates is sufficient. For the pull windows, we will translate the guest to the target value, and for the release windows, we will use the coordinates from the final pull window."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "import shutil"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "In window p000 we will translate the guest 0.0 Angstroms.\n",
      "In window p001 we will translate the guest 1.1 Angstroms.\n",
      "In window p002 we will translate the guest 2.1 Angstroms.\n",
      "In window p003 we will translate the guest 3.2 Angstroms.\n",
      "In window p004 we will translate the guest 4.2 Angstroms.\n",
      "In window p005 we will translate the guest 5.3 Angstroms.\n",
      "In window p006 we will translate the guest 6.4 Angstroms.\n",
      "In window p007 we will translate the guest 7.4 Angstroms.\n",
      "In window p008 we will translate the guest 8.5 Angstroms.\n",
      "In window p009 we will translate the guest 9.5 Angstroms.\n",
      "In window p010 we will translate the guest 10.6 Angstroms.\n",
      "In window p011 we will translate the guest 11.6 Angstroms.\n",
      "In window p012 we will translate the guest 12.7 Angstroms.\n",
      "In window p013 we will translate the guest 13.8 Angstroms.\n",
      "In window p014 we will translate the guest 14.8 Angstroms.\n",
      "In window p015 we will translate the guest 15.9 Angstroms.\n",
      "In window p016 we will translate the guest 16.9 Angstroms.\n",
      "In window p017 we will translate the guest 18.0 Angstroms.\n"
     ]
    }
   ],
   "source": [
    "for window in window_list:\n",
    "    if window[0] == \"a\":\n",
    "        shutil.copy(os.path.join(complex_dir, f\"{base_name}.prmtop\"),\n",
    "                    os.path.join(work_dir, window, f\"{base_name}.prmtop\"))\n",
    "        shutil.copy(os.path.join(complex_dir, f\"{base_name}.rst7\"),\n",
    "                    os.path.join(work_dir, window, f\"{base_name}.rst7\"))\n",
    "\n",
    "    elif window[0] == \"p\":\n",
    "        structure = pmd.load_file(\n",
    "            os.path.join(complex_dir, f\"{base_name}.prmtop\"), \n",
    "            os.path.join(complex_dir, f\"{base_name}.rst7\"), \n",
    "            structure = True\n",
    "        )\n",
    "        target_difference = guest_restraints[0].phase['pull']['targets'][int(window[1:])] -\\\n",
    "                            guest_restraints[0].pull['target_initial']\n",
    "        print(f\"In window {window} we will translate the guest {target_difference.magnitude:0.1f}.\")\n",
    "        \n",
    "        for atom in structure.atoms:\n",
    "            if atom.residue.name == \"BUT\":\n",
    "                atom.xz += target_difference.magnitude\n",
    "                \n",
    "        structure.save(os.path.join(work_dir, window, f\"{base_name}.prmtop\"), overwrite=True)\n",
    "        structure.save(os.path.join(work_dir, window, f\"{base_name}.rst7\"), overwrite=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Add dummy atom restraints\n",
    "Here we will add positional restraints on the dummy atoms. The reference positions are extracted from the structure files we just created in each window."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "for window in window_list:\n",
    "    folder = os.path.join(work_dir, window)\n",
    "    \n",
    "    structure = pmd.load_file(\n",
    "        os.path.join(folder, f\"{base_name}.prmtop\"), \n",
    "        os.path.join(folder, f\"{base_name}.rst7\"), \n",
    "        structure = True\n",
    "    )\n",
    "    \n",
    "    colvars.add_dummy_atom_restraints(structure, window)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Simulation\n",
    "<a id='simulate'></a>\n",
    "Now that we have prepared the structure and `colvars.tcl` file for each window, we have everything ready to go. **pAPRika** comes with a python wrapper for `NAMD` that can help set up default parameters for the simulation. There are some high level options that we set directly, like `simulation.path`, and then we call the function `config_gb_min()` to setup reasonable default simulation parameters for a minimization and production run. The temperature of the system will be controlled with the `Langevin` thermostat, which is the default. For running simulations with a different thermostat parse the thermostat of choice using the `Thermostat` Enum class to `config_gb_md()`. For example, if we want to run our simulations with the *Lowe-Anderson* thermostat:\n",
    "\n",
    "```python\n",
    "simulation = NAMD()\n",
    "...\n",
    "simulation.config_gb_md(NAMD.Thermostat.LoweAnderson)\n",
    "simulation.run()\n",
    "```\n",
    "\n",
    "See the [pAPRika](<https://www.ks.uiuc.edu/Research/namd/2.14/ug/node38.html>) documentation for more details on the different temperature control algorithm implemented.\n",
    "\n",
    "<font size='3'><ins>**Note**</ins></font>: `NAMD` version 2 is similar to `GROMACS` in that more than one thread for every GPU is preferred to deliver better speed. We set the number of processors with `simulation.n_threads = 8` and the GPU device with `simulation.gpu_devices = 0`. Starting with `NAMD` version 3 (currently in alpha phase) a smaller number of `n_threads` is required per GPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "from paprika.simulate import NAMD"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Initialize logger"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "from importlib import reload\n",
    "reload(logging)\n",
    "\n",
    "logger = logging.getLogger()\n",
    "logging.basicConfig(\n",
    "    format='%(asctime)s %(message)s',\n",
    "    datefmt='%Y-%m-%d %I:%M:%S %p',\n",
    "    level=logging.INFO\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Energy Minimization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2020-10-09 11:31:02 AM Running minimization in window a000...\n",
      "2020-10-09 11:31:07 AM Running minimization in window a001...\n",
      "2020-10-09 11:31:13 AM Running minimization in window a002...\n",
      "2020-10-09 11:31:20 AM Running minimization in window a003...\n",
      "2020-10-09 11:31:25 AM Running minimization in window a004...\n",
      "2020-10-09 11:31:30 AM Running minimization in window a005...\n",
      "2020-10-09 11:31:34 AM Running minimization in window a006...\n",
      "2020-10-09 11:31:39 AM Running minimization in window a007...\n",
      "2020-10-09 11:31:43 AM Running minimization in window a008...\n",
      "2020-10-09 11:31:48 AM Running minimization in window a009...\n",
      "2020-10-09 11:31:53 AM Running minimization in window a010...\n",
      "2020-10-09 11:31:59 AM Running minimization in window a011...\n",
      "2020-10-09 11:32:03 AM Running minimization in window a012...\n",
      "2020-10-09 11:32:08 AM Running minimization in window a013...\n",
      "2020-10-09 11:32:12 AM Running minimization in window p000...\n",
      "2020-10-09 11:32:16 AM Running minimization in window p001...\n",
      "2020-10-09 11:32:21 AM Running minimization in window p002...\n",
      "2020-10-09 11:32:25 AM Running minimization in window p003...\n",
      "2020-10-09 11:32:30 AM Running minimization in window p004...\n",
      "2020-10-09 11:32:34 AM Running minimization in window p005...\n",
      "2020-10-09 11:32:39 AM Running minimization in window p006...\n",
      "2020-10-09 11:32:43 AM Running minimization in window p007...\n",
      "2020-10-09 11:32:48 AM Running minimization in window p008...\n",
      "2020-10-09 11:32:56 AM Running minimization in window p009...\n",
      "2020-10-09 11:33:01 AM Running minimization in window p010...\n",
      "2020-10-09 11:33:05 AM Running minimization in window p011...\n",
      "2020-10-09 11:33:10 AM Running minimization in window p012...\n",
      "2020-10-09 11:33:14 AM Running minimization in window p013...\n",
      "2020-10-09 11:33:19 AM Running minimization in window p014...\n",
      "2020-10-09 11:33:23 AM Running minimization in window p015...\n",
      "2020-10-09 11:33:28 AM Running minimization in window p016...\n",
      "2020-10-09 11:33:32 AM Running minimization in window p017...\n"
     ]
    }
   ],
   "source": [
    "for window in window_list:\n",
    "    simulation = NAMD()\n",
    "    simulation.executable = \"namd2\"\n",
    "    simulation.n_threads = 4\n",
    "#     simulation.gpu_devices = 0\n",
    "    simulation.path = os.path.join(work_dir, window)\n",
    "    \n",
    "    simulation.topology = f\"{base_name}.prmtop\"\n",
    "    simulation.coordinates = f\"{base_name}.rst7\"\n",
    "    simulation.prefix = \"minimize\"\n",
    "    simulation.colvars_file = \"colvars.tcl\"\n",
    "    \n",
    "    simulation.config_gb_min()\n",
    "    \n",
    "    logging.info(f\"Running minimization in window {window}...\")\n",
    "    simulation.run(overwrite=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Thermalization\n",
    "Here, we will do a slow heating process to demonstrate the `custom_run_command` feature of the `NAMD` wrapper. We will heat the system from 50 K to 300 K in steps of 50 K and will do a total of 500 MD steps for each temperature."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2020-10-09 11:33:37 AM Running thermalization in window a000...\n",
      "2020-10-09 11:33:39 AM Running thermalization in window a001...\n",
      "2020-10-09 11:33:40 AM Running thermalization in window a002...\n",
      "2020-10-09 11:33:42 AM Running thermalization in window a003...\n",
      "2020-10-09 11:33:44 AM Running thermalization in window a004...\n",
      "2020-10-09 11:33:46 AM Running thermalization in window a005...\n",
      "2020-10-09 11:33:47 AM Running thermalization in window a006...\n",
      "2020-10-09 11:33:49 AM Running thermalization in window a007...\n",
      "2020-10-09 11:33:51 AM Running thermalization in window a008...\n",
      "2020-10-09 11:33:53 AM Running thermalization in window a009...\n",
      "2020-10-09 11:33:54 AM Running thermalization in window a010...\n",
      "2020-10-09 11:33:56 AM Running thermalization in window a011...\n",
      "2020-10-09 11:33:58 AM Running thermalization in window a012...\n",
      "2020-10-09 11:34:00 AM Running thermalization in window a013...\n",
      "2020-10-09 11:34:02 AM Running thermalization in window p000...\n",
      "2020-10-09 11:34:06 AM Running thermalization in window p001...\n",
      "2020-10-09 11:34:09 AM Running thermalization in window p002...\n",
      "2020-10-09 11:34:13 AM Running thermalization in window p003...\n",
      "2020-10-09 11:34:15 AM Running thermalization in window p004...\n",
      "2020-10-09 11:34:17 AM Running thermalization in window p005...\n",
      "2020-10-09 11:34:18 AM Running thermalization in window p006...\n",
      "2020-10-09 11:34:20 AM Running thermalization in window p007...\n",
      "2020-10-09 11:34:22 AM Running thermalization in window p008...\n",
      "2020-10-09 11:34:24 AM Running thermalization in window p009...\n",
      "2020-10-09 11:34:25 AM Running thermalization in window p010...\n",
      "2020-10-09 11:34:27 AM Running thermalization in window p011...\n",
      "2020-10-09 11:34:29 AM Running thermalization in window p012...\n",
      "2020-10-09 11:34:30 AM Running thermalization in window p013...\n",
      "2020-10-09 11:34:32 AM Running thermalization in window p014...\n",
      "2020-10-09 11:34:34 AM Running thermalization in window p015...\n",
      "2020-10-09 11:34:36 AM Running thermalization in window p016...\n",
      "2020-10-09 11:34:37 AM Running thermalization in window p017...\n"
     ]
    }
   ],
   "source": [
    "for window in window_list:\n",
    "    simulation = NAMD()\n",
    "    simulation.executable = \"namd2\"\n",
    "    simulation.n_threads = 4\n",
    "#     simulation.gpu_devices = 0\n",
    "    simulation.path = os.path.join(work_dir, window)\n",
    "    \n",
    "    simulation.topology = f\"{base_name}.prmtop\"\n",
    "    simulation.coordinates = f\"{base_name}.rst7\"\n",
    "    simulation.prefix = \"thermalization\"\n",
    "    simulation.checkpoint = \"minimize\"\n",
    "    simulation.colvars_file = \"colvars.tcl\"\n",
    "    \n",
    "    simulation.config_gb_md()\n",
    "    \n",
    "    # Slow heating\n",
    "    simulation.custom_run_commands = [\n",
    "        \"for {set temp 50} {$temp <= 300} {incr temp 50} {\",\n",
    "        \"langevinTemp $temp\",\n",
    "        \"reinitvels $temp\",\n",
    "        \"run 500\",\n",
    "        \"}\",\n",
    "    ]\n",
    "    \n",
    "    logging.info(f\"Running thermalization in window {window}...\")\n",
    "    simulation.run(overwrite=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will skip the equilibration step and go straight to production"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Production Run\n",
    "We will run the production phase for 5,000 steps.\n",
    "\n",
    "<font size='3'><ins>**Note**</ins></font>: a proper production run would require much longer simulation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2020-10-09 11:34:39 AM Running production in window a000...\n",
      "2020-10-09 11:34:41 AM Running production in window a001...\n",
      "2020-10-09 11:34:44 AM Running production in window a002...\n",
      "2020-10-09 11:34:46 AM Running production in window a003...\n",
      "2020-10-09 11:34:49 AM Running production in window a004...\n",
      "2020-10-09 11:34:52 AM Running production in window a005...\n",
      "2020-10-09 11:34:54 AM Running production in window a006...\n",
      "2020-10-09 11:34:57 AM Running production in window a007...\n",
      "2020-10-09 11:35:00 AM Running production in window a008...\n",
      "2020-10-09 11:35:02 AM Running production in window a009...\n",
      "2020-10-09 11:35:05 AM Running production in window a010...\n",
      "2020-10-09 11:35:10 AM Running production in window a011...\n",
      "2020-10-09 11:35:15 AM Running production in window a012...\n",
      "2020-10-09 11:35:17 AM Running production in window a013...\n",
      "2020-10-09 11:35:21 AM Running production in window p000...\n",
      "2020-10-09 11:35:25 AM Running production in window p001...\n",
      "2020-10-09 11:35:29 AM Running production in window p002...\n",
      "2020-10-09 11:35:33 AM Running production in window p003...\n",
      "2020-10-09 11:35:36 AM Running production in window p004...\n",
      "2020-10-09 11:35:40 AM Running production in window p005...\n",
      "2020-10-09 11:35:42 AM Running production in window p006...\n",
      "2020-10-09 11:35:46 AM Running production in window p007...\n",
      "2020-10-09 11:35:51 AM Running production in window p008...\n",
      "2020-10-09 11:35:54 AM Running production in window p009...\n",
      "2020-10-09 11:35:58 AM Running production in window p010...\n",
      "2020-10-09 11:36:06 AM Running production in window p011...\n",
      "2020-10-09 11:36:15 AM Running production in window p012...\n",
      "2020-10-09 11:36:21 AM Running production in window p013...\n",
      "2020-10-09 11:36:27 AM Running production in window p014...\n",
      "2020-10-09 11:36:32 AM Running production in window p015...\n",
      "2020-10-09 11:36:37 AM Running production in window p016...\n",
      "2020-10-09 11:36:41 AM Running production in window p017...\n"
     ]
    }
   ],
   "source": [
    "for window in window_list:\n",
    "    simulation = NAMD()\n",
    "    simulation.executable = \"namd2\"\n",
    "    simulation.n_threads = 4\n",
    "#     simulation.gpu_devices = 0\n",
    "    simulation.path = os.path.join(work_dir, window)\n",
    "    \n",
    "    simulation.topology = f\"{base_name}.prmtop\"\n",
    "    simulation.coordinates = f\"{base_name}.rst7\"\n",
    "    simulation.prefix = \"production\"\n",
    "    simulation.checkpoint = \"thermalization\"\n",
    "    simulation.colvars_file = \"colvars.tcl\"\n",
    "    \n",
    "    simulation.config_gb_md()\n",
    "    simulation.control[\"run\"] = 5000\n",
    "    \n",
    "    logging.info(f\"Running production in window {window}...\")\n",
    "    simulation.run(overwrite=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Analysis"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once the simulation is completed, we can using the `Analysis` module to determine the binding free energy. We supply the location of the parameter information, a string or list for the file names (wildcards supported), the location of the windows, and the restraints on the guest.\n",
    "\n",
    "In this example, we use the method `ti-block` which determines the free energy using **t**hermodynamic **i**integration and then estimates the standard error of the mean at each data point using blocking analysis. Bootstrapping it used to determine the uncertainty of the  full thermodynamic integral for each phase.\n",
    "\n",
    "After running `compute_free_energy()`, a dictionary called `results` will be populated, that contains the free energy and SEM for each phase of the simulation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "import paprika.analysis as analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "free_energy = analysis.fe_calc()\n",
    "free_energy.topology = \"cb6-but-dum.prmtop\"\n",
    "free_energy.trajectory = 'production*.dcd'\n",
    "free_energy.path = work_dir\n",
    "free_energy.restraint_list = guest_restraints\n",
    "free_energy.collect_data()\n",
    "free_energy.methods = ['ti-block']\n",
    "free_energy.ti_matrix = \"full\"\n",
    "free_energy.boot_cycles = 1000\n",
    "free_energy.compute_free_energy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We also need to calculate the free-energy cost of releasing the restraints on the guest molecule."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "free_energy.compute_ref_state_work([\n",
    "    guest_restraints[0], guest_restraints[1], None, \n",
    "    None, guest_restraints[2], None\n",
    "])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then we add the free-energies together and combine the uncertainties to get the binding-free energy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "binding_affinity = -1 * (\n",
    "    free_energy.results[\"attach\"][\"ti-block\"][\"fe\"] + \\\n",
    "    free_energy.results[\"pull\"][\"ti-block\"][\"fe\"] + \\\n",
    "    free_energy.results[\"ref_state_work\"]\n",
    ")\n",
    "\n",
    "sem = np.sqrt(\n",
    "    free_energy.results[\"attach\"][\"ti-block\"][\"sem\"]**2 + \\\n",
    "    free_energy.results[\"pull\"][\"ti-block\"][\"sem\"]**2\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The binding affinity of butane to cucurbit[6]uril = -6.16 +/- 4.40 kcal/mol\n"
     ]
    }
   ],
   "source": [
    "print(f\"The binding affinity of butane to cucurbit[6]uril = {binding_affinity.magnitude:0.2f} +/- {sem.magnitude:0.2f} kcal/mol\")"
   ]
  }
 ],
 "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.8"
  },
  "pycharm": {
   "stem_cell": {
    "cell_type": "raw",
    "metadata": {
     "collapsed": false
    },
    "source": []
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}