SynaptotagminC2A.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42847660",
   "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 torch.optim as optim\n",
    "import matplotlib.gridspec as gridspec\n",
    "\n",
    "import h5py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43117557",
   "metadata": {},
   "source": [
    "### Hyperparameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f87cbf8",
   "metadata": {},
   "outputs": [],
   "source": [
    "stride = 1\n",
    "\n",
    "tau = 100//stride \n",
    "\n",
    "output_sizes = [8,8]\n",
    "number_subsystems = len(output_sizes)\n",
    "# tau list for timescales estimation\n",
    "tau_list = [1,2,4,8]\n",
    "\n",
    "# Batch size for Stochastic Gradient descent\n",
    "batch_size = 20000\n",
    "# Which trajectory points percentage is used as validation and testing, the rest is for training\n",
    "valid_ratio = 0.3\n",
    "test_ratio = 0.0001\n",
    "# How many hidden layers the network chi has\n",
    "network_depth = 3\n",
    "\n",
    "# Width of every layer of chi\n",
    "layer_width = 100\n",
    "# create a list with the number of nodes for each layer\n",
    "nodes = [layer_width]*network_depth\n",
    "# data preparation\n",
    "# how many residues are skipped for distance calculation\n",
    "skip_res = 6\n",
    "# Size of the windows for attention mechanism\n",
    "patchsize = 8\n",
    "# How many residues are skipped before defining a new window\n",
    "skip_over = 4\n",
    "\n",
    "# How strong the fake subsystem is\n",
    "factor_fake = 2.\n",
    "# How large the noise in the mask for regularization is\n",
    "noise = 2.\n",
    "# Threshold after which the attention weight is set to zero\n",
    "cutoff=0.9\n",
    "# Learning rate\n",
    "learning_rate=0.001\n",
    "# epsilon\n",
    "epsilon=1e-5\n",
    "# score method\n",
    "score_mode='regularize' # one of ('trunc', 'regularize', 'clamp', 'old')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d10b4cf1",
   "metadata": {},
   "source": [
    "### Load data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "588c7ec9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# data set has a total length of 184 µs with a 1 ns resolution (total of 184000 frames)\n",
    "\n",
    "data_trajs = []\n",
    "hdf5_names = []\n",
    "loaded_data_stride = 100\n",
    "exclude_list = []\n",
    "with h5py.File(f\"/group/ag_cmb/scratch/deeptime_data/syt/syt_0cal_internal1by1_stride{loaded_data_stride}.hdf5\", \"r\") as f: # 1 frame = 1 ns\n",
    "    #print(\"datasets:\", f.keys())\n",
    "    for n, name in enumerate(f.keys()):\n",
    "        if n not in exclude_list:\n",
    "            hdf5_names.append(name)\n",
    "            dset = f[name]\n",
    "            dat = dset[...].astype('float32')\n",
    "\n",
    "            data_trajs.append(1/np.exp(dat[::stride]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "51a0e9fd",
   "metadata": {},
   "source": [
    "### Define dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f45ee7fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "from deeptime.util.data import TrajectoriesDataset\n",
    "\n",
    "dataset = TrajectoriesDataset.from_numpy(lagtime=tau, data=data_trajs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22c7b97e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# split into training/validation/test set\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])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d6acabf",
   "metadata": {},
   "source": [
    "### Define networks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "634d4b2b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from masks import Mask_proteins\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",
    "train_mean = np.concatenate(train_data.dataset.trajectories, axis=0).mean(0)\n",
    "train_std = np.concatenate(train_data.dataset.trajectories, axis=0).std(0)\n",
    "print(device)\n",
    "input_size = train_data.dataset.trajectories[0].shape[-1]\n",
    "mask = Mask_proteins(input_size, number_subsystems, skip_res=skip_res, patchsize=patchsize, skip=skip_over, mean=torch.Tensor(train_mean),\n",
    "            std=torch.Tensor(train_std), factor_fake=factor_fake, noise=noise, 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)         "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea28cbbb",
   "metadata": {},
   "source": [
    "### Create iVAMPnets estimator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05aabb2d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from ivampnets import iVAMPnet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0325caf5",
   "metadata": {},
   "outputs": [],
   "source": [
    "ivampnet = iVAMPnet(lobes, mask, device, learning_rate=learning_rate, epsilon=epsilon, score_mode=score_mode)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "967887b4",
   "metadata": {},
   "source": [
    "### Plot mask before training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dba6b102",
   "metadata": {},
   "outputs": [],
   "source": [
    "from examples import plot_mask\n",
    "plot_mask(mask, skip=10)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f03dd31d",
   "metadata": {},
   "source": [
    "### Create data loader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "06a08456",
   "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": "c8a820a0",
   "metadata": {},
   "source": [
    "### Create a tensorboard writer to observe performance during training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7038f8fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "tensorboard_installed = False\n",
    "if tensorboard_installed:\n",
    "    from torch.utils.tensorboard import SummaryWriter\n",
    "    writer = SummaryWriter(log_dir='./runs/Syt/')\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",
   "id": "229a10ff",
   "metadata": {},
   "source": [
    "### Fit the model on the training data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bfaff88e",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = ivampnet.fit(loader_train, n_epochs=50, validation_loader=loader_val, mask=True, lam_decomp=20., \n",
    "                     lam_trace=1., start_mask=0, end_trace=20, tb_writer=writer, clip=False).fetch_model()\n",
    "\n",
    "plot_mask(mask, skip=10)\n",
    "mask.noise=5.\n",
    "model = ivampnet.fit(loader_train, n_epochs=150, validation_loader=loader_val, mask=True, lam_decomp=50., \n",
    "                     lam_trace=0., start_mask=0, end_trace=0, tb_writer=writer, clip=False).fetch_model()\n",
    "plot_mask(mask, skip=10)\n",
    "mask.noise=10.\n",
    "model = ivampnet.fit(loader_train, n_epochs=150, validation_loader=loader_val, mask=True, lam_decomp=100., \n",
    "                     lam_trace=0., start_mask=0, end_trace=0, tb_writer=writer, clip=False).fetch_model()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "474445e3",
   "metadata": {},
   "source": [
    "### Plot training and validation scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a16599e",
   "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": "markdown",
   "id": "2e65c886",
   "metadata": {},
   "source": [
    "### Plot the mask after training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66a61989",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_mask(mask, skip=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8437a367",
   "metadata": {},
   "outputs": [],
   "source": [
    "from examples import plot_protein_its, plot_protein_mask\n",
    "plot_protein_mask(mask, skip_start=4)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f01dfcd4",
   "metadata": {},
   "source": [
    "### Finally train without noise"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92c01b2a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# the noise is only important to make the training of the mask meaningfull.\n",
    "# Here, the mask should be well trained, so we disable the mask training from here on.\n",
    "mask.noise=0.\n",
    "model = ivampnet.fit(loader_train, n_epochs=300, validation_loader=loader_val, mask=False, lam_decomp=100., \n",
    "                     lam_trace=0., start_mask=0, end_trace=0, tb_writer=writer, clip=False).fetch_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ddc0b3ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "# In principle you can then also train the model without enforcing the decomposition score anymore\n",
    "# However, you should observe if the independence score rise significantly, then you need to reverse the progress\n",
    "# You can use the save_criteria parameter to control it.\n",
    "model = ivampnet.fit(loader_train, n_epochs=300, validation_loader=loader_val, mask=False, lam_decomp=0., \n",
    "                     lam_trace=0., start_mask=0, end_trace=0, tb_writer=writer, clip=False, save_criteria=0.012).fetch_model()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "29e2875c",
   "metadata": {},
   "source": [
    "### Estimate implied timescales"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "036bd72b",
   "metadata": {},
   "outputs": [],
   "source": [
    "runs = 5\n",
    "its = [[] for _ in range(runs)]\n",
    "# cheap error estimation, instead of retraining chi, evaluate the model on different trajectories\n",
    "percentage = 0.9\n",
    "N_trajs = len(dataset.trajectories)\n",
    "indexes_traj = np.arange(N_trajs)\n",
    "n_val = int(N_trajs * percentage)\n",
    "msmlags=np.array([1,2,4,6,10,15,20,25])*10\n",
    "for run in range(runs):\n",
    "    for tau_i in msmlags:\n",
    "        np.random.shuffle(indexes_traj)\n",
    "        indexes_used = indexes_traj[:n_val]\n",
    "        data_t = np.concatenate([dataset.trajectories[a][:-tau_i] for a in indexes_used], axis=0)\n",
    "        data_tau = np.concatenate([dataset.trajectories[a][tau_i:] for a in indexes_used], axis=0)\n",
    "        its[run].append(model.timescales(data_t, data_tau, tau_i, batchsize=10000))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e16ad8b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# reorder its, subsystems can have different outputsizes!\n",
    "its_reorder = [np.zeros((runs,len(msmlags), output_sizes[n]-1)) for n in range(number_subsystems)]\n",
    "for n in range(number_subsystems):\n",
    "    for run in range(runs):\n",
    "        for lag in range(len(msmlags)):\n",
    "            its_reorder[n][run,lag] = its[run][lag][n]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57c8d5a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "axes, fig = plot_protein_its(its_reorder, msmlags, ylog=True, multiple_runs=True, percent=0.9)\n",
    "x_ticks = np.array([1,5,10,20,40])*10\n",
    "x_ticks_labels = x_ticks*stride # for estimating the right units!\n",
    "y_ticks = np.array([1000,10000, 100000])/stride\n",
    "y_ticks_labels = y_ticks*stride/1000\n",
    "for n in range(number_subsystems):\n",
    "    ax=axes[n]\n",
    "    ax.plot(msmlags,msmlags, 'k')\n",
    "    ax.fill_between(msmlags, msmlags[0], msmlags, color = 'k', alpha = 0.2)\n",
    "    ax.set_xlabel('Lagtime [ns]', fontsize=16)\n",
    "    if n==0:\n",
    "        ax.set_ylabel('Implied Timescales [$\\mu$s]', fontsize=16)\n",
    "    ax.legend(fontsize=14, loc='lower right')\n",
    "    ax.set_xticks(x_ticks)\n",
    "    ax.set_xticklabels(x_ticks_labels, fontsize=14)\n",
    "    ax.set_yticks(y_ticks)\n",
    "    ax.set_yticklabels(y_ticks_labels, fontsize=14)\n",
    "    ax.tick_params(direction='out', length=6, width=2, colors='k',\n",
    "                   grid_color='k', grid_alpha=0.5)\n",
    "    ax.set_xlim(10,250)\n",
    "    ax.set_ylim(0.01*1000, 200*1000)\n",
    "    # fig.savefig('./Syt_its.pdf', bbox_inches='tight')\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9df5595e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# reproduces Fig. 5b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e6ec4ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ivampnet.save_params('./Syt_params')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f51b21f6",
   "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.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}