{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "## Lecture 17: ODEs 3: Runge-Kutta\n",
    "### Phys 503, J.R. Gladden, University of Mississippi"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "### Runge-Kutta 4^th order\n",
    "4th order R-K is the \"work horse\" of ODE solvers. It has good accuracy and is stable.  There are more sophisticated methods that may be more efficient, but many solvers are based on R-K. See slides for details of the definition and logic.\n",
    "\n",
    "** Exponential Example - compare to Euler **"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "%matplotlib qt4\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "##################\n",
    "\n",
    "\n",
    "def rk4(f,x0,T,n=100):\n",
    "    dt=T/float(n)\n",
    "    times=np.linspace(0,T,n)\n",
    "    x=np.zeros(n)\n",
    "    x[0]=x0\n",
    "    for k in range(n-1):\n",
    "        F1=f(x[k])\n",
    "        F2=f(x[k]+0.5*dt*F1)\n",
    "        F3=f(x[k]+0.5*dt*F2)\n",
    "        F4=f(x[k]+dt*F3)\n",
    "        x[k+1]=x[k]+dt/6.0*(F1 + 2*F2 + 2*F3 +F4)\n",
    "    return times, x\n",
    "\n",
    "def euler(f,x0,T,n=100):\n",
    "    dt=T/float(n)\n",
    "    times=np.linspace(0,T,n)\n",
    "    x=np.zeros(n)\n",
    "    x[0]=x0\n",
    "    for k in range(n-1):\n",
    "        x[k+1]=x[k]+dt*f(x[k])\n",
    "    return times,x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "### Usage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "def run_expen():\n",
    "    def funct(x): return -x\n",
    "    \n",
    "    t,x_rk=rk4(funct,10,10,n=100)\n",
    "    t2,x_euler=euler(funct,10,10,n=100)\n",
    "    \n",
    "    exact=10*np.exp(-t2)\n",
    "    plt.plot(t2,exact,'k-',label='Exact')\n",
    "    \n",
    "    plt.plot(t,x_rk,'b-',label='RK4')\n",
    "    plt.plot(t2,x_euler,'g-',label='Euler')\n",
    "    \n",
    "    plt.legend()\n",
    "    plt.ylabel('X(t)')\n",
    "    plt.xlabel('Time')\n",
    "\n",
    "    plt.figure()\n",
    "    plt.plot(t,abs(x_rk - exact),label = 'RK4 Error')\n",
    "    plt.plot(t,abs(x_euler - exact),label = 'Euler Error')\n",
    "    plt.legend()\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "run_expen()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "### ODE solvers: scipy.integrate.odeint\n",
    "There are several mature ODE solvers supplied by different modules.  'scipy' is a common and well developed module with many useful tools.  These solvers are more sophisticated than what we have learned and many are based on VERY mature FORTRAN libraries such as 'odepack'."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": true,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "from scipy.integrate import odeint\n",
    "plt.close('all')\n",
    "t=np.linspace(0,10,100)\n",
    "\n",
    "def epidemic(state,t,r,a):\n",
    "    s=state[0]\n",
    "    i=state[1]\n",
    "    sdot = -r*s*i\n",
    "    idot = r*s*i - a*i\n",
    "    return [sdot,idot]\n",
    "\n",
    "state0=[200,2]\n",
    "r=0.02\n",
    "a=0.4\n",
    "\n",
    "soln, info = odeint(epidemic,state0,t,args=(r,a),full_output=True)\n",
    "\n",
    "s_soln = soln[:,0]\n",
    "i_soln = soln[:,1]\n",
    "\n",
    "plt.plot(t,s_soln)\n",
    "plt.plot(t,i_soln)\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'hu': array([ 0.02116528,  0.0391886 ,  0.02919062,  0.02919062,  0.01891494,\n",
       "         0.01880723,  0.03010075,  0.03010075,  0.03010075,  0.03880257,\n",
       "         0.03174594,  0.03174594,  0.03174594,  0.03701998,  0.03701998,\n",
       "         0.03701998,  0.04444459,  0.03579745,  0.03579745,  0.03982774,\n",
       "         0.03982774,  0.03982774,  0.03982774,  0.03982774,  0.03982774,\n",
       "         0.04533338,  0.04533338,  0.04533338,  0.06620826,  0.06620826,\n",
       "         0.06620826,  0.06620826,  0.06620826,  0.06620826,  0.06620826,\n",
       "         0.06620826,  0.06620826,  0.06620826,  0.06620826,  0.06620826,\n",
       "         0.07450483,  0.07450483,  0.07450483,  0.07450483,  0.07450483,\n",
       "         0.07450483,  0.09274909,  0.09274909,  0.09274909,  0.09274909,\n",
       "         0.09274909,  0.09274909,  0.09274909,  0.11853904,  0.11853904,\n",
       "         0.11853904,  0.11853904,  0.11853904,  0.11853904,  0.11853904,\n",
       "         0.11853904,  0.11853904,  0.15656804,  0.15656804,  0.15656804,\n",
       "         0.15656804,  0.15656804,  0.15656804,  0.15656804,  0.15656804,\n",
       "         0.15656804,  0.15656804,  0.15656804,  0.15656804,  0.15656804,\n",
       "         0.21377751,  0.21377751,  0.21377751,  0.21377751,  0.21377751,\n",
       "         0.21377751,  0.21377751,  0.21377751,  0.21377751,  0.21377751,\n",
       "         0.21377751,  0.21377751,  0.21377751,  0.21377751,  0.21377751,\n",
       "         0.21377751,  0.21377751,  0.3024592 ,  0.3024592 ,  0.3024592 ,\n",
       "         0.3024592 ,  0.3024592 ,  0.3024592 ,  0.3024592 ]),\n",
       " 'imxer': -1,\n",
       " 'leniw': 22,\n",
       " 'lenrw': 52,\n",
       " 'message': 'Integration successful.',\n",
       " 'mused': array([1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "        1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "        1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "        1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "        1, 1, 1, 1, 1, 1, 1], dtype=int32),\n",
       " 'nfe': array([ 35,  45,  53,  59,  73,  87,  99, 105, 111, 121, 127, 135, 141,\n",
       "        147, 153, 159, 163, 171, 177, 183, 187, 193, 197, 203, 207, 213,\n",
       "        217, 221, 227, 229, 233, 235, 239, 241, 245, 247, 251, 253, 257,\n",
       "        259, 263, 265, 267, 271, 273, 275, 279, 281, 283, 285, 287, 289,\n",
       "        291, 293, 295, 297, 299, 299, 301, 303, 305, 307, 309, 309, 311,\n",
       "        313, 313, 315, 317, 317, 319, 319, 321, 323, 323, 325, 325, 327,\n",
       "        327, 329, 329, 331, 331, 333, 333, 335, 335, 337, 337, 339, 339,\n",
       "        339, 341, 341, 341, 343, 343, 343, 345], dtype=int32),\n",
       " 'nje': array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "        0, 0, 0, 0, 0, 0, 0], dtype=int32),\n",
       " 'nqu': array([5, 5, 5, 5, 4, 5, 6, 6, 6, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7,\n",
       "        7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7,\n",
       "        7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7,\n",
       "        7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7,\n",
       "        7, 7, 7, 7, 7, 7, 7], dtype=int32),\n",
       " 'nst': array([ 16,  21,  24,  27,  32,  38,  44,  47,  50,  54,  56,  60,  63,\n",
       "         65,  68,  71,  73,  76,  79,  82,  84,  87,  89,  92,  94,  97,\n",
       "         99, 101, 104, 105, 107, 108, 110, 111, 113, 114, 116, 117, 119,\n",
       "        120, 122, 123, 124, 126, 127, 128, 130, 131, 132, 133, 134, 135,\n",
       "        136, 137, 138, 139, 140, 140, 141, 142, 143, 144, 145, 145, 146,\n",
       "        147, 147, 148, 149, 149, 150, 150, 151, 152, 152, 153, 153, 154,\n",
       "        154, 155, 155, 156, 156, 157, 157, 158, 158, 159, 159, 160, 160,\n",
       "        160, 161, 161, 161, 162, 162, 162, 163], dtype=int32),\n",
       " 'tcur': array([  0.10280228,   0.22665199,   0.32422184,   0.41179371,\n",
       "          0.51182119,   0.6092666 ,   0.7334035 ,   0.82370576,\n",
       "          0.91400802,   1.04311285,   1.11366135,   1.2406451 ,\n",
       "          1.33588291,   1.41809488,   1.52915483,   1.64021478,\n",
       "          1.72167936,   1.82907171,   1.93646406,   2.0478867 ,\n",
       "          2.12754218,   2.24702539,   2.32668087,   2.44616408,\n",
       "          2.52581956,   2.65631406,   2.74698083,   2.83764759,\n",
       "          2.99452262,   3.06073087,   3.19314739,   3.25935565,\n",
       "          3.39177216,   3.45798042,   3.59039694,   3.65660519,\n",
       "          3.78902171,   3.85522997,   3.98764648,   4.05385474,\n",
       "          4.2028644 ,   4.27736924,   4.35187407,   4.50088373,\n",
       "          4.57538856,   4.6498934 ,   4.83539157,   4.92814066,\n",
       "          5.02088975,   5.11363883,   5.20638792,   5.29913701,\n",
       "          5.3918861 ,   5.51042513,   5.62896417,   5.74750321,\n",
       "          5.86604224,   5.86604224,   5.98458128,   6.10312032,\n",
       "          6.22165935,   6.34019839,   6.49676643,   6.49676643,\n",
       "          6.65333447,   6.80990251,   6.80990251,   6.96647056,\n",
       "          7.1230386 ,   7.1230386 ,   7.27960664,   7.27960664,\n",
       "          7.43617468,   7.59274272,   7.59274272,   7.80652023,\n",
       "          7.80652023,   8.02029774,   8.02029774,   8.23407525,\n",
       "          8.23407525,   8.44785275,   8.44785275,   8.66163026,\n",
       "          8.66163026,   8.87540777,   8.87540777,   9.08918527,\n",
       "          9.08918527,   9.30296278,   9.30296278,   9.30296278,\n",
       "          9.60542198,   9.60542198,   9.60542198,   9.90788118,\n",
       "          9.90788118,   9.90788118,  10.21034038]),\n",
       " 'tolsf': array([  5.72938864e-313,   5.72938864e-313,   5.72938864e-313,\n",
       "          5.72938864e-313,   5.72938864e-313,   5.72938864e-313,\n",
       "          5.72938864e-313,   5.72938864e-313,   5.72938864e-313,\n",
       "          5.72938864e-313,   5.72938864e-313,   5.72938864e-313,\n",
       "          5.72938864e-313,   5.72938864e-313,   5.72938864e-313,\n",
       "          5.72938864e-313,   5.72938864e-313,   5.72938864e-313,\n",
       "          5.72938864e-313,   5.72938864e-313,   5.72938864e-313,\n",
       "          5.72938864e-313,   5.72938864e-313,   5.72938864e-313,\n",
       "          5.72938864e-313,   5.72938864e-313,   5.72938864e-313,\n",
       "          5.72938864e-313,   5.72938864e-313,   5.72938864e-313,\n",
       "          5.72938864e-313,   5.72938864e-313,   5.72938864e-313,\n",
       "          5.72938864e-313,   5.72938864e-313,   5.72938864e-313,\n",
       "          5.72938864e-313,   5.72938864e-313,   5.72938864e-313,\n",
       "          5.72938864e-313,   5.72938864e-313,   5.72938864e-313,\n",
       "          5.72938864e-313,   5.72938864e-313,   5.72938864e-313,\n",
       "          5.72938864e-313,   5.72938864e-313,   5.72938864e-313,\n",
       "          5.72938864e-313,   5.72938864e-313,   5.72938864e-313,\n",
       "          5.72938864e-313,   5.72938864e-313,   5.72938864e-313,\n",
       "          5.72938864e-313,   5.72938864e-313,   5.72938864e-313,\n",
       "          5.72938864e-313,   5.72938864e-313,   5.72938864e-313,\n",
       "          5.72938864e-313,   5.72938864e-313,   5.72938864e-313,\n",
       "          5.72938864e-313,   5.72938864e-313,   5.72938864e-313,\n",
       "          5.72938864e-313,   5.72938864e-313,   5.72938864e-313,\n",
       "          5.72938864e-313,   5.72938864e-313,   5.72938864e-313,\n",
       "          5.72938864e-313,   5.72938864e-313,   5.72938864e-313,\n",
       "          5.72938864e-313,   5.72938864e-313,   5.72938864e-313,\n",
       "          5.72938864e-313,   5.72938864e-313,   5.72938864e-313,\n",
       "          5.72938864e-313,   5.72938864e-313,   5.72938864e-313,\n",
       "          5.72938864e-313,   5.72938864e-313,   5.72938864e-313,\n",
       "          5.72938864e-313,   5.72938864e-313,   5.72938864e-313,\n",
       "          5.72938864e-313,   5.72938864e-313,   5.72938864e-313,\n",
       "          5.72938864e-313,   5.72938864e-313,   5.72938864e-313,\n",
       "          5.72938864e-313,   5.72938864e-313,   5.72938864e-313]),\n",
       " 'tsw': array([ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,\n",
       "         0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,\n",
       "         0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,\n",
       "         0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,\n",
       "         0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,\n",
       "         0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,\n",
       "         0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,\n",
       "         0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.])}"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "info\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "### Pendulum Example\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "def run_pend(r0=np.pi/10.0,v0=0,L=1.0,m=1.0):\n",
    "    def pend_acc(r): return -g/L*np.sin(r)\n",
    "    \n",
    "    def rk4(f,r0,v0,T,n=5000):\n",
    "        dt=T/float(n)\n",
    "        times=np.linspace(0,T,n)\n",
    "        r=np.zeros(n)\n",
    "        v=np.zeros(n)\n",
    "        ke=np.zeros(n)\n",
    "        pe=np.zeros(n)\n",
    "        r[0]=r0\n",
    "        v[0]=v0\n",
    "        ke[0]=0.5*m*L**2*v0**2\n",
    "        pe[0]=m*g*L*(1.0-np.cos(r0))\n",
    "\n",
    "        for k in range(n-1):\n",
    "            #coupled equations we need to do them together\n",
    "            vc=v[k]\n",
    "            rc=r[k]\n",
    "            F1= f(rc)*dt  \n",
    "            Fp1 = vc*dt\n",
    "            F2=f(rc+0.5*Fp1*dt) \n",
    "            Fp2 = vc+F1*dt/2.0\n",
    "            F3=f(rc+0.5*Fp2*dt)\n",
    "            Fp3 = vc+F2*dt/2.0\n",
    "            F4=f(rc+dt*Fp3)\n",
    "            Fp4 = vc+F3*dt\n",
    "\n",
    "            v[k+1]=vc+dt/6.0*(F1 + 2*F2 + 2*F3 +F4)\n",
    "            r[k+1] = rc+ dt/6.0 * (Fp1 +2*Fp2+2*Fp3+Fp4)\n",
    "\t\t\t\n",
    "            ke[k+1] = 0.5*m*L**2*v[k+1]**2\n",
    "            pe[k+1] = m*g*L*(1.0-np.cos(r[k+1]))\n",
    "        return times,r,v,ke,pe\n",
    "    \n",
    "    g=9.8\n",
    "    T = 2*np.pi*np.sqrt(L/g)\n",
    "    print 'Period is: %2.2f' % T\n",
    "\n",
    "    t,r_rk,v_rk,ke,pe = rk4(pend_acc,r0,v0,5*T)\n",
    "    #t,x_euler = euler(pend_acc,1,10)\n",
    "    \n",
    "    x_sma = r0*np.cos(2*np.pi*t/T)\n",
    "\n",
    "    plt.plot(t,r_rk,'b-',label='RK4')\n",
    "    #plot(t,x_euler,'g-',label='euler')\n",
    "    plt.plot(t,x_sma,'k-',label='Small Angle')\n",
    "    plt.xlabel('Time (s)')\n",
    "    plt.ylabel('Angle (rad)')\n",
    "    plt.legend()\n",
    "\t\n",
    "    plt.figure(2)\n",
    "    plt.plot(t,v_rk,'b-')\n",
    "    plt.xlabel('Time (s)')\n",
    "    plt.ylabel('Velocity (rad/sec)')\n",
    "\n",
    "    plt.figure(3)\n",
    "    plt.plot(t,ke,'b-',label='KE')\n",
    "    plt.plot(t,pe,'g-',label='PE')\n",
    "    plt.plot(t,ke+pe,'k-',label='Total E')\n",
    "    plt.ylabel('Energy (J)')\n",
    "    plt.xlabel('Time (s)')\n",
    "    plt.legend()\n",
    "\n",
    "    plt.figure(4)\n",
    "    plt.plot(r_rk,m*v_rk, lw=2)\n",
    "    plt.fill_between(r_rk,m*v_rk, color='r',alpha=0.7)\n",
    "    plt.xlabel('Angle')\n",
    "    plt.ylabel('Momentum')\n",
    "    plt.title('The Action')\n",
    "    plt.show()\n",
    "\n",
    "plt.close('all')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Period is: 2.84\n"
     ]
    }
   ],
   "source": [
    "run_pend(r0=np.pi/100.,L=2.0)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "### Usage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}