notebooks/00-pentapeptide-showcase.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 00 - Showcase pentapeptide: a PyEMMA walkthrough\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",
    "In this notebook, we introduce the most basic features of PyEMMA.\n",
    "Overall, the notebook serves as an example workflow for analyzing molecular dynamics trajectories.\n",
    "We keep the details to an absolute minimum and refer to the more specialized notebooks.\n",
    "There, you can find more details, exercises, and some theory.\n",
    "\n",
    "Maintainers: [@cwehmeyer](https://github.com/cwehmeyer), [@marscher](https://github.com/marscher), [@thempel](https://github.com/thempel), [@psolsson](https://github.com/psolsson)\n",
    "\n",
    "**Remember**:\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>;\n",
    "- you can find the full documentation at [PyEMMA.org](http://www.pyemma.org).\n",
    "\n",
    "---\n",
    "\n",
    "We show the PyEMMA workflow using 25 independent simulations of a pentapeptide:\n",
    "\n",
    "<img src=\"static/pentapeptide-structure.png\" width=\"320\" height=\"171\" />\n",
    "\n",
    "The pentapeptide was simulated with implicit solvent and the trajectories were saved with a $0.1$ ns timestep."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib as mpl\n",
    "import numpy as np\n",
    "import mdshare\n",
    "import pyemma\n",
    "from pyemma.util.contexts import settings"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data input and featurization  [➜ 📓](01-data-io-and-featurization.ipynb) \n",
    "### Loading data\n",
    "\n",
    "We start our short walkthrough tutorial by loading a topology file (in this case, a PDB) and the trajectory data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "pdb = mdshare.fetch('pentapeptide-impl-solv.pdb', working_directory='data')\n",
    "files = mdshare.fetch('pentapeptide-*-500ns-impl-solv.xtc', working_directory='data')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The above cell has downloaded the data from our servers;\n",
    "in general, only strings of file paths need to be provided."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(pdb)\n",
    "print(files)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As it is unknown which feature describes the system best, we start with a broad systematic analysis.\n",
    "For the sake of simplicity, we are only interested in modeling the backbone kinetics.\n",
    "Thus, we only consider features describing the backbone and not the side chains.\n",
    "\n",
    "In PyEMMA, the `featurizer` is a central object that incorporates the system's topology.\n",
    "Features are easily computed by adding the target feature, e.g.,\n",
    "with `featurizer.add_backbone_torsions()`.\n",
    "We will load backbone torsion angles, backbone heavy atom positions, and backbone heavy atom distances. \n",
    "\n",
    "⚠️ Please note that the structures have been aligned before.\n",
    "Since in that case we loose track of the periodic box,\n",
    "we have to switch off the `periodic` flag for the distance and torsion angle computations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torsions_feat = pyemma.coordinates.featurizer(pdb)\n",
    "torsions_feat.add_backbone_torsions(cossin=True, periodic=False)\n",
    "torsions_data = pyemma.coordinates.load(files, features=torsions_feat)\n",
    "labels = ['backbone\\ntorsions']\n",
    "\n",
    "positions_feat = pyemma.coordinates.featurizer(pdb)\n",
    "positions_feat.add_selection(positions_feat.select_Backbone())\n",
    "positions_data = pyemma.coordinates.load(files, features=positions_feat)\n",
    "labels += ['backbone atom\\npositions']\n",
    "\n",
    "distances_feat = pyemma.coordinates.featurizer(pdb)\n",
    "distances_feat.add_distances(\n",
    "    distances_feat.pairs(distances_feat.select_Backbone(), excluded_neighbors=2), periodic=False)\n",
    "distances_data = pyemma.coordinates.load(files, features=distances_feat)\n",
    "labels += ['backbone atom\\ndistances']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Feature selection\n",
    "We will now rank the three featurizations by means of the VAMP2 score which measures the kinetic variance contained in these features <a id=\"ref-1\" href=\"#cite-gmrq\">mcgibbon-15</a>, <a id=\"ref-2\" href=\"#cite-vamp-preprint\">wu-17</a>, <a id=\"ref-3\" href=\"#cite-vampnet\">mardt-17</a>.\n",
    "The minimum value of this score is $1$, which corresponds to the invariant measure or equilibrium.\n",
    "\n",
    "As we are comparing featurizations with different dimensionality,\n",
    "we use the dimension parameter to provide an upper bound for the number of dynamic processes included in the scoring."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def score_cv(data, dim, lag, number_of_splits=10, validation_fraction=0.5):\n",
    "    \"\"\"Compute a cross-validated VAMP2 score.\n",
    "    \n",
    "    We randomly split the list of independent trajectories into\n",
    "    a training and a validation set, compute the VAMP2 score,\n",
    "    and repeat this process several times.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    data : list of numpy.ndarrays\n",
    "        The input data.\n",
    "    dim : int\n",
    "        Number of processes to score; equivalent to the dimension\n",
    "        after projecting the data with VAMP2.\n",
    "    lag : int\n",
    "        Lag time for the VAMP2 scoring.\n",
    "    number_of_splits : int, optional, default=10\n",
    "        How often do we repeat the splitting and score calculation.\n",
    "    validation_fraction : int, optional, default=0.5\n",
    "        Fraction of trajectories which should go into the validation\n",
    "        set during a split.\n",
    "    \"\"\"\n",
    "    # we temporarily suppress very short-lived progress bars\n",
    "    with pyemma.util.contexts.settings(show_progress_bars=False):\n",
    "        nval = int(len(data) * validation_fraction)\n",
    "        scores = np.zeros(number_of_splits)\n",
    "        for n in range(number_of_splits):\n",
    "            ival = np.random.choice(len(data), size=nval, replace=False)\n",
    "            vamp = pyemma.coordinates.vamp(\n",
    "                [d for i, d in enumerate(data) if i not in ival], lag=lag, dim=dim)\n",
    "            scores[n] = vamp.score([d for i, d in enumerate(data) if i in ival])\n",
    "    return scores\n",
    "\n",
    "\n",
    "dim = 10\n",
    "\n",
    "fig, axes = plt.subplots(1, 3, figsize=(12, 3), sharey=True)\n",
    "for ax, lag in zip(axes.flat, [5, 10, 20]):\n",
    "    torsions_scores = score_cv(torsions_data, lag=lag, dim=dim)\n",
    "    scores = [torsions_scores.mean()]\n",
    "    errors = [torsions_scores.std()]\n",
    "    positions_scores = score_cv(positions_data, lag=lag, dim=dim)\n",
    "    scores += [positions_scores.mean()]\n",
    "    errors += [positions_scores.std()]\n",
    "    distances_scores = score_cv(distances_data, lag=lag, dim=dim)\n",
    "    scores += [distances_scores.mean()]\n",
    "    errors += [distances_scores.std()]\n",
    "    ax.bar(labels, scores, yerr=errors, color=['C0', 'C1', 'C2'])\n",
    "    ax.set_title(r'lag time $\\tau$={:.1f}ns'.format(lag * 0.1))\n",
    "    if lag == 5:\n",
    "        # save for later\n",
    "        vamp_bars_plot = dict(\n",
    "            labels=labels, scores=scores, errors=errors, dim=dim, lag=lag)\n",
    "axes[0].set_ylabel('VAMP2 score')\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "⚠️ Please note that the VAMP-2 score is **not** suitable to select the proper lag time as scores for different lag times are not comparable.\n",
    "\n",
    "Here, we just compare different features at each given lag time separately.\n",
    "We consistently find backbone torsions to have a slight advantage.\n",
    "\n",
    "Thus, we add a more detailed VAMP2 score analysis for backbone torsions by varying the dimension parameter for several lag times:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lags = [1, 2, 5, 10, 20]\n",
    "dims = [i + 1 for i in range(10)]\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "for i, lag in enumerate(lags):\n",
    "    scores_ = np.array([score_cv(torsions_data, dim, lag)\n",
    "                        for dim in dims])\n",
    "    scores = np.mean(scores_, axis=1)\n",
    "    errors = np.std(scores_, axis=1, ddof=1)\n",
    "    color = 'C{}'.format(i)\n",
    "    ax.fill_between(dims, scores - errors, scores + errors, alpha=0.3, facecolor=color)\n",
    "    ax.plot(dims, scores, '--o', color=color, label='lag={:.1f}ns'.format(lag * 0.1))\n",
    "ax.legend()\n",
    "ax.set_xlabel('number of dimensions')\n",
    "ax.set_ylabel('VAMP2 score')\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We observe that, for lag times above 0.5 ns, using more than four dimensions does not increase the score, i.e.,\n",
    "the first four dimensions contain all relevant information of the slow dynamics.\n",
    "\n",
    "Based on this result, we try a TICA projection with lag time 0.5 ns (5 steps).\n",
    "Please note that this is a modeler's choice based on the best heuristic that is currently available to our knowledge. \n",
    "It might be necessary to re-adjust the TICA lag time after the MSM estimation. \n",
    "\n",
    "\n",
    "## Coordinate transform and discretization  [➜ 📓](02-dimension-reduction-and-discretization.ipynb)\n",
    "\n",
    "### TICA\n",
    "\n",
    "The goal of the next step is to find a function that maps the usually high-dimensional input space into some lower-dimensional space that captures the important dynamics.\n",
    "The recommended way of doing so is a time-lagged independent component analysis (TICA), <a id=\"ref-4\" href=\"#cite-tica2\">molgedey-94</a>, <a id=\"ref-5\" href=\"#cite-tica\">perez-hernandez-13</a>.\n",
    "We perform TICA (with kinetic map scaling) using the lag time obtained from the VAMP-2 score.\n",
    "\n",
    "By using the tica() function's default parameters, we will use as many dimensions as necessary in order to preserve $95\\%$ of the kinetic variance.\n",
    "By default, tica() also applies a kinetic map scaling.\n",
    "This scaling ensures that Euclidean distances in the projected space approximate kinetic distances,\n",
    "which is beneficial during the subsequent discretization.\n",
    "\n",
    "Please note that the general `PyEMMA` API is consistant for all estimators.\n",
    "By calling the TICA estimator with the data (`tica = pyemma.coordinates.tica(torsions_data)`),\n",
    "the estimation is done and an estimator instance returned (`tica`);\n",
    "this object contains all the information about the specific transformation.\n",
    "For small systems, we can access the transformed data by calling `tica.get_output()`.\n",
    "For large systems, we recommend to pass the `tica` object itself into the subsequent stages, e.g.,\n",
    "clustering, in order to avoid loading all transformed data into memory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tica = pyemma.coordinates.tica(torsions_data, lag=5)\n",
    "tica_output = tica.get_output()\n",
    "tica_concatenated = np.concatenate(tica_output)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "⚠️ Please note that the above variable `tica_concatenated` is a concatenation of independent trajectories and as such **not suitable** for any kind of analysis which **evaluates transitions**.\n",
    "We use this variable purely for visualization purposes which require concatenated data.\n",
    "Throughout all notebooks, we will use the `_concatenated` postfix to denote data which is not safe to run a TICA or MSM estimation on.\n",
    "\n",
    "We visualize the marginal and joint distributions of our TICA components by simple histograming:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1, 2, figsize=(10, 4))\n",
    "pyemma.plots.plot_feature_histograms(\n",
    "    tica_concatenated,\n",
    "    ax=axes[0],\n",
    "    feature_labels=['IC1', 'IC2', 'IC3', 'IC4'],\n",
    "    ylog=True)\n",
    "pyemma.plots.plot_density(*tica_concatenated[:, :2].T, ax=axes[1], logscale=True)\n",
    "axes[1].set_xlabel('IC 1')\n",
    "axes[1].set_ylabel('IC 2')\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We note that the projection yields defined clusters of high density, which are most likely to be identified as metastable basins. \n",
    "\n",
    "Let’s have a look at one of the trajectories and what it looks like in the space of the first four TICA components.\n",
    "We can see that the TICA components nicely resolve the slow transitions as discrete jumps.\n",
    "Thus, metastability is well-described in this projection."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(4, 1, figsize=(12, 5), sharex=True)\n",
    "x = 0.1 * np.arange(tica_output[0].shape[0])\n",
    "for i, (ax, tic) in enumerate(zip(axes.flat, tica_output[0].T)):\n",
    "    ax.plot(x, tic)\n",
    "    ax.set_ylabel('IC {}'.format(i + 1))\n",
    "axes[-1].set_xlabel('time / ns')\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Discretization\n",
    "\n",
    "The TICA coordinates will now be clustered into a number of discrete states using the $k$-means algorithm.\n",
    "The $k$-means algorithm requires as input the desired number of clusters.\n",
    "The trajectories are automatically assigned to the cluster centers by calling `cluster.dtrajs`. \n",
    "\n",
    "⚠️ It is a priori unclear what the optimal number of cluster centers $k$ is.\n",
    "It largely depends on the distribution of our data and on the number of dimensions we use.\n",
    "\n",
    "In the following, we will estimate unvalidated Markov models using different numbers of cluster centers and use the VAMP-2 score (using cross validation) as a heuristic.\n",
    "This approach requires us to guess the MSM lag time, which we set to the TICA lag time of $5$ steps (or $0.5$ ns).\n",
    "Since the clustering algorithm is stochastic,\n",
    "we conduct multiple rounds of discretization at each number of cluster centers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "n_clustercenters = [5, 10, 30, 75, 200, 450]\n",
    "\n",
    "scores = np.zeros((len(n_clustercenters), 5))\n",
    "for n, k in enumerate(n_clustercenters):\n",
    "    for m in range(5):\n",
    "        with pyemma.util.contexts.settings(show_progress_bars=False):\n",
    "            _cl = pyemma.coordinates.cluster_kmeans(\n",
    "                tica_output, k=k, max_iter=50, stride=50)\n",
    "            _msm = pyemma.msm.estimate_markov_model(_cl.dtrajs, 5)\n",
    "            scores[n, m] = _msm.score_cv(\n",
    "                _cl.dtrajs, n=1, score_method='VAMP2', score_k=min(10, k))\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "lower, upper = pyemma.util.statistics.confidence_interval(scores.T.tolist(), conf=0.9)\n",
    "ax.fill_between(n_clustercenters, lower, upper, alpha=0.3)\n",
    "ax.plot(n_clustercenters, np.mean(scores, axis=1), '-o')\n",
    "ax.semilogx()\n",
    "ax.set_xlabel('number of cluster centers')\n",
    "ax.set_ylabel('VAMP-2 score')\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We find that the VAMP-2 score is already saturated at $75$ states.\n",
    "We will use this number for further analysis.\n",
    "\n",
    "As already stated above, the score has been generated using MSMs that were not validated,\n",
    "meaning that the above plot is really just a heuristic.\n",
    "Besides having an optimal score, we want to obtain a model that describes physically interesting states.\n",
    "Thus, the number of states $k$ is often re-adjusted after model inspection."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "cluster = pyemma.coordinates.cluster_kmeans(\n",
    "    tica_output, k=75, max_iter=50, stride=10, fixed_seed=1)\n",
    "dtrajs_concatenated = np.concatenate(cluster.dtrajs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "⚠️ Please note that $k$-means clustering is a stochastic procedure and,\n",
    "hence, individual runs might yield different results.\n",
    "In order to have reproducible results throughout the remainder of this notebook for each run,\n",
    "we are fixing the random seed (not advised in general). \n",
    "\n",
    "We can check the location of our discrete states by plotting them onto the density of our data in the first two TICA dimensions.\n",
    "The cluster centers are contained in the `cluster` object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(4, 4))\n",
    "pyemma.plots.plot_density(\n",
    "    *tica_concatenated[:, :2].T, ax=ax, cbar=False, alpha=0.3)\n",
    "ax.scatter(*cluster.clustercenters[:, :2].T, s=5, c='C1')\n",
    "ax.set_xlabel('IC 1')\n",
    "ax.set_ylabel('IC 2')\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We note that the states are well distributed in the low-dimensional TICA subspace.\n",
    "\n",
    "## MSM estimation and validation  [➜ 📓](03-msm-estimation-and-validation.ipynb)\n",
    "\n",
    "### Implied timescales\n",
    "\n",
    "The first validation that is usually done when estimating a Markov model is the estimation of implied timescales (ITS) $t_i$.\n",
    "They are computed from the eigenvalues $\\lambda_i$ of the Markov transition matrix by\n",
    "\n",
    "$$ t_i = \\frac{-\\tau}{\\ln\\left|\\lambda_i(\\tau)\\right|} $$ \n",
    "\n",
    "with $\\tau$ being the lag time.\n",
    "The ITS $t_i$ approximates the decorrelation time of the $i$-th process;\n",
    "the latter being independent of the model's hyper-parameter $\\tau$.\n",
    "Thus, we look for ITS convergence and choose $\\tau$ accordingly, i.e.,\n",
    "within a range where the ITS are approximately invariant.\n",
    "Please note that the lag time also represents the time resolution limit of the estimated Markov model <a id=\"ref-6\" href=\"#cite-msm-jhp\">prinz-11</a>, <a id=\"ref-7\" href=\"#cite-swope-its\">swope-04</a>.\n",
    "\n",
    "Now, we calculate the ITS with `pyemma.msm.its()` up to a lagtime of $50$ steps,\n",
    "which is equivalent to $5$ ns for this dataset ($\\Delta t=0.1$ ns).\n",
    "The uncertainty of the implied timescales is quantified based upon Markov models sampled according to a Bayesian scheme.\n",
    "\n",
    "⚠️ If this this is too time-consuming,\n",
    "maximum likelihood MSMs can be used instead by setting the `errors` keyword argument to `None`.\n",
    "\n",
    "Please note that, instead of a single number `lags=50`,\n",
    "an array can be passed to compute the ITS at defined lag times.\n",
    "When we pass an integer $K$ as this value,\n",
    "a set of lag times starting from $\\Delta t$ to $K\\Delta t$ will be generated,\n",
    "using a multiplier of $1.5$ between successive lag times."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "its = pyemma.msm.its(cluster.dtrajs, lags=50, nits=10, errors='bayes')\n",
    "pyemma.plots.plot_implied_timescales(its, units='ns', dt=0.1);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The solid lines correspond to the ITS of maximum likelihood MSMs.\n",
    "The confidence intervals are depicted by the shaded areas;\n",
    "they contain 95% of the samples generated by the Bayesian MSM.\n",
    "The sample means are given by dashed lines.\n",
    "\n",
    "The implied timescales converge quickly.\n",
    "Above $0.5$ ns, the implied timescales of the slowest processes are constant within error.\n",
    "We thus select a lag time of $5$ steps ($0.5$ ns) to build a Markov model.\n",
    "As a quick check, we print the fraction of states and counts that are in the active set. \n",
    "\n",
    "Please note the similarity of the `msm` object to the `tica` object.\n",
    "Both are estimator instances and contain all the relevant information from the estimation and methods for validation and further analysis. \n",
    "\n",
    "In order to keep track of our trajectory time step, a `dt_traj` keyword argument can be passed that contains the trajectory time step unit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "msm = pyemma.msm.bayesian_markov_model(cluster.dtrajs, lag=5, dt_traj='0.1 ns')\n",
    "print('fraction of states used = {:.2f}'.format(msm.active_state_fraction))\n",
    "print('fraction of counts used = {:.2f}'.format(msm.active_count_fraction))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Chapman-Kolmogorov test\n",
    "The model is validated with a Chapman-Kolmogorov test. It compares the right and the left side of the Chapman-Kolmogorov equation\n",
    "\n",
    "$$ \\mathbf{P}(k \\tau) = \\mathbf{P}^k(\\tau) $$\n",
    "\n",
    "with $\\mathbf{P}(\\tau)$ being the transition matrix and lag time $\\tau$.\n",
    "PyEMMA automatically estimates a new MSM transition matrix at lag time $k \\tau$ and propagates the original transition matrix by the $k$-th power.\n",
    "The highest $k$ can be adjusted using the `mlags` keyword argument of `msm.cktest()`.\n",
    "\n",
    "Since we can only inspect the result for a small number of (macro-) states,\n",
    "we use the implied timescales plot as a heuristic to estimate a number of metastable states to test for. \n",
    "We can resolve $4$ slow processes up to lag times of $2.5$ ns.\n",
    "Since the Chapman-Kolmogorov test involves estimations at higher lag times,\n",
    "we will attempt to capture those processes choosing $5$ metastable states."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nstates = 5\n",
    "cktest = msm.cktest(nstates, mlags=6)\n",
    "pyemma.plots.plot_cktest(cktest, dt=0.1, units='ns');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Assuming $5$ metastable states in the data yields a passing Chapman-Kolmogorov test.\n",
    "\n",
    "## MSM spectral analysis [➜ 📓](04-msm-analysis.ipynb)\n",
    "\n",
    "From the MSM object `msm`, various properties can be obtained. We start the spectral analysis by examining the implied timescales."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def its_separation_err(ts, ts_err):\n",
    "    \"\"\"\n",
    "    Error propagation from ITS standard deviation to timescale separation.\n",
    "    \"\"\"\n",
    "    return ts[:-1] / ts[1:] * np.sqrt(\n",
    "        (ts_err[:-1] / ts[:-1])**2 + (ts_err[1:] / ts[1:])**2)\n",
    "\n",
    "\n",
    "nits = 15\n",
    "\n",
    "timescales_mean = msm.sample_mean('timescales', k=nits)\n",
    "timescales_std = msm.sample_std('timescales', k=nits)\n",
    "\n",
    "fig, axes = plt.subplots(1, 2, figsize=(10, 4))\n",
    "\n",
    "axes[0].errorbar(\n",
    "    range(1, nits + 1),\n",
    "    timescales_mean, \n",
    "    yerr=timescales_std, \n",
    "    fmt='.', markersize=10)\n",
    "axes[1].errorbar(\n",
    "    range(1, nits),\n",
    "    timescales_mean[:-1] / timescales_mean[1:], \n",
    "    yerr=its_separation_err(\n",
    "        timescales_mean, \n",
    "        timescales_std), \n",
    "    fmt='.', \n",
    "    markersize=10,\n",
    "    color='C0')\n",
    "\n",
    "for i, ax in enumerate(axes):\n",
    "    ax.set_xticks(range(1, nits + 1))\n",
    "    ax.grid(True, axis='x', linestyle=':')\n",
    "    \n",
    "axes[0].axhline(msm.lag * 0.1, lw=1.5, color='k')\n",
    "axes[0].axhspan(0, msm.lag * 0.1, alpha=0.3, color='k')\n",
    "axes[0].set_xlabel('implied timescale index')\n",
    "axes[0].set_ylabel('implied timescales / ns')\n",
    "axes[1].set_xticks(range(1, nits))\n",
    "axes[1].set_xticklabels(\n",
    "    [\"{:d}/{:d}\".format(k, k + 1) for k in range(1, nits + 2)],\n",
    "    rotation=45)\n",
    "axes[1].set_xlabel('implied timescale indices')\n",
    "axes[1].set_ylabel('timescale separation')\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we see, PyEMMA sorts the implied timescales (and their corresponding eigenfunctions) in descending order.\n",
    "From the timescale separation we can see that a comparably large timescale gap,\n",
    "within the time resolution of the model, is between the 4th and 5th process,\n",
    "suggesting that $5$ metastable states may be a good choice for coarse graining.\n",
    "We discuss this further below.\n",
    "\n",
    "We go on by analyzing the stationary distribution and the free energy computed over the first two TICA coordinates.\n",
    "The stationary distribution, $\\pi$, is stored in `msm.pi` or (as an alias) `msm.stationary_distribution`.\n",
    "We compute the free energy landscape by re-weighting the trajectory frames with stationary probabilities from the MSM (returned by `msm.trajectory_weights()`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1, 2, figsize=(10, 4), sharex=True, sharey=True)\n",
    "pyemma.plots.plot_contour(\n",
    "    *tica_concatenated[:, :2].T,\n",
    "    msm.pi[dtrajs_concatenated],\n",
    "    ax=axes[0],\n",
    "    mask=True,\n",
    "    cbar_label='stationary distribution')\n",
    "pyemma.plots.plot_free_energy(\n",
    "    *tica_concatenated[:, :2].T,\n",
    "    weights=np.concatenate(msm.trajectory_weights()),\n",
    "    ax=axes[1],\n",
    "    legacy=False)\n",
    "for ax in axes.flat:\n",
    "    ax.set_xlabel('IC 1')\n",
    "axes[0].set_ylabel('IC 2')\n",
    "axes[0].set_title('Stationary distribution', fontweight='bold')\n",
    "axes[1].set_title('Reweighted free energy surface', fontweight='bold')\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The eigenvectors corresponding to the slowest processes (largest implied timescales) contain information about what configurational changes are happening on which timescales.\n",
    "We analyze the slowest processes by inspecting the value of the first four eigenfunctions projected on two the first TICA coordinates.\n",
    "As the first right eigenvector corresponds to the stationary process (equilibrium), it is constant at $1$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eigvec = msm.eigenvectors_right()\n",
    "print('The first eigenvector is one: {} (min={}, max={})'.format(\n",
    "    np.allclose(eigvec[:, 0], 1, atol=1e-15), eigvec[:, 0].min(), eigvec[:, 0].max()))\n",
    "\n",
    "fig, axes = plt.subplots(1, 4, figsize=(15, 3), sharex=True, sharey=True)\n",
    "for i, ax in enumerate(axes.flat):\n",
    "    pyemma.plots.plot_contour(\n",
    "        *tica_concatenated[:, :2].T,\n",
    "        eigvec[dtrajs_concatenated, i + 1],\n",
    "        ax=ax,\n",
    "        cmap='PiYG',\n",
    "        cbar_label='{}. right eigenvector'.format(i + 2),\n",
    "        mask=True)\n",
    "    ax.set_xlabel('IC 1')\n",
    "axes[0].set_ylabel('IC 2')\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The eigenvectors of the MSM contain information about the conformational changes are happening,\n",
    "governed by the correponding implied timescales.\n",
    "Specifically, the minimum and maximum-valued components of an eigenvector indicate the states between a process shifts probability.\n",
    "The relaxation timescale of this exchange process is exactly the implied timescale.\n",
    "\n",
    "Since the eigenvectors were internally sorted according to their eigenvalue,\n",
    "the above visualizations depict the four slowest processes of the implied timescale plot.\n",
    "We see that, indeed, the slowest processes occur inbetween the dense clusters in the TICA projection.\n",
    "\n",
    "## PCCA & TPT [➜ 📓](05-pcca-tpt.ipynb)\n",
    "\n",
    "### Perron cluster cluster analysis\n",
    "\n",
    "**Note:** We will assign the integer numbers $1 \\dots $ `nstates` to PCCA++ metastable states. As PyEMMA is written in Python, it internally indexes states starting from $0$. In consequence, numbers in the code cells differ by $-1$ from the plot labels and markdown text. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "msm.pcca(nstates)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The PCCA++ algorithm computes so called memberships, i.e., the probability of each microstate to belong to a given macrostate.\n",
    "In other words, PCCA++ does a fuzzy assignment of the microstates to macrostates which is encoded in the memberships.\n",
    "We can visualize the $5$ membership distributions over the first $2$ TICA dimensions as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1, 5, figsize=(15, 3), sharex=True, sharey=True)\n",
    "for i, ax in enumerate(axes.flat):\n",
    "    pyemma.plots.plot_contour(\n",
    "        *tica_concatenated[:, :2].T,\n",
    "        msm.metastable_distributions[i][dtrajs_concatenated],\n",
    "        ax=ax,\n",
    "        cmap='afmhot_r', \n",
    "        mask=True,\n",
    "        cbar_label='metastable distribution {}'.format(i + 1))\n",
    "    ax.set_xlabel('IC 1')\n",
    "axes[0].set_ylabel('IC 2')\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we see, the membership probabilities roughly match to the basins of the free energy landscape presented above.\n",
    "\n",
    "In some cases, it might be useful to convert these distributions into crisp assignments.\n",
    "This can be computed by taking the argmax of each microstate's memberships to the macrostates.\n",
    "They are contained in `msm.metastable_assignments`. Let's see what this looks like in the first two TICA projections."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "metastable_traj = msm.metastable_assignments[dtrajs_concatenated]\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(5, 4))\n",
    "_, _, misc = pyemma.plots.plot_state_map(\n",
    "    *tica_concatenated[:, :2].T, metastable_traj, ax=ax)\n",
    "ax.set_xlabel('IC 1')\n",
    "ax.set_ylabel('IC 2')\n",
    "misc['cbar'].set_ticklabels([r'$\\mathcal{S}_%d$' % (i + 1)\n",
    "                             for i in range(nstates)])\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As expected, PCCA++ has nicely separated our state space within the first two TICA components.\n",
    "\n",
    "At this point, we usually want to investigate which molecular structures the identified metastable structures correspond to.\n",
    "We generate a number of representative sample structures for each macrostate and store them in a trajectory file for visual inspection.\n",
    "The following cell writes trajectory files to hard disc. They can be loaded and analyzed with external software packages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pcca_samples = msm.sample_by_distributions(msm.metastable_distributions, 10)\n",
    "torsions_source = pyemma.coordinates.source(files, features=torsions_feat)\n",
    "pyemma.coordinates.save_trajs(\n",
    "    torsions_source,\n",
    "    pcca_samples,\n",
    "    outfiles=['./data/pcca{}_10samples.pdb'.format(n + 1)\n",
    "              for n in range(msm.n_metastable)])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Alternatively, one can visualize the structures in this notebook with NGLView.\n",
    "For this, we need to provide a custom function that defines the representations of our molecule:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def visualize_metastable(samples, cmap, selection='not element H'):\n",
    "    \"\"\" visualize metastable states\n",
    "    Parameters\n",
    "    ----------\n",
    "    samples: list of mdtraj.Trajectory objects\n",
    "        each element contains all samples for one metastable state.\n",
    "    cmap: matplotlib.colors.ListedColormap\n",
    "        color map used to visualize metastable states before.\n",
    "    selection: str\n",
    "        which part of the molecule to selection for visualization. For details have a look here:\n",
    "        http://mdtraj.org/latest/examples/atom-selection.html#Atom-Selection-Language\n",
    "    \"\"\"\n",
    "    import nglview\n",
    "    from matplotlib.colors import to_hex\n",
    "\n",
    "    widget = nglview.NGLWidget()\n",
    "    widget.clear_representations()\n",
    "    ref = samples[0]\n",
    "    for i, s in enumerate(samples):\n",
    "        s = s.superpose(ref, atom_indices=s.top.select('resid 2 3 and mass > 2'))\n",
    "        s = s.atom_slice(s.top.select(selection))\n",
    "        comp = widget.add_trajectory(s)\n",
    "        comp.add_licorice()\n",
    "\n",
    "    # this has to be done in a separate loop for whatever reason...\n",
    "    x = np.linspace(0, 1, num=len(samples))\n",
    "    for i, x_ in enumerate(x):\n",
    "        c = to_hex(cmap(x_))\n",
    "        widget.update_licorice(color=c, component=i, repr_index=i)\n",
    "        widget.remove_cartoon(component=i)\n",
    "    return widget"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Similar to saving the trajectories to disc, we now create a list with `mdtraj.Trajectory` objects that contain samples of our metastable structures. They are visualized with NGLView as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "my_samples = [pyemma.coordinates.save_traj(files, idist, outfile=None, top=pdb)\n",
    "              for idist in msm.sample_by_distributions(msm.metastable_distributions, 50)]\n",
    "\n",
    "cmap = mpl.cm.get_cmap('viridis', nstates)\n",
    "visualize_metastable(my_samples, cmap)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This coarse-grained representation of the dynamics is more directly amenable to human interpretation.\n",
    "Nevertheless, as for the conventional MSM, we can still compute several interesting properties.\n",
    "We start with the stationary distribution which encodes the free energy of the states.\n",
    "This can be achieved by summing all the  contributions to a coarse-grained state $\\mathcal{S}_i$:\n",
    "\n",
    "$$ G_i = - \\textrm{k}_\\textrm{B} T \\ln \\sum_{j\\in \\mathcal{S}_i} \\pi_j $$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('state\\tπ\\t\\tG/kT')\n",
    "for i, s in enumerate(msm.metastable_sets):\n",
    "    p = msm.pi[s].sum()\n",
    "    print('{}\\t{:f}\\t{:f}'.format(i + 1, p, -np.log(p)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Knowing PCCA++ metastable states, we can also extract mean first passage times (MFPTs) between them:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from itertools import product\n",
    "\n",
    "mfpt = np.zeros((nstates, nstates))\n",
    "for i, j in product(range(nstates), repeat=2):\n",
    "    mfpt[i, j] = msm.mfpt(\n",
    "        msm.metastable_sets[i],\n",
    "        msm.metastable_sets[j])\n",
    "\n",
    "from pandas import DataFrame\n",
    "print('MFPT / ns:')\n",
    "DataFrame(np.round(mfpt, decimals=2), index=range(1, nstates + 1), columns=range(1, nstates + 1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As it will become clear in further sections, metastable state $1$ could be distinguished from the other states by experiment.\n",
    "We can extract the MFPT into (out of) this state from (into) any other state using the Bayesian sample as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "A = msm.metastable_sets[0]\n",
    "B = np.concatenate(msm.metastable_sets[1:])\n",
    "print('MFPT 1 -> other: ({:6.1f} ± {:5.1f}) ns'.format(\n",
    "    msm.sample_mean('mfpt', A, B), msm.sample_std('mfpt', A, B)))\n",
    "print('MFPT other -> 1: ({:.1f} ± {:5.1f}) ns'.format(\n",
    "    msm.sample_mean('mfpt', B, A), msm.sample_std('mfpt', B, A)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The MFPT from state $1$ to any other state is very short compared to the other direction, i.e., this state has a comparably short lifetime. \n",
    "\n",
    "### Transition path theory\n",
    "\n",
    "The flux between metastable states can be computed and coarse-grained as follows.\n",
    "As an example, we compute the flux between metastable states $2$ and $4$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "start, final = 1, 3\n",
    "A = msm.metastable_sets[start]\n",
    "B = msm.metastable_sets[final]\n",
    "flux = pyemma.msm.tpt(msm, A, B)\n",
    "\n",
    "cg, cgflux = flux.coarse_grain(msm.metastable_sets)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The committor as projected onto the first two TICA dimensions can be displayed with a filled contour plot:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(5, 4))\n",
    "\n",
    "pyemma.plots.plot_contour(\n",
    "    *tica_concatenated[:, :2].T,\n",
    "    flux.committor[dtrajs_concatenated],\n",
    "    cmap='brg',\n",
    "    ax=ax,\n",
    "    mask=True,\n",
    "    cbar_label=r'committor $\\mathcal{S}_%d \\to \\mathcal{S}_%d$' % (\n",
    "        start + 1, final + 1))\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We find that the committor is constant within the metastable sets defined above.\n",
    "Transition regions can be identified by committor values $\\approx 0.5$.\n",
    "\n",
    "## Computing experimental observables [➜ 📓](06-expectations-and-observables.ipynb)\n",
    "\n",
    "Having thoroughly constructed, validated, and analyzed our MSM,\n",
    "we may want to take the next step and compare our model to experimental data.\n",
    "PyEMMA enables computation of stationary as well as dynamic experimental observables;\n",
    "below we give give some examples of this.\n",
    "We will make use of some external library functionality provided by MDTraj <a id=\"ref-9\" href=\"#cite-mdtraj\">mcgibbon-15</a>.\n",
    "\n",
    "[Notebook 06](06-expectations-and-observables.ipynb) includes a brief summary of the theory used in below.\n",
    "More in-depth descriptions of the theory and their applications to various data can be found in the following references:\n",
    "- <a id=\"ref-10\" href=\"#cite-simon-amm\">olsson-17</a>\n",
    "- <a id=\"ref-11\" href=\"#cite-noe-fingerprints\">noe-11</a>\n",
    "- <a id=\"ref-12\" href=\"#cite-simon-mech-mod-nmr\">olsson-16</a>\n",
    "- <a id=\"ref-13\" href=\"#cite-noe-dy-neut-scatt\">lindner-13</a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from mdtraj import shrake_rupley, compute_rg\n",
    "\n",
    "#We compute a maximum likelihood MSM for comparison\n",
    "mlmsm = pyemma.msm.estimate_markov_model(cluster.dtrajs, lag=5, dt_traj='0.1 ns')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The way PyEMMA computes all experimental observables relies on us having a value of the observable for each of the Markov states in our MSM or frames in our trajectories.\n",
    "Computing these quantities depends on the experiment in question and may require specific domain knowledege or involve time-consuming computations.\n",
    "Here we will pre-compute these experimental observables in the following cells.\n",
    "\n",
    "First, we sample $20$ representative configurations from each Markov state.\n",
    "Note that $20$ representative configurations is not a universal number:\n",
    "a specific observable of interest may require substantially more representative configurations to converge."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "markov_samples = [smpl for smpl in msm.sample_by_state(20)]\n",
    "\n",
    "reader = pyemma.coordinates.source(files, top=pdb)\n",
    "samples = [pyemma.coordinates.save_traj(reader, smpl, outfile=None, top=pdb)\n",
    "            for smpl in markov_samples]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we back-compute two experimental observables, the solvent accessible surface area (SASA) and the radius of gyration,\n",
    "of each of the sampled configurations using MDTraj and average each of these observables over the Markov states.\n",
    "Each of these averages corresponds to our Markov state proxy of the macroscopic experimental observable we want to predict."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute solvent accessible surface area for all samples\n",
    "markov_sasa_all = [shrake_rupley(sample, mode='residue')\n",
    "                   for sample in samples]\n",
    "\n",
    "# Compute radius of gyration for all samples\n",
    "markov_rg_all = [compute_rg(sample) for sample in samples]\n",
    "\n",
    "# Average over Markov states for both observables.\n",
    "markov_average_trp_sasa = np.array(markov_sasa_all).mean(axis=1)[:, 0]\n",
    "markov_average_rg = np.array(markov_rg_all).mean(axis=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Radius of gyration\n",
    "The radius of gyration $r_g(x)$ is a measure of the over-all size of a molecule in a certain configuration $x$.\n",
    "It is a quantity often extracted from light-scattering experiments.\n",
    "In the context of proteins and nucleic acids these experiments often happen on bulk samples and the observable is therefore stationary and averaged by the Boltzmann distribution:\n",
    "\n",
    "$$ R_{g, \\mathrm{obs}} = \\mathcal{Z}^{-1}\\int_{\\Omega} \\mathrm{d}x\\, r_g(x) \\exp\\left(-\\frac{E(x)}{kT}\\right). $$\n",
    "\n",
    "This value is also called the _expectation value of $r^2_g(x)$ with respect to the Boltzmann distribution_.\n",
    "Since we have access to the stationary distribution from the Markov model,\n",
    "we can approximate this continuous integral by the sum:\n",
    "\n",
    "$$ R_{g, \\mathrm{obs}} \\approx \\sum_i \\pi_i \\mathbf{r_g}_i $$\n",
    "with\n",
    "$$ \\mathbf{r_g}_i  = \\frac{1}{\\pi_i\\mathcal{Z}} \\int_{x\\in S_i} \\mathrm{d}x\\, r_g(x)\\exp\\left(-\\frac{E(x)}{kT}\\right). $$\n",
    "\n",
    "The last expression constitutes the average of the experimental observable $r_g(x)$ confined to each of our Markov states individually. \n",
    "\n",
    "Recall that we pre-computed the vector $\\mathbf{r_g}=\\{\\mathbf{r_g}_i\\}$ as `markov_average_rg` above.\n",
    "Using this vector and our estimated Markov model,\n",
    "we can compute the expectation value using the `expectation()` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('The average radius of gyration of penta-peptide is'\n",
    "      ' {:.3f} nm'.format(msm.expectation(markov_average_rg)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since we have a Bayesian MSM estimated,\n",
    "we can also compute the uncertainty in our prediction of the observable as standard deviations or confidence intervals.\n",
    "To achieve this, we use the methods `sample_std` and `sample_conf`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('The standard deviation of our prediction of the average radius of gyration'\n",
    "      ' of pentapeptide is {:.9f} nm'.format(\n",
    "          msm.sample_std('expectation', markov_average_rg)))\n",
    "print('The {:d}% CI of our prediction of the average radius of gyration of'\n",
    "      ' pentapeptide have the bounds ({:.5f}, {:.5f})'.format(\n",
    "          int(msm.conf * 100), *msm.sample_conf('expectation', markov_average_rg)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So, our model is very confident in the prediction of the radius of gyration.\n",
    "However, this does not guarantee it to be accurate &mdash; that is, agree with experimental measurements.\n",
    "If we lack quantitative agreement with experiments,\n",
    "we can estimate MSMs which optimally balance experimental data and simulation data using the Augmented Markov model (AMM) procedure.\n",
    "A dedicated AMM tutorial can be [found here](http://www.emma-project.org/latest/generated/augmented_markov_model_walkthrough.html).\n",
    "\n",
    "### Trp-flourescene auto-correlation\n",
    "\n",
    "Fluctuations in tryptophan flourescence can be measured using spectroscopic techniques.\n",
    "These fluctuations depend, among other things,\n",
    "on the solvent accessible surface area (SASA) of tryptophan residues.\n",
    "Above, we pre-computed the SASA using the Shrake-Rupley algorithm which we show here projected onto the first two TICA dimensions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(5, 4))\n",
    "pyemma.plots.plot_contour(\n",
    "    *tica_concatenated[:, :2].T,\n",
    "    markov_average_trp_sasa[dtrajs_concatenated],\n",
    "    ax=ax,\n",
    "    mask=True,\n",
    "    cbar_label=r'Trp-1 SASA / nm$^2$')\n",
    "ax.set_xlabel('IC 1')\n",
    "ax.set_ylabel('IC 2')\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As for the stationary expectation (ensemble averages) considered above,\n",
    "we can use our pre-computed SASA vector to compute the auto-correlation function of trypotophan flourescene using the MSM `correlation()` method (See [Notebook 06](06-expectations-and-observables.ipynb#Dynamic/kinetic-experimental-observables) for details):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eq_time_ml, eq_acf_ml = mlmsm.correlation(markov_average_trp_sasa, maxtime=150)\n",
    "\n",
    "eq_time_bayes, eq_acf_bayes = msm.sample_mean(\n",
    "    'correlation',\n",
    "    np.array(markov_average_trp_sasa),\n",
    "    maxtime=150)\n",
    "\n",
    "eq_acf_bayes_ci_l, eq_acf_bayes_ci_u = msm.sample_conf(\n",
    "    'correlation',\n",
    "    np.array(markov_average_trp_sasa),\n",
    "    maxtime=150)\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "ax.plot(eq_time_ml, eq_acf_ml, '-o', color='C1', label='ML MSM')\n",
    "ax.plot(\n",
    "    eq_time_bayes,\n",
    "    eq_acf_bayes,\n",
    "    '--x',\n",
    "    color='C0',\n",
    "    label='Bayes sample mean')\n",
    "ax.fill_between(\n",
    "    eq_time_bayes,\n",
    "    eq_acf_bayes_ci_l[1],\n",
    "    eq_acf_bayes_ci_u[1],\n",
    "    facecolor='C0',\n",
    "    alpha=0.3)\n",
    "ax.semilogx()\n",
    "\n",
    "ax.set_xlim((eq_time_ml[1], eq_time_ml[-1]))\n",
    "ax.set_xlabel(r'time / $\\mathrm{ns}$')\n",
    "ax.set_ylabel(r'Trp-1 SASA ACF / $\\mathrm{nm}^4$')\n",
    "\n",
    "ax.legend()\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note the scale on the $y$-axis: this amplitude is likely too small to be experimentally measurable considering experimental uncertainty.\n",
    "\n",
    "However, using more advanced experimental setups such as stopped flow, T-jump, P-jump, and others,\n",
    "we can prepare our ensemble in a non-equilibrium initial condition. (See [Notebook 06](06-expectations-and-observables.ipynb#Dynamic/kinetic-experimental-observables) for details)\n",
    "\n",
    "Let us say we can experimentally prepare a sample to only be in metastable state $\\mathcal{S}_1$.\n",
    "In this case, the initial condition will be given by the metastable distribution of metastable state $\\mathcal{S}_1$, $p_0$.\n",
    "With PyEMMA, we can simulate the relaxation from such a non-equilibrium initial condition $p_0$ back to equilibrium using the `relaxation()` method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eq_time_ml, eq_relax_ml = mlmsm.relaxation(\n",
    "    msm.metastable_distributions[0],\n",
    "    markov_average_trp_sasa,\n",
    "    maxtime=150)\n",
    "\n",
    "eq_time_bayes, eq_relax_bayes = msm.sample_mean(\n",
    "    'relaxation',\n",
    "    msm.metastable_distributions[0],\n",
    "    np.array(markov_average_trp_sasa),\n",
    "    maxtime=150)\n",
    "\n",
    "eq_relax_bayes_CI_l, eq_relax_bayes_CI_u = msm.sample_conf(\n",
    "    'relaxation', \n",
    "    msm.metastable_distributions[0],\n",
    "    np.array(markov_average_trp_sasa),\n",
    "    maxtime=150)\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "ax.plot(eq_time_ml, eq_relax_ml, '-o', color='C1', label='ML MSM')\n",
    "ax.plot(\n",
    "    eq_time_bayes,\n",
    "    eq_relax_bayes,\n",
    "    '--x',\n",
    "    color='C0',\n",
    "    label='Bayes sample mean')\n",
    "ax.fill_between(\n",
    "    eq_time_bayes,\n",
    "    eq_relax_bayes_CI_l[1],\n",
    "    eq_relax_bayes_CI_u[1],\n",
    "    facecolor='C0',\n",
    "    alpha=0.3)\n",
    "ax.semilogx()\n",
    "\n",
    "ax.set_xlim((eq_time_ml[1], eq_time_ml[-1]))\n",
    "ax.set_xlabel(r'time / $\\mathrm{ns}$')\n",
    "ax.set_ylabel(r'Average Trp-1 SASA / $\\mathrm{nm}^2$')\n",
    "\n",
    "ax.legend()\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This signal is much stronger than what we got from the auto-correlation at equilibrium above!\n",
    "\n",
    "If we compute the average observable of each of our metastable states and compare them to the ensemble average,\n",
    "we will get an idea of which initial distribution would give us the strongest signal.\n",
    "We compare these values and report their absolute differences below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "state2ensemble = np.abs(msm.expectation(markov_average_trp_sasa) - \n",
    "                        msm.metastable_distributions.dot(np.array(markov_average_trp_sasa)))\n",
    "DataFrame(np.round(state2ensemble, 3), index=range(1, nstates + 1), columns=[''])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note how preparing our system in metastable state $\\mathcal{S}_1$ will give us the strongest signal signified by the largest absolute difference between the observable in one metastable state to the global ensemble average.\n",
    "However, note that there may be other initial conditions which may give us even stronger signals. \n",
    "\n",
    "Using methods such as the ones shown above,\n",
    "we can imagine how to use MSMs to design experiments to aid validation and testing our model.\n",
    "\n",
    "## Hidden Markov models [➜ 📓](07-hidden-markov-state-models.ipynb)\n",
    "\n",
    "Another way of approaching metastable dynamics is with hidden Markov models (HMMs) <a id=\"ref-14\" href=\"#cite-noe-proj-hid-msm\">noe-15</a>.\n",
    "HMMs model the dynamics between so-called hidden states which we initiate from the metastable states found by PCCA++.\n",
    "The estimation is less prone to discretization errors as we do not assume Markovianity in the space of our cluster centers.\n",
    "It further provides a natural coarse graining into a given number of hidden states,\n",
    "thus yielding a means to generate a closed-form Markov model of the metastable dynamics.\n",
    "\n",
    "In [Notebook 07 📓](07-hidden-markov-state-models.ipynb), after introducing HMMs,\n",
    "you have the chance to estimate an HMM for the pentapeptide in the exercises section."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Assembling manuscript figures\n",
    "\n",
    "In the following, we use the results from this notebook to compile figures for the manuscript in this repository.\n",
    "We start by setting new defaults which allow us to use the resulting figures without scaling in the manuscript."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib.gridspec import GridSpec\n",
    "from matplotlib.ticker import LogLocator\n",
    "from matplotlib.cm import get_cmap\n",
    "from pyemma.plots.markovtests import _add_ck_subplot\n",
    "\n",
    "mpl.rcParams['axes.titlesize'] = 6\n",
    "mpl.rcParams['axes.labelsize'] = 6\n",
    "mpl.rcParams['legend.fontsize'] = 5\n",
    "mpl.rcParams['xtick.labelsize'] = 5\n",
    "mpl.rcParams['ytick.labelsize'] = 5\n",
    "mpl.rcParams['xtick.minor.pad'] = 2\n",
    "mpl.rcParams['xtick.major.pad'] = 3\n",
    "mpl.rcParams['ytick.minor.pad'] = 2\n",
    "mpl.rcParams['ytick.major.pad'] = 3\n",
    "mpl.rcParams['axes.labelpad'] = 1\n",
    "mpl.rcParams['lines.markersize'] = 4"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is Figure 3 (a,b,c,d) which sketches the system and coordinates part:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(3.47, 4.65))\n",
    "gw = int(np.floor(0.5 + 1000 * fig.get_figwidth()))\n",
    "gh = int(np.floor(0.5 + 1000 * fig.get_figheight()))\n",
    "gs = plt.GridSpec(gh, gw)\n",
    "gs.update(hspace=0.0, wspace=0.0, left=0.0, right=1.0, bottom=0.0, top=1.0)\n",
    "\n",
    "ax_box = fig.add_subplot(gs[:, :])\n",
    "ax_box.set_axis_off()\n",
    "ax_box.text(0.00, 0.95, '(a)', size=10, zorder=1)\n",
    "ax_box.text(0.00, 0.58, '(b)', size=10)\n",
    "ax_box.text(0.55, 0.58, '(c)', size=10)\n",
    "ax_box.text(0.00, 0.22, '(d)', size=10)\n",
    "\n",
    "ax_mol = fig.add_subplot(gs[:1600, -2820:-400])\n",
    "ax_mol.set_axis_off()\n",
    "ax_mol.imshow(plt.imread('static/pentapeptide-structure.png'))\n",
    "\n",
    "ax_feat = fig.add_subplot(gs[2000:3150, 400:1800])\n",
    "ax_feat.bar(\n",
    "    vamp_bars_plot['labels'],\n",
    "    vamp_bars_plot['scores'],\n",
    "    yerr=vamp_bars_plot['errors'],\n",
    "    color=['C0', 'C1', 'C2'])\n",
    "ax_feat.set_ylabel('VAMP2 score')\n",
    "ax_feat.set_title(r'lag time $\\tau$ = {:.1f} ns'.format(vamp_bars_plot['lag'] * 0.1))\n",
    "ax_feat.set_ylim(2.75, 3.65)\n",
    "ax_feat.tick_params(axis='x', labelrotation=20)\n",
    "\n",
    "ax_sample_free_energy = fig.add_subplot(gs[2000:3150, 2200:3350])\n",
    "_, _, misc = pyemma.plots.plot_free_energy(\n",
    "    *tica_concatenated.T[:2],\n",
    "    ax=ax_sample_free_energy,\n",
    "    cax=fig.add_subplot(gs[1900:1950, 2200:3350]),\n",
    "    cbar_orientation='horizontal',\n",
    "    legacy=False)\n",
    "misc['cbar'].set_label('sample free energy / kT')\n",
    "misc['cbar'].set_ticks(np.arange(9))\n",
    "misc['cbar'].ax.xaxis.set_ticks_position('top')\n",
    "misc['cbar'].ax.xaxis.set_label_position('top')\n",
    "ax_sample_free_energy.set_xlabel('IC 1')\n",
    "ax_sample_free_energy.set_ylabel('IC 2')\n",
    "\n",
    "x = 0.1 * np.arange(tica_output[0].shape[0])\n",
    "ax_tic1 = fig.add_subplot(gs[3650:4000, 400:3350])\n",
    "ax_tic2 = fig.add_subplot(gs[4000:4350, 400:3350])\n",
    "\n",
    "ax_tic1.plot(x, tica_output[0][:, 0], linewidth=0.25)\n",
    "ax_tic2.plot(x, tica_output[0][:, 1], linewidth=0.25)\n",
    "ax_tic1.set_ylabel('IC 1')\n",
    "ax_tic2.set_ylabel('IC 2')\n",
    "ax_tic2.set_xlabel('time / ns')\n",
    "\n",
    "fig.savefig('data/figure_3.pdf', dpi=300)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next is Figure 4 (a,b) which shows estimation and validation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(3.47, 2.60))\n",
    "gw = int(np.floor(0.5 + 1000 * fig.get_figwidth()))\n",
    "gh = int(np.floor(0.5 + 1000 * fig.get_figheight()))\n",
    "gs = plt.GridSpec(gh, gw)\n",
    "gs.update(hspace=0.0, wspace=0.0, left=0.0, right=1.0, bottom=0.0, top=1.0)\n",
    "\n",
    "ax_box = fig.add_subplot(gs[:, :])\n",
    "ax_box.set_axis_off()\n",
    "ax_box.text(0.00, 0.95, '(a)', size=10)\n",
    "ax_box.text(0.00, 0.36, '(b)', size=10)\n",
    "\n",
    "ax_its = fig.add_subplot(gs[50:1300, 400:3350])\n",
    "pyemma.plots.plot_implied_timescales(its, units='ns', dt=0.1, ax=ax_its, nits=4, ylog=True)\n",
    "ax_its.set_ylim(1, ax_its.get_ylim()[1])\n",
    "ax_its.set_xlabel(r'lag time $\\tau$ / ns')\n",
    "\n",
    "ax_ck = [\n",
    "    fig.add_subplot(gs[1730:2300, 400:970]),\n",
    "    fig.add_subplot(gs[1730:2300, 995:1565]),\n",
    "    fig.add_subplot(gs[1730:2300, 1590:2160]),\n",
    "    fig.add_subplot(gs[1730:2300, 2185:2755]),\n",
    "    fig.add_subplot(gs[1730:2300, 2780:3350])]\n",
    "\n",
    "for k, ax in enumerate(ax_ck):\n",
    "    lest, lpred = _add_ck_subplot(\n",
    "        cktest, ax, k, k, ipos=cktest.nsets - 1, dt=0.1, units='ns', linewidth=0.7)\n",
    "    if k > 0:\n",
    "        ax.set_yticks([])\n",
    "predlabel = 'predict ({:3.1f}%)'.format(100.0 * cktest.conf)\n",
    "estlabel = 'estimate'\n",
    "ax_ck[-1].legend(\n",
    "    (lest[0], lpred[0]),\n",
    "    (estlabel, predlabel),\n",
    "    frameon=False,\n",
    "    loc='upper left',\n",
    "    bbox_to_anchor=(-0.4, 1.45))\n",
    "\n",
    "fig.savefig('data/figure_4.pdf', dpi=300)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Figure 5 (a,b,c,d) highlights the basic analysis part using map plots:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(3.47, 3.95))\n",
    "gw = int(np.floor(0.5 + 1000 * fig.get_figwidth()))\n",
    "gh = int(np.floor(0.5 + 1000 * fig.get_figheight()))\n",
    "gs = plt.GridSpec(gh, gw)\n",
    "gs.update(hspace=0.0, wspace=0.0, left=0.0, right=1.0, bottom=0.0, top=1.0)\n",
    "\n",
    "ax_box = fig.add_subplot(gs[:, :])\n",
    "ax_box.set_axis_off()\n",
    "ax_box.text(0.03, 0.97, '(a)', size=10)\n",
    "ax_box.text(0.53, 0.97, '(b)', size=10)\n",
    "ax_box.text(0.03, 0.50, '(c)', size=10)\n",
    "ax_box.text(0.53, 0.50, '(d)', size=10)\n",
    "\n",
    "ax_fe = fig.add_subplot(gs[400:1750, 400:1750])\n",
    "_, _, misc = pyemma.plots.plot_free_energy(\n",
    "    *tica_concatenated.T[:2],\n",
    "    weights=np.concatenate(msm.trajectory_weights()),\n",
    "    ax=ax_fe,\n",
    "    cax=fig.add_subplot(gs[300:350, 400:1750]),\n",
    "    cbar_orientation='horizontal',\n",
    "    legacy=False)\n",
    "misc['cbar'].set_ticks(np.linspace(0, 8, 5))\n",
    "misc['cbar'].ax.xaxis.set_ticks_position('top')\n",
    "misc['cbar'].ax.xaxis.set_label_position('top')\n",
    "misc['cbar'].set_label(r'free energy / $\\mathrm{k}_\\mathrm{B}T$')\n",
    "ax_fe.set_ylabel('IC 2')\n",
    "ax_fe.set_xticklabels([])\n",
    "\n",
    "ax_state = fig.add_subplot(gs[400:1750, 2000:3350])\n",
    "_, _, misc = pyemma.plots.plot_state_map(\n",
    "    *tica_concatenated.T[:2],\n",
    "    metastable_traj,\n",
    "    ax=ax_state,\n",
    "    cax=fig.add_subplot(gs[300:350, 2000:3350]),\n",
    "    cbar_label='metastable state',\n",
    "    cbar_orientation='horizontal')\n",
    "misc['cbar'].ax.xaxis.set_ticks_position('top')\n",
    "misc['cbar'].ax.xaxis.set_label_position('top')\n",
    "misc['cbar'].set_ticklabels([r'$\\mathcal{S}_%d$' % (i + 1)\n",
    "                             for i in range(nstates)])\n",
    "ax_state.set_xticklabels([])\n",
    "ax_state.set_yticklabels([])\n",
    "ax_state.text(0.70, 6.30, '$\\mathcal{S}_1$', size=10)\n",
    "ax_state.text(3.00, 4.00, '$\\mathcal{S}_2$', size=10)\n",
    "ax_state.text(2.50, 2.00, '$\\mathcal{S}_3$', size=10)\n",
    "ax_state.text(5.20, -0.50, '$\\mathcal{S}_4$', size=10)\n",
    "ax_state.text(-0.20, 0.00, '$\\mathcal{S}_5$', size=10)\n",
    "\n",
    "evec_idx = 1\n",
    "ax_eig = fig.add_subplot(gs[2300:3650, 400:1750])\n",
    "_, _, misc = pyemma.plots.plot_contour(\n",
    "    *tica_concatenated.T[:2],\n",
    "    eigvec[dtrajs_concatenated, evec_idx],\n",
    "    cmap='PiYG',\n",
    "    ax=ax_eig,\n",
    "    mask=True,\n",
    "    cax=fig.add_subplot(gs[2200:2250, 400:1750]),\n",
    "    cbar_label='{}. right eigenvector'.format(evec_idx + 1),\n",
    "    cbar_orientation='horizontal')\n",
    "misc['cbar'].set_ticks(np.linspace(*misc['cbar'].get_clim(), 3))\n",
    "misc['cbar'].ax.xaxis.set_ticks_position('top')\n",
    "misc['cbar'].ax.xaxis.set_label_position('top')\n",
    "ax_eig.set_xlabel('IC 1')\n",
    "ax_eig.set_ylabel('IC 2')\n",
    "\n",
    "ax_flux = fig.add_subplot(gs[2300:3650, 2000:3350])\n",
    "_, _, misc = pyemma.plots.plot_contour(\n",
    "    *tica_concatenated.T[:2],\n",
    "    flux.committor[dtrajs_concatenated],\n",
    "    cmap='brg',\n",
    "    ax=ax_flux,\n",
    "    mask=True,\n",
    "    cax=fig.add_subplot(gs[2200:2250, 2000:3350]),\n",
    "    cbar_label=r'committor $\\mathcal{S}_%d \\to \\mathcal{S}_%d$' % (\n",
    "        start + 1, final + 1),\n",
    "    cbar_orientation='horizontal')\n",
    "misc['cbar'].set_ticks(np.linspace(0, 1, 3))\n",
    "misc['cbar'].set_ticklabels(['start', 'transition state', 'final'])\n",
    "misc['cbar'].ax.xaxis.set_ticks_position('top')\n",
    "misc['cbar'].ax.xaxis.set_label_position('top')\n",
    "ax_flux.set_xlabel('IC 1')\n",
    "ax_flux.set_yticklabels([])\n",
    "\n",
    "fig.savefig('data/figure_5.pdf', dpi=300)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Figure 6 visualizes the $\\mathcal{S}_2\\to\\mathcal{S}_4$ committor:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(3.47, 1.804))\n",
    "gw = int(np.floor(0.5 + 1000 * fig.get_figwidth()))\n",
    "gh = int(np.floor(0.5 + 1000 * fig.get_figheight()))\n",
    "gs = plt.GridSpec(gh, gw)\n",
    "gs.update(hspace=0.0, wspace=0.0, left=0.0, right=1.0, bottom=0.0, top=1.0)\n",
    "\n",
    "ax_box = fig.add_subplot(gs[:, :])\n",
    "ax_box.set_axis_off()\n",
    "\n",
    "posx = cgflux.forward_committor\n",
    "posy = 0.1 * np.array([4.5, 2.3, 2.3, 7, 4.5])\n",
    "pos = np.vstack([posx, posy]).T\n",
    "\n",
    "sizes = [[385, 432], [388, 526], [347, 500], [367, 348], [260, 374]]\n",
    "anchors = [[1120, 950], [220, 500], [1630, 830], [2800, 600], [1600, 50]]\n",
    "for i, ((w, h), (x, y)) in enumerate(zip(sizes, anchors)):\n",
    "    ax_ = fig.add_subplot(gs[y:y+h, x:x+w])\n",
    "    ax_.set_axis_off()\n",
    "    ax_.text(0, 500, r'$\\mathcal{S}_' + '{}$'.format(i + 1))\n",
    "    ax_.imshow(plt.imread('static/hmm-backbone-{}-{}x{}.png'.format(i + 1, w, h)))\n",
    "\n",
    "flux_ax = fig.add_subplot(gs[:1500, :])\n",
    "pyemma.plots.plot_flux(\n",
    "    cgflux,\n",
    "    pos=pos,\n",
    "    ax=flux_ax,\n",
    "    state_sizes=np.array([1.5 for _ in range(len(cgflux.stationary_distribution))]),\n",
    "    state_colors='None',\n",
    "    max_width=15,\n",
    "    max_height=15,\n",
    "    minflux=2e-5,\n",
    "    arrow_scale=1.,\n",
    "    size=5,\n",
    "    state_labels=None,\n",
    "    show_committor=True)\n",
    "flux_ax.set_xticks(np.arange(0, 1.2, .2))\n",
    "\n",
    "fig.savefig('data/figure_6.pdf', dpi=300)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And, finally, Figure 7 (a,b) depicts the Trp-1 autocorrelation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(3.47, 2.5))\n",
    "gw = int(np.floor(0.5 + 1000 * fig.get_figwidth()))\n",
    "gh = int(np.floor(0.5 + 1000 * fig.get_figheight()))\n",
    "gs = plt.GridSpec(gh, gw)\n",
    "gs.update(hspace=0.0, wspace=0.0, left=0.0, right=1.0, bottom=0.0, top=1.0)\n",
    "\n",
    "ax_box = fig.add_subplot(gs[:, :])\n",
    "ax_box.set_axis_off()\n",
    "ax_box.text(0.01, 0.95, '(a)', size=10)\n",
    "ax_box.text(0.01, 0.50, '(b)', size=10)\n",
    "\n",
    "ax_acf = fig.add_subplot(gs[50:1075, 500:3400])\n",
    "ax_acf.plot(eq_time_ml, eq_acf_ml, '-', color='C1', label='ML MSM')\n",
    "ax_acf.plot(\n",
    "    eq_time_bayes,\n",
    "    eq_acf_bayes,\n",
    "    '--',\n",
    "    color='C0',\n",
    "    label='Bayesian MSM')\n",
    "ax_acf.fill_between(\n",
    "    eq_time_bayes,\n",
    "    eq_acf_bayes_ci_l[1],\n",
    "    eq_acf_bayes_ci_u[1],\n",
    "    facecolor='C0',\n",
    "    alpha=0.3)\n",
    "ax_acf.semilogx()\n",
    "ax_acf.set_xlim((eq_time_ml[1], eq_time_ml[-1]))\n",
    "ax_acf.set_xticks([])\n",
    "ax_acf.set_ylabel(r'ACF / nm$^4$')\n",
    "ax_acf.legend()\n",
    "\n",
    "ax_rlx = fig.add_subplot(gs[1125:2150, 500:3400])\n",
    "ax_rlx.plot(eq_time_ml, eq_relax_ml, '-', color='C1', label='ML MSM')\n",
    "ax_rlx.plot(\n",
    "    eq_time_bayes,\n",
    "    eq_relax_bayes,\n",
    "    '--',\n",
    "    color='C0',\n",
    "    label='Bayesian MSM')\n",
    "ax_rlx.fill_between(\n",
    "    eq_time_bayes,\n",
    "    eq_relax_bayes_CI_l[1],\n",
    "    eq_relax_bayes_CI_u[1],\n",
    "    facecolor='C0',\n",
    "    alpha=0.3)\n",
    "ax_rlx.semilogx()\n",
    "ax_rlx.set_ylabel(r'Average / nm$^2$')\n",
    "ax_rlx.set_xlim((eq_time_ml[1], eq_time_ml[-1]))\n",
    "ax_rlx.set_xlabel(r'time / ns')\n",
    "\n",
    "fig.savefig('data/figure_7.pdf', dpi=300)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Wrapping up\n",
    "\n",
    "In this notebook, we have illustrated how to conduct an MD analysis with Markov models using PyEMMA.\n",
    "For the next steps, we recommend to work through the remaining eight notebooks, ideally in order,\n",
    "to strengthen your understanding of the PyEMMA software.\n",
    "\n",
    "1. [Data-I/O and featurization 📓](01-data-io-and-featurization.ipynb) \n",
    "2. [Dimension reduction and discretization 📓](02-dimension-reduction-and-discretization.ipynb) \n",
    "3. [MSM estimation and validation 📓](03-msm-estimation-and-validation.ipynb)\n",
    "4. [MSM analysis 📓](04-msm-analysis.ipynb)\n",
    "5. [PCCA and TPT analysis 📓](05-pcca-tpt.ipynb)\n",
    "6. [Expectations and observables 📓](06-expectations-and-observables.ipynb)\n",
    "7. [Hidden Markov state models (HMMs) 📓](07-hidden-markov-state-models.ipynb).\n",
    "8. [Common problems & bad data situations 📓](08-common-problems.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## References\n",
    "\n",
    "<a id=\"cite-gmrq\"/><sup><a href=#ref-1>[^]</a></sup>Robert T. McGibbon and Vijay S. Pande. 2015. _Variational cross-validation of slow dynamical modes in molecular kinetics_. [URL](https://doi.org/10.1063/1.4916292)\n",
    "\n",
    "<a id=\"cite-vamp-preprint\"/><sup><a href=#ref-2>[^]</a></sup>Wu, H. and Noé, F.. 2017. _Variational approach for learning Markov processes from time series data_. [URL](https://arxiv.org/pdf/1707.04659.pdf)\n",
    "\n",
    "<a id=\"cite-vampnet\"/><sup><a href=#ref-3>[^]</a></sup>Andreas Mardt and Luca Pasquali and Hao Wu and Frank Noé. 2018. _VAMPnets for deep learning of molecular kinetics_. [URL](https://doi.org/10.1038/s41467-017-02388-1)\n",
    "\n",
    "<a id=\"cite-tica2\"/><sup><a href=#ref-4>[^]</a></sup>Molgedey, L. and Schuster, H. G.. 1994. _Separation of a mixture of independent signals using time delayed correlations_. [URL](http://dx.doi.org/10.1103/PhysRevLett.72.3634)\n",
    "\n",
    "<a id=\"cite-tica\"/><sup><a href=#ref-5>[^]</a></sup>Guillermo Pérez-Hernández and Fabian Paul and Toni Giorgino and Gianni De Fabritiis and Frank Noé. 2013. _Identification of slow molecular order parameters for Markov model construction_. [URL](https://doi.org/10.1063/1.4811489)\n",
    "\n",
    "<a id=\"cite-msm-jhp\"/><sup><a href=#ref-6>[^]</a></sup>Prinz, Jan-Hendrik and Wu, Hao and Sarich, Marco and Keller, Bettina and Senne, Martin and Held, Martin and Chodera, John D. and Schütte, Christof and Noé, Frank. 2011. _Markov models of molecular kinetics: Generation and validation_. [URL](http://scitation.aip.org/content/aip/journal/jcp/134/17/10.1063/1.3565032)\n",
    "\n",
    "<a id=\"cite-swope-its\"/><sup><a href=#ref-7>[^]</a></sup>William C. Swope and Jed W. Pitera and Frank Suits. 2004. _Describing Protein Folding Kinetics by Molecular Dynamics Simulations. 1. Theory\\textdagger_. [URL](https://doi.org/10.1021/jp037421y)\n",
    "\n",
    "<a id=\"cite-pcca_plus_plus\"/><sup><a href=#ref-8>[^]</a></sup>Susanna Röblitz and Marcus Weber. 2013. _Fuzzy spectral clustering by PCCA+: application to Markov state models and data classification_. [URL](https://doi.org/10.1007/s11634-013-0134-6)\n",
    "\n",
    "<a id=\"cite-mdtraj\"/><sup><a href=#ref-9>[^]</a></sup>McGibbon, Robert T. and Beauchamp, Kyle A. and Harrigan, Matthew P. and Klein, Christoph and Swails, Jason M. and Hernández, Carlos X. and Schwantes, Christian R. and Wang, Lee-Ping and Lane, Thomas J. and Pande, Vijay S.. 2015. _MDTraj: A Modern Open Library for the Analysis of Molecular Dynamics Trajectories_.\n",
    "\n",
    "<a id=\"cite-simon-amm\"/><sup><a href=#ref-10>[^]</a></sup>Simon Olsson and Hao Wu and Fabian Paul and Cecilia Clementi and Frank Noé. 2017. _Combining experimental and simulation data of molecular processes via augmented Markov models_. [URL](https://doi.org/10.1073/pnas.1704803114)\n",
    "\n",
    "<a id=\"cite-noe-fingerprints\"/><sup><a href=#ref-11>[^]</a></sup>F. Noe and S. Doose and I. Daidone and M. Lollmann and M. Sauer and J. D. Chodera and J. C. Smith. 2011. _Dynamical fingerprints for probing individual relaxation processes in biomolecular dynamics with simulations and kinetic experiments_. [URL](https://doi.org/10.1073/pnas.1004646108)\n",
    "\n",
    "<a id=\"cite-simon-mech-mod-nmr\"/><sup><a href=#ref-12>[^]</a></sup>Simon Olsson and Frank Noé. 2016. _Mechanistic Models of Chemical Exchange Induced Relaxation in Protein NMR_. [URL](https://doi.org/10.1021/jacs.6b09460)\n",
    "\n",
    "<a id=\"cite-noe-dy-neut-scatt\"/><sup><a href=#ref-13>[^]</a></sup>Benjamin Lindner and Zheng Yi and Jan-Hendrik Prinz and Jeremy C. Smith and Frank Noé. 2013. _Dynamic neutron scattering from conformational dynamics. I. Theory and Markov models_. [URL](https://doi.org/10.1063/1.4824070)\n",
    "\n",
    "<a id=\"cite-noe-proj-hid-msm\"/><sup><a href=#ref-14>[^]</a></sup>Frank Noé and Hao Wu and Jan-Hendrik Prinz and Nuria Plattner. 2013. _Projected and hidden Markov models for calculating kinetics and metastable states of complex molecules_. [URL](https://doi.org/10.1063/1.4828816)\n",
    "\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.6.6"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": false,
   "sideBar": true,
   "skip_h1_title": true,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}