{
 "metadata": {
  "name": "05_Animations"
 },
 "nbformat": 3,
 "nbformat_minor": 0,
 "worksheets": [
  {
   "cells": [
    {
     "cell_type": "heading",
     "level": 1,
     "metadata": {},
     "source": [
      "Animations in Matplotlib"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "%pylab"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "We showed previously how to create basic animations in matplotlib\n",
      "using a simple canvas update:"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "fig, ax = plt.subplots()\n",
      "x = np.linspace(0, 10, 1000)\n",
      "lines = ax.plot(x, np.sin(x), lw=2)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "for i in range(100):\n",
      "    lines[0].set_ydata(np.sin(x + 0.1 * i))\n",
      "    fig.canvas.draw()"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "The problem with this approach is that it can operate unpredictably\n",
      "between different backends, and may not even run on some!  Also, if\n",
      "you'd like to save the animation to file, you must do a lot of things\n",
      "by hand.  Thankfully, the Matplotlib team put together a very nice\n",
      "animation package in the version 1.1 release.  Creating animations is\n",
      "fairly simple with this tool, as we'll see now.\n",
      "\n",
      "Note that this framework is not without its own issues: in particular,\n",
      "animations will not work within the MacOSX backend (see, for example,\n",
      "https://github.com/matplotlib/matplotlib/issues/531)"
     ]
    },