{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 8 Modeling a Drone Swinging in a Trifilar Pendulum\n",
    "\n",
    "A trifilar pendulum is a common tool for determining the inertia of a rigid body. In the video below a small quadcopter drone is hung from a trifilar pendulum and set into an oscillation about the vertical axis. The frequency (or period) of this oscillation correlates with the inertia of the drone about the axis of rotation. In this notebook you will learn how to create a mathematical model of this system."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import YouTubeVideo\n",
    "YouTubeVideo('HrG7xhrLbWo')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import sympy as sm\n",
    "sm.init_printing()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib notebook"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# System parameters\n",
    "\n",
    "We introduced a number of variable to describe the geometry in the problem:\n",
    "\n",
    "- $m$: total mass of the drone\n",
    "- $l$: length of the very stiff strings\n",
    "- $I$: rotational inertia of the drone about a vertical axis and with respect to the center of mass\n",
    "- $g$: acceleration due to gravity\n",
    "- $r$: the radial distance from the center of mass to the location of the string attachments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "m, I, g, l, r = sm.symbols('m, I, g, l, r', real=True, positive=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The goal is to write the equations of motion in terms of the generalized coordinate. In our case are going to use the rotational angle of the drone, $\\theta$. This variable should be a function of time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t = sm.symbols('t')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "theta = sm.Function('theta')(t)\n",
    "theta"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With $\\theta$ being a function of time the first and second derivatives can be easily taken with `.diff()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "theta.diff(t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "theta.diff(t, 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Getting the geometry and velocity right\n",
    "\n",
    "The height of the drone from its equilibrium resting position is written in terms of $\\theta$ below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "h = l - sm.sqrt(l**2 - (2 * r * sm.sin(theta / 2))**2)\n",
    "h"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The time derivative of $h$ is also needed to compute the linear kinetic energy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "v = h.diff(t)\n",
    "v"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This can be simplified a bit:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "v = sm.trigsimp(v)\n",
    "v"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Forming the Lagrangian\n",
    "\n",
    "Now the total kinetic energy can be formed by summing the linear and rotational kinetic energies:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "T = m * v**2 / 2 + I * theta.diff()**2 / 2\n",
    "T"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The total potential energy is:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "U = m * g * h\n",
    "U"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The Lagrangian:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "L = T - U\n",
    "L"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Forming and evaluating Lagrange's Equation\n",
    "\n",
    "Now Lagrange's equation of the second kind can be formed with respect to the single generalized coordinate.\n",
    "\n",
    "$$\\frac{d}{dt}\\left(\\frac{\\partial L}{\\partial \\dot{\\theta}}\\right) - \\frac{\\partial L}{\\partial \\theta} = 0$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "zero = L.diff(theta.diff(t)).diff(t) - L.diff(theta)\n",
    "zero"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Some trigonometric simplification can clean it up a bit:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "zero = sm.trigsimp(zero)\n",
    "zero"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Putting the equations of motion into first order form\n",
    "\n",
    "The equation of motion is linear with respect to $\\ddot{\\theta}$ and the coefficient can be extracted with:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "acc_coeff = zero.coeff(theta.diff(t, 2))\n",
    "acc_coeff"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This system can be put into first order form by introducing a new variable $\\omega = \\dot{\\theta}$ (this relationship is also the first of the two ODEs). The second is then:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "omega = sm.Function('omega')(t)\n",
    "omega_dot = sm.simplify(-(zero.subs({theta.diff(): omega}) - acc_coeff * omega.diff()) / acc_coeff)\n",
    "sm.Eq(omega.diff(), omega_dot)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Simulating the nonlinear equations of motion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from resonance.nonlinear_systems import SingleDoFNonLinearSystem"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nl_sys = SingleDoFNonLinearSystem()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nl_sys.constants['m'] = 1  # kg\n",
    "nl_sys.constants['r'] = 0.3  # m\n",
    "nl_sys.constants['l'] = 0.75  # m\n",
    "nl_sys.constants['g'] = 9.81  # m/s**2\n",
    "nl_sys.constants['I'] = 0.3**2  # kg m**2\n",
    "\n",
    "nl_sys.coordinates['theta'] = np.deg2rad(15) # rad\n",
    "nl_sys.speeds['omega'] = 0.0  # rad/s"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The 2nd order differential equations for a non-linear system have to be specificied in explicit first order form for most numerical integration routines (which is what is used behind the scenes in resonance). This means we need to introduce a new variable, a generalized speed, that defines the first of the two first order equations. We also have to move everything to the right hand side of the equations except for terms with derivatives.:\n",
    "\n",
    "$$\n",
    "\\dot{\\theta} = \\omega \\\\\n",
    "\\dot{\\omega} = f(\\omega, \\theta, m, r, l, g, I, t)\n",
    "$$\n",
    "\n",
    "You have to create a function that computes the time derivatives of the coordinates and speeds (the states) given the current values of the coordinates and speeds. The function will have arguments that match your variable names defined above and it **must output the time derivatives that correspond to the order of `sys.states`.** See below how to translate this mathematical expression into a \"right hand side\" function."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Option 1) Copy code form SymPy output and modify by hand:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "str(omega_dot)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def eval_rhs(theta, omega, I, m, r, l, g):\n",
    "    theta_dot = omega\n",
    "    omega_dot = (-m*r**2*(2*g*(l**2 + 2*r**2*np.cos(theta) -\n",
    "                               2*r**2)**4*np.sin(theta) +\n",
    "                 2*r**4*(l**2 + 2*r**2*np.cos(theta) -\n",
    "                         2*r**2)**(5/2)*omega**2*np.sin(theta)**3 +\n",
    "                 r**2*(l**2 + 2*r**2*np.cos(theta) -\n",
    "                       2*r**2)**(7/2)*omega**2*np.sin(2*theta))/\n",
    "                 (2*(I*(l**2 + 2*r**2*np.cos(theta) - 2*r**2) +\n",
    "                     m*r**4*np.sin(theta)**2)*(l**2 + 2*r**2*np.cos(theta) -\n",
    "                                               2*r**2)**(7/2)))\n",
    "    return theta_dot, omega_dot\n",
    "\n",
    "nl_sys.diff_eq_func = eval_rhs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Try out the function with some arbitrary values:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_rhs(0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Option 2) Use lambdify to create the function with no manual copying."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_omega_dot = sm.lambdify((theta, omega, I, m, r, l, g), omega_dot)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can see that the function works the same as the above function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_omega_dot(0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def eval_rhs(theta, omega, I, m, r, l, g):\n",
    "    theta_dot = omega\n",
    "    omega_dot = eval_omega_dot(theta, omega, I, m, r, l, g)\n",
    "    return theta_dot, omega_dot\n",
    "\n",
    "nl_sys.diff_eq_func = eval_rhs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_rhs(0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nl_traj = nl_sys.free_response(20.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nl_traj.plot(subplots=True);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Linearizing the Equations of Motion"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Assume $\\theta$ is small, therefore:\n",
    "\n",
    "- $\\sin\\theta = \\theta$\n",
    "- $\\cos(\\theta) = 1$\n",
    "- $\\theta \\cdot \\theta = 0$\n",
    "- $\\dot{\\theta} \\cdot \\dot{\\theta} = 0$\n",
    "- $\\sin^2\\theta = \\theta^2 = 0$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "zero = zero.subs({theta.diff(): omega})\n",
    "zero_lin = zero.subs({sm.sin(theta): theta,\n",
    "                      sm.cos(theta): 1,\n",
    "                      sm.sin(theta / 2): theta / 2,\n",
    "                      sm.cos(theta / 2): 1,\n",
    "                      theta**2: 0,\n",
    "                      omega**2: 0,\n",
    "                     }).subs(theta**2, 0)  # have to get this last one because order of subs matters\n",
    "sm.Eq(zero_lin, 0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Simulating the linear equations of motion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from resonance.linear_systems import SingleDoFLinearSystem"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sys = SingleDoFLinearSystem()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sys.constants['m'] = 1  # kg\n",
    "sys.constants['r'] = 0.3  # m\n",
    "sys.constants['l'] = 0.7  # m\n",
    "sys.constants['g'] = 9.81  # m/s**2\n",
    "sys.constants['I'] = 0.3**2  # kg m**2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sys.coordinates['theta'] = np.deg2rad(15)  # rad\n",
    "sys.speeds['omega'] = 0.0  # rad/s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def coeffs(m, r, l, g, I):\n",
    "    k = g * m * r**2 / l\n",
    "    return I, 0, k"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sys.canonical_coeffs_func = coeffs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "traj = sys.free_response(5.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "traj.plot(subplots=True);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Compare the nonlinear and linear motion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "initial_angle = 35  # degrees\n",
    "\n",
    "nl_sys.coordinates['theta'] = np.deg2rad(initial_angle)\n",
    "sys.coordinates['theta'] = np.deg2rad(initial_angle)\n",
    "\n",
    "nl_traj = nl_sys.free_response(10)\n",
    "traj = sys.free_response(10)\n",
    "\n",
    "axes = nl_traj.plot(subplots=True, color='blue')\n",
    "traj.plot(subplots=True, ax=axes, color='red');"
   ]
  }
 ],
 "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.6.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}