Observable_training.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Example notebook to build a deep MSM with observables\n",
    "\n",
    "This notebook aims to be a template for users trying to build a deep reversible Markov State model with additional experimental observables.\n",
    "We simulate the situation by firstly estimate observables on the whole data set. These values will represent the true \"experimental\" observables. Afterwards, we will disturb our data in order to mimic the situation that the simulation has a systematic bias, e.g. through the force field. We will estimate again a deep reversible MSM with the additional observables in the hope to recover the kinetics of the undisturbed data.\n",
    "\n",
    "The code is based on the package deeptime, where this code should soon be integrated. This should be seen as a \n",
    "pre-alpha version.\n",
    "\n",
    "In case you have real experimental data available, you must skip the part with the data disturbance. However, it seems recommendable to train a baseline model without the additional information to check if the performance improves."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "# import mdshare  # for trajectory data\n",
    "\n",
    "from tqdm.notebook import tqdm  # progress bar"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if torch.cuda.is_available():\n",
    "    device = torch.device(\"cuda\")\n",
    "    torch.backends.cudnn.benchmark = True\n",
    "else:\n",
    "    device = torch.device(\"cpu\")\n",
    "torch.set_num_threads(12)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Example how to load your data (of course you can change the features to your likings)\n",
    "def load_trajectories(files, pdb):\n",
    "    '''\n",
    "    Loads all the trajectories specified in files and directly estimates the residue_mindist features \n",
    "    used in the paper.\n",
    "    You can change the features: see http://www.emma-project.org/latest/api/generated/pyemma.coordinates.featurizer.html#\n",
    "    ------\n",
    "    Inputs:\n",
    "    files: list of strings. List of locations of all trajectories which should be loaded.\n",
    "    pdb: string. location of the corresponding pdb file.\n",
    "    \n",
    "    Returns:\n",
    "    data: list of np.array. List of the residue_mindist features for all trajectories specified in files.\n",
    "            If only one trajectory is supplied, it returns directly the np.array.\n",
    "    '''\n",
    "    \n",
    "    import pyemma\n",
    "    feat = pyemma.coordinates.featurizer(pdb)\n",
    "    feat.add_residue_mindist(residue_pairs='all', scheme='closest-heavy', ignore_nonprotein=True, threshold=None, periodic=True)\n",
    "    data = pyemma.coordinates.load(files, features=feat)\n",
    "    \n",
    "    return data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the data\n",
    "# with the example code from above\n",
    "files = ['/path/to/file1', 'path/to/file2']\n",
    "pdb = '/path/to/pdb.pdb'\n",
    "output_all_files = load_trajectories(files, pdb)\n",
    "# you can then save the processed data, to directly load the interesting features\n",
    "# np.save('/path/to/save', output_all_files)\n",
    "# output_all_files = np.load('/path/to/save.npy')\n",
    "# output_all_files = np.load('/path/to/save.npy') # need to specify where your data lies. \n",
    "# You can use the lines above to prepare your own data\n",
    "# output_all_files = np.load('/srv/public/andreas/data/desres/2f4k/villin_skip1.npy') # this line is for checking\n",
    "traj_whole = output_all_files\n",
    "\n",
    "traj_data_points, input_size = traj_whole[0].shape\n",
    "# Skip data to make the data less correlated\n",
    "skip=1\n",
    "data = [traj_whole[0][::skip]]\n",
    "\n",
    "n_residues = int(-1/2 + np.sqrt(1/4+input_size*2) + 3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Hyperparameter definitions, should be adapted for specific problems\n",
    "\n",
    "# number of output nodes/states of the MSM or Koopman model, therefore also nodes of chi\n",
    "# The list defines how the output will be coarse grained from first to last entry\n",
    "output_size = 3\n",
    "\n",
    "# Tau, how much is the timeshift of the two datasets in the default training\n",
    "# tau_chi for pretraining the vampnet usually smaller than the tau for the deepMSM\n",
    "tau = 50*25//skip # 5, 20\n",
    "tau_chi = 25//skip\n",
    "\n",
    "# Batch size for Stochastic Gradient descent\n",
    "batch_size = 512\n",
    "# Larger batch size for fine tuning weights at the end of training\n",
    "batch_size_large = 20000\n",
    "\n",
    "# Which trajectory points percentage is used as training, validation, and rest for test\n",
    "valid_ratio = 0.3\n",
    "test_ratio = 0.3\n",
    "\n",
    "# How many hidden layers the network chi has\n",
    "network_depth = 4\n",
    "\n",
    "# Width of every layer of chi\n",
    "layer_width = 30\n",
    "\n",
    "# Mask hyperparameter\n",
    "mask_const=False # if the trained attention mask is constant over time\n",
    "patchsize=4 # size of the sliding window\n",
    "mask_depth=4 # if time dependent how many hidden layers has the attention network\n",
    "mask_width=30 # the width of the attention hidden layers\n",
    "factor_att=True # if to use a factor which scales the input on average back to input\n",
    "regularizer_noise = 1.0 # noise to regularize\n",
    "\n",
    "# Learning rate used for the ADAM optimizer\n",
    "learning_rate = 5e-4\n",
    "\n",
    "# create a list with the number of nodes for each layer\n",
    "nodes = [layer_width]*network_depth\n",
    "\n",
    "# epsilon for numerical inversion of correlation matrices\n",
    "epsilon = np.array(1e-7).astype('float32')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Split the data into train, validation, and test set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from deeptime.util.data import TrajectoryDataset\n",
    "\n",
    "dataset = TrajectoryDataset(lagtime=tau_chi, trajectory=data[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_val = int(len(dataset)*valid_ratio)\n",
    "n_test = int(len(dataset)*test_ratio)\n",
    "train_data, val_data, test_data = torch.utils.data.random_split(dataset, [len(dataset) - n_val - n_test, n_val, n_test])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define the structure of the VAMPnet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from helper import Mean_std_layer, pred_batchwise, get_its, get_ck, estimate_mu\n",
    "\n",
    "normalizer = Mean_std_layer(input_size, mean=torch.Tensor(train_data.dataset.data.mean(0)),\n",
    "                           std=torch.Tensor(train_data.dataset.data.std(0)))\n",
    "\n",
    "lobe = nn.Sequential(\n",
    "    normalizer,\n",
    "    nn.Linear(data[0].shape[1], layer_width), nn.ELU(),\n",
    "    nn.Linear(layer_width, layer_width), nn.ELU(),\n",
    "    nn.Linear(layer_width, layer_width), nn.ELU(),\n",
    "    nn.Linear(layer_width, layer_width), nn.ELU(),\n",
    "    nn.Linear(layer_width, layer_width), nn.ELU(),\n",
    "    nn.Linear(layer_width, output_size),\n",
    "    nn.Softmax(dim=1)  # obtain fuzzy probability distribution over output states\n",
    ")\n",
    "from copy import deepcopy\n",
    "lobe_timelagged = deepcopy(lobe).to(device=device)\n",
    "lobe = lobe.to(device=device)\n",
    "\n",
    "print(lobe)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define the estimators"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from deeptime.decomposition.deep import VAMPNet\n",
    "from deepmsm import DeepMSM\n",
    "\n",
    "vampnet = VAMPNet(lobe=lobe, learning_rate=5e-4, device=device) # for pretraining the VAMPnet without mask\n",
    "deepmsm = DeepMSM(lobe=lobe, output_dim=output_size, learning_rate=5e-4, device=device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create DataLoader for validation and training data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import DataLoader\n",
    "\n",
    "loader_train = DataLoader(train_data, batch_size=batch_size, shuffle=True)\n",
    "loader_val = DataLoader(val_data, batch_size=len(val_data), shuffle=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Can be used to record the training performance with tensorboard\n",
    "# it is not necessary for training or using the methods\n",
    "# if you do not wish to install the additional package just leave the flag to false!\n",
    "tensorboard_installed = True\n",
    "if tensorboard_installed:\n",
    "    from torch.utils.tensorboard import SummaryWriter\n",
    "    writer = SummaryWriter()\n",
    "    input_model, _ = next(iter(loader_train))\n",
    "    writer.add_graph(lobe, input_to_model=input_model.to(device))\n",
    "else:\n",
    "    writer=None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train the vampnet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = vampnet.fit(loader_train, n_epochs=10,\n",
    "                    validation_loader=loader_val, progress=tqdm).fetch_model()\n",
    "plt.loglog(*vampnet.train_scores.T, label='training')\n",
    "plt.loglog(*vampnet.validation_scores.T, label='validation')\n",
    "plt.xlabel('step')\n",
    "plt.ylabel('score')\n",
    "plt.legend();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "state_probabilities = model.transform(data[0])\n",
    "for ix, (mini, maxi) in enumerate(zip(np.min(state_probabilities, axis=0),\n",
    "                                      np.max(state_probabilities, axis=0))):\n",
    "    print(f\"State {ix+1}: [{mini}, {maxi}]\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Extract the parameters of the trained vampnet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "state_dict_vampnet = vampnet.lobe.state_dict()\n",
    "vampnet.lobe.load_state_dict(state_dict_vampnet)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train for the deepMSM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# train only for the matrix S\n",
    "deepmsm.fit(loader_train, n_epochs=100, validation_loader=loader_val, train_mode='s', tb_writer=writer)\n",
    "plt.loglog(*deepmsm.train_scores.T, label='training')\n",
    "plt.loglog(*deepmsm.validation_scores.T, label='validation')\n",
    "plt.xlabel('step')\n",
    "plt.ylabel('score')\n",
    "plt.legend();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train for S and u\n",
    "deepmsm.fit(loader_train, n_epochs=1000, validation_loader=loader_val, train_mode='us', tb_writer=writer)\n",
    "plt.loglog(*deepmsm.train_scores.T, label='training')\n",
    "plt.loglog(*deepmsm.validation_scores.T, label='validation')\n",
    "plt.xlabel('step')\n",
    "plt.ylabel('score')\n",
    "plt.legend();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Train for chi, u, and S in an iterative manner\n",
    "deepmsm.fit_routine(loader_train, n_epochs=5, validation_loader=loader_val, rel=0.001, reset_u=False, \n",
    "                    max_iter=1000, tb_writer=writer)\n",
    "plt.loglog(*deepmsm.train_scores.T, label='training')\n",
    "plt.loglog(*deepmsm.validation_scores.T, label='validation')\n",
    "plt.xlabel('step')\n",
    "plt.ylabel('score')\n",
    "plt.legend();\n",
    "# reset u in order to escape possible local minima\n",
    "deepmsm.fit_routine(loader_train, n_epochs=5, validation_loader=loader_val, rel=0.001, reset_u=True, \n",
    "                    max_iter=1000, tb_writer=writer)\n",
    "plt.loglog(*deepmsm.train_scores.T, label='training')\n",
    "plt.loglog(*deepmsm.validation_scores.T, label='validation')\n",
    "plt.xlabel('step')\n",
    "plt.ylabel('score')\n",
    "plt.legend();\n",
    "deepmsm.fit_routine(loader_train, n_epochs=5, validation_loader=loader_val, rel=0.001, reset_u=False, \n",
    "                    max_iter=1000, tb_writer=writer)\n",
    "plt.loglog(*deepmsm.train_scores.T, label='training')\n",
    "plt.loglog(*deepmsm.validation_scores.T, label='validation')\n",
    "plt.xlabel('step')\n",
    "plt.ylabel('score')\n",
    "plt.legend();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Save the final deepMSM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "deepmsm.save_params('./test_params')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Extract the model and estimate the transition matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_msm = deepmsm.fetch_model()\n",
    "T = model_msm.get_transition_matrix(test_data.dataset.data, test_data.dataset.data_lagged)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.linalg.eigvals(T)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Estimate the transition matrix for different tau values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define tau values\n",
    "steps = 8\n",
    "tau_msm = tau\n",
    "tau_ck = np.arange(1,(steps+1))*tau_msm\n",
    "tau_its = np.concatenate([np.array([1, 3, 5]), tau_ck])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "deepmsm.load_params('./test_params.npz')\n",
    "T_results = np.ones((len(tau_its) ,output_size, output_size))\n",
    "its_all_vamp = []\n",
    "for i, tau_i in enumerate(tau_its):\n",
    "    if i==0: # T for this tau was already evaluated\n",
    "        T_results[i]=T\n",
    "    else:\n",
    "        # split the data with the new tau\n",
    "        dataset = TrajectoryDataset(lagtime=tau_i, trajectory=data[0])\n",
    "        n_val = int(len(dataset)*valid_ratio)\n",
    "        n_test = int(len(dataset)*test_ratio)\n",
    "        train_data, val_data, test_data = torch.utils.data.random_split(dataset, [len(dataset) - n_val - n_test, n_val, n_test])\n",
    "        loader_train = DataLoader(train_data, batch_size=batch_size, shuffle=True)\n",
    "        loader_val = DataLoader(val_data, batch_size=len(val_data), shuffle=False)\n",
    "        # reset u and S to be retrained for the new tau\n",
    "        deepmsm.reset_u_S(loader_train)\n",
    "        # reset the optimizers for u and S\n",
    "        deepmsm.reset_opt_u_S(lr=1)\n",
    "        # train for S\n",
    "        for _ in range(5):\n",
    "            model_msm_i = deepmsm.fit(loader_train, n_epochs=1000, validation_loader=loader_val, train_mode='s').fetch_model()\n",
    "            # train for u and S\n",
    "            model_msm_i = deepmsm.fit(loader_train, n_epochs=100, validation_loader=loader_val, train_mode='us').fetch_model()\n",
    "        # retrieve the transition matrix for the specific tau\n",
    "        T_results[i]  = model_msm_i.get_transition_matrix(test_data.dataset.data, test_data.dataset.data_lagged)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Estimate implied timescales"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "its = get_its(T_results, tau_its, calculate_K=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fac = 200.*skip*1e-6  # factor to change from frames into mikroseconds, adapt for your data!!!\n",
    "# fac = 0.0002\n",
    "\n",
    "plt.figure(figsize=(6,4));\n",
    "\n",
    "label_x = np.array([.1,0.3,1, 2, 5,10,100,1000])/fac # array is in microsecond\n",
    "label_y = np.array([.1,1, 2, 5,10, 100, 1000])/fac\n",
    "# fig = plt.figure(figsize = (8,8))\n",
    "for j in range(0,output_size-1):\n",
    "    plt.semilogy(tau_its, its[::-1][j], lw=5)\n",
    "#     plt.fill_between(tau_its, all_its_vamp_min[::-1][j], all_its_vamp_max[::-1][j], alpha = 0.3)\n",
    "plt.semilogy(tau_its,tau_its, 'k')\n",
    "plt.xlabel('lag [$\\mu$s]', fontsize=26)\n",
    "plt.xticks(label_x, label_x*fac, fontsize=22)\n",
    "plt.ylabel('timescale [$\\mu$s]', fontsize=26)\n",
    "plt.yticks(label_y, np.round(label_y*fac, decimals=1), fontsize=22)\n",
    "plt.fill_between(tau_its,tau_its,0.1,alpha = 0.2,color='k');\n",
    "plt.ylim(0.01/fac, 3/fac)\n",
    "plt.xlim(tau_its[0], 1/fac)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Estimate CK-test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "predicted, estimated = get_ck(T_results[3:], tau_ck)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.gridspec as gridspec\n",
    "fig = plt.figure(figsize = (16,16))\n",
    "gs1 = gridspec.GridSpec(output_size, output_size)\n",
    "gs1.update(wspace=0.1, hspace=0.05)\n",
    "states = output_size\n",
    "for index_i in range(states):\n",
    "    for index_j in range(states):\n",
    "        ax = plt.subplot(gs1[index_i*output_size+index_j])\n",
    "        ax.plot(tau_ck, predicted[index_i, index_j], color='b', lw=4)\n",
    "        ax.plot(tau_ck, estimated[index_i, index_j], color = 'r', lw=4, linestyle = '--')\n",
    "#         ax.fill_between(tau_ck,lx_min[index_i, index_j],lx_max[index_i, index_j], alpha = 0.25 )\n",
    "#         ax.errorbar(tau_ck, rx_mean[index_i, index_j], yerr= np.array([rx_mean[index_i][index_j]-rx_min[index_i][index_j], rx_max[index_i][index_j]-rx_mean[index_i][index_j]]), color = 'r', lw=4, linestyle = '--')\n",
    "        title = str(index_i+1)+ '->' +str(index_j+1)\n",
    "        \n",
    "        ax.text(.75,.8, title,\n",
    "            horizontalalignment='center',\n",
    "            transform=ax.transAxes,  fontdict = {'size':26})\n",
    "    \n",
    "        ax.set_ylim((-0.1,1.1));\n",
    "        ax.set_xlim((0, tau_ck[-1]+5));\n",
    "        \n",
    "        if (index_j == 0):\n",
    "            ax.axes.get_yaxis().set_ticks([0, 1])\n",
    "            ax.yaxis.set_tick_params(labelsize=32)\n",
    "        \n",
    "        else:\n",
    "            ax.axes.get_yaxis().set_ticks([])\n",
    "        \n",
    "        if (index_i == output_size -1):\n",
    "            \n",
    "            xticks = np.array([2000,6000])\n",
    "            float_formatter = lambda x: np.array([(\"%.1f\" % y if y > 0.001 else \"0\") for y in x])\n",
    "            \n",
    "            ax.xaxis.set_ticks(xticks);\n",
    "            ax.xaxis.set_ticklabels(((xticks*fac*1000).astype('int')/1000));\n",
    "            ax.xaxis.set_tick_params(labelsize=32)\n",
    "        else:\n",
    "            ax.axes.get_xaxis().set_ticks([])\n",
    "            \n",
    "        if (index_i == output_size - 1 and index_j == output_size - 4):\n",
    "            ax.text(2.16, -0.4, \"[$\\mu$s]\",\n",
    "                horizontalalignment='center',\n",
    "                transform=ax.transAxes,  fontdict = {'size':28})\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Retrain the model for tau_msm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "deepmsm.load_params('./test_params.npz')\n",
    "dataset = TrajectoryDataset(lagtime=tau_msm, trajectory=data[0])\n",
    "n_val = int(len(dataset)*valid_ratio)\n",
    "n_test = int(len(dataset)*test_ratio)\n",
    "train_data, val_data, test_data = torch.utils.data.random_split(dataset, [len(dataset) - n_val - n_test, n_val, n_test])\n",
    "loader_train = DataLoader(train_data, batch_size=batch_size, shuffle=True)\n",
    "loader_val = DataLoader(val_data, batch_size=len(val_data), shuffle=False)\n",
    "# reset u and S to be retrained for the new tau\n",
    "deepmsm.reset_u_S(loader_train)\n",
    "# reset the optimizers for u and S\n",
    "deepmsm.reset_opt_u_S(lr=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for _ in range(5):\n",
    "    model_msm_final = deepmsm.fit(loader_train, n_epochs=1000, validation_loader=loader_val, train_mode='s', tb_writer=writer).fetch_model()\n",
    "    # train for u and S\n",
    "    model_msm_final = deepmsm.fit(loader_train, n_epochs=100, validation_loader=loader_val, train_mode='us', tb_writer=writer).fetch_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_msm_final.timescales(test_data.dataset.data, test_data.dataset.data_lagged, tau_msm)*fac"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "T_true = model_msm_final.get_transition_matrix(test_data.dataset.data, test_data.dataset.data_lagged)\n",
    "all_eigval_true = np.sort(np.linalg.eigvals(T_true))[:2]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### We have now trained our reference model which uses the true data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define observable \n",
    "Here we define an observable as a specific contact being formed or not. Firstly, we will estimate all observables which we want to use as the true \"experimental\" ones."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from helper import TimeLaggedDatasetObs\n",
    "from deeptime.markov.tools.analysis import mfpt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# return the model\n",
    "state_dict_true = deepmsm.state_dict()\n",
    "# Define which contact you wanna use\n",
    "contact_obs = 41 # index in the input array of the contact we use as microscopic observable\n",
    "distances = - np.log(traj_whole[0][:,contact_obs])\n",
    "contacts = (distances <0.45).astype('int')\n",
    "\n",
    "obs_values = np.array([contacts, (contacts-1)*-1]).T\n",
    "chi_true = model_msm_final(traj_whole[0])\n",
    "mu_true = model_msm_final.get_mu(traj_whole[0][tau_msm:])\n",
    "states_mu_true = estimate_mu(mu_true, chi_true, np.arange(tau_msm, chi_true.shape[0]))\n",
    "\n",
    "dataset = TimeLaggedDatasetObs.from_trajectory(lagtime=tau_msm, data=traj_whole[0], \n",
    "                                               data_obs_ev=obs_values, data_obs_ac=obs_values)\n",
    "# we know that the unfolded state is the most probable one, folded second for Villin, adapt for your data!!!\n",
    "sort_id = np.argsort(states_mu_true)\n",
    "index_unfolded = sort_id[-1] \n",
    "index_folded = sort_id[-2]\n",
    "# If we wanna identify the folding process we need to define which states are the ones most extreme to that particular process\n",
    "obs_true = model_msm_final.observables(dataset.data, dataset.data_lagged, dataset.data_obs_ev, dataset.data_obs_ac, state1=[index_unfolded], state2=[index_folded])\n",
    "# extract the true values\n",
    "ev_true = obs_true[0]\n",
    "ac_true = obs_true[1]\n",
    "eigval_true = obs_true[2] \n",
    "# estimate folding and unfolding rates\n",
    "mfpt_fold_true = mfpt(T_true, index_folded, index_unfolded) * tau_msm * fac\n",
    "mfpt_unfold_true = mfpt(T_true, index_unfolded, index_folded) * tau_msm * fac"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Manipulate the data\n",
    "\n",
    "\n",
    "We seek to remove transitions between the folded and unfolded state to mimic a force field which overestimates the energy barrier between these two states.\n",
    "\n",
    "if you have real experimental data skip this part. Always work on your whole simulation data!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We seek to manipulate the folding process, so first we identify the process\n",
    "T = model_msm_final.get_transition_matrix(test_data.dataset.data, test_data.dataset.data_lagged)\n",
    "eigvals, eigvecs = np.linalg.eig(T)\n",
    "sort_id = np.argsort(eigvals)\n",
    "eigvals_sort = eigvals[sort_id]\n",
    "eigvecs_sort = eigvecs[:,sort_id]\n",
    "\n",
    "# Estimate the eigenfunction corresponding to the folding process which is the fastest one for our model\n",
    "eigfunc = chi_true[::skip] @ eigvecs_sort[:,-3]\n",
    "min_eigfunc = eigfunc.min()\n",
    "max_eigfunc = eigfunc.max()\n",
    "\n",
    "# Now find the transitions\n",
    "\n",
    "# Find data points which are close to the folded and unfolded state\n",
    "starting_points = np.where(eigfunc < (min_eigfunc + 0.05))[0]\n",
    "end_points = np.where(eigfunc > (max_eigfunc - 0.05))[0]\n",
    "if starting_points[0] > end_points[0]:\n",
    "    temp = starting_points\n",
    "    starting_points = end_points\n",
    "    end_points = temp\n",
    "transition_forward = []\n",
    "transition_backward = []\n",
    "flag = True\n",
    "counter = 0\n",
    "while flag:\n",
    "    \n",
    "    if counter%2==0: # if forward transition\n",
    "        # find the last frame before changing to end state\n",
    "        last_frame = np.where((starting_points - end_points[0])<0)[0][-1] \n",
    "        # The transition already starts tau_msm before\n",
    "        transition_forward.append(np.arange(starting_points[last_frame]-tau_msm, end_points[0]))\n",
    "        \n",
    "        starting_points = starting_points[last_frame+1:]\n",
    "    else: # backward direction\n",
    "        \n",
    "        # find the last frame before changing to end state\n",
    "        last_frame = np.where((end_points - starting_points[0])<0)[0][-1] \n",
    "        # The transition already starts tau_msm before\n",
    "        transition_backward.append(np.arange(end_points[last_frame]-tau_msm, starting_points[0]))\n",
    "        \n",
    "        end_points = end_points[last_frame+1:]\n",
    "        \n",
    "        \n",
    "    counter +=1\n",
    "    if len(end_points)==0 or len(starting_points)==0:\n",
    "        flag=False\n",
    "# now find the frames which are not part of a transition\n",
    "non_transition = np.arange(eigfunc.shape[0]-tau_msm)\n",
    "non_transition = np.setdiff1d(non_transition, np.concatenate(transition_forward+ transition_forward))        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# we can check our results by plotting the classifications\n",
    "skip_fig = 1\n",
    "plt.plot(non_transition, eigfunc[non_transition][::skip_fig], '.')\n",
    "forwards = np.concatenate(transition_forward, axis=0)\n",
    "plt.plot(forwards[::skip_fig], eigfunc[forwards][::skip_fig], '.')\n",
    "backwards = np.concatenate(transition_backward, axis=0)\n",
    "plt.plot(backwards[::skip_fig], eigfunc[backwards][::skip_fig], '.')\n",
    "plt.xlim(0,20000)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The non_transition frames are definitely in the manipulated data set, we leave these untouched\n",
    "ind_train = []\n",
    "ind_valid = []\n",
    "ind_test = []\n",
    "# we assign them randomly into training/validation/test set\n",
    "non_length = non_transition.shape[0]//3\n",
    "np.random.shuffle(non_transition)\n",
    "ind_train.append(non_transition[:non_length])\n",
    "ind_valid.append(non_transition[non_length:2*non_length])\n",
    "ind_test.append(non_transition[2*non_length:])\n",
    "\n",
    "# now take only a percentage of forward and backward events into the data\n",
    "p_for = 0.25 # percentage of how many forward events end into the data\n",
    "p_back = 0.25 # percentage for the backward event\n",
    "\n",
    "nr_for = int(p_for*len(transition_forward)//3) # Number of transitions in each data set (training, validation, test)\n",
    "print('Number of forward transitions in each data set: {}'.format(nr_for))\n",
    "ind_trajs_temp = np.arange(len(transition_forward))\n",
    "np.random.shuffle(ind_trajs_temp) # shuffle where they end in\n",
    "for i in range(nr_for):\n",
    "    ind_train.append(transition_forward[ind_trajs_temp[i]])\n",
    "\n",
    "    ind_valid.append(transition_forward[ind_trajs_temp[i+nr_for]])\n",
    "\n",
    "    ind_test.append(transition_forward[ind_trajs_temp[i+2*nr_for]])\n",
    "# the same for the unfolding\n",
    "p_back = 0.25\n",
    "\n",
    "nr_back = int(p_back*len(transition_backward)//3)\n",
    "print('Number of backward transitions in each data set: {}'.format(nr_for))\n",
    "ind_trajs_temp = np.arange(len(transition_backward))\n",
    "np.random.shuffle(ind_trajs_temp)\n",
    "for i in range(nr_back):\n",
    "    ind_train.append(transition_backward[ind_trajs_temp[i]])\n",
    "\n",
    "    ind_valid.append(transition_backward[ind_trajs_temp[i+nr_back]])\n",
    "\n",
    "    ind_test.append(transition_backward[ind_trajs_temp[i+2*nr_back]])\n",
    "\n",
    "ind_train = np.concatenate(ind_train)\n",
    "ind_valid = np.concatenate(ind_valid)\n",
    "ind_test = np.concatenate(ind_test)\n",
    "\n",
    "np.random.shuffle(ind_train)\n",
    "np.random.shuffle(ind_valid)\n",
    "np.random.shuffle(ind_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Prepare the data for training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# we have now the frame indexes which will be included into train/validation and test set\n",
    "train_data = TimeLaggedDatasetObs.from_frames(lagtime=tau_msm, data=traj_whole[0], frames=ind_train,\n",
    "                                               data_obs_ev=None, data_obs_ac=None)\n",
    "val_data =  TimeLaggedDatasetObs.from_frames(lagtime=tau_msm, data=traj_whole[0], frames=ind_valid,\n",
    "                                               data_obs_ev=None, data_obs_ac=None) \n",
    "test_data = TimeLaggedDatasetObs.from_frames(lagtime=tau_msm, data=traj_whole[0], frames=ind_test,\n",
    "                                               data_obs_ev=None, data_obs_ac=None)\n",
    "# This is now the data which is manipulated.\n",
    "loader_train = DataLoader(train_data, batch_size=batch_size, shuffle=True)\n",
    "loader_val = DataLoader(val_data, batch_size=len(val_data), shuffle=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train a model on the manipulated data but without including additional information about the observables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Since we start again with the result of the VAMPnet, we need to load that first\n",
    "vampnet.lobe.load_state_dict(state_dict_vampnet)\n",
    "# and train it on the new data\n",
    "model = vampnet.fit(loader_train, n_epochs=10,\n",
    "                    validation_loader=loader_val, progress=tqdm).fetch_model()\n",
    "plt.loglog(*vampnet.train_scores.T, label='training')\n",
    "plt.loglog(*vampnet.validation_scores.T, label='validation')\n",
    "plt.xlabel('step')\n",
    "plt.ylabel('score')\n",
    "plt.legend();\n",
    "plt.show()\n",
    "state_dict_vampnet = vampnet.lobe.state_dict()\n",
    "# reset u and S to be retrained for the new tau\n",
    "deepmsm.set_rev_var(loader_train)\n",
    "# deepmsm.reset_u_S_wo()\n",
    "# deepmsm.reset_u_S(loader_train)\n",
    "# reset the optimizers for u and S\n",
    "# deepmsm.reset_opt_u_S(lr=1)\n",
    "deepmsm.reset_opt_all(lr=1)\n",
    "deepmsm.reset_scores()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# train only for the matrix S\n",
    "deepmsm.fit(loader_train, n_epochs=100, validation_loader=loader_val, train_mode='s', tb_writer=writer)\n",
    "plt.loglog(*deepmsm.train_scores.T, label='training')\n",
    "plt.loglog(*deepmsm.validation_scores.T, label='validation')\n",
    "plt.xlabel('step')\n",
    "plt.ylabel('score')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# train for u and S\n",
    "deepmsm.fit(loader_train, n_epochs=1000, validation_loader=loader_val, train_mode='us', tb_writer=writer)\n",
    "plt.loglog(*deepmsm.train_scores.T, label='training')\n",
    "plt.loglog(*deepmsm.validation_scores.T, label='validation')\n",
    "plt.xlabel('step')\n",
    "plt.ylabel('score')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train for chi, u, and S in an iterative manner\n",
    "deepmsm.fit_routine(loader_train, n_epochs=5, validation_loader=loader_val, rel=0.001, reset_u=False, \n",
    "                    max_iter=1000, tb_writer=writer)\n",
    "plt.loglog(*deepmsm.train_scores.T, label='training')\n",
    "plt.loglog(*deepmsm.validation_scores.T, label='validation')\n",
    "plt.xlabel('step')\n",
    "plt.ylabel('score')\n",
    "plt.legend();\n",
    "deepmsm.fit_routine(loader_train, n_epochs=5, validation_loader=loader_val, rel=0.001, reset_u=True, \n",
    "                    max_iter=1000, tb_writer=writer)\n",
    "plt.loglog(*deepmsm.train_scores.T, label='training')\n",
    "plt.loglog(*deepmsm.validation_scores.T, label='validation')\n",
    "plt.xlabel('step')\n",
    "plt.ylabel('score')\n",
    "plt.legend();\n",
    "deepmsm.fit_routine(loader_train, n_epochs=5, validation_loader=loader_val, rel=0.001, reset_u=False, \n",
    "                    max_iter=1000, tb_writer=writer)\n",
    "plt.loglog(*deepmsm.train_scores.T, label='training')\n",
    "plt.loglog(*deepmsm.validation_scores.T, label='validation')\n",
    "plt.xlabel('step')\n",
    "plt.ylabel('score')\n",
    "plt.legend();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create the data with observables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data = TimeLaggedDatasetObs.from_frames(lagtime=tau_msm, data=traj_whole[0], frames=ind_train,\n",
    "                                               data_obs_ev=obs_values, data_obs_ac=obs_values)\n",
    "val_data =  TimeLaggedDatasetObs.from_frames(lagtime=tau_msm, data=traj_whole[0], frames=ind_valid,\n",
    "                                               data_obs_ev=obs_values, data_obs_ac=obs_values) \n",
    "test_data = TimeLaggedDatasetObs.from_frames(lagtime=tau_msm, data=traj_whole[0], frames=ind_test,\n",
    "                                               data_obs_ev=obs_values, data_obs_ac=obs_values)\n",
    "# This is now the data which is manipulated.\n",
    "loader_train = DataLoader(train_data, batch_size=batch_size, shuffle=True)\n",
    "loader_val = DataLoader(val_data, batch_size=len(val_data), shuffle=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Estimate the observables with the model without additional informations of experimental values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "obs_before = model_msm_final.observables(test_data.data, test_data.data_lagged, test_data.data_obs_ev, test_data.data_obs_ac, state1=[index_unfolded], state2=[index_folded])\n",
    "# extract the true values\n",
    "ev_before = obs_before[0]\n",
    "ac_before = obs_before[1]\n",
    "eigval_before = obs_before[2] "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Compare the defined observables\n",
    "The values of before should be different from the true ones, otherwise our data manipulation did not have an effect on that particular observables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ev_true, ev_before"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ac_true, ac_before"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eigval_true, eigval_before"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# check all the eigenvalues of the transition matrix\n",
    "T_before = model_msm_final.get_transition_matrix(test_data.data, test_data.data_lagged)\n",
    "all_eigval_before = np.sort(np.linalg.eigvals(T_before))[:2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check the estimation of the defined states\n",
    "mu_before = model_msm_final.get_mu(test_data.data_lagged)\n",
    "states_mu_before = estimate_mu(mu_before, chi_true[::skip], ind_test+tau_msm)\n",
    "print(states_mu_true, states_mu_before)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# estimate folding and unfolding rates\n",
    "mfpt_fold_before = mfpt(T_before, index_folded, index_unfolded) * tau_msm * fac\n",
    "mfpt_unfold_before = mfpt(T_before, index_unfolded, index_folded) * tau_msm * fac"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(mfpt_fold_before, mfpt_fold_true)\n",
    "print(mfpt_unfold_before, mfpt_unfold_true)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save the weights\n",
    "state_dict_before = deepmsm.state_dict()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Now train with the the observables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Since we start again with the result of the VAMPnet, we need to load that first\n",
    "vampnet.lobe.load_state_dict(state_dict_vampnet)\n",
    "# reset u and S to be retrained for the new tau\n",
    "# deepmsm.set_rev_var(loader_train)\n",
    "deepmsm.reset_u_S_wo()\n",
    "# deepmsm.reset_u_S(loader_train)\n",
    "# reset the optimizers for u and S\n",
    "# deepmsm.reset_opt_u_S(lr=1)\n",
    "deepmsm.reset_opt_all(lr=0.1)\n",
    "deepmsm.reset_scores()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define the regularization parameter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# you can turn them off by switching a value to 0\n",
    "xi_ev = np.array([1.,1.])*10.\n",
    "xi_ac = np.array([1.,1.])*10.\n",
    "xi_its = np.array([1.])*10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "deepmsm.fit_obs(loader_train, 1000, validation_loader=loader_val, train_mode='s', \n",
    "               exp_ev=ev_true, exp_ac=ac_true, exp_its=eigval_true,\n",
    "               xi_ev=xi_ev, xi_ac=xi_ac, xi_its=xi_its,\n",
    "               its_state1=[index_unfolded], its_state2=[index_folded], tb_writer=writer)\n",
    "plt.loglog(*np.abs(deepmsm.train_scores.T), label='training')\n",
    "plt.loglog(*np.abs(deepmsm.validation_scores.T), label='validation')\n",
    "plt.xlabel('step')\n",
    "plt.ylabel('score')\n",
    "plt.legend();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Training loop\n",
    "score_temp = deepmsm.validation_scores[-1,-1]\n",
    "weights_temp = deepmsm.state_dict()\n",
    "for i in range(10):\n",
    "    # Train for everything\n",
    "    deepmsm.fit_obs(loader_train, 5, validation_loader=loader_val, train_mode='all', \n",
    "               exp_ev=ev_true, exp_ac=ac_true, exp_its=eigval_true,\n",
    "               xi_ev=xi_ev, xi_ac=xi_ac, xi_its=xi_its,\n",
    "               its_state1=[index_unfolded], its_state2=[index_folded], tb_writer=writer)\n",
    "    # train only for u and s\n",
    "    deepmsm.fit_obs(loader_train, 100, validation_loader=loader_val, train_mode='us', \n",
    "               exp_ev=ev_true, exp_ac=ac_true, exp_its=eigval_true,\n",
    "               xi_ev=xi_ev, xi_ac=xi_ac, xi_its=xi_its,\n",
    "               its_state1=[index_unfolded], its_state2=[index_folded], tb_writer=writer)\n",
    "    # save the weights if the validation score is better!\n",
    "    if deepmsm.validation_scores[-1,-1]< score_temp:\n",
    "        score_temp = deepmsm.validation_scores[-1,-1]\n",
    "        print('new score: {:.3}'.format(score_temp))\n",
    "        weights_temp = deepmsm.state_dict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "obs_after = model_msm_final.observables(test_data.data, test_data.data_lagged, test_data.data_obs_ev, test_data.data_obs_ac, state1=[index_unfolded], state2=[index_folded])\n",
    "# extract the true values\n",
    "ev_after = obs_after[0]\n",
    "ac_after = obs_after[1]\n",
    "eigval_after = obs_after[2] "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ev_after, ev_true"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ac_after, ac_true"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eigval_after, eigval_true"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# check all the eigenvalues of the transition matrix\n",
    "T_after = model_msm_final.get_transition_matrix(test_data.data, test_data.data_lagged)\n",
    "all_eigval_after = np.sort(np.linalg.eigvals(T_after))[:2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# estimate the stationary distribution of predefined states\n",
    "mu_after = model_msm_final.get_mu(test_data.data_lagged)\n",
    "states_mu_after = estimate_mu(mu_after, chi_true[::skip], ind_test+tau_msm)\n",
    "print(states_mu_true, states_mu_after)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# estimate folding and unfolding rates\n",
    "mfpt_fold_after = mfpt(T_after, index_folded, index_unfolded) * tau_msm * fac\n",
    "mfpt_unfold_after = mfpt(T_after, index_unfolded, index_folded) * tau_msm * fac"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(mfpt_fold_after, mfpt_fold_true)\n",
    "print(mfpt_unfold_true, mfpt_unfold_after)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot the final comparison"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "labels = ['True', '', '']\n",
    "for i in range(2):\n",
    "    plt.hlines(ev_true[i], i-0.25, i+0.25,'k', '--',label=labels[i])\n",
    "plt.plot(ev_after, 'o', ms=10, label='+ Obs')\n",
    "plt.plot(ev_before, 'o', ms=10, label='without')\n",
    "plt.xlabel('Contact', fontsize=18)\n",
    "plt.ylabel('Value [%]', fontsize=18)\n",
    "plt.xticks([0,1], ['Formed', 'Unformed'], fontsize=16)\n",
    "plt.title('Expectation Values', fontsize=18)\n",
    "plt.legend(fontsize=14)\n",
    "plt.show()\n",
    "\n",
    "for i in range(2):\n",
    "    plt.hlines(ac_true[i], i-0.25, i+0.25,'k', '--',label=labels[i])\n",
    "plt.plot(ac_after, 'o', ms=10, label='+ Obs')\n",
    "plt.plot(ac_before, 'o', ms=10, label='without')\n",
    "plt.xlabel('Contact staying', fontsize=18)\n",
    "plt.ylabel('Value [%]', fontsize=18)\n",
    "plt.xticks([0,1], ['Formed', 'Unformed'], fontsize=16)\n",
    "plt.title('Autocorrelation Values', fontsize=18)\n",
    "plt.legend(fontsize=14)\n",
    "plt.show()\n",
    "\n",
    "for i in range(2):\n",
    "    plt.hlines(all_eigval_true[i], i-0.25, i+0.25,'k', '--',label=labels[i])\n",
    "plt.plot(all_eigval_after, 'o', ms=10, label='+ Obs')\n",
    "plt.plot(all_eigval_before, 'o', ms=10, label='without')\n",
    "plt.xlabel('Eigenvalue', fontsize=18)\n",
    "plt.ylabel('Value', fontsize=18)\n",
    "plt.xticks([0,1], ['Folding', 'Misfolding'], fontsize=16)\n",
    "plt.title('Eigenvalue', fontsize=18)\n",
    "plt.legend(fontsize=14)\n",
    "plt.show()\n",
    "\n",
    "for i in range(3):\n",
    "    plt.hlines(np.sort(states_mu_true)[i], i-0.25, i+0.25,'k', '--',label=labels[i])\n",
    "plt.plot(np.sort(states_mu_after), 'o', ms=10, label='+ Obs')\n",
    "plt.plot(np.sort(states_mu_before), 'o', ms=10, label='without')\n",
    "plt.xlabel('State', fontsize=18)\n",
    "plt.ylabel('Probability [%]', fontsize=18)\n",
    "plt.xticks([0,1,2], ['Misfolded', 'Folded', 'Unfolded'], fontsize=16)\n",
    "plt.title('Stationary Distribution', fontsize=18)\n",
    "plt.legend(fontsize=14)\n",
    "plt.show()\n",
    "\n",
    "for i in range(2):\n",
    "    plt.hlines([mfpt_fold_true, mfpt_unfold_true][i], i-0.25, i+0.25,'k', '--',label=labels[i])\n",
    "plt.plot([mfpt_fold_after, mfpt_unfold_after], 'o', ms=10, label='+ Obs')\n",
    "plt.plot([mfpt_fold_before, mfpt_unfold_before], 'o', ms=10, label='without')\n",
    "plt.xlabel('Process', fontsize=18)\n",
    "plt.ylabel('MFPT [\\mu s]', fontsize=18)\n",
    "plt.xticks([0,1], ['Unfolding', 'Folding'], fontsize=16)\n",
    "plt.title('Mean First Passage Time', fontsize=18)\n",
    "plt.legend(fontsize=14)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}