Lorenz Attractor.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import numpy as np, matplotlib.pyplot as plt, glob, os\n",
    "import IPython.display as IPdisplay, matplotlib.font_manager as fm\n",
    "from scipy.integrate import odeint\n",
    "from mpl_toolkits.mplot3d.axes3d import Axes3D\n",
    "from PIL import Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# define the fonts to use for plots\n",
    "family = 'Myriad Pro'\n",
    "title_font = fm.FontProperties(family=family, style='normal', size=20, weight='normal', stretch='normal')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "save_folder = 'images/lorenz-animate'\n",
    "if not os.path.exists(save_folder):\n",
    "    os.makedirs(save_folder)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# define the initial system state (aka x, y, z positions in space)\n",
    "initial_state = [0.1, 0, 0]\n",
    "\n",
    "# define the system parameters sigma, rho, and beta\n",
    "sigma = 10.\n",
    "rho   = 28.\n",
    "beta  = 8./3.\n",
    "\n",
    "# define the time points to solve for, evenly spaced between the start and end times\n",
    "start_time = 1\n",
    "end_time = 60\n",
    "interval = 100\n",
    "time_points = np.linspace(start_time, end_time, end_time * interval)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# define the lorenz system\n",
    "def lorenz_system(current_state, t):\n",
    "    x, y, z = current_state\n",
    "    dx_dt = sigma * (y - x)\n",
    "    dy_dt = x * (rho - z) - y\n",
    "    dz_dt = x * y - beta * z\n",
    "    return [dx_dt, dy_dt, dz_dt]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# plot the system in 3 dimensions\n",
    "def plot_lorenz(xyz, n):\n",
    "    fig = plt.figure(figsize=(12, 9))\n",
    "    ax = fig.gca(projection='3d')\n",
    "    ax.xaxis.set_pane_color((1,1,1,1))\n",
    "    ax.yaxis.set_pane_color((1,1,1,1))\n",
    "    ax.zaxis.set_pane_color((1,1,1,1))\n",
    "    x = xyz[:, 0]\n",
    "    y = xyz[:, 1]\n",
    "    z = xyz[:, 2]\n",
    "    ax.plot(x, y, z, color='g', alpha=0.7, linewidth=0.7)\n",
    "    ax.set_xlim((-30,30))\n",
    "    ax.set_ylim((-30,30))\n",
    "    ax.set_zlim((0,50))\n",
    "    ax.set_title('Lorenz system attractor', fontproperties=title_font)\n",
    "    \n",
    "    plt.savefig('{}/{:03d}.png'.format(save_folder, n), dpi=60, bbox_inches='tight', pad_inches=0.1)\n",
    "    plt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# return a list in iteratively larger chunks\n",
    "def get_chunks(full_list, size):\n",
    "    size = max(1, size)\n",
    "    chunks = [full_list[0:i] for i in range(1, len(full_list) + 1, size)]\n",
    "    return chunks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# get incrementally larger chunks of the time points, to reveal the attractor one frame at a time\n",
    "chunks = get_chunks(time_points, size=20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# get the points to plot, one chunk of time steps at a time, by integrating the system of equations\n",
    "points = [odeint(lorenz_system, initial_state, chunk) for chunk in chunks]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# plot each set of points, one at a time, saving each plot\n",
    "for n, point in enumerate(points):\n",
    "    plot_lorenz(point, n)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Animate it\n",
    "Create an animated gif of all the plots then display it inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# create a tuple of display durations, one for each frame\n",
    "first_last = 100 #show the first and last frames for 100 ms\n",
    "standard_duration = 5 #show all other frames for 5 ms\n",
    "durations = tuple([first_last] + [standard_duration] * (len(points) - 2) + [first_last])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# load all the static images into a list\n",
    "images = [Image.open(image) for image in glob.glob('{}/*.png'.format(save_folder))]\n",
    "gif_filepath = 'images/animated-lorenz-attractor.gif'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save as an animated gif\n",
    "gif = images[0]\n",
    "gif.info['duration'] = durations #ms per frame\n",
    "gif.info['loop'] = 0 #how many times to loop (0=infinite)\n",
    "gif.save(fp=gif_filepath, format='gif', save_all=True, append_images=images[1:])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# verify that the number of frames in the gif equals the number of image files and durations\n",
    "Image.open(gif_filepath).n_frames == len(images) == len(durations)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<img src=\"images/animated-lorenz-attractor.gif\"/>"
      ],
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "IPdisplay.Image(url=gif_filepath)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "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
}