notebooks/02-io-features-hands-on.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data input, featurization and coordinate transforms in PyEMMA\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",
    "## Loading MD example data from our FTP server\n",
    "Ingredients:\n",
    "- Topology file: PDB\n",
    "- Trajectory data: List of .XTC files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from mdshare import fetch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "topfile = fetch('alanine-dipeptide-nowater.pdb', working_directory='data')\n",
    "traj_list = [fetch('alanine-dipeptide-%d-250ns-nowater.xtc' % i, working_directory='data') for i in range(3)]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `fetch` function fetches the data from our servers. **Do not use `mdshare` for your own data!**\n",
    "\n",
    "## Import PyEMMA & friends"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pyemma\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "plt.matplotlib.rcParams.update({'font.size': 16})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Several ways of processing the same data\n",
    "### Backbone torsions\n",
    "- The best possible discription for Ala2\n",
    "- Two dimensions that discribe the full dynamics\n",
    "- A priori known\n",
    "\n",
    "#### Exercise: Define the featurizer and add backbone torsions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "solution2": "hidden",
    "solution2_first": true
   },
   "outputs": [],
   "source": [
    "bbtorsion_feat = # FIXME\n",
    "# FIXME"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "solution2": "hidden"
   },
   "outputs": [],
   "source": [
    "bbtorsion_feat = pyemma.coordinates.featurizer(topfile)\n",
    "bbtorsion_feat.add_backbone_torsions()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Exercise: Load the data into memory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "solution2": "hidden",
    "solution2_first": true
   },
   "outputs": [],
   "source": [
    "bbtorsions = # FIXME"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "solution2": "hidden"
   },
   "outputs": [],
   "source": [
    "bbtorsions = pyemma.coordinates.load(traj_list, bbtorsion_feat)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyemma.plots.plot_free_energy(np.concatenate(bbtorsions)[:, 0], np.concatenate(bbtorsions)[:, 1])\n",
    "plt.xlabel('$\\Phi$ / rad') \n",
    "plt.ylabel('$\\Psi$ / rad');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### heavy atom distances\n",
    "- without prior knowledge usually a good choice\n",
    "- very high dimensional even for this system\n",
    "\n",
    "#### Exercise: define a second featurizer object and add heavy atom distances:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "solution2": "hidden",
    "solution2_first": true
   },
   "outputs": [],
   "source": [
    "heavy_atom_dist_feat =  # FIXME\n",
    "heavy_atom_indices =  # FIXME\n",
    "# FIXME"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "solution2": "hidden"
   },
   "outputs": [],
   "source": [
    "heavy_atom_dist_feat = pyemma.coordinates.featurizer(topfile)\n",
    "heavy_atom_indices = heavy_atom_dist_feat.select_Heavy()\n",
    "\n",
    "heavy_atom_dist_feat.add_distances(heavy_atom_indices, periodic=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(heavy_atom_indices)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "heavy_atom_dist_feat.dimension()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "heavy_atom_distances = pyemma.coordinates.load(traj_list, heavy_atom_dist_feat)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "solution2": "hidden",
    "solution2_first": true
   },
   "source": [
    "#### Exercise: Visualize the heavy atom distances."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "solution2": "hidden"
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(10, 14))\n",
    "pyemma.plots.plot_feature_histograms(np.concatenate(heavy_atom_distances), feature_labels=heavy_atom_dist_feat, ax=ax)\n",
    "ax.set_xlabel('heavy atom distance')\n",
    "ax.set_title('distance histograms per dimension (normalized)');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## VAMP-scoring: Which features are best?\n",
    "We already learned that two dimensions are a good choice for our data. Now, we want to compare different input features with the VAMP-2 score.\n",
    "Please complete the next task at the following lag times:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dim = 2\n",
    "lags = [10, 100, 1000]  # ps"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "solution2": "hidden",
    "solution2_first": true
   },
   "source": [
    "#### Exercise: Perform cross-validated VAMP-scoring for backbone torsions and heavy-atom distances."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "solution2": "hidden"
   },
   "outputs": [],
   "source": [
    "from deeptime.decomposition import TICA, vamp_score_cv\n",
    "\n",
    "fig, axes = plt.subplots(1, 3, figsize=(12, 3), sharey=True)\n",
    "\n",
    "labels = ['backbone\\ntorsions', 'heavy Atom\\ndistances']\n",
    "\n",
    "tica_estimator = TICA(lagtime=lags[0], dim=dim)\n",
    "\n",
    "for ax, lag in zip(axes.flat, lags):\n",
    "    tica_estimator.lagtime = lag\n",
    "    torsions_scores = vamp_score_cv(tica_estimator, trajs=bbtorsions, blocksplit=False, n=3)\n",
    "    scores = [torsions_scores.mean()]\n",
    "    errors = [torsions_scores.std()]\n",
    "    distances_scores = vamp_score_cv(tica_estimator, trajs=heavy_atom_distances, blocksplit=False, n=3)\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$={}ps'.format(lag))\n",
    "\n",
    "axes[0].set_ylabel('VAMP2 score')\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Discussion:\n",
    "Which feature looks best and why?\n",
    "\n",
    "## TICA projection of heavy atom distances\n",
    "#### Exercise: Do a TICA projection of the heavy atom distances"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "solution2": "hidden",
    "solution2_first": true
   },
   "outputs": [],
   "source": [
    "tica = # FIXME"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "solution2": "hidden"
   },
   "outputs": [],
   "source": [
    "tica_estimator = TICA(lagtime=10, var_cutoff=0.95)\n",
    "tica = tica_estimator.fit_fetch(heavy_atom_distances)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tica.output_dimension"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tics = tica.transform(heavy_atom_distances)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "pyemma.plots.plot_free_energy(np.concatenate(tics)[:, 0], np.concatenate(tics)[:, 1])\n",
    "plt.xlabel('TIC 1') \n",
    "plt.ylabel('TIC 2');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Exercise: Perform a PCA projection of heavy atom distances"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "solution2": "hidden",
    "solution2_first": true
   },
   "outputs": [],
   "source": [
    "pca = pyemma.coordinates.pca()  # FIXME"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "solution2": "hidden"
   },
   "outputs": [],
   "source": [
    "pca = pyemma.coordinates.pca(heavy_atom_distances, dim=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pcs = [pca.transform(traj) for traj in heavy_atom_distances]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyemma.plots.plot_free_energy(np.concatenate(pcs)[:, 0], np.concatenate(pcs)[:, 1])\n",
    "plt.xlabel('IC 1') \n",
    "plt.ylabel('IC 2');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1, 3, figsize=(15, 4))\n",
    "t = ['backbone torsions', 'TICs', 'PCs']\n",
    "for n, _y in enumerate([bbtorsions, tics, pcs]):\n",
    "    pyemma.plots.plot_free_energy(np.concatenate(_y)[:, 0], np.concatenate(_y)[:, 1], ax=axes[n], cbar=False)\n",
    "    axes[n].set_title(t[n])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Discussion:\n",
    "What do you think are the differences between the plots in terms of the dynamics they describe?\n",
    "\n",
    "## Different ways of discretizing the output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y = bbtorsions  # if you want, you can change this later and try e.g. the TICA transformed data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Exercise: Perform k-means clustering and plot the cluster centers into the free energy landscape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "solution2": "hidden",
    "solution2_first": true
   },
   "outputs": [],
   "source": [
    "clustering_kmeans =  # FIXME"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "solution2": "hidden"
   },
   "outputs": [],
   "source": [
    "from deeptime.clustering import KMeans\n",
    "from tqdm.notebook import tqdm\n",
    "\n",
    "kmeans_estimator = KMeans(75, max_iter=30, progress=tqdm)\n",
    "stride = 10\n",
    "clustering_kmeans = kmeans_estimator.fit_fetch(np.concatenate(y)[::stride])\n",
    "# different k, stride, max_iter can be used!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "solution2": "hidden",
    "solution2_first": true
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "# FIXME\n",
    "pyemma.plots.plot_free_energy(*np.concatenate(y).T, ax=ax)\n",
    "ax.set_xlabel('$\\Phi$ / rad') \n",
    "ax.set_ylabel('$\\Psi$ / rad');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "solution2": "hidden"
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "ax.plot(*clustering_kmeans.cluster_centers.T, 'ko')\n",
    "pyemma.plots.plot_free_energy(*np.concatenate(y).T, ax=ax)\n",
    "ax.set_xlabel('$\\Phi$ / rad') \n",
    "ax.set_ylabel('$\\Psi$ / rad');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Exercise: Do the same with regular space clustering"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "solution2": "hidden",
    "solution2_first": true
   },
   "outputs": [],
   "source": [
    "clustering_regspace = # FIXME\n",
    "clustering_regspace.n_clusters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "solution2": "hidden"
   },
   "outputs": [],
   "source": [
    "from deeptime.clustering import RegularSpace\n",
    "\n",
    "regspace_estimator = RegularSpace(dmin=0.4)\n",
    "clustering_regspace = regspace_estimator.fit_fetch(np.concatenate(y))\n",
    "clustering_regspace.n_clusters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "ax.plot(*clustering_regspace.cluster_centers.T, 'ko')\n",
    "pyemma.plots.plot_free_energy(*np.concatenate(y).T, ax=ax)\n",
    "ax.set_xlabel('$\\Phi$ / rad') \n",
    "ax.set_ylabel('$\\Psi$ / rad');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Discussion:\n",
    "In your group, discuss the differences between the two clustering algorithms. Which one do you think is better? Which one is faster?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Add-on: A quick MSM estimate to check our work\n",
    "If you are already familiar with Markov state modeling, have a look at the following plots. It tells us which combination of features/projection/clustering conserves the slowest process in the system. Further, we might find that in some cases, MSM implied timescales converge faster than in others."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from deeptime.markov.msm import MaximumLikelihoodMSM\n",
    "from deeptime.util.validation import implied_timescales\n",
    "from deeptime.plots import plot_implied_timescales"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(3, 3, figsize=(15, 13))\n",
    "t = ['backbone torsions', 'TICs', 'PCs']\n",
    "for n, _y in enumerate([bbtorsions, tics, pcs]):\n",
    "    pyemma.plots.plot_free_energy(*np.concatenate(_y).T, ax=axes[0][n], cbar=False)\n",
    "    axes[0][n].set_title(t[n], fontweight='bold')\n",
    "\n",
    "    data = np.concatenate(_y)[::100]\n",
    "    clusterings = [\n",
    "        KMeans(75, max_iter=30).fit(data).fetch_model(),\n",
    "        RegularSpace(dmin=0.4 if n==0 else .4 / (2.2 * n)).fit(data).fetch_model()\n",
    "    ]\n",
    "    for cl_n, cl_obj in enumerate(clusterings):\n",
    "        axes[0][n].plot(*cl_obj.cluster_centers.T, 'ko' if cl_n == 0 else 'rs', alpha=.8)\n",
    "        dtrajs = [cl_obj.transform(traj) for traj in _y]\n",
    "        models = []\n",
    "        for lag in [1, 2, 4, 6, 8]:\n",
    "            models.append(MaximumLikelihoodMSM(lagtime=lag).fit_fetch(dtrajs))\n",
    "        its = implied_timescales(models)\n",
    "        plot_implied_timescales(its, n_its=4, ax=axes[cl_n+1][n])\n",
    "        axes[cl_n+1][n].set_yscale('log')\n",
    "        # its = implied_timescales_msm(dtrajs, lagtimes=[1, 2, 4, 6, 8], nits=4, bayesian=False)\n",
    "        # pyemma.plots.plot_implied_timescales(its, ax=axes[cl_n+1][n])\n",
    "        axes[cl_n+1][n].set_ylim(1e-1, 3e3)\n",
    "        axes[cl_n+1][n].set_ylabel('')\n",
    "axes[1][0].set_ylabel('k-means clustering', fontweight='bold')\n",
    "axes[2][0].set_ylabel('regspace clustering', fontweight='bold')\n",
    "\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "deeptime",
   "language": "python",
   "name": "deeptime"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.10"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}