{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Introduction"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this last notebook, we will create a full state feedback controller that will allow the human model to balance. There isn't time to give a good overview of controls, but we will use a very common optimal control method to quickly create a controller for our system."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setup"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Bring in the results from the previous notebooks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from __future__ import print_function\n",
    "from solution.visualization import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will be working with numerical matrices so we will need several functions from NumPy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from numpy import array, zeros, eye, asarray, dot, rad2deg\n",
    "from numpy.linalg import inv"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will be simulating the model again with the controller in place and will make some plots of the state trajectories as before so bring in some useful functions from matplotlib. First enable inline plotting:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from matplotlib.pyplot import plot, xlabel, ylabel, legend, rcParams"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "rcParams['figure.figsize'] = (14, 8)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Bring in some SymPy functions for simpilfication and printing:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from sympy import simplify, Matrix, matrix2numpy\n",
    "from sympy.physics.vector import init_vprinting, vlatex"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "init_vprinting(use_latex='mathjax', pretty_print=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will be displaying numerical matrices and the IPython magic `%precision` can be used to adjust how many decimal places are shown for nice compact printing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "%precision 3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Linearization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We current have symbolic equations of motion in this form:\n",
    "\n",
    "$$ \\mathbf{M}(\\mathbf{x}) \\dot{\\mathbf{x}} = \\mathbf{f}(\\mathbf{x}, \\mathbf{u}, t) $$\n",
    "\n",
    "where $\\mathbf{M}$ is the mass matrix, $\\mathbf{x}$ is the state vector made up of the generalized coordinates and speeds, and $\\mathbf{f}(\\mathbf{x}, \\mathbf{u})$ is a vector of non-linear functions of the states, $\\mathbf{x}$, and the specified inputs, $\\mathbf{u}$.\n",
    "\n",
    "This model is relatively easy to linearize and we will use a quick and dirty method to do so. In general, the non-linear equations of motion (for an unconstrained system) can be transformed into a linear form by expanding the Taylor series about small perturbations and truncating any non-linear terms. For this problem, this amounts to computing the [Jacobian](https://en.wikipedia.org/wiki/Jacobian) of the forcing vector, $\\mathbf{f}$. The linear equations take this form:\n",
    "\n",
    "$$ \\mathbf{M} \\Delta \\dot{\\mathbf{x}} = \\mathbf{f}_A \\Delta \\mathbf{x} + \\mathbf{f}_B \\Delta \\mathbf{u} $$\n",
    "\n",
    "where $\\mathbf{f}_A$ and $\\mathbf{f}_B$ are the Jacobians of $\\mathbf{f}$ with respect to the states and specified inputs, respectively:\n",
    "\n",
    "$$ \\mathbf{f}_A = \\left[ \\frac{\\delta \\mathbf{f} }{ \\delta x_1}, \\ldots, \\frac{\\delta \\mathbf{f} }{ \\delta x_n} \\right] $$\n",
    "\n",
    "$$ \\mathbf{f}_B = \\left[ \\frac{\\delta \\mathbf{f} }{ \\delta u_1}, \\ldots, \\frac{\\delta \\mathbf{f} }{ \\delta u_n} \\right] $$\n",
    "\n",
    "and $\\Delta \\mathbf{x}$ and $\\Delta \\mathbf{u}$ are _differential_ changes in $\\mathbf{x}$ and $\\mathbf{u}$ respectively.\n",
    "\n",
    "Once $\\mathbf{f}_A$ and $\\mathbf{f}_B$ are found, then the standard state space form of the linear equations can be formed by transforming the linear equations.\n",
    "\n",
    "$$ \\Delta \\dot{x} = \\mathbf{M}^{-1} \\mathbf{f}_A \\Delta \\mathbf{x} + \\mathbf{M}^{-1} \\mathbf{f}_B \\Delta \\mathbf{u} $$\n",
    "\n",
    "which can be rewritten as\n",
    "\n",
    "$$ \\Delta \\dot{x} = \\mathbf{A} \\Delta \\mathbf{x} + \\mathbf{B} \\Delta \\mathbf{u} $$\n",
    "   \n",
    "where $\\mathbf{A}$ is the state matrix and $\\mathbf{B}$ is the input matrix."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The first step is to specify the equilbrium point about which to linearize. Our system has two unique equilibrium points, one when the person is upright, $\\mathbf{x} = [0, \\ldots, 0]$, and one when the person is upside down, $\\mathbf{x} = [\\pi, \\ldots, \\pi]$ (180 degrees). We are only concerned with the former because we want the model to \"stand up\"."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "equilibrium_point = zeros(len(coordinates + speeds))\n",
    "equilibrium_point"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will use SymPy's `sub()` method to convert the linear equations of motion to numerical arrays, so we create two dictionaries containing the numerical values of the equilibrium point and the constants."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "equilibrium_dict = dict(zip(coordinates + speeds, equilibrium_point))\n",
    "equilibrium_dict"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As an exercise, create a `parameter_dict` that maps the symbolic constants to the numerical constants. We have variable names for each already available from the previous notebooks (`constants` and `numerical_constants`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "%load exercise_solutions/n09_control_parameter-dict.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `KanesMethod` object from earlier contains methods for linearizing the equations of motion. First, we'll create a `Linearizer` object:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "linearizer = kane.to_linearizer()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This creates an object that holds a *generalized form* of the equations of motion, which can be used to linearize the system in a robust way. When it does this, it finds the inputs `r`, and orders them alphabetically. We have our own order we want to use from before, so we need to change this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "linearizer.r == Matrix(specified)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "linearizer.r = Matrix(specified)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To perform the actual linearization, we'll use the `linearize` method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "help(linearizer.linearize)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Because we want to get the equations in the form\n",
    "\n",
    "$$ \\Delta \\dot{x} = \\mathbf{A} \\Delta \\mathbf{x} + \\mathbf{B} \\Delta \\mathbf{u} $$\n",
    "\n",
    "we'll set `A_and_B=True`. We can also perform the variable to numeric substitution in this step by passing in both the `equilibrium_dict` and the `parameter_dict` to the operating point."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "A, B = linearizer.linearize(op_point=[equilibrium_dict, parameter_dict], A_and_B=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can convert from symbolic to numeric numpy matrices, and display our results:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "A = matrix2numpy(A, dtype=float)\n",
    "B = matrix2numpy(B, dtype=float)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "A"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "B"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Controller Design"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before we try to create a controller, we need to check to see if it is actually [controllable](https://en.wikipedia.org/wiki/Controlability). You can import a function from the `utils` module included with these notebooks to check."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from utils import controllable"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "controllable(A, B)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So yes, the system is controllable. Now we will use [optimal control theory](https://en.wikipedia.org/wiki/Optimal_control) to create a controller that compute the necessary joint torques at each time step to balance the person.\n",
    "\n",
    "It turns out that if you setup a cost function in which you'd like to minimize the deviation from the desired state, 0, and minimize the effort in the joint torques, there is an elegant solution to find a matrix $\\mathbf{K}$ which can be used to compute the joint torques given the current state value.\n",
    "\n",
    "Given the cost function:\n",
    "\n",
    "$$J = \\frac{1}{2} \\int_0^\\infty\\left[\\mathbf{x}^T \\mathbf{Q} \\mathbf{x} + \\mathbf{u}^T \\mathbf{R} \\mathbf{u}\\right]dt$$\n",
    "\n",
    "There is a matrix $\\mathbf{K}$ that computes the inputs $\\mathbf{u}(t)$ given the states $\\mathbf{x}$:\n",
    "\n",
    "$$\\mathbf{u}(t) = -\\mathbf{K} \\mathbf{x}(t)$$\n",
    "\n",
    "where $\\mathbf{K}$ can be found by the solution to the algebraic [Riccati Equation](https://en.wikipedia.org/wiki/Riccati_equation) such that:\n",
    "\n",
    "$$\\mathbf{K} = \\mathbf{R}^{-1} \\mathbf{B}^T \\mathbf{S}$$\n",
    "\n",
    "and the Riccati equation is:\n",
    "\n",
    "$$\\textbf{0} = -\\textbf{S}\\textbf{A}-\\textbf{A}^{\\text{T}}\\textbf{S}+\\textbf{S}\\textbf{B}\\textbf{R}^{-1}\\textbf{B}^{\\text{T}}\\textbf{S}-\\textbf{Q}$$\n",
    "\n",
    "SciPy provides a function that can compute the solution, $\\mathbf{S}$, of the continous algebraic Riccati Equation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from scipy.linalg import solve_continuous_are"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The weighting matrices, $\\mathbf{Q}$ and $\\mathbf{R}$, allow you to control how much weight is applied to minimizing the error in the states and the effort in the inputs. A typical starting point for these matrices is to set them equal to the identity matrix, which can be created with the NumPy `eye()` function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "Q = eye(6)\n",
    "Q"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "R = eye(3)\n",
    "R"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$\\mathbf{K}$ can now be computed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "S = solve_continuous_are(A, B, Q, R)\n",
    "K = dot(dot(inv(R), B.T),  S)\n",
    "K"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Remember that the `right_hand_side` function we generated earlier can take a function for the specified values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "help(right_hand_side)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a function named `controller` that computes the specified inputs given the current state and time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def controller():\n",
    "    \"\"\"Returns the output of the controller, i.e. the joint torques, given the current state.\"\"\"\n",
    "    return"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "%load exercise_solutions/n09_control_controller-function.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can integrate the equations of motion again, but this time with the controller in place."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Integrate the equations of motion and plot the coordinates and speeds as we did in the simulation notebook and see if the system behaves as expected."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "%load exercise_solutions/n09_control_integrate-eom.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Plot Trajectories"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Use the same plot functions as were used in the simulation notebook to inspect the state trajectories."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Plot the angles\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "%load exercise_solutions/n09_control_plot-angles.py\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Plot the angular velocities\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "%load exercise_solutions/n09_control_plot-velocities.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What do these plots tell you about the controller? Does it work?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Visualization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, regenerate the visualization to see the 3D animation of the system under the influence of the controller."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "scene.states_trajectories = y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "scene.display_ipython()"
   ]
  }
 ],
 "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.5.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}