
    {
     "cell_type": "heading",
     "level": 2,
     "metadata": {},
     "source": [
      "Creating a Basic Animation"
     ]
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "Creating a basic animation is a matter of initializing the plot,\n",
      "creating functions to update the frames, and passing these functions\n",
      "to an animation object.\n",
      "\n",
      "We start by setting up a normal matplotlib figure:"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "import matplotlib.pyplot as plt\n",
      "import numpy as np\n",
      "\n",
      "# First set up the figure, the axis, and the plot element we want to animate\n",
      "fig = plt.figure()\n",
      "ax = plt.axes(xlim=(0, 2), ylim=(-2, 2))\n",
      "line, = ax.plot([], [], lw=2)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "Next we'll create an ``init()`` function and an ``animate()`` function.\n",
      "\n",
      "The purpose of the ``init()`` function is to set the background of the\n",
      "animation: it should essentially hide any plot elements that you don't\n",
      "want to be shown in every frame.\n",
      "\n",
      "The purpose of the ``animate()`` function is to update the plot elements\n",
      "for each frame."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "# initialization function: plot the background of each frame\n",
      "def init():\n",
      "    line.set_data([], [])\n",
      "    return line,"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "# animation function.  This is called sequentially\n",
      "x = np.linspace(0, 2, 1000)\n",
      "\n",
      "def animate(i):\n",
      "    y = np.sin(2 * np.pi * (x - 0.01 * i))\n",
      "    line.set_data(x, y)\n",
      "    return line,"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "Creating the animation now is a matter of passing these initialization and\n",
      "frame-step functions to the animator:"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "from matplotlib import animation\n",
      "\n",
      "# call the animator.  blit=True means only re-draw the parts that have changed.\n",
      "anim = animation.FuncAnimation(fig, animate, init_func=init,\n",
      "                               frames=200, interval=20, blit=True)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "heading",
     "level": 3,
     "metadata": {},
     "source": [
      "Saving the Animation"
     ]
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "To save the animation, you'll need one of several additional\n",
      "programs available on your system.  The animation library\n",
      "works with both `ffmpeg``, ``mencoder``, and ``avconv``,\n",
      "and the current master of matplotlib also works with \n",
      "``imagemagick`` to natively create animated gifs.\n",
      "\n",
      "For more information on these, see http://matplotlib.sourceforge.net/api/animation_api.html\n",
      "\n",
      "The save command looks something like this:"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "anim.save('basic_animation.mp4', fps=30)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "heading",
     "level": 2,
     "metadata": {},
     "source": [
      "3D Animation: The Lorenz Equations"
     ]
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "The Lorenz equations are an interesting set of nonlinear differential\n",
      "equations which pop up in a surprising number of physical systems.\n",
      "Here we'll use scipy to solve the equations, and matplotlib to animate\n",
      "the results in 3D.  The full script can be found in ``examples/lorenz_animation.py``"
     ]
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "The Lorenz Equations are as follows:\n",
      "\n",
      "$dx/dt = \\sigma(y\u2212x)$\n",
      "\n",
      "$dy/dt = x(\\rho\u2212z)\u2212y$\n",
      "\n",
      "$dz/dt = x y \u2212 \\beta z$\n",
      "\n",
      "with free parameters $\\sigma$, $\\rho$, and $\\beta$.\n",
      "We can use scipy to solve these equations relatively easily:"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "import numpy as np\n",
      "from scipy import integrate\n",
      "\n",
      "# define the lorenz derivatives\n",
      "def lorenz_deriv((x, y, z), t0, sigma=10., beta=8./3, rho=28.0):\n",
      "    \"\"\"Compute the time-derivative of a Lorenz system.\"\"\"\n",
      "    return [sigma * (y - x), x * (rho - z) - y, x * y - beta * z]\n",
      "\n",
      "# Choose random starting points, uniformly distributed from -15 to 15\n",
      "N_trajectories = 20\n",
      "np.random.seed(1)\n",
      "x0 = -15 + 30 * np.random.random((N_trajectories, 3))\n",
      "\n",
      "# Solve for the trajectories\n",
      "t = np.linspace(0, 4, 1000)\n",
      "x_t = np.asarray([integrate.odeint(lorenz_deriv, x0i, t)\n",
      "                  for x0i in x0])"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "We can plot the resulting trajectories in 3D using the ``mplot3d`` toolkit:"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "from mpl_toolkits.mplot3d import Axes3D\n",
      "\n",
      "fig = plt.figure()\n",
      "ax = fig.add_axes([0, 0, 1, 1], projection='3d')\n",
      "\n",
      "# choose a different color for each trajectory\n",
      "colors = plt.cm.jet(np.linspace(0, 1, N_trajectories))\n",
      "\n",
      "for i in range(N_trajectories):\n",
      "    x, y, z = x_t[i].T\n",
      "    ax.plot(x, y, z, c=colors[i])"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "heading",
     "level": 3,
     "metadata": {},
     "source": [
      "Animating the Result"
     ]
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "We can use this machinery the same way we have previously to animate these trajectories in 3D.  You should recognize most of this from what we've\n",
      "seen earlier:"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "from mpl_toolkits.mplot3d import Axes3D\n",
      "\n",
      "# Set up figure & 3D axis for animation\n",
      "fig = plt.figure()\n",
      "ax = fig.add_axes([0, 0, 1, 1], projection='3d')\n",
      "ax.axis('off')\n",
      "\n",
      "# choose a different color for each trajectory\n",
      "colors = plt.cm.jet(np.linspace(0, 1, N_trajectories))\n",
      "\n",
      "# set up lines and points\n",
      "lines = []\n",
      "pts = []\n",
      "for c in colors:\n",
      "    lines += ax.plot([], [], [], '-', c=c)\n",
      "    pts += ax.plot([], [], [], 'o', c=c)\n",
      "\n",
      "# prepare the axes limits\n",
      "ax.set_xlim((-25, 25))\n",
      "ax.set_ylim((-35, 35))\n",
      "ax.set_zlim((5, 55))\n",
      "\n",
      "# set point-of-view: specified by (altitude degrees, azimuth degrees)\n",
      "ax.view_init(30, 0)\n",
      "\n",
      "# initialization function: plot the background of each frame\n",
      "def init():\n",
      "    for line, pt in zip(lines, pts):\n",
      "        line.set_data([], [])\n",
      "        line.set_3d_properties([])\n",
      "\n",
      "        pt.set_data([], [])\n",
      "        pt.set_3d_properties([])\n",
      "    return lines + pts\n",
      "\n",
      "# animation function.  This will be called sequentially with the frame number\n",
      "def animate(i):\n",
      "    # we'll step two time-steps per frame.  This leads to nice results.\n",
      "    i = (2 * i) % x_t.shape[1]\n",
      "\n",
      "    for line, pt, xi in zip(lines, pts, x_t):\n",
      "        x, y, z = xi[:i].T\n",
      "        line.set_data(x, y)\n",
      "        line.set_3d_properties(z)\n",
      "\n",
      "        pt.set_data(x[-1:], y[-1:])\n",
      "        pt.set_3d_properties(z[-1:])\n",
      "\n",
      "    # rotate the point of view\n",
      "    ax.view_init(30, 0.3 * i)\n",
      "    fig.canvas.draw()\n",
      "    return lines + pts\n",
      "\n",
      "# instantiate the animator.\n",
      "anim = animation.FuncAnimation(fig, animate, init_func=init,\n",
      "                               frames=500, interval=30, blit=True)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "The following command can be run to save this as an mp4:"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "# Save as mp4. This requires mplayer or ffmpeg to be installed\n",
      "#anim.save('lorenz_attractor.mp4', fps=15, extra_args=['-vcodec', 'libx264'])"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "heading",
     "level": 2,
     "metadata": {},
     "source": [
      "Other Fun Examples"
     ]
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "I've included several other fun animation examples with this tutorial,\n",
      "which you can run from the examples directory and inspect via the\n",
      "``%load`` command:"
     ]
    },
    {
     "cell_type": "heading",
     "level": 3,
     "metadata": {},
     "source": [
      "Double Penulum"
     ]
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "This is an example I adapted from the matplotlib documentation,\n",
      "and wrote about here: http://jakevdp.github.com/blog/2012/08/18/matplotlib-animation-tutorial/"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "%run ../examples/double_pendulum.py"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "heading",
     "level": 3,
     "metadata": {},
     "source": [
      "Particles in a Box"
     ]
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "This is another example I created and showed here: http://jakevdp.github.com/blog/2012/08/18/matplotlib-animation-tutorial/"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "%run ../examples/particle_box.py"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "heading",
     "level": 3,
     "metadata": {},
     "source": [
      "Schrodinger Equation"
     ]
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "This example is slightly more involved, and uses a half-step\n",
      "Schrodinger Equation algorithm to solve for the evolution of\n",
      "a wave function within a given potential.  Details on the algorithm\n",
      "can be seen here: http://jakevdp.github.com/blog/2012/09/05/quantum-python/"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "%run ../examples/animate_schrodinger.py"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    }
   ],
   "metadata": {}
  }
 ]
}