docs/tutorials/03-tutorial-k-cl.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# pAPRika tutorial 3 - K-Cl dissociation\n",
    "\n",
    "In this example, we will setup and simulate KCl dissociation using a single distance restraints. This tutorial will\n",
    "demonstrate the use of `pAPRika` besides host-guest systems."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-02-19T19:41:31.793751Z",
     "start_time": "2019-02-19T19:41:31.766398Z"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "\n",
    "import numpy as np\n",
    "import parmed as pmd\n",
    "\n",
    "from paprika.analysis import fe_calc\n",
    "from paprika.build.system import TLeap\n",
    "from paprika.io import NumpyEncoder\n",
    "from paprika.restraints.restraints import DAT_restraint\n",
    "from paprika.restraints.amber import amber_restraint_line",
    "from paprika.restraints.utils import create_window_list",
    "from paprika.simulate import AMBER"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Specify directory for data\n",
    "\n",
    "In this case, we just need some initial coordinates. In other cases, we might need `mol2` or `frcmod` files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-02-19T19:39:41.142698Z",
     "start_time": "2019-02-19T19:39:41.138350Z"
    }
   },
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "cwd = Path().resolve()\n",
    "k_cl_pdb = os.path.abspath(os.path.join(cwd, \"../../paprika/data/k-cl/k-cl.pdb\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup the calculation\n",
    "\n",
    "### Build the vacuum `prmtop` and `inpcrd` files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-02-19T19:41:54.354236Z",
     "start_time": "2019-02-19T19:41:54.168055Z"
    }
   },
   "outputs": [],
   "source": [
    "# Build the model in vacuum\n",
    "\n",
    "system = TLeap()\n",
    "system.template_lines = [\n",
    "    \"source leaprc.water.tip3p\",\n",
    "    \"loadamberparams frcmod.ionsjc_tip3p\",\n",
    "    f\"model = loadpdb {k_cl_pdb}\",\n",
    "]\n",
    "system.output_path = \"tmp\"\n",
    "system.output_prefix = \"k-cl\"\n",
    "system.pbc_type = None\n",
    "system.target_waters = None\n",
    "system.neutralize = False\n",
    "system.build()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-02-19T19:35:41.478380Z",
     "start_time": "2019-02-19T19:35:41.142816Z"
    }
   },
   "source": [
    "### Specify the number of windows for the umbrella sampling\n",
    "\n",
    "These are overkill; I have been testing how much data we need to converge this calculation and how quickly we can run a stripped-down version on Travis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-02-19T19:42:01.825732Z",
     "start_time": "2019-02-19T19:42:01.815923Z"
    }
   },
   "outputs": [],
   "source": [
    "attach_fractions = np.linspace(0, 1.0, 25)\n",
    "initial_distance = 2.65\n",
    "pull_distances = np.linspace(0 + initial_distance, 16.0 + initial_distance, 40)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Add a single distance restraint between K+ and Cl-"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-02-19T19:42:02.333845Z",
     "start_time": "2019-02-19T19:42:02.317742Z"
    }
   },
   "outputs": [],
   "source": [
    "restraint = DAT_restraint()\n",
    "restraint.continuous_apr = True\n",
    "restraint.amber_index = True\n",
    "restraint.topology = k_cl_pdb\n",
    "restraint.mask1 = \"@K+\"\n",
    "restraint.mask2 = \"@Cl-\"\n",
    "\n",
    "restraint.attach[\"target\"] = initial_distance\n",
    "restraint.attach[\"fraction_list\"] = attach_fractions\n",
    "restraint.attach[\"fc_final\"] = 10.0\n",
    "restraint.pull[\"fc\"] = restraint.attach[\"fc_final\"]\n",
    "restraint.pull[\"target_list\"] = pull_distances\n",
    "restraint.initialize()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Optionally, add a \"wall restraint\" to define the bound state and speed convergence\n",
    "\n",
    "This will apply a harmonic potential that keeps the \"guest\" Cl- within 3.5 Angstroms."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-02-19T19:42:03.287625Z",
     "start_time": "2019-02-19T19:42:03.240205Z"
    }
   },
   "outputs": [],
   "source": [
    "wall = DAT_restraint()\n",
    "wall.auto_apr = False\n",
    "wall.amber_index = True\n",
    "wall.topology = k_cl_pdb\n",
    "wall.mask1 = \"@K+\"\n",
    "wall.mask2 = \"@Cl-\"\n",
    "\n",
    "wall.attach[\"fc_initial\"] = 1.0\n",
    "wall.attach[\"fc_final\"] = 1.0\n",
    "\n",
    "wall.custom_restraint_values[\"rk2\"] = 1.0\n",
    "wall.custom_restraint_values[\"rk3\"] = 1.0\n",
    "wall.custom_restraint_values[\"r1\"] = 0.0\n",
    "wall.custom_restraint_values[\"r2\"] = 3.5\n",
    "wall.custom_restraint_values[\"r3\"] = 3.5\n",
    "wall.custom_restraint_values[\"r4\"] = 999\n",
    "\n",
    "wall.attach[\"target\"] = 3.5\n",
    "wall.attach[\"num_windows\"] = len(attach_fractions)\n",
    "\n",
    "wall.initialize()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create the directories for each window and write the AMBER-style restraint input file\n",
    "\n",
    "This makes it easy to run each window in parallel as a separate simulation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-02-19T19:42:04.734063Z",
     "start_time": "2019-02-19T19:42:04.382757Z"
    }
   },
   "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 0.4 Angstroms.\n",
      "In window p002 we will translate the guest 0.8 Angstroms.\n",
      "In window p003 we will translate the guest 1.2 Angstroms.\n",
      "In window p004 we will translate the guest 1.6 Angstroms.\n",
      "In window p005 we will translate the guest 2.1 Angstroms.\n",
      "In window p006 we will translate the guest 2.5 Angstroms.\n",
      "In window p007 we will translate the guest 2.9 Angstroms.\n",
      "In window p008 we will translate the guest 3.3 Angstroms.\n",
      "In window p009 we will translate the guest 3.7 Angstroms.\n",
      "In window p010 we will translate the guest 4.1 Angstroms.\n",
      "In window p011 we will translate the guest 4.5 Angstroms.\n",
      "In window p012 we will translate the guest 4.9 Angstroms.\n",
      "In window p013 we will translate the guest 5.3 Angstroms.\n",
      "In window p014 we will translate the guest 5.7 Angstroms.\n",
      "In window p015 we will translate the guest 6.2 Angstroms.\n",
      "In window p016 we will translate the guest 6.6 Angstroms.\n",
      "In window p017 we will translate the guest 7.0 Angstroms.\n",
      "In window p018 we will translate the guest 7.4 Angstroms.\n",
      "In window p019 we will translate the guest 7.8 Angstroms.\n",
      "In window p020 we will translate the guest 8.2 Angstroms.\n",
      "In window p021 we will translate the guest 8.6 Angstroms.\n",
      "In window p022 we will translate the guest 9.0 Angstroms.\n",
      "In window p023 we will translate the guest 9.4 Angstroms.\n",
      "In window p024 we will translate the guest 9.8 Angstroms.\n",
      "In window p025 we will translate the guest 10.3 Angstroms.\n",
      "In window p026 we will translate the guest 10.7 Angstroms.\n",
      "In window p027 we will translate the guest 11.1 Angstroms.\n",
      "In window p028 we will translate the guest 11.5 Angstroms.\n",
      "In window p029 we will translate the guest 11.9 Angstroms.\n",
      "In window p030 we will translate the guest 12.3 Angstroms.\n",
      "In window p031 we will translate the guest 12.7 Angstroms.\n",
      "In window p032 we will translate the guest 13.1 Angstroms.\n",
      "In window p033 we will translate the guest 13.5 Angstroms.\n",
      "In window p034 we will translate the guest 13.9 Angstroms.\n",
      "In window p035 we will translate the guest 14.4 Angstroms.\n",
      "In window p036 we will translate the guest 14.8 Angstroms.\n",
      "In window p037 we will translate the guest 15.2 Angstroms.\n",
      "In window p038 we will translate the guest 15.6 Angstroms.\n",
      "In window p039 we will translate the guest 16.0 Angstroms.\n"
     ]
    }
   ],
   "source": [
    "# Create folder for the working directory\n",
    "window_list = create_window_list([restraint])\n",
    "for window in window_list:\n",
    "    os.makedirs(f\"tmp/windows/{window}\", exist_ok=True)\n",
    "\n",
    "# Write the AMBER restraint files for each window\n",
    "for window in window_list:\n",
    "    with open(f\"tmp/windows/{window}/disang.rest\", \"a\") as file:\n",
    "        if window[0] == \"a\":\n",
    "            for r in [restraint, wall]:\n",
    "                string = amber_restraint_line(r, window)\n",
    "                if string is not None:\n",
    "                    file.write(string)\n",
    "        else:\n",
    "            string = amber_restraint_line(restraint, window)\n",
    "            file.write(string)\n",
    "\n",
    "# Generate the topology and coordinates file for each window \n",
    "for window in window_list:\n",
    "    if window[0] == \"a\":\n",
    "        structure = pmd.load_file(\"tmp/k-cl.prmtop\", \"tmp/k-cl.rst7\", structure=True)\n",
    "        for atom in structure.atoms:\n",
    "            if atom.name == \"Cl-\":\n",
    "                atom.xz = 2.65\n",
    "        structure.save(f\"tmp/windows/{window}/k-cl.prmtop\", overwrite=True)\n",
    "        structure.save(f\"tmp/windows/{window}/k-cl.rst7\", overwrite=True)\n",
    "\n",
    "    elif window[0] == \"p\":\n",
    "        structure = pmd.load_file(\"tmp/k-cl.prmtop\", \"tmp/k-cl.rst7\", structure=True)\n",
    "        target_difference = (\n",
    "            restraint.phase[\"pull\"][\"targets\"][int(window[1:])]\n",
    "            - restraint.phase[\"pull\"][\"targets\"][0]\n",
    "        )\n",
    "        print(\n",
    "            f\"In window {window} we will translate the guest {target_difference.magnitude:0.1f}.\"\n",
    "        )\n",
    "        for atom in structure.atoms:\n",
    "            if atom.name == \"Cl-\":\n",
    "                atom.xz += target_difference.magnitude\n",
    "        structure.save(f\"tmp/windows/{window}/k-cl.prmtop\", overwrite=True)\n",
    "        structure.save(f\"tmp/windows/{window}/k-cl.rst7\", overwrite=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Optionally, tweak some parameters, like changing the charge of K+ to 1.3 and Cl- to -1.3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-02-19T19:42:35.566455Z",
     "start_time": "2019-02-19T19:42:35.251649Z"
    }
   },
   "outputs": [],
   "source": [
    "for window in window_list:\n",
    "    structure = pmd.load_file(\n",
    "        f\"tmp/windows/{window}/k-cl.prmtop\",\n",
    "        f\"tmp/windows/{window}/k-cl.rst7\",\n",
    "        structure=True,\n",
    "    )\n",
    "    for atom in structure.atoms:\n",
    "        if atom.name == \"Cl-\":\n",
    "            atom.charge = -1.3\n",
    "        elif atom.name == \"K+\":\n",
    "            atom.charge = 1.3\n",
    "    structure.save(f\"tmp/windows/{window}/k-cl.prmtop\", overwrite=True)\n",
    "    structure.save(f\"tmp/windows/{window}/k-cl.rst7\", overwrite=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Solvate the structure in each window to the same number of waters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-02-19T19:46:31.525642Z",
     "start_time": "2019-02-19T19:43:05.990956Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Solvating window a000...\n",
      "Solvating window a001...\n",
      "Solvating window a002...\n",
      "Solvating window a003...\n",
      "Solvating window a004...\n",
      "Solvating window a005...\n",
      "Solvating window a006...\n",
      "Solvating window a007...\n",
      "Solvating window a008...\n",
      "Solvating window a009...\n",
      "Solvating window a010...\n",
      "Solvating window a011...\n",
      "Solvating window a012...\n",
      "Solvating window a013...\n",
      "Solvating window a014...\n",
      "Solvating window a015...\n",
      "Solvating window a016...\n",
      "Solvating window a017...\n",
      "Solvating window a018...\n",
      "Solvating window a019...\n",
      "Solvating window a020...\n",
      "Solvating window a021...\n",
      "Solvating window a022...\n",
      "Solvating window a023...\n",
      "Solvating window p000...\n",
      "Solvating window p001...\n",
      "Solvating window p002...\n",
      "Solvating window p003...\n",
      "Solvating window p004...\n",
      "Solvating window p005...\n",
      "Solvating window p006...\n",
      "Solvating window p007...\n",
      "Solvating window p008...\n",
      "Solvating window p009...\n",
      "Solvating window p010...\n",
      "Solvating window p011...\n",
      "Solvating window p012...\n",
      "Solvating window p013...\n",
      "Solvating window p014...\n",
      "Solvating window p015...\n",
      "Solvating window p016...\n",
      "Solvating window p017...\n",
      "Solvating window p018...\n",
      "Solvating window p019...\n",
      "Solvating window p020...\n",
      "Solvating window p021...\n",
      "Solvating window p022...\n",
      "Solvating window p023...\n",
      "Solvating window p024...\n",
      "Solvating window p025...\n",
      "Solvating window p026...\n",
      "Solvating window p027...\n",
      "Solvating window p028...\n",
      "Solvating window p029...\n",
      "Solvating window p030...\n",
      "Solvating window p031...\n",
      "Solvating window p032...\n",
      "Solvating window p033...\n",
      "Solvating window p034...\n",
      "Solvating window p035...\n",
      "Solvating window p036...\n",
      "Solvating window p037...\n",
      "Solvating window p038...\n",
      "Solvating window p039...\n"
     ]
    }
   ],
   "source": [
    "for window in window_list:\n",
    "    print(f\"Solvating window {window}...\")\n",
    "\n",
    "    if os.path.exists(f\"tmp/windows/{window}/k-cl-sol.prmtop\"):\n",
    "        print(\"Skipping...\")\n",
    "        continue\n",
    "\n",
    "    structure = pmd.load_file(\n",
    "        f\"tmp/windows/{window}/k-cl.prmtop\", f\"tmp/windows/{window}/k-cl.rst7\"\n",
    "    )\n",
    "\n",
    "    if not os.path.exists(f\"tmp/windows/{window}/k-cl.pdb\"):\n",
    "        structure.save(f\"tmp/windows/{window}/k-cl.pdb\")\n",
    "\n",
    "    system = TLeap()\n",
    "    system.output_path = os.path.join(\"tmp\", \"windows\", window)\n",
    "    system.output_prefix = \"k-cl-sol\"\n",
    "\n",
    "    system.target_waters = 2000\n",
    "    system.neutralize = False\n",
    "    system.template_lines = [\n",
    "        \"source leaprc.water.tip3p\",\n",
    "        \"model = loadpdb k-cl.pdb\"\n",
    "    ]\n",
    "    system.build()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run the calculation\n",
    "\n",
    "We have a few helper functions -- like `config_pbc_min()` and `config_pbc_md()` -- that help setup some smart defaults for AMBER. (I'll make a note to work on adding this for the OpenMM side of things.) The simulations can either be run directly, as indicated below, with `simulation.run()` or the input file can be written using `_amber_write_input_file()` and then wrapped using a cluster script (like PBS or whatever)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "start_time": "2019-02-19T19:47:47.281Z"
    },
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Energy Minimization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "for window in window_list:\n",
    "    simulation = AMBER()\n",
    "    simulation.executable = \"pmemd.cuda\"\n",
    "\n",
    "    simulation.path = f\"tmp/windows/{window}/\"\n",
    "    simulation.prefix = \"minimize\"\n",
    "\n",
    "    simulation.topology = \"k-cl-sol.prmtop\"\n",
    "    simulation.coordinates = \"k-cl-sol.rst7\"\n",
    "    simulation.ref = \"k-cl-sol.rst7\"\n",
    "    simulation.restraint_file = \"disang.rest\"\n",
    "\n",
    "    simulation.config_pbc_min()\n",
    "    simulation.cntrl[\"ntr\"] = 1\n",
    "    simulation.cntrl[\"restraint_wt\"] = 50.0\n",
    "    simulation.cntrl[\"restraintmask\"] = \"':1-2'\"\n",
    "    print(f\"Running minimization in window {window}...\")\n",
    "    simulation.run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Production Run"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Simulate\n",
    "for window in window_list:\n",
    "    simulation = AMBER()\n",
    "    simulation.executable = \"pmemd.cuda\"\n",
    "\n",
    "    simulation.path = f\"tmp/windows/{window}/\"\n",
    "    simulation.prefix = \"production\"\n",
    "\n",
    "    simulation.topology = \"k-cl-sol.prmtop\"\n",
    "    simulation.coordinate = \"minimize.rst7\"\n",
    "    simulation.ref = \"k-cl-sol.rst7\"\n",
    "    simulation.restraint_file = \"disang.rest\"\n",
    "\n",
    "    simulation.config_pbc_md()\n",
    "    simulation.cntrl[\"nstlim\"] = 50000\n",
    "\n",
    "    print(f\"Running production in window {window}...\")\n",
    "    simulation.run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Analyze the simulation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setup the analysis\n",
    "\n",
    "The analysis needs to know about:\n",
    "\n",
    "- The parameter file that was used for the molecules,\n",
    "- The simulation path,\n",
    "- The trajectories, and\n",
    "- The method to do the analysis (e.g., TI for the free energy with blocking analysis for the SEM)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "start_time": "2019-02-19T19:50:14.029Z"
    }
   },
   "outputs": [],
   "source": [
    "free_energy = fe_calc()\n",
    "free_energy.topology = \"k-cl-sol.prmtop\"\n",
    "free_energy.trajectory = \"production*.nc\"\n",
    "free_energy.path = \"tmp/windows\"\n",
    "free_energy.restraint_list = [restraint]\n",
    "free_energy.collect_data()\n",
    "free_energy.methods = [\"ti-block\", \"mbar-block\"]\n",
    "free_energy.ti_matrix = \"full\"\n",
    "free_energy.boot_cycles = 1000"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run the analysis and save the results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "free_energy.compute_free_energy()\n",
    "free_energy.compute_ref_state_work([restraint, None, None, None, None, None])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"./tmp/results.json\", \"w\") as f:\n",
    "    dumped = json.dumps(free_energy.results, cls=NumpyEncoder)\n",
    "    f.write(dumped)\n",
    "\n",
    "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",
    ")\n",
    "\n",
    "print(\n",
    "    f\"The binding affinity for K+ (+1.3) and Cl- (-1.3) = {binding_affinity.magnitude:0.2f} +/- {sem.magnitude:0.2f} kcal/mol\"\n",
    ")"
   ]
  }
 ],
 "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}