notebooks/07-hidden-markov-state-models.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 07 - Hidden Markov state models (HMMs)\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 will learn about hidden Markov state models and how to use them to deal with poor discretization.\n",
    "We further explain how to obtain a coarse-grained model based on an initial MSM analysis.\n",
    "\n",
    "The prequisites for this tutorial are\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",
    "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",
    "When estimating a regular MSM, we assume that the dynamics between microstates defined by our clustering algorithm is Markovian.\n",
    "Hidden Markov models (HMMs), in comparison, only make this assumption in the space of so-called hidden states that have not been directly observed.\n",
    "This means that instead of finding a maximum likelihood (ML) transition matrix between microstates,\n",
    "an ML transition matrix is estimated between hidden states.\n",
    "In the same step, the probability of a microstate to belong to a certain hidden state is estimated.\n",
    "\n",
    "⚠️ Please note that the PyEMMA implementation internally initiates the estimation with a regular MSM and PCCA++.\n",
    "\n",
    "There are two major use-cases for HMMs in PyEMMA:\n",
    "- HMMs are more robust to poor space discretization and can be used to overcome difficult clustering situations\n",
    "- HMMs offer a coarse graining into metastable (hidden) states\n",
    "\n",
    "In this notebook, we will demonstrate how to estimate HMMs and how they behave in comparison to MSMs.\n",
    "\n",
    "Literature:\n",
    "- <a id=\"ref-1\" href=\"#cite-hmm-baum-welch-alg\">baum-1970</a>\n",
    "- <a id=\"ref-2\" href=\"#cite-noe-proj-hid-msm\">noe-13</a>\n",
    "- <a id=\"ref-3\" href=\"#cite-hmm-tutorial\">rabiner-89</a>\n",
    "\n",
    "⚠️ We have assigned the integer numbers $1 \\dots $ `nstates` to HMM metastable states.\n",
    "As PyEMMA is written in Python, it internally indexes states starting from $0$.\n",
    "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": [
    "%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",
    "In this example, we are going to demonstrate the robustness of HMMs against poor discretization and show some of its properties.\n",
    "We start by loading the two-dimensional data as well as the true discrete trajectory from an archive using numpy:"
   ]
  },
  {
   "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",
    "    good_dtraj = fh['discrete_trajectory']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now estimate a reference (regular) MSM from the well-discretized data which is shown in the next panel (left).\n",
    "We include an implied timescales plot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1, 2, figsize=(10, 4))\n",
    "pyemma.plots.plot_state_map(*data.T, good_dtraj, ax=axes[0])\n",
    "axes[0].scatter(*np.asarray([[0, -1], [0, 1]]).T, s=15, c='C1')\n",
    "\n",
    "axes[0].set_xlabel('$x$')\n",
    "axes[0].set_xlim(-4, 4)\n",
    "axes[0].set_ylim(-4, 4)\n",
    "axes[0].set_aspect('equal')\n",
    "axes[0].set_ylabel('$y$')\n",
    "axes[0].set_title('discretization')\n",
    "\n",
    "lags = [i for i in range(1, 10)]\n",
    "pyemma.plots.plot_implied_timescales(\n",
    "    pyemma.msm.its(good_dtraj, lags=lags, errors='bayes'), ylog=False, ax=axes[1])\n",
    "axes[1].set_title('MSM with good discretization')\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We note that for this very good discretization, the implied timescales are converged  from lagtime $1$ step.\n",
    "We continue to build an MSM object and perform the Chapman-Kolmogorov test:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "reference_msm = pyemma.msm.estimate_markov_model(good_dtraj, lag=1)\n",
    "pyemma.plots.plot_cktest(reference_msm.cktest(2));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The Chapman-Kolmogorov test shows excellent agreement between higher lagtime estimation and model prediction.\n",
    "We thus take this model as a reference.\n",
    "\n",
    "Let's now deliberately choose a very bad discretization..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "poor_clustercenters = np.asarray([[-2.5, -1.4], \n",
    "                                  [0.3, 1.2], \n",
    "                                  [2.7, -0.6]])\n",
    "poor_dtraj = pyemma.coordinates.assign_to_centers(data, centers=poor_clustercenters)[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "... and repeat the ITS estimation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1, 2, figsize=(10, 4))\n",
    "pyemma.plots.plot_state_map(*data.T, poor_dtraj, ax=axes[0])\n",
    "axes[0].scatter(*poor_clustercenters.T, s=15, c='C1')\n",
    "\n",
    "axes[0].set_xlabel('$x$')\n",
    "axes[0].set_xlim(-4, 4)\n",
    "axes[0].set_ylim(-4, 4)\n",
    "axes[0].set_aspect('equal')\n",
    "axes[0].set_ylabel('$y$')\n",
    "axes[0].set_title('discretization')\n",
    "\n",
    "pyemma.plots.plot_implied_timescales(\n",
    "    pyemma.msm.its(poor_dtraj, lags=lags, errors='bayes'), ylog=False, ax=axes[1])\n",
    "axes[1].set_title('MSM with poor discretization')\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Obviously, the discretization is very poor and does not mirror the basins of the double-well potential anymore.\n",
    "All three discrete states include data points from the two metastable regions (left panel) and,\n",
    "as the right panel shows, this discretization error cannot be fixed by using a large lagtime for a regular MSM estimation.\n",
    "Thus, the MSM clearly is not able to resolve the slow process connecting the two basins.\n",
    "\n",
    "⚠️ We do not see any ITS above the lag time horizon and, hence, cannot estimate any MSM with this discretization.\n",
    "\n",
    "Let us now repeat both estimations, using well and poorly discretized data, with hidden Markov models instead of regular MSMs.\n",
    "We begin with the implied timescale convergence using the `pyemma.msm.timescales_hmsm()` function and two hidden states:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "its_hmm_poor = pyemma.msm.timescales_hmsm(poor_dtraj, 2, lags=lags, errors='bayes')\n",
    "its_hmm_good = pyemma.msm.timescales_hmsm(good_dtraj, 2, lags=lags, errors='bayes')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We go on visualizing the results as with regular MSM implied timescales and include, as a dotted line,\n",
    "the result from our previously estimated reference MSM. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1, 2, figsize=(10, 4), sharey=True)\n",
    "pyemma.plots.plot_implied_timescales(\n",
    "    its_hmm_poor,\n",
    "    ylog=False, ax=axes[0])\n",
    "pyemma.plots.plot_implied_timescales(\n",
    "    its_hmm_good,\n",
    "    ylog=False, ax=axes[1])\n",
    "axes[0].set_title('HMM with poor discretization')\n",
    "axes[1].set_title('HMM with good discretization')\n",
    "\n",
    "for n, ax in enumerate(axes.flat):\n",
    "    ax.set_ylim(-0.5, 12.5)\n",
    "    ax.hlines(reference_msm.timescales()[0], *ax.get_xlim(), linestyle=':', label='reference MSM' if n == 0 else None)\n",
    "fig.legend(loc=9)\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In contrast to a regular MSM, both discretizations give us converged implied timescales from the very start (lagtime $1$ step).\n",
    "The only difference is that the poor discretization yields a larger error and we loose the process faster.\n",
    "As the HMM computes the implied timescales of a process between two hidden states,\n",
    "we do not assume Markovianity in the original state space.\n",
    "Thus, the deliberate discretization error we made is compensated by the algorithm,\n",
    "making it robust against poor clustering.\n",
    "\n",
    "In order to validate this claim, we estimate HMMs using both discretizations at lagtime $1$ step and two hidden states..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "poor_hmm = pyemma.msm.estimate_hidden_markov_model(poor_dtraj, 2, lag=1)\n",
    "good_hmm = pyemma.msm.estimate_hidden_markov_model(good_dtraj, 2, lag=1)\n",
    "\n",
    "print('MSM (ref):  1. implied timescale = {:.2f} steps'.format(reference_msm.timescales()[0]))\n",
    "print('HMM (poor): 1. implied timescale = {:.2f} steps'.format(poor_hmm.timescales()[0]))\n",
    "print('HMM (good): 1. implied timescale = {:.2f} steps'.format(good_hmm.timescales()[0]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "... and obtain nearly identical estimates for the first implied timescale that agree with the reference MSM.\n",
    "\n",
    "We observe that HMMs, unlike MSMs, seem to be somewhat resistant to discretization errors.\n",
    "\n",
    "Regarding the CK test, we again see that the `poor_hmm`..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyemma.plots.plot_cktest(poor_hmm.cktest(2));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "... and the `good_hmm`..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyemma.plots.plot_cktest(good_hmm.cktest(2));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "... are in perfect agreement and the Chapman-Kolmogorov test has passed.\n",
    "\n",
    "The HMM, even for the poor discretization, has learned an assignment between the microstates to the hidden states that reflects the true dynamics.\n",
    "We can extract this information from the HMM object using its `hmm.hidden_state_probabilities` property.\n",
    "It contains the probabilities for each microstate to be in a given hidden state over time,\n",
    "for each trajectory (which is why we have to take the $0$-th element from this list). \n",
    "\n",
    "Please note that there is also a time-independent property that contains the observation probabilities of a hidden state as a function of the microstates.\n",
    "It is stored in `hmm.observation_probabilities` but not demonstrated here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(2, 1, figsize=(10, 4), sharex=True)\n",
    "axes[0].plot(data[:100, 1], 'r.-', label='cont. traj')\n",
    "axes[0].step(np.arange(100), poor_clustercenters[poor_dtraj[:100], 0], label='discrete traj', linewidth=1, c='k')\n",
    "axes[0].set_ylabel('y coordinate \\ncluster center')\n",
    "\n",
    "for n, hidden_state_probability_traj in enumerate(poor_hmm.hidden_state_probabilities[0].T):\n",
    "    axes[1].plot(hidden_state_probability_traj[:100], '.-', label='P(state {})'.format(n + 1))\n",
    "    \n",
    "axes[1].set_ylabel('hidden state prob')\n",
    "fig.legend();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we see in the trajectory excerpt of the first $100$ steps,\n",
    "the microstate assignment probabilities to hidden states over time reflect the actual dynamics very well,\n",
    "even though the discrete trajectory has basically lost notion of the original double-well basins.\n",
    "\n",
    "Let's now plot the hidden state probabilities in the original space:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1, 2, figsize=(10, 4))\n",
    "for i, ax in enumerate(axes.flat):\n",
    "    pyemma.plots.plot_contour(\n",
    "        *data.T,\n",
    "        poor_hmm.hidden_state_probabilities[0][:, i], # index 0: only 1 trajectory\n",
    "        ax=ax,\n",
    "        cmap='afmhot_r', \n",
    "        mask=True,\n",
    "        cbar_label='P(state {})'.format(i+1))\n",
    "    ax.set_xlabel('$y')\n",
    "axes[0].set_ylabel('$x$')\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Even though we find some artifacts of the initial clustering,\n",
    "the HMM has basically overcome the poor discretization and found hidden states that,\n",
    "with high certainty, mirror the original double-well basins.\n",
    "\n",
    "Let's finish this example by comparing the hidden state trajectories with the discrete trajectories that were used for estimating the reference MSM (the \"true\" discrete trajectories)..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('HMM (good): Hidden state trajectory consistency: '\n",
    "      '{:.3f}'.format(sum(good_hmm.hidden_state_trajectories[0] == good_dtraj)/good_dtraj.shape[0]))\n",
    "print('HMM (poor): Hidden state trajectory consistency: '\n",
    "      '{:.3f}'.format(sum(poor_hmm.hidden_state_trajectories[0] == good_dtraj)/good_dtraj.shape[0]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "... and the stationary distributions of both HMMs to the reference MSM:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('MSM (ref):  stationary distribution = {}'.format(np.round(reference_msm.pi, 4)))\n",
    "print('HMM (poor): stationary distribution = {}'.format(np.round(poor_hmm.pi, 4)))\n",
    "print('HMM (good): stationary distribution = {}'.format(np.round(good_hmm.pi, 4)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We find that even with a very poor discretization,\n",
    "HMMs are capable of recovering the kinetics of the underlying process with very little error. \n",
    "\n",
    "## Case 2: low-dimensional molecular dynamics data (alanine dipeptide)\n",
    "\n",
    "We are now illustrating a typical use case of hidden markov state models:\n",
    "estimating an MSM that is used as a heuristics for the number of slow processes or hidden states,\n",
    "and estimating an HMM (to overcome potential discretization issues and to resolve faster processes than an MSM).\n",
    "\n",
    "We fetch the alanine dipeptide data set, load the backbone torsions into memory..."
   ]
  },
  {
   "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(data, k=75, max_iter=50, stride=10)\n",
    "dtrajs_concatenated = np.concatenate(cluster.dtrajs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "... 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 to help selecting a suitable lag time:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "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(data_concatenated, feature_labels=['$\\Phi$', '$\\Psi$'], ax=axes[0])\n",
    "pyemma.plots.plot_density(*data_concatenated.T, ax=axes[1], cbar=False, alpha=0.3)\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": [
    "Based on the implied timescale convergence plot, we choose a lagtime of $10$ ps.\n",
    "We further find $3$ slow processes in the implied timescales plot,\n",
    "meaning that we can expect $4$ metastable sets or hidden states.\n",
    "\n",
    "First, we estimate a Bayesian MSM, and show the results of a CK test:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bayesian_msm = pyemma.msm.bayesian_markov_model(cluster.dtrajs, lag=10, dt_traj='1 ps')\n",
    "\n",
    "nstates = 4\n",
    "pyemma.plots.plot_cktest(bayesian_msm.cktest(nstates), units='ps');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "At this point, we have a (bayesian) MSM with $75$ discrete states and basic validation.\n",
    "To obtain an HMM with only four states (the number for which we have validated our MSM),\n",
    "we compute the implied timescales for HMMs with this number of hidden states. \n",
    "\n",
    "We repeat the ITS convergence analysis using (bayesian) HMMs and small lagtimes for a $4$-state HMM.\n",
    "For demonstration purposes, we add the same analysis with a $6$-state HMM to visualize what happens if the number of states is not as clear as in this example: \n",
    "\n",
    "⚠️ Executing the following cell might take a few minutes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1, 2, figsize=(10, 4), sharey=True)\n",
    "pyemma.plots.plot_implied_timescales(\n",
    "    pyemma.msm.timescales_hmsm(\n",
    "        cluster.dtrajs, 4, lags=[1, 2, 3, 6], errors='bayes', nsamples=50),\n",
    "    ax=axes[0], units='ps')\n",
    "pyemma.plots.plot_implied_timescales(\n",
    "    pyemma.msm.timescales_hmsm(\n",
    "        cluster.dtrajs, 6, lags=[1, 2, 3, 6], errors='bayes', nsamples=50),\n",
    "    ax=axes[1], units='ps')\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The left panel shows that an HMM with four hidden states yields converged implied timescales from lagtime $1$.\n",
    "\n",
    "The right panel, however, shows that an HMM with six hidden states and lagtime $1$ can resolve two additional processes.\n",
    "\n",
    "Let us follow up on this and perform a CK test for a four state HMM at lagtime $1$..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hmm_4 = pyemma.msm.bayesian_hidden_markov_model(cluster.dtrajs, 4, lag=1, dt_traj='1 ps', nsamples=50)\n",
    "pyemma.plots.plot_cktest(hmm_4.cktest(mlags=6), units='ps');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "... and than the six state HMM at laggtime $1$ (we use `mlags=2` because we would loose the two fast processes at lagtimes $\\geq3$):\n",
    "\n",
    "⚠️ Executing the following cell might take a few minutes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hmm_6 = pyemma.msm.bayesian_hidden_markov_model(cluster.dtrajs, 6, lag=1, dt_traj='1 ps', nsamples=50)\n",
    "pyemma.plots.plot_cktest(hmm_6.cktest(mlags=2), units='ps');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In both cases, the CK test is passed.\n",
    "\n",
    "If we now compare both metastable membership plots..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1, 2, figsize=(10, 4), sharex=True, sharey=True)\n",
    "for hmm, ax in zip([hmm_4, hmm_6], axes.flat):\n",
    "    _, _, misc = pyemma.plots.plot_state_map(\n",
    "        *data_concatenated.T,\n",
    "        hmm.metastable_assignments[dtrajs_concatenated], \n",
    "        ax=ax)\n",
    "    ax.set_title('HMM with {} hidden states'.format(hmm.nstates))\n",
    "    ax.set_xlabel('$\\Phi$')\n",
    "    misc['cbar'].set_ticklabels(range(1, hmm.nstates + 1))\n",
    "axes[0].set_ylabel('$\\Psi$')\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "... we see that the six state HMM is able to subdivide the two largest metastable states of the four state HMM and,\n",
    "thus, gives us a more detailed view on the underlying system.\n",
    "As one would have expected from the implied timescale plot,\n",
    "the metastable dynamics is already well-described with $4$ hidden states.\n",
    "\n",
    "⚠️ Due to the low sensibility to discretization errors,\n",
    "we can additionally afford to estimate HMMs at smaller lagtimes than MSMs and,\n",
    "thus, resolve more processes.\n",
    "\n",
    "Like with classical MSMs, we can further analyze properties of the HMM.\n",
    "As an example, have a look at the transition paths and committor probabilities below. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "A = [0]\n",
    "B = [3]\n",
    "flux = pyemma.msm.tpt(hmm_4, A, B)\n",
    "\n",
    "highest_membership = hmm_4.metastable_distributions.argmax(1)\n",
    "coarse_state_centers = cluster.clustercenters[hmm_4.observable_set[highest_membership]]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Please note one important difference to MSMs:\n",
    "since HMMs operate directly on the hidden states,\n",
    "we must not compute the flux between the `msm.metastable_sets` but between the lists of macrostate numbers,\n",
    "e.g., instead of `A = msm.metastable_sets[0]` we set `A = [0]`. \n",
    "\n",
    "Let's now visualize the committor as before.\n",
    "Does it look familiar?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(10, 7))\n",
    "\n",
    "pyemma.plots.plot_contour(\n",
    "    *data_concatenated.T,\n",
    "    flux.committor[hmm_4.metastable_assignments[dtrajs_concatenated]],\n",
    "    cmap='brg',\n",
    "    ax=ax,\n",
    "    mask=True,\n",
    "    cbar_label=r'committor 0 $\\to$ 3',\n",
    "    alpha=0.8,\n",
    "    zorder=-1)\n",
    "\n",
    "pyemma.plots.plot_flux(\n",
    "    flux,\n",
    "    coarse_state_centers,\n",
    "    flux.stationary_distribution,\n",
    "    ax=ax,\n",
    "    show_committor=False,\n",
    "    figpadding=0,\n",
    "    show_frame=True,\n",
    "    state_labels=['A','' ,'', 'B'], \n",
    "    arrow_label_format='%2.e / ps');\n",
    "\n",
    "ax.set_xlabel('$\\Phi$')\n",
    "ax.set_ylabel('$\\Psi$')\n",
    "ax.set_xlim(data_concatenated[:, 0].min(), data_concatenated[:, 0].max())\n",
    "ax.set_ylim(data_concatenated[:, 1].min(), data_concatenated[:, 1].max())\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we see, in addition to the properties described above, HMMs provide the same analysis tools as MSMs.\n",
    "For example, eigenvectors and mean first passage times can be extracted as described in [Notebook 04 ➜ 📓](04-msm-analysis.ipynb) and [Notebook 05 ➜ 📓](05-pcca-tpt.ipynb). \n",
    "\n",
    "Let us now repeat this approach again for another featurization:\n",
    "we already know that it is possible to resolve six metastable states (five slow processes) using an HMM estimated on a discretization of the backbone torsions.\n",
    "Can you achieve the same level of resolution using heavy atom distances and a suitable TICA projection?\n",
    "\n",
    "#### Exercise 1\n",
    "\n",
    "Obtain the heavy atom distances, use TICA for dimension reduction, and discretize using a method of your choice."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "solution2": "hidden",
    "solution2_first": true
   },
   "outputs": [],
   "source": [
    "feat = #FIXME\n",
    "feat. #FIXME\n",
    "data = #FIXME\n",
    "\n",
    "tica = #FIXME\n",
    "tica_concatenated = np.concatenate(tica.get_output())\n",
    "\n",
    "cluster = #FIXME\n",
    "dtrajs_concatenated = #FIXME\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(tica_concatenated, ax=axes[0])\n",
    "pyemma.plots.plot_density(*tica_concatenated[:, :2].T, ax=axes[1], cbar=False, alpha=0.1)\n",
    "axes[1].scatter(*cluster.clustercenters[:, :2].T, s=15, c='C1')\n",
    "axes[1].set_xlabel('IC 1')\n",
    "axes[1].set_ylabel('IC 2')\n",
    "pyemma.plots.plot_implied_timescales(its, ax=axes[2], units='ps')\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "solution2": "hidden"
   },
   "source": [
    "###### Solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "solution2": "hidden"
   },
   "outputs": [],
   "source": [
    "feat = pyemma.coordinates.featurizer(pdb)\n",
    "pairs = feat.pairs(feat.select_Heavy())\n",
    "feat.add_distances(pairs, periodic=False)\n",
    "data = pyemma.coordinates.load(files, features=feat)\n",
    "\n",
    "tica = pyemma.coordinates.tica(data, lag=3)\n",
    "tica_concatenated = np.concatenate(tica.get_output())\n",
    "\n",
    "cluster = pyemma.coordinates.cluster_kmeans(tica, k=75, 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(tica_concatenated, ax=axes[0])\n",
    "pyemma.plots.plot_density(*tica_concatenated[:, :2].T, ax=axes[1], cbar=False, alpha=0.1)\n",
    "axes[1].scatter(*cluster.clustercenters[:, :2].T, s=15, c='C1')\n",
    "axes[1].set_xlabel('IC 1')\n",
    "axes[1].set_ylabel('IC 2')\n",
    "pyemma.plots.plot_implied_timescales(its, ax=axes[2], units='ps')\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Exercise 2\n",
    "Let's see if your discretized data is suitable to converge five slow implied timescales using a bayesian HMM."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "solution2": "hidden",
    "solution2_first": true
   },
   "outputs": [],
   "source": [
    "pyemma.plots.plot_implied_timescales #FIXME"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "solution2": "hidden"
   },
   "source": [
    "###### Solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "solution2": "hidden"
   },
   "outputs": [],
   "source": [
    "pyemma.plots.plot_implied_timescales(\n",
    "    pyemma.msm.timescales_hmsm(\n",
    "        cluster.dtrajs, 6, lags=[1, 2, 3, 4], errors='bayes'), units='ps');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Exercise 3\n",
    "Estimate a bayesian HMM and perform a CK test."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "solution2": "hidden",
    "solution2_first": true
   },
   "outputs": [],
   "source": [
    "hmm = #FIXME\n",
    "pyemma.plots. #FIXME"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "solution2": "hidden"
   },
   "source": [
    "###### Solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "solution2": "hidden"
   },
   "outputs": [],
   "source": [
    "hmm = pyemma.msm.bayesian_hidden_markov_model(cluster.dtrajs, 6, lag=1, dt_traj='1 ps')\n",
    "pyemma.plots.plot_cktest(hmm.cktest(mlags=2), units='ps');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Exercise 4\n",
    "Now that you have a model, be creative and visualize the metastable regions in your projected space."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "solution2": "hidden",
    "solution2_first": true
   },
   "outputs": [],
   "source": [
    "#FIXME"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "solution2": "hidden"
   },
   "source": [
    "###### Solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "solution2": "hidden"
   },
   "outputs": [],
   "source": [
    "def draw_panel(ax, i, j):\n",
    "    _, _, misc = pyemma.plots.plot_state_map(\n",
    "        *tica_concatenated[:, [i, j]].T,\n",
    "        hmm.metastable_assignments[dtrajs_concatenated],\n",
    "        ax=ax)\n",
    "    ax.set_xlabel('IC {}'.format(i + 1))\n",
    "    ax.set_ylabel('IC {}'.format(j + 1))\n",
    "    misc['cbar'].set_ticklabels(range(1, hmm.nstates + 1))\n",
    "\n",
    "fig, axes = plt.subplots(2, 2, figsize=(10, 8))\n",
    "draw_panel(axes[0, 0], 0, 2)\n",
    "draw_panel(axes[0, 1], 1, 2)\n",
    "draw_panel(axes[1, 0], 0, 1)\n",
    "axes[1, 1].set_axis_off()\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Case 3: another molecular dynamics data set (pentapeptide)\n",
    "\n",
    "Let us now repeat the analysis of our pentapeptide using an HMM.\n",
    "\n",
    "We fetch the pentapeptide data set and prepare the discrete trajectories as explained in [Notebook 00 ➜ 📓](00-pentapeptide-showcase.ipynb).\n",
    "There, we already learned that $5$ metastable states are a good choice for our model.\n",
    "According to our implied timescales plot, we can resolve four processes for up to $2.5$ ns in our data. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pdb = mdshare.fetch('pentapeptide-impl-solv.pdb', working_directory='data')\n",
    "files = mdshare.fetch('pentapeptide-*-500ns-impl-solv.xtc', working_directory='data')\n",
    "\n",
    "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",
    "\n",
    "tica = pyemma.coordinates.tica(torsions_data, lag=5)\n",
    "tica_output = tica.get_output()\n",
    "tica_concatenated = np.concatenate(tica_output)\n",
    "\n",
    "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": [
    "#### Exercise 5\n",
    "\n",
    "Please go ahead and estimate an HMM with $5$ hidden states by probing for timescales convergence for up to $2.5$ ns with Bayesian error estimation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "solution2": "hidden",
    "solution2_first": true
   },
   "outputs": [],
   "source": [
    "nstates = #FIXME\n",
    "its_hmm = #FIXME\n",
    "pyemma.plots.plot_implied_timescales(its_hmm, units='ns', dt=0.1);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "solution2": "hidden"
   },
   "source": [
    "###### Solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "solution2": "hidden"
   },
   "outputs": [],
   "source": [
    "nstates = 5\n",
    "its_hmm = pyemma.msm.timescales_hmsm(cluster.dtrajs, nstates, lags=25, errors='bayes')\n",
    "pyemma.plots.plot_implied_timescales(its_hmm, units='ns', dt=0.1);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We note that the ITS are converged within the error from $0.1$ ns, earlier than regular MSMs due to the mentioned compensation of discretization error. \n",
    "\n",
    "\n",
    "#### Exercise 6\n",
    "\n",
    "Estimate a Bayesian HMM at this lag time and conduct the CK-validation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "solution2": "hidden",
    "solution2_first": true
   },
   "outputs": [],
   "source": [
    "hmm = #FIXME\n",
    "\n",
    "cktest = hmm.cktest(mlags=6)\n",
    "pyemma.plots.plot_cktest(cktest, dt=0.1, units='ns');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "solution2": "hidden"
   },
   "source": [
    "###### Solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "solution2": "hidden"
   },
   "outputs": [],
   "source": [
    "hmm = pyemma.msm.bayesian_hidden_markov_model(cluster.dtrajs, nstates, lag=1, dt_traj='0.1 ns')\n",
    "\n",
    "cktest = hmm.cktest(mlags=6)\n",
    "pyemma.plots.plot_cktest(cktest, dt=0.1, units='ns');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Exercise 7\n",
    "\n",
    "Plot the state map of the hidden states:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "solution2": "hidden",
    "solution2_first": true
   },
   "outputs": [],
   "source": [
    "metastable_traj = #FIXME\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(10, 7))\n",
    "_, _, misc = pyemma.plots.plot_state_map(\n",
    "    *tica_concatenated[:, [0, 1]].T, metastable_traj, ax=ax)\n",
    "\n",
    "misc['cbar'].set_ticklabels(range(1, hmm.nstates + 1))\n",
    "ax.set_xlabel('IC 1')\n",
    "ax.set_ylabel('IC 2')\n",
    "\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "solution2": "hidden"
   },
   "source": [
    "###### Solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "solution2": "hidden"
   },
   "outputs": [],
   "source": [
    "metastable_traj = hmm.metastable_assignments[dtrajs_concatenated]\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(10, 7))\n",
    "_, _, misc = pyemma.plots.plot_state_map(\n",
    "    *tica_concatenated[:, [0, 1]].T, metastable_traj, ax=ax)\n",
    "\n",
    "misc['cbar'].set_ticklabels(range(1, hmm.nstates + 1))\n",
    "ax.set_xlabel('IC 1')\n",
    "ax.set_ylabel('IC 2')\n",
    "\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We find state definitions that are similar than the PCCA++ ones but not the same.\n",
    "Especially, the largest basin in the 2D TICA projection has become separated into different hidden states.\n",
    "The reason might be that several metastable states are projected onto the same region by TICA and,\n",
    "after clustering, could not be resolved by a regular MSM anymore.\n",
    "\n",
    "Another benefit of HMMs is that transition matrices between the hidden states are computed as part of the algorithm.\n",
    "Any property derived from it, e.g. the stationary distribution, can directly be computed.\n",
    "\n",
    "#### Exercise 8:\n",
    "\n",
    "Visualize the transition probabilities and stationary probabilities.\n",
    "It might be useful to check out `pyemma.plots.plot_markov_model()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "solution2": "hidden",
    "solution2_first": true
   },
   "outputs": [],
   "source": [
    "#FIXME"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "solution2": "hidden"
   },
   "source": [
    "###### Solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "solution2": "hidden"
   },
   "outputs": [],
   "source": [
    "pyemma.plots.plot_markov_model(hmm);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Exercise 9:\n",
    "\n",
    "In the last step, save some (e.g. $50$) representative structures of the pentapeptide to disk.\n",
    "You might want to use your favorite molecular dynamics viewer and have a look at the structures\n",
    "(it might be necessary to align them, though)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "solution2": "hidden",
    "solution2_first": true
   },
   "outputs": [],
   "source": [
    "hmm_samples = #FIXME\n",
    "\n",
    "data_source = pyemma.coordinates.source(files, features=torsions_feat)\n",
    "\n",
    "pyemma.coordinates.save_trajs(\n",
    "    data_source,\n",
    "    hmm_samples,\n",
    "    outfiles=['./data/hmm_{}.pdb'.format(n + 1)\n",
    "              for n in range(hmm.nstates)])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "solution2": "hidden"
   },
   "source": [
    "###### Solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "solution2": "hidden"
   },
   "outputs": [],
   "source": [
    "hmm_samples = hmm.sample_by_observation_probabilities(50)\n",
    "\n",
    "data_source = pyemma.coordinates.source(files, features=torsions_feat)\n",
    "\n",
    "pyemma.coordinates.save_trajs(\n",
    "    data_source,\n",
    "    hmm_samples,\n",
    "    outfiles=['./data/hmm_{}.pdb'.format(n + 1)\n",
    "              for n in range(hmm.nstates)])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Wrapping up\n",
    "\n",
    "In this notebook, we have learned how to use a hidden Markov state model (HMM) and how they differ from an MSM.\n",
    "In detail, we have used\n",
    "- `pyemma.msm.timescales_hmsm()` function to obtain an implied timescale object for HMMs,\n",
    "- `pyemma.msm.estimate_hidden_markov_model()` to estimate a regular HMM,\n",
    "- `pyemma.msm.bayesian_hidden_markov_model()` to estimate a Bayesian HMM, \n",
    "- the `metastable_assignments` attribute of an HMM object to access the metastable membership of discrete states, \n",
    "- the `hidden_state_probabilities` attribute to assess probabilities of hidden states over time, and\n",
    "- the `hidden_state_trajectories` attribute that extracts the most likely trajectory in hidden state space."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## References\n",
    "\n",
    "<a id=\"cite-hmm-baum-welch-alg\"/><sup><a href=#ref-1>[^]</a></sup>Leonard E. Baum and Ted Petrie and George Soules and Norman Weiss. 1970. _A Maximization Technique Occurring in the Statistical Analysis of Probabilistic Functions of Markov Chains_. [URL](http://www.jstor.org/stable/2239727)\n",
    "\n",
    "<a id=\"cite-noe-proj-hid-msm\"/><sup><a href=#ref-2>[^]</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",
    "<a id=\"cite-hmm-tutorial\"/><sup><a href=#ref-3>[^]</a></sup>L.R. Rabiner. 1989. _A tutorial on hidden Markov models and selected applications in speech recognition_. [URL](https://doi.org/10.1109/5.18626)\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
}