notebooks/Tutorial3_Analysis.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tutorial 3: Analysis\n",
    "\n",
    "<a id='trajanalysis'></a>\n",
    "\n",
    "<a rel=\"license\" href=\"http://creativecommons.org/licenses/by/4.0/\"><img alt=\"Creative Commons Licence\" style=\"border-width:0\" src=\"https://i.creativecommons.org/l/by/4.0/88x31.png\" title='This work is licensed under a Creative Commons Attribution 4.0 International License.' align=\"right\"/></a>\n",
    "\n",
    "This notebook is adapted from materials developed for the [2018 Workshop/Hackathon](https://github.com/MDAnalysis/WorkshopHackathon2018) and the [MDAnalysis User Guide](https://userguide.mdanalysis.org/stable/examples/analysis/custom_trajectory_analysis.html#Creating-an-analysis-from-a-function).\n",
    "\n",
    "### Learning outcomes \n",
    "\n",
    "- Understand how analysis methods that act on `Atomgroups`, such as root mean square deviation (RMSD), work\n",
    "- Creating your own analysis class with MDAnalysis\n",
    "\n",
    "#### Additional resources\n",
    "\n",
    " - During the workshop, feel free to ask questions at any time\n",
    " - For more on how to use MDAnalysis, see the [User Guide](https://userguide.mdanalysis.org/stable/index.html) and [documentation](https://docs.mdanalysis.org/stable/index.html)\n",
    " - Ask questions on the [GitHub Discussions forum](https://github.com/MDAnalysis/mdanalysis/discussions) or on [Discord](https://discord.gg/fXTSfDJyxE)\n",
    " - Report bugs on [GitHub](https://github.com/MDAnalysis/mdanalysis/issues?)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Google Colab package installs\n",
    "\n",
    "This installs the necessary packages for Google Colab. Please only run these if you are using Colab."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NBVAL_SKIP\n",
    "!pip install condacolab\n",
    "import condacolab\n",
    "condacolab.install()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NBVAL_SKIP\n",
    "import condacolab\n",
    "condacolab.check()\n",
    "!mamba install -c conda-forge mdanalysis mdanalysistests mdanalysisdata nglview rdkit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NBVAL_SKIP\n",
    "# enable third party jupyter widgets\n",
    "from google.colab import output\n",
    "output.enable_custom_widget_manager()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **Jupyter cheat sheet**\n",
    "- to run the currently highlighted cell, hold <kbd>&#x21E7; Shift</kbd> and press <kbd>&#x23ce; Enter</kbd>;\n",
    "- to get help for a specific function, place the cursor within the function's brackets, hold <kbd>&#x21E7; Shift</kbd>, and press <kbd>&#x21E5; Tab</kbd>;"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-warning\"><b> REMEMBER: variables persist between cells</b> \n",
    "    \n",
    "Be aware that it is the order of execution of cells that is important in a Jupyter notebook, not the <em>order</em> in which they appear. Python will remember <em>all</em> the code that was run previously, including any variables you have defined, irrespective of the order in the notebook. Therefore if you define variables lower down the notebook and then (re)run cells further up, those defined further down will still be present. </div> "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\") \n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "import numpy as np\n",
    "import nglview as nv\n",
    "import MDAnalysis as mda\n",
    "from MDAnalysis.lib import distances \n",
    "from MDAnalysis.analysis import rdf\n",
    "import MDAnalysisData as data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Root mean square deviation (RMSD)\n",
    "\n",
    "Here we calculate the RMSD of domains in adenylate kinase (AdK), a phosophotransferase enzyme, as it transitions from an open to closed structure."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from MDAnalysis.tests.datafiles import PSF, DCD, CRD\n",
    "from MDAnalysis.analysis import rms\n",
    "\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The trajectory DCD samples a transition from a closed to an open conformation. AdK has three domains:\n",
    "\n",
    "- CORE\n",
    "- LID: an ATP-binding domain\n",
    "- NMP: an AMP-binding domain\n",
    "\n",
    "The LID and NMP domains move around the stable CORE as the enzyme transitions between the opened and closed conformations. One way to quantify this movement is by calculating the root mean square deviation (RMSD) of atomic positions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "u = mda.Universe(PSF, DCD)  # closed AdK (PDB ID: 1AKE)\n",
    "ref = mda.Universe(PSF, CRD)  # open AdK (PDB ID: 4AKE)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Background: RMSD\n",
    "\n",
    "The root mean square deviation (RMSD) of particle coordinates is one measure of distance, or dissimilarity, between molecular conformations. Each structure should have matching elementwise atoms $i$ in the same order, as the distance between them is calculated and summed for the final result. It is calculated between coordinate arrays $\\mathbf{x}$ and $\\mathbf{x}^{\\text{ref}}$ according to the equation below:\n",
    "\n",
    "$$ \\text{RMSD}(\\mathbf{x}, \\mathbf{x}^{\\text{ref}}) = \\sqrt{\\frac{1}{n} \\sum_{i=1}^{n}{|\\mathbf{x}_i-\\mathbf{x}_i^{\\text{ref}}|^2}} $$\n",
    "\n",
    "As molecules can move around, the structure $\\mathbf{x}$ is usually translated by a vector $\\mathbf{t}$ and rotated by a matrix $\\mathsf{R}$ to align with the reference $\\mathbf{x}^{\\text{ref}}$ such that the RMSD is minimised. The RMSD after this optimal superposition can be expressed as follows:\n",
    "\n",
    "$$ \\text{RMSD}(\\mathbf{x}, \\mathbf{x}^{\\text{ref}}) = \\min_{\\mathsf{R}, \\mathbf{t}} %\n",
    "  \\sqrt{\\frac{1}{N} \\sum_{i=1}^{N} \\left[ %\n",
    "      (\\mathsf{R}\\cdot\\mathbf{x}_{i}(t) + \\mathbf{t}) - \\mathbf{x}_{i}^{\\text{ref}} \\right]^{2}}$$\n",
    "\n",
    "The RMSD between one reference state and a trajectory of structures is often calculated as a way to measure the dissimilarity of the trajectory conformational ensemble to the reference. This reference is frequently the first frame of the trajectory (the default in MDAnalysis), in which case it can provide insight into the overall movement from the initial starting point. W\n",
    "\n",
    "Typically not all coordinates in a structures are included in an RMSD analysis. With proteins, the fluctuation of the residue side-chains is not representative of overall conformational change. Therefore when RMSD analyses are performed to investigate large-scale movements in proteins, the atoms are usually restricted only to the backbone atoms (forming the amide-bond chain) or the alpha-carbon atoms. \n",
    "\n",
    "MDAnalysis provides functions and classes to calculate the RMSD between coordinate arrays, and `Universes` or `AtomGroups`.  The contribution of each particle $i$ to the final RMSD value can also be weighted.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### RMSD between two sets of coordinates\n",
    "\n",
    "The MDAnalysis.analysis.rms.rmsd function returns the root mean square deviation (in Angstrom) between two sets of coordinates. Here, we calculate the RMSD between the backbone atoms of the open and closed conformations of AdK. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# RMSD Function Implementing the Above Equation\n",
    "backbone_ag = u.select_atoms('backbone')\n",
    "ref_ag = ref.select_atoms('backbone')\n",
    "rms.rmsd(backbone_ag.positions,  # coordinates to align\n",
    "         ref_ag.positions,  # reference coordinates\n",
    "         center=True,  # subtract the center of geometry\n",
    "         superposition=True)  # superimpose coordinates"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-success\"> <b>Exercise 1</b> \n",
    "\n",
    "Write your own trajectory iteration and get the RMSD of backbone atoms relative to the reference over time.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Exercise 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Exercise 1 solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Exercise 1 solution\n",
    "\n",
    "rmsd_results = np.zeros(u.trajectory.n_frames)\n",
    "\n",
    "backbone_ag = u.select_atoms('backbone')\n",
    "ref_ag = ref.select_atoms('backbone')\n",
    "\n",
    "for i, ts in enumerate(u.trajectory):\n",
    "    rmsd_results[i] = rms.rmsd(backbone_ag.positions,  # coordinates to align\n",
    "                               ref_ag.positions,  # reference coordinates\n",
    "                               center=True,  # subtract the center of geometry\n",
    "                               superposition=True)  # superimpose coordinates\n",
    "\n",
    "print(rmsd_results)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### RMSD analysis class\n",
    "\n",
    "It is more efficient to use the `MDAnalysis.analysis.rms.RMSD` class to calculate the RMSD of an entire trajectory to a single reference point, than to use the the `MDAnalysis.analysis.rms.rmsd` function.\n",
    "\n",
    "The `rms.RMSD` class first performs a rotational and translational alignment of the target trajectory to the reference universe at `ref_frame`, using the atoms in select to determine the transformation. The RMSD of the select selection is calculated."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-success\"> <b>Exercise 2</b> \n",
    "\n",
    "Now use `RMSD` class to achieve the same goal.\n",
    "\n",
    "Hint: refer to https://docs.mdanalysis.org/stable/documentation_pages/analysis/rms.html#MDAnalysis.analysis.rms.RMSD or use the magic command `rms.RMSD?`\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Exercise 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Exercise 2 solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Exercise 2 solution\n",
    "\n",
    "backbone_ag = u.select_atoms('backbone')\n",
    "ref_ag = ref.select_atoms('backbone')\n",
    "R = rms.RMSD(atomgroup=backbone_ag,\n",
    "             reference=ref_ag)\n",
    "R.run()\n",
    "\n",
    "R_results = R.results.rmsd.T[2] # the third column is the rmsd value\n",
    "\n",
    "# Alternative solution\n",
    "\n",
    "R_2 = rms.RMSD(atomgroup=u,\n",
    "             reference=ref,\n",
    "             select='backbone')\n",
    "R_2.run()\n",
    "\n",
    "R_2_results = R_2.results.rmsd.T[2] # the third column is the rmsd value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(R_results)\n",
    "plt.plot(R_2_results)\n",
    "plt.plot(rmsd_results)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### RMSD with multiple selections"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-success\">\n",
    "  <b>Exercise 3</b> \n",
    "\n",
    "  I now want to see the relative movement of each important domain (CORE, LID, NMP defined below) based on the RMSD value relative to its **starting position**. What could I do?\n",
    "\n",
    "  Hint 1: Align the system based on all the backbone atoms.\n",
    "\n",
    "  Hint 2: Use `groupselections`.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "CORE = 'backbone and (resid 1-29 or resid 60-121 or resid 160-214)'\n",
    "LID = 'backbone and resid 122-159'\n",
    "NMP = 'backbone and resid 30-59'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Exercise 3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Exercise 3 solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "R = rms.RMSD(u,  # universe to align\n",
    "             u,  # reference universe or atomgroup\n",
    "             select='backbone',  # group to superimpose and calculate RMSD\n",
    "             groupselections=[CORE, LID, NMP],  # groups for RMSD\n",
    "             ref_frame=0)  # frame index of the reference\n",
    "R.run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plotting the data\n",
    "\n",
    "We can easily plot this data using the common data analysis package `pandas`. We turn the `R.rmsd` array into a `DataFrame` and label each column below.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.DataFrame(R.results.rmsd,\n",
    "                  columns=['Frame', 'Time (ns)',\n",
    "                           'Backbone', 'CORE',\n",
    "                           'LID', 'NMP'])\n",
    "\n",
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ax = df.plot(x='Frame', y=['Backbone', 'CORE', 'LID', 'NMP'],\n",
    "             kind='line')\n",
    "ax.set_ylabel(r'RMSD ($\\AA$)')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-info\"> <b> Reminder: </b> \n",
    "\n",
    "\n",
    "Some analysis tools (`radius_of_gyration`, `center_of_mass` etc) act on `positions` (which are properties of a single `timestep`). To calculate the property of interest for each frame, we need to iterate over the whole trajectory. \n",
    "\n",
    "Others (`RDF`, `rms.RMSD` etc) act on `AtomGroups`, and can iterate over the whole trajectory for us.\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Writing your own analysis class\n",
    "\n",
    "We will now demonstrate how to write your own analysis class.\n",
    "\n",
    "Note: there are methods e.g. `AnalysisFromFunction`, `analysis_class` that can convert simple function into a functional analysis class.\n",
    "\n",
    "**If you implement your own analysis method, please consider [contributing it to the MDAnalysis codebase!](https://www.mdanalysis.org/UserGuide/contributing.html) or creating an MDAKit!**\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-06-09T12:13:02.845041605Z",
     "start_time": "2023-06-09T12:13:00.083014779Z"
    },
    "execution": {
     "iopub.execute_input": "2021-05-19T05:45:44.674135Z",
     "iopub.status.busy": "2021-05-19T05:45:44.673007Z",
     "iopub.status.idle": "2021-05-19T05:45:46.290616Z",
     "shell.execute_reply": "2021-05-19T05:45:46.291341Z"
    }
   },
   "outputs": [],
   "source": [
    "from MDAnalysis.tests.datafiles import PSF, DCD\n",
    "from MDAnalysis.analysis.base import AnalysisBase\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Radius of gyration\n",
    "\n",
    "Let's start off by defining a standalone analysis function.\n",
    "\n",
    "The radius of gyration of a structure measures how compact it is. In [GROMACS](http://manual.gromacs.org/documentation/2019-rc1/reference-manual/analysis/radius-of-gyration.html), it is calculated as follows: \n",
    "\n",
    "$$ R_g = \\sqrt{\\frac{\\sum_i m_i \\mathbf{r}_i^2}{\\sum_i m_i}}$$\n",
    "\n",
    "where $m_i$ is the mass of atom $i$ and $\\mathbf{r}_i$ is the position of atom $i$, relative to the center-of-mass of the selection.\n",
    "\n",
    "The radius of gyration around each axis can also be determined separately. For example, the radius of gyration around the x-axis:\n",
    "\n",
    "$$ R_{i, x} = \\sqrt{\\frac{\\sum_i m_i [r_{i, y}^2 + r_{i, z}^2]}{\\sum_i m_i}}$$\n",
    "\n",
    "Below, we define a function that takes an AtomGroup and calculates the radii of gyration. We could write this function to only need the AtomGroup. However, we also add in a `masses` argument and a `total_mass` keyword to avoid recomputing the mass and total mass for each frame."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-06-09T12:13:02.913378701Z",
     "start_time": "2023-06-09T12:13:02.846193220Z"
    },
    "execution": {
     "iopub.execute_input": "2021-05-19T05:45:46.297926Z",
     "iopub.status.busy": "2021-05-19T05:45:46.297404Z",
     "iopub.status.idle": "2021-05-19T05:45:46.299247Z",
     "shell.execute_reply": "2021-05-19T05:45:46.299751Z"
    }
   },
   "outputs": [],
   "source": [
    "def radgyr(atomgroup, masses, total_mass=None):\n",
    "    # coordinates change for each frame\n",
    "    coordinates = atomgroup.positions\n",
    "    center_of_mass = atomgroup.center_of_mass()\n",
    "    \n",
    "    # get squared distance from center\n",
    "    ri_sq = (coordinates-center_of_mass)**2\n",
    "    # sum the unweighted positions\n",
    "    sq = np.sum(ri_sq, axis=1)\n",
    "    sq_x = np.sum(ri_sq[:,[1,2]], axis=1) # sum over y and z\n",
    "    sq_y = np.sum(ri_sq[:,[0,2]], axis=1) # sum over x and z\n",
    "    sq_z = np.sum(ri_sq[:,[0,1]], axis=1) # sum over x and y\n",
    "    \n",
    "    # make into array\n",
    "    sq_rs = np.array([sq, sq_x, sq_y, sq_z])\n",
    "    \n",
    "    # weight positions\n",
    "    rog_sq = np.sum(masses*sq_rs, axis=1)/total_mass\n",
    "    # square root and return\n",
    "    return np.sqrt(rog_sq)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "I want to calculate the radius of gyration of my protein overall and along three dimensions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "u = mda.Universe(PSF, DCD)\n",
    "protein = u.select_atoms('protein')\n",
    "protein_masses = protein.masses\n",
    "total_mass = np.sum(protein_masses)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "radgyr_results = np.zeros([u.trajectory.n_frames, 4])\n",
    "\n",
    "for i, ts in enumerate(u.trajectory):\n",
    "    radgyr_results[i] = radgyr(atomgroup=protein,\n",
    "                          masses=protein_masses,\n",
    "                          total_mass=total_mass)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "labels = ['all', 'x-axis', 'y-axis', 'z-axis']\n",
    "for col, label in zip(radgyr_results.T, labels):\n",
    "    plt.plot(col, label=label)\n",
    "plt.legend()\n",
    "plt.ylabel('Radius of gyration (Å)')\n",
    "plt.xlabel('Frame')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "view = nv.show_mdanalysis(protein)\n",
    "view.add_representation('axes')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "view"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-success\">\n",
    "  <b>Exercise 4</b> \n",
    "\n",
    "  I want to create a `RadiusOfGyration` class that inherits from `AnalysisBase` to add this functionality to my MD toolbox. Recall the three methods that I need to implement:\n",
    "  - `__init__`\n",
    "  - `_prepare`\n",
    "  - `_single_frame`\n",
    "\n",
    "  Here’s a reference for how the `run` method in `AnalysisBase` works:\n",
    "  ```python\n",
    "  class AnalysisBase:\n",
    "     ...\n",
    "     def run(self, start, stop, step, ...):\n",
    "          self._setup_frames()\n",
    "          self._prepare()\n",
    "          for ts in self.trajectory:\n",
    "              self._single_frame()\n",
    "          self._conclude()\n",
    "          return self\n",
    "  ```\n",
    "\n",
    "  Hint 1: Use `self._frame_index` to retrieve the current frame index.\n",
    "\n",
    "  Hint 2: Use the `radgyr` function in `_single_frame`.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Exercise 4\n",
    "class RadiusOfGyration(AnalysisBase):\n",
    "    def __init__(self, atomgroup, **kwargs):\n",
    "        # set parameters\n",
    "        ...\n",
    "        super().__init__(atomgroup.universe.trajectory,\n",
    "                         **kwargs)\n",
    "\n",
    "    def _prepare(self):\n",
    "        # called before iteration on the trajectory has begun.\n",
    "        # initialize results\n",
    "        ...\n",
    "\n",
    "    def _single_frame(self):\n",
    "        # called after the trajectory is moved onto each new frame.\n",
    "        ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rog = RadiusOfGyration(atomgroup=protein)\n",
    "rog.run()\n",
    "\n",
    "rog_results = rog.results.radgyr\n",
    "\n",
    "labels = ['all', 'x-axis', 'y-axis', 'z-axis']\n",
    "for col, label in zip(rog_results.T, labels):\n",
    "    plt.plot(col, label=label)\n",
    "plt.legend()\n",
    "plt.ylabel('Radius of gyration (Å)')\n",
    "plt.xlabel('Frame')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Exercise 4 solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Exercise 4 solution\n",
    "class RadiusOfGyration(AnalysisBase):\n",
    "    def __init__(self, atomgroup, **kwargs):\n",
    "        self.atomgroup = atomgroup\n",
    "        self.atom_masses = atomgroup.masses\n",
    "        self.total_mass = np.sum(atomgroup.masses)\n",
    "        super().__init__(atomgroup.universe.trajectory,\n",
    "                         **kwargs)\n",
    "\n",
    "    def _prepare(self):\n",
    "        # called before iteration on the trajectory has begun.\n",
    "        # initialize results\n",
    "        self.results.radgyr = np.zeros([u.trajectory.n_frames, 4])\n",
    "\n",
    "    def _single_frame(self):\n",
    "        # called after the trajectory is moved onto each new frame.\n",
    "        self.results.radgyr[self._frame_index] = radgyr(self.atomgroup, self.atom_masses, self.total_mass)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rog = RadiusOfGyration(atomgroup=protein)\n",
    "rog.run()\n",
    "\n",
    "rog_results = rog.results.radgyr\n",
    "\n",
    "labels = ['all', 'x-axis', 'y-axis', 'z-axis']\n",
    "for col, label in zip(rog_results.T, labels):\n",
    "    plt.plot(col, label=label)\n",
    "plt.legend()\n",
    "plt.ylabel('Radius of gyration (Å)')\n",
    "plt.xlabel('Frame')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sharing your analysis methods with the community!\n",
    "\n",
    "If you think that you will want to reuse your new analysis, or that others might find it helpful, please consider creating an MDAKit for it! Making your code open-source can have many benefits; others may notice an unexpected bug or suggest ways to optimise your code."
   ]
  }
 ],
 "metadata": {
  "hide_input": false,
  "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.11.9"
  },
  "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": {
    "height": "calc(100% - 180px)",
    "left": "10px",
    "top": "150px",
    "width": "288px"
   },
   "toc_section_display": true,
   "toc_window_display": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}