{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c017196f",
   "metadata": {},
   "source": [
    "# PVTOL LQR + EQF example\n",
    "RMM, 14 Feb 2022\n",
    "\n",
    "This notebook illustrates the implementation of an extended Kalman filter and the use of the estimated state for LQR feedback."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "544525ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.patches as patches\n",
    "import control as ct\n",
    "import ctrlutil as ct_"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "859834cf",
   "metadata": {},
   "source": [
    "## System definition\n",
    "The dynamics of the system\n",
    "with disturbances on the $x$ and $y$ variables is given by\n",
    "$$\n",
    "  \\begin{aligned}\n",
    "    m \\ddot x &= F_1 \\cos\\theta - F_2 \\sin\\theta - c \\dot x + d_x, \\\\\n",
    "    m \\ddot y &= F_1 \\sin\\theta + F_2 \\cos\\theta - c \\dot y - m g + d_y, \\\\\n",
    "    J \\ddot \\theta &= r F_1.\n",
    "  \\end{aligned}\n",
    "$$\n",
    "The measured values of the system are the position and orientation,\n",
    "with added noise $n_x$, $n_y$, and $n_\\theta$:\n",
    "$$\n",
    "  \\vec y = \\begin{bmatrix} x \\\\ y \\\\ \\theta \\end{bmatrix} + \n",
    "  \\begin{bmatrix} n_x \\\\ n_y \\\\ n_z \\end{bmatrix}.\n",
    "$$\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ffafed74",
   "metadata": {},
   "outputs": [],
   "source": [
    "# pvtol = nominal system (no disturbances or noise)\n",
    "# noisy_pvtol = pvtol w/ process disturbances and sensor noise\n",
    "from pvtol import pvtol, noisy_pvtol, plot_results\n",
    "\n",
    "# Find the equiblirum point corresponding to the origin\n",
    "xe, ue = ct.find_eqpt(\n",
    "    pvtol, np.zeros(pvtol.nstates),\n",
    "    np.zeros(pvtol.ninputs), [0, 0, 0, 0, 0, 0],\n",
    "    iu=range(2, pvtol.ninputs), iy=[0, 1])\n",
    "\n",
    "x0, u0 = ct.find_eqpt(\n",
    "    pvtol, np.zeros(pvtol.nstates),\n",
    "    np.zeros(pvtol.ninputs), np.array([2, 1, 0, 0, 0, 0]),\n",
    "    iu=range(2, pvtol.ninputs), iy=[0, 1])\n",
    "\n",
    "# Extract the linearization for use in LQR design\n",
    "pvtol_lin = pvtol.linearize(xe, ue)\n",
    "A, B = pvtol_lin.A, pvtol_lin.B\n",
    "\n",
    "print(pvtol, \"\\n\")\n",
    "print(noisy_pvtol)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e1ee7c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Disturbance and noise intensities\n",
    "Qv = np.diag([1e-2, 1e-2])\n",
    "Qw = np.array([[2e-4, 0, 1e-5], [0, 2e-4, 1e-5], [1e-5, 1e-5, 1e-4]])\n",
    "Qwinv = np.linalg.inv(Qw)\n",
    "\n",
    "# Initial state covariance\n",
    "P0 = np.eye(pvtol.nstates)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e4c52c73",
   "metadata": {},
   "source": [
    "## Control system design"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3647bf15",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the disturbance input and measured output matrices\n",
    "F = np.array([[0, 0], [0, 0], [0, 0], [1, 0], [0, 1], [0, 0]])\n",
    "C = np.eye(3, 6)\n",
    "\n",
    "# Estimator update law\n",
    "def estimator_update(t, x, u, params):\n",
    "    # Extract the states of the estimator\n",
    "    xhat = x[0:pvtol.nstates]\n",
    "    P = x[pvtol.nstates:].reshape(pvtol.nstates, pvtol.nstates)\n",
    "\n",
    "    # Extract the inputs to the estimator\n",
    "    y = u[0:3]                  # just grab the first three outputs\n",
    "    u = u[3:5]                  # get the inputs that were applied as well\n",
    "\n",
    "    # Compute the linearization at the current state\n",
    "    A = pvtol.A(xhat, u)        # A matrix depends on current state\n",
    "    # A = pvtol.A(xe, ue)       # Fixed A matrix (for testing/comparison)\n",
    "    \n",
    "    # Compute the optimal again\n",
    "    L = P @ C.T @ Qwinv\n",
    "\n",
    "    # Update the state estimate\n",
    "    xhatdot = pvtol.updfcn(t, xhat, u, params) - L @ (C @ xhat - y)\n",
    "\n",
    "    # Update the covariance\n",
    "    Pdot = A @ P + P @ A.T - P @ C.T @ Qwinv @ C @ P + F @ Qv @ F.T\n",
    "\n",
    "    # Return the derivative\n",
    "    return np.hstack([xhatdot, Pdot.reshape(-1)])\n",
    "\n",
    "estimator = ct.NonlinearIOSystem(\n",
    "    estimator_update, lambda t, x, u, params: x[0:pvtol.nstates],\n",
    "    states=pvtol.nstates + pvtol.nstates**2,\n",
    "    inputs= noisy_pvtol.state_list[0:3] \\\n",
    "        + noisy_pvtol.input_list[0:pvtol.ninputs],\n",
    "    outputs=[f'xh{i}' for i in range(pvtol.nstates)],\n",
    ")\n",
    "print(estimator)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9787db61",
   "metadata": {},
   "outputs": [],
   "source": [
    "#\n",
    "# LQR design w/ physically motivated weighting\n",
    "#\n",
    "# Shoot for 1 cm error in x, 10 cm error in y.  Try to keep the angle\n",
    "# less than 5 degrees in making the adjustments.  Penalize side forces\n",
    "# due to loss in efficiency.\n",
    "#\n",
    "\n",
    "Qx = np.diag([100, 10, (180/np.pi) / 5, 0, 0, 0])\n",
    "Qu = np.diag([10, 1])\n",
    "K, _, _ = ct.lqr(A, B, Qx, Qu)\n",
    "\n",
    "#\n",
    "# Control system construction: combine LQR w/ EKF\n",
    "#\n",
    "# Use the linearization around the origin to design the optimal gains\n",
    "# to see how they compare to the final value of P for the EKF\n",
    "#\n",
    "\n",
    "# Construct the state feedback controller with estimated state as input\n",
    "statefbk, _ = ct_.create_statefbk_iosystem(pvtol, K, estimator=estimator)\n",
    "print(statefbk, \"\\n\")\n",
    "\n",
    "# Reconstruct the control system with the noisy version of the process\n",
    "# Create a closed loop system around the controller\n",
    "clsys = ct.interconnect(\n",
    "    [noisy_pvtol, statefbk, estimator],\n",
    "    inplist = statefbk.input_list[0:pvtol.ninputs + pvtol.nstates] + \\\n",
    "        noisy_pvtol.input_list[pvtol.ninputs:],\n",
    "    inputs = statefbk.input_list[0:pvtol.ninputs + pvtol.nstates] + \\\n",
    "        noisy_pvtol.input_list[pvtol.ninputs:],\n",
    "    outlist = pvtol.output_list + statefbk.output_list + estimator.output_list,\n",
    "    outputs = pvtol.output_list + statefbk.output_list + estimator.output_list\n",
    ")\n",
    "print(clsys)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7bf558a0",
   "metadata": {},
   "source": [
    "## Simulations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2583a0e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the time vector for the simulation\n",
    "Tf = 10\n",
    "T = np.linspace(0, Tf, 1000)\n",
    "\n",
    "# Create representative process disturbance and sensor noise vectors\n",
    "np.random.seed(117)           # avoid figures changing from run to run\n",
    "V = ct_.white_noise(T, Qv) # smaller disturbances and noise then design\n",
    "W = ct_.white_noise(T, Qw)\n",
    "plt.plot(T, V[0], label=\"V[0]\")\n",
    "plt.plot(T, W[0], label=\"W[0]\")\n",
    "plt.legend();"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d944709",
   "metadata": {},
   "source": [
    "### LQR with EKF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad7a9750",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Put together the input for the system\n",
    "U = np.vstack([\n",
    "    np.outer(xe, np.ones_like(T)),      # xd\n",
    "    np.outer(ue, np.ones_like(T)),      # ud\n",
    "    V, W                                # disturbances and noise\n",
    "])\n",
    "X0 = np.hstack([x0, np.zeros(pvtol.nstates), P0.reshape(-1)])\n",
    "\n",
    "# Initial condition response\n",
    "resp = ct.input_output_response(clsys, T, U, X0)\n",
    "\n",
    "# Plot the response\n",
    "plot_results(T, resp.states, resp.outputs[pvtol.nstates:])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5f24119",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Response of the first two states, including internal estimates\n",
    "plt.figure()\n",
    "h1, = plt.plot(resp.time, resp.outputs[0], 'b-', linewidth=0.75)\n",
    "h2, = plt.plot(resp.time, resp.outputs[1], 'r-', linewidth=0.75)\n",
    "\n",
    "# Add on the internal estimator states\n",
    "xh0 = clsys.find_output('xh0')\n",
    "xh1 = clsys.find_output('xh1')\n",
    "h3, = plt.plot(resp.time, resp.outputs[xh0], 'k--')\n",
    "h4, = plt.plot(resp.time, resp.outputs[xh1], 'k--')\n",
    "\n",
    "plt.plot([0, 10], [0, 0], 'k--', linewidth=0.5)\n",
    "plt.ylabel(\"Position $x$, $y$ [m]\")\n",
    "plt.xlabel(\"Time $t$ [s]\")\n",
    "plt.legend(\n",
    "    [h1, h2, h3, h4], ['$x$', '$y$', '$\\hat{x}$', '$\\hat{y}$'], \n",
    "    loc='upper right', frameon=False, ncol=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0c0d5c99",
   "metadata": {},
   "source": [
    "### Full state feedback"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b6a1f1c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute the full state feedback solution\n",
    "lqr_ctrl, _ = ct_.create_statefbk_iosystem(pvtol, K)\n",
    "\n",
    "lqr_clsys = ct.interconnect(\n",
    "    [noisy_pvtol, lqr_ctrl],\n",
    "    inplist = lqr_ctrl.input_list[0:pvtol.ninputs + pvtol.nstates] + \\\n",
    "        noisy_pvtol.input_list[pvtol.ninputs:],\n",
    "    inputs = lqr_ctrl.input_list[0:pvtol.ninputs + pvtol.nstates] + \\\n",
    "        noisy_pvtol.input_list[pvtol.ninputs:],\n",
    "    outlist = pvtol.output_list + lqr_ctrl.output_list,\n",
    "    outputs = pvtol.output_list + lqr_ctrl.output_list\n",
    ")\n",
    "\n",
    "# Put together the input for the system\n",
    "U = np.vstack([\n",
    "    np.outer(xe, np.ones_like(T)),      # xd\n",
    "    np.outer(ue, np.ones_like(T)),      # ud\n",
    "    V, W * 0                            # disturbances and noise\n",
    "])\n",
    "\n",
    "# Run a simulation with full state feedback\n",
    "lqr_resp = ct.input_output_response(lqr_clsys, T, U, x0)\n",
    "\n",
    "# Compare the results\n",
    "plt.plot(resp.states[0], resp.states[1], 'b-', label=\"Extended KF\")\n",
    "plt.plot(lqr_resp.states[0], lqr_resp.states[1], 'r-', label=\"Full state\")\n",
    "\n",
    "plt.xlabel('$x$ [m]')\n",
    "plt.ylabel('$y$ [m]')\n",
    "plt.axis('equal')\n",
    "plt.legend(frameon=False);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc86067c",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}