CH07/CH07_SEC03_SINDY_Lorenz.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import rcParams\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "from scipy import integrate\n",
    "\n",
    "\n",
    "rcParams.update({'font.size': 18})\n",
    "plt.rcParams['figure.figsize'] = [12, 12]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Simulate the Lorenz System\n",
    "\n",
    "dt = 0.01\n",
    "T = 50\n",
    "t = np.arange(dt,T+dt,dt)\n",
    "beta = 8/3\n",
    "sigma = 10\n",
    "rho = 28\n",
    "n = 3\n",
    "\n",
    "def lorenz_deriv(x_y_z, t0, sigma=sigma, beta=beta, rho=rho):\n",
    "    x, y, z = x_y_z\n",
    "    return [sigma * (y - x), x * (rho - z) - y, x * y - beta * z]\n",
    "\n",
    "np.random.seed(123)\n",
    "x0 = (-8,8,27)\n",
    "\n",
    "x = integrate.odeint(lorenz_deriv, x0, t,rtol=10**(-12),atol=10**(-12)*np.ones_like(x0))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Compute Derivative\n",
    "dx = np.zeros_like(x)\n",
    "for j in range(len(t)):\n",
    "    dx[j,:] = lorenz_deriv(x[j,:],0,sigma,beta,rho)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [],
   "source": [
    "## SINDy Function Definitions\n",
    "\n",
    "def poolData(yin,nVars,polyorder):\n",
    "    n = yin.shape[0]\n",
    "    yout = np.zeros((n,1))\n",
    "    \n",
    "    # poly order 0\n",
    "    yout[:,0] = np.ones(n)\n",
    "    \n",
    "    # poly order 1\n",
    "    for i in range(nVars):\n",
    "        yout = np.append(yout,yin[:,i].reshape((yin.shape[0],1)),axis=1)\n",
    "    \n",
    "    # poly order 2\n",
    "    if polyorder >= 2:\n",
    "        for i in range(nVars):\n",
    "            for j in range(i,nVars):\n",
    "                yout = np.append(yout,(yin[:,i]*yin[:,j]).reshape((yin.shape[0],1)),axis=1)\n",
    "                \n",
    "    # poly order 3\n",
    "    if polyorder >= 3:\n",
    "        for i in range(nVars):\n",
    "            for j in range(i,nVars):\n",
    "                for k in range(j,nVars):\n",
    "                    yout = np.append(yout,(yin[:,i]*yin[:,j]*yin[:,k]).reshape((yin.shape[0],1)),axis=1)\n",
    "    \n",
    "    return yout\n",
    "\n",
    "def sparsifyDynamics(Theta,dXdt,lamb,n):\n",
    "    Xi = np.linalg.lstsq(Theta,dXdt,rcond=None)[0] # Initial guess: Least-squares\n",
    "    \n",
    "    for k in range(10):\n",
    "        smallinds = np.abs(Xi) < lamb # Find small coefficients\n",
    "        Xi[smallinds] = 0                          # and threshold\n",
    "        for ind in range(n):                       # n is state dimension\n",
    "            biginds = smallinds[:,ind] == 0\n",
    "            # Regress dynamics onto remaining terms to find sparse Xi\n",
    "            Xi[biginds,ind] = np.linalg.lstsq(Theta[:,biginds],dXdt[:,ind],rcond=None)[0]\n",
    "            \n",
    "    return Xi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[  0.           0.           0.        ]\n",
      " [-10.          28.           0.        ]\n",
      " [ 10.          -1.           0.        ]\n",
      " [  0.           0.          -2.66666667]\n",
      " [  0.           0.           0.        ]\n",
      " [  0.           0.           1.        ]\n",
      " [  0.          -1.           0.        ]\n",
      " [  0.           0.           0.        ]\n",
      " [  0.           0.           0.        ]\n",
      " [  0.           0.           0.        ]\n",
      " [  0.           0.           0.        ]\n",
      " [  0.           0.           0.        ]\n",
      " [  0.           0.           0.        ]\n",
      " [  0.           0.           0.        ]\n",
      " [  0.           0.           0.        ]\n",
      " [  0.           0.           0.        ]\n",
      " [  0.           0.           0.        ]\n",
      " [  0.           0.           0.        ]\n",
      " [  0.           0.           0.        ]\n",
      " [  0.           0.           0.        ]]\n"
     ]
    }
   ],
   "source": [
    "Theta = poolData(x,n,3) # Up to third order polynomials\n",
    "lamb = 0.025 # sparsification knob lambda\n",
    "Xi = sparsifyDynamics(Theta,dx,lamb,n)\n",
    "\n",
    "print(Xi)"
   ]
  },
  {
   "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}