{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> This is one of the 100 recipes of the [IPython Cookbook](http://ipython-books.github.io/), the definitive guide to high-performance scientific computing and data science in Python.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 12.3. Simulating an Ordinary Differential Equation with SciPy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "1. Let's import NumPy, SciPy (`integrate` package), and matplotlib."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import scipy.integrate as spi\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "2. We define a few parameters appearing in our model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "m = 1.  # particle's mass\n",
    "k = 1.  # drag coefficient\n",
    "g = 9.81  # gravity acceleration"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "3. We have two variables: `x` and `y` (two dimensions). We note $\\mathbf{u}=(x,y)$. The ODE we are going to simulate is:\n",
    "\n",
    "$$\\ddot{\\mathbf{u}} = -\\frac{k}{m} \\dot{\\mathbf{u}} + \\mathbf{g}$$\n",
    "\n",
    "where $\\mathbf{g}$ is the gravity acceleration vector. In order to simulate this second-order ODE with SciPy, we can convert it to a first-order ODE (another option would be to solve $\\dot{\\mathbf{u}}$ first before integrating the solution). To do this, we consider two 2D variables: $\\mathbf{u}$ and $\\dot{\\mathbf{u}}$. We note $\\mathbf{v} = (\\mathbf{u}, \\dot{\\mathbf{u}})$. We can express $\\dot{\\mathbf{v}}$ as a function of $\\mathbf{v}$. Now, we create the initial vector $\\mathbf{v}_0$ at time $t=0$: it has four components."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# The initial position is (0, 0).\n",
    "v0 = np.zeros(4)\n",
    "# The initial speed vector is oriented\n",
    "# to the top right.\n",
    "v0[2] = 4.\n",
    "v0[3] = 10."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "4. We need to create a Python function $f$ that takes the current vector $\\mathbf{v}(t_0)$ and a time $t_0$ as argument (with optional parameters), and that returns the derivative $\\dot{\\mathbf{v}}(t_0)$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def f(v, t0, k):\n",
    "    # v has four components: v=[u, u'].\n",
    "    u, udot = v[:2], v[2:]\n",
    "    # We compute the second derivative u'' of u.\n",
    "    udotdot = -k/m * udot\n",
    "    udotdot[1] -= g\n",
    "    # We return v'=[u', u''].\n",
    "    return np.r_[udot, udotdot]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "3. Now, we simulate the system for different values of $k$. We use the SciPy function `odeint`, defined in the `scipy.integrate` package."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(6,3));\n",
    "# We want to evaluate the system on 30 linearly\n",
    "# spaced times between t=0 and t=3.\n",
    "t = np.linspace(0., 3., 30)\n",
    "# We simulate the system for different values of k.\n",
    "for k in np.linspace(0., 1., 5):\n",
    "    # We simulate the system and evaluate $v$ on the \n",
    "    # given times.\n",
    "    v = spi.odeint(f, v0, t, args=(k,))\n",
    "    # We plot the particle's trajectory.\n",
    "    plt.plot(v[:,0], v[:,1], 'o-', mew=1, ms=8, mec='w',\n",
    "                label='k={0:.1f}'.format(k));\n",
    "plt.legend();\n",
    "plt.xlim(0, 12);\n",
    "plt.ylim(0, 6);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The most outward trajectory (blue) corresponds to drag-free motion (without air resistance). It is a parabola. In the other trajectories, we can observe the increasing effect of air resistance, parameterized with $k$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> You'll find all the explanations, figures, references, and much more in the book (to be released later this summer).\n",
    "\n",
    "> [IPython Cookbook](http://ipython-books.github.io/), by [Cyrille Rossant](http://cyrille.rossant.net), Packt Publishing, 2014 (500 pages)."
   ]
  }
 ],
 "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.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}