Koopman RL 3-26.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Koopman RL Report: 03.26.2021"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## OpenAI Gym Koopman Prediction Application: CartPole\n",
    "To see how well (and quickly) the Koopman operator could predict future states, I used a Q-learning agent that learned a near-optimal policy for the CartPole environment and fit an approximate Koopman operator to it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import numpy as np\n",
    "import scipy as sp\n",
    "import pykoopman as pk\n",
    "from pydmd import OptDMD\n",
    "import gym\n",
    "from sklearn.preprocessing import KBinsDiscretizer\n",
    "from typing import Tuple"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = np.load('/state-action-inputs.npy') # 20,000 entries\n",
    "X = X[:int(X.shape[0]*0.0015)] # 30 points!\n",
    "\n",
    "# Fit Koopman operator using closed-form solution to DMD\n",
    "optdmd = OptDMD(svd_rank=15)\n",
    "model_optdmd = pk.Koopman(regressor=optdmd)\n",
    "model_optdmd.fit(X)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env = gym.make('CartPole-v0')\n",
    "koopEnv = gym.make('CartPole-v0')\n",
    "\n",
    "Q_table = np.load('Q_table.npy')\n",
    "\n",
    "n_bins = ( 6, 12 )\n",
    "lower_bounds = [ env.observation_space.low[2], -math.radians(50) ]\n",
    "upper_bounds = [ env.observation_space.high[2], math.radians(50) ]\n",
    "\n",
    "def discretizer( _, __, angle, pole_velocity ) -> Tuple[int,...]:\n",
    "    \"\"\"Convert continuous state into a discrete state\"\"\"\n",
    "    est = KBinsDiscretizer(n_bins=n_bins, encode='ordinal', strategy='uniform')\n",
    "    est.fit([ lower_bounds, upper_bounds ])\n",
    "    return tuple( map( int, est.transform([[ angle, pole_velocity ]])[0] ) )\n",
    "\n",
    "def policy(state: tuple):\n",
    "    \"\"\" Choosing an action on epsilon-greedy policy \"\"\"\n",
    "    return np.argmax(Q_table[state])\n",
    "\n",
    "current_state = discretizer(*env.reset())\n",
    "current_stateK = discretizer(*koopEnv.reset())\n",
    "action = policy(current_state)\n",
    "actionK = policy(current_state)\n",
    "\n",
    "num_steps = 200\n",
    "q_learner_reward = 0\n",
    "koopman_reward = 0\n",
    "\n",
    "for i in range(num_steps):\n",
    "    # environment details\n",
    "    observation, reward, done, _ = env.step(action)\n",
    "    observationK, rewardK, doneK, _ = koopEnv.step(actionK)\n",
    "\n",
    "    # keep track of rewards\n",
    "    q_learner_reward += reward\n",
    "    koopman_reward += rewardK\n",
    "\n",
    "    # discretize state - hoping generator won't have to!\n",
    "    new_state = discretizer(*observation)\n",
    "    new_stateK = discretizer(*observationK)\n",
    "\n",
    "    # get actions\n",
    "    next_action = policy(new_state)\n",
    "    prediction = model_optdmd.predict(np.array([*list(current_stateK), actionK]))\n",
    "    prediction = np.round(np.real(prediction))\n",
    "    next_actionK = int(prediction[-1])\n",
    "\n",
    "    # update environments\n",
    "    action = next_action\n",
    "    actionK = next_actionK\n",
    "    current_state = new_state\n",
    "    current_stateK = new_stateK\n",
    "\n",
    "print(\"Q rewards:\", q_learner_reward)\n",
    "print(\"K rewards:\", koopman_reward)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see that the rewards are both 200 which means that the Koopman predictor works very well given good data, though there are plenty of papers on the subject. One thing we may want to look into is how well we can learn a controller from the Koopman operator, but since we are focused on the Generator operator, what we really want to see now is the predictive power of the Koopman Generator and how it can be used for control!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Stochastic Koopman Generator Analysis"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We simulated some paths from a standard Brownian Motion (drift coefficient 0 and diffusion coefficient 1) and then tried 3 different methods from two papers: \n",
    "+ Klus et al 2020 <https://arxiv.org/pdf/1909.10638.pdf>\n",
    "+ Li and Duan 2020 <https://arxiv.org/ftp/arxiv/papers/2005/2005.03769.pdf>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Simlulation of Brownian Data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We simulated 20 paths of standard BM each with 5000 steps in a time interval of size 5000 so that the time step was 1. We took each of these 20 paths to be a state variable in our state vector. Our state vector is thus comprised of 20 iid BMs. Formally, our state vector dynamics have the form\n",
    "$$\n",
    "    \\text{d}\\tilde X_t = b\\text{d}t + \\sigma\\text{d}W_t\n",
    "$$\n",
    "where $b=0$ is a $n=20$ dimensional vector of 0s and $\\sigma = I_{n\\times n}$ is a $n\\times n$ identity matrix."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Fitting the BM data using Generator EDMD (gEDMD)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from brownian import brownian\n",
    "\n",
    "# The Diffusion process parameter.\n",
    "sigma = 1\n",
    "# Total time.\n",
    "T = 10000\n",
    "# Number of steps.\n",
    "N = T # 10000\n",
    "# Time step size\n",
    "dt = T/N\n",
    "# Number of realizations to generate.\n",
    "m = 20\n",
    "# Create an empty array to store the realizations.\n",
    "X = np.empty((m, N+1))\n",
    "# Initial values of x.\n",
    "X[:, 0] = 50\n",
    "brownian(X[:, 0], N, dt, sigma, out=X[:, 1:])\n",
    "Z = np.roll(X,-1)[:, :-1]\n",
    "X = X[:, :-1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Fitting the BM data using Generator EDMD (gEDMD) from Klus et al. 2020"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "From the learner's point of view, we assume that we are in the class of continuous Markov processes and thus that the generator is of the form\n",
    "$$\n",
    "    \\mathcal{L}f = b\\cdot\\nabla_{\\tilde x}f + \\frac{1}{2}a:\\nabla^2_{\\tilde x}f \n",
    "    %= \\sum_{i=1}^n b_i\\frac{\\partial f}{\\partial \\tilde{x}_i} + \\frac{1}{2} \\sum_{i=1}^n \\sum_{j=1}^n a_{ij} \\frac{\\partial^2 f}{\\partial \\tilde{x}_i \\partial \\tilde{x}_j},\n",
    "$$\n",
    "where $a = \\sigma \\sigma^\\top$, $\\nabla^2_x$ denotes the Hessian, and $:$ denotes the double dot product. Applying the generator to each dictionary function \\psi_k and assuming that we have access to a single ergodic sample with time step $dt= 1$, we can use the following finite difference estimator of $d\\psi_k$:\n",
    "$$\n",
    "\\widehat{\\text{d}\\psi_k}(\\tilde{\\mathbf{x}}_l) = \\frac{1}{t}(\\tilde{\\mathbf{x}}_{l+1} - \\tilde{\\mathbf{x}}_l) \\cdot \\nabla\\psi_k(\\tilde{\\mathbf{x}}_l) + \\frac{1}{2t} \\Big[(\\tilde{\\mathbf{x}}_{l+1} - \\tilde{\\mathbf{x}}_l)(\\tilde{\\mathbf{x}}_{l+1} - \\tilde{\\mathbf{x}}_l)^\\top\\Big] : \\nabla^2 \\psi_k(\\tilde{\\mathbf{x}}_l)\n",
    "$$\n",
    "Note that we are adopting Klus's notation here only for reference. The stochastic total differential $d\\psi_k$ is a different object that the generator of the Koopman operator they are related in that the drift of the stochastic total differential is the same thing as the generator.\n",
    "\n",
    "Next, we set up matrices for the dictionary and the generator applied to it: \n",
    "$$\n",
    "    \\Psi_{\\mathbf{X}} = \\begin{bmatrix} \n",
    "                        \\psi_1(\\tilde{\\mathbf{x}}_1) & \\dots & \\psi_1(\\tilde{\\mathbf{x}}_m) \\\\\n",
    "                        \\vdots & \\ddots & \\vdots \\\\\n",
    "                        \\psi_k(\\tilde{\\mathbf{x}}_1) & \\dots & \\psi_k(\\tilde{\\mathbf{x}}_m) \n",
    "                    \\end{bmatrix},\n",
    "$$\n",
    "$$\n",
    "    \\text{d}\\Psi_{\\mathbf{X}} = \\begin{bmatrix} \n",
    "                        \\text{d}\\psi_1(\\tilde{\\mathbf{x}}_1) & \\dots & \\text{d}\\psi_1(\\tilde{\\mathbf{x}}_m) \\\\\n",
    "                        \\vdots & \\ddots & \\vdots \\\\\n",
    "                        \\text{d}\\psi_k(\\tilde{\\mathbf{x}}_1) & \\dots & \\text{d}\\psi_k(\\tilde{\\mathbf{x}}_m) \n",
    "                    \\end{bmatrix}\n",
    "$$\n",
    "\n",
    "The idea behind generator EDMD is that we assume that the genertor applied to the the dictionary functions can be (\"approximately\") expressed as a linear combination of the dictionary functions and find the coeficients of those linear combinations by minimizing $|| \\text{d}\\Psi_{\\tilde{\\mathbf{X}}} - M\\Psi_{\\tilde{\\mathbf{X}}} ||_F$ which leads to the least-squares approximation $$ M = \\text{d}\\Psi_{\\tilde{\\mathbf{X}}} \\Psi^{+}_{\\tilde{\\mathbf{X}}} = (\\text{d}\\Psi_{\\tilde{\\mathbf{X}}}\\Psi_{\\tilde{\\mathbf{X}}}^\\top)(\\Psi_{\\tilde{\\mathbf{X}}}\\Psi_{\\tilde{\\mathbf{X}}}^\\top)^+ $$\n",
    "\n",
    "Thus, we obtain the empirical estimate $L=M^T$ of the Koopman generator $\\mathcal{L}$.\n",
    "\n",
    "For the dictionary space, we chose monomials of up to order 2. We also tried monomials of order 1, which is not sufficient to pick up the diffusion term, but was successful at picking up the drift quicker. We hypothesize that this is because there are fewer terms in the regression over the dictionary functions.\n",
    "\n",
    "Note: we used Klus's [d3s repo](https://github.com/sklus/d3s) to create the PsiX and dPsiX objects\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import observables\n",
    "# from sympy import symbols\n",
    "# from sympy.polys.monomials import itermonomials, monomial_count\n",
    "# from sympy.polys.orderings import monomial_key\n",
    "\n",
    "# Construct B matrix as seen in 3.1.2 of the reference paper\n",
    "def constructB(d, k):\n",
    "    Bt = np.zeros((d, k))\n",
    "    if k == 1:\n",
    "        Bt[0,0] = 1\n",
    "    else:\n",
    "        num = np.arange(d)\n",
    "        Bt[num, num+1] = 1\n",
    "    B = Bt.T\n",
    "    return B\n",
    "\n",
    "# Construct similar B matrix as above, but for second order monomials\n",
    "def constructSecondOrderB(s, k):\n",
    "    Bt = np.zeros((s, k))\n",
    "    if k == 1:\n",
    "        Bt[0,0] = 1\n",
    "    else:\n",
    "        row = 0\n",
    "        for i in range(d+1, d+1+s):\n",
    "            Bt[row,i] = 1\n",
    "            row += 1\n",
    "    B = Bt.T\n",
    "    return B\n",
    "\n",
    "d = X.shape[0]\n",
    "m = X.shape[1]\n",
    "s = int(d*(d+1)/2) # number of second order poly terms\n",
    "rtoler=1e-02\n",
    "atoler=1e-02\n",
    "psi = observables.monomials(2)\n",
    "\n",
    "# x_str = \"\"\n",
    "# for i in range(d):\n",
    "#     x_str += 'x_' + str(i) + ', '\n",
    "# x_syms = symbols(x_str)\n",
    "# M = itermonomials(x_syms, 2)\n",
    "# sortedM = sorted(M, key=monomial_key('grlex', np.flip(x_syms)))\n",
    "# print(sortedM)\n",
    "\n",
    "Psi_X = psi(X)\n",
    "Psi_X_T = Psi_X.T\n",
    "nablaPsi = psi.diff(X)\n",
    "nabla2Psi = psi.ddiff(X)\n",
    "# print(\"nablaPsi Shape\", nablaPsi.shape)\n",
    "k = Psi_X.shape[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This computes dpsi_k(x) exactly as in the paper TODO: Put in exact reference here\n",
    "# t = 1 is a placeholder time step, not really sure what it should be\n",
    "def dpsi(k, l, t=1):\n",
    "    difference = (X[:, l+1] - X[:, l])\n",
    "    term_1 = (1/t) * (difference)\n",
    "    term_2 = nablaPsi[k, :, l]\n",
    "    term_3 = (1/(2*t)) * (difference.reshape(-1, 1) @ difference.reshape(1, -1))\n",
    "    term_4 = nabla2Psi[k, :, :, l]\n",
    "    return np.dot(term_1, term_2) + np.tensordot(term_3, term_4)\n",
    "vectorized_dpsi = np.vectorize(dpsi)\n",
    "\n",
    "# Construct \\text{d}\\Psi_X matrix\n",
    "dPsi_X = np.empty((k, m))\n",
    "for column in range(m-1):\n",
    "    dPsi_X[:, column] = vectorized_dpsi(range(k), column)\n",
    "\n",
    "# Calculate Koopman generator approximation\n",
    "train = int(m * 0.8)\n",
    "test = m - train\n",
    "M = dPsi_X[:, :train] @ np.linalg.pinv(Psi_X[:, :train]) # \\widehat{L}^\\top\n",
    "L = M.T # estimate of Koopman generator"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we outline 2 methods for how to identify the drift term $b$ from the diffusion equation using gEDMD which vary only in the last step of the procedure outlined below\n",
    "\n",
    "1. Get eigenfunctions $\\xi_i$ of $\\widehat{L}$ and store them in $\\Xi = \\{ \\xi_1, \\cdots, \\xi_n \\}$.\n",
    "\n",
    "2. Express eigenfunctions in terms of the dictionary $\\psi(x) = \\Xi^\\top \\psi(x)$\n",
    "\n",
    "3. Calculate Koopman modes $V = B^\\top \\Xi^{-\\top}$\n",
    "\n",
    "4. We can then use the Koopman generator approximation $\\widehat{L}$ directly and get drift $b$:\n",
    "    $$ \\mathcal{L}g(x) = b(x) \\approx (LB)^\\top  \\psi(x) $$\n",
    "    \n",
    "4. Alternatively, express drift in terms of reduced order eigenexpansion (we decided to set our default cutoff to be one tenth of the total eigenfunctions)\n",
    "    $$ (\\mathcal{L}g) (x) = b(x) \\approx \\sum_{\\ell = 1}^{cutoff} \\lambda_\\ell \\psi_\\ell(x) v_\\ell $$\n",
    "    where $v_\\ell$ is the $\\ell$th column vector of $V= B^\\top \\Xi$ where B is as above and $\\Xi$ is the matrix with column vectors as eigenvectors of $\\hat L$ and $\\lambda$'s are their associated eigenvalues.\n",
    "\n",
    "\n",
    "Note that the above method should work for any state dependent drift, including as a particular case, our constant drift vector $b(x) = b$.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Eigen decomposition\n",
    "eig_vals, eig_vecs = sp.sparse.linalg.eigs(L) if sp.sparse.issparse(L) else sp.linalg.eig(L)\n",
    "# Compute eigenfunction matrix\n",
    "eig_funcs = (eig_vecs).T @ Psi_X\n",
    "\n",
    "# Construct estimates of drift vector (b) and diffusion matrix (a) using two methods:\n",
    "# 1. Directly from dictionary functions without dimension reduction  \n",
    "# 2. Construct eigendecomposition and restrict its order\n",
    "\n",
    "# Construct B matrix that selects first-order monomials (except 1) when multiplied by list of dictionary functions\n",
    "B = constructB(d, k)\n",
    "# Construct second order B matrix (selects second-order monomials)\n",
    "second_orderB = constructSecondOrderB(s, k)\n",
    "\n",
    "# Computed b function (sometimes denoted by \\mu) without dimension reduction\n",
    "L_times_B_transposed = (L @ B).T\n",
    "def b(l):\n",
    "    return L_times_B_transposed @ Psi_X[:, l] # (k,)\n",
    "\n",
    "# Calculate Koopman modes\n",
    "V_v1 = B.T @ np.linalg.inv((eig_vecs).T)\n",
    "\n",
    "# The b_v2 function allows for heavy dimension reduction\n",
    "# default is reducing by 90% (taking the first k/10 eigen-parts)\n",
    "# TODO: Figure out correct place to take reals\n",
    "def b_v2(l, num_dims=k//10, V=V_v1):\n",
    "    res = 0\n",
    "    for ell in range(k-1, k-num_dims, -1):\n",
    "        res += eig_vals[ell] * eig_funcs[ell, l] * V[:, ell] #.reshape(-1, 1)\n",
    "    return np.real(res)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By calculating the b's separetely and printing out their values, we could tell that they were (relatively) close to zero (numbers were around the +/- 1e-2 area) which is exactly what we were hoping to find. It is nice that the heavily reduced b_v2 performed well because it means that we can actually take advantage of the computational benefits."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The two methods to identify $\\sigma(x)$ are as follows\n",
    "\n",
    "1. Use the observable $g(x)\\in \\mathbb{R}^{m}$ where $m = d(d+1)/2$ as all monomials of order 2, i.e. $g(x) = [x_1^2,\\, x_1x_2,\\, x_2^2,\\, x_1x_3,\\, x_2x_3,\\, x_3^2,\\, \\cdots]^\\top$\n",
    "2. Find $B \\in \\mathbb{R}^{k \\times s}$ s.t. $B^\\top \\psi(x) = g(x)$\n",
    "3.  \n",
    "\\begin{align}\n",
    "    \\mathcal{L}g(x) &= \\nabla g(x) b(x) + \\text{vecUpTri}(a)\\\\\n",
    "    &= \\nabla(B^\\top \\psi(x)) b(x) + \\text{vecUpTri}(a)\\\\\n",
    "    &= B^\\top \\nabla \\psi(x) b(x) + \\text{vecUpTri}(a)\n",
    "\\end{align}$\\;$\n",
    "> where vecUpTri($a$) is a vectorized (flattened) upper triangle of the a matrix (including the diagonal)         which follows the same indexing convention as $g(x)$ above.\n",
    "       \n",
    "4. Given an estimate for $\\widehat{b}(x)$, and using $\\mathcal{L}g(x) \\approx (\\widehat{L} B)^\\top \\psi(x)$ we can estimate the upper triangular entries (including the diagonal) of matrix $a$ as\n",
    "$$ \\implies \\text{vecUpTri}(\\hat a) \\approx (\\widehat{L}B)^\\top \\psi(x) - B^\\top \\nabla \\psi(x) \\widehat{b}(x) $$ \n",
    "Since $a = \\sigma\\sigma^\\top$, it is symmetric, and having estimates of the upper triangular entries means that we have an estimate for $a$ as well. \n",
    "\n",
    "5. Alternatively, we can use similar eigendecomposition reduction trick as we did with the drift above. To do do, we express the generator applied to the second order monomials $g(x)$ in terms of reduced order eigenexpansion (we decided to set our default cutoff to be one tenth of the total eigenfunctions)\n",
    "    $$ (\\mathcal{L}g) (x) \\approx \\sum_{\\ell = 1}^{cutoff} \\lambda_\\ell \\psi_\\ell(x) v_\\ell $$\n",
    "    where $v_\\ell$ is the $\\ell$th column vector of $V = B^\\top \\Xi^{-\\top}$ where $B$ is as above and $\\Xi$ is the matrix with column vectors as eigenvectors of $\\hat L$ and $\\lambda$'s are their associated eigenvalues.\n",
    "\n",
    "    With this, we can estimate the upper triangle of matrix $a$ as: $\\;$\n",
    "    $$ \\text{vecUpTri}(\\hat a) \\approx \\sum_{\\ell = 1}^{cutoff} \\lambda_\\ell \\psi_\\ell(x) v_\\ell  - B^\\top \\nabla \\psi(x) \\widehat{b}(x) $$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "L_times_second_orderB_transpose = (L @ second_orderB).T\n",
    "\n",
    "def a(l):\n",
    "    return (L_times_second_orderB_transpose @ Psi_X[:, l]) - \\\n",
    "        (second_orderB.T @ nablaPsi[:, :, l] @ b(l))\n",
    "\n",
    "V_v2 = second_orderB.T @ np.linalg.inv((eig_vecs).T)\n",
    "def a_v2(l):\n",
    "    return (b_v2(l, V=V_v2)) - \\\n",
    "        (second_orderB.T @ nablaPsi[:, :, l] @ b_v2(l))\n",
    "\n",
    "# Reshape a vector as matrix and perform some tests\n",
    "def covarianceMatrix(a_func, l):\n",
    "    a_l = a_func(l)\n",
    "    covariance = np.zeros((d, d))\n",
    "    row = 0\n",
    "    col = 0\n",
    "    covariance[row, col] = a_l[0]\n",
    "    col += 1\n",
    "    n = 1\n",
    "    while col < d:\n",
    "        covariance[row, col] = a_l[n]\n",
    "        covariance[col, row] = a_l[n]\n",
    "        if row == col: \n",
    "            col += 1\n",
    "            row = 0\n",
    "        else:\n",
    "            row += 1\n",
    "        n +=1\n",
    "    return covariance"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Unfortunately, the results of both a and a_v2 are far from identity matrices. In fact, most of the entries were on the scale of +/- 1e1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We tried one final appoach that was adapted from Lu and Duan 2021. They exploit the Perron-Frobenius operator (PF)(solution to the Fokker-Planck equation) to analyze systems with non-Gaussian Levy noise in addition to Gaussian noise. We restricted their method to analyze our system which just has Guassian noise. We will describe the estimation procedure here, but for the theory, please see their paper.\n",
    "\n",
    "First, we assume that the drift can be approximated as a linear combination of the dictionary functions, $b(x) \\approx \\psi(x)^\\top C$ where $C$ is a $k\\times d$ matrix. Next, we collect a sample of the finite differences between observations $S\\in \\mathbb{R}^{d\\times m}$ where the jth column vector is\n",
    "$$\n",
    "S_j = (x_j - z_j)/h\n",
    "$$\n",
    "where $x_j$ is the jth snapshot vector and $z_j=x_{j+h}$ is the associated future value where we take time step $h=1$. To find the coeficient matrix $C$, we solve the following least squares problem\n",
    "$$\n",
    "min_{C} ||\\Psi_X^\\top C - S^\\top||_F\n",
    "$$\n",
    "which leads to the solution\n",
    "$$\n",
    "\\hat C = (\\Psi_X \\Psi_X)^{-1}(\\Psi_XB^\\top)\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "B = np.zeros((d, m))\n",
    "m_range = np.arange(m)\n",
    "B = X[:, m_range] - Z[:, m_range]\n",
    "# print(\"B shape:\", B.shape)\n",
    "# print(\"Psi_X transpose shape:\", Psi_X_T.shape)\n",
    "PsiMult = sp.linalg.inv(Psi_X @ Psi_X_T) @ Psi_X\n",
    "C = PsiMult @ B.T\n",
    "# Each col of matric C represents the coeficients in a linear combo of the dictionary functions that makes up each component of the drift vector. So each c_{} \n",
    "# print(\"C shape:\", C.shape)\n",
    "\n",
    "b_v3 = C.T @ Psi_X\n",
    "# for l in range(5):\n",
    "#     print(b_v3[:, l])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This new version of b reached good results, but not quite as good as either b or b_v2 from earlier. The entries were on the scale of +/- 1e-1."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, to calculate the diffusion matrix using Li and Duan's approach, we make the same assumption about each element of the diffusion matrix as being approximated as a linear combination of the dictionary functions \n",
    "$$\n",
    "a_{ij}(x) \\approx d_ij \\cdot \\psi(x)\n",
    "$$\n",
    "where $d_ij = [d_{ij,1},\\, \\cdots,\\, d_{ij,k}]^\\top$. For the diffusion terms, we need to collect samples of the (cross) quadratic variation\n",
    "$$\n",
    "B_{ij} = h^{-1}[(x_{i,1}-z_{i,1})(x_{j,1}-z_{j,1})]\n",
    "$$\n",
    "where $i$ and $j$ represent the $i$-th and $j$-th components of the snapshot vectors. The resulting least squares problem to find the coeficients $d_{ij}$ for each ($i$, $j$) pair is\n",
    "$$\n",
    "||\\Psi_X^\\top d_{ij} - B_{ij}||\n",
    "$$\n",
    "which results in the estimate \n",
    "$$\n",
    "\\hat{d}_{ij} = (\\Psi_X\\Psi_X^\\top)^{-1}(\\Psi_X B_{ij})\n",
    "$$\n",
    "TODO: Figure out how to vectorize/matrix-ize for faster computation!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def a_v3(l):\n",
    "    diffusionDictCoefs = np.empty((d, d, k))\n",
    "    diffusionMat = np.empty((d, d))\n",
    "    for i in range(d):\n",
    "        for j in range(d):\n",
    "            Bij = B[i]*B[j]\n",
    "            diffusionDictCoefs[i, j] = PsiMult @ Bij\n",
    "            diffusionMat[i, j] = np.dot(diffusionDictCoefs[i, j], Psi_X[:,l])\n",
    "    return diffusionMat"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This version of a  achieved the results we were hoping for and returned near identity matrices with numbers of about +/- 1e-1 in the entries outside the diagonal."
   ]
  }
 ],
 "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}