{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8bf0c3d7",
   "metadata": {},
   "source": [
    "# Double pendulum\n",
    "\n",
    "This example uses the `sequence()` and `subdraw()` animations to animate `plot` plot objects.\n",
    "\n",
    "This example is a Diplotocus implementation of [this animation from the matplotlib docs](https://matplotlib.org/stable/gallery/animation/double_pendulum.html).\n",
    "\n",
    "We first define a function `get_pendulum_positions()` that returns the double pendulum nodes positions for a range of times."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "ef3943c0",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "def get_pendulum_positions(G,L1,L2,M1,M2,t_stop):\n",
    "    def derivs(state):\n",
    "        dydx = np.zeros_like(state)\n",
    "\n",
    "        dydx[0] = state[1]\n",
    "\n",
    "        delta = state[2] - state[0]\n",
    "        den1 = (M1+M2) * L1 - M2 * L1 * np.cos(delta) * np.cos(delta)\n",
    "        dydx[1] = ((M2 * L1 * state[1] * state[1] * np.sin(delta) * np.cos(delta)\n",
    "                    + M2 * G * np.sin(state[2]) * np.cos(delta)\n",
    "                    + M2 * L2 * state[3] * state[3] * np.sin(delta)\n",
    "                    - (M1+M2) * G * np.sin(state[0]))\n",
    "                / den1)\n",
    "\n",
    "        dydx[2] = state[3]\n",
    "\n",
    "        den2 = (L2/L1) * den1\n",
    "        dydx[3] = ((- M2 * L2 * state[3] * state[3] * np.sin(delta) * np.cos(delta)\n",
    "                    + (M1+M2) * G * np.sin(state[0]) * np.cos(delta)\n",
    "                    - (M1+M2) * L1 * state[1] * state[1] * np.sin(delta)\n",
    "                    - (M1+M2) * G * np.sin(state[2]))\n",
    "                / den2)\n",
    "\n",
    "        return dydx\n",
    "\n",
    "    # create a time array from 0..t_stop sampled at 0.01 second steps\n",
    "    dt = 0.01\n",
    "    t = np.arange(0, t_stop, dt)\n",
    "\n",
    "    # th1 and th2 are the initial angles (degrees)\n",
    "    # w10 and w20 are the initial angular velocities (degrees per second)\n",
    "    th1 = 120.0\n",
    "    w1 = 0.0\n",
    "    th2 = -10.0\n",
    "    w2 = 0.0\n",
    "\n",
    "    # initial state\n",
    "    state = np.radians([th1, w1, th2, w2])\n",
    "\n",
    "    # integrate the ODE using Euler's method\n",
    "    y = np.empty((len(t), 4))\n",
    "    y[0] = state\n",
    "    for i in range(1, len(t)):\n",
    "        y[i] = y[i - 1] + derivs(y[i - 1]) * dt\n",
    "\n",
    "    x0 = np.zeros(len(t))\n",
    "    y0 = np.zeros(len(t))\n",
    "\n",
    "    x1 = L1*np.sin(y[:, 0])\n",
    "    y1 = -L1*np.cos(y[:, 0])\n",
    "\n",
    "    x2 = L2*np.sin(y[:, 2]) + x1\n",
    "    y2 = -L2*np.cos(y[:, 2]) + y1\n",
    "\n",
    "    return x0,x1,x2,y0,y1,y2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ca684cb",
   "metadata": {},
   "source": [
    "Then we generate the positions, create a `plot()` object `p` that displays the double pendulum and another `plot()` object `trace` that displays the last 5 positions of its end node:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8ecc773",
   "metadata": {},
   "outputs": [],
   "source": [
    "import diplotocus as dpl\n",
    "\n",
    "G = 9.8  # acceleration due to gravity, in m/s^2\n",
    "L1 = 1.0  # length of pendulum 1 in m\n",
    "L2 = 1.0  # length of pendulum 2 in m\n",
    "L = L1 + L2 # maximal length of the combined pendulum\n",
    "M1 = 1.0  # mass of pendulum 1 in kg\n",
    "M2 = 1.0  # mass of pendulum 2 in kg\n",
    "t_stop = 10  # how many seconds to simulate\n",
    "\n",
    "x0,x1,x2,y0,y1,y2 = get_pendulum_positions(G,L1,L2,M1,M2,t_stop)\n",
    "\n",
    "xs = np.stack((x0,x1,x2),axis=-1)\n",
    "ys = np.stack((y0,y1,y2),axis=-1)\n",
    "\n",
    "tl = dpl.Timeline(xlim=(-L,L),ylim=(-L,1))\n",
    "\n",
    "p = dpl.plot(xs,ys,marker='o',lw=2,ls='-',zorder=2)\n",
    "p.sequence(600)\n",
    "\n",
    "trace = dpl.plot(x2,y2,lw=4,ls='-',zorder=1)\n",
    "trace.subdraw(-5,0,600)\n",
    "\n",
    "tl.animate((p,trace))\n",
    "tl.save_video('../../_static/examples/pendulum.mp4')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "af93efc6",
   "metadata": {
    "tags": [
     "remove-input"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "    <video width=\"640\" height=\"460\" autoplay loop muted playsinline>\n",
       "    <source src=\"../../_static/examples/pendulum.mp4?randId={}\" type=\"video/mp4\">\n",
       "    </video>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from IPython.display import HTML, display\n",
    "display(HTML(\"\"\"\n",
    "    <video width=\"640\" height=\"460\" autoplay loop muted playsinline>\n",
    "    <source src=\"../../_static/examples/pendulum.mp4?randId={}\" type=\"video/mp4\">\n",
    "    </video>\n",
    "\"\"\"))"
   ]
  }
 ],
 "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.13.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}