{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Centering a trajectory in the box\n",
    "\n",
    "Here we use MDAnalysis transformations to make a protein whole, center it in the box, and then wrap the water back into the box. We then look at how to do this on-the-fly.\n",
    "\n",
    "**Last updated:** December 2022 with MDAnalysis 2.4.0-dev0\n",
    "\n",
    "**Minimum version of MDAnalysis:** 1.0.0\n",
    "\n",
    "**Packages required:**\n",
    "    \n",
    "* MDAnalysis (<a data-cite=\"michaud-agrawal_mdanalysis_2011\" href=\"https://doi.org/10.1002/jcc.21787\">Michaud-Agrawal *et al.*, 2011</a>, <a data-cite=\"gowers_mdanalysis_2016\" href=\"https://doi.org/10.25080/Majora-629e541a-00e\">Gowers *et al.*, 2016</a>)\n",
    "* MDAnalysisTests\n",
    "\n",
    "**Optional packages for visualisation:**\n",
    "* [nglview](http://nglviewer.org/nglview/latest/) (<a data-cite=\"nguyen_nglviewinteractive_2018\" href=\"https://doi.org/10.1093/bioinformatics/btx789\">Nguyen *et al.*, 2018</a>)\n",
    "\n",
    "Throughout this tutorial we will include cells for visualising Universes with the [NGLView](http://nglviewer.org/nglview/latest/api.html) library. However, these will be commented out, and we will show the expected images generated instead of the interactive widgets.\n",
    "\n",
    "\n",
    "**See also:**\n",
    "\n",
    "* [On-the-fly transformations](../../trajectories/transformations.rst)\n",
    "* [On-the-fly transformations (blog post)](https://www.mdanalysis.org/2020/03/09/on-the-fly-transformations/)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-05-19T05:45:49.493629Z",
     "iopub.status.busy": "2021-05-19T05:45:49.492806Z",
     "iopub.status.idle": "2021-05-19T05:45:50.118657Z",
     "shell.execute_reply": "2021-05-19T05:45:50.119019Z"
    }
   },
   "outputs": [],
   "source": [
    "import MDAnalysis as mda\n",
    "from MDAnalysis.tests.datafiles import TPR, XTC\n",
    "import numpy as np\n",
    "# import nglview as nv"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loading files\n",
    "\n",
    "The test files we will be working with here are trajectories of a adenylate kinase (AdK), a phosophotransferase enzyme. (<a data-cite=\"beckstein_zipping_2009\" href=\"https://doi.org/10.1016/j.jmb.2009.09.009\">Beckstein *et al.*, 2009</a>)\n",
    "\n",
    "For the step-by-step transformations, we need to load the trajectory into memory so that our changes to the coordinates persist. If your trajectory is too large for that, see the [on-the-fly transformation](#Doing-all-this-on-the-fly) section for how to do this out-of-memory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-05-19T05:45:50.122672Z",
     "iopub.status.busy": "2021-05-19T05:45:50.122181Z",
     "iopub.status.idle": "2021-05-19T05:45:50.786263Z",
     "shell.execute_reply": "2021-05-19T05:45:50.786947Z"
    }
   },
   "outputs": [],
   "source": [
    "u = mda.Universe(TPR, XTC, in_memory=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Before transformation\n",
    "\n",
    "If you have NGLView installed, you can view the trajectory as it currently is."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-05-19T05:45:50.794872Z",
     "iopub.status.busy": "2021-05-19T05:45:50.793969Z",
     "iopub.status.idle": "2021-05-19T05:45:58.850848Z",
     "shell.execute_reply": "2021-05-19T05:45:58.818353Z"
    }
   },
   "outputs": [],
   "source": [
    "# view = nv.show_mdanalysis(u)\n",
    "# view.add_representation('point', 'resname SOL')\n",
    "# view.center()\n",
    "# view"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from nglview.contrib.movie import MovieMaker\n",
    "# movie = MovieMaker(\n",
    "#     view,\n",
    "#     step=2,\n",
    "#     render_params={\"factor\": 2},  # average quality render\n",
    "#     output='original.gif',\n",
    "# )\n",
    "# movie.make()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Otherwise, we embed a gif of it below.\n",
    "\n",
    "![original](original.gif)\n",
    "\n",
    "For easier analysis and nicer visualisation, we want to center this protein in the box."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Unwrapping the protein\n",
    "\n",
    "The first step is to \"unwrap\" the protein from the border of the box, to make the protein whole. MDAnalysis provides the `AtomGroup.unwrap` function to do this easily. Note that this function requires your universe to have bonds in it.\n",
    "\n",
    "We loop over the trajectory to unwrap for each frame."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-05-19T05:45:58.854539Z",
     "iopub.status.busy": "2021-05-19T05:45:58.853991Z",
     "iopub.status.idle": "2021-05-19T05:46:04.129889Z",
     "shell.execute_reply": "2021-05-19T05:46:04.129259Z"
    }
   },
   "outputs": [],
   "source": [
    "protein = u.select_atoms('protein')\n",
    "\n",
    "for ts in u.trajectory:\n",
    "    protein.unwrap(compound='fragments')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you can see, the protein is now whole, but not centered."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-05-19T05:46:04.141371Z",
     "iopub.status.busy": "2021-05-19T05:46:04.138141Z",
     "iopub.status.idle": "2021-05-19T05:46:11.057520Z",
     "shell.execute_reply": "2021-05-19T05:46:11.057135Z"
    }
   },
   "outputs": [],
   "source": [
    "# unwrapped = nv.show_mdanalysis(u)\n",
    "# unwrapped.add_representation('point', 'resname SOL')\n",
    "# unwrapped.center()\n",
    "# unwrapped"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Over the course of the trajectory it leaves the box.\n",
    "\n",
    "![unwrapped](unwrapped.gif)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Centering in the box\n",
    "\n",
    "The next step is to center the protein in the box. We calculate the center-of-mass of the protein and the center of the box for each timestep. We then move *all* the atoms so that the protein center-of-mass is in the center of the box.\n",
    "\n",
    "If you don't have masses in your trajectory, try using the `center_of_geometry`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-05-19T05:46:11.061760Z",
     "iopub.status.busy": "2021-05-19T05:46:11.061155Z",
     "iopub.status.idle": "2021-05-19T05:46:11.084322Z",
     "shell.execute_reply": "2021-05-19T05:46:11.084963Z"
    }
   },
   "outputs": [],
   "source": [
    "for ts in u.trajectory:\n",
    "    protein_center = protein.center_of_mass(wrap=True)\n",
    "    dim = ts.triclinic_dimensions\n",
    "    box_center = np.sum(dim, axis=0) / 2\n",
    "    u.atoms.translate(box_center - protein_center)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The protein is now in the center of the box, but the solvent is likely outside it, as we have just moved all the atoms."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-05-19T05:46:11.095271Z",
     "iopub.status.busy": "2021-05-19T05:46:11.092061Z",
     "iopub.status.idle": "2021-05-19T05:46:17.896781Z",
     "shell.execute_reply": "2021-05-19T05:46:17.896317Z"
    }
   },
   "outputs": [],
   "source": [
    "# centered = nv.show_mdanalysis(u)\n",
    "# centered.add_representation('point', 'resname SOL')\n",
    "# centered.center()\n",
    "# centered"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![centered](centered.gif)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Wrapping the solvent back into the box\n",
    "\n",
    "Luckily, MDAnalysis also has `AtomGroup.wrap` to wrap molecules back into the box. Our trajectory has dimensions defined, which the function will find automatically. If your trajectory does not, or you wish to use a differently sized box, you can pass in a ``box`` with dimensions in the form ``[lx, ly, lz, alpha, beta, gamma]``."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-05-19T05:46:17.900481Z",
     "iopub.status.busy": "2021-05-19T05:46:17.899971Z",
     "iopub.status.idle": "2021-05-19T05:46:21.558588Z",
     "shell.execute_reply": "2021-05-19T05:46:21.559143Z"
    }
   },
   "outputs": [],
   "source": [
    "not_protein = u.select_atoms('not protein')\n",
    "\n",
    "for ts in u.trajectory:\n",
    "    not_protein.wrap(compound='residues')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And now it is centered!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-05-19T05:46:21.568256Z",
     "iopub.status.busy": "2021-05-19T05:46:21.567471Z",
     "iopub.status.idle": "2021-05-19T05:46:28.762115Z",
     "shell.execute_reply": "2021-05-19T05:46:28.761729Z"
    }
   },
   "outputs": [],
   "source": [
    "# wrapped = nv.show_mdanalysis(u)\n",
    "# wrapped.add_representation('point', 'resname SOL')\n",
    "# wrapped.center()\n",
    "# wrapped"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![wrapped](wrapped.gif)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Doing all this on-the-fly\n",
    "\n",
    "Running all the transformations above can be difficult if your trajectory is large, or your computational resources are limited. Use on-the-fly transformations to keep your data out-of-memory.\n",
    "\n",
    "Some common transformations are defined in `MDAnalysis.transformations`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-05-19T05:46:28.765508Z",
     "iopub.status.busy": "2021-05-19T05:46:28.765022Z",
     "iopub.status.idle": "2021-05-19T05:46:28.833133Z",
     "shell.execute_reply": "2021-05-19T05:46:28.833588Z"
    }
   },
   "outputs": [],
   "source": [
    "import MDAnalysis.transformations as trans"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We re-load our universe."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-05-19T05:46:28.840377Z",
     "iopub.status.busy": "2021-05-19T05:46:28.839266Z",
     "iopub.status.idle": "2021-05-19T05:46:29.289898Z",
     "shell.execute_reply": "2021-05-19T05:46:29.290449Z"
    }
   },
   "outputs": [],
   "source": [
    "u2 = mda.Universe(TPR, XTC)\n",
    "\n",
    "protein2 = u2.select_atoms('protein')\n",
    "not_protein2 = u2.select_atoms('not protein')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "From version 1.0.0 onwards, the `MDAnalysis.transformations` module contains `wrap` and `unwrap` functions that correspond with the `AtomGroup` versions above. You can only use `add_transformations` *once*, so pass them all at the same time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-05-19T05:46:29.463298Z",
     "iopub.status.busy": "2021-05-19T05:46:29.332748Z",
     "iopub.status.idle": "2021-05-19T05:46:31.122496Z",
     "shell.execute_reply": "2021-05-19T05:46:31.122884Z"
    }
   },
   "outputs": [],
   "source": [
    "transforms = [trans.unwrap(protein2),\n",
    "              trans.center_in_box(protein2, wrap=True),\n",
    "              trans.wrap(not_protein2)]\n",
    "\n",
    "u2.trajectory.add_transformations(*transforms)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-05-19T05:46:31.137999Z",
     "iopub.status.busy": "2021-05-19T05:46:31.133162Z",
     "iopub.status.idle": "2021-05-19T05:46:40.016868Z",
     "shell.execute_reply": "2021-05-19T05:46:39.993052Z"
    }
   },
   "outputs": [],
   "source": [
    "# otf = nv.show_mdanalysis(u2)\n",
    "# otf.add_representation('point', 'resname SOL')\n",
    "# otf.center()\n",
    "# otf"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![on the fly](on_the_fly.gif)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## References\n",
    "\n",
    "[1] Oliver Beckstein, Elizabeth&nbsp;J. Denning, Juan&nbsp;R. Perilla, and Thomas&nbsp;B. Woolf.\n",
    "Zipping and <span class=\"bibtex-protected\">Unzipping</span> of <span class=\"bibtex-protected\">Adenylate</span> <span class=\"bibtex-protected\">Kinase</span>: <span class=\"bibtex-protected\">Atomistic</span> <span class=\"bibtex-protected\">Insights</span> into the <span class=\"bibtex-protected\">Ensemble</span> of <span class=\"bibtex-protected\">Open</span>↔<span class=\"bibtex-protected\">Closed</span> <span class=\"bibtex-protected\">Transitions</span>.\n",
    "<em>Journal of Molecular Biology</em>, 394(1):160–176, November 2009.\n",
    "00107.\n",
    "URL: <a href=\"https://linkinghub.elsevier.com/retrieve/pii/S0022283609011164\">https://linkinghub.elsevier.com/retrieve/pii/S0022283609011164</a>, <a href=\"https://doi.org/10.1016/j.jmb.2009.09.009\">doi:10.1016/j.jmb.2009.09.009</a>.\n",
    "\n",
    "[2] Richard&nbsp;J. Gowers, Max Linke, Jonathan Barnoud, Tyler J.&nbsp;E. Reddy, Manuel&nbsp;N. Melo, Sean&nbsp;L. Seyler, Jan Domański, David&nbsp;L. Dotson, Sébastien Buchoux, Ian&nbsp;M. Kenney, and Oliver Beckstein.\n",
    "<span class=\"bibtex-protected\">MDAnalysis</span>: <span class=\"bibtex-protected\">A</span> <span class=\"bibtex-protected\">Python</span> <span class=\"bibtex-protected\">Package</span> for the <span class=\"bibtex-protected\">Rapid</span> <span class=\"bibtex-protected\">Analysis</span> of <span class=\"bibtex-protected\">Molecular</span> <span class=\"bibtex-protected\">Dynamics</span> <span class=\"bibtex-protected\">Simulations</span>.\n",
    "<em>Proceedings of the 15th Python in Science Conference</em>, pages 98–105, 2016.\n",
    "00152.\n",
    "URL: <a href=\"https://conference.scipy.org/proceedings/scipy2016/oliver_beckstein.html\">https://conference.scipy.org/proceedings/scipy2016/oliver_beckstein.html</a>, <a href=\"https://doi.org/10.25080/Majora-629e541a-00e\">doi:10.25080/Majora-629e541a-00e</a>.\n",
    "\n",
    "[3] Naveen Michaud-Agrawal, Elizabeth&nbsp;J. Denning, Thomas&nbsp;B. Woolf, and Oliver Beckstein.\n",
    "<span class=\"bibtex-protected\">MDAnalysis</span>: <span class=\"bibtex-protected\">A</span> toolkit for the analysis of molecular dynamics simulations.\n",
    "<em>Journal of Computational Chemistry</em>, 32(10):2319–2327, July 2011.\n",
    "00778.\n",
    "URL: <a href=\"http://doi.wiley.com/10.1002/jcc.21787\">http://doi.wiley.com/10.1002/jcc.21787</a>, <a href=\"https://doi.org/10.1002/jcc.21787\">doi:10.1002/jcc.21787</a>.\n",
    "\n",
    "[4] Hai Nguyen, David&nbsp;A Case, and Alexander&nbsp;S Rose.\n",
    "<span class=\"bibtex-protected\">NGLview</span>–interactive molecular graphics for <span class=\"bibtex-protected\">Jupyter</span> notebooks.\n",
    "<em>Bioinformatics</em>, 34(7):1241–1242, April 2018.\n",
    "00024.\n",
    "URL: <a href=\"https://academic.oup.com/bioinformatics/article/34/7/1241/4721781\">https://academic.oup.com/bioinformatics/article/34/7/1241/4721781</a>, <a href=\"https://doi.org/10.1093/bioinformatics/btx789\">doi:10.1093/bioinformatics/btx789</a>."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.2"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": false,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  },
  "vscode": {
   "interpreter": {
    "hash": "7b52aa17ef4e9358c0e91ff3f0bf50d10667bb8b55636d4b9fb967a2ff94bd8c"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}