notebooks/06-expectations-and-observables.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 06 - Expectations and observables\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 covers the calculation of stationary and dynamic expectation values.\n",
    "These quantities include ensemble averaged observables measurable by experiments,\n",
    "correlation functions, and spectral densities.\n",
    "\n",
    "For this notebook, we assume familiarity with \n",
    "- data loading/visualization ([Notebook 01 ➜ 📓](01-data-io-and-featurization.ipynb))\n",
    "- dimension reduction ([Notebook 02 ➜ 📓](02-dimension-reduction-and-discretization.ipynb))\n",
    "- MSM estimation/validation ([Notebook 03 ➜ 📓](03-msm-estimation-and-validation.ipynb))\n",
    "- basic analyis techniques ([Notebook 04 ➜ 📓](04-msm-analysis.ipynb) and [Notebook 05 ➜ 📓](05-pcca-tpt.ipynb))\n",
    "\n",
    "Literature:\n",
    "\n",
    "- <a id=\"ref-1\" href=\"#cite-simon-amm\">olsson-17</a>\n",
    "- <a id=\"ref-2\" href=\"#cite-noe-fingerprints\">noe-11</a>\n",
    "- <a id=\"ref-3\" href=\"#cite-simon-mech-mod-nmr\">olsson-16</a>\n",
    "- <a id=\"ref-4\" href=\"#cite-noe-dy-neut-scatt\">lindner-13</a>\n",
    "\n",
    "\n",
    "Maintainers: [@psolsson](https://github.com/psolsson), [@cwehmeyer](https://github.com/cwehmeyer)\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 recall that a (stationary) ensemble average $\\langle O \\rangle$ is defined by a Boltzmann-weighted average with the potential energy function $E(\\cdot)$:\n",
    "\n",
    "$$ \\langle O \\rangle = \\mathcal{Z}^{-1}\\int_{\\Omega}\\;\\mathrm{d}x\\, o(x)\\exp(-\\beta E(x)) $$\n",
    "\n",
    "where $x \\in \\Omega$ is a molecular configuration and $o(\\cdot)$ is a function which relates the atomic coordinates of a configuration to a time-independent, microscopic observable.\n",
    "The function $o(\\cdot)$ is often called a 'forward model';\n",
    "it frequently involves computing distances or angles between particular atoms.\n",
    "For MSMs we have discretized our configurational space $\\Omega$, which simplifies the expression to the sum\n",
    "\n",
    "$$ \\langle O \\rangle = \\sum_i \\pi_i \\mathbf{o}_i $$\n",
    "\n",
    "where $\\pi_i$ corresponds to the integrated probability density of the segment of configuration space assigned to Markov state $i$.\n",
    "The stationary distribution $\\boldsymbol{\\pi}$ is computed as the left-eigenvector corresponding to the eigenvalue of $1$ of the MSM transition  matrix.\n",
    "The mapping of the experimental observable is mapped onto the Markov states as the vector $o$ with elements,\n",
    "\n",
    "$$ \\mathbf{o}_i = \\frac{1}{\\pi_i\\mathcal{Z}} \\int_{x\\in S_i} \\mathrm{d}x\\, o(x)\\exp(-\\beta E(x)). $$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import mdshare\n",
    "import pyemma"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Case 1: preprocessed, two-dimensional data (toy model)\n",
    "\n",
    "Again we estimate the MSM for the two-dimensional toy model to perform further analysis on."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "file = mdshare.fetch('hmm-doublewell-2d-100k.npz', working_directory='data')\n",
    "with np.load(file) as fh:\n",
    "    data = fh['trajectory']\n",
    "\n",
    "cluster = pyemma.coordinates.cluster_kmeans(data, k=50, max_iter=50)\n",
    "dtrajs_concatenated = cluster.dtrajs[0]\n",
    "\n",
    "its = pyemma.msm.its(\n",
    "    cluster.dtrajs, lags=[1, 2, 3, 5, 7, 10], nits=3, errors='bayes')\n",
    "\n",
    "fig, axes = plt.subplots(1, 3, figsize=(12, 3))\n",
    "pyemma.plots.plot_feature_histograms(\n",
    "    data, feature_labels=['$x$', '$y$'], ax=axes[0])\n",
    "pyemma.plots.plot_density(*data.T, ax=axes[1], cbar=False, alpha=0.1)\n",
    "axes[1].scatter(*cluster.clustercenters.T, s=15, c='C1')\n",
    "axes[1].set_xlabel('$x$')\n",
    "axes[1].set_ylabel('$y$')\n",
    "axes[1].set_xlim(-4, 4)\n",
    "axes[1].set_ylim(-4, 4)\n",
    "axes[1].set_aspect('equal')\n",
    "pyemma.plots.plot_implied_timescales(its, ylog=False, ax=axes[2])\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "msm = pyemma.msm.estimate_markov_model(cluster.dtrajs, lag=1)\n",
    "bayesian_msm = pyemma.msm.bayesian_markov_model(cluster.dtrajs, lag=1)\n",
    "\n",
    "print('fraction of states used = {:.2f}'.format(msm.active_state_fraction))\n",
    "print('fraction of counts used = {:.2f}'.format(msm.active_count_fraction))\n",
    "\n",
    "nstates = 2\n",
    "pyemma.plots.plot_cktest(msm.cktest(nstates));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### An imaginary experiment\n",
    "Let us first take a look at stationary ensemble averages for this system.\n",
    "Say the simulation data above represents a protein switching between two metastable configurations,\n",
    "perfectly separated by the $y$-coordinate.\n",
    "Imagine then that by inspection of the two metastable configurations,\n",
    "we have designed an experiment which allows us to measure an observable defined by\n",
    "\n",
    "$$ o(x,y) = 0.5x + y + 4. $$\n",
    "\n",
    "We compute the observables for the entire simulation trajectory,\n",
    "inspect the empirical histograms and the histograms of the metastable sets identified by a two-state PCCA++ analysis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_observable(x, y): \n",
    "    return 0.5 * x + y + 4\n",
    "\n",
    "\n",
    "pcca = msm.pcca(2)\n",
    "observable_traj = compute_observable(*data.T)\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "ax.hist(\n",
    "    observable_traj,\n",
    "    bins=128,\n",
    "    histtype='step',\n",
    "    label='all data',\n",
    "    color='k',\n",
    "    density=True)\n",
    "for num, metastable_set in enumerate(pcca.metastable_sets):\n",
    "    traj_indices = np.where(np.isin(dtrajs_concatenated, metastable_set))[0]\n",
    "    ax.hist(\n",
    "        observable_traj[traj_indices],\n",
    "        bins=64,\n",
    "        density=True,\n",
    "        histtype='step',\n",
    "        label='metastable {}'.format(num + 1))\n",
    "ax.legend()\n",
    "ax.set_xlabel(r'$o(x, y)$')\n",
    "ax.set_ylabel('empirical probability')\n",
    "ax.set_xlim(-0.5, 7.5)\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the case of nonequilibrium data, one might wish to compute a reweighted instead of the empirical distribution.\n",
    "To do so, we have to pass the output of `msm.trajectory_weights()` as weights for the histogramming as shown in the left panel below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "msm_trajectory_weights = msm.trajectory_weights()[0]\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "ax.hist(\n",
    "    observable_traj,\n",
    "    bins=128,\n",
    "    weights=msm_trajectory_weights,\n",
    "    histtype='step',\n",
    "    label='reweighted',\n",
    "    color='C0',\n",
    "    lw=2,\n",
    "    density=True)\n",
    "ax.hist(\n",
    "    observable_traj,\n",
    "    bins=128,\n",
    "    histtype='step',\n",
    "    label='empirical',\n",
    "    color='k',\n",
    "    density=True)\n",
    "ax.set_xlabel('$o(x, y)$')\n",
    "ax.legend()\n",
    "ax.set_ylabel('probability');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the present case, the data is well equilibrated and we see no difference between the empirical and the properly reweighted distribution. \n",
    "\n",
    "In the experiment we have measured this observable to be equal to $3.5$ a.u.\n",
    "In order to compute this experimental observable from our MSM from above, we need to compute the experimental observable for each of our Markov states by taking the average within the Markov state:\n",
    "\n",
    "$$ \\mathbf{o}_i = \\frac{1}{N_{a \\in S_i} } \\sum_{a \\in S_i} o(a_x,a_y). $$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "observable_mapped = np.array([observable_traj[dtrajs_concatenated == i].mean()\n",
    "                              for i in range(cluster.n_clusters)])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us visualize the observable on the cluster centers and compare to the stationary distribution of our Markov model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(2, 2, figsize=(10, 8), sharex=True, sharey=True)\n",
    "pyemma.plots.plot_contour(\n",
    "    *data.T,\n",
    "    observable_traj,\n",
    "    ax=axes[0, 0],\n",
    "    cbar_label='empirical observable',\n",
    "    mask=True)\n",
    "pyemma.plots.plot_density(\n",
    "    *data.T,\n",
    "    ax=axes[0, 1],\n",
    "    cbar_label='empirical distribution')\n",
    "pyemma.plots.plot_contour(\n",
    "    *data.T,\n",
    "    observable_mapped[dtrajs_concatenated],\n",
    "    ax=axes[1, 0],\n",
    "    cbar_label='discretized observable',\n",
    "    mask=True)\n",
    "pyemma.plots.plot_contour(\n",
    "    *data.T,\n",
    "    msm.stationary_distribution[dtrajs_concatenated],\n",
    "    ax=axes[1, 1],\n",
    "    cbar_label='MSM stationary distribution',\n",
    "    mask=True)\n",
    "for ax in axes.flat:\n",
    "    ax.set_xlim(-4, 4)\n",
    "    ax.set_ylim(-4, 4)\n",
    "    ax.set_aspect('equal')\n",
    "for ax in axes[1, :]:\n",
    "    ax.scatter(*cluster.clustercenters.T, s=15, c='C1')\n",
    "    ax.set_xlabel('$x$')\n",
    "for ax in axes[:, 0]:\n",
    "    ax.set_ylabel('$y$')\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "These plots illustrate how we lose some resolution in our observable when clustering the data.\n",
    "The reason is that the observable is considered constant at the sample average of all simulation frames assigned to a given cluster.\n",
    "With our vector of observables for our experiment we can compute the ensemble average by using the `expectation` method of our MSM:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Experimental result {:.3f}'.format(3.5))\n",
    "print('Prediction from MSM {:.3f}'.format(msm.expectation(observable_mapped)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We observe that our Markov model has a small deviation from the \"experimental\" value which we, in a practical case,\n",
    "would have to gauge against imprecisions in the prediction of the observable\n",
    "(some contributions include: approximations involving the forward model, the dicretization/projection error and limited sampling)\n",
    "and the experimental uncertainty.\n",
    "\n",
    "### Computing observables with error-bars\n",
    "\n",
    "If we have estimated a Bayesian MSM or HMM,\n",
    "we can compute limited sampling contribution to our experimental observable prediction by computing sample averages and confidence intervals."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "observable_sample_mean = bayesian_msm.sample_mean('expectation', observable_mapped)\n",
    "print('Bayesian Markov state model confidence interval {:.0f}%'.format(\n",
    "    bayesian_msm.conf * 100))\n",
    "\n",
    "observable_ci = bayesian_msm.sample_conf('expectation', observable_mapped)\n",
    "print('Observable prediction: {:.2f}, CI: [{:.2f}, {:.2f}]'.format(\n",
    "    observable_sample_mean,  *observable_ci))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have designed a new complementary experiment which is now given by\n",
    "\n",
    "$$o_2(x,y) = (y+2)^2 + x. $$\n",
    "\n",
    "#### Exercise 1\n",
    "Compute the difference of our MSM prediction of this observable to the experimental value of $5$ a.u.\n",
    "Compute the sample average and confidence interval from our Bayesian MSM.\n",
    "Is the experimental value within within the $95\\%$ confidence interval?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "solution2": "hidden",
    "solution2_first": true
   },
   "outputs": [],
   "source": [
    "def compute_new_observable(x, y):\n",
    "    return #FIXME\n",
    "\n",
    "\n",
    "new_observable_traj = #FIXME\n",
    "new_observable_mapped = np.array([observable_traj[dtrajs_concatenated == i].mean()\n",
    "                                  for i in range(cluster.n_clusters)])\n",
    "\n",
    "\n",
    "def is_within(x, a, b):\n",
    "    return ' not ' if x < a or x > b else ' '\n",
    "\n",
    "\n",
    "experiment2 = 5.0\n",
    "msm_prediction = #FIXME\n",
    "bayesian_msm_sample_mean = #FIXME\n",
    "bayesian_msm_ci95 = #FIXME\n",
    "\n",
    "print('Difference between MSM and experiment {:.2f}'.format(experiment2 - msm_prediction))\n",
    "\n",
    "print('Expl. {:.1f}, Sample mean {:.2f} and CI [{:.2f}, {:.2f}]'.format(\n",
    "    experiment2,\n",
    "    bayesian_msm_sample_mean,\n",
    "    *bayesian_msm_ci95))\n",
    "print('Experiment is{}within the 95% CI of the Bayesian MSM'.format(\n",
    "    is_within(experiment2, *bayesian_msm_ci95)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "solution2": "hidden"
   },
   "source": [
    "###### Solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "solution2": "hidden"
   },
   "outputs": [],
   "source": [
    "def compute_new_observable(x, y):\n",
    "    return (2 + y)**2 + x\n",
    "\n",
    "\n",
    "new_observable_traj = compute_new_observable(*data.T)\n",
    "new_observable_mapped = np.array([observable_traj[dtrajs_concatenated == i].mean()\n",
    "                                  for i in range(cluster.n_clusters)])\n",
    "\n",
    "\n",
    "def is_within(x, a, b):\n",
    "    return ' not ' if x < a or x > b else ' '\n",
    "\n",
    "\n",
    "experiment2 = 5.0\n",
    "msm_prediction = msm.expectation(new_observable_mapped)\n",
    "bayesian_msm_sample_mean = bayesian_msm.sample_mean('expectation', new_observable_mapped)\n",
    "bayesian_msm_ci95 = bayesian_msm.sample_conf('expectation', new_observable_mapped)\n",
    "\n",
    "print('Difference between MSM and experiment {:.2f}'.format(experiment2 - msm_prediction))\n",
    "\n",
    "print('Expl. {:.1f}, Sample mean {:.2f} and CI [{:.2f}, {:.2f}]'.format(\n",
    "    experiment2,\n",
    "    bayesian_msm_sample_mean,\n",
    "    *bayesian_msm_ci95))\n",
    "print('Experiment is{}within the 95% CI of the Bayesian MSM'.format(\n",
    "    is_within(experiment2, *bayesian_msm_ci95)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In cases like these, we can use the Augmented Markov models (AMMs) to estimate MSMs which optimally balance experimental and simulation data <a id=\"ref-5\" href=\"#cite-simon-amm\">olsson-16</a>.\n",
    "A walkthrough tutorial can be found here: http://www.emma-project.org/latest/generated/augmented_markov_model_walkthrough.html\n",
    "\n",
    "## Case 2: low-dimensional molecular dynamics data (alanine dipeptide)\n",
    "\n",
    "We fetch the alanine dipeptide data set, load the backbone torsions into memory,\n",
    "discretize the full space using $k$-means clustering,\n",
    "visualize the marginal and joint distributions of both components as well as the cluster centers,\n",
    "and show the ITS convergence:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pdb = mdshare.fetch('alanine-dipeptide-nowater.pdb', working_directory='data')\n",
    "files = mdshare.fetch('alanine-dipeptide-*-250ns-nowater.xtc', working_directory='data')\n",
    "\n",
    "feat = pyemma.coordinates.featurizer(pdb)\n",
    "feat.add_backbone_torsions(periodic=False)\n",
    "data = pyemma.coordinates.load(files, features=feat)\n",
    "data_concatenated = np.concatenate(data)\n",
    "\n",
    "cluster = pyemma.coordinates.cluster_kmeans(\n",
    "    data, k=100, max_iter=50, stride=10)\n",
    "dtrajs_concatenated = np.concatenate(cluster.dtrajs)\n",
    "\n",
    "its = pyemma.msm.its(\n",
    "    cluster.dtrajs, lags=[1, 2, 5, 10, 20, 50], nits=4, errors='bayes')\n",
    "\n",
    "fig, axes = plt.subplots(1, 3, figsize=(12, 3))\n",
    "pyemma.plots.plot_feature_histograms(\n",
    "    data_concatenated, feature_labels=['$\\Phi$', '$\\Psi$'], ax=axes[0])\n",
    "pyemma.plots.plot_density(*data_concatenated.T, ax=axes[1], cbar=False, alpha=0.1)\n",
    "axes[1].scatter(*cluster.clustercenters.T, s=15, c='C1')\n",
    "axes[1].set_xlabel('$\\Phi$')\n",
    "axes[1].set_ylabel('$\\Psi$')\n",
    "pyemma.plots.plot_implied_timescales(its, ax=axes[2], units='ps')\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We estimate a Markov model at a lagtime of $10$ ps and do a Chapman-Kolmogorov validation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "msm = pyemma.msm.estimate_markov_model(\n",
    "    cluster.dtrajs, lag=10, dt_traj='0.001 ns')\n",
    "\n",
    "print('fraction of states used = {:f}'.format(msm.active_state_fraction))\n",
    "print('fraction of counts used = {:f}'.format(msm.active_count_fraction))\n",
    "\n",
    "bayesian_msm = pyemma.msm.bayesian_markov_model(\n",
    "    cluster.dtrajs, lag=10, dt_traj='0.001 ns')\n",
    "\n",
    "nstates = 4\n",
    "pyemma.plots.plot_cktest(bayesian_msm.cktest(nstates), units='ps')\n",
    "\n",
    "msm.pcca(nstates)\n",
    "bayesian_msm.pcca(nstates);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using PyEMMA `coordinates` module to aid computation of molecular observables\n",
    "\n",
    "In this case we want to make use of features computed from molecular trajectory with the help of the PyEMMA `coordinates` module.\n",
    "For this small molecular systems we will focus on the $^1H^N-^1H^{\\alpha}$ $^3J$-coupling,\n",
    "a NMR parameter which is sensitive to the dihedral between the plane spanned by the inter-atomic vectors of H-N and N-C$\\alpha$ and the plane spanned by H$\\alpha$-C$\\alpha$ and C$\\alpha$-N.\n",
    "A common forward model to back-calculate $^3J$-couplings from a molecular configurations is called the Karplus equation\n",
    "\n",
    "$$ ^3J(\\theta) = A\\cos^2{\\theta}+B\\cos{\\theta} + C, $$\n",
    "\n",
    "where the (Karplus) parameters $A$, $B$ and $C$ are empirical constants which depend on the properties of the atoms involved in the observable.\n",
    "Here, we use the values $A=8.754 \\, \\mathrm{Hz}$, $B=−1.222\\, \\mathrm{Hz}$ and $C=0.111\\, \\mathrm{Hz}$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "feat2 = pyemma.coordinates.featurizer(pdb)\n",
    "feat2.add_dihedrals([[7, 6, 8, 9]], periodic=False)  # add relevant dihedral\n",
    "print('Atoms involving dihedral: {}'.format(\n",
    "    [feat2.topology.atom(i) for i in [7, 6, 8, 9]]))\n",
    "\n",
    "dihedral_trajs = pyemma.coordinates.load(files, features=feat2)  # load to memory\n",
    "\n",
    "\n",
    "def Karplus(theta):\n",
    "    \"\"\"The forward model.\"\"\"\n",
    "    return 8.754 * np.cos(theta)**2 - 1.222 * np.cos(theta) + 0.111\n",
    "\n",
    "\n",
    "# evaluate forward model on all dihedral trajectories\n",
    "observable_trajs = [Karplus(traj.ravel()) for traj in dihedral_trajs]\n",
    "observable_trajs_concatenated = np.concatenate(observable_trajs)\n",
    "\n",
    "# compute observable for Markov states\n",
    "jcoupl_markov = np.array([\n",
    "    observable_trajs_concatenated[dtrajs_concatenated == i].mean()\n",
    "    for i in range(cluster.n_clusters)])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us visualize this observable along with the stationary distribution as above:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1, 2, figsize=(10, 4))\n",
    "pyemma.plots.plot_contour(\n",
    "    *data_concatenated.T,\n",
    "    jcoupl_markov[dtrajs_concatenated],\n",
    "    ax=axes[0],\n",
    "    cbar_label='$^3J$-coupling / Hz',\n",
    "    mask=True)\n",
    "pyemma.plots.plot_contour(\n",
    "    *data_concatenated.T,\n",
    "    msm.pi[dtrajs_concatenated],\n",
    "    ax=axes[1],\n",
    "    cbar_label='MSM stationary distribution',\n",
    "    mask=True)\n",
    "for ax in axes.flat:\n",
    "    ax.scatter(*cluster.clustercenters.T, s=5, c='C1')\n",
    "    ax.set_xlabel('$\\Phi$')\n",
    "    ax.set_aspect('equal')\n",
    "axes[0].set_ylabel('$\\Psi$')\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Exercise 2\n",
    "Predict the value of the $^3J$-coupling using the `msm` instance of 2-Ala."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "skip": true,
    "solution2": "hidden",
    "solution2_first": true
   },
   "outputs": [],
   "source": [
    "predicted_3j #FIXME\n",
    "print('Predicted value of 3J coupling {:.3f} Hz'.format(predicted_3j))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "solution2": "hidden"
   },
   "source": [
    "###### Solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "solution2": "hidden"
   },
   "outputs": [],
   "source": [
    "predicted_3j = msm.expectation(jcoupl_markov)\n",
    "print('Predicted value of 3J coupling {:.3f} Hz'.format(predicted_3j))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Exercise 3\n",
    "\n",
    "Compute the sample mean and sample standard deviation of the $^3J$-coupling using the `bayesian_msm` instance of 2-Ala.\n",
    "\n",
    "**Hint**: use the `sample_std` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "skip": true,
    "solution2": "hidden",
    "solution2_first": true
   },
   "outputs": [],
   "source": [
    "bayesian_sample_mean_3j = #FIXME\n",
    "bayesian_ci_3j = #FIXME\n",
    "\n",
    "print('Predicted 3J coupling {:.3f} ({:.3f}) Hz'.format(\n",
    "    bayesian_sample_mean_3j, bayesian_ci_3j))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "solution2": "hidden"
   },
   "source": [
    "###### Solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "solution2": "hidden"
   },
   "outputs": [],
   "source": [
    "bayesian_sample_mean_3j = bayesian_msm.sample_mean('expectation', jcoupl_markov)\n",
    "bayesian_ci_3j = bayesian_msm.sample_std('expectation', jcoupl_markov)\n",
    "\n",
    "print('Predicted 3J coupling {:.3f} ({:.3f}) Hz'.format(\n",
    "    bayesian_sample_mean_3j, bayesian_ci_3j))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dynamic/kinetic experimental observables\n",
    "\n",
    "We can measure kinetic information experimentally through time-correlation functions or spectral densities.\n",
    "Markov models are probabilistic models of the dynamics, i.e.,\n",
    "if we have an ergodic dynamics described by a Markov model,\n",
    "we can compute the equilibrium time-correlation functions using\n",
    "\n",
    "$$ \\mathbb{E}[o(x_{t})o(x_{t+k\\tau})] = \\underbrace{\\mathbf{o}^T}_{\\text{transpose of vector of observables in Markov states}} \\overbrace{\\boldsymbol{\\Pi}}^{\\text{diagonal matrix of stationary distribution}} \\underbrace{\\mathbf{P}^k(\\tau)}_{\\text{transition matrix}} \\overbrace{\\mathbf{o}}^{\\text{vector of observables in Markov states}} $$\n",
    "\n",
    "We can recast this expression to\n",
    "\n",
    "$$ \\mathbb{E}[o(x_{t})o(x_{t+k\\tau})] = (\\mathbf{o}^T\\boldsymbol{\\pi})^2 + \\sum^N_{i=2} \\exp\\left(-\\frac{k\\tau}{t_i}\\right)(\\mathbf{o}^T\\boldsymbol{\\phi}_i)^2 $$\n",
    "\n",
    "by using the spectral decomposition of the transition matrix ([1](#References)).\n",
    "$\\phi_i$ is the $i$th left eigenvector of $\\mathbf{P}(\\tau)$, with the associated implied timescale $t_i$.\n",
    "We see that the auto-correlation functions take the form of a multi-exponential decay. \n",
    "\n",
    "We now consider relaxation from a nonequilibrium state $\\mathbf{p_0}$, which can be achieved by T-jump or P-jump experiments.\n",
    "In this case, the initial distribution and final ensembles are not the same ($\\mathbf{p_0}$ and $\\boldsymbol{\\pi}$) and so the auto-correlation becomes:\n",
    "\n",
    "$$ \\mathbb{E}[o(x_{t})o(x_{t+k\\tau})] = (\\mathbf{\\hat p_0}^T\\boldsymbol{\\pi})(\\mathbf{l}^T\\boldsymbol{\\pi}) + \\sum^N_{i=2} \\exp\\left(-\\frac{k\\tau}{\\lambda_i}\\right)(\\mathbf{o}^T\\boldsymbol{\\phi}_i)(\\mathbf{\\hat p_0}^T\\boldsymbol{\\phi}_i) $$\n",
    "\n",
    "with $\\mathbf{\\hat p_0}=\\boldsymbol{\\Pi}^{-1}\\mathbf{p_0}$.\n",
    "\n",
    "\n",
    "In PyEMMA, the MSM and HMM objects have the methods `correlation` and `relaxation` which implement calculation of auto-correlation in equilibrium and relaxation from a specified nonequilibrium distribution.\n",
    "Many experimental observables do not measure correlation functions directly but some quantity which depends on a correlation function, or its Fourier transform ([2-3](#References)).\n",
    "In many of these cases it is possible to obtain analytical expressions of the observable which depends only on the amplitudes $(\\mathbf{o}^T\\boldsymbol{\\phi}_i)^2$ and time-scales $t_i$,\n",
    "quantities which can be computed given $\\mathbf{o}$ using the `fingerprint_correlation` and `fingerprint_relaxation` methods of the MSM or HMM objects. \n",
    "\n",
    "We here compute the auto-correlation function of the distance between the amide and alpha protons."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "feat3 = pyemma.coordinates.featurizer(pdb)\n",
    "feat3.add_distances([[7, 9]], periodic=False)  # add relevant distance\n",
    "print('Distance between atoms: {}'.format(\n",
    "    [feat3.topology.atom(i) for i in [7, 9]]))\n",
    "\n",
    "proton_distance_trajs = pyemma.coordinates.load(files, features=feat3)  # load to memory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dist_markov = np.array([\n",
    "    np.concatenate(proton_distance_trajs).flatten()[dtrajs_concatenated == i].mean() \n",
    "    for i in range(cluster.n_clusters)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eq_time_ml, eq_acf_ml = msm.correlation(dist_markov, maxtime=15)\n",
    "\n",
    "eq_time_bayes, eq_acf_bayes = bayesian_msm.sample_mean(\n",
    "    'correlation',\n",
    "    dist_markov,\n",
    "    maxtime=15)\n",
    "\n",
    "eq_acf_bayes_ci_l, eq_acf_bayes_ci_u = bayesian_msm.sample_conf(\n",
    "    'correlation',\n",
    "    dist_markov,\n",
    "    maxtime=15)\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_ylim([x * 0.999 if i==0 else x * 1.001 for i, x in enumerate(ax.get_ylim())])\n",
    "ax.set_xlabel(r'time / $\\mathrm{ns}$')\n",
    "ax.set_ylabel(r'd($^{\\alpha}$H, $^{\\mathrm{N}}$H) ACF / $\\mathrm{nm}^2$')\n",
    "ax.legend()\n",
    "fig.tight_layout()\n",
    "\n",
    "print(eq_time_bayes)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The amplitude of the computed auto-correlation functions is very small.\n",
    "In experiments, we can often prepare a nonequilibrium ensemble, and measure the relaxation back to the steady-state as discussed above.\n",
    "In PyEMMA, we simulate such experiments using the `relaxation` method of the `msm` object.\n",
    "In addition to the arguments required by `correlation`, the `relaxation` computation requires an initial condition $p_0$.\n",
    "Let's say we can prepare the 2-ALA ensemble such that it coincides with metastable distribution $1$,\n",
    "as identified by the PCCA++ analysis carried out above:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eq_time_ml, eq_acf_ml = msm.relaxation(\n",
    "    msm.metastable_distributions[1],\n",
    "    dist_markov,\n",
    "    maxtime=15)\n",
    "\n",
    "eq_time_bayes, eq_acf_bayes = bayesian_msm.sample_mean(\n",
    "    'relaxation',\n",
    "    msm.metastable_distributions[1],\n",
    "    dist_markov,\n",
    "    maxtime=15)\n",
    "\n",
    "eq_acf_bayes_CI_l, eq_acf_bayes_CI_u = bayesian_msm.sample_conf(\n",
    "    'relaxation', \n",
    "    msm.metastable_distributions[1],\n",
    "    dist_markov,\n",
    "    maxtime=15)\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "ax.semilogx(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'Avg. d($^{\\alpha}$H, $^{\\mathrm{N}}$H) / $\\mathrm{nm}$')\n",
    "\n",
    "ax.legend()\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Much better! This signal is orders of magnitude stronger than the one observed for the auto-correlation function at equilibrium.\n",
    "\n",
    "### Dynamic fingerprints\n",
    "\n",
    "A new experiment allows us to measure the quantity:\n",
    "\n",
    "$$ \\sum_{i=2}^N \\overbrace{t_i}^{\\text{implied timescale }i} \\underbrace{(\\mathbf{o}^T\\boldsymbol{\\phi}_i)}_{\\text{Fingerprint amplitude } i}$$\n",
    "\n",
    "where $\\mathbf{o}$ is the distance amide and alpha protons mapped onto the Markov states.\n",
    "\n",
    "#### Exercise 4\n",
    "\n",
    "Using the `fingerprint_correlation` method, compute the timescales and amplitudes of this new observable.\n",
    "\n",
    "**Hint**: the vector $\\mathbf{o}$ was computed above and is stored in the variable `dist_markov`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "skip": true,
    "solution2": "hidden",
    "solution2_first": true
   },
   "outputs": [],
   "source": [
    "def compute_new_observable(timescales, amplitudes):\n",
    "    return timescales[1:].dot(amplitudes[1:])\n",
    "\n",
    "\n",
    "timescales, amplitudes = #FIXME\n",
    "\n",
    "print('Value of new observable {:.3e} nm ns'.format(\n",
    "    compute_new_observable(timescales, amplitudes)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "solution2": "hidden"
   },
   "source": [
    "###### Solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "solution2": "hidden"
   },
   "outputs": [],
   "source": [
    "def compute_new_observable(timescales, amplitudes):\n",
    "    return timescales[1:].dot(amplitudes[1:])\n",
    "\n",
    "\n",
    "timescales, amplitudes = msm.fingerprint_correlation(dist_markov)\n",
    "\n",
    "print('Value of new observable {:.3e} nm ns'.format(\n",
    "    compute_new_observable(timescales, amplitudes)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Wrapping up\n",
    "\n",
    "In this notebook, we have covered how to compute ensemble averaged properties/observables from Markov state models and hidden Markov state models.\n",
    "We have further covered the quantification of uncertainty in these predictions.\n",
    "We have specifically discussed the MSM/HMM object methods\n",
    "* `expectation()` computes the stationary ensemble average of a property.\n",
    "* `correlation()` computes the equilibrium auto/cross-correlation function of a/two property/properties.\n",
    "* `relaxation()` computes the relaxation auto/cross-correlation function of a/two property/properties given an initial state distribution.\n",
    "* `fingerprint_correlation()` computes the amplitudes and time-scales of the MSM/HMM auto/cross-correlation function.\n",
    "* `fingerprint_relaxation()` computes the amplitudes and time-scales of the MSM/HMM relaxation auto/cross-correlation function.\n",
    "We have further covered the following object methods to compute statistics from Bayesian MSM/HMMs:\n",
    "* `sample_mean()` computes the mean of a property over the sampled MSMs/HMMs in a Bayesian model. \n",
    "* `sample_conf()` computes the confidence interval of a property over the sampled MSMs/HMMs in a Bayesian model.\n",
    "* `sample_std()` computes the standard deviation of a property over the sampled MSMs/HMMs in a Bayesian model.\n",
    "\n",
    "Finally, we have shown how to use these methods together with precomputed observables."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## References\n",
    "\n",
    "<a id=\"cite-simon-amm\"/><sup><a href=#ref-1>[^]</a><a href=#ref-5>[^]</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-2>[^]</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-3>[^]</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-4>[^]</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"
   ]
  }
 ],
 "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.5"
  },
  "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
}