double-pendulum.ipynb

{
 "metadata": {
  "name": ""
 },
 "nbformat": 3,
 "nbformat_minor": 0,
 "worksheets": [
  {
   "cells": [
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "# comment these lines if you want interactive mode,\n",
      "# i.e., if you want to see the animation in real time.\n",
      "import matplotlib\n",
      "matplotlib.use('Agg')"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 1
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "import matplotlib.pyplot as plt\n",
      "import numpy as np\n",
      "import os\n",
      "from scipy.integrate import ode"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 2
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "def equations(t, y, args):\n",
      "    \"\"\" the equations for the double pendulum \"\"\"\n",
      "    x1 = y[0] # x1 = theta1, angle\n",
      "    x2 = y[1] # x2 = theta2, angle\n",
      "    p1 = y[2] # p1 = omega1, angular velocity\n",
      "    p2 = y[3] # p2 = omega2, angular velocity\n",
      "    l1,l2,m1,m2,g = args\n",
      "    x1_eq = p1\n",
      "    x2_eq = p2\n",
      "    p1_eq = -((g*(2*m1+m2)*np.sin(x1)+m2*(g*np.sin(x1-2*x2)+2*(l2*p2**2+l1*p1**2*np.cos(x1-x2))*np.sin(x1-x2)))/(2*l1*(m1+m2-m2*(np.cos(x1-x2))**2)))\n",
      "    p2_eq = ((l1*(m1+m2)*p1**2+g*(m1+m2)*np.cos(x1)+l2*m2*p2**2*np.cos(x1-x2))*np.sin(x1-x2))/(l2*(m1+m2-m2*(np.cos(x1-x2))**2))\n",
      "    return [x1_eq, x2_eq, p1_eq, p2_eq]\n",
      "\n",
      "def calculate_trajectory(args,time,y0):\n",
      "    \"\"\" uses scipy's ode itegrator to simulate the equations \"\"\"\n",
      "    t0,t1,dt = time\n",
      "    r = ode(equations).set_integrator('dopri5')\n",
      "    r.set_initial_value(y0, t0).set_f_params(args)\n",
      "    data=[[t0, y0[0], y0[1], y0[2], y0[3] ]]\n",
      "    while r.successful() and r.t < t1:\n",
      "        r.integrate(r.t+dt)\n",
      "        data.append([r.t, r.y[0], r.y[1], r.y[2], r.y[3] ])\n",
      "    return np.array(data)\n",
      "\n",
      "def from_angle_to_xy(args,angles):\n",
      "    \"\"\" converts angles into xy positions \"\"\"\n",
      "    l1,l2,m1,m2,g = args\n",
      "    time,theta1,theta2 = angles.T\n",
      "    x1 =  l1*np.sin(theta1)\n",
      "    y1 = -l1*np.cos(theta1)\n",
      "    x2 =  l2*np.sin(theta2) + x1\n",
      "    y2 = -l2*np.cos(theta2) + y1\n",
      "    return np.array([time,x1,y1,x2,y2]).T"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 3
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "l1 = 1.0 # length of arms\n",
      "l2 = 1.0\n",
      "m1 = 1.0 # mass of the pendulum\n",
      "m2 = 1.0\n",
      "g  = 10.0 # acceleration of gravity\n",
      "args = [l1,l2,m1,m2,g]\n",
      "time = [0.0,5.0,0.05] # start, finish, dt\n",
      "ic   = [np.pi, np.pi*0.9, 0.0, 0.0] # almost upright initial condition\n",
      "\n",
      "# here the magic happens\n",
      "d = calculate_trajectory(args,time,ic)\n",
      "data_TXY = from_angle_to_xy(args,d[:,:3])"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 4
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "# let's plot stuff, and make a nice movie\n",
      "make_movie=True\n",
      "plt.ion()\n",
      "fig = plt.figure()\n",
      "ax = fig.add_subplot(111)\n",
      "\n",
      "def plot_last_seconds(data,index):\n",
      "    \"\"\" Plots a line with the trajectory of the tip of pendulum 2 (x2,y2)\n",
      "    \"\"\"\n",
      "    how_long = 1.0 # seconds\n",
      "    n = int(how_long/time[2])\n",
      "    to_plot = data[:index,:]\n",
      "    if index < n:\n",
      "        prepend =  np.tile(data[0],(n-index,1))\n",
      "        to_plot = np.vstack([prepend,to_plot])\n",
      "        index = n\n",
      "    colormap = plt.cm.Greys_r\n",
      "    colors = [colormap(i) for i in np.linspace(0.0, 1.0, n-1)]\n",
      "    plots = []\n",
      "    for j in np.arange(n-1):\n",
      "        p, = ax.plot(to_plot[index-j-1:index-j+1,3],to_plot[index-j-1:index-j+1,4],\n",
      "                color=colors[j])\n",
      "        plots.append(p)\n",
      "    return plots\n",
      "\n",
      "# \"plot\" returns a tuple of line objects, thus the comma\n",
      "t,x1,y1,x2,y2 = data_TXY[0]\n",
      "line1, = ax.plot([0.0,x1], [0.0,y1], 'r-')\n",
      "line2, = ax.plot([x1,x2], [y1,y2], 'r-')\n",
      "circ1, = ax.plot([x1], [y1], 'ro',markersize=10)\n",
      "circ2, = ax.plot([x2], [y2], 'ro',markersize=10)\n",
      "ax.axis([-2.5,2.5,-2.5,2.5])\n",
      "plt.axes().set_aspect('equal')\n",
      "\n",
      "frame_names = []\n",
      "for i,v in enumerate(data_TXY):\n",
      "    # plot_last_seconds considerably slows down the simulation,\n",
      "    # but makes it much prettier...\n",
      "    pls = plot_last_seconds(data_TXY,i)\n",
      "    t,x1,y1,x2,y2 = v\n",
      "    # print(\"t={:.2f}\".format(t)) # you might want to know how things are going...\n",
      "    line1.set_data([0.0,x1],[0.0,y1])\n",
      "    line2.set_data([x1,x2],[y1,y2])\n",
      "    circ1.set_data([x1],[y1])\n",
      "    circ2.set_data([x2],[y2])\n",
      "    ax.set_title(\"t={:.3f}\".format(t))\n",
      "    fig.canvas.draw()\n",
      "    if make_movie:\n",
      "        fname = \"_tmp{:05d}.png\".format(i)\n",
      "        frame_names.append(fname)\n",
      "        fig.savefig(fname,bbox_inches='tight')\n",
      "    for k in pls:\n",
      "        k.remove()\n",
      "\n",
      "if make_movie:\n",
      "    frames = \"_tmp%5d.png\"\n",
      "    movie_command = \"mencoder mf://*.png -mf fps=24:type=png --ovc lavc -lavcopts vcodec=mpeg4:mbd=2:trell -oac copy -o double_pend.avi\"\n",
      "    \n",
      "    # we might have other .png figures in the directory\n",
      "    # in this case, use the code below\n",
      "    f = open(\"png_list.txt\", \"w\")\n",
      "    for i in frame_names:\n",
      "        f.write(i+\"\\n\")\n",
      "    f.close()\n",
      "    movie_command = \"mencoder mf://@png_list.txt -mf fps=24:type=png -ovc lavc -lavcopts vcodec=mpeg4:mbd=2:trell -oac copy -o double_pend.avi\"\n",
      "    \n",
      "    err=os.system(movie_command)\n",
      "    if err!=0:\n",
      "        raise RuntimeError(\"Couldn't run mencoder.  Data in tmp*.png files\")\n",
      "    for fname in frame_names:\n",
      "        os.remove(fname)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 5
    }
   ],
   "metadata": {}
  }
 ]
}