Duffing_model.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e71c6a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib nbagg\n",
    "import math\n",
    "import numpy as np\n",
    "import scipy as sp\n",
    "from scipy import linalg\n",
    "from scipy import integrate\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import cm\n",
    "\n",
    "import pickle\n",
    "\n",
    "from actuated_Koopman_EM_4 import ActuatedKoopmanEM, save_EM_model, load_EM_model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f33c3ba5",
   "metadata": {},
   "source": [
    "# Duffing equation\n",
    "\n",
    "We consider the unforced Duffing equation,\n",
    "\\begin{equation}\n",
    "    \\ddot{y} = - \\delta \\dot{y} - y(\\beta + \\alpha y^2),\n",
    "\\end{equation}\n",
    "with parameters $\\delta=0.5$, $\\beta=-1$, and $\\alpha=1$, which were used in: \n",
    "\n",
    "M. O. Williams, I. G. Kevrekidis, and C. W. Rowley, ``A data-driven approximation of the Koopman operator: Extending dynamic mode decomposition'', J. Nonlinear Science, 2015\n",
    "\n",
    "and in\n",
    "\n",
    "S. E. Otto and C. W. Rowley, ``Linearly recurrent autoencoder networks for learning dynamics'', SIAM J. Applied Dynamical Systems, 2019\n",
    "\n",
    "With these parameters, the Duffing equation has stable fixed points at $x=\\pm 1$ with eigenvalues $\\lambda_{1,2} = \\frac{1}{4} \\left(-1 \\pm \\sqrt{31} i\\right)$ and a saddle point at the origin with eigenvalues $\\lambda_{3,4} = \\frac{1}{4} \\left(-1 \\pm \\sqrt{17}\\right)$. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae431511",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Duffing equation state dimension\n",
    "n_states = 2\n",
    "\n",
    "# number of output observables\n",
    "n_outputs = 2\n",
    "\n",
    "# paramters\n",
    "delta = 0.5\n",
    "beta = -1.0\n",
    "alpha = 1.0\n",
    "\n",
    "# Duffing equation as 1st order system\n",
    "# x[0] = y, x[1] = dy/dt\n",
    "f_sys = lambda t, x: (x[1], -delta*x[1] - x[0]*(beta + alpha*np.square(x[0])))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55a81298",
   "metadata": {},
   "outputs": [],
   "source": [
    "lam_sys = np.array([-1+np.sqrt(31)*1j, \\\n",
    "                    -1-np.sqrt(31)*1j, \\\n",
    "                    -1+np.sqrt(17), \\\n",
    "                    -1-np.sqrt(17)]) / 4.0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "264e548a",
   "metadata": {},
   "source": [
    "## Generate data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab0b9153",
   "metadata": {},
   "outputs": [],
   "source": [
    "# length of trajectories and sampling interval\n",
    "n_t = 801\n",
    "t_data = np.linspace(0, 16, n_t)\n",
    "dt = t_data[1] - t_data[0]\n",
    "\n",
    "# initial condition domain\n",
    "ICbox = [-2.0, 2.0, -2.0, 2.0]\n",
    "\n",
    "# number of trajectories to sample\n",
    "n_traj = 50\n",
    "\n",
    "X_data = np.zeros((n_traj, n_t, n_states))\n",
    "Y_data = np.zeros((n_traj, n_t, n_outputs))\n",
    "U_data = np.ones((n_traj, n_t-1, 1))\n",
    "for k in range(n_traj):\n",
    "    # initial condition\n",
    "    x0 = np.zeros(n_states)\n",
    "    x0[0] = np.random.uniform(ICbox[0], ICbox[1])\n",
    "    x0[1] = np.random.uniform(ICbox[2], ICbox[3])\n",
    "    \n",
    "    sol = sp.integrate.solve_ivp(fun = f_sys, \\\n",
    "                                 t_span = (t_data[0], t_data[-1]), \\\n",
    "                                 y0 = x0, \\\n",
    "                                 t_eval = t_data)\n",
    "    \n",
    "    X_data[k, :, :] = sol.y.T\n",
    "    # Y_data[k, :, 0] = sol.y[0, :]\n",
    "    Y_data[k, :, :] = sol.y[:n_outputs, :].T\n",
    "\n",
    "# noise to be added to observation data\n",
    "Noise_data = 0.01*np.random.randn(n_traj, n_t, n_outputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4871caac",
   "metadata": {},
   "source": [
    "## Save data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75400895",
   "metadata": {},
   "outputs": [],
   "source": [
    "fname = 'models/Duffing/eval_data_fullobs_fine_tmp.p'\n",
    "\n",
    "data = {'X_data': X_data, \\\n",
    "        'Y_data': Y_data, \\\n",
    "        'Noise_data': Noise_data, \\\n",
    "        'U_data': U_data, \\\n",
    "        't_data': t_data, \\\n",
    "        'n_traj': n_traj, \\\n",
    "        'n_t': n_t, \\\n",
    "        'dt': dt, \\\n",
    "        'ICbox': ICbox}\n",
    "\n",
    "with open( fname, \"wb\" ) as f:\n",
    "    pickle.dump(data, f)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "036e7fbe",
   "metadata": {},
   "source": [
    "## Load data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "97aeb965",
   "metadata": {},
   "outputs": [],
   "source": [
    "fname = 'models/Duffing/train_data_fullobs_fine.p'\n",
    "\n",
    "with open(fname,'rb') as f:\n",
    "    data = pickle.load(f)\n",
    "\n",
    "X_train = data['X_data']\n",
    "Y_train = data['Y_data']\n",
    "Noise_train = data['Noise_data']\n",
    "U_train = data['U_data']\n",
    "t_train = data['t_data']\n",
    "n_traj_train = data['n_traj']\n",
    "n_t_train = data['n_t']\n",
    "dt_train = data['dt']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd70eeb3",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "for k in range(n_traj_train):\n",
    "    plt.plot(X_train[k,:,0], X_train[k,:,1], 'b.-')\n",
    "    \n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "09fe7253",
   "metadata": {},
   "source": [
    "# EDMD with Legendre polynomials"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4865c02b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def SpecialPolyDict(X, degree, prnt=False):\n",
    "    # Construct a dictionary of special polynomials.\n",
    "    # X is an n_dim by m_samples matrix\n",
    "    # degree is the maximum power of each feature in X\n",
    "    # returns Phi, an (degree+1)^n_dim by m_samples matrix\n",
    "    #     whose rows contain all special polynomial products\n",
    "    #     of the features in X up to the given degree\n",
    "    \n",
    "    n_dim = X.shape[0]\n",
    "    m_samples = X.shape[1]\n",
    "    N_dict = (degree+1)**n_dim\n",
    "    \n",
    "    Phi = np.ones((N_dict, m_samples))\n",
    "    \n",
    "    # powers is a vector containing the powers on each variable\n",
    "    # the entries of powers are found by counting in base (degree + 1)\n",
    "    powers = np.zeros(n_dim, dtype=int)\n",
    "    if prnt:\n",
    "        print(powers)\n",
    "    \n",
    "    for k in range(1, N_dict):\n",
    "        # increment the elements of powers (by counting)\n",
    "        powers[0] = powers[0] + 1\n",
    "        # carry the ones\n",
    "        for i in range(n_dim-1):\n",
    "            if powers[i] == degree + 1:\n",
    "                # carry the one\n",
    "                powers[i] = 0\n",
    "                powers[i+1] = powers[i+1] + 1\n",
    "        \n",
    "        if prnt:\n",
    "            print(powers)\n",
    "        \n",
    "        # construct kth special monomial with powers on each variable\n",
    "        # in X given by powers\n",
    "        for i in range(n_dim):\n",
    "            # Phi[k,:] = Phi[k,:] * sp.special.eval_hermite(powers[i], X[i,:])\n",
    "            Phi[k,:] = Phi[k,:] * sp.special.eval_legendre(powers[i], X[i,:])\n",
    "    \n",
    "    return Phi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ed50785",
   "metadata": {},
   "outputs": [],
   "source": [
    "# data matrices\n",
    "X = np.vstack(X_train).T\n",
    "X_Phi = np.vstack(X_train).T\n",
    "X_Phi_IC = np.vstack(X_train[:,0,:]).T\n",
    "\n",
    "# time-shifted data matrices\n",
    "X_Phi0 = np.vstack(X_train[:,:-1,:]).T\n",
    "X_Phi1 = np.vstack(X_train[:,1:,:]).T\n",
    "\n",
    "# construct feature matrices\n",
    "deg = 3\n",
    "Phi = SpecialPolyDict(X_Phi/2.0, degree=deg, prnt=False)\n",
    "Phi_IC = SpecialPolyDict(X_Phi_IC/2.0, degree=deg)\n",
    "Phi0 = SpecialPolyDict(X_Phi0/2.0, degree=deg)\n",
    "Phi1 = SpecialPolyDict(X_Phi1/2.0, degree=deg)\n",
    "\n",
    "print(Phi.shape[0])\n",
    "\n",
    "## Non-truncated EDMD\n",
    "\n",
    "# find matrix approximation of Koopman generator\n",
    "V_EDMD, residuals, _, _ = np.linalg.lstsq(Phi0.T, (Phi1.T - Phi0.T)/dt_train, rcond=None)\n",
    "V_EDMD[:,0] = 0.0\n",
    "\n",
    "mse = sum(residuals)/Phi0.shape[1]\n",
    "print(mse)\n",
    "\n",
    "# output reconstruction map\n",
    "C_EDMD = np.zeros((2,Phi.shape[0]))\n",
    "C_EDMD[0,1] = 2.0\n",
    "C_EDMD[1,deg+1] = 2.0\n",
    "\n",
    "# ## SVD-truncated EDMD\n",
    "# r = 15\n",
    "# U_Phi, s_Phi, VT_Phi = sp.linalg.svd(Phi[1:,:], full_matrices=False, compute_uv=True)\n",
    "# Phi_r = np.vstack([Phi[0,:], np.reshape(s_Phi[:r-1], (-1,1))*VT_Phi[:r-1,:]])\n",
    "# Phi_IC_r = np.vstack([Phi_IC[0,:], np.dot(U_Phi[:,:r-1].T, Phi_IC[1:,:])])\n",
    "# Phi0_r = np.vstack([Phi0[0,:], np.dot(U_Phi[:,:r-1].T, Phi0[1:,:])])\n",
    "# Phi1_r = np.vstack([Phi1[0,:], np.dot(U_Phi[:,:r-1].T, Phi1[1:,:])])\n",
    "\n",
    "# # find matrix approximation of Koopman generator\n",
    "# V_EDMD_r, residuals, _, _ = np.linalg.lstsq(Phi0_r.T, (Phi1_r.T - Phi0_r.T)/dt_train, rcond=None)\n",
    "# V_EDMD_r[:,0] = 0.0\n",
    "\n",
    "# mse = sum(residuals)/Phi0_r.shape[1]\n",
    "# print(mse)\n",
    "\n",
    "# # output reconstruction map\n",
    "# C_EDMD_r = np.transpose(np.linalg.lstsq(Phi_r.T, np.vstack(Y_train), rcond=None)[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45d103f2",
   "metadata": {},
   "source": [
    "## Initial spectrum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8126b975",
   "metadata": {},
   "outputs": [],
   "source": [
    "def eig_transform(A):\n",
    "    lam, Vl, Vr = sp.linalg.eig(A, left=True, right=True)\n",
    "    w = np.diag(np.dot(np.conj(Vr).T, Vl))\n",
    "    Vl = Vl/w\n",
    "    return lam, Vl, Vr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6e96a2d",
   "metadata": {},
   "outputs": [],
   "source": [
    "lam_EDMD, Vl_EDMD, Vr_EDMD = eig_transform(V_EDMD.T)\n",
    "# lam_EDMD, Vl_EDMD, Vr_EDMD = eig_transform(V_EDMD_r.T)\n",
    "\n",
    "plt.figure()\n",
    "plt.plot(np.real(lam_EDMD), np.imag(lam_EDMD), 'bo')\n",
    "plt.plot(np.real(lam_sys), np.imag(lam_sys), 'kx')\n",
    "plt.axvline(x=0, color='k')\n",
    "plt.axis('equal')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f3d4950",
   "metadata": {},
   "source": [
    "### Eigenfunction values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "84c78ffc",
   "metadata": {},
   "outputs": [],
   "source": [
    "e_fun_vals_EDMD = np.dot(np.conj(Vl_EDMD).T, Phi)\n",
    "# e_fun_vals_EDMD = np.dot(np.conj(Vl_EDMD).T, Phi_r)\n",
    "\n",
    "efun_idcs = np.argsort(np.real(lam_EDMD))[::-1]\n",
    "\n",
    "for i in range(len(efun_idcs)):\n",
    "    print('lam_{:d} = {:f} + {:f}i'.format(i, \\\n",
    "                                           lam_EDMD[efun_idcs[i]].real, \\\n",
    "                                           lam_EDMD[efun_idcs[i]].imag))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "238d8d31",
   "metadata": {},
   "outputs": [],
   "source": [
    "ix = 1\n",
    "efun_idx = efun_idcs[ix]\n",
    "\n",
    "# color scaling\n",
    "vmean = np.mean(np.reshape(np.real(e_fun_vals_EDMD[efun_idx,:]), (-1)))\n",
    "vscl = np.max(np.absolute(np.reshape(np.real(e_fun_vals_EDMD[efun_idx,:]), (-1)) - vmean))\n",
    "# vscl = np.max(np.reshape(np.absolute(np.real(e_fun_vals_EDMD[efun_idx,:])), (-1)))\n",
    "\n",
    "plt.figure()\n",
    "\n",
    "plt.scatter(X_Phi[0,:], \\\n",
    "            X_Phi[1,:], \\\n",
    "            c = np.real(e_fun_vals_EDMD[efun_idx,:]), \\\n",
    "            cmap = cm.seismic, \\\n",
    "            vmin = vmean-vscl, \\\n",
    "            vmax = vmean+vscl, \\\n",
    "            s=10)\n",
    "\n",
    "plt.title('phi_{:d}, lam_{:d} = {:.5f} + {:.5f}i'.format(ix, ix, lam_EDMD[efun_idx].real, lam_EDMD[efun_idx].imag))\n",
    "\n",
    "# plt.savefig('figures/Duffing/EDMD_dim{:d}_efun{:d}_real.png'.format(V_EDMD.shape[0], ix), dpi=None, facecolor='w', edgecolor='w', \\\n",
    "#             transparent=False, bbox_inches=None, pad_inches=0.0)\n",
    "# plt.savefig('figures/Duffing/EDMD_dim{:d}_efun{:d}_real.eps'.format(V_EDMD.shape[0], ix), dpi=None, facecolor='w', edgecolor='w', \\\n",
    "#             transparent=False, bbox_inches=None, pad_inches=0.0)\n",
    "\n",
    "plt.show()\n",
    "\n",
    "\n",
    "# color scaling\n",
    "vscl = np.max(np.reshape(np.absolute(e_fun_vals_EDMD[efun_idx,:]), (-1)))\n",
    "\n",
    "plt.figure()\n",
    "\n",
    "plt.scatter(X_Phi[0,:], \\\n",
    "            X_Phi[1,:], \\\n",
    "            c = np.absolute(e_fun_vals_EDMD[efun_idx,:]), \\\n",
    "            cmap = cm.inferno, \\\n",
    "            vmin = 0, \\\n",
    "            vmax = vscl, \\\n",
    "            s=10)\n",
    "\n",
    "plt.title('|phi_{:d}|, lam_{:d} = {:.5f} + {:.5f}i'.format(ix, ix, lam_EDMD[efun_idx].real, lam_EDMD[efun_idx].imag))\n",
    "\n",
    "# plt.savefig('figures/Duffing/EDMD_dim{:d}_efun{:d}_mag.png'.format(V_EDMD.shape[0], ix), dpi=None, facecolor='w', edgecolor='w', \\\n",
    "#             transparent=False, bbox_inches=None, pad_inches=0.0)\n",
    "# plt.savefig('figures/Duffing/EDMD_dim{:d}_efun{:d}_mag.eps'.format(V_EDMD.shape[0], ix), dpi=None, facecolor='w', edgecolor='w', \\\n",
    "#             transparent=False, bbox_inches=None, pad_inches=0.0)\n",
    "\n",
    "plt.show()\n",
    "\n",
    "# plot angle of complex eigenfunction\n",
    "plt.figure()\n",
    "\n",
    "plt.scatter(X_Phi[0,:], \\\n",
    "            X_Phi[1,:], \\\n",
    "            c = np.angle(e_fun_vals_EDMD[efun_idx,:]), \\\n",
    "            cmap = cm.twilight, \\\n",
    "            vmin = -np.pi, \\\n",
    "            vmax = np.pi, \\\n",
    "            s=10)\n",
    "\n",
    "plt.title('arg(phi_{:d}), lam_{:d} = {:.5f} + {:.5f}i'.format(ix, ix, lam_EDMD[efun_idx].real, lam_EDMD[efun_idx].imag))\n",
    "\n",
    "# plt.savefig('figures/Duffing/EDMD_dim{:d}_efun{:d}_arg.png'.format(V_EDMD.shape[0], ix), dpi=None, facecolor='w', edgecolor='w', \\\n",
    "#             transparent=False, bbox_inches=None, pad_inches=0.0)\n",
    "# plt.savefig('figures/Duffing/EDMD_dim{:d}_efun{:d}_arg.eps'.format(V_EDMD.shape[0], ix), dpi=None, facecolor='w', edgecolor='w', \\\n",
    "#             transparent=False, bbox_inches=None, pad_inches=0.0)\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2404d4f5",
   "metadata": {},
   "source": [
    "# Train Koopman EM model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f012be55",
   "metadata": {},
   "source": [
    "## Initialize model with EDMD"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef43943f",
   "metadata": {},
   "outputs": [],
   "source": [
    "## initial generator matrix\n",
    "V_init = V_EDMD.T\n",
    "\n",
    "lam_init, _, _ = eig_transform(V_init)\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "ax.plot(np.real(lam_EDMD), np.imag(lam_EDMD), 'bo')\n",
    "ax.plot(np.real(lam_init), np.imag(lam_init), 'gd')\n",
    "ax.plot(np.real(lam_sys), np.imag(lam_sys), 'kx')\n",
    "ax.axvline(x=0, color='k')\n",
    "circle = plt.Circle((-1.0/dt_train, 0), radius=1.0/dt_train, edgecolor='k', facecolor='None')\n",
    "ax.add_patch(circle)\n",
    "plt.axis('equal')\n",
    "plt.draw()\n",
    "\n",
    "## estimate process noise covariance\n",
    "E = Phi0 + dt_train*np.dot(V_init, Phi0) - Phi1\n",
    "Sig_v_init = np.dot(E[1:,:], E[1:,:].T)/E.shape[1]\n",
    "\n",
    "## estimate distribution of lifted initial states\n",
    "\n",
    "# initialize state mean\n",
    "Z = Phi_IC[1:,:].T\n",
    "mu_0_EDMD = np.mean(Z, axis=0)\n",
    "\n",
    "# initialize state covariance\n",
    "Zc = Z - mu_0_EDMD\n",
    "Sig_0_EDMD = np.dot(Zc.T, Zc)/Zc.shape[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02299a91",
   "metadata": {},
   "outputs": [],
   "source": [
    "y_dim = 2\n",
    "u_dim = 1\n",
    "z_dim = V_init.shape[0]\n",
    "\n",
    "# initialize generator\n",
    "V_matrices = np.zeros((u_dim,z_dim,z_dim))\n",
    "V_matrices[0,:,:] = V_init\n",
    "\n",
    "\n",
    "# initialize process noise covariance\n",
    "Sig_v = Sig_v_init\n",
    "\n",
    "# observation matrix (fixed)\n",
    "C = C_EDMD[:,1:]\n",
    "\n",
    "# observation offset\n",
    "h = C_EDMD[:,0]\n",
    "\n",
    "# initialize observation noise covariance\n",
    "Sig_w = 0.01*np.eye(y_dim)\n",
    "\n",
    "# initialize initial state mean\n",
    "mu_0 = mu_0_EDMD\n",
    "\n",
    "# initialize initial state covariance\n",
    "Sig_0 = Sig_0_EDMD\n",
    "\n",
    "EM_model = ActuatedKoopmanEM(Y_train+Noise_train, dt_train*U_train, \\\n",
    "                             V_matrices, C, h, Sig_v, Sig_w, mu_0, Sig_0, \\\n",
    "                             compute_log_likelihood=True)\n",
    "\n",
    "# save the model\n",
    "fname = 'models/Duffing/model_dim{:d}'.format(z_dim)\n",
    "save_EM_model(EM_model, fname, Y_train+0.0*Noise_train, dt_train*U_train)\n",
    "\n",
    "fname = 'models/Duffing/EDMD_model_dim{:d}'.format(z_dim)\n",
    "save_EM_model(EM_model, fname, Y_train+0.0*Noise_train, dt_train*U_train)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "75f603d1",
   "metadata": {},
   "source": [
    "## Train model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07dcd00f",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "delta_L_tol = 1.0e-2\n",
    "\n",
    "# number of steps between log likelihood calculations\n",
    "n_skip_L_calc = 5\n",
    "\n",
    "# model dimension\n",
    "# z_dim = 16\n",
    "\n",
    "# Load the model\n",
    "fname = 'models/Duffing/model_dim{:d}'.format(z_dim)\n",
    "EM_model = ActuatedKoopmanEM()\n",
    "EM_model, _, _ = load_EM_model(EM_model, fname)\n",
    "\n",
    "# log likelihood of the model\n",
    "L = EM_model.L\n",
    "delta_L = np.inf\n",
    "\n",
    "## if the log likelihood has not converged continue running, \n",
    "# run the optimization\n",
    "while delta_L > delta_L_tol:\n",
    "    \n",
    "    print('dimension {:d} iteration {:d} \\n'.format(z_dim, EM_model.EM_iter+1))\n",
    "    if EM_model.EM_iter%n_skip_L_calc == 0 and EM_model.EM_iter > 0:\n",
    "        # run EM step and compute log likelihood\n",
    "        EM_model.run_EM_step(explicit_time_step = True, optimize_IC=True, \n",
    "                             optimize_observation_map=False, optimize_process_noise=True, \\\n",
    "                             compute_log_likelihood=True)\n",
    "        \n",
    "        # average change in log likelihood\n",
    "        delta_L = (EM_model.L - L)/n_skip_L_calc\n",
    "        print('delta_L = {:3e}\\n'.format(delta_L))\n",
    "        \n",
    "        # update the log likelihood\n",
    "        L = EM_model.L\n",
    "        \n",
    "        ## Save the model\n",
    "        if delta_L > 0:\n",
    "            save_EM_model(EM_model, fname, Y_train+0.0*Noise_train, dt_train*U_train)\n",
    "        \n",
    "        # plot eigenvalues\n",
    "        lam_EM, Vl_EM, Vr_EM = eig_transform(EM_model.V_matrices[0,:,:].T)\n",
    "        \n",
    "        fig, ax = plt.subplots()\n",
    "        ax.plot(np.real(lam_EDMD), np.imag(lam_EDMD), 'bo')\n",
    "        ax.plot(np.real(lam_EM), np.imag(lam_EM), 'gd')\n",
    "        ax.plot(np.real(lam_sys), np.imag(lam_sys), 'kx')\n",
    "        ax.axvline(x=0, color='k')\n",
    "        circle = plt.Circle((-1.0/dt_train, 0), radius=1.0/dt_train, edgecolor='k', facecolor='None')\n",
    "        ax.add_patch(circle)\n",
    "        plt.axis('equal')\n",
    "        plt.draw()\n",
    "        \n",
    "    else:\n",
    "        ## run an EM step without computing log likelihood\n",
    "        EM_model.run_EM_step(explicit_time_step = True, optimize_IC=True, \\\n",
    "                             optimize_observation_map=False, optimize_process_noise=True, \\\n",
    "                             compute_log_likelihood=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb8549b9",
   "metadata": {},
   "source": [
    "## Model spectrum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "813694c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "z_dim = 16\n",
    "\n",
    "fname = 'models/Duffing/model_dim{:d}'.format(z_dim)\n",
    "EM_model = ActuatedKoopmanEM()\n",
    "EM_model, _, _ = load_EM_model(EM_model, fname)\n",
    "\n",
    "lam_EM, Vl_EM, Vr_EM = eig_transform(EM_model.V_matrices[0,:,:])\n",
    "\n",
    "# lam_EM, Vl_EM, Vr_EM = eig_transform(V_init)\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "ax.plot(np.real(lam_EDMD), np.imag(lam_EDMD), 'gd', label='EDMD')\n",
    "ax.plot(np.real(lam_EM), np.imag(lam_EM), 'bo', label='EM model')\n",
    "ax.plot(np.real(lam_sys), np.imag(lam_sys), 'kx', label='Duffing fixed points')\n",
    "ax.axvline(x=0, color='k')\n",
    "circle = plt.Circle((-1.0/dt_train, 0), radius=1.0/dt_train, \\\n",
    "                    edgecolor='k', \\\n",
    "                    linestyle='--', \\\n",
    "                    facecolor='None', \\\n",
    "                    label='Explicit Euler stability')\n",
    "ax.add_patch(circle)\n",
    "# ax.set(xlim=(-8, 2), ylim=(-6, 6))\n",
    "plt.legend()\n",
    "plt.axis('equal')\n",
    "\n",
    "# plt.savefig('figures/Duffing/eigenvalues_dim{:d}.png'.format(z_dim), dpi=None, facecolor='w', edgecolor='w', \\\n",
    "#             transparent=False, bbox_inches=None, pad_inches=0.0)\n",
    "# plt.savefig('figures/Duffing/eigenvalues_dim{:d}.eps'.format(z_dim), dpi=None, facecolor='w', edgecolor='w', \\\n",
    "#             transparent=False, bbox_inches=None, pad_inches=0.0)\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a43d4382",
   "metadata": {},
   "outputs": [],
   "source": [
    "# model state estimates on training data\n",
    "Z_EM = np.vstack(EM_model.muhats).T\n",
    "Z_EM = np.vstack([np.ones((1,Z_EM.shape[1])), Z_EM])\n",
    "\n",
    "# transform into eigen-coordinates of drift generator\n",
    "e_fun_vals_EM = np.dot(np.conj(Vl_EM).T, Z_EM)\n",
    "\n",
    "efun_idcs = np.argsort(np.real(lam_EM))[::-1]\n",
    "\n",
    "for i in range(len(efun_idcs)):\n",
    "    print('lam_{:d} = {:f} + {:f}i'.format(i, \\\n",
    "                                           lam_EM[efun_idcs[i]].real, \\\n",
    "                                           lam_EM[efun_idcs[i]].imag))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df4478f3",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "ix = 0\n",
    "efun_idx = efun_idcs[ix]\n",
    "\n",
    "# color scaling\n",
    "vmean = np.mean(np.reshape(np.real(e_fun_vals_EM[efun_idx,:]), (-1)))\n",
    "vscl = np.max(np.absolute(np.reshape(np.real(e_fun_vals_EM[efun_idx,:]), (-1)) - vmean))\n",
    "#vscl = np.max(np.reshape(np.absolute(np.real(e_fun_vals_EM[efun_idx,:])), (-1)))\n",
    "\n",
    "plt.figure()\n",
    "\n",
    "plt.scatter(X[0,:], \\\n",
    "            X[1,:], \\\n",
    "            c = np.real(e_fun_vals_EM[efun_idx,:]), \\\n",
    "            cmap = cm.seismic, \\\n",
    "            vmin = vmean-vscl, \\\n",
    "            vmax = vmean+vscl, \\\n",
    "            s=10)\n",
    "\n",
    "plt.title('phi_{:d}, lam_{:d} = {:.5f} + {:.5f}i'.format(ix, ix, lam_EM[efun_idx].real, lam_EM[efun_idx].imag))\n",
    "\n",
    "# plt.savefig('figures/Duffing/EM_dim{:d}_efun{:d}_real.png'.format(z_dim, ix), dpi=None, facecolor='w', edgecolor='w', \\\n",
    "#             transparent=False, bbox_inches=None, pad_inches=0.0)\n",
    "# plt.savefig('figures/Duffing/EM_dim{:d}_efun{:d}_real.eps'.format(z_dim, ix), dpi=None, facecolor='w', edgecolor='w', \\\n",
    "#             transparent=False, bbox_inches=None, pad_inches=0.0)\n",
    "\n",
    "plt.show()\n",
    "\n",
    "\n",
    "# color scaling\n",
    "vscl = np.max(np.reshape(np.absolute(e_fun_vals_EM[efun_idx,:]), (-1)))\n",
    "\n",
    "plt.figure()\n",
    "\n",
    "plt.scatter(X[0,:], \\\n",
    "            X[1,:], \\\n",
    "            c = np.absolute(e_fun_vals_EM[efun_idx,:]), \\\n",
    "            cmap = cm.inferno, \\\n",
    "            vmin = 0, \\\n",
    "            vmax = vscl, \\\n",
    "            s=10)\n",
    "\n",
    "plt.title('|phi_{:d}|, lam_{:d} = {:.5f} + {:.5f}i'.format(ix, ix, lam_EM[efun_idx].real, lam_EM[efun_idx].imag))\n",
    "\n",
    "# plt.savefig('figures/Duffing/EM_dim{:d}_efun{:d}_mag.png'.format(z_dim, ix), dpi=None, facecolor='w', edgecolor='w', \\\n",
    "#             transparent=False, bbox_inches=None, pad_inches=0.0)\n",
    "# plt.savefig('figures/Duffing/EM_dim{:d}_efun{:d}_mag.eps'.format(z_dim, ix), dpi=None, facecolor='w', edgecolor='w', \\\n",
    "#             transparent=False, bbox_inches=None, pad_inches=0.0)\n",
    "\n",
    "plt.show()\n",
    "\n",
    "# plot angle of complex eigenfunction\n",
    "plt.figure()\n",
    "\n",
    "plt.scatter(X[0,:], \\\n",
    "            X[1,:], \\\n",
    "            c = np.angle(e_fun_vals_EM[efun_idx,:]), \\\n",
    "            cmap = cm.twilight, \\\n",
    "            vmin = -np.pi, \\\n",
    "            vmax = np.pi, \\\n",
    "            s=10)\n",
    "\n",
    "plt.title('arg(phi_{:d}), lam_{:d} = {:.5f} + {:.5f}i'.format(ix, ix, lam_EM[efun_idx].real, lam_EM[efun_idx].imag))\n",
    "\n",
    "# plt.savefig('figures/Duffing/EM_dim{:d}_efun{:d}_arg.png'.format(z_dim, ix), dpi=None, facecolor='w', edgecolor='w', \\\n",
    "#             transparent=False, bbox_inches=None, pad_inches=0.0)\n",
    "# plt.savefig('figures/Duffing/EM_dim{:d}_efun{:d}_arg.eps'.format(z_dim, ix), dpi=None, facecolor='w', edgecolor='w', \\\n",
    "#             transparent=False, bbox_inches=None, pad_inches=0.0)\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3a17f11",
   "metadata": {},
   "source": [
    "# Prediction accuracy"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17a598e9",
   "metadata": {},
   "source": [
    "## Load data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1e34a07",
   "metadata": {},
   "outputs": [],
   "source": [
    "fname = 'models/Duffing/eval_data_fullobs_fine.p'\n",
    "\n",
    "with open(fname,'rb') as f:\n",
    "    data = pickle.load(f)\n",
    "\n",
    "X_eval = data['X_data']\n",
    "Y_eval = data['Y_data']\n",
    "Noise_eval = data['Noise_data']\n",
    "U_eval = data['U_data']\n",
    "t_eval = data['t_data']\n",
    "n_traj_eval = data['n_traj']\n",
    "n_t_eval = data['n_t']\n",
    "dt_eval = data['dt']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3bf94748",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "for k in range(n_traj_eval):\n",
    "    plt.plot(X_eval[k,:,0], X_eval[k,:,1], 'b.-')\n",
    "    \n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0aab3f50",
   "metadata": {},
   "source": [
    "## Split trajectories into observation and prediction segments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe5eebec",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(n_t_eval)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "517f7cb2",
   "metadata": {},
   "outputs": [],
   "source": [
    "n_t_eval_obs = 200\n",
    "n_t_eval_pred = n_t_eval - n_t_eval_obs\n",
    "\n",
    "t_eval_obs = t_eval[:n_t_eval_obs]\n",
    "t_eval_pred = t_eval[n_t_eval_obs:]\n",
    "\n",
    "X_eval_obs = X_eval[:,:n_t_eval_obs,:]\n",
    "X_eval_pred = X_eval[:,n_t_eval_obs:,:]\n",
    "\n",
    "Y_eval_obs = Y_eval[:,:n_t_eval_obs,:]\n",
    "Y_eval_pred = Y_eval[:,n_t_eval_obs:,:]\n",
    "\n",
    "Noise_eval_obs = Noise_eval[:,:n_t_eval_obs,:]\n",
    "\n",
    "U_eval_obs = U_eval[:,:n_t_eval_obs-1,:]\n",
    "U_eval_pred = U_eval[:,n_t_eval_obs-1:,:]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0084ec99",
   "metadata": {},
   "source": [
    "## Load EDMD model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b717d177",
   "metadata": {},
   "outputs": [],
   "source": [
    "fname = 'models/Duffing/EDMD_model_dim{:d}'.format(z_dim)\n",
    "EDMD_model = ActuatedKoopmanEM()\n",
    "EDMD_model, _, _ = load_EM_model(EDMD_model, fname)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a48f7b20",
   "metadata": {},
   "source": [
    "## Infer latent states on observation data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6430a52c",
   "metadata": {},
   "outputs": [],
   "source": [
    "muhats, Sighats = EM_model.infer_latent_states(Y_eval_obs+Noise_eval_obs, dt_eval*U_eval_obs)\n",
    "\n",
    "# muhats_EDMD, Sighats_EDMD = EDMD_model.infer_latent_states(Y_eval_obs+Noise_eval_obs, dt_eval*U_eval_obs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a4958e8",
   "metadata": {},
   "source": [
    "## Predict dynamics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46a7cffe",
   "metadata": {},
   "outputs": [],
   "source": [
    "## EM predictions\n",
    "Y_pred, Sig_y_pred, Z_pred, Sig_z_pred = EM_model.predict_dynamics(muhats[:,-1,:], \\\n",
    "                                                                   Sighats[:,-1,:,:], \\\n",
    "                                                                   dt_eval*U_eval_pred)\n",
    "\n",
    "## EDMD predictions\n",
    "# Y_pred_EDMD, Sig_y_pred_EDMD, Z_pred_EDMD, Sig_z_pred_EDMD = EDMD_model.predict_dynamics(muhats_EDMD[:,-1,:], \\\n",
    "#                                                                    Sighats_EDMD[:,-1,:,:], \\\n",
    "#                                                                    dt_eval*U_eval_pred)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "995fe099",
   "metadata": {},
   "outputs": [],
   "source": [
    "## EDMD predictions\n",
    "\n",
    "# initial conditions\n",
    "Phi_eval = SpecialPolyDict(np.vstack(X_eval_obs[:,-1,:]).T/2.0, degree=deg)\n",
    "\n",
    "# make predictions\n",
    "Y_pred_EDMD = np.zeros(Y_eval_pred.shape)\n",
    "for k in range(n_t_eval_pred):\n",
    "    Phi_eval = Phi_eval + dt_eval*np.dot(V_EDMD.T, Phi_eval)\n",
    "    Y_pred_EDMD[:,k,:] = np.transpose(np.dot(C_EDMD, Phi_eval))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dba9d4d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "traj_idx = np.random.randint(n_traj_eval)\n",
    "\n",
    "plt.figure()\n",
    "\n",
    "plt.plot(t_eval_pred, Y_eval_pred[traj_idx,:,0], 'k-')\n",
    "\n",
    "plt.plot(t_eval_pred, Y_pred[traj_idx,:,0], 'b-')\n",
    "plt.plot(t_eval_pred, Y_pred[traj_idx,:,0]+2*Sig_y_pred[traj_idx,:,0,0], 'b--')\n",
    "plt.plot(t_eval_pred, Y_pred[traj_idx,:,0]-2*Sig_y_pred[traj_idx,:,0,0], 'b--')\n",
    "\n",
    "plt.plot(t_eval_pred, Y_pred_EDMD[traj_idx,:,0], 'g-')\n",
    "# plt.plot(t_eval_pred, Y_pred_EDMD[traj_idx,:,0]+2*Sig_y_pred_EDMD[traj_idx,:,0,0], 'g--')\n",
    "# plt.plot(t_eval_pred, Y_pred_EDMD[traj_idx,:,0]-2*Sig_y_pred_EDMD[traj_idx,:,0,0], 'g--')\n",
    "\n",
    "plt.ylim([-2,2])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "459e9c50",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(2)\n",
    "\n",
    "# error in x\n",
    "axs[0].plot(t_eval_pred, np.mean(np.square(Y_pred_EDMD[:,:,0] - Y_eval_pred[:,:,0]), axis=0), 'g-', label='EDMD')\n",
    "axs[0].plot(t_eval_pred, np.mean(np.square(Y_pred[:,:,0] - Y_eval_pred[:,:,0]), axis=0), 'b-', label='EM model')\n",
    "axs[0].legend(loc='lower right')\n",
    "\n",
    "# error in dx/dt\n",
    "axs[1].plot(t_eval_pred, np.mean(np.square(Y_pred_EDMD[:,:,1] - Y_eval_pred[:,:,1]), axis=0), 'g-', label='EDMD')\n",
    "axs[1].plot(t_eval_pred, np.mean(np.square(Y_pred[:,:,1] - Y_eval_pred[:,:,1]), axis=0), 'b-', label='EM model')\n",
    "\n",
    "fig.tight_layout()\n",
    "\n",
    "# plt.savefig('figures/Duffing/dim{:d}_pred_mse.png'.format(z_dim), dpi=None, facecolor='w', edgecolor='w', \\\n",
    "#             transparent=False, bbox_inches=None, pad_inches=0.0)\n",
    "# plt.savefig('figures/Duffing/dim{:d}_pred_mse.eps'.format(z_dim), dpi=None, facecolor='w', edgecolor='w', \\\n",
    "#             transparent=False, bbox_inches=None, pad_inches=0.0)\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d72d8b89",
   "metadata": {},
   "outputs": [],
   "source": [
    "pred_errors_EM = np.square(np.linalg.norm(Y_pred - Y_eval_pred, axis=2))\n",
    "pred_errors_EDMD = np.square(np.linalg.norm(Y_pred_EDMD - Y_eval_pred, axis=2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11bc1384",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "\n",
    "plt.plot(t_eval_pred, np.mean(pred_errors_EM, axis=0), 'b-', label='EM model')\n",
    "plt.plot(t_eval_pred, np.mean(pred_errors_EDMD, axis=0), 'g-', label='EDMD')\n",
    "plt.legend(loc='lower right')\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f4b7ebd",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}