Toymodel_2Systems.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6882cf61",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch\n",
    "from torch.utils import data\n",
    "import matplotlib.gridspec as gridspec"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b23855bb",
   "metadata": {},
   "source": [
    "### Hyperparameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8089031",
   "metadata": {},
   "outputs": [],
   "source": [
    "tau = 1\n",
    "output_sizes = [3,2]\n",
    "\n",
    "# Batch size for Stochastic Gradient descent\n",
    "batch_size = 10000\n",
    "\n",
    "# How many hidden layers the network chi has\n",
    "network_depth = 3\n",
    "\n",
    "# Width of every layer of chi\n",
    "layer_width = 30\n",
    "\n",
    "# Learning rate used for the ADAM optimizer\n",
    "\n",
    "# create a list with the number of nodes for each layer\n",
    "nodes = [layer_width]*network_depth\n",
    "\n",
    "number_subsystems = len(output_sizes)\n",
    "\n",
    "# How strong the fake subsystem is\n",
    "factor_fake = 0.\n",
    "# How large the noise in the mask for regularization is\n",
    "noise = 1.\n",
    "# Threshold after which the attention weight is set to zero\n",
    "cutoff=0.9\n",
    "# Learning rate\n",
    "learning_rate=0.005\n",
    "# epsilon\n",
    "epsilon=1e-8\n",
    "# score method\n",
    "score_mode='regularize' # one of ('trunc', 'regularize', 'clamp', 'old')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6eea92f",
   "metadata": {},
   "source": [
    "### Create toymodel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7365b8cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "from examples import Toymodel_2Systems"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7deff63d",
   "metadata": {},
   "outputs": [],
   "source": [
    "eps_list = [.025, .125, .05, .1]\n",
    "toymodel = Toymodel_2Systems(eps_list)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11585e56",
   "metadata": {},
   "source": [
    "### Sample hidden and observable trajectory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "511eeeb8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# training data with 100000 steps\n",
    "hidden_state_traj, observable_traj = toymodel.generate_traj(100000)\n",
    "\n",
    "# validation data with 10000 steps\n",
    "hidden_state_traj_valid, observable_traj_valid = toymodel.generate_traj(10000)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d05ab91",
   "metadata": {},
   "source": [
    "### Plot trajectory and true global eigenfunctions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2f05cbc",
   "metadata": {},
   "outputs": [],
   "source": [
    "toymodel.plot_toymodel(hidden_state_traj_valid, observable_traj_valid)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4dac631b",
   "metadata": {},
   "outputs": [],
   "source": [
    "toymodel.plot_eigfunc(hidden_state_traj_valid, observable_traj_valid)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ebf9daa7",
   "metadata": {},
   "source": [
    "### Define training and validation set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a93fd590",
   "metadata": {},
   "outputs": [],
   "source": [
    "from deeptime.util.data import TrajectoryDataset\n",
    "\n",
    "train_data = TrajectoryDataset(lagtime=tau, trajectory=observable_traj.astype('float32'))\n",
    "val_data = TrajectoryDataset(lagtime=tau, trajectory=observable_traj_valid.astype('float32'))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "433c8d8f",
   "metadata": {},
   "source": [
    "### Define networks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7702a34",
   "metadata": {},
   "outputs": [],
   "source": [
    "from masks import Mask\n",
    "from collections import OrderedDict\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "# Assuming that we are on a CUDA machine, this should print a CUDA device:\n",
    "\n",
    "print(device)\n",
    "input_size = observable_traj.shape[1] \n",
    "mask = Mask(input_size, number_subsystems, mean=torch.Tensor(train_data.data.mean(0)),\n",
    "            std=torch.Tensor(train_data.data.std(0)), factor_fake=factor_fake, noise=noise, \n",
    "            device=device, cutoff=cutoff)\n",
    "mask.to(device=device)\n",
    "lobes = []\n",
    "for output_size in output_sizes:\n",
    "    lobe_dict = OrderedDict([('Layer_input', nn.Linear(input_size, layer_width)),\n",
    "                            ('Elu_input', nn.ELU())])\n",
    "    for d in range(network_depth):\n",
    "        lobe_dict['Layer'+str(d)]=nn.Linear(layer_width, layer_width)\n",
    "        lobe_dict['Elu'+str(d)]=nn.ELU()\n",
    "    lobe_dict['Layer_output']=nn.Linear(layer_width, output_size)\n",
    "    lobe_dict['Softmax']=nn.Softmax(dim=1) # obtain fuzzy probability distribution over output states\n",
    "    \n",
    "    lobe = nn.Sequential(\n",
    "        lobe_dict \n",
    "    )\n",
    "    lobes.append(lobe.to(device=device))\n",
    "\n",
    "print(mask)\n",
    "print(lobes)\n",
    "                 "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6533c8b0",
   "metadata": {},
   "source": [
    "### Create iVAMPnets estimator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f55db7a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from ivampnets import iVAMPnet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6c568ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "ivampnet = iVAMPnet(lobes, mask, device, learning_rate=learning_rate, epsilon=epsilon, score_mode=score_mode, learning_rate_mask=learning_rate/2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "60929530",
   "metadata": {},
   "source": [
    "### Plot mask before training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e14a6358",
   "metadata": {},
   "outputs": [],
   "source": [
    "from examples import plot_mask\n",
    "plot_mask(mask)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7875dc2",
   "metadata": {},
   "source": [
    "### Create data loader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7ccc0bf",
   "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": "markdown",
   "id": "8f98d914",
   "metadata": {},
   "source": [
    "### Create a tensorboard writer to observe performance during training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e100e00",
   "metadata": {},
   "outputs": [],
   "source": [
    "tensorboard_installed = False\n",
    "if tensorboard_installed:\n",
    "    from torch.utils.tensorboard import SummaryWriter\n",
    "    writer = SummaryWriter('./runs/Toy2/')\n",
    "    input_model, _ = next(iter(loader_train))\n",
    "else:\n",
    "    writer=None"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "26c4993c",
   "metadata": {},
   "source": [
    "### Fit the model on the training data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d2532ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "# model = ivampnet.fit(loader_train, n_epochs=10, validation_loader=loader_val, mask=False, lam_decomp=2., \n",
    "#                      lam_trace=1., start_mask=0, end_trace=2, tb_writer=writer, clip=False).fetch_model()\n",
    "# model = ivampnet.fit(loader_train, n_epochs=10, validation_loader=loader_val, mask=True, lam_decomp=0., \n",
    "#                      lam_trace=1., start_mask=0, end_trace=2, tb_writer=writer, clip=False).fetch_model()\n",
    "lam_pen_perc=0.05\n",
    "lam_pen_C00=0. \n",
    "lam_pen_C11=0. \n",
    "lam_pen_C01=0.\n",
    "model = ivampnet.fit(loader_train, n_epochs=10, validation_loader=loader_val, mask=False, lam_decomp=2., \n",
    "                     lam_trace=0., start_mask=0, end_trace=2, tb_writer=writer, clip=False,\n",
    "                     lam_pen_perc=lam_pen_perc, lam_pen_C00=lam_pen_C00, lam_pen_C11=lam_pen_C11,\n",
    "                     lam_pen_C01=lam_pen_C01).fetch_model()\n",
    "model = ivampnet.fit(loader_train, n_epochs=10, validation_loader=loader_val, mask=True, lam_decomp=0., \n",
    "                     lam_trace=0., start_mask=0, end_trace=2, tb_writer=writer, clip=False,\n",
    "                     lam_pen_perc=lam_pen_perc, lam_pen_C00=lam_pen_C00, lam_pen_C11=lam_pen_C11,\n",
    "                     lam_pen_C01=lam_pen_C01).fetch_model()\n",
    "if ivampnet.train_pen_scores[-1,1]>0.02:\n",
    "    print('The model does not seem to be converged to an independent solution!')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "edcbc766",
   "metadata": {},
   "outputs": [],
   "source": [
    "# execution time (cpu): ~ 30 sec"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "375d23fc",
   "metadata": {},
   "source": [
    "### Plot the training and validation scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41e73e2f",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.loglog(*ivampnet.train_scores.T, label='training')\n",
    "plt.loglog(*ivampnet.validation_scores.T, label='validation')\n",
    "plt.xlabel('step')\n",
    "plt.ylabel('score')\n",
    "plt.legend();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dbbed897",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.loglog(*ivampnet.train_pen_C01.T, label='training')\n",
    "plt.loglog(*ivampnet.validation_pen_C01.T, label='validation')\n",
    "plt.xlabel('step')\n",
    "plt.ylabel('score')\n",
    "plt.legend();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c19cc279",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.loglog(*ivampnet.train_pen_scores.T, label='training')\n",
    "plt.loglog(*ivampnet.validation_pen_scores.T, label='validation')\n",
    "plt.xlabel('step')\n",
    "plt.ylabel('score')\n",
    "plt.legend();"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a20f6c9c",
   "metadata": {},
   "source": [
    "### Plot the mask after the training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b1e2ed6",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_mask(mask)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ed8483a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# reproduces Fig. 3b (or permutation of it)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0935fa5",
   "metadata": {},
   "source": [
    "### Compare the eigenvalues from the true and the estimated transition matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8b68c6a",
   "metadata": {},
   "outputs": [],
   "source": [
    "mask.noise=0 # set the noise to zero. \n",
    "T_list = model.get_transition_matrix(val_data.data, val_data.data_lagged)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f891e262",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Estimated eigenvalues\n",
    "for T in T_list:\n",
    "    print(np.linalg.eigvals(T))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5181f619",
   "metadata": {},
   "outputs": [],
   "source": [
    "# True eigenvalues\n",
    "print(np.linalg.eigvals(toymodel.T1)), np.linalg.eigvals(toymodel.T2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7bd6fa2a",
   "metadata": {},
   "source": [
    "### Plot the state assignment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d91b1e36",
   "metadata": {},
   "outputs": [],
   "source": [
    "from examples import plot_states\n",
    "plot_states(model, val_data.data)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "91ae1a08",
   "metadata": {},
   "source": [
    "### Estimate the eigenfunctions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01a49a12",
   "metadata": {},
   "outputs": [],
   "source": [
    "from examples import plot_eigfuncs\n",
    "plot_eigfuncs(model, val_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d8049c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# reproduces Fig. 3c (possibly with permutation of state assignments)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "rseed",
   "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.8.15"
  },
  "vscode": {
   "interpreter": {
    "hash": "ae7d9bbee64574db456ae08990e902edc98caa89a93077ee255475653bb8dd96"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}