{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "524d2c27",
   "metadata": {},
   "source": [
    "# Inverted double pendulum nonlinear estimation example\n",
    "\n",
    "Ioannis Mandralis, 26 Feb 2023\n",
    "\n",
    "In this example we work through estimation of the state of the inverted double pendulum system with a sensor capable of measuring the full state. First we look at estimating the noisy state subject to process and measurement noise using a perfect model. Then we assume we have a slightly perturbed model and try to estimate the state again (in practical applications such a model might have been obtained using machine learning for example).\n",
    "\n",
    "All calculations are done in discrete time, first using the Extended Kalman Filter in Predictor-Corrector form and then comparing to the Unscented Kalman Filter."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f72ba5f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import scipy as sp\n",
    "import matplotlib.pyplot as plt\n",
    "import control as ct\n",
    "\n",
    "# Define some line styles for later use\n",
    "ebarstyle = {'elinewidth': 0.5, 'capsize': 2}\n",
    "xdstyle = {'color': 'k', 'linestyle': '--', 'linewidth': 0.5, \n",
    "           'marker': '+', 'markersize': 4}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d6e932c",
   "metadata": {},
   "source": [
    "## System description\n",
    "The state of the inverted double pendulum is $(x,\\theta_1,\\theta_2,\\dot x,\\dot\\theta_1,\\dot\\theta_2)$. The control input is the linear force applied to the cart $u$. \n",
    "\n",
    "<img src=\"https://fbswiki.org/wiki/images/c/c9/Dbpend-geometry.png\"  width=\"50%\" height=\"30%\">\n",
    "The dynamics of the inverted double pendulum are given by:\n",
    "\n",
    "$$\n",
    "    D(q)\\ddot q + C(q,\\dot q)\\dot q + G(q) = H u\n",
    "$$\n",
    "\n",
    "where $q=(x,\\theta_1,\\theta_2)$ is the generalized coordinate. This system has been transformed into a 6 state first order ODE and the matrices $D,C,G,H$ have been implemented in the associated python file. Note: in the image above $\\theta_0$ is meant to correspond to $x$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4134411e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dinvpend import dbpend_noisy, dbpend_noisy_perturbed, plot_results, plot_estimate\n",
    "from ukf import UKF\n",
    "print(dbpend_noisy)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3132cdaf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Find the equilibrium point corresponding to both links in the upright position\n",
    "xeq, ueq = np.zeros(6), np.zeros(4)\n",
    "print(f\"xeq:{xeq}\")\n",
    "print(f\"ueq:{ueq}\")\n",
    "\n",
    "linsys = dbpend_noisy.linearize(xeq,ueq)\n",
    "Q,R = np.diag([1,10,10,1,1,1]), 100\n",
    "K,_,_ = ct.lqr(linsys.A,np.expand_dims(linsys.B[:,0],axis=1),Q,R)\n",
    "\n",
    "def controller_output(t,x,z,params):\n",
    "    xd = z[:6]\n",
    "    ud = z[6:7]\n",
    "    x  = z[7:]\n",
    "    \n",
    "    return ud + K @ (xd - x)\n",
    "\n",
    "ctrl = ct.NonlinearIOSystem(\n",
    "    None, controller_output, name='ctrl',\n",
    "    inputs=['xd', 'th1d','th2d','xdotd','th1dotd','th2dotd','ud',\n",
    "            'x','th1','th2','xdot','th1dot','th2dot'], \n",
    "    outputs=['u']\n",
    ")\n",
    "\n",
    "# Create the complete control system\n",
    "clsys = ct.interconnect((dbpend_noisy, ctrl), name='clsys',\n",
    "        inputs=['xd','th1d','th2d','xdotd','th1dotd','th2dotd','ud','vx','vth1','vth2'], \n",
    "        outputs=['x','th1','th2','xdot','th1dot','th2dot','u']\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6414ef06",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Simulate the closed loop system \n",
    "\n",
    "# Start the pendulum in a \"jackknife\" position with no input and no velocity\n",
    "x0 = (0.0,np.deg2rad(15),np.deg2rad(-15),0.0,0.0,0.0)\n",
    "u0 = 0.0\n",
    "\n",
    "# Plot the step response with respect to the reference input\n",
    "Tf = 2.0\n",
    "Ts = 0.0001\n",
    "tvec = np.arange(0, Tf+Ts, Ts)\n",
    "xd = xeq\n",
    "ud = ueq[0]\n",
    "\n",
    "# Add disturbances to the state \n",
    "np.random.seed(117)\n",
    "Rv = np.diag([10, 10, 10])*Ts\n",
    "V = ct.white_noise(tvec, Rv, dt=Ts)\n",
    "\n",
    "U = np.vstack([np.outer(xd,np.ones_like(tvec)), np.outer(ud,np.ones_like(tvec)), V])\n",
    "timepts, output = ct.input_output_response(clsys, tvec, U, X0=x0)\n",
    "x,u = output[:6],output[6]\n",
    "plot_results(timepts,x,u)\n",
    "\n",
    "# let the desired trajectory be xd,ud\n",
    "xd,ud = x,u"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c142a5c",
   "metadata": {},
   "source": [
    "## Sensor Model\n",
    "First assume we can measure all states (subject to noise) and further assume that there is a disturbance entering in $\\ddot x, \\ddot\\theta_1, \\ddot\\theta_2$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a166c6e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Disturbance and noise intensities\n",
    "Rw = np.diag([0.1,0.01,0.01,0.1,0.01,0.01])\n",
    "\n",
    "# Create process disturbance and sensor noise vectors\n",
    "W = ct.white_noise(timepts, Rw, dt=Ts)\n",
    "\n",
    "# Create the noisy inputs for our estimators by corrupting the state with noise \n",
    "# and also giving the commanded inputs\n",
    "Y = xd + W\n",
    "U = np.vstack([Y, ud])\n",
    "\n",
    "# Plot the noisy output signal\n",
    "plt.subplot(3, 1, 1)\n",
    "plt.plot()\n",
    "plt.scatter(timepts, Y[0])\n",
    "plt.plot(timepts, xd[0], **xdstyle)\n",
    "plt.xlabel(\"$t$ [sec]\")\n",
    "plt.ylabel(r\"$x$ position of cart[m]\")\n",
    "plt.title(\"Noisy measurements\")\n",
    "\n",
    "plt.subplot(3, 1, 2)\n",
    "plt.plot()\n",
    "plt.scatter(timepts, Y[1])\n",
    "plt.plot(timepts, xd[1], **xdstyle)\n",
    "plt.xlabel(\"$t$ [sec]\")\n",
    "plt.ylabel(r\"$\\theta_1$ [rad]\")\n",
    "plt.title(\"Noisy measurements\")\n",
    "\n",
    "plt.subplot(3, 1, 3)\n",
    "plt.plot()\n",
    "plt.scatter(timepts, Y[2])\n",
    "plt.plot(timepts, xd[2], **xdstyle)\n",
    "plt.xlabel(\"$t$ [sec]\")\n",
    "plt.ylabel(r\"$\\theta_2$ [rad]\")\n",
    "plt.title(\"Noisy measurements\")\n",
    "\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41758afd",
   "metadata": {},
   "source": [
    "## Discrete Time System\n",
    "\n",
    "In order to implement the Extended Kalman Filter in discrete time we need to first approximate the continuous dynamics by a discrete time system. Our system is of the form:\n",
    "\n",
    "$$\n",
    "\\dot x = f(x,u)\n",
    "$$\n",
    "\n",
    "To convert to discrete time use the forward Euler method to approximate the discrete state dynamics\n",
    "\n",
    "$$\n",
    "x[k+1] = x[k] + T_s f(x[k],v[k],u[k])=: F(x[k],v[k],u[k])\n",
    "$$\n",
    "\n",
    "$F$ is defined below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f92c5deb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the discrete time update function\n",
    "def F(x, v, u):\n",
    "    f_ = dbpend_noisy.updfcn(0.0,x,np.hstack([u,v]),{\n",
    "        'L1': 0.5,               # length of link 1\n",
    "        'L2': 0.5,               # length of link 2\n",
    "        'm1': 1.0,               # mass of link 1\n",
    "        'm2': 1.0,               # mass of link 2\n",
    "        'm0': 4.0,               # mass of cart\n",
    "        'g' : 9.81               # gravitational acceleration\n",
    "    })\n",
    "    return x + Ts*f_"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d84c585f",
   "metadata": {},
   "source": [
    "## Extended Kalman Filter \n",
    "\n",
    "We now construct a discrete time Extended Kalman Filter to estimate the system state given the noisy measurements and the commanded control inputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "909345b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the disturbance input and measured output matrices\n",
    "C = np.eye(6)\n",
    "\n",
    "# Create an array to store the results\n",
    "xhat = np.zeros((dbpend_noisy.nstates, timepts.size))\n",
    "P = np.zeros((dbpend_noisy.nstates, dbpend_noisy.nstates, timepts.size))\n",
    "\n",
    "# Update the estimates at each time\n",
    "for i, t in enumerate(timepts):\n",
    "    # Prediction step\n",
    "    if i == 0:\n",
    "        # Use the initial condition\n",
    "        xkkm1 = xd[:, 0]\n",
    "        Pkkm1 = np.eye(dbpend_noisy.nstates)\n",
    "    else:\n",
    "        # Prediction step \n",
    "        linsys = dbpend_noisy.linearize(xkk,np.hstack([ud[i-1],np.array([0,0,0])]))\n",
    "        A = np.eye(dbpend_noisy.nstates) + Ts*linsys.A \n",
    "        F_ = Ts * linsys.B[:,1:]\n",
    "        xkkm1 = F(xkk,np.array([0,0,0]),np.array([ud[i-1]]))\n",
    "        Pkkm1 = A @ Pkk @ A.T + F_ @ Rv @ F_.T\n",
    "    \n",
    "    # Correction step\n",
    "    L = Pkkm1 @ C.T @ np.linalg.inv(Rw + C @ Pkkm1 @ C.T)\n",
    "    xkk = xkkm1 - L @ (C @ xkkm1 - Y[:, i])\n",
    "    Pkk = Pkkm1 - L @ C @ Pkkm1\n",
    "\n",
    "    # Save the state estimate and covariance for later plotting\n",
    "    xhat[:, i], P[:, :, i] = xkkm1, Pkkm1  # For comparison to Kalman form\n",
    "\n",
    "plt.figure(figsize=(10,6))\n",
    "plot_estimate(timepts,xhat,P,xd,'EKF','r',obsv=[1,1,1,1,1,1])\n",
    "plt.legend()\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e4292a21",
   "metadata": {},
   "source": [
    "## Unscented Kalman Filter \n",
    "\n",
    "We now construct an Unscented Kalman Filter for the system."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "994592ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the output function\n",
    "def H(x,w):\n",
    "    return C@x+w\n",
    "\n",
    "# Initialize the unscented kalman filter\n",
    "ukf = UKF(dim_x=6, dim_z=6, Q=Rv, R=Rw, kappa=0.0)\n",
    "\n",
    "# Create an array to store the results\n",
    "xhat_ukf = np.zeros((dbpend_noisy.nstates, timepts.size))\n",
    "P_ukf = np.zeros((dbpend_noisy.nstates, dbpend_noisy.nstates, timepts.size))\n",
    "\n",
    "# Update the estimates at each time\n",
    "for i, t in enumerate(timepts):\n",
    "    # Prediction step\n",
    "    if i == 0:\n",
    "        # Use the initial condition\n",
    "        xkkm1 = xd[:,0]\n",
    "        Pkkm1 = np.eye(dbpend_noisy.nstates)\n",
    "    else:\n",
    "        # Prediction step\n",
    "        xkkm1, Pkkm1, _ = ukf.predict(F, xkk, Pkk, np.array([ud[i-1]]))\n",
    "\n",
    "    # Correction step \n",
    "    xkk, Pkk, _ = ukf.correct(H, xkkm1, Pkkm1, Y[:, i])\n",
    "\n",
    "    # Save the state estimate and covariance for later plotting\n",
    "    xhat_ukf[:, i], P_ukf[:, :, i] = xkkm1, Pkkm1\n",
    "    \n",
    "plt.figure(figsize=(10,6))\n",
    "plot_estimate(timepts,xhat_ukf,P_ukf,xd,'UKF','b',obsv=[1,1,1,1,1,1])\n",
    "# plot_estimate(timepts,xhat,P,xd,'EKF','r',obsv=[1,1,1,1,1,1])\n",
    "plt.legend()\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3f61ac7",
   "metadata": {},
   "source": [
    "## Perturbed model\n",
    "Now we try to estimate the system with a model which is slightly wrong"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a556a44",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dinvpend import dbpend_noisy_perturbed\n",
    "\n",
    "# Define the discrete time update function\n",
    "def F(x, v, u):\n",
    "    f_ = dbpend_noisy_perturbed.updfcn(0.0,x,np.hstack([u,v]),{\n",
    "        'L1': 0.5,               # length of link 1\n",
    "        'L2': 0.5,               # length of link 2\n",
    "        'm1': 1.0,               # mass of link 1\n",
    "        'm2': 1.0,               # mass of link 2\n",
    "        'm0': 4.0,               # mass of cart\n",
    "        'g' : 9.81               # gravitational acceleration\n",
    "    })\n",
    "    return x + Ts*f_"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a26a8b70",
   "metadata": {},
   "source": [
    "## Extended Kalman Filter with perturbed model\n",
    "\n",
    "We now construct a discrete time Extended Kalman Filter to estimate the system state given the noisy measurements and the commanded control inputs using the perturbed model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "81d46344",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the disturbance input and measured output matrices\n",
    "C = np.eye(6)\n",
    "\n",
    "# Create an array to store the results\n",
    "xhat = np.zeros((dbpend_noisy_perturbed.nstates, timepts.size))\n",
    "P = np.zeros((dbpend_noisy_perturbed.nstates, dbpend_noisy_perturbed.nstates, timepts.size))\n",
    "\n",
    "# Update the estimates at each time\n",
    "for i, t in enumerate(timepts):\n",
    "    # Prediction step\n",
    "    if i == 0:\n",
    "        # Use the initial condition\n",
    "        xkkm1 = xd[:, 0]\n",
    "        Pkkm1 = np.eye(dbpend_noisy_perturbed.nstates)\n",
    "    else:\n",
    "        # Prediction step \n",
    "        linsys = dbpend_noisy_perturbed.linearize(xkk,np.hstack([ud[i-1],np.array([0,0,0])]))\n",
    "        A = np.eye(dbpend_noisy_perturbed.nstates) + Ts*linsys.A \n",
    "        F_ = Ts * linsys.B[:,1:]\n",
    "        xkkm1 = F(xkk,np.array([0,0,0]),np.array([ud[i-1]]))\n",
    "        Pkkm1 = A @ Pkk @ A.T + F_ @ Rv @ F_.T\n",
    "    \n",
    "    # Correction step\n",
    "    L = Pkkm1 @ C.T @ np.linalg.inv(Rw + C @ Pkkm1 @ C.T)\n",
    "    xkk = xkkm1 - L @ (C @ xkkm1 - Y[:, i])\n",
    "    Pkk = Pkkm1 - L @ C @ Pkkm1\n",
    "\n",
    "    # Save the state estimate and covariance for later plotting\n",
    "    xhat[:, i], P[:, :, i] = xkkm1, Pkkm1  # For comparison to Kalman form\n",
    "\n",
    "plt.figure(figsize=(10,6))\n",
    "plot_estimate(timepts,xhat,P,xd,'EKF','r',obsv=[1,1,1,1,1,1])\n",
    "plt.legend()\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dce8fd60",
   "metadata": {},
   "source": [
    "## Unscented Kalman Filter with perturbed model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05b8752f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize the unscented kalman filter\n",
    "ukf = UKF(dim_x=6, dim_z=6, Q=Rv, R=Rw, kappa=0.0)\n",
    "\n",
    "# Create an array to store the results\n",
    "xhat_ukf = np.zeros((dbpend_noisy_perturbed.nstates, timepts.size))\n",
    "P_ukf = np.zeros((dbpend_noisy_perturbed.nstates, dbpend_noisy_perturbed.nstates, timepts.size))\n",
    "\n",
    "# Update the estimates at each time\n",
    "for i, t in enumerate(timepts):\n",
    "    # Prediction step\n",
    "    if i == 0:\n",
    "        # Use the initial condition\n",
    "        xkkm1 = xd[:,0]\n",
    "        Pkkm1 = np.eye(dbpend_noisy_perturbed.nstates)\n",
    "    else:\n",
    "        # Prediction step\n",
    "        xkkm1, Pkkm1, _ = ukf.predict(F, xkk, Pkk, np.array([ud[i-1]]))\n",
    "\n",
    "    # Correction step \n",
    "    xkk, Pkk, _ = ukf.correct(H, xkkm1, Pkkm1, Y[:, i])\n",
    "\n",
    "    # Save the state estimate and covariance for later plotting\n",
    "    xhat_ukf[:, i], P_ukf[:, :, i] = xkkm1, Pkkm1\n",
    "    \n",
    "plt.figure(figsize=(10,6))\n",
    "plot_estimate(timepts,xhat_ukf,P_ukf,xd,'UKF','b',obsv=[1,1,1,1,1,1])\n",
    "# plot_estimate(timepts,xhat,P,xd,'EKF','r',obsv=[1,1,1,1,1,1])\n",
    "plt.legend()\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "282cfeca",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}