{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Exploring the Lorenz System of Differential Equations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this Notebook we explore the Lorenz system of differential equations:\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "\\dot{x} & = \\sigma(y-x) \\\\\n",
    "\\dot{y} & = \\rho x - y - xz \\\\\n",
    "\\dot{z} & = -\\beta z + xy\n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "This is one of the classic systems in non-linear differential equations. It exhibits a range of different behaviors as the parameters ($\\sigma$, $\\beta$, $\\rho$) are varied."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, we import the needed things from IPython, NumPy, Matplotlib and SciPy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from ipywidgets import interact, interactive\n",
    "from IPython.display import clear_output, display, HTML"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from scipy import integrate\n",
    "\n",
    "from matplotlib import pyplot as plt\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "from matplotlib.colors import cnames\n",
    "from matplotlib import animation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Computing the trajectories and plotting the result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We define a function that can integrate the differential equations numerically and then plot the solutions. This function has arguments that control the parameters of the differential equation ($\\sigma$, $\\beta$, $\\rho$), the numerical integration (`N`, `max_time`) and the visualization (`angle`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def solve_lorenz(N=10, angle=0.0, max_time=4.0, sigma=10.0, beta=8./3, rho=28.0):\n",
    "\n",
    "    fig = plt.figure()\n",
    "    ax = fig.add_axes([0, 0, 1, 1], projection='3d')\n",
    "    ax.axis('off')\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",
    "    def lorenz_deriv(x_y_z, t0, sigma=sigma, beta=beta, rho=rho):\n",
    "        \"\"\"Compute the time-derivative of a Lorenz system.\"\"\"\n",
    "        x, y, z = x_y_z\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",
    "    np.random.seed(1)\n",
    "    x0 = -15 + 30 * np.random.random((N, 3))\n",
    "\n",
    "    # Solve for the trajectories\n",
    "    t = np.linspace(0, max_time, int(250*max_time))\n",
    "    x_t = np.asarray([integrate.odeint(lorenz_deriv, x0i, t)\n",
    "                      for x0i in x0])\n",
    "    \n",
    "    # choose a different color for each trajectory\n",
    "    colors = plt.cm.jet(np.linspace(0, 1, N))\n",
    "\n",
    "    for i in range(N):\n",
    "        x, y, z = x_t[i,:,:].T\n",
    "        lines = ax.plot(x, y, z, '-', c=colors[i])\n",
    "        plt.setp(lines, linewidth=2)\n",
    "\n",
    "    ax.view_init(30, angle)\n",
    "    plt.show()\n",
    "\n",
    "    return t, x_t"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's call the function once to view the solutions. For this set of parameters, we see the trajectories swirling around two points, called attractors. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "t, x_t = solve_lorenz(angle=0, N=10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using IPython's `interactive` function, we can explore how the trajectories behave as we change the various parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "w = interactive(solve_lorenz, angle=(0.,360.), N=(0,50), sigma=(0.0,50.0), rho=(0.0,50.0))\n",
    "display(w)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The object returned by `interactive` is a `Widget` object and it has attributes that contain the current result and arguments:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "t, x_t = w.result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "w.kwargs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After interacting with the system, we can take the result and perform further computations. In this case, we compute the average positions in $x$, $y$ and $z$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "xyz_avg = x_t.mean(axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "xyz_avg.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Creating histograms of the average positions (across different trajectories) show that on average the trajectories swirl about the attractors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "plt.hist(xyz_avg[:,0])\n",
    "plt.title('Average $x(t)$')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "plt.hist(xyz_avg[:,1])\n",
    "plt.title('Average $y(t)$')"
   ]
  }
 ],
 "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.4.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
