Villin_training_with_experimental_observables.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Notebook to construct a deep reversible Markov State Model where additional information of experimental values can be included\n",
    "\n",
    "The training steps include:\n",
    "1. Pretraining of an ordinary VAMPnet (can be used to compare the results)\n",
    "2. Initiate the values for the paramter matrix $\\mathbf{S}$ and the reweighting vector $\\mathbf{u}$ and train for matching specific observables\n",
    "3. Training for $\\mathbf{u}$ and $\\mathbf{S}$ and train for matching specific observables\n",
    "4. Training for $\\boldsymbol{\\chi}$ and $\\mathbf{u}$ and $\\mathbf{S}$ and train for matching specific observables\n",
    "\n",
    "The analysis consists of:\n",
    "1. Compare the performance on predefined observables not included in the training against a model which is trained without the additional experimental information\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import all the relevant packages\n",
    "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check for cuda supported device\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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define where the data lies on the local machine\n",
    "# Needs to be adapted for own data!!!\n",
    "test_system = '2F4K'\n",
    "pdb_system = '2f4k_villin.pdb'\n",
    "\n",
    "root = '/group/ag_cmb/simulation-data/DESRES-Science2011-FastProteinFolding/DESRES-Trajectory_{0}-0-protein/{0}-0-protein/{0}-0-protein'.format(test_system)"
   ]
  },
  {
   "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=1"
   ]
  },
  {
   "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",
    "output_sizes = [3,2]\n",
    "# tau list for timescales estimation\n",
    "tau_list = [25,50,100,150,200]\n",
    "number_taus = len(tau_list)\n",
    "tau_list = np.array(tau_list)//skip\n",
    "\n",
    "# Tau, how much is the timeshift of the two datasets in the default training\n",
    "tau = 2*25//skip # 5, 20\n",
    "tau_chi = 2*25//skip\n",
    "\n",
    "# Batch size for Stochastic Gradient descent\n",
    "batch_size = 10000\n",
    "\n",
    "# Which trajectory points percentage is used as training\n",
    "train_ratio = 0.33\n",
    "valid_ratio = 0.33\n",
    "\n",
    "# How many hidden layers the network chi has\n",
    "network_depth = 5\n",
    "\n",
    "# Width of every layer of chi\n",
    "layer_width = 100\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": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "traj_whole_new = [traj_whole[0][::skip]]\n",
    "input_size = traj_whole_new[0].shape[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# function for estimating the Koopman operator\n",
    "def estimate_koopman_op(trajs, tau, force_symmetric = False):\n",
    "    '''Estimates the koopman operator for a given trajectory at the lag time\n",
    "        specified. The formula for the estimation is:\n",
    "            K = C00 ^ -1 @ C01\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    traj: numpy array with size [traj_timesteps, traj_dimensions]\n",
    "        Trajectory described by the returned koopman operator\n",
    "\n",
    "    tau: int\n",
    "        Time shift at which the koopman operator is estimated\n",
    "        \n",
    "    force_symmetric: boolean, default = False\n",
    "        if true, calculates the symmetrized version of K instead\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    koopman_op: numpy array with shape [traj_dimensions, traj_dimensions]\n",
    "        Koopman operator estimated at timeshift tau\n",
    "\n",
    "    '''\n",
    "    # if tau larger 0, interpret trajs as either a list of trajectories\n",
    "    # or a single trajectory\n",
    "    # otherwise interpret trajs as a list of the data and time-lagged data\n",
    "    # possibly in random order\n",
    "    if tau > 0:\n",
    "        if type(trajs) == list:\n",
    "            traj = np.concatenate([t[:-tau] for t in trajs], axis = 0)\n",
    "            traj_lag = np.concatenate([t[tau:] for t in trajs], axis = 0)\n",
    "        else:\n",
    "            traj = trajs[:-tau]\n",
    "            traj_lag = trajs[tau:]\n",
    "    else:\n",
    "        traj = trajs[0]\n",
    "        traj_lag = trajs[1]\n",
    "        \n",
    "    koopman_op = np.eye(traj.shape[1])\n",
    "\n",
    "    c_0 = traj.T @ traj\n",
    "    c_tau = traj.T @ traj_lag\n",
    "    \n",
    "    # if you want to symmetrize the correlation matrices\n",
    "    if force_symmetric:\n",
    "        c_0 = c_0 + traj_lag.T @ traj_lag\n",
    "        c_tau = c_tau + traj_lag.T @ traj\n",
    "\n",
    "    eigv_all, eigvec_all = np.linalg.eig(c_0)\n",
    "    include = eigv_all > epsilon\n",
    "    eigv = eigv_all[include]\n",
    "    eigvec = eigvec_all[:,include]\n",
    "    c0_inv = eigvec @ np.diag(1/eigv) @ np.transpose(eigvec)\n",
    "\n",
    "    koopman_op = c0_inv @ c_tau\n",
    "\n",
    "    return koopman_op\n",
    "\n",
    "\n",
    "\n",
    "# utility function for plotting implied timescales\n",
    "# the hyperparameters allow for it to calculate errorbars and work from different input data types\n",
    "def get_its(data, lags, calculate_K = True, multiple_runs = False):\n",
    "    \n",
    "    def get_single_its(data):\n",
    "\n",
    "        if type(data) == list:\n",
    "            outputsize = data[0].shape[1]\n",
    "        else:\n",
    "            outputsize = data.shape[1]\n",
    "\n",
    "        single_its = np.zeros((outputsize-1, len(lags)))\n",
    "\n",
    "        for t, tau_lag in enumerate(lags):\n",
    "            if calculate_K:\n",
    "                koopman_op = self.estimate_koopman_op(data, tau_lag)\n",
    "            else:\n",
    "                koopman_op = data[t]\n",
    "            k_eigvals, k_eigvec = np.linalg.eig(np.real(koopman_op))\n",
    "            k_eigvals = np.sort(np.absolute(k_eigvals))\n",
    "            k_eigvals = k_eigvals[:-1]\n",
    "            single_its[:,t] = (-tau_lag / np.log(k_eigvals))\n",
    "\n",
    "        return np.array(single_its)\n",
    "\n",
    "\n",
    "    if not multiple_runs:\n",
    "\n",
    "        its = get_single_its(data)\n",
    "\n",
    "    else:\n",
    "\n",
    "        its = []\n",
    "        for data_run in data:\n",
    "            its.append(get_single_its(data_run))\n",
    "\n",
    "    return its\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Helper functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _inv(x, ret_sqrt=False):\n",
    "    '''Utility function that returns the inverse of a matrix, with the\n",
    "    option to return the square root of the inverse matrix.\n",
    "    Parameters\n",
    "    ----------\n",
    "    x: numpy array with shape [m,m]\n",
    "        matrix to be inverted\n",
    "\n",
    "    ret_sqrt: bool, optional, default = False\n",
    "        if True, the square root of the inverse matrix is returned instead\n",
    "    Returns\n",
    "    -------\n",
    "    x_inv: numpy array with shape [m,m]\n",
    "        inverse of the original matrix\n",
    "    '''\n",
    "\n",
    "    # Calculate eigvalues and eigvectors\n",
    "    eigval_all, eigvec_all = torch.symeig(x, eigenvectors=True)\n",
    "\n",
    "    # Filter out eigvalues below threshold and corresponding eigvectors\n",
    "    eig_th = torch.Tensor(epsilon)\n",
    "    index_eig = eigval_all > eig_th\n",
    "    eigval = eigval_all[index_eig]\n",
    "    eigvec = eigvec_all[:,index_eig]\n",
    "\n",
    "    # Build the diagonal matrix with the filtered eigenvalues or square\n",
    "    # root of the filtered eigenvalues according to the parameter\n",
    "    if ret_sqrt:\n",
    "        diag = torch.diag(torch.sqrt(1/eigval))\n",
    "    else:\n",
    "        diag = torch.diag(1/eigval)\n",
    "    # Rebuild the square root of the inverse matrix\n",
    "    x_inv = torch.matmul(eigvec, torch.matmul(diag, eigvec.T))\n",
    "\n",
    "    return x_inv\n",
    "\n",
    "\n",
    "def _prep_data(data_t, data_tau):\n",
    "    '''Utility function that transorms the input data from a tensorflow - \n",
    "    viable format to a structure used by the following functions in the\n",
    "    pipeline.\n",
    "    Parameters\n",
    "    ----------\n",
    "    data: tensorflow tensor with shape [b, 2*o]\n",
    "        original format of the data\n",
    "    Returns\n",
    "    -------\n",
    "    x: tensorflow tensor with shape [o, b]\n",
    "        transposed, mean-free data corresponding to the left, lag-free lobe\n",
    "        of the network\n",
    "\n",
    "    y: tensorflow tensor with shape [o, b]\n",
    "        transposed, mean-free data corresponding to the right, lagged lobe\n",
    "        of the network\n",
    "\n",
    "    b: tensorflow float32\n",
    "        batch size of the data\n",
    "\n",
    "    o: int\n",
    "        output size of each lobe of the network\n",
    "\n",
    "    '''\n",
    "\n",
    "\n",
    "    # Subtract the mean\n",
    "    x = data_t - torch.mean(data_t, dim=0, keepdim=True)\n",
    "    y = data_tau - torch.mean(data_tau, dim=0, keepdim=True)\n",
    "\n",
    "    return x, y\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Definition of the classes for the model and the attention network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Mask(torch.nn.Module):\n",
    "    ''' Attention mask either independent from the time point (mask_const=True) or dependent.\n",
    "    If dependent the attention is estimated via a NN with depth and width given as input, which are \n",
    "    otherwise ignored. \n",
    "    The attention mechanism assumes that distances are used. skip_res is number of residues skiped when estimating\n",
    "    the distance. \n",
    "    \n",
    "    ---------\n",
    "    Inputs\n",
    "    ---------\n",
    "    input_size: int\n",
    "            the dimension of the input array, which are expected to be distances\n",
    "    mask_const: bool\n",
    "            if True a constant mask is used independent from the actual time frame, which results\n",
    "            in trainable weights of vector with the dimension of the number of windows\n",
    "            if False a time dependent attention mechanism is used with a NN with depth and width\n",
    "    depth: int\n",
    "            the number of hidden layers for the attention mechanism if mask_const=True.\n",
    "    width: int\n",
    "            the width of hidden layers for the attention mechanism if mask_const=True.\n",
    "    heads: int\n",
    "            Number of how many weights per residue which will be averaged.\n",
    "    sensitivity: float\n",
    "            Prefactor in the softmax function to regulate how hard the assignment tends to be.\n",
    "    fac: bool\n",
    "            if True The weight for each residue is multiplied with the number of residues. \n",
    "            This circumvents that the product of the two weights do not scale too fast to zero. The weights\n",
    "            are then distributed around 1 instead of 1/number_residues.\n",
    "              \n",
    "            \n",
    "    \n",
    "    \n",
    "    '''\n",
    "    def __init__(self, input_size, mask_const, depth=0, width=100 , heads=1, sensitivity=1, fac=True):\n",
    "        super(Mask, self).__init__()\n",
    "        \n",
    "#         self.alpha = torch.Tensor(1, input_size, N, heads).fill_(0)\n",
    "\n",
    "        skip_res = 2\n",
    "        self.n_residues = int(-1/2 + np.sqrt(1/4+input_size*2) + skip_res)\n",
    "        self.residues_1 = []\n",
    "        self.residues_2 = []\n",
    "        # estimate the pairs\n",
    "        for n1 in range(self.n_residues-skip_res):\n",
    "            for n2 in range(n1+skip_res, self.n_residues):\n",
    "                self.residues_1.append(n1)\n",
    "                self.residues_2.append(n2)\n",
    "        self.mask_const = mask_const\n",
    "        if mask_const:        \n",
    "            self.alpha = torch.randn((1, self.n_residues, heads)) * 0.5\n",
    "            self.weight = torch.nn.Parameter(data=self.alpha, requires_grad=True)\n",
    "        else:\n",
    "            nodes = [input_size]\n",
    "            for i in range(depth):\n",
    "                nodes.append(width)\n",
    "            \n",
    "            \n",
    "            self.hfc = [nn.Linear(nodes[i], nodes[i+1]) for i in range(len(nodes)-1)]\n",
    "            self.softmax = nn.Linear(nodes[-1], self.n_residues * heads)\n",
    "            \n",
    "        self.sensitivity = sensitivity\n",
    "        if fac:\n",
    "            self.fac = self.n_residues * self.n_residues\n",
    "        else:\n",
    "            self.fac = 1.\n",
    "        self.heads = heads\n",
    "        \n",
    "    def forward(self, x):\n",
    "        \n",
    "        # weights for each residue\n",
    "        if self.mask_const:\n",
    "            weight_sf = torch.sum(F.softmax(self.weight*self.sensitivity, dim=1), dim=2)/self.heads\n",
    "        else:\n",
    "            y = x\n",
    "            for layer in self.hfc:\n",
    "                y = layer(y)\n",
    "            y = self.softmax(y)\n",
    "            y = torch.reshape(y, (-1, self.n_residues, heads))\n",
    "            weight_sf = torch.sum(F.softmax(y*self.sensitivity, dim=1), dim=2)/self.heads\n",
    "            \n",
    "            \n",
    "        # estimate for applying it\n",
    "        weight_1 = weight_sf[:,self.residues_1]\n",
    "        weight_2 = weight_sf[:,self.residues_2]\n",
    "        \n",
    "        masked_x = x * weight_1 * weight_2 * self.fac# include factor\n",
    "        \n",
    "        return masked_x\n",
    "    \n",
    "    def get_softmax(self, x=None):\n",
    "        if self.mask_const:\n",
    "            weight_sf = torch.sum(F.softmax(self.weight*self.sensitivity, dim=1), dim=2)/self.heads\n",
    "        else:\n",
    "            \n",
    "            y = x\n",
    "            for layer in self.hfc:\n",
    "                y = layer(y)\n",
    "            y = self.softmax(y)\n",
    "            y = torch.reshape(y, (-1, n_residues, heads))\n",
    "            weight_sf = torch.sum(F.softmax(y*self.sensitivity, dim=1), dim=2)/self.heads\n",
    "\n",
    "        \n",
    "        return weight_sf\n",
    "    \n",
    "    def set_weights(self, weights):\n",
    "        if self.mask_const:\n",
    "            alpha = torch.Tensor(weights)\n",
    "\n",
    "            self.weight = torch.nn.Parameter(data=alpha, requires_grad=True)\n",
    "        else:\n",
    "            print('not implemented yet. You need to define all layers in the mask')\n",
    "\n",
    "            \n",
    "\n",
    "            \n",
    "class Coarse_grain(torch.nn.Module):\n",
    "    ''' Definition of the coarse-graining layer, which has a trainable matrix which maps the input \n",
    "        classification from intput_dim dimensions on output_dim dimensions. Where output_dim<input_dim.\n",
    "        \n",
    "        -------\n",
    "        Inputs\n",
    "        -------\n",
    "        input_dim: int\n",
    "                The number of states of the classification which should be coarse-grained\n",
    "        output_dim: int\n",
    "                The number of states input_dim should be coarse-grained to, which should be strictly\n",
    "                smaller\n",
    "        sen: int\n",
    "                The sensitifity of the softmax function. Can be interpreted as a inverse temperature.\n",
    "                Where a larger value results in harder assignments and vice versa.\n",
    "    \n",
    "    \n",
    "    '''\n",
    "    def __init__(self, input_dim, output_dim, sen=1):\n",
    "        super(Coarse_grain, self).__init__()\n",
    "\n",
    "        self.N = input_dim\n",
    "        self.M = output_dim\n",
    "        self.sen = sen\n",
    "        \n",
    "        self.alpha = torch.randn((self.N, self.M)) * 0.5\n",
    "        self.weight = torch.nn.Parameter(data=self.alpha, requires_grad=True)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        \n",
    "        \n",
    "        kernel = F.softmax(self.sen * self.weight, dim=1)\n",
    "        \n",
    "        ret = x @ kernel\n",
    "\n",
    "        return ret\n",
    "    \n",
    "    def get_softmax(self):\n",
    "        \n",
    "        return F.softmax(self.sen * self.weight, dim=1)\n",
    "            \n",
    "    def reset_params(self):\n",
    "        \n",
    "        with torch.no_grad():\n",
    "            \n",
    "            \n",
    "            self.weight.copy_(torch.randn((self.N, self.M)) * 0.5) \n",
    "        \n",
    "\n",
    "        \n",
    "class U_layer(torch.nn.Module):\n",
    "    ''' Layer to define the trainable u vector for a deep reversible MSM, which fulfills all\n",
    "        necessary constraints.\n",
    "        \n",
    "        --------\n",
    "        Inputs\n",
    "        --------\n",
    "        output_dim: int\n",
    "                Number of states of the corresponding chi function\n",
    "        activation: function\n",
    "                Activation function which will be applied to the weights. For a reversible deep MSM\n",
    "                this function should be monoton and non-negative.\n",
    "    \n",
    "    \n",
    "    '''\n",
    "    def __init__(self, output_dim, activation):\n",
    "        super(U_layer, self).__init__()\n",
    "\n",
    "        self.M = output_dim\n",
    "        \n",
    "        self.alpha = torch.Tensor(1, self.M).fill_(1/self.M)\n",
    "        self.u_kernel = torch.nn.Parameter(data=self.alpha, requires_grad=True)\n",
    "        self.acti = activation\n",
    "        \n",
    "    def forward(self, chi_t, chi_tau):\n",
    "        \n",
    "        # we need batchsize to stack the outputs later so that it fullfills keras' requirements\n",
    "        batchsize = chi_t.shape[0]\n",
    "\n",
    "        # note: corr_tau is the correlation matrix of the time-shifted data\n",
    "        # presented in the paper at page 6, \"Normalization of transition density\"\n",
    "        corr_tau = 1./batchsize * torch.matmul(chi_tau.T, chi_tau)\n",
    "        chi_mean = torch.mean(chi_tau, dim=0, keepdim=True)\n",
    "\n",
    "        kernel_u = self.acti(self.u_kernel)\n",
    "\n",
    "        # u is the normalized and transformed kernel of this layer\n",
    "        u = kernel_u / torch.sum(chi_mean * kernel_u, dim=1, keepdim=True)\n",
    "\n",
    "        v = torch.matmul(corr_tau, u.T)\n",
    "\n",
    "        mu = 1./batchsize * torch.matmul(chi_tau, u.T)\n",
    "        \n",
    "        Sigma =  torch.matmul((chi_tau * mu).T, chi_tau)\n",
    "        \n",
    "        chi_mean_t = torch.mean(chi_t, dim=0, keepdim=True)\n",
    "        u_t = kernel_u / torch.sum(chi_mean_t * kernel_u, dim=1, keepdim=True)\n",
    "        mu_t = 1./batchsize * torch.matmul(chi_t, u_t.T)\n",
    "        Sigma_t =  torch.matmul((chi_t * mu_t).T, chi_t)\n",
    "\n",
    "        gamma = chi_tau * (torch.matmul(chi_tau, u.T))\n",
    "\n",
    "        C_00 = 1./batchsize * torch.matmul(chi_t.T, chi_t)\n",
    "        C_11 = 1./batchsize * torch.matmul(gamma.T, gamma)\n",
    "        C_01 = 1./batchsize * torch.matmul(chi_t.T, gamma)\n",
    "\n",
    "\n",
    "        ret = [\n",
    "            v,\n",
    "            C_00,\n",
    "            C_11,\n",
    "            C_01,\n",
    "            Sigma,\n",
    "            mu_t,\n",
    "            Sigma_t\n",
    "        ]\n",
    "        \n",
    "        return ret\n",
    "    \n",
    "class S_layer(torch.nn.Module):\n",
    "    ''' Layer to define the trainable S matrix for a deep reversible MSM, which fulfills all\n",
    "        necessary constraints.\n",
    "        \n",
    "        --------\n",
    "        Inputs\n",
    "        --------\n",
    "        output_dim: int\n",
    "                Number of states of the corresponding chi function\n",
    "        activation: function\n",
    "                Activation function which will be applied to the weights. For a reversible deep MSM\n",
    "                this function should be monoton and non-negative.\n",
    "        renorm: bool\n",
    "                if True, S will be strictly non-negative, which is necessary for a deep reversible MSM\n",
    "                However, if a deep reversible Koopman model is supposed to be trained it can be set\n",
    "                False.\n",
    "    \n",
    "    \n",
    "    '''\n",
    "    def __init__(self, output_dim, activation, renorm=True):\n",
    "        super(S_layer, self).__init__()\n",
    "\n",
    "        self.M = output_dim\n",
    "        \n",
    "        self.alpha = torch.Tensor(self.M, self.M).fill_(0.1)\n",
    "        self.S_kernel = torch.nn.Parameter(data=self.alpha, requires_grad=True)\n",
    "        self.acti = activation\n",
    "        self.renorm = renorm\n",
    "        \n",
    "    def forward(self, v, C_00, C_11, C_01, Sigma):\n",
    "        \n",
    "        # we need batchsize to stack the outputs later so that it fullfills keras' requirements\n",
    "            \n",
    "        batchsize = v.shape[0]\n",
    "\n",
    "        # transform the kernel weights\n",
    "        kernel_w = self.acti(self.S_kernel)\n",
    "        \n",
    "        # enforce symmetry\n",
    "        W1 = kernel_w + kernel_w.T\n",
    "\n",
    "        # normalize the weights\n",
    "        norm = W1 @ v\n",
    "        \n",
    "        if self.renorm:\n",
    "            # make sure that the largest value of norm is < 1\n",
    "            quasi_inf_norm = lambda x: torch.sum((x**20))**(1./20)\n",
    "#             print(norm, quasi_inf_norm(norm))\n",
    "            W1 = W1 / quasi_inf_norm(norm)\n",
    "            norm = W1 @ v\n",
    "\n",
    "        \n",
    "        w2 = (1 - torch.squeeze(norm)) / torch.squeeze(v)\n",
    "        S = W1 + torch.diag(w2)\n",
    "\n",
    "        # calculate K\n",
    "        K = S @ Sigma\n",
    "\n",
    "        # VAMP-E matrix for the computation of the loss\n",
    "        VampE_matrix = S.T @ C_00 @ S @ C_11 - 2*S.T @ C_01\n",
    "\n",
    "        # stack outputs so that the first dimension is = batchsize, keras requirement\n",
    "        \n",
    "        ret = [VampE_matrix, K, S]\n",
    "        \n",
    "        return ret\n",
    "            \n",
    "class VampNet(nn.Module):\n",
    "    ''' VAMPnet class\n",
    "    \n",
    "        --------\n",
    "        Inputs\n",
    "        --------\n",
    "        input_dim: int\n",
    "                Number of input features.\n",
    "        output_sizes: list\n",
    "                Number of states of the corresponding chi function and possible following coarse graining\n",
    "                layers. The first value is the number of output nodes of the VAMPnet. All possible values\n",
    "                afterwards will be interpreted as coarse-graining layers, mapping from the previous value \n",
    "                to the one afterwards. Therefore, the int values in the list should be in decreasing order.\n",
    "        nodes: list\n",
    "                List of nodes for each hidden layer.\n",
    "        train_mean: np.array\n",
    "                The mean of all input features over the whole training set. The mean will be substracted\n",
    "                before the data is passed into first hidden layer. We prefer this way over\n",
    "                Batchnormalization layers since the output during training does not depend which other\n",
    "                training data is in the batch.\n",
    "        train_std: np.array\n",
    "                The standard deviation of all input features over the whole training set. The input wil be\n",
    "                devided by it before passing it into the first hidden layer. Together with the substraction\n",
    "                of train_mean it should lead to approx. normal distributed input data.\n",
    "        valid_T: bool\n",
    "                If true: Constraints are active to ensure a valid transition matrix.\n",
    "                else: the constraints are inactive.\n",
    "        reversible: bool\n",
    "                If true: Constraints are active to ensure reversibility of the model.\n",
    "                else: the constraints are inactive.\n",
    "        mask_const: bool\n",
    "                If true: The attention mechanism is not time dependent and therefore only representeted\n",
    "                by a trainable vector with the same size as the output of an attention network\n",
    "                else: An attention neural network is used which results into time dependent attention weights.\n",
    "                The hyperparameters are given by mask_depth and mask_width.\n",
    "        mask_depth: int\n",
    "                Determines the number of hidden layers of the attention NN.\n",
    "        mask_width: int\n",
    "                Determines the number of hidden nodes for each of the layers for the attention NN.\n",
    "        heads: int\n",
    "                Number of how many weights per residue which will be averaged.\n",
    "        sensitivity: float\n",
    "                Prefactor in the softmax function to regulate how hard the assignment tends to be.\n",
    "        fac: bool\n",
    "                if True The weight for each residue is multiplied with the number of residues. \n",
    "                This circumvents that the product of the two weights do not scale too fast to zero. The weights\n",
    "                are then distributed around 1 instead of 1/number_residues.\n",
    "        softmax_fac: float\n",
    "                A prefactor in the last softmax layer, which can be interprested as an inverse temperature\n",
    "                for the classification of the states. The larger the value the harder the assignment will be.\n",
    "    \n",
    "    \n",
    "    '''\n",
    "    def __init__(self, input_size, output_sizes, nodes, train_mean, train_std, \n",
    "                 valid_T=False, reversible=False,\n",
    "                 mask_const=True, mask_depth=0, mask_width=0, heads=1, sensitivity=1, fac=True,\n",
    "                 softmax_fac=1.):\n",
    "        super(VampNet, self).__init__()\n",
    "        \n",
    "        # which physical constraints are enacted which can need more parameters with different contraints\n",
    "        # Furthermore the best learning practice is different\n",
    "        self.valid_T = valid_T\n",
    "        self.reversible = reversible\n",
    "        if valid_T:\n",
    "            model='DMSM'\n",
    "        else:\n",
    "            model='VAMPnet'\n",
    "        if reversible:\n",
    "            rev='rev'\n",
    "        else:\n",
    "            rev=''\n",
    "        print('The trained model is a '+rev+model) \n",
    "        \n",
    "        if valid_T and not reversible:\n",
    "            self.gamma = True\n",
    "        else:\n",
    "            self.gamma = False\n",
    "        if reversible:\n",
    "            if valid_T:\n",
    "                self.factor_S = 1\n",
    "                self.factor_u = 0.01\n",
    "                self.renorm = True\n",
    "                acti_S = lambda x: self.factor_S * torch.exp(x)\n",
    "                acti_u = lambda x: self.factor_u * torch.exp(x)\n",
    "                \n",
    "            else: \n",
    "                self.factor_S = 0.001\n",
    "                self.factor_u = .000001\n",
    "\n",
    "                linear_S = lambda x: self.factor_S * x\n",
    "                linear_u = lambda x: self.factor_u * x\n",
    "                self.renorm = False\n",
    "                acti_S = linear_S\n",
    "                acti_u = linear_u\n",
    "            \n",
    "            \n",
    "            self.u_layers = [U_layer(o, acti_u) for o in output_sizes]\n",
    "            self.S_layers = [S_layer(o, acti_S, self.renorm) for o in output_sizes]\n",
    "                \n",
    "# activations for reversible VAMPnets\n",
    "\n",
    "\n",
    "#activations for revDMSM\n",
    "\n",
    "        \n",
    "        \n",
    "        self.N = len(output_sizes)\n",
    "        self.output_sizes = output_sizes\n",
    "        \n",
    "        self.softmax_fac = softmax_fac\n",
    "        self.fac = fac\n",
    "        self.nodes = nodes\n",
    "        self.mask_const = mask_const\n",
    "        self.Mask = Mask(input_size, mask_const, mask_depth, mask_width, heads, sensitivity, fac=fac)\n",
    "        \n",
    "        self.hfc = [nn.Linear(input_size, nodes[0])]\n",
    "        for i in range(len(nodes)-1):\n",
    "            self.hfc.append(nn.Linear(nodes[i], nodes[i+1]))\n",
    "            \n",
    "        \n",
    "        self.fc_softmax = nn.Linear(nodes[-1], output_sizes[0])\n",
    "        \n",
    "        if self.gamma:\n",
    "            self.gamma_layer = nn.Linear(nodes[-1], output_sizes[0])\n",
    "            # is just needed once, not for coarse graining, since it convert the same as chi\n",
    "        \n",
    "        self.coarse_grain_layer = [Coarse_grain(output_sizes[n], output_sizes[n+1]) for n in range(len(output_sizes)-1)]\n",
    "        \n",
    "        self.train_mean = torch.Tensor(train_mean)\n",
    "        self.train_std = torch.Tensor(train_std)\n",
    "    \n",
    "    def forward_before_sm(self, x):\n",
    "        \n",
    "        x = (x-self.train_mean)/self.train_std\n",
    "        shape = x.shape\n",
    "        b = shape[0]\n",
    "        \n",
    "        x = self.Mask(x) # b x input_size\n",
    "\n",
    "        \n",
    "        for layer_list_i in self.hfc:\n",
    "\n",
    "            x = F.elu(layer_list_i(x))\n",
    "            \n",
    "        return x\n",
    "        \n",
    "    def forward(self, x):\n",
    "        \n",
    "        x = self.forward_before_sm(x)\n",
    "\n",
    "        x_output = F.softmax(self.softmax_fac*self.fc_softmax(x), dim=1)\n",
    "        \n",
    "        return x_output\n",
    "    \n",
    "    def forward_cg(self, x, id):\n",
    "        \n",
    "        x_output = self.coarse_grain_layer[id](x)\n",
    "        \n",
    "        return x_output\n",
    "    \n",
    "    def forward_all(self, x):\n",
    "        outputs = []\n",
    "        \n",
    "        x = self.forward(x)\n",
    "        \n",
    "        outputs.append(x)\n",
    "        \n",
    "        for layer in self.coarse_grain_layer:\n",
    "            x = layer(x)\n",
    "            outputs.append(x)\n",
    "            \n",
    "        return outputs\n",
    "    \n",
    "    def forward_gamma(self, x):\n",
    "        \n",
    "        x_gamma = F.elu(self.gamma_layer(x)) + 1. # plus 1 so it is always positive\n",
    "        \n",
    "        return x_gamma\n",
    "    \n",
    "        \n",
    "    def get_attention(self, x=None):\n",
    "        return self.Mask.get_softmax(x=x)\n",
    "    \n",
    "    def set_soft_fac(self, new_value):\n",
    "        self.softmax_fac = torch.Tensor(new_value)\n",
    "        \n",
    "    def set_soft_fac_cg(self, id_cg, new_value):\n",
    "        \n",
    "        self.coarse_grain_layer[id_cg].sen = new_value\n",
    "    \n",
    "    def get_params_vamp(self):\n",
    "        if self.mask_const:\n",
    "            for param in self.Mask.parameters():\n",
    "                yield param\n",
    "        else:\n",
    "            for layer in self.Mask.hfc:\n",
    "                for param in layer.parameters():\n",
    "                    yield param\n",
    "            for param in self.Mask.softmax.parameters():\n",
    "                yield param\n",
    "                    \n",
    "        for layer in self.hfc:\n",
    "            \n",
    "            for param in layer.parameters():\n",
    "                yield param\n",
    "        \n",
    "        for param in self.fc_softmax.parameters():\n",
    "            yield param\n",
    "    \n",
    "    def get_params_DMSM(self, all=True):\n",
    "        if all:\n",
    "            if self.mask_const:\n",
    "                for param in self.Mask.parameters():\n",
    "                    yield param\n",
    "            else:\n",
    "                for layer in self.Mask.hfc:\n",
    "                    for param in layer.parameters():\n",
    "                        yield param\n",
    "                for param in self.Mask.softmax.parameters():\n",
    "                    yield param\n",
    "\n",
    "            for layer in self.hfc:\n",
    "\n",
    "                for param in layer.parameters():\n",
    "                    yield param\n",
    "\n",
    "            for param in self.fc_softmax.parameters():\n",
    "                yield param\n",
    "        for param in self.gamma_layer.parameters():\n",
    "            yield param\n",
    "\n",
    "    def get_params_rev(self, all=True, u_flag=True, S_flag=True):\n",
    "        if all:\n",
    "            if self.mask_const:\n",
    "                for param in self.Mask.parameters():\n",
    "                    yield param\n",
    "            else:\n",
    "                for layer in self.Mask.hfc:\n",
    "                    for param in layer.parameters():\n",
    "                        yield param\n",
    "                for param in self.Mask.softmax.parameters():\n",
    "                    yield param\n",
    "\n",
    "            for layer in self.hfc:\n",
    "\n",
    "                for param in layer.parameters():\n",
    "                    yield param\n",
    "\n",
    "            for param in self.fc_softmax.parameters():\n",
    "                yield param\n",
    "        if u_flag:\n",
    "            for param in self.u_layers[0].parameters():\n",
    "                yield param\n",
    "        if S_flag:\n",
    "            for param in self.S_layers[0].parameters():\n",
    "                yield param\n",
    "    \n",
    "    def get_params_all(self):\n",
    "        if self.mask_const:\n",
    "            for param in self.Mask.parameters():\n",
    "                yield param\n",
    "        else:\n",
    "            for layer in self.Mask.hfc:\n",
    "                for param in layer.parameters():\n",
    "                    yield param\n",
    "            for param in self.Mask.softmax.parameters():\n",
    "                yield param\n",
    "                    \n",
    "        for layer in self.hfc:\n",
    "            \n",
    "            for param in layer.parameters():\n",
    "                yield param\n",
    "        \n",
    "        for param in self.fc_softmax.parameters():\n",
    "            yield param\n",
    "        for layer in self.coarse_grain_layer:\n",
    "            for param in layer.parameters():\n",
    "                yield param\n",
    "        \n",
    "    \n",
    "    def get_params_wo_mask(self):\n",
    "        for layer in self.hfc:\n",
    "            for param in layer.parameters():\n",
    "                yield param\n",
    "       \n",
    "        for param in self.fc_softmax.parameters():\n",
    "            yield param\n",
    "    \n",
    "    def get_params_softmax(self):\n",
    "        \n",
    "        for param in self.fc_softmax.parameters():\n",
    "            yield param\n",
    "            \n",
    "    \n",
    "    def get_params_mask(self):\n",
    "        if self.mask_const:\n",
    "            for param in self.Mask.parameters():\n",
    "                yield param\n",
    "        else:\n",
    "            for layer in self.Mask.hfc:\n",
    "                for param in layer.parameters():\n",
    "                    yield param\n",
    "            for param in self.Mask.softmax.parameters():\n",
    "                yield param\n",
    "            \n",
    "    def get_params_cg(self, index=[]):\n",
    "        for id in index:\n",
    "            for param in self.coarse_grain_layer[id].parameters():\n",
    "                yield param\n",
    "    \n",
    "    def get_params_cg_rev(self, index=[], all=True):\n",
    "        \n",
    "        for id in index:\n",
    "            if all:\n",
    "                for param in self.coarse_grain_layer[id].parameters():\n",
    "                    yield param\n",
    "            for param in self.u_layers[id+1].parameters():\n",
    "                yield param\n",
    "            for param in self.S_layers[id+1].parameters():\n",
    "                yield param\n",
    "                \n",
    "                \n",
    "    def set_rev_var(self, layer_id=0, S=True):\n",
    "        \n",
    "        chi_t = self.forward(tensor_train_X1)\n",
    "        chi_tau = self.forward(tensor_train_X2)\n",
    "    \n",
    "        for i in range(layer_id):\n",
    "            chi_t = self.forward_cg(chi_t, i)\n",
    "            chi_tau = self.forward_cg(chi_tau, i)\n",
    "            \n",
    "        Data_chi_X = chi_t.detach().numpy()\n",
    "        Data_chi_Y = chi_tau.detach().numpy()\n",
    "        fullbatch = Data_chi_X.shape[0]\n",
    "\n",
    "\n",
    "        c_0 = 1/fullbatch * Data_chi_X.T @ Data_chi_X\n",
    "        c_tau = 1/fullbatch * Data_chi_X.T @ Data_chi_Y\n",
    "        c_1 = 1/fullbatch * Data_chi_Y.T @ Data_chi_Y\n",
    "\n",
    "        eigv_all, eigvec_all = np.linalg.eigh(c_0)\n",
    "        include = eigv_all > epsilon\n",
    "        eigv = eigv_all[include]\n",
    "        eigvec = eigvec_all[:,include]\n",
    "        c0_inv = eigvec @ np.diag(1/eigv) @ np.transpose(eigvec)\n",
    "\n",
    "        K_vamp = c0_inv @ c_tau\n",
    "\n",
    "        # estimate pi, the stationary distribution vector\n",
    "        eigv, eigvec = np.linalg.eig(K_vamp.T)\n",
    "        ind_pi = np.argmin((eigv-1)**2)\n",
    "\n",
    "        pi_vec = eigvec[:,ind_pi]\n",
    "        pi = pi_vec / np.sum(pi_vec, keepdims=True)\n",
    "\n",
    "        # reverse the consruction of u \n",
    "        u_optimal = c0_inv @ pi\n",
    "        \n",
    "        if self.valid_T:\n",
    "            u_kernel = np.log(np.abs(u_optimal/self.factor_u))\n",
    "        else:\n",
    "            u_kernel = u_optimal / self.factor_u\n",
    "        \n",
    "        with torch.no_grad():\n",
    "            for param in self.u_layers[layer_id].parameters():\n",
    "            \n",
    "                param.copy_(torch.Tensor(u_kernel[None,:]))  \n",
    "            \n",
    "        if S:\n",
    "            \n",
    "            _, _, _, _, Sigma_input, _, _ = self.u_layers[layer_id](chi_t, chi_tau)\n",
    "            Sigma = Sigma_input.detach().numpy()\n",
    "            \n",
    "            eigv_all, eigvec_all = np.linalg.eigh(Sigma)\n",
    "            include = eigv_all > epsilon\n",
    "            eigv = eigv_all[include]\n",
    "            eigvec = eigvec_all[:,include]\n",
    "            sigma_inv = eigvec @ np.diag(1/eigv) @ np.transpose(eigvec)\n",
    "\n",
    "            # reverse the construction of S\n",
    "            S_nonrev = K_vamp @ sigma_inv\n",
    "            S_rev_add = 1/2 * (S_nonrev + S_nonrev.T)\n",
    "            if self.valid_T:\n",
    "                kernel_S = S_rev_add / 2.\n",
    "                kernel_S = np.log(np.abs(kernel_S/self.factor_S))\n",
    "            else:\n",
    "                kernel_S = S_rev_add / 2. / self.factor_S\n",
    "\n",
    "            with torch.no_grad():\n",
    "                for param in self.S_layers[layer_id].parameters():\n",
    "\n",
    "                    param.copy_(torch.Tensor(kernel_S)) \n",
    "            \n",
    "    def get_weights(self):\n",
    "        \n",
    "        weights_dict = {}\n",
    "        weights_dict['hfc'] = []\n",
    "        weights_dict['sm'] = []\n",
    "\n",
    "        for layer in  self.hfc:\n",
    "            for param in layer.parameters():\n",
    "                weights_dict['hfc'].append(param.detach().numpy().copy())\n",
    "\n",
    "        for param in self.fc_softmax.parameters():\n",
    "            weights_dict['sm'].append(param.detach().numpy().copy())\n",
    "\n",
    "        if self.reversible:\n",
    "            weights_dict['S'] = [param.detach().numpy().copy() for param in self.S_layers[0].parameters()]\n",
    "            weights_dict['u'] = [param.detach().numpy().copy() for param in self.u_layers[0].parameters()]\n",
    "\n",
    "        if self.mask_const:\n",
    "            weights_dict['Mask'] = [param.detach().numpy().copy() for param in self.Mask.parameters()]\n",
    "        else:\n",
    "            weights_dict['Mask_hf'] = []\n",
    "            weights_dict['Mask_sm'] = []\n",
    "\n",
    "            for layer in self.Mask.hfc:\n",
    "                for param in layer.parameters():\n",
    "                    weights_dict['Mask_hf'].append(param.detach().numpy().copy())\n",
    "            for param in self.Mask.softmax.parameters():\n",
    "                weights_dict['Mask_sm'].append(param.detach().numpy().copy())\n",
    "\n",
    "        weights_dict['cg'] = []\n",
    "        weights_dict['S_cg'] = []\n",
    "        weights_dict['u_cg'] = []\n",
    "        for i, layer in enumerate(self.coarse_grain_layer):\n",
    "            for param in layer.parameters():\n",
    "                weights_dict['cg'].append(param.detach().numpy().copy())\n",
    "            weights_dict['S_cg'].append([param.detach().numpy().copy() for param in self.S_layers[i+1].parameters()][0])\n",
    "            weights_dict['u_cg'].append([param.detach().numpy().copy() for param in self.u_layers[i+1].parameters()][0])\n",
    "        \n",
    "        return weights_dict\n",
    "    \n",
    "    def set_weights(self, weights_dict):\n",
    "        with torch.no_grad():\n",
    "            i = 0\n",
    "            for layer in self.hfc:\n",
    "                for param in layer.parameters():    \n",
    "                    param.copy_(torch.Tensor(weights_dict['hfc'][i])) \n",
    "                    i+=1\n",
    "            i = 0\n",
    "            for param in self.fc_softmax.parameters():\n",
    "                param.copy_(torch.Tensor(weights_dict['sm'][i]))\n",
    "                i+=1\n",
    "\n",
    "            if self.reversible:\n",
    "                i=0\n",
    "                for param in self.S_layers[0].parameters():\n",
    "                    param.copy_(torch.Tensor(weights_dict['S'][i]))\n",
    "                    i+=1\n",
    "                i=0\n",
    "                for param in self.u_layers[0].parameters():\n",
    "                    param.copy_(torch.Tensor(weights_dict['u'][i]))\n",
    "                    i+=1\n",
    "\n",
    "            if self.mask_const:\n",
    "                i=0\n",
    "                for param in self.Mask.parameters():\n",
    "                    param.copy_(torch.Tensor(weights_dict['Mask'][i]))\n",
    "                    i+=1\n",
    "            else:\n",
    "                i=0\n",
    "                for layer in self.Mask.hfc:\n",
    "                    for param in layer.parameters():\n",
    "                        param.copy_(torch.Tensor(weights_dict['Mask_hf'][i]))\n",
    "                        i+=1\n",
    "                i=0\n",
    "                for param in self.Mask.softmax.parameters():\n",
    "                    param.copy_(torch.Tensor(weights_dict['Mask_sm'][i]))\n",
    "                    i+=1\n",
    "            i=0\n",
    "            for layer in self.coarse_grain_layer:\n",
    "\n",
    "                for param in layer.parameters():\n",
    "                    param.copy_(torch.Tensor(weights_dict['cg'][i]))\n",
    "                for param in self.S_layers[i+1].parameters():\n",
    "                    param.copy_(torch.Tensor(weights_dict['S_cg'][i]))\n",
    "                for param in self.u_layers[i+1].parameters():\n",
    "                    param.copy_(torch.Tensor(weights_dict['u_cg'][i]))\n",
    "                i+=1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def VAMP_score(chi_t, chi_tau, corr=False):\n",
    "    '''Calculates the VAMP-2 score with respect to the network lobes while \n",
    "    symmetrizing the correlation matrices. Can be used as a loss function\n",
    "    for keras models.\n",
    "    Parameters\n",
    "    ----------\n",
    "    chi_t: tensorflow tensor.\n",
    "        parameter not needed for the calculation, added to comply with Keras\n",
    "        rules for loss fuctions format.\n",
    "\n",
    "    chi_tau: tensorflow tensor with shape [batch_size, 2 * output_size]\n",
    "        output of the two lobes of the network\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    loss_score: tensorflow tensor with shape [1].\n",
    "    '''\n",
    "    shape = chi_t.shape\n",
    "    \n",
    "        \n",
    "    batch_size = shape[0]\n",
    "\n",
    "\n",
    "\n",
    "    x, y = _prep_data(chi_t, chi_tau) \n",
    "\n",
    "    # Calculate the covariance matrices\n",
    "    cov_00 = 1/(batch_size - 1) * torch.matmul(x.T, x) \n",
    "    cov_11 = 1/(batch_size - 1) * torch.matmul(y.T, y)\n",
    "    cov_01 = 1/(batch_size - 1) * torch.matmul(x.T, y)\n",
    "\n",
    "    # Calculate the inverse of the self-covariance matrices\n",
    "    cov_00_inv = _inv(cov_00, ret_sqrt = True)\n",
    "    cov_11_inv = _inv(cov_11, ret_sqrt = True)\n",
    "    \n",
    "\n",
    "    # Estimate Vamp-matrix\n",
    "    vamp_matrix = torch.matmul(cov_00_inv, torch.matmul(cov_01, cov_11_inv))\n",
    "    \n",
    "    \n",
    "    vamp_score = torch.norm(vamp_matrix)\n",
    "    score = (vamp_score**2).unsqueeze(0)\n",
    "    if corr:\n",
    "        return score, torch.trace(cov_00)\n",
    "    else:\n",
    "        return score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def vampe_loss(chi, gamma):\n",
    "    ''' Estimates the vampe_loss for a deepMSM'''\n",
    "    b = chi.shape[0]\n",
    "    \n",
    "    c00 = 1/b*(torch.matmul(chi.T, chi))\n",
    "    c11 = 1/b*(torch.matmul(gamma.T, gamma))\n",
    "    c01 = 1/b*(torch.matmul(chi.T, gamma))\n",
    "\n",
    "    gamma_dia_inv = torch.diag(1/(torch.mean(gamma, dim=0)))  # add something so no devide by zero\n",
    "\n",
    "    first_term = c00 @ gamma_dia_inv @ c11 @ gamma_dia_inv\n",
    "    second_term = 2 * (c01 @ gamma_dia_inv)\n",
    "    vampe_arg = first_term - second_term\n",
    "    vampe = torch.trace(vampe_arg)\n",
    "    \n",
    "    return -vampe"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def vampe_loss_rev(chi_t, chi_tau, layer_id=0, return_mu=False, return_mu_K_Sigma=False):\n",
    "    '''Estimates the VAMP-E loss for a reversible model.'''\n",
    "    v, C_00, C_11, C_01, Sigma, mu_t, Sigma_t = Full_net.u_layers[layer_id](chi_t, chi_tau)\n",
    "    matrix, K, _ = Full_net.S_layers[layer_id](v, C_00, C_11, C_01, Sigma)\n",
    "    vampe = torch.trace(matrix)\n",
    "    \n",
    "    if return_mu:\n",
    "        return -vampe, mu_t\n",
    "    \n",
    "    elif return_mu_K_Sigma:\n",
    "        return -vampe, mu_t, K, Sigma_t\n",
    "    \n",
    "    else:\n",
    "        return -vampe\n",
    "    \n",
    "def vampe_loss_rev_only_S(v, C_00, C_11, C_01, Sigma, layer_id=0):\n",
    "    '''Estimates the VAMP-E loss for a reversible model in the case if only S is supposed to be trained.'''\n",
    "    matrix, K, _ = Full_net.S_layers[layer_id](v, C_00, C_11, C_01, Sigma)\n",
    "#     print(K)\n",
    "    vampe = torch.trace(matrix)\n",
    "    \n",
    "    return -vampe"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fac_bf_sm = True\n",
    "valid_T=True # if valid transition matrix is enforced\n",
    "reversible=True # if reversibility is enforced\n",
    "\n",
    "# attention stuff\n",
    "mask_const=True # if the trained attention mask is constant over time\n",
    "heads=1\n",
    "mask_depth=2 # if time dependent how many hidden layers has the attention network\n",
    "mask_width=100 # the width of the attention hidden layers\n",
    "sensitivity=1. # factor before attention softmax for clearer assignment\n",
    "factor_att=True # if to use a factor which scales the input on average back to input\n",
    "\n",
    "softmax_fac=1. # factor before classification softmax\n",
    "\n",
    "train_mean = traj_whole[0].mean(0)\n",
    "train_std = traj_whole[0].std(0)\n",
    "\n",
    "Full_net = VampNet(input_size, output_sizes, nodes, train_mean, train_std, \n",
    "                 valid_T=valid_T, reversible=reversible,\n",
    "                 mask_const=mask_const, mask_depth=mask_depth, mask_width=mask_width, heads=heads, sensitivity=sensitivity,\n",
    "                 fac=factor_att,\n",
    "                 softmax_fac=softmax_fac)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "weights_dict = Full_net.get_weights()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer_vamp = optim.Adam(Full_net.get_params_vamp(), lr=learning_rate*10)\n",
    "optimizer_full = optim.Adam(Full_net.get_params_all(), lr=learning_rate) # also for all coarse grain\n",
    "\n",
    "if Full_net.reversible:\n",
    "    optimizer_rev = optim.Adam(Full_net.get_params_rev(), lr=learning_rate/10)\n",
    "    optimizer_rev_u_S = optim.Adam(Full_net.get_params_rev(all=False), lr=learning_rate)\n",
    "    optimizer_rev_S = optim.Adam(Full_net.get_params_rev(all=False, u_flag=False), lr=learning_rate*10)\n",
    "    optimizer_rev_u = optim.Adam(Full_net.get_params_rev(all=False, S_flag=False), lr=learning_rate*10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Estimation of the real value of the observables and Data manipulation\n",
    "\n",
    "1. Load a model trained on the full data set and estimate the real observable values\n",
    "2. Identify a good microscopic observable which could be possibly measured in an experiment (here a contact)\n",
    "3. Manipulate the data by removing folding and unfolding events from the data by observing the eigenfunction corresponding to the folding process and identify the frames which represent the crossing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# First load the model estimated on the whole data set\n",
    "weights_msm = np.load('./dicts/tau2500skip1msm.npy', allow_pickle=True).item()\n",
    "# We will use a pre trained network on the whole trajectory to estimate the true observable values\n",
    "train_mean = np.load('./dicts/tau2500skip1msm_mean.npy')\n",
    "train_std = np.load('./dicts/tau2500skip1msm_std.npy')\n",
    "\n",
    "Full_net.set_weights(weights_msm)\n",
    "Full_net.train_mean=train_mean\n",
    "Full_net.train_std=train_std\n",
    "tau = 50*2*25//skip\n",
    "chi_true = Full_net.forward(torch.Tensor(traj_whole_new[0])).detach()\n",
    "chi_whole = chi_true.numpy()\n",
    "distances = - np.log(traj_whole_new[0])\n",
    "contacts = (distances <0.45).astype('int')\n",
    "# Define which distance is the microscopic observable\n",
    "contact_obs = [41]\n",
    "# Define helper functions to estimate the observables\n",
    "def plot_mu(tensor_train_X1, tensor_train_X2, frames):\n",
    "    ''' Estimates the state probability of a reference model. The stationary distribution\n",
    "    mu is estimated from the current model, but the state assignment stems from the\n",
    "    reference model. This makes it comparable over several models.\n",
    "    '''\n",
    "    chi_t = Full_net(tensor_train_X1)\n",
    "    chi_tau = Full_net(tensor_train_X2)\n",
    "    _, _, _, _, _, mu_t, _ = Full_net.u_layers[0](chi_t, chi_tau)\n",
    "    mu = mu_t\n",
    "    state_prob = torch.sum(mu * chi_true[frames], dim=0).detach().numpy()\n",
    "#     plt.plot(state_prob, '.')\n",
    "#     plt.show()\n",
    "    return state_prob\n",
    "# True stationary distribution of the three states\n",
    "prob_states_true = plot_mu(torch.Tensor(traj_whole_new[0][:-tau]), torch.Tensor(traj_whole_new[0][tau:]), np.arange(traj_whole_new[0].shape[0]-tau))\n",
    "# functions to estimate the difference of the true observable value and the one from the current model\n",
    "def obs_loss(obs_value, mu, exp_value):\n",
    "    \n",
    "    exp_value_estimated = torch.sum(obs_value * mu)\n",
    "    \n",
    "    error = torch.abs(exp_value - exp_value_estimated)\n",
    "    \n",
    "    return error\n",
    "\n",
    "def obs_time_loss(obs_value, mu, chi, K, Sigma, exp_value):\n",
    "\n",
    "    state_weight = mu*chi\n",
    "\n",
    "    pi = torch.sum(state_weight, dim=0) # prob to be in a state\n",
    "    # obs value within a state, the weighting factor needs to be normalized for each state\n",
    "    ai = torch.sum(state_weight*obs_value, dim=0, keepdims=True) / torch.sum(state_weight, dim=0, keepdims=True)\n",
    "    # prob to observe an unconditional jump state i to j\n",
    "    X = Sigma @ K\n",
    "    a_sim = torch.matmul(ai, torch.matmul(X, ai.T)) # shape 1x1\n",
    "        \n",
    "    error = torch.abs(a_sim - exp_value) \n",
    "           \n",
    "    return error\n",
    "\n",
    "def vampe_loss_rev_obs(chi_t, chi_tau, layer_id=0, return_mu=False, all_eigvals=False):\n",
    "    \n",
    "    v, C_00, C_11, C_01, Sigma, mu_t, Sigma_t = Full_net.u_layers[layer_id](chi_t, chi_tau)\n",
    "    matrix, K, S = Full_net.S_layers[layer_id](v, C_00, C_11, C_01, Sigma)\n",
    "    vampe = torch.trace(matrix)\n",
    "    \n",
    "    \n",
    "    eigval_all, eigvec_all = torch.symeig(Sigma, eigenvectors=True)\n",
    "\n",
    "    # Filter out eigvalues below threshold and corresponding eigvectors\n",
    "    eig_th = torch.Tensor(epsilon)\n",
    "    index_eig = eigval_all > eig_th\n",
    "#     print(index_eig)\n",
    "    eigval = eigval_all[index_eig]\n",
    "    eigvec = eigvec_all[:,index_eig]\n",
    "\n",
    "    # Build the diagonal matrix with the filtered eigenvalues or square\n",
    "    # root of the filtered eigenvalues according to the parameter\n",
    "\n",
    "    diag = torch.diag(torch.sqrt(eigval))\n",
    "\n",
    "#     print(diag.shape, eigvec.shape)    \n",
    "    # Rebuild the square root of the inverse matrix\n",
    "    Sigma_sqrt = torch.matmul(eigvec, torch.matmul(diag, eigvec.T))\n",
    "    \n",
    "    S_similar = Sigma_sqrt @ S @ Sigma_sqrt\n",
    "    \n",
    "    eigval_all, eigvec_all = torch.symeig(S_similar, eigenvectors=True)\n",
    "    \n",
    "#     sort_eigvals = torch.sort(eigval_all)\n",
    "    if all_eigvals:\n",
    "        ret = [-vampe, S_similar, K, eigval_all]\n",
    "    else:\n",
    "        ret = [-vampe, S_similar, K, eigval_all[0]]\n",
    "    \n",
    "    return ret\n",
    "\n",
    "chi_t = chi_true[:-tau]\n",
    "chi_tau = chi_true[tau:]\n",
    "\n",
    "\n",
    "_, S_similar, K_test, eigval_fold = vampe_loss_rev_obs(chi_t, chi_tau)\n",
    "# True eigenvalue of the folding process\n",
    "true_value = eigval_fold.detach()\n",
    "\n",
    "def get_K_rev(tensor_t=torch.Tensor(traj_whole_new[0][:-tau]), tensor_tau=torch.Tensor(traj_whole_new[0][tau:])):\n",
    "    \n",
    "    chi_t = Full_net(tensor_t)\n",
    "    chi_tau = Full_net(tensor_tau)\n",
    "    v, C_00, C_11, C_01, Sigma, _, _ = Full_net.u_layers[0](chi_t, chi_tau)\n",
    "    _, K, _ = Full_net.S_layers[0](v, C_00, C_11, C_01, Sigma)\n",
    "    \n",
    "    return K.detach()\n",
    "# microscopic observable value for each time frame, if a contact is formed or not\n",
    "obs_value_time = [(contacts[:,i][:,None]).astype('float32') for i in contact_obs]\n",
    "# and the reverse observable\n",
    "obs_value_time.append(((obs_value_time[0]-1.)*-1.).astype('float32'))\n",
    "\n",
    "\n",
    "\n",
    "chi_t = torch.Tensor(chi_whole[:-tau])\n",
    "chi_tau = torch.Tensor(chi_whole[tau:])\n",
    "# Estimate the true auto correlation value of the contact staying formed\n",
    "exp_value_time = []\n",
    "for obs_v in obs_value_time:\n",
    "    obs_value_tensor = torch.Tensor(obs_v[:-tau])\n",
    "    exp_value_tensor = torch.Tensor(np.array([0.]))\n",
    "\n",
    "    score_curr, mu_t, K, Sigma = vampe_loss_rev(chi_t, chi_tau, return_mu_K_Sigma=True)\n",
    "    error_obs = obs_time_loss(obs_value_tensor, mu_t, chi_t, K, Sigma, exp_value_tensor)\n",
    "    exp_value_time.append(error_obs.detach())\n",
    "print(exp_value_time)\n",
    "# Estimate the expectation value that the contact is formed\n",
    "obs_exp_value = [(contacts[:,i][:,None]).astype('float32') for i in contact_obs]\n",
    "exp_exp_value = []\n",
    "for obs_v in obs_exp_value:\n",
    "    obs_value_tensor = torch.Tensor(obs_v[:-tau])\n",
    "    exp_value_tensor = torch.Tensor(np.array([0.]))\n",
    "\n",
    "    score_curr, mu_t = vampe_loss_rev(chi_t, chi_tau, return_mu=True)\n",
    "    error_obs = obs_loss(obs_value_tensor, mu_t, exp_value_tensor)\n",
    "    exp_exp_value.append(error_obs.detach())\n",
    "print(exp_exp_value)\n",
    "\n",
    "for obs in obs_value_time:\n",
    "#     plt.plot(obs)\n",
    "#     plt.show()\n",
    "    print(np.sum(obs[:-tau]*obs[tau:])/(obs[tau:].shape[0]))\n",
    "# remove data depending on where you land\n",
    "\n",
    "K_exp = get_K_rev()\n",
    "np.linalg.eigvals(K_exp)\n",
    "eigvals, eigvec = np.linalg.eig(K_exp)\n",
    "sort_ind = np.argsort(eigvals)\n",
    "eigvals_sort = eigvals[sort_ind]\n",
    "eigvec_sort = eigvec[:,sort_ind]\n",
    "\n",
    "# Manipulate the data\n",
    "data_corrupt = 'barrier' \n",
    "if data_corrupt == 'barrier':\n",
    "    K = get_K_rev(torch.Tensor(traj_whole[0][:-tau]), torch.Tensor(traj_whole[0][tau:]))\n",
    "    # K = estimate_koopman_op(chi_whole, tau)\n",
    "    np.linalg.eigvals(K)\n",
    "    eigvals, eigvecs = np.linalg.eig(K)\n",
    "    sort_id = np.argsort(eigvals)\n",
    "    eigvals_sort = eigvals[sort_id]\n",
    "    eigvecs_sort = eigvecs[:,sort_id]\n",
    "    print(eigvals_sort[-3])\n",
    "    # Estimate the eigenfunction corresponding to the folding process\n",
    "    eigfunc = chi_whole @ eigvecs_sort[:,-3]\n",
    "    min_eigfunc = eigfunc.min()\n",
    "    max_eigfunc = eigfunc.max()\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",
    "    \n",
    "    # truncate the whole trajectory into folding and unfolding events\n",
    "    transition_forward_ind = []\n",
    "    transition_backward_ind = []\n",
    "\n",
    "    for j, s in enumerate(starting_points):\n",
    "\n",
    "        distance_end = end_points - s\n",
    "\n",
    "        where_positive = np.where(distance_end>0)[0]\n",
    "        if len(where_positive)>0:\n",
    "            e = end_points[where_positive[0]]\n",
    "            if j+1 < starting_points.shape[0]:\n",
    "                next_s = starting_points[j+1]\n",
    "                if next_s > e:\n",
    "                    s_new = s-tau\n",
    "                    if s_new < 0:\n",
    "                        s_new = 0\n",
    "                    transition_forward_ind.append(np.arange(s_new, e))\n",
    "                    print('Found new transition {}, {}'.format(s_new, e))\n",
    "\n",
    "\n",
    "    for j, s in enumerate(end_points):\n",
    "\n",
    "        distance_end = starting_points - s\n",
    "\n",
    "        where_positive = np.where(distance_end>0)[0]\n",
    "        if len(where_positive)>0:\n",
    "            e = starting_points[where_positive[0]]\n",
    "            if j+1 < end_points.shape[0]:\n",
    "                next_s = end_points[j+1]\n",
    "                if next_s > e:\n",
    "                    s_new = s-tau\n",
    "                    if s_new < 0:\n",
    "                        s_new = 0\n",
    "                    transition_backward_ind.append(np.arange(s_new, e))\n",
    "                    print('Found new back transition {}, {}'.format(s_new, e))\n",
    "\n",
    "    # Find the frames which are not part of the folding/unfolding\n",
    "    non_transition_ind = []\n",
    "\n",
    "    len_for = len(transition_forward_ind)\n",
    "    len_back = len(transition_backward_ind)\n",
    "\n",
    "    if len_for > len_back:\n",
    "        length_total = len_for\n",
    "    else:\n",
    "        length_total = len_back\n",
    "    s = 0\n",
    "    for i in range(length_total):\n",
    "        if i < len_for:\n",
    "            forward = transition_forward_ind[i]\n",
    "        if i < len_back:\n",
    "            backward = transition_backward_ind[i]\n",
    "\n",
    "        if forward[0] < backward[0]:\n",
    "            e = forward[0]\n",
    "            non_transition_ind.append(np.arange(s,e))\n",
    "            s = forward[-1]+1\n",
    "\n",
    "            if s < backward[0]:\n",
    "                e = backward[0]\n",
    "                non_transition_ind.append(np.arange(s,e))\n",
    "            s = backward[-1]\n",
    "        else:\n",
    "            e = backward[0]\n",
    "            non_transition_ind.append(np.arange(s,e))\n",
    "            s = backward[-1]+1\n",
    "\n",
    "            if s < forward[0]:\n",
    "                e = forward[0]\n",
    "                non_transition_ind.append(np.arange(s,e))\n",
    "            s = forward[-1]\n",
    "\n",
    "    non_transition_ind.append(np.arange(s+tau,traj_whole[0].shape[0]-tau))\n",
    "\n",
    "\n",
    "    non_transition_ind = np.concatenate(non_transition_ind)\n",
    "    test_for = np.concatenate(transition_forward_ind)\n",
    "    test_back = np.concatenate(transition_backward_ind)\n",
    "    print(test_for.shape, test_back.shape, non_transition_ind.shape)\n",
    "    plt.plot(eigfunc,'.')\n",
    "    plt.plot(non_transition_ind, eigfunc[non_transition_ind], '.')\n",
    "    plt.plot(test_for, eigfunc[test_for], '.')\n",
    "    plt.plot(test_back, eigfunc[test_back], '.')\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "    ind_train = []\n",
    "    ind_valid = []\n",
    "    ind_test = []\n",
    "    non_length = non_transition_ind.shape[0]//3\n",
    "    np.random.shuffle(non_transition_ind)\n",
    "    ind_train.append(non_transition_ind[:non_length])\n",
    "    ind_valid.append(non_transition_ind[non_length:2*non_length])\n",
    "    ind_test.append(non_transition_ind[2*non_length:])\n",
    "    # add folding events with a probability of p_for, which influences how much less often a folding\n",
    "    # will occur in the fake trajectory\n",
    "    p_for = 0.25\n",
    "\n",
    "    nr_for = int(p_for*len(transition_forward_ind)//3)\n",
    "    print('Number of transitions: {}'.format(nr_for))\n",
    "    ind_trajs_temp = np.arange(len(transition_forward_ind))\n",
    "    np.random.shuffle(ind_trajs_temp)\n",
    "    for i in range(nr_for):\n",
    "        ind_train.append(transition_forward_ind[ind_trajs_temp[i]])\n",
    "\n",
    "        ind_valid.append(transition_forward_ind[ind_trajs_temp[i+nr_for]])\n",
    "\n",
    "        ind_test.append(transition_forward_ind[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_ind)//3)\n",
    "    print('Number of transitions: {}'.format(nr_for))\n",
    "    ind_trajs_temp = np.arange(len(transition_backward_ind))\n",
    "    np.random.shuffle(ind_trajs_temp)\n",
    "    for i in range(nr_back):\n",
    "        ind_train.append(transition_backward_ind[ind_trajs_temp[i]])\n",
    "\n",
    "        ind_valid.append(transition_backward_ind[ind_trajs_temp[i+nr_back]])\n",
    "\n",
    "        ind_test.append(transition_backward_ind[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",
    "    plt.plot(ind_train, eigfunc[ind_train], '.')\n",
    "    plt.plot(ind_valid, eigfunc[ind_valid], '.')\n",
    "    plt.plot(ind_test, eigfunc[ind_test], '.')\n",
    "    plt.show()\n",
    "    \n",
    "    np.random.shuffle(ind_train)\n",
    "    np.random.shuffle(ind_valid)\n",
    "    np.random.shuffle(ind_test)\n",
    "    def get_data_for_tau_corrupt(single_traj, obs_value_time, tau, ind_train, ind_valid, ind_test):\n",
    "\n",
    "\n",
    "\n",
    "        obs_ord = []\n",
    "        for ob in obs_value_time:\n",
    "            obs_ord.append(ob[:-tau])\n",
    "        traj_ord = single_traj[:-tau]\n",
    "        traj_ord_lag = single_traj[tau:]\n",
    "\n",
    "        traj_data_train = traj_ord[ind_train]\n",
    "        traj_data_train_lag = traj_ord_lag[ind_train]\n",
    "        obs_train = []\n",
    "        for ob in obs_ord:\n",
    "            obs_train.append(ob[ind_train].astype('float32'))\n",
    "        \n",
    "        traj_data_valid = traj_ord[ind_valid]\n",
    "        traj_data_valid_lag = traj_ord_lag[ind_valid]\n",
    "        \n",
    "        obs_valid = []\n",
    "        for ob in obs_ord:\n",
    "            obs_valid.append(ob[ind_valid].astype('float32'))\n",
    "        # Input of the first network\n",
    "        \n",
    "        \n",
    "        traj_data_test = traj_ord[ind_test]\n",
    "        traj_data_test_lag = traj_ord_lag[ind_test]\n",
    "        obs_test = []\n",
    "        for ob in obs_ord:\n",
    "            obs_test.append(ob[ind_test].astype('float32'))\n",
    "        \n",
    "        X1_train = traj_data_train.astype('float32')\n",
    "        X2_train  = traj_data_train_lag.astype('float32')\n",
    "\n",
    "        # Input for validation\n",
    "        X1_vali = traj_data_valid.astype('float32')\n",
    "        X2_vali = traj_data_valid_lag.astype('float32')\n",
    "    \n",
    "        X1_test = traj_data_test.astype('float32')\n",
    "        X2_test = traj_data_test_lag.astype('float32')\n",
    "\n",
    "        return X1_train, X2_train, obs_train, X1_vali, X2_vali, obs_valid, X1_test, X2_test, obs_test\n",
    "\n",
    "\n",
    "    X1_train_cor, X2_train_cor, obs_train, X1_vali_cor, X2_vali_cor, obs_valid, X1_test_cor, X2_test_cor, obs_test = get_data_for_tau_corrupt(traj_whole_new[0], obs_exp_value, tau, ind_train, ind_valid, ind_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('The true eigenvalue of the folding process: {:.3}'.format(true_value))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The indeces for the test trajectory for evaluating the performance on\n",
    "frames_test = ind_test"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Get data for training chi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if data_corrupt == 'barrier':     \n",
    "    def get_data_for_tau(single_traj, tau, ind_train, ind_valid, ind_test):\n",
    "\n",
    "\n",
    "\n",
    "    \n",
    "        traj_ord = single_traj[:-tau]\n",
    "        traj_ord_lag = single_traj[tau:]\n",
    "\n",
    "\n",
    "        traj_data_train = traj_ord[ind_train]\n",
    "        traj_data_train_lag = traj_ord_lag[ind_train]\n",
    "        \n",
    "        traj_data_valid = traj_ord[ind_valid]\n",
    "        traj_data_valid_lag = traj_ord_lag[ind_valid]\n",
    "        \n",
    "        \n",
    "        traj_data_test = traj_ord[ind_test]\n",
    "        traj_data_test_lag = traj_ord_lag[ind_test]\n",
    "        \n",
    "        X1_train = traj_data_train.astype('float32')\n",
    "        X2_train  = traj_data_train_lag.astype('float32')\n",
    "\n",
    "        # Input for validation\n",
    "        X1_vali = traj_data_valid.astype('float32')\n",
    "        X2_vali = traj_data_valid_lag.astype('float32')\n",
    "    \n",
    "        X1_test = traj_data_test.astype('float32')\n",
    "        X2_test = traj_data_test_lag.astype('float32')\n",
    "        \n",
    "\n",
    "        return X1_train, X2_train, X1_vali, X2_vali, X1_test, X2_test\n",
    "\n",
    "\n",
    "    X1_train, X2_train, X1_vali, X2_vali, X1_test, X2_test = get_data_for_tau(traj_whole_new[0], tau_chi, ind_train, ind_valid, ind_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tensor_train_X1 = torch.Tensor(X1_train)\n",
    "tensor_train_X2 = torch.Tensor(X2_train) # transform to torch tensor\n",
    "tensor_valid_X1 = torch.Tensor(X1_vali)\n",
    "tensor_valid_X2 = torch.Tensor(X2_vali)\n",
    "tensor_test_X1 = torch.Tensor(X1_test)\n",
    "tensor_test_X2 = torch.Tensor(X2_test)\n",
    "\n",
    "trainset = data.TensorDataset(tensor_train_X1, tensor_train_X2) # create your datset\n",
    "trainloader = data.DataLoader(trainset, batch_size=batch_size,\n",
    "                              shuffle=True, num_workers=2)\n",
    "\n",
    "trainloader_full = data.DataLoader(trainset, batch_size=300000,\n",
    "                              shuffle=True, num_workers=2)\n",
    "\n",
    "testset = data.TensorDataset(tensor_valid_X1, tensor_valid_X2) # create your datset\n",
    "testloader = data.DataLoader(testset, batch_size=batch_size,\n",
    "                             shuffle=True, num_workers=2)\n",
    "\n",
    "full_batch = False\n",
    "if full_batch:\n",
    "    train_l = trainloader_full\n",
    "else:\n",
    "    train_l = trainloader"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Trianing loops for the different models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_for_VAMPnet(runs, plot_mask_every=10, verbose=True, plot_training=True, corr=False, best_weights_flag=False):\n",
    "    \n",
    "    epoch_loss = np.zeros(runs)\n",
    "    epoch_loss_corr = np.zeros(runs)\n",
    "    epoch_loss_valid = np.zeros(runs)\n",
    "    epoch_loss_valid_corr = np.zeros(runs)\n",
    "#     sen = np.linspace(.1,2,10)\n",
    "    if best_weights_flag:\n",
    "        best_score=0.\n",
    "        best_weights = Full_net.get_weights()\n",
    "    for epoch in range(runs):  # loop over the dataset multiple times\n",
    "\n",
    "#         Full_net.\n",
    "        opt = optimizer_vamp\n",
    "    #     sen_temp = sen_set[epoch//sen_every]\n",
    "\n",
    "#         Full_net.set_soft_fac(sen[[epoch]])\n",
    "\n",
    "\n",
    "        running_epoch_loss = []\n",
    "        running_epoch_corr = []\n",
    "        for i, data_batch in enumerate(train_l, 0):\n",
    "            # get the inputs; data is a list of [inputs, labels]\n",
    "            inputs_t, inputs_tau = data_batch\n",
    "\n",
    "            # zero the parameter gradients\n",
    "\n",
    "            opt.zero_grad()\n",
    "\n",
    "            # estimate weights\n",
    "            chi_t = Full_net(inputs_t)\n",
    "            chi_tau = Full_net(inputs_tau)\n",
    "\n",
    "            score_list = VAMP_score(chi_t, chi_tau, corr=corr)\n",
    "\n",
    "            if corr:\n",
    "                score_curr = score_list[0]\n",
    "                score_corr = score_list[1]\n",
    "                loss = -score_curr - score_corr\n",
    "                running_epoch_corr.append(score_corr.item())\n",
    "            else:\n",
    "                score_curr = score_list\n",
    "                loss = - score_curr\n",
    "\n",
    "            loss.backward()\n",
    "\n",
    "            opt.step()\n",
    "\n",
    "            running_epoch_loss.append(score_curr.item())\n",
    "\n",
    "        # validation\n",
    "        chi_t_vali = Full_net(tensor_valid_X1).detach()\n",
    "        chi_tau_vali = Full_net(tensor_valid_X2).detach()\n",
    "        loss_vali_list = VAMP_score(chi_t_vali, chi_tau_vali, corr=corr)\n",
    "        del chi_t_vali\n",
    "        del chi_tau_vali\n",
    "        del chi_t\n",
    "        del chi_tau\n",
    "        if corr:\n",
    "            loss_vali = loss_vali_list[0].detach()\n",
    "            loss_vali_corr = loss_vali_list[1].detach()\n",
    "            epoch_loss_valid_corr[epoch] = loss_vali_corr.item()\n",
    "        else:\n",
    "            loss_vali = loss_vali_list.detach()\n",
    "        epoch_loss_valid[epoch] = loss_vali.item()\n",
    "        if best_weights_flag:\n",
    "            if epoch_loss_valid[epoch]>best_score:\n",
    "                print('Better validation score, save weights')\n",
    "                best_weights = Full_net.get_weights()\n",
    "                best_score = epoch_loss_valid[epoch]\n",
    "            \n",
    "        epoch_loss[epoch] = np.mean(running_epoch_loss)\n",
    "        epoch_loss_corr[epoch] = np.mean(running_epoch_corr)\n",
    "        if verbose:\n",
    "            print('Run {}, total loss: {:.3}, valid loss: {:.3}'.format(epoch+1, epoch_loss[epoch], epoch_loss_valid[epoch]))\n",
    "\n",
    "    if plot_training:\n",
    "        plt.plot(np.arange(1,runs+1), epoch_loss, label='VAMP_loss')\n",
    "        plt.plot(np.arange(1,runs+1), epoch_loss_valid, label='VAMP_loss_valid')\n",
    "        plt.legend()\n",
    "        plt.show()\n",
    "        if corr:\n",
    "            plt.plot(np.arange(1,runs+1), epoch_loss_corr, label='Corr')\n",
    "            plt.plot(np.arange(1,runs+1), epoch_loss_corr, label='Corr_valid')\n",
    "            plt.legend()\n",
    "            plt.show()\n",
    "    if best_weights_flag:\n",
    "        print('Set best weights')\n",
    "        Full_net.set_weights(best_weights)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Put the parameters to the value before we loaded the true model\n",
    "Full_net.set_weights(weights_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Pretrain a VAMPnet which will be used as initialization for both models w/o exp. observables\n",
    "This model only sees the manipulated data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "optimizer_vamp = optim.Adam(Full_net.get_params_vamp(), lr=learning_rate)\n",
    "train_for_VAMPnet(30, corr=True)\n",
    "\n",
    "optimizer_vamp = optim.Adam(Full_net.get_params_vamp(), lr=learning_rate)\n",
    "train_for_VAMPnet(30)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_for_VAMPnet(100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "weights_dict_chi = Full_net.get_weights()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Helper functions to train for the reversible DMSM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_for_S(runs=100, verbose=True, plot_training=True):\n",
    "    \n",
    "    chi_t = Full_net(tensor_train_X1)\n",
    "    chi_tau = Full_net(tensor_train_X2)\n",
    "    v, C_00, C_11, C_01, Sigma, _, _ = Full_net.u_layers[0](chi_t, chi_tau)\n",
    "    \n",
    "    v = v.detach()\n",
    "    C_00 = C_00.detach()\n",
    "    C_11 = C_11.detach()\n",
    "    C_01 = C_01.detach()\n",
    "    Sigma = Sigma.detach()\n",
    "    \n",
    "    chi_t_valid = Full_net(tensor_valid_X1)\n",
    "    chi_tau_valid = Full_net(tensor_valid_X2)\n",
    "    v_valid, C_00_valid, C_11_valid, C_01_valid, Sigma_valid, _, _ = Full_net.u_layers[0](chi_t_valid, chi_tau_valid)\n",
    "    \n",
    "    v_valid = v_valid.detach()\n",
    "    C_00_valid = C_00_valid.detach()\n",
    "    C_11_valid = C_11_valid.detach()\n",
    "    C_01_valid = C_01_valid.detach()\n",
    "    Sigma_valid = Sigma_valid.detach()\n",
    "    \n",
    "    epoch_loss = np.zeros(runs)\n",
    "    epoch_loss_valid = np.zeros(runs)\n",
    "    opt = optimizer_rev_S\n",
    "    for epoch in range(runs):  # loop over the dataset multiple times\n",
    "\n",
    "        opt.zero_grad()\n",
    "\n",
    "\n",
    "        score_curr = vampe_loss_rev_only_S(v, C_00, C_11, C_01, Sigma)\n",
    "\n",
    "\n",
    "\n",
    "        loss = - score_curr\n",
    "\n",
    "        loss.backward()\n",
    "\n",
    "        opt.step()\n",
    "\n",
    "\n",
    "        epoch_loss[epoch] = np.mean(-loss.item())\n",
    "        \n",
    "        score_curr_valid = vampe_loss_rev_only_S(v_valid, C_00_valid, C_11_valid, C_01_valid, Sigma_valid)\n",
    "        epoch_loss_valid[epoch] = np.mean(score_curr_valid.item())\n",
    "        if verbose:\n",
    "            print('Run {}, total loss: {:.3}, valid loss: {:.3}'.format(epoch+1, epoch_loss[epoch], epoch_loss_valid[epoch]))\n",
    "\n",
    "    if plot_training:\n",
    "        plt.plot(np.arange(1,runs+1), epoch_loss, label='VAMP_loss')\n",
    "        plt.plot(np.arange(1,runs+1), epoch_loss_valid, label='VAMP_loss_valid')\n",
    "        plt.legend()\n",
    "        plt.show()    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_for_u_S(runs=100, verbose=True, plot_training=True):\n",
    "    \n",
    "    chi_t = Full_net(tensor_train_X1).detach()\n",
    "    chi_tau = Full_net(tensor_train_X2).detach()\n",
    "    \n",
    "    chi_t_valid = Full_net(tensor_valid_X1).detach()\n",
    "    chi_tau_valid = Full_net(tensor_valid_X2).detach()\n",
    "    \n",
    "    \n",
    "    \n",
    "    \n",
    "    epoch_loss = np.zeros(runs)\n",
    "    epoch_loss_valid = np.zeros(runs)\n",
    "    opt = optimizer_rev_u_S\n",
    "    for epoch in range(runs):  # loop over the dataset multiple times\n",
    "\n",
    "        opt.zero_grad()\n",
    "\n",
    "\n",
    "        score_curr = vampe_loss_rev(chi_t, chi_tau)\n",
    "\n",
    "        loss = - score_curr\n",
    "\n",
    "        loss.backward()\n",
    "\n",
    "        opt.step()\n",
    "\n",
    "        score_curr_valid = vampe_loss_rev(chi_t_valid, chi_tau_valid)\n",
    "        epoch_loss[epoch] = np.mean(-loss.item())\n",
    "        epoch_loss_valid[epoch] = np.mean(score_curr_valid.item())\n",
    "        if verbose:\n",
    "            print('Run {}, total loss: {:.3}, valid loss: {:.3}'.format(epoch+1, epoch_loss[epoch], epoch_loss_valid[epoch]))\n",
    "\n",
    "    if plot_training:\n",
    "        plt.plot(np.arange(1,runs+1), epoch_loss, label='VAMP_loss')\n",
    "        plt.plot(np.arange(1,runs+1), epoch_loss_valid, label='VAMP_loss_valid')\n",
    "        plt.legend()\n",
    "        plt.show()    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train a revDMSM without observable knowledge to compare with"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Full_net.set_weights(weights_dict_chi)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "chi_t = Full_net.forward(tensor_train_X1).detach().numpy()\n",
    "chi_tau = Full_net.forward(tensor_train_X2).detach().numpy()\n",
    "K_vamp = estimate_koopman_op([chi_t, chi_tau], 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if Full_net.reversible:\n",
    "    optimizer_rev = optim.Adam(Full_net.get_params_rev(), lr=learning_rate/10)\n",
    "    \n",
    "    optimizer_rev_S = optim.Adam(Full_net.get_params_rev(all=False, u_flag=False), lr=learning_rate*10)\n",
    "    optimizer_rev_u = optim.Adam(Full_net.get_params_rev(all=False, S_flag=False), lr=learning_rate*10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# First initialize u and S with the estimates of the VAMPnet, then train for S\n",
    "if Full_net.reversible:\n",
    "    Full_net.set_rev_var(S=True)\n",
    "    train_for_S(runs=3000, verbose=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "weights_after_S = Full_net.get_weights()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_for_rev(runs, plot_mask_every=10, verbose=True, plot_training=True):\n",
    "    \n",
    "    epoch_loss = np.zeros(runs)\n",
    "    epoch_loss_valid = np.zeros(runs)\n",
    "\n",
    "    for epoch in range(runs):  # loop over the dataset multiple times\n",
    "\n",
    "\n",
    "        opt = optimizer_rev\n",
    " \n",
    "\n",
    "        running_epoch_loss = []\n",
    "        \n",
    "        for i, data_batch in enumerate(trainloader_full, 0):\n",
    "            # get the inputs; data is a list of [inputs, labels]\n",
    "            inputs_t, inputs_tau = data_batch\n",
    "\n",
    "            # zero the parameter gradients\n",
    "\n",
    "            opt.zero_grad()\n",
    "\n",
    "            # estimate weights\n",
    "            chi_t = Full_net(inputs_t)\n",
    "            chi_tau = Full_net(inputs_tau)\n",
    "\n",
    "            score_curr = vampe_loss_rev(chi_t, chi_tau)\n",
    "\n",
    "\n",
    "\n",
    "            loss = - score_curr\n",
    "\n",
    "            loss.backward()\n",
    "\n",
    "            opt.step()\n",
    "\n",
    "            running_epoch_loss.append(score_curr.item())\n",
    "\n",
    "        # validation\n",
    "        chi_t_valid = Full_net(tensor_valid_X1)\n",
    "        chi_tau_valid = Full_net(tensor_valid_X2)\n",
    "        score_valid = vampe_loss_rev(chi_t_valid, chi_tau_valid)\n",
    "        \n",
    "        epoch_loss[epoch] = np.mean(running_epoch_loss)\n",
    "        epoch_loss_valid[epoch] = np.mean(score_valid.item())\n",
    "        if verbose:\n",
    "            print('Run {}, total loss: {:.3}, valid loss: {:.3}'.format(epoch+1, epoch_loss[epoch], epoch_loss_valid[epoch]))\n",
    "\n",
    "    if plot_training:\n",
    "        plt.plot(np.arange(1,runs+1), epoch_loss, label='VAMP_loss')\n",
    "        plt.plot(np.arange(1,runs+1), epoch_loss_valid, label='VAMP_loss_valid')\n",
    "        plt.legend()\n",
    "        plt.show()    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if Full_net.reversible:\n",
    "    optimizer_rev = optim.Adam(Full_net.get_params_rev(), lr=learning_rate/100)\n",
    "    optimizer_rev_u_S = optim.Adam(Full_net.get_params_rev(all=False), lr=learning_rate)\n",
    "    optimizer_rev_S = optim.Adam(Full_net.get_params_rev(all=False, u_flag=False), lr=learning_rate*10)\n",
    "    optimizer_rev_u = optim.Adam(Full_net.get_params_rev(all=False, S_flag=False), lr=learning_rate*1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "weights_before_rec = Full_net.get_weights()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training loop until the validation score is converged"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def recursive_opt_rev(verbose, plot_training, reset=True):\n",
    "    \n",
    "    chi_t = Full_net(tensor_valid_X1)\n",
    "    chi_tau = Full_net(tensor_valid_X2)\n",
    "    score_old = vampe_loss_rev(chi_t, chi_tau).item()\n",
    "    \n",
    "    \n",
    "    weights_all = Full_net.get_weights()\n",
    "    score = score_old\n",
    "\n",
    "\n",
    "    while (score >= score_old):\n",
    "        if reset:\n",
    "            Full_net.set_rev_var(S=False)\n",
    "\n",
    "        train_for_S(runs=1000, verbose=verbose, plot_training=plot_training)\n",
    "\n",
    "        train_for_u_S(runs=1000, verbose=verbose, plot_training=plot_training)\n",
    "\n",
    "        train_for_rev(runs=10, verbose=verbose, plot_training=plot_training)\n",
    "\n",
    "        chi_t = Full_net(tensor_valid_X1)\n",
    "        chi_tau = Full_net(tensor_valid_X2)\n",
    "        score = vampe_loss_rev(chi_t, chi_tau).item()\n",
    "\n",
    "        print('Old score {}, new score {}'.format(score_old, score))\n",
    "        if (score>score_old):\n",
    "            print('Score is better and weights are saved')\n",
    "            score_old = score\n",
    "            weights_all = Full_net.get_weights()\n",
    "\n",
    "\n",
    "    Full_net.set_weights(weights_all)\n",
    "\n",
    "    \n",
    "    weights_all = Full_net.get_weights()\n",
    "    return weights_all"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Full_net.set_weights(weights_after_S)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "weights_after = recursive_opt_rev(verbose=False, plot_training=True, reset=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if Full_net.reversible:\n",
    "    optimizer_rev = optim.Adam(Full_net.get_params_rev(), lr=learning_rate/100)\n",
    "    optimizer_rev_u_S = optim.Adam(Full_net.get_params_rev(all=False), lr=learning_rate)\n",
    "    optimizer_rev_S = optim.Adam(Full_net.get_params_rev(all=False, u_flag=False), lr=learning_rate*10)\n",
    "    optimizer_rev_u = optim.Adam(Full_net.get_params_rev(all=False, S_flag=False), lr=learning_rate*10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Trianing loop until the validation score is converged but with resetting the parameters of u"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def recursive_opt_rev(verbose, plot_training, reset=True):\n",
    "    \n",
    "    chi_t = Full_net(tensor_valid_X1)\n",
    "    chi_tau = Full_net(tensor_valid_X2)\n",
    "    score_old = vampe_loss_rev(chi_t, chi_tau).item()\n",
    "    \n",
    "    \n",
    "    weights_all = Full_net.get_weights()\n",
    "    score = score_old\n",
    "\n",
    "\n",
    "    while (score >= score_old):\n",
    "        if reset:\n",
    "            Full_net.set_rev_var(S=False)\n",
    "        train_for_rev(runs=10, verbose=verbose, plot_training=plot_training)\n",
    "        \n",
    "        train_for_S(runs=1000, verbose=verbose, plot_training=plot_training)\n",
    "\n",
    "        train_for_u_S(runs=100, verbose=verbose, plot_training=plot_training)\n",
    "\n",
    "        \n",
    "\n",
    "        chi_t = Full_net(tensor_valid_X1)\n",
    "        chi_tau = Full_net(tensor_valid_X2)\n",
    "        score = vampe_loss_rev(chi_t, chi_tau).item()\n",
    "\n",
    "        print('Old score {}, new score {}'.format(score_old, score))\n",
    "        if (score>score_old):\n",
    "            print('Score is better and weights are saved')\n",
    "            score_old = score\n",
    "            weights_all = Full_net.get_weights()\n",
    "\n",
    "\n",
    "    Full_net.set_weights(weights_all)\n",
    "\n",
    "    \n",
    "    weights_all = Full_net.get_weights()\n",
    "    return weights_all"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "weights_after = recursive_opt_rev(verbose=False, plot_training=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "K = get_K_rev()\n",
    "# compare the eigenvalues of the VAMPnet result with the revDMSM result\n",
    "np.linalg.eigvals(K), np.linalg.eigvals(K_vamp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "weights_after_rec = Full_net.get_weights()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train u and S of the DMSM for larger tau, the data was prepared above"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tau = 50*2*25//skip\n",
    "\n",
    "tensor_train_X1 = torch.Tensor(X1_train_cor)\n",
    "tensor_train_X2 = torch.Tensor(X2_train_cor) # transform to torch tensor\n",
    "tensor_valid_X1 = torch.Tensor(X1_vali_cor)\n",
    "tensor_valid_X2 = torch.Tensor(X2_vali_cor)\n",
    "tensor_test_X1 = torch.Tensor(X1_test_cor)\n",
    "tensor_test_X2 = torch.Tensor(X2_test_cor)\n",
    "\n",
    "trainset = data.TensorDataset(tensor_train_X1, tensor_train_X2) # create your datset\n",
    "trainloader = data.DataLoader(trainset, batch_size=batch_size,\n",
    "                              shuffle=True, num_workers=2)\n",
    "\n",
    "trainloader_full = data.DataLoader(trainset, batch_size=300000,\n",
    "                              shuffle=True, num_workers=2)\n",
    "\n",
    "testset = data.TensorDataset(tensor_valid_X1, tensor_valid_X2) # create your datset\n",
    "testloader = data.DataLoader(testset, batch_size=batch_size,\n",
    "                             shuffle=True, num_workers=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Full_net.set_weights(weights_after_rec)\n",
    "with torch.no_grad():\n",
    "    for param in Full_net.S_layers[0].parameters():\n",
    "        param.copy_(torch.Tensor(np.ones((output_sizes[0], output_sizes[0]))))\n",
    "with torch.no_grad():\n",
    "    for param in Full_net.u_layers[0].parameters():\n",
    "        param.copy_(torch.Tensor(np.ones((1, output_sizes[0]))))\n",
    "        \n",
    "optimizer_rev_u_S = optim.Adam(Full_net.get_params_rev(all=False), lr=learning_rate)\n",
    "optimizer_rev_S = optim.Adam(Full_net.get_params_rev(all=False, u_flag=False), lr=learning_rate*10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for _ in range(5):\n",
    "    train_for_S(runs=1000, verbose=False, plot_training=True)\n",
    "\n",
    "    train_for_u_S(runs=1000, verbose=False, plot_training=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "K_msm = get_K_rev(tensor_test_X1, tensor_test_X2)\n",
    "chi_t = Full_net.forward(tensor_train_X1).detach().numpy()\n",
    "chi_tau = Full_net.forward(tensor_train_X2).detach().numpy()\n",
    "K_vamp_msm = estimate_koopman_op([chi_t, chi_tau], 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compare the eigenvalues of VAMPnet and revDMSM, they should be similar\n",
    "np.linalg.eigvals(K_msm), np.linalg.eigvals(K_vamp_msm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "weights_msm = Full_net.get_weights()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Simple renaming of the data\n",
    "tensor_train_X1_cor = torch.Tensor(X1_train_cor)\n",
    "tensor_train_X2_cor = torch.Tensor(X2_train_cor) # transform to torch tensor\n",
    "tensor_valid_X1_cor = torch.Tensor(X1_vali_cor)\n",
    "tensor_valid_X2_cor = torch.Tensor(X2_vali_cor)\n",
    "tensor_test_X1_cor = torch.Tensor(X1_test_cor)\n",
    "tensor_test_X2_cor = torch.Tensor(X2_test_cor)\n",
    "\n",
    "trainset_cor = data.TensorDataset(tensor_train_X1_cor, tensor_train_X2_cor) # create your datset\n",
    "trainloader_cor = data.DataLoader(trainset_cor, batch_size=batch_size,\n",
    "                              shuffle=True, num_workers=2)\n",
    "\n",
    "trainloader_full_cor = data.DataLoader(trainset_cor, batch_size=X1_train_cor.shape[0],\n",
    "                              shuffle=True, num_workers=2)\n",
    "\n",
    "testset_cor = data.TensorDataset(tensor_valid_X1_cor, tensor_valid_X2_cor) # create your datset\n",
    "testloader_cor = data.DataLoader(testset_cor, batch_size=X1_vali_cor.shape[0],\n",
    "                             shuffle=True, num_workers=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Estimate the observable values and compare to the true one\n",
    "\n",
    "These values should be different, since we manipulated the data and did not use any information from the true simulation yet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "chi_t = Full_net(torch.Tensor(X1_train_cor))\n",
    "chi_tau = Full_net(torch.Tensor(X2_train_cor))\n",
    "\n",
    "\n",
    "_, S_similar, K_test, eigval_fold = vampe_loss_rev_obs(chi_t, chi_tau)\n",
    "eigval_fold_before = eigval_fold.detach()\n",
    "\n",
    "exp_value_time_before = []\n",
    "for obs_v in obs_train:\n",
    "    obs_value_tensor = torch.Tensor(obs_v)\n",
    "    exp_value_tensor = torch.Tensor(np.array([0.]))\n",
    "\n",
    "    score_curr, mu_t, K, Sigma = vampe_loss_rev(chi_t, chi_tau, return_mu_K_Sigma=True)\n",
    "    error_obs = obs_time_loss(obs_value_tensor, mu_t, chi_t, K, Sigma, exp_value_tensor)\n",
    "    exp_value_time_before.append(error_obs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "exp_value_time_before, exp_exp_value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eigval_fold_before, true_value"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Estimate model with information of some experimental observables\n",
    "Here we use the expectation value of the contact being formed and the timescale of the folding process, but it can be easily adapted to the other cases presented in the paper"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Loss functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def timescale_loss(obs_value, exp_value):\n",
    "    \n",
    "    error = torch.sum(torch.abs(exp_value - obs_value))\n",
    "    \n",
    "    return error\n",
    "\n",
    "    \n",
    "def obs_loss(obs_value, mu, exp_value):\n",
    "    \n",
    "    exp_value_estimated = torch.sum(obs_value * mu)\n",
    "    \n",
    "    error = torch.abs(exp_value - exp_value_estimated)\n",
    "    \n",
    "    return error\n",
    "\n",
    "def obs_time_loss(obs_value, mu, chi, K, Sigma, exp_value):\n",
    "\n",
    "    state_weight = mu*chi\n",
    "\n",
    "    pi = torch.sum(state_weight, dim=0) # prob to be in a state\n",
    "    # obs value within a state, the weighting factor needs to be normalized for each state\n",
    "    ai = torch.sum(state_weight*obs_value, dim=0, keepdims=True) / torch.sum(state_weight, dim=0, keepdims=True)\n",
    "    # prob to observe an unconditional jump state i to j\n",
    "    X = Sigma @ K\n",
    "    a_sim = torch.matmul(ai, torch.matmul(X, ai.T)) # shape 1x1\n",
    "        \n",
    "    error = torch.abs(a_sim - exp_value) \n",
    "           \n",
    "    return error"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Since we want to match eigenvalues of a specific process we need to define, what the process is\n",
    "Our approach: We pretrained a VAMPnet, where we can identify the states which are the folded and which is the unfolded state. This helps us to identify the process which changes the most between these two states."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get frames for folded and unfolded structure as extreme values of the eigenfunction of the folding process\n",
    "frame1_fu, frame2_fu = np.argmax(eigfunc), np.argmin(eigfunc) \n",
    "# get the corresponding states of the new chi\n",
    "pred_fu = Full_net.forward(torch.Tensor(traj_whole[0][[frame1_fu,frame2_fu]]))\n",
    "state1_fu = np.argmax(pred_fu[0].detach().numpy())\n",
    "state2_fu = np.argmax(pred_fu[1].detach().numpy())\n",
    "def vampe_loss_rev_obs(chi_t, chi_tau, layer_id=0, return_mu=False, all_eigvals=False):\n",
    "    \n",
    "    v, C_00, C_11, C_01, Sigma, mu_t, Sigma_t = Full_net.u_layers[layer_id](chi_t, chi_tau)\n",
    "    matrix, K, S = Full_net.S_layers[layer_id](v, C_00, C_11, C_01, Sigma)\n",
    "    vampe = torch.trace(matrix)\n",
    "    \n",
    "    \n",
    "    eigval_all, eigvec_all = torch.symeig(Sigma, eigenvectors=True)\n",
    "\n",
    "    # Filter out eigvalues below threshold and corresponding eigvectors\n",
    "    eig_th = torch.Tensor(epsilon)\n",
    "    index_eig = eigval_all > eig_th\n",
    "#     print(index_eig)\n",
    "    eigval = eigval_all[index_eig]\n",
    "    eigvec = eigvec_all[:,index_eig]\n",
    "\n",
    "    # Build the diagonal matrix with the filtered eigenvalues or square\n",
    "    # root of the filtered eigenvalues according to the parameter\n",
    "\n",
    "    diag = torch.diag(torch.sqrt(eigval))\n",
    "    diag_inv = torch.diag(1/torch.sqrt(eigval))\n",
    "    #     print(diag.shape, eigvec.shape)    \n",
    "    # Rebuild the square root of the inverse matrix\n",
    "    Sigma_sqrt = torch.matmul(eigvec, torch.matmul(diag, eigvec.T))\n",
    "    Sigma_sqrt_inv = torch.matmul(eigvec, torch.matmul(diag_inv, eigvec.T))\n",
    "\n",
    "    S_similar = Sigma_sqrt @ S @ Sigma_sqrt\n",
    "\n",
    "    eigval_all, eigvec_all = torch.symeig(S_similar, eigenvectors=True)\n",
    "    eigvecs_K = Sigma_sqrt_inv @ eigvec_all\n",
    "\n",
    "    # Find the process which is the folding\n",
    "    process_id = torch.argmin(eigvecs_K[state1_fu,:]*eigvecs_K[state2_fu,:]).detach()\n",
    "    \n",
    "#     sort_eigvals = torch.sort(eigval_all)\n",
    "    if all_eigvals:\n",
    "        ret = [-vampe, S_similar, K, eigval_all]\n",
    "    else:\n",
    "        ret = [-vampe, S_similar, K, eigval_all[process_id], mu_t, Sigma_t]\n",
    "    \n",
    "    return ret"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training function for updating u and S while matching the observable"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_for_obs_timescale(exp_value_ts, obs_value, exp_value, weight_loss_ts=1, weight_loss=1, runs=100, S_flag=True, u_flag=False, verbose=False, plot_training=False):\n",
    "        \n",
    "        \n",
    "    chi_t = Full_net(tensor_train_X1_cor).detach()\n",
    "    chi_tau = Full_net(tensor_train_X2_cor).detach()\n",
    "    \n",
    "    \n",
    "    \n",
    "    \n",
    "    \n",
    "    epoch_loss = np.zeros(runs)\n",
    "    epoch_vampe = np.zeros(runs)\n",
    "    epoch_error = np.zeros(runs)\n",
    "    epoch_error1 = np.zeros(runs)\n",
    "#     epoch_error2 = np.zeros(runs)\n",
    "    \n",
    "    \n",
    "    \n",
    "    if u_flag and S_flag:\n",
    "        opt = optimizer_rev_u_S_cor\n",
    "    elif u_flag:\n",
    "        opt = optimizer_rev_u_cor\n",
    "    else:\n",
    "        opt = optimizer_rev_S_cor\n",
    "    \n",
    "    obs_value_tensor = []\n",
    "    exp_value_tensor = []\n",
    "    for ob in obs_value:\n",
    "        obs_value_tensor.append(torch.Tensor(ob))\n",
    "    for ex in exp_value:\n",
    "        exp_value_tensor.append(ex.detach())\n",
    "        \n",
    "    \n",
    "    for epoch in range(runs):  # loop over the dataset multiple times\n",
    "    \n",
    "        opt.zero_grad()\n",
    "\n",
    "        # estimate weights\n",
    "#             print(inputs_t)\n",
    "\n",
    "#             print(chi_t)\n",
    "        score_curr, S_similar, K, eigval_fold, mu_t, Sigma_t = vampe_loss_rev_obs(chi_t, chi_tau)\n",
    "        \n",
    "        error_obs = timescale_loss(eigval_fold, exp_value_ts)\n",
    "        \n",
    "        error_obs1 = obs_loss(obs_value_tensor[0], mu_t, exp_value_tensor[0])\n",
    "        \n",
    "#         error_obs1 = obs_time_loss(obs_value_tensor[0], mu_t, chi_t, K, Sigma_t, exp_value_tensor[0])\n",
    "#         error_obs2 = obs_time_loss(obs_value_tensor[1], mu_t, chi_t, K, Sigma_t, exp_value_tensor[1])\n",
    "        \n",
    "        loss = - score_curr + weight_loss * error_obs1 + weight_loss_ts * error_obs\n",
    "\n",
    "        loss.backward()\n",
    "\n",
    "        opt.step()\n",
    "    \n",
    "        epoch_loss[epoch] = np.mean(-loss.item())\n",
    "        epoch_vampe[epoch] = np.mean(score_curr.item())\n",
    "        epoch_error[epoch] = np.mean(error_obs.item())\n",
    "        epoch_error1[epoch] = np.mean(error_obs1.item())\n",
    "#         epoch_error2[epoch] = np.mean(error_obs2.item())\n",
    "        if verbose:\n",
    "            print('Run {}, total loss: {:.3}'.format(epoch+1, epoch_loss[epoch]))\n",
    "\n",
    "    if plot_training:\n",
    "        plt.plot(np.arange(1,runs+1), epoch_loss, label='Total_loss')\n",
    "        plt.legend()\n",
    "        plt.show() \n",
    "        \n",
    "        plt.plot(np.arange(1,runs+1), epoch_vampe, label='VAMP_loss')\n",
    "        plt.legend()\n",
    "        plt.show() \n",
    "        \n",
    "        plt.plot(np.arange(1,runs+1), epoch_error, label='Error_loss_ts')\n",
    "        plt.legend()\n",
    "        plt.show() \n",
    "        \n",
    "        plt.plot(np.arange(1,runs+1), epoch_error1, label='Error_loss')\n",
    "#         plt.plot(np.arange(1,runs+1), epoch_error2, label='Error_loss2')\n",
    "        plt.legend()\n",
    "        plt.show() "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training function to update all parameters while matching the observables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_for_obs_timescale_all(exp_value_ts, obs_value, exp_value, weight_loss_ts=1, weight_loss=1, runs=100, verbose=False, plot_training=False):\n",
    "        \n",
    "         \n",
    "    \n",
    "    epoch_loss = np.zeros(runs)\n",
    "    epoch_vampe = np.zeros(runs)\n",
    "    epoch_error = np.zeros(runs)\n",
    "    epoch_error1 = np.zeros(runs)\n",
    "#     epoch_error2 = np.zeros(runs)\n",
    "    \n",
    "    opt = optimizer_rev_cor\n",
    "    obs_value_tensor = []\n",
    "    exp_value_tensor = []\n",
    "    for ob in obs_value:\n",
    "        obs_value_tensor.append(torch.Tensor(ob))\n",
    "    for ex in exp_value:\n",
    "        exp_value_tensor.append(ex.detach())\n",
    "        \n",
    "    \n",
    "    for epoch in range(runs):  # loop over the dataset multiple times\n",
    "    \n",
    "        opt.zero_grad()\n",
    "\n",
    "        # estimate weights\n",
    "#             print(inputs_t)\n",
    "\n",
    "#             print(chi_t)\n",
    "        chi_t = Full_net(tensor_train_X1_cor)\n",
    "        chi_tau = Full_net(tensor_train_X2_cor)\n",
    "        score_curr, S_similar, K, eigval_fold, mu_t, Sigma_t = vampe_loss_rev_obs(chi_t, chi_tau)\n",
    "        \n",
    "        error_obs = timescale_loss(eigval_fold, exp_value_ts)\n",
    "        \n",
    "        error_obs1 = obs_loss(obs_value_tensor[0], mu_t, exp_value_tensor[0])\n",
    "        \n",
    "#         error_obs1 = obs_time_loss(obs_value_tensor[0], mu_t, chi_t, K, Sigma_t, exp_value_tensor[0])\n",
    "#         error_obs2 = obs_time_loss(obs_value_tensor[1], mu_t, chi_t, K, Sigma_t, exp_value_tensor[1])\n",
    "        \n",
    "        loss = - score_curr + weight_loss * error_obs1 + weight_loss_ts * error_obs\n",
    "\n",
    "        loss.backward()\n",
    "\n",
    "        opt.step()\n",
    "    \n",
    "        epoch_loss[epoch] = np.mean(-loss.item())\n",
    "        epoch_vampe[epoch] = np.mean(score_curr.item())\n",
    "        epoch_error[epoch] = np.mean(error_obs.item())\n",
    "        epoch_error1[epoch] = np.mean(error_obs1.item())\n",
    "#         epoch_error2[epoch] = np.mean(error_obs2.item())\n",
    "        if verbose:\n",
    "            print('Run {}, total loss: {:.3}'.format(epoch+1, epoch_loss[epoch]))\n",
    "\n",
    "    if plot_training:\n",
    "        plt.plot(np.arange(1,runs+1), epoch_loss, label='Total_loss')\n",
    "        plt.legend()\n",
    "        plt.show() \n",
    "        \n",
    "        plt.plot(np.arange(1,runs+1), epoch_vampe, label='VAMP_loss')\n",
    "        plt.legend()\n",
    "        plt.show() \n",
    "        \n",
    "        plt.plot(np.arange(1,runs+1), epoch_error, label='Error_loss')\n",
    "        plt.legend()\n",
    "        plt.show() \n",
    "        \n",
    "        plt.plot(np.arange(1,runs+1), epoch_error1, label='Error_loss')\n",
    "#         plt.plot(np.arange(1,runs+1), epoch_error2, label='Error_loss2')\n",
    "        plt.legend()\n",
    "        plt.show() \n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train only for the expectation value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_for_obs(obs_value, exp_value, weight_loss=1, runs=100, S_flag=True, u_flag=False, verbose=False, plot_training=False):\n",
    "        \n",
    "        \n",
    "    chi_t = Full_net(tensor_train_X1_cor).detach()\n",
    "    chi_tau = Full_net(tensor_train_X2_cor).detach()\n",
    "    \n",
    "    epoch_loss = np.zeros(runs)\n",
    "    epoch_vampe = np.zeros(runs)\n",
    "    epoch_error1 = np.zeros(runs)\n",
    "#     epoch_error2 = np.zeros(runs)\n",
    "    if u_flag and S_flag:\n",
    "        opt = optimizer_rev_u_S_cor\n",
    "    elif u_flag:\n",
    "        opt = optimizer_rev_u_cor\n",
    "    else:\n",
    "        opt = optimizer_rev_S_cor\n",
    "    obs_value_tensor = []\n",
    "    exp_value_tensor = []\n",
    "    for ob in obs_value:\n",
    "        obs_value_tensor.append(torch.Tensor(ob))\n",
    "    for ex in exp_value:\n",
    "        exp_value_tensor.append(ex.detach())\n",
    "    \n",
    "    for epoch in range(runs):  # loop over the dataset multiple times\n",
    "    \n",
    "        opt.zero_grad()\n",
    "\n",
    "        # estimate weights\n",
    "#             print(inputs_t)\n",
    "\n",
    "#             print(chi_t)\n",
    "        score_curr, mu_t = vampe_loss_rev(chi_t, chi_tau, return_mu=True)\n",
    "        \n",
    "        error_obs1 = obs_loss(obs_value_tensor[0], mu_t, exp_value_tensor[0])\n",
    "#         error_obs2 = obs_time_loss(obs_value_tensor[1], mu_t, chi_t, K, Sigma_t, exp_value_tensor[1])\n",
    "        \n",
    "        loss = - score_curr + weight_loss * error_obs1# + weight_loss * error_obs2\n",
    "\n",
    "        loss.backward()\n",
    "\n",
    "        opt.step()\n",
    "    \n",
    "        epoch_loss[epoch] = np.mean(-loss.item())\n",
    "        epoch_vampe[epoch] = np.mean(score_curr.item())\n",
    "        epoch_error1[epoch] = np.mean(error_obs1.item())\n",
    "#         epoch_error2[epoch] = np.mean(error_obs2.item())\n",
    "        if verbose:\n",
    "            print('Run {}, total loss: {:.3}'.format(epoch+1, epoch_loss[epoch]))\n",
    "\n",
    "    if plot_training:\n",
    "        plt.plot(np.arange(1,runs+1), epoch_loss, label='Total_loss')\n",
    "        plt.legend()\n",
    "        plt.show() \n",
    "        \n",
    "        plt.plot(np.arange(1,runs+1), epoch_vampe, label='VAMP_loss')\n",
    "        plt.legend()\n",
    "        plt.show() \n",
    "        \n",
    "        plt.plot(np.arange(1,runs+1), epoch_error1, label='Error_loss')\n",
    "#         plt.plot(np.arange(1,runs+1), epoch_error2, label='Error_loss2')\n",
    "        plt.legend()\n",
    "        plt.show() "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_for_obs_all(obs_value, exp_value, weight_loss=1, runs=100, verbose=False, plot_training=False):\n",
    "        \n",
    "\n",
    "    epoch_loss = np.zeros(runs)\n",
    "    epoch_vampe = np.zeros(runs)\n",
    "    epoch_error1 = np.zeros(runs)\n",
    "#     epoch_error2 = np.zeros(runs)\n",
    "    opt = optimizer_rev_cor\n",
    "    \n",
    "    obs_value_tensor = []\n",
    "    exp_value_tensor = []\n",
    "    for ob in obs_value:\n",
    "        obs_value_tensor.append(torch.Tensor(ob))\n",
    "    for ex in exp_value:\n",
    "        exp_value_tensor.append(ex.detach())\n",
    "    \n",
    "    for epoch in range(runs):  # loop over the dataset multiple times\n",
    "    \n",
    "        opt.zero_grad()\n",
    "\n",
    "        # estimate weights\n",
    "#             print(inputs_t)\n",
    "\n",
    "#             print(chi_t)\n",
    "        chi_t = Full_net(tensor_train_X1_cor)\n",
    "        chi_tau = Full_net(tensor_train_X2_cor)\n",
    "        score_curr, mu_t = vampe_loss_rev(chi_t, chi_tau, return_mu=True)\n",
    "        \n",
    "        error_obs1 = obs_loss(obs_value_tensor[0], mu_t, exp_value_tensor[0])\n",
    "#         error_obs2 = obs_time_loss(obs_value_tensor[1], mu_t, chi_t, K, Sigma_t, exp_value_tensor[1])\n",
    "        \n",
    "        loss = - score_curr + weight_loss * error_obs1# + weight_loss * error_obs2\n",
    "\n",
    "        loss.backward()\n",
    "\n",
    "        opt.step()\n",
    "    \n",
    "        epoch_loss[epoch] = np.mean(-loss.item())\n",
    "        epoch_vampe[epoch] = np.mean(score_curr.item())\n",
    "        epoch_error1[epoch] = np.mean(error_obs1.item())\n",
    "#         epoch_error2[epoch] = np.mean(error_obs2.item())\n",
    "        if verbose:\n",
    "            print('Run {}, total loss: {:.3}'.format(epoch+1, epoch_loss[epoch]))\n",
    "\n",
    "    if plot_training:\n",
    "        plt.plot(np.arange(1,runs+1), epoch_loss, label='Total_loss')\n",
    "        plt.legend()\n",
    "        plt.show() \n",
    "        \n",
    "        plt.plot(np.arange(1,runs+1), epoch_vampe, label='VAMP_loss')\n",
    "        plt.legend()\n",
    "        plt.show() \n",
    "        \n",
    "        plt.plot(np.arange(1,runs+1), epoch_error1, label='Error_loss')\n",
    "        plt.plot(np.arange(1,runs+1), epoch_error2, label='Error_loss2')\n",
    "        plt.legend()\n",
    "        plt.show() "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def eval_obs(tensor_X1, tensor_X2, exp_value_ts, obs_value, exp_value, weight_loss_ts=1., weight_loss=1.):\n",
    "    obs_value_tensor = []\n",
    "    exp_value_tensor = []\n",
    "    for ob in obs_value:\n",
    "        obs_value_tensor.append(torch.Tensor(ob))\n",
    "    for ex in exp_value:\n",
    "        exp_value_tensor.append(ex.detach())\n",
    "        \n",
    "    \n",
    "    chi_t = Full_net(tensor_X1)\n",
    "    chi_tau = Full_net(tensor_X2)\n",
    "    score_curr, mu_t = vampe_loss_rev(chi_t, chi_tau, return_mu=True)\n",
    "        \n",
    "    error_obs1 = obs_loss(obs_value_tensor[0], mu_t, exp_value_tensor[0])\n",
    "#         error_obs2 = obs_time_loss(obs_value_tensor[1], mu_t, chi_t, K, Sigma_t, exp_value_tensor[1])\n",
    "        \n",
    "    # + weight_loss * error_obs2\n",
    "\n",
    "#     error_obs1 = obs_time_loss(obs_value_tensor[0], mu_t, chi_t, K, Sigma_t, exp_value_tensor[0])\n",
    "#     error_obs2 = obs_time_loss(obs_value_tensor[1], mu_t, chi_t, K, Sigma_t, exp_value_tensor[1])\n",
    "    _, S_similar, K_test, eigval_fold, _, _ = vampe_loss_rev_obs(chi_t, chi_tau)\n",
    "    eigval_fold = eigval_fold.detach()\n",
    "    error_obs = timescale_loss(eigval_fold, exp_value_ts)\n",
    "    loss = - score_curr + weight_loss * error_obs1 + weight_loss_ts * error_obs\n",
    "#     loss = - score_curr + weight_loss * error_obs1 + weight_loss * error_obs2 + weight_loss_ts * error_obs\n",
    "    \n",
    "    return -loss.item()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Full_net.set_weights(weights_msm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prob_states_before = plot_mu(tensor_test_X1_cor, tensor_test_X2_cor, frames_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "chi_t_test = Full_net(torch.Tensor(X1_test_cor))\n",
    "chi_tau_test = Full_net(torch.Tensor(X2_test_cor))\n",
    "exp_value_time_before = []\n",
    "for obs_v in obs_test:\n",
    "    obs_value_tensor = torch.Tensor(obs_v)\n",
    "    exp_value_tensor = torch.Tensor(np.array([0.]))\n",
    "\n",
    "    score_curr, mu_t = vampe_loss_rev(chi_t_test, chi_tau_test, return_mu=True)\n",
    "    error_obs = obs_loss(obs_value_tensor, mu_t, exp_value_tensor)\n",
    "    exp_value_time_before.append(error_obs.detach().numpy())\n",
    "exp_value_time_before = np.array(exp_value_time_before).squeeze()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Now train also for the observable"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "change_S = True\n",
    "retrain_chi = True\n",
    "reset = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Reset the weights\n",
    "Full_net.set_weights(weights_dict_chi)\n",
    "\n",
    "if reset:\n",
    "    with torch.no_grad():\n",
    "        for param in Full_net.u_layers[0].parameters():\n",
    "            param.copy_(torch.Tensor(np.ones((1, output_sizes[0]))))\n",
    "    \n",
    "                \n",
    "                \n",
    "optimizer_rev_cor = optim.Adam(Full_net.get_params_rev(), lr=learning_rate)\n",
    "optimizer_rev_u_S_cor = optim.Adam(Full_net.get_params_rev(all=False), lr=learning_rate)\n",
    "optimizer_rev_S_cor = optim.Adam(Full_net.get_params_rev(all=False, u_flag=False), lr=learning_rate*1)\n",
    "optimizer_rev_u_cor = optim.Adam(Full_net.get_params_rev(all=False, S_flag=False), lr=learning_rate*1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# how much should the exp values be enforced\n",
    "weight_loss_list= np.linspace(0.1,3,10)\n",
    "weight_loss = 10.\n",
    "weight_loss_ts = 10.#10."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if change_S:\n",
    "    if reset:\n",
    "        with torch.no_grad():\n",
    "            for param in Full_net.S_layers[0].parameters():\n",
    "                param.copy_(torch.Tensor(np.ones((output_sizes[0], output_sizes[0]))))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### First train u and S"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "score_best_valid = -5#eval_obs(tensor_valid_X1_cor, tensor_valid_X2_cor, true_value, weight_loss=weight_loss)\n",
    "score_best_train = -5#eval_obs(tensor_train_X1_cor, tensor_train_X2_cor, true_value, weight_loss=weight_loss)\n",
    "for _ in range(150):\n",
    "    if change_S:\n",
    "        train_for_obs_timescale(true_value, obs_train, exp_exp_value, weight_loss=weight_loss, weight_loss_ts=weight_loss_ts, runs=100, S_flag=True, u_flag=True, verbose=False, plot_training=False)\n",
    "        score_temp_valid = eval_obs(tensor_valid_X1_cor, tensor_valid_X2_cor, true_value, obs_valid, exp_exp_value, weight_loss_ts=weight_loss_ts, weight_loss=weight_loss)\n",
    "        score_temp_train = eval_obs(tensor_train_X1_cor, tensor_train_X2_cor, true_value, obs_train, exp_exp_value, weight_loss=weight_loss, weight_loss_ts=weight_loss_ts)\n",
    "    \n",
    "        K_sim_before3 = get_K_rev(tensor_train_X1_cor, tensor_train_X2_cor)\n",
    "        print(np.linalg.eigvals(K_sim_before3), np.linalg.eigvals(K_exp) )\n",
    "        print('The score for training is: {:.5}, and for valid: {:.5}'.format(score_temp_train, score_temp_valid))\n",
    "        if score_temp_valid > score_best_valid:\n",
    "            print('Better validation score, save weights')\n",
    "            score_best_valid = score_temp_valid\n",
    "            weights_temp_best_valid = Full_net.get_weights()\n",
    "\n",
    "        if score_temp_train > score_best_train:\n",
    "            print('Better training score, save weights')\n",
    "            score_best_train = score_temp_train\n",
    "            weights_temp_best_train = Full_net.get_weights()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Full_net.set_weights(weights_temp_best_valid)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train everything"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if (change_S and retrain_chi):\n",
    "    for c in range(150):\n",
    "       \n",
    "        train_for_obs_timescale_all(true_value, obs_train, exp_exp_value, weight_loss_ts=weight_loss_ts, weight_loss=weight_loss, runs=10, verbose=False, plot_training=False)\n",
    "        score_temp_valid = eval_obs(tensor_valid_X1_cor, tensor_valid_X2_cor, true_value, obs_valid, exp_exp_value, weight_loss_ts=weight_loss_ts, weight_loss=weight_loss)\n",
    "        score_temp_train = eval_obs(tensor_train_X1_cor, tensor_train_X2_cor, true_value, obs_train, exp_exp_value, weight_loss=weight_loss, weight_loss_ts=weight_loss_ts)\n",
    "        print('The score for training is: {:.5}, and for valid: {:.5}'.format(score_temp_train, score_temp_valid))\n",
    "        if score_temp_valid > score_best_valid:\n",
    "            print('Better validation score, save weights')\n",
    "            score_best_valid = score_temp_valid\n",
    "            weights_temp_best_valid = Full_net.get_weights()\n",
    "            print(plot_mu(tensor_test_X1_cor, tensor_test_X2_cor, frames_test), prob_states_true)\n",
    "        if score_temp_train > score_best_train:\n",
    "            print('Better training score, save weights')\n",
    "            score_best_train = score_temp_train\n",
    "            weights_temp_best_train = Full_net.get_weights()\n",
    "        \n",
    "        train_for_obs_timescale(true_value, obs_train, exp_exp_value, weight_loss_ts=weight_loss_ts, weight_loss=weight_loss, runs=100, S_flag=True, u_flag=True, verbose=False, plot_training=False)\n",
    "\n",
    "        K_sim_before3 = get_K_rev(tensor_train_X1_cor, tensor_train_X2_cor)\n",
    "        print(np.linalg.eigvals(K_sim_before3), np.linalg.eigvals(K_exp) )\n",
    "        score_temp_valid = eval_obs(tensor_valid_X1_cor, tensor_valid_X2_cor, true_value, obs_valid, exp_exp_value, weight_loss_ts=weight_loss_ts, weight_loss=weight_loss)\n",
    "        score_temp_train = eval_obs(tensor_train_X1_cor, tensor_train_X2_cor, true_value, obs_train, exp_exp_value, weight_loss=weight_loss, weight_loss_ts=weight_loss_ts)\n",
    "        print('The score for training is: {:.5}, and for valid: {:.5}'.format(score_temp_train, score_temp_valid))\n",
    "        if score_temp_valid > score_best_valid:\n",
    "            print('Better validation score, save weights')\n",
    "            score_best_valid = score_temp_valid\n",
    "            weights_temp_best_valid = Full_net.get_weights()\n",
    "            print(plot_mu(tensor_test_X1_cor, tensor_test_X2_cor, frames_test), prob_states_true)\n",
    "        if score_temp_train > score_best_train:\n",
    "            print('Better training score, save weights')\n",
    "            score_best_train = score_temp_train\n",
    "            weights_temp_best_train = Full_net.get_weights()\n",
    "            \n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Full_net.set_weights(weights_temp_best_valid)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Final training for only u and S"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for _ in range(100):\n",
    "    if change_S:\n",
    "        train_for_obs_timescale(true_value, obs_train, exp_exp_value, weight_loss_ts=weight_loss_ts, weight_loss=weight_loss, runs=100, S_flag=True, u_flag=True, verbose=False, plot_training=False)\n",
    "        score_temp_valid = eval_obs(tensor_valid_X1_cor, tensor_valid_X2_cor, true_value, obs_valid, exp_exp_value, weight_loss_ts=weight_loss_ts, weight_loss=weight_loss)\n",
    "        score_temp_train = eval_obs(tensor_train_X1_cor, tensor_train_X2_cor, true_value, obs_train, exp_exp_value, weight_loss=weight_loss, weight_loss_ts=weight_loss_ts)\n",
    "        K_sim_before3 = get_K_rev(tensor_train_X1_cor, tensor_train_X2_cor)\n",
    "        print(np.linalg.eigvals(K_sim_before3), np.linalg.eigvals(K_exp) )\n",
    "        print('The score for training is: {:.5}, and for valid: {:.5}'.format(score_temp_train, score_temp_valid))\n",
    "        if score_temp_valid > score_best_valid:\n",
    "            print('Better validation score, save weights')\n",
    "            score_best_valid = score_temp_valid\n",
    "            weights_temp_best_valid = Full_net.get_weights()\n",
    "            print(plot_mu(tensor_test_X1_cor, tensor_test_X2_cor, frames_test), prob_states_true)\n",
    "            \n",
    "        if score_temp_train > score_best_train:\n",
    "            print('Better training score, save weights')\n",
    "            score_best_train = score_temp_train\n",
    "            weights_temp_best_train = Full_net.get_weights()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "weights_temp_end = Full_net.get_weights()\n",
    "Full_net.set_weights(weights_temp_best_valid)\n",
    "# Full_net.set_weights(weights_temp_best_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "chi_t = Full_net(torch.Tensor(X1_test_cor))\n",
    "chi_tau = Full_net(torch.Tensor(X2_test_cor))\n",
    "\n",
    "\n",
    "_, S_similar, K_test, eigval_fold, _, _ = vampe_loss_rev_obs(chi_t, chi_tau)\n",
    "\n",
    "\n",
    "exp_value_time_after = []\n",
    "for obs_v in obs_test:\n",
    "    obs_value_tensor = torch.Tensor(obs_v)\n",
    "    exp_value_tensor = torch.Tensor(np.array([0.]))\n",
    "\n",
    "    score_curr, mu_t = vampe_loss_rev(chi_t, chi_tau, return_mu=True)\n",
    "    error_obs = obs_loss(obs_value_tensor, mu_t, exp_value_tensor)\n",
    "    exp_value_time_after.append(error_obs.detach().numpy())\n",
    "exp_value_time_after = np.array(exp_value_time_after).squeeze()\n",
    "\n",
    "print(eigval_fold)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "true_value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "K_sim_after = get_K_rev(tensor_test_X1_cor, tensor_test_X2_cor)\n",
    "np.linalg.eigvals(K_sim_after), np.linalg.eigvals(K_exp) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Estimate stationary distribution of states\n",
    "prob_states_after = plot_mu(tensor_test_X1_cor, tensor_test_X2_cor, frames_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.linalg.eigvals(K_msm)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define plotting functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fac_unit = 200.*skip*1e-6\n",
    "def plot_obs(ts=True):\n",
    "    \n",
    "    eigvals_true = np.linalg.eigvals(K_exp)\n",
    "    eigvals_true = np.sort(eigvals_true)[:-1]\n",
    "    if ts:\n",
    "        its_true = - tau / np.log(eigvals_true) * fac_unit\n",
    "    else:\n",
    "        its_true = eigvals_true\n",
    "        \n",
    "    eigvals_before = np.linalg.eigvals(K_msm)\n",
    "    eigvals_before = np.sort(eigvals_before)[:-1]\n",
    "    if ts:\n",
    "        its_before = - tau / np.log(np.abs(eigvals_before)) * fac_unit\n",
    "    else:\n",
    "        its_before = eigvals_before\n",
    "        \n",
    "    eigvals_after = np.linalg.eigvals(K_sim_after)\n",
    "    eigvals_after = np.sort(eigvals_after)[:-1]\n",
    "    if ts:\n",
    "        its_after = - tau / np.log(eigvals_after) * fac_unit\n",
    "    else:\n",
    "        its_after = eigvals_after\n",
    "    \n",
    "    for i, it in enumerate(its_true):\n",
    "        if i==0:\n",
    "            label='True'\n",
    "        else:\n",
    "            label=''\n",
    "        plt.hlines(its_true[i], i-0.25, i+0.25, ls='--', colors='k', label=label, lw=2)\n",
    "    \n",
    "    plt.plot(its_before,'v', label='Before', ms=10)\n",
    "    plt.plot(its_after,'o', label='After', ms=10)\n",
    "    plt.title('Implied Timescales')\n",
    "    indicator = ['I', 'II']\n",
    "    plt.xticks([0,1],indicator, fontsize=14)\n",
    "    plt.xlabel('Process', fontsize=16)\n",
    "    plt.ylabel('Implied Timescale [$\\mu$s]', fontsize=16)\n",
    "    plt.legend()\n",
    "    plt.show()\n",
    "    \n",
    "    \n",
    "def plot_prob_states():\n",
    "    \n",
    "    for i, pi_i in enumerate(prob_states_true):\n",
    "        if i==0:\n",
    "            label='True'\n",
    "        else:\n",
    "            label=''\n",
    "        plt.hlines(pi_i*100, i-0.25, i+0.25, ls='--', colors='k', label=label, lw=2)\n",
    "        \n",
    "    plt.plot(prob_states_before*100,'v', label='Before', ms=10)\n",
    "    plt.plot(prob_states_after*100,'o', label='After', ms=10)\n",
    "    plt.title('State Probability')\n",
    "    indicator = ['I', 'II', 'III']\n",
    "    plt.xticks([0,1,2],indicator, fontsize=14)\n",
    "    plt.xlabel('State', fontsize=16)\n",
    "    plt.ylabel('Probability [$\\%$]', fontsize=16)\n",
    "    plt.legend()\n",
    "    plt.show()\n",
    "    \n",
    "    \n",
    "def plot_obs_train():\n",
    "    \n",
    "    for i, pi_i in enumerate(exp_exp_value):\n",
    "        if i==0:\n",
    "            label='True'\n",
    "        else:\n",
    "            label=''\n",
    "        \n",
    "        plt.hlines(pi_i.numpy(), i-0.25, i+0.25, ls='--', colors='k', label=label, lw=2)\n",
    "        \n",
    "    \n",
    "    plt.plot(exp_value_time_before,'v', label='Before', ms=10)\n",
    "    plt.plot(exp_value_time_after,'o', label='After', ms=10)\n",
    "    plt.title('Observable')\n",
    "    indicator = ['I']\n",
    "    plt.xticks([0],indicator, fontsize=14)\n",
    "    plt.xlabel('Observable', fontsize=16)\n",
    "    plt.ylabel('Value', fontsize=16)\n",
    "    plt.legend()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "plot_obs()\n",
    "plot_prob_states()\n",
    "plot_obs_train()"
   ]
  },
  {
   "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
}