{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Double Pendulum Example\n",
    "\n",
    "A double pendulum is a dynamic system which has chaotic mostion for many initial conditions.\n",
    "\n",
    "The derivation of the derivative function for this example is taken from https://myphysicslab.com/pendulum/double-pendulum-en.html\n",
    "\n",
    "![asdf](static/Double-Pendulum.svg \"Double Pendulum\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When the system is defined as a series of first-order differential equations the state vector is\n",
    "\n",
    "$$X = \\begin{bmatrix}\n",
    "\\theta_1 \\\\ \\theta_2 \\\\ \\omega_1 \\\\ \\omega_2\n",
    "\\end{bmatrix}$$\n",
    "and the derivative of the system is\n",
    "\n",
    "$$\\dot{X} = \\begin{bmatrix}\n",
    "\\dot{\\theta_1} \\\\ \\dot{\\theta_2} \\\\ \\dot{\\omega_1} \\\\ \\dot{\\omega_2}\n",
    "\\end{bmatrix} = \\begin{bmatrix}\n",
    "\\omega_1 \\\\ \\omega_2 \\\\\n",
    "\\frac{-g (2 m_1 + m_2) \\sin{\\theta_1} - m_2 g \\sin(\\theta_1 - 2 \\theta_2) - 2 \\sin(\\theta_1 - \\theta_2) m_2 (\\omega_2^2 L_2 + \\omega_1^2 L_1 \\cos(\\theta_1 - \\theta_2))}{L_1 (2 m_1 + m_2 - m_2 \\cos(2 \\theta_1 - 2 \\theta_2))} \\\\\n",
    "\\frac{2 \\sin(\\theta_1 - \\theta_2)(\\omega_1^2 L_1 (m_1 + m_2) + g (m_1 + m_2) \\cos{\\theta_1} + \\omega_2^2 L_2 m_2 \\cos(\\theta_1 - \\theta_2))}{L_2 (2 m_1 + m_2 - m_2 \\cos(2 \\theta_1 - 2 \\theta_2))}\n",
    "\\end{bmatrix}.\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cd .."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import ode\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.animation import FuncAnimation\n",
    "from IPython.display import HTML\n",
    "\n",
    "def doublependulum(t,X):\n",
    "    th1, th2, om1, om2 = X\n",
    "    g = 9.81\n",
    "    m1 = 1\n",
    "    m2 = 1\n",
    "    l1 = 1\n",
    "    l2 = 1\n",
    "    k1 = -g * ((2 * m1) + m2) * np.sin(th1)\n",
    "    k2 = m2 * g * np.sin(th1 - (2 * th2))\n",
    "    k3 = 2 * np.sin(th1 - th2) * m2\n",
    "    k4 = ((om2**2) * l2) + ((om1**2) * l1 * np.cos(th1 - th2))\n",
    "    k5 = m2 * np.cos((2 * th1) - (2 * th2))\n",
    "    k6 = 2 * np.sin(th1 - th2)\n",
    "    k7 = ((om1**2) * l1 * (m1 + m2))\n",
    "    k8 = g * (m1 + m2) * np.cos(th1)\n",
    "    k9 = (om2**2) * l2 * m2 * np.cos(th1 - th2)\n",
    "    dX = np.array([\n",
    "        om1,\n",
    "        om2,\n",
    "        (k1 - k2 - (k3 * k4)) / (l1 * ((2 * m1) + m2 - k5)),\n",
    "        (k6 * (k7 + k8 + k9)) / (l2 * ((2 * m1) + m2 - k5))\n",
    "    ])\n",
    "    return dX"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def angles2xy(th1, th2, om1, om2):\n",
    "    l1 = 1\n",
    "    l2 = 1\n",
    "    x1 = l1 * np.sin(th1)\n",
    "    y1 = -l1 * np.cos(th1)\n",
    "    x2 = x1 + (l2 * np.sin(th2))\n",
    "    y2 = y1 - (l2 * np.cos(th2))\n",
    "    return x1, y1, x2, y2\n",
    "    \n",
    "\n",
    "et, ex = ode.euler(doublependulum, [np.pi/2,np.pi,0,0], [0,5], .025)\n",
    "th1, th2, om1, om2 = ex\n",
    "ex1, ey1, ex2, ey2 = zip(*[angles2xy(*xi) for xi in zip(th1, th2, om1, om2)])\n",
    "\n",
    "bet, bex = ode.backwardeuler(doublependulum, [np.pi/2,np.pi,0,0], [0,5], .025)\n",
    "bex1, bey1, bex2, bey2 = zip(*[angles2xy(*xi) for xi in zip(*bex)])\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "ax.plot(ex2, ey2, 'r', bex2, bey2, 'g')\n",
    "ax.set_aspect('equal', 'datalim')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#plt.rcParams['figure.dpi'] = 120\n",
    "fig, ax = plt.subplots()\n",
    "ax.axis([-2, 2, -2, 2])\n",
    "ax.set_aspect('equal')\n",
    "trace_euler, = plt.plot([], [], 'r', animated=True)\n",
    "pend_euler, = plt.plot([], [], 'ro-', animated=True)\n",
    "trace_backeuler, = plt.plot([], [], 'g', animated=True)\n",
    "pend_backeuler, = plt.plot([], [], 'go-', animated=True)\n",
    "anim_ex2 = []\n",
    "anim_ey2 = []\n",
    "anim_bex2 = []\n",
    "anim_bey2 = []\n",
    "\n",
    "def update(frame):\n",
    "    anim_ex2.append(ex2[frame])\n",
    "    anim_ey2.append(ey2[frame])\n",
    "    pend_euler_x = [0, ex1[frame], ex2[frame]]\n",
    "    pend_euler_y = [0, ey1[frame], ey2[frame]]\n",
    "    anim_bex2.append(bex2[frame])\n",
    "    anim_bey2.append(bey2[frame])\n",
    "    pend_backeuler_x = [0, bex1[frame], bex2[frame]]\n",
    "    pend_backeuler_y = [0, bey1[frame], bey2[frame]]\n",
    "    trace_euler.set_data(anim_ex2, anim_ey2)\n",
    "    pend_euler.set_data(pend_euler_x, pend_euler_y)\n",
    "    trace_backeuler.set_data(anim_bex2, anim_bey2)\n",
    "    pend_backeuler.set_data(pend_backeuler_x, pend_backeuler_y)\n",
    "    return trace_euler, pend_euler, trace_backeuler, pend_backeuler\n",
    "\n",
    "%matplotlib inline\n",
    "ani = FuncAnimation(fig, update, frames=range(len(et)), blit=True)\n",
    "\n",
    "HTML_ani = HTML(ani.to_jshtml())\n",
    "#HTML_ani = HTML(ani.to_html5_video())\n",
    "#ani.save('DoublePendulumExample.gif', writer='imagemagick', fps=60)\n",
    "HTML_ani"
   ]
  }
 ],
 "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}