Attention_and_coarse_graining.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Example notebook to build an hierarchical deep MSM\n",
    "\n",
    "This notebook aims to be a template for users trying to build a deep reversible Markov State model with coarse-graining layers and an additional attention mechanism in order to find important residues for the kinetics.\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."
   ]
  },
  {
   "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=25\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_sizes = [4,3,2]\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 = 20*2*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.2\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, Mask, pred_batchwise, plot_mask, get_its, get_ck, plot_cg\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",
    "mask = Mask(data[0].shape[1],mask_const, mask_depth, mask_width, patchsize, fac=factor_att, noise=regularizer_noise, device=device)\n",
    "\n",
    "lobe = nn.Sequential(\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_sizes[0]),\n",
    "    nn.Softmax(dim=1)  # obtain fuzzy probability distribution over output states\n",
    ")\n",
    "lobe_vampnet = nn.Sequential(\n",
    "    normalizer,\n",
    "    lobe  # obtain fuzzy probability distribution over output states\n",
    ")\n",
    "lobe_msm = nn.Sequential(\n",
    "    lobe)\n",
    "lobe_mask = nn.Sequential(\n",
    "    normalizer,\n",
    "    mask)\n",
    "lobe_vampnet_mask = nn.Sequential(\n",
    "    lobe_mask,\n",
    "    lobe)\n",
    "from copy import deepcopy\n",
    "lobe_timelagged = deepcopy(lobe).to(device=device)\n",
    "lobe = lobe.to(device=device)\n",
    "lobe_vampnet.to(device=device)\n",
    "lobe_msm.to(device=device)\n",
    "lobe_vampnet_mask.to(device=device)\n",
    "lobe_mask.to(device=device)\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_vampnet, learning_rate=5e-4, device=device) # for pretraining the VAMPnet without mask\n",
    "vampnet_mask = VAMPNet(lobe=lobe_vampnet_mask, learning_rate=5e-4, device=device)\n",
    "deepmsm = DeepMSM(lobe=lobe, output_dim=output_sizes[0], coarse_grain=output_sizes[1:], mask=lobe_mask, 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 = False\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": [
    "### Plot the mask before training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_mask(data=data[0], lobe=lobe_vampnet, mask=mask, mask_const=mask_const, device=device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train the vampnet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = vampnet.fit(loader_train, n_epochs=30,\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": [
    "model = vampnet_mask.fit(loader_train, n_epochs=30,\n",
    "                    validation_loader=loader_val, progress=tqdm).fetch_model()\n",
    "plt.loglog(*vampnet_mask.train_scores.T, label='training')\n",
    "plt.loglog(*vampnet_mask.validation_scores.T, label='validation')\n",
    "plt.xlabel('step')\n",
    "plt.ylabel('score')\n",
    "plt.legend();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot the mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_mask(data=data[0], lobe=lobe_vampnet_mask, mask=mask, mask_const=mask_const, device=device)"
   ]
  },
  {
   "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_mask.lobe.state_dict()\n",
    "vampnet_mask.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=50, validation_loader=loader_val, rel=0.001, reset_u=False, 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": [
    "### Plot the mask of the trained deepMSM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_mask(data=data[0], lobe=lobe_msm, mask=mask, mask_const=mask_const, device=device, vmax=0.25)"
   ]
  },
  {
   "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_sizes[0], output_sizes[0]))\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(2):\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 \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_sizes[0]-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], 0.3/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",
    "output_size = output_sizes[0]\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([20,60])\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));\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": [
    "model_msm_final.get_transition_matrix(test_data.dataset.data, test_data.dataset.data_lagged)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Estimate coarse-grain model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialize the coarse-graining layer with pcca\n",
    "deepmsm.reset_cg(idx=0, lr=0.1)\n",
    "deepmsm.initialize_cg_layer(idx=0, data_loader=loader_train, factor=1.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# train the first coarse-graining layer\n",
    "model_cg_1 = deepmsm.fit_cg(loader_train, n_epochs=3000, validation_loader=loader_val, train_mode='single', idx=0).fetch_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_cg_1.timescales_cg(test_data.dataset.data, test_data.dataset.data_lagged, tau=tau_msm, idx=0)*fac"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot the learned coarse-graining\n",
    "plot_cg(deepmsm.cg_list[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialize the second coarse-graining layer\n",
    "deepmsm.reset_cg(idx=1, lr=0.1)\n",
    "deepmsm.initialize_cg_layer(idx=1, data_loader=loader_train, factor=1.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# train the second coarse-graining layer\n",
    "model_cg_2 = deepmsm.fit_cg(loader_train, n_epochs=3000, validation_loader=loader_val, train_mode='single', idx=1).fetch_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_cg_2.timescales_cg(test_data.dataset.data, test_data.dataset.data_lagged, tau=tau_msm, idx=1)*fac"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot the learned coarse-graining\n",
    "plot_cg(deepmsm.cg_list[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# train for all respresentations at the same time\n",
    "model_cg_all = deepmsm.fit_cg(loader_train, n_epochs=1000, validation_loader=loader_val, train_mode='all').fetch_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(model_msm_final.timescales(test_data.dataset.data, test_data.dataset.data_lagged, tau_msm)*fac)\n",
    "print(model_cg_1.timescales_cg(test_data.dataset.data, test_data.dataset.data_lagged, tau=tau_msm, idx=0)*fac)\n",
    "print(model_cg_2.timescales_cg(test_data.dataset.data, test_data.dataset.data_lagged, tau=tau_msm, idx=1)*fac)"
   ]
  }
 ],
 "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
}