ch09-code-listing.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Chapter 9: Ordinary differential equations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Robert Johansson\n",
    "\n",
    "Source code listings for [Numerical Python - A Practical Techniques Approach for Industry](http://www.apress.com/9781484205549) (ISBN 978-1-484205-54-9).\n",
    "\n",
    "The source code listings can be downloaded from http://www.apress.com/9781484205549"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "%config InlineBackend.figure_format='retina'\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib as mpl\n",
    "# mpl.rcParams['text.usetex'] = True\n",
    "mpl.rcParams['mathtext.fontset'] = 'stix'\n",
    "mpl.rcParams['font.family'] = 'serif'\n",
    "mpl.rcParams['font.sans-serif'] = 'stix'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import sympy\n",
    "sympy.init_printing()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from scipy import integrate"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Symbolic ODE solving with SymPy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Newton's law of cooling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "t, k, T0, Ta = sympy.symbols(\"t, k, T_0, T_a\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "T = sympy.Function(\"T\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "ode = T(t).diff(t) + k*(T(t) - Ta)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "ode"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "ode_sol = sympy.dsolve(ode)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "ode_sol"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "ode_sol.lhs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "ode_sol.rhs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "ics = {T(0): T0}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "ics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "C_eq = ode_sol.subs(t, 0).subs(ics)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "C_eq"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "C_sol = sympy.solve(C_eq)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "C_sol"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "ode_sol.subs(C_sol[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Function for applying initial conditions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def apply_ics(sol, ics, x, known_params):\n",
    "    \"\"\"\n",
    "    Apply the initial conditions (ics), given as a dictionary on\n",
    "    the form ics = {y(0): y0: y(x).diff(x).subs(x, 0): yp0, ...}\n",
    "    to the solution of the ODE with indepdendent variable x.\n",
    "    The undetermined integration constants C1, C2, ... are extracted\n",
    "    from the free symbols of the ODE solution, excluding symbols in\n",
    "    the known_params list.\n",
    "    \"\"\"\n",
    "    free_params = sol.free_symbols - set(known_params)\n",
    "    eqs = [(sol.lhs.diff(x, n) - sol.rhs.diff(x, n)).subs(x, 0).subs(ics)\n",
    "           for n in range(len(ics))]\n",
    "    sol_params = sympy.solve(eqs, free_params)\n",
    "    return sol.subs(sol_params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "ode_sol"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "apply_ics(ode_sol, ics, t, [k, Ta])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "ode_sol = apply_ics(ode_sol, ics, t, [k, Ta]).simplify()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "ode_sol"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "y_x = sympy.lambdify((t, k), ode_sol.rhs.subs({T0: 5, Ta: 1}), 'numpy')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(8, 4))\n",
    "\n",
    "x = np.linspace(0, 4, 100)\n",
    "\n",
    "for k in [1, 2, 3]:\n",
    "    ax.plot(x, y_x(x, k), label=r\"$k=%d$\" % k)\n",
    "\n",
    "ax.set_title(r\"$%s$\" % sympy.latex(ode_sol), fontsize=18)\n",
    "ax.set_xlabel(r\"$x$\", fontsize=18)\n",
    "ax.set_ylabel(r\"$y$\", fontsize=18)\n",
    "ax.legend()\n",
    "\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Damped harmonic oscillator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "t, omega0 = sympy.symbols(\"t, omega_0\", positive=True)\n",
    "gamma = sympy.symbols(\"gamma\", complex=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "x = sympy.Function(\"x\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "ode = x(t).diff(t, 2) + 2 * gamma * omega0 * x(t).diff(t) + omega0**2 * x(t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "ode"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "ode_sol = sympy.dsolve(ode)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "ode_sol"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "ics = {x(0): 1, x(t).diff(t).subs(t, 0): 0}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "ics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "x_t_sol = apply_ics(ode_sol, ics, t, [omega0, gamma])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "x_t_sol"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "x_t_critical = sympy.limit(x_t_sol.rhs, gamma, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "x_t_critical"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(8, 4))\n",
    "\n",
    "tt = np.linspace(0, 3, 250)\n",
    "for g in [0.1, 0.5, 1, 2.0, 5.0]:\n",
    "    if g == 1:\n",
    "        x_t = sympy.lambdify(t, x_t_critical.subs({omega0: 2.0 * sympy.pi}), 'numpy')\n",
    "    else:\n",
    "        x_t = sympy.lambdify(t, x_t_sol.rhs.subs({omega0: 2.0 * sympy.pi, gamma: g}), 'numpy')\n",
    "    ax.plot(tt, x_t(tt).real, label=r\"$\\gamma = %.1f$\" % g)\n",
    "\n",
    "ax.set_xlabel(r\"$t$\", fontsize=18)\n",
    "ax.set_ylabel(r\"$x(t)$\", fontsize=18)\n",
    "ax.set_xlim(0, 3)\n",
    "ax.legend()\n",
    "\n",
    "fig.tight_layout()\n",
    "fig.savefig('ch9-harmonic-oscillator.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# example of equation that sympy cannot solve"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = sympy.symbols(\"x\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y = sympy.Function(\"y\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f = y(x)**2 + x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sympy.Eq(y(x).diff(x, x), f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# sympy is unable to solve this differential equation\n",
    "sympy.dsolve(y(x).diff(x, x) - f)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Direction fields"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def  plot_direction_field(x, y_x, f_xy, x_lim=(-5, 5), y_lim=(-5, 5), ax=None):\n",
    "    \n",
    "    f_np = sympy.lambdify((x, y_x), f_xy, 'numpy')\n",
    "    \n",
    "    x_vec = np.linspace(x_lim[0], x_lim[1], 20)\n",
    "    y_vec = np.linspace(y_lim[0], y_lim[1], 20)\n",
    "    \n",
    "    if ax is None:\n",
    "        _, ax = plt.subplots(figsize=(4, 4))\n",
    "\n",
    "    dx = x_vec[1] - x_vec[0]\n",
    "    dy = y_vec[1] - y_vec[0]\n",
    "\n",
    "    for m, xx in enumerate(x_vec):\n",
    "        for n, yy in enumerate(y_vec):\n",
    "            Dy = f_np(xx, yy) * dx\n",
    "            Dx = 0.8 * dx**2 / np.sqrt(dx**2 + Dy**2)\n",
    "            Dy = 0.8 * Dy*dy / np.sqrt(dx**2 + Dy**2)\n",
    "            ax.plot([xx - Dx/2, xx + Dx/2],\n",
    "                    [yy - Dy/2, yy + Dy/2], 'b', lw=0.5)\n",
    "    ax.axis('tight')\n",
    "\n",
    "    ax.set_title(r\"$%s$\" %\n",
    "                 (sympy.latex(sympy.Eq(y(x).diff(x), f_xy))),\n",
    "                 fontsize=18)\n",
    "    \n",
    "    return ax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "x = sympy.symbols(\"x\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "y = sympy.Function(\"y\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1, 3, figsize=(12, 4))\n",
    "\n",
    "plot_direction_field(x, y(x), y(x)**2 + x, ax=axes[0])\n",
    "plot_direction_field(x, y(x), -x / y(x), ax=axes[1])\n",
    "plot_direction_field(x, y(x), y(x)**2 / x, ax=axes[2])\n",
    "\n",
    "fig.tight_layout()\n",
    "fig.savefig('ch9-direction-field.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Inexact solutions to ODEs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "x = sympy.symbols(\"x\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "y = sympy.Function(\"y\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "f = y(x)**2 + x\n",
    "#f = sympy.cos(y(x)**2) + x\n",
    "#f = sympy.sqrt(y(x)) * sympy.cos(x**2)\n",
    "#f = y(x)**2 / x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "sympy.Eq(y(x).diff(x), f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "ics = {y(0): 0}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "ode_sol = sympy.dsolve(y(x).diff(x) - f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "ode_sol"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "ode_sol = apply_ics(ode_sol, {y(0): 0}, x, [])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "ode_sol"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "ode_sol = sympy.dsolve(y(x).diff(x) - f, ics=ics)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "ode_sol"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1, 2, figsize=(8, 4))\n",
    "\n",
    "plot_direction_field(x, y(x), f, ax=axes[0])\n",
    "x_vec = np.linspace(-3, 3, 100)\n",
    "axes[0].plot(x_vec, sympy.lambdify(x, ode_sol.rhs.removeO())(x_vec), 'b', lw=2)\n",
    "axes[0].set_ylim(-5, 5)\n",
    "\n",
    "plot_direction_field(x, y(x), f, ax=axes[1])\n",
    "x_vec = np.linspace(-1, 1, 100)\n",
    "axes[1].plot(x_vec, sympy.lambdify(x, ode_sol.rhs.removeO())(x_vec), 'b', lw=2)\n",
    "\n",
    "ode_sol_m = ode_sol_p = ode_sol\n",
    "dx = 0.125\n",
    "for x0 in np.arange(1, 2., dx):\n",
    "    x_vec = np.linspace(x0, x0 + dx, 100)\n",
    "    ics = {y(x0): ode_sol_p.rhs.removeO().subs(x, x0)}\n",
    "    ode_sol_p = sympy.dsolve(y(x).diff(x) - f, ics=ics, n=6)\n",
    "    axes[1].plot(x_vec, sympy.lambdify(x, ode_sol_p.rhs.removeO())(x_vec), 'r', lw=2)\n",
    "\n",
    "for x0 in np.arange(-1, -5, -dx):\n",
    "    x_vec = np.linspace(x0, x0 - dx, 100)\n",
    "    ics = {y(x0): ode_sol_m.rhs.removeO().subs(x, x0)}\n",
    "    ode_sol_m = sympy.dsolve(y(x).diff(x) - f, ics=ics, n=6)\n",
    "    axes[1].plot(x_vec, sympy.lambdify(x, ode_sol_m.rhs.removeO())(x_vec), 'r', lw=2)\n",
    "    \n",
    "fig.tight_layout()\n",
    "fig.savefig(\"ch9-direction-field-and-approx-sol.pdf\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Laplace transformation method"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "t = sympy.symbols(\"t\", positive=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "s, Y = sympy.symbols(\"s, Y\", real=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "y = sympy.Function(\"y\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "ode = y(t).diff(t, 2) + 2 * y(t).diff(t) + 10 * y(t) - 2 * sympy.sin(3*t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "ode"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "L_y = sympy.laplace_transform(y(t), t, s)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "L_y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "L_ode = sympy.laplace_transform(ode, t, s, noconds=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "L_ode"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def laplace_transform_derivatives(e):\n",
    "    \"\"\"\n",
    "    Evaluate the laplace transforms of derivatives of functions\n",
    "    \"\"\"\n",
    "    if isinstance(e, sympy.LaplaceTransform):\n",
    "        if isinstance(e.args[0], sympy.Derivative):\n",
    "            d, t, s = e.args\n",
    "            n = len(d.args) - 1\n",
    "            return ((s**n) * sympy.LaplaceTransform(d.args[0], t, s) - \n",
    "                    sum([s**(n-i) * sympy.diff(d.args[0], t, i-1).subs(t, 0)\n",
    "                         for i in range(1, n+1)]))\n",
    "        \n",
    "    if isinstance(e, (sympy.Add, sympy.Mul)):\n",
    "        t = type(e)\n",
    "        return t(*[laplace_transform_derivatives(arg) for arg in e.args])\n",
    "    \n",
    "    return e"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "L_ode_2 = laplace_transform_derivatives(L_ode)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "L_ode_2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "L_ode_3 = L_ode_2.subs(L_y, Y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "L_ode_3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "ics = {y(0): 1, y(t).diff(t).subs(t, 0): 0}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "ics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "L_ode_4 = L_ode_3.subs(ics)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "L_ode_4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "Y_sol = sympy.solve(L_ode_4, Y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "Y_sol"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "sympy.apart(Y_sol[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "y_sol = sympy.inverse_laplace_transform(Y_sol[0], s, t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "sympy.simplify(y_sol)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "y_t = sympy.lambdify(t, y_sol, 'numpy')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(8, 4))\n",
    "\n",
    "tt = np.linspace(0, 10, 500)\n",
    "ax.plot(tt, y_t(tt).real)\n",
    "ax.set_xlabel(r\"$t$\", fontsize=18)\n",
    "ax.set_ylabel(r\"$y(t)$\", fontsize=18)\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Numerical integration of ODEs using SciPy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "x = sympy.symbols(\"x\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "y = sympy.Function(\"y\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "f = y(x)**2 + x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "f_np = sympy.lambdify((y(x), x), f, 'math')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "y0 = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "xp = np.linspace(0, 1.9, 100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "xp.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "yp = integrate.odeint(f_np, y0, xp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "xm = np.linspace(0, -5, 100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "ym = integrate.odeint(f_np, y0, xm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1, figsize=(4, 4))\n",
    "plot_direction_field(x, y(x), f, ax=ax)\n",
    "ax.plot(xm, ym, 'b', lw=2)\n",
    "ax.plot(xp, yp, 'r', lw=2)\n",
    "fig.savefig('ch9-odeint-single-eq-example.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Lotka-Volterra equations for predator/pray populations\n",
    "\n",
    "$$\n",
    "x'(t) = a x - b x y\n",
    "$$\n",
    "\n",
    "$$\n",
    "y'(t) = c x y - d y\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "a, b, c, d = 0.4, 0.002, 0.001, 0.7"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def f(xy, t):\n",
    "    x, y = xy\n",
    "    return [a * x - b * x * y,\n",
    "            c * x * y - d * y]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "xy0 = [600, 400]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "t = np.linspace(0, 50, 250)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "xy_t = integrate.odeint(f, xy0, t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "xy_t.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1, 2, figsize=(8, 4))\n",
    "\n",
    "axes[0].plot(t, xy_t[:,0], 'r', label=\"Prey\")\n",
    "axes[0].plot(t, xy_t[:,1], 'b', label=\"Predator\")\n",
    "axes[0].set_xlabel(\"Time\")\n",
    "axes[0].set_ylabel(\"Number of animals\")\n",
    "axes[0].legend()\n",
    "\n",
    "axes[1].plot(xy_t[:,0], xy_t[:,1], 'k')\n",
    "axes[1].set_xlabel(\"Number of prey\")\n",
    "axes[1].set_ylabel(\"Number of predators\")\n",
    "fig.tight_layout()\n",
    "fig.savefig('ch9-lokta-volterra.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Lorenz equations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$\n",
    "x'(t) = \\sigma(y - x)\n",
    "$$\n",
    "$$\n",
    "y'(t) = x(\\rho - z) - y \n",
    "$$\n",
    "$$\n",
    "z'(t) = x y - \\beta z\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def f(xyz, t, rho, sigma, beta):\n",
    "    x, y, z = xyz\n",
    "    return [sigma * (y - x),\n",
    "            x * (rho - z) - y,\n",
    "            x * y - beta * z]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "rho = 28\n",
    "sigma = 8\n",
    "beta = 8/3.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "t = np.linspace(0, 25, 10000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "xyz0 = [1.0, 1.0, 1.0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "xyz1 = integrate.odeint(f, xyz0, t, args=(rho, sigma, beta))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "xyz2 = integrate.odeint(f, xyz0, t, args=(rho, sigma, 0.6*beta))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "xyz3 = integrate.odeint(f, xyz0, t, args=(rho, 2*sigma, 0.6*beta))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "xyz3.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from mpl_toolkits.mplot3d.axes3d import Axes3D"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(12, 3.5), subplot_kw={'projection': '3d'})\n",
    "\n",
    "for ax, xyz, c in [(ax1, xyz1, 'r'), (ax2, xyz2, 'b'), (ax3, xyz3, 'g')]:\n",
    "    ax.plot(xyz[:,0], xyz[:,1], xyz[:,2], c, alpha=0.5)\n",
    "    ax.set_xlabel('$x$', fontsize=16)\n",
    "    ax.set_ylabel('$y$', fontsize=16)\n",
    "    ax.set_zlabel('$z$', fontsize=16)\n",
    "    ax.set_xticks([-15, 0, 15])\n",
    "    ax.set_yticks([-20, 0, 20])\n",
    "    ax.set_zticks([0, 20, 40])\n",
    "\n",
    "fig.tight_layout()\n",
    "fig.savefig('ch9-lorenz-equations.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Coupled damped springs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As second-order equations:\n",
    "\n",
    "\\begin{eqnarray}\n",
    "m_1 x_1''(t) + \\gamma_1 x_1'(t) + k_1 (x_1(t) - l_1) - k_2 (x_2(t) - x_1(t) - l_2) &=& 0\\\\\n",
    "m_2 x_2''(t) + \\gamma_2 x_2' + k_2 (x_2 - x_1 - l_2) &=& 0\n",
    "\\end{eqnarray}\n",
    "\n",
    "On standard form:\n",
    "\n",
    "\\begin{align}\n",
    "y_1'(t) &= y_2(t) \\\\\n",
    "y_2'(t) &= -\\gamma_1/m_1 y_2(t) - k_1/m_1 (y_1(t) - l_1) + k_2 (y_3(t) - y_1(t) - l_2)/m_1 \\\\\n",
    "y_3'(t) &= y_4(t) \\\\\n",
    "y_4'(t) &= - \\gamma_2 y_4(t)/m_2 - k_2 (y_3(t) - y_1(t) - l_2)/m_2 \\\\\n",
    "\\end{align}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def f(t, y, args):\n",
    "    m1, k1, g1, m2, k2, g2 = args\n",
    "    \n",
    "    return [y[1], \n",
    "            - k1/m1 * y[0] + k2/m1 * (y[2] - y[0]) - g1/m1 * y[1], \n",
    "            y[3], \n",
    "            - k2/m2 * (y[2] - y[0]) - g2/m2 * y[3] ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "m1, k1, g1 = 1.0, 10.0, 0.5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "m2, k2, g2 = 2.0, 40.0, 0.25"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "args = (m1, k1, g1, m2, k2, g2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "y0 = [1.0, 0, 0.5, 0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "t = np.linspace(0, 20, 1000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "r = integrate.ode(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "r.set_integrator('lsoda');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "r.set_initial_value(y0, t[0]);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "r.set_f_params(args);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "dt = t[1] - t[0]\n",
    "y = np.zeros((len(t), len(y0)))\n",
    "idx = 0\n",
    "while r.successful() and r.t < t[-1]:\n",
    "    y[idx, :] = r.y\n",
    "    r.integrate(r.t + dt)\n",
    "    idx += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(10, 4))\n",
    "ax1 = plt.subplot2grid((2, 5), (0, 0), colspan=3)\n",
    "ax2 = plt.subplot2grid((2, 5), (1, 0), colspan=3)\n",
    "ax3 = plt.subplot2grid((2, 5), (0, 3), colspan=2, rowspan=2)\n",
    "\n",
    "ax1.plot(t, y[:, 0], 'r')\n",
    "ax1.set_ylabel('$x_1$', fontsize=18)\n",
    "ax1.set_yticks([-1, -.5, 0, .5, 1])\n",
    "\n",
    "ax2.plot(t, y[:, 2], 'b')\n",
    "ax2.set_xlabel('$t$', fontsize=18)\n",
    "ax2.set_ylabel('$x_2$', fontsize=18)\n",
    "ax2.set_yticks([-1, -.5, 0, .5, 1])\n",
    "\n",
    "ax3.plot(y[:, 0], y[:, 2], 'k')\n",
    "ax3.set_xlabel('$x_1$', fontsize=18)\n",
    "ax3.set_ylabel('$x_2$', fontsize=18)\n",
    "ax3.set_xticks([-1, -.5, 0, .5, 1])\n",
    "ax3.set_yticks([-1, -.5, 0, .5, 1])\n",
    "\n",
    "fig.tight_layout()\n",
    "fig.savefig('ch9-coupled-damped-springs.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Same calculation as above, but with specifying the Jacobian as well:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def jac(t, y, args):\n",
    "    m1, k1, g1, m2, k2, g2 = args\n",
    "    \n",
    "    return [[0, 1, 0, 0], \n",
    "            [- k1/m1 - k2/m1, - g1/m1 * y[1], k2/m1, 0],\n",
    "            [0, 0, 1, 0],\n",
    "            [k2/m2, 0, - k2/m2, - g2/m2]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "r = integrate.ode(f, jac).set_f_params(args).set_jac_params(args).set_initial_value(y0, t[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "dt = t[1] - t[0]\n",
    "y = np.zeros((len(t), len(y0)))\n",
    "idx = 0\n",
    "while r.successful() and r.t < t[-1]:\n",
    "    y[idx, :] = r.y\n",
    "    r.integrate(r.t + dt)\n",
    "    idx += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(10, 4))\n",
    "ax1 = plt.subplot2grid((2, 5), (0, 0), colspan=3)\n",
    "ax2 = plt.subplot2grid((2, 5), (1, 0), colspan=3)\n",
    "ax3 = plt.subplot2grid((2, 5), (0, 3), colspan=2, rowspan=2)\n",
    "\n",
    "ax1.plot(t, y[:, 0], 'r')\n",
    "ax1.set_ylabel('$x_1$', fontsize=18)\n",
    "ax1.set_yticks([-1, -.5, 0, .5, 1])\n",
    "\n",
    "ax2.plot(t, y[:, 2], 'b')\n",
    "ax2.set_xlabel('$t$', fontsize=18)\n",
    "ax2.set_ylabel('$x_2$', fontsize=18)\n",
    "ax2.set_yticks([-1, -.5, 0, .5, 1])\n",
    "\n",
    "ax3.plot(y[:, 0], y[:, 2], 'k')\n",
    "ax3.set_xlabel('$x_1$', fontsize=18)\n",
    "ax3.set_ylabel('$x_2$', fontsize=18)\n",
    "ax3.set_xticks([-1, -.5, 0, .5, 1])\n",
    "ax3.set_yticks([-1, -.5, 0, .5, 1])\n",
    "\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Same calculation, but using SymPy to setup the problem for SciPy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "L1 = L2 = 0\n",
    "t = sympy.symbols(\"t\")\n",
    "m1, k1, b1 = sympy.symbols(\"m_1, k_1, b_1\")\n",
    "m2, k2, b2 = sympy.symbols(\"m_2, k_2, b_2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "x1 = sympy.Function(\"x_1\")\n",
    "x2 = sympy.Function(\"x_2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "ode1 = sympy.Eq(m1 * x1(t).diff(t,t,) + b1 * x1(t).diff(t) + k1*(x1(t)-L1) - k2*(x2(t)-x1(t) - L2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "ode2 = sympy.Eq(m2 * x2(t).diff(t,t,) + b2 * x2(t).diff(t) + k2*(x2(t)-x1(t)-L2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "params = {m1: 1.0, k1: 10.0, b1: 0.5,\n",
    "          m2: 2.0, k2: 40.0, b2: 0.25}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "args"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "y1 = sympy.Function(\"y_1\")\n",
    "y2 = sympy.Function(\"y_2\")\n",
    "y3 = sympy.Function(\"y_3\")\n",
    "y4 = sympy.Function(\"y_4\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "varchange = {x1(t).diff(t, t): y2(t).diff(t), \n",
    "             x1(t): y1(t),\n",
    "             x2(t).diff(t, t): y4(t).diff(t), \n",
    "             x2(t): y3(t)}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "(ode1.subs(varchange).lhs, ode2.subs(varchange).lhs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "ode3 = y1(t).diff(t) - y2(t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "ode4 = y3(t).diff(t) - y4(t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "vcsol = sympy.solve((ode1.subs(varchange), ode2.subs(varchange), ode3, ode4),\n",
    "                    (y1(t).diff(t), y2(t).diff(t), y3(t).diff(t), y4(t).diff(t)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "vcsol"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "ode_rhs = sympy.Matrix([y1(t).diff(t), y2(t).diff(t), y3(t).diff(t), y4(t).diff(t)]).subs(vcsol)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "y = sympy.Matrix([y1(t), y2(t), y3(t), y4(t)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "sympy.Eq(y.diff(t), ode_rhs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ode_rhs.subs(params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_f_np = sympy.lambdify((t, y), ode_rhs.subs(params), 'numpy')\n",
    "f_np = lambda _x, _y, *args: _f_np(_x, _y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "y0 = [1.0, 0, 0.5, 0]\n",
    "tt = np.linspace(0, 20, 1000)\n",
    "\n",
    "r = integrate.ode(f_np)\n",
    "r.set_integrator('lsoda');\n",
    "r.set_initial_value(y0, tt[0]);\n",
    "\n",
    "dt = tt[1] - tt[0]\n",
    "yy = np.zeros((len(tt), len(y0)))\n",
    "idx = 0\n",
    "while r.successful() and r.t < tt[-1]:\n",
    "    yy[idx, :] = r.y\n",
    "    r.integrate(r.t + dt)\n",
    "    idx += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(10, 4))\n",
    "ax1 = plt.subplot2grid((2, 5), (0, 0), colspan=3)\n",
    "ax2 = plt.subplot2grid((2, 5), (1, 0), colspan=3)\n",
    "ax3 = plt.subplot2grid((2, 5), (0, 3), colspan=2, rowspan=2)\n",
    "\n",
    "ax1.plot(tt, yy[:, 0], 'r')\n",
    "ax1.set_ylabel('$x_1$', fontsize=18)\n",
    "ax1.set_yticks([-1, -.5, 0, .5, 1])\n",
    "\n",
    "ax2.plot(tt, yy[:, 2], 'b')\n",
    "ax2.set_xlabel('$t$', fontsize=18)\n",
    "ax2.set_ylabel('$x_2$', fontsize=18)\n",
    "ax2.set_yticks([-1, -.5, 0, .5, 1])\n",
    "\n",
    "ax3.plot(yy[:, 0], yy[:, 2], 'k')\n",
    "ax3.set_xlabel('$x_1$', fontsize=18)\n",
    "ax3.set_ylabel('$x_2$', fontsize=18)\n",
    "ax3.set_xticks([-1, -.5, 0, .5, 1])\n",
    "ax3.set_yticks([-1, -.5, 0, .5, 1])\n",
    "\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Doube pendulum"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "http://scienceworld.wolfram.com/physics/DoublePendulum.html"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$\n",
    "(m_1+m_2) l_1\\theta_1'' + m_2l_2\\theta_2''\\cos(\\theta_1-\\theta_2)\n",
    "+ m_2l_2(\\theta_2')^2\\sin(\\theta_1-\\theta_2)+g(m_1+m_2)\\sin(\\theta_1) = 0\n",
    "$$\n",
    "\n",
    "$$\n",
    "m_2l_2\\theta_2'' + m_2l_1\\theta_1''\\cos(\\theta_1-\\theta_2) - m_2l_1 (\\theta_1')^2 \\sin(\\theta_1-\\theta_2)\n",
    "+m_2g\\sin(\\theta_2) = 0\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "t, g, m1, l1, m2, l2 = sympy.symbols(\"t, g, m_1, l_1, m_2, l_2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "theta1, theta2 = sympy.symbols(\"theta_1, theta_2\", cls=sympy.Function)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "ode1 = sympy.Eq((m1+m2)*l1 * theta1(t).diff(t,t) +\n",
    "                m2*l2 * theta2(t).diff(t,t) * sympy.cos(theta1(t)-theta2(t))  +\n",
    "                m2*l2 * theta2(t).diff(t)**2 * sympy.sin(theta1(t)-theta2(t)) + \n",
    "                g*(m1+m2) * sympy.sin(theta1(t)))\n",
    "ode1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "ode2 = sympy.Eq(m2*l2 * theta2(t).diff(t,t) +\n",
    "                m2*l1 * theta1(t).diff(t,t) * sympy.cos(theta1(t)-theta2(t)) -\n",
    "                m2*l1 * theta1(t).diff(t)**2 * sympy.sin(theta1(t) - theta2(t)) +\n",
    "                m2*g * sympy.sin(theta2(t)))\n",
    "ode2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# this is fruitless, sympy cannot solve these ODEs\n",
    "try:\n",
    "    sympy.dsolve(ode1, ode2)\n",
    "except Exception as e:\n",
    "    print(e)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "y1, y2, y3, y4 = sympy.symbols(\"y_1, y_2, y_3, y_4\", cls=sympy.Function)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "varchange = {theta1(t).diff(t, t): y2(t).diff(t), \n",
    "             theta1(t): y1(t),\n",
    "             theta2(t).diff(t, t): y4(t).diff(t), \n",
    "             theta2(t): y3(t)}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "ode1_vc = ode1.subs(varchange)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "ode2_vc = ode2.subs(varchange)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "ode3 = y1(t).diff(t) - y2(t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "ode4 = y3(t).diff(t) - y4(t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "y = sympy.Matrix([y1(t), y2(t), y3(t), y4(t)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "vcsol = sympy.solve((ode1_vc, ode2_vc, ode3, ode4), y.diff(t), dict=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "f = y.diff(t).subs(vcsol[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "sympy.Eq(y.diff(t), f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "params = {m1: 5.0, l1: 2.0,\n",
    "          m2: 1.0, l2: 1.0, g: 10.0}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "_f_np = sympy.lambdify((t, y), f.subs(params), 'numpy')\n",
    "f_np = lambda _t, _y, *args: _f_np(_t, _y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "jac = sympy.Matrix([[fj.diff(yi) for yi in y] for fj in f])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "_jac_np = sympy.lambdify((t, y), jac.subs(params), 'numpy')\n",
    "jac_np = lambda _t, _y, *args: _jac_np(_t, _y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "y0 = [2.0, 0, 0.0, 0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "tt = np.linspace(0, 20, 1000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "r = integrate.ode(f_np, jac_np).set_initial_value(y0, tt[0]);\n",
    " \n",
    "dt = tt[1] - tt[0]\n",
    "yy = np.zeros((len(tt), len(y0)))\n",
    "idx = 0\n",
    "while r.successful() and r.t < tt[-1]:\n",
    "    yy[idx, :] = r.y\n",
    "    r.integrate(r.t + dt)\n",
    "    idx += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(10, 4))\n",
    "ax1 = plt.subplot2grid((2, 5), (0, 0), colspan=3)\n",
    "ax2 = plt.subplot2grid((2, 5), (1, 0), colspan=3)\n",
    "ax3 = plt.subplot2grid((2, 5), (0, 3), colspan=2, rowspan=2)\n",
    "\n",
    "ax1.plot(tt, yy[:, 0], 'r')\n",
    "ax1.set_ylabel(r'$\\theta_1$', fontsize=18)\n",
    "\n",
    "ax2.plot(tt, yy[:, 2], 'b')\n",
    "ax2.set_xlabel('$t$', fontsize=18)\n",
    "ax2.set_ylabel(r'$\\theta_2$', fontsize=18)\n",
    "\n",
    "ax3.plot(yy[:, 0], yy[:, 2], 'k')\n",
    "ax3.set_xlabel(r'$\\theta_1$', fontsize=18)\n",
    "ax3.set_ylabel(r'$\\theta_2$', fontsize=18)\n",
    "\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "theta1_np, theta2_np = yy[:, 0], yy[:, 2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "x1 = params[l1] * np.sin(theta1_np)\n",
    "y1 = -params[l1] * np.cos(theta1_np)\n",
    "x2 = x1 + params[l2] * np.sin(theta2_np)\n",
    "y2 = y1 - params[l2] * np.cos(theta2_np)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(10, 4))\n",
    "ax1 = plt.subplot2grid((2, 5), (0, 0), colspan=3)\n",
    "ax2 = plt.subplot2grid((2, 5), (1, 0), colspan=3)\n",
    "ax3 = plt.subplot2grid((2, 5), (0, 3), colspan=2, rowspan=2)\n",
    "\n",
    "ax1.plot(tt, x1, 'r')\n",
    "ax1.plot(tt, y1, 'b')\n",
    "ax1.set_ylabel('$x_1, y_1$', fontsize=18)\n",
    "ax1.set_yticks([-3, 0, 3])\n",
    "\n",
    "ax2.plot(tt, x2, 'r')\n",
    "ax2.plot(tt, y2, 'b')\n",
    "ax2.set_xlabel('$t$', fontsize=18)\n",
    "ax2.set_ylabel('$x_2, y_2$', fontsize=18)\n",
    "ax2.set_yticks([-3, 0, 3])\n",
    "\n",
    "ax3.plot(x1, y1, 'r', lw=2.0, label=\"trajectory of pendulum 1\")\n",
    "ax3.plot(x2, y2, 'b', lw=0.5, label=\"trajectory of pendulum 2\")\n",
    "ax3.set_xlabel('$x$', fontsize=18)\n",
    "ax3.set_ylabel('$y$', fontsize=18)\n",
    "ax3.set_xticks([-3, 0, 3])\n",
    "ax3.set_yticks([-3, 0, 3])\n",
    "ax3.legend()\n",
    "\n",
    "fig.tight_layout()\n",
    "fig.savefig('ch9-double-pendulum.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Versions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "%reload_ext version_information"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "%version_information numpy, scipy, sympy, matplotlib"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.6 (NPM ed2)",
   "language": "python",
   "name": "py3.6-npm-e2"
  },
  "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}