{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c5a1858a",
   "metadata": {},
   "source": [
    "# PVTOL full controller stack\n",
    "\n",
    "RMM, 3 Mar 2023\n",
    "\n",
    "Consider the PVTOL system (\\python{pvtol_noisy}) used throughout the course:\n",
    "\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} \\qquad\n",
    "  \\vec Y = \n",
    "    \\begin{bmatrix} x \\\\ y \\\\ \\theta \\end{bmatrix} + \n",
    "    \\begin{bmatrix} N_x \\\\ N_y \\\\ N_z \\end{bmatrix}.\n",
    "$$\n",
    "\n",
    "Assume that the input disturbances are modeled by independent, first\n",
    "order Markov (Ornstein-Uhlenbeck) processes with\n",
    "$Q_D = \\text{diag}(0.01, 0.01)$ and $\\omega_0 = 1$ and that the noise\n",
    "is modeled as white noise with covariance matrix\n",
    "\n",
    "$$\n",
    "  Q_N = \\begin{bmatrix}\n",
    "    2 \\times 10^{-4} & 0 & 1 \\times 10^{-5} \\\\\n",
    "    0 & 2 \\times 10^{-4} & 1 \\times 10^{-5} \\\\\n",
    "    1 \\times 10^{-5} & 1 \\times 10^{-5} & 1 \\times 10^{-4}\n",
    "  \\end{bmatrix}.\n",
    "$$\n",
    "\n",
    "Design a controller consisting of a trajectory generation module, a\n",
    "gain-scheduled, trajectory tracking module, and a state estimation\n",
    "module the moves the system from the origin to the equilibrum point\n",
    "point $x_\\text{f}$, $y_\\text{f}$ = 10, 0 while satisfying the\n",
    "constraint $0.5 \\sin(\\pi x / 10) - 0.1 \\leq y \\leq 1$.\n",
    "\n",
    "(A template that can be used to generate the plots required for this problem is availalble on the course website.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1be7545a",
   "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",
    "import control.optimal as opt\n",
    "import control.flatsys as fs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c32ec3f8",
   "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, pvtol_noisy, plot_results\n",
    "print(pvtol)\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])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "081764e0",
   "metadata": {},
   "source": [
    "## Estimator\n",
    "\n",
    "We start by designing an estimator for the systme (see W7 Python lecture for some starting code that you might use)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "778fb908",
   "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": "code",
   "execution_count": null,
   "id": "ee2dd9d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add code for your solution here.  Your code should create an I/O system `estimator`\n",
    "# that takes as input the measured output Y and the applied input u."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "922631c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(estimator)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46d8463d",
   "metadata": {},
   "source": [
    "## Gain scheduled controller\n",
    "\n",
    "We next design our (gain scheduled) controller for the system.  We did this in week 4 of the class.  (Note: be careful with some of the modifications that we made to the controller for the purposes of illustration in the lecture.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1fd734f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add code for your solution here.  Your code should create an I/O system `gs_ctrl`\n",
    "# that takes as input the desired trajectory xd, ud and the *estimated* state and \n",
    "# produces the input forces to the PVTOL system."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0857171",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(gs_ctrl)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ecd28a73",
   "metadata": {},
   "source": [
    "## Trajectory generation\n",
    "\n",
    "Finally, we need to design the trajectory that we want to follow.  The specifications that you are given are very similar to those that were used in HW #4, where you designed an LQR controller for the PVTOL system.  \n",
    "\n",
    "The code below defines the initial conditions, final conditions, and constraints."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5eb12bfa",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the initial and final conditions\n",
    "x_delta = np.array([10, 0, 0, 0, 0, 0])\n",
    "x0, u0 = ct.find_eqpt(pvtol, np.zeros(6), np.zeros(2), np.zeros(6), iy=[0, 1])\n",
    "xf, uf = ct.find_eqpt(pvtol, x0 + x_delta, u0, x0 + x_delta, iy=[0, 1])\n",
    "\n",
    "# Define the time horizon for the maneuver\n",
    "Tf = 5\n",
    "timepts = np.linspace(0, Tf, 100, endpoint=False)\n",
    "\n",
    "# Create a constraint corresponding to the obstacle\n",
    "from scipy.optimize import NonlinearConstraint\n",
    "from math import sin, pi\n",
    "ceiling = (NonlinearConstraint, lambda x, u: x[1], [-1], [1])\n",
    "nicolas = (NonlinearConstraint,\n",
    "           lambda x, u: x[1] - (0.5 * sin(pi * x[0] / 10) - 0.1), [0], [1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87c4bfc7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add code for your solution here.  Your code should create a trajectory xd, ud \n",
    "# that satisfies the specifications for the problem.\n",
    "\n",
    "# You may need to adjust some of the numbers above to account for system \n",
    "# uncertainties.  If you make changes, make sure to include comments on what you\n",
    "# changed (and better to make the changes below rather than editing the cell above)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e59ddc29",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Extend the trajectory to hold the final position for Tf seconds\n",
    "holdpts = np.arange(Tf, Tf + Tf, timepts[1]-timepts[0])\n",
    "xd = np.hstack([xd, np.outer(xf, np.ones_like(holdpts))])\n",
    "ud = np.hstack([ud, np.outer(uf, np.ones_like(holdpts))])\n",
    "timepts = np.hstack([timepts, holdpts])\n",
    "\n",
    "# Plot the desired trajectory\n",
    "plot_results(timepts, xd, ud)\n",
    "plt.suptitle('Desired Trajectory')\n",
    "\n",
    "# Add the constraints to the plot\n",
    "plt.subplot(2, 1, 1)\n",
    "plt.plot([0, 10], [1, 1], 'r--')\n",
    "x_nic = np.linspace(0, 10, 50)\n",
    "y_nic = 0.5 * np.sin(pi * x_nic / 10) - 0.1\n",
    "plt.plot(x_nic, y_nic, 'r--')\n",
    "plt.text(5, 0, 'Nicolas Petit', ha='center')\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "affe55fa",
   "metadata": {},
   "source": [
    "## Final Control System Implementation\n",
    "\n",
    "We now put together the final control system and simulate it.  If you have named your inputs and outputs to each of the subsystems properly, the code below should connect everything up correctly.  If you get errors about inputs or outputs that are not connected to anything, check the names of your inputs and outputs in the various\n",
    "systems above and make sure everything lines up as it should."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50dff557",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the interconnected system\n",
    "clsys = ct.interconnect(\n",
    "    [pvtol_noisy, gs_ctrl, estimator],\n",
    "    inputs=gs_clsys.input_labels + ['Dx', 'Dy'] + ['Nx', 'Ny', 'Nth'],\n",
    "    outputs=gs_clsys.output_labels\n",
    ")\n",
    "print(clsys)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f24e6f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate disturbance and noise vectors\n",
    "V = ct.white_noise(timepts, Qv)\n",
    "W = ct.white_noise(timepts, Qw)\n",
    "for i in range(V.shape[0]):\n",
    "    plt.subplot(2, 3, i+1)\n",
    "    plt.plot(timepts, V[i])\n",
    "    plt.ylabel(f'V[{i}]')\n",
    "\n",
    "for i in range(W.shape[0]):\n",
    "    plt.subplot(2, 3, i+4)\n",
    "    plt.plot(timepts, W[i])\n",
    "    plt.ylabel(f'W[{i}]')\n",
    "    plt.xlabel('Time $t$ [s]')\n",
    "\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f63091cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Simulate the opten loop system and plot the results (+ state trajectory)\n",
    "resp = ct.input_output_response(clsys, timepts, [xd, ud, V, W], x0)\n",
    "plot_results(resp.time, resp.outputs[0:6], resp.outputs[6:8])\n",
    "\n",
    "# Add the constraints to the plot\n",
    "plt.subplot(2, 1, 1)\n",
    "plt.plot([0, 10], [1, 1], 'r--')\n",
    "x_nic = np.linspace(0, 10, 50)\n",
    "y_nic = 0.5 * np.sin(pi * x_nic / 10) - 0.1\n",
    "plt.plot(x_nic, y_nic, 'r--')\n",
    "plt.text(5, 0, 'Nicolas Petit', ha='center')\n",
    "plt.suptitle(\"Measured Trajectory\")\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89221230",
   "metadata": {},
   "source": [
    "Comment on your design here.  If you were Nicolas Petit, would you be OK with the controller that you designed?  If not, try to modify the controller to make Nicolas less nervous and explain what you modified and why."
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}