{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "## Physics 503: Scientific Computing\n",
    "### Lecture 16: Ordinary Differential Equations 3 - Systems of ODEs\n",
    "### J.R. Gladden, University of Mississippi"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "### Coupled ODEs\n",
    "Often physical systems have multiple variables the depend on each other. Each variable can have it's own ODE which depends on other variables.  A clasic example is from **Epidemiology** - the study of how a disease propagates through a population.  The variables here are the number of infected people $I$ and the number of susceptible people $S$. A simple model which describes such a system is:\n",
    "\n",
    "$$ \\frac{dS}{dt} = - r S I  $$\n",
    "$$\\frac{dI}{dt} = r S I -a I $$\n",
    "where $r$ is a measure of how infectious the disease is and $a$ is a measure of how quick the recovery period is.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "#### Imports "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "%matplotlib wx\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true,
    "deletable": true,
    "editable": true
   },
   "source": [
    "#### Problem setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "def epidemic(s0,i0,T,n,r,a):\n",
    "    dt=T/float(n)\n",
    "    t=np.linspace(0,T,n+1)\n",
    "    s=np.zeros(n+1)\n",
    "    i=np.zeros(n+1)\n",
    "    s[0]=s0\n",
    "    i[0]=i0\n",
    "    for k in range(n):\n",
    "        s[k+1]=s[k] - dt*r*s[k]*i[k]\n",
    "        i[k+1]=i[k] + dt*(r*s[k]*i[k] - a*i[k])\n",
    "    return t,s,i\n",
    "    \n",
    " #Model parameters\n",
    "T=10\n",
    "n=500\n",
    "s0=200\n",
    "i0=1\n",
    "r=0.03\n",
    "a=0.6\n",
    "\n",
    "t,s,i = epidemic(s0,i0,T,n,r,a)\n",
    "\n",
    "plt.plot(t,s,'b-',lw=2,label='Suseptibles')\n",
    "plt.plot(t,i,'g-',lw=2,label='Infectives')\n",
    "plt.xlabel('Time')\n",
    "plt.ylabel('Population')\n",
    "plt.title('Epidemic model \\n r=%g and a=%g'%(r,a))\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "#### Lotka - Volterra Predator-Prey model\n",
    "This is a classic model from biology which describes the populations of a ecosystem of predators ($y$) and prey ($x$) (say wolves and rabbits).  The equations are:\n",
    "$$ \\frac{dx}{dt} = x(a- by - cx)$$\n",
    "$$ \\frac{dy}{dt} = -y (d - e x) $$\n",
    "where the parameters are a measure of:\n",
    "- a: prey birth rate\n",
    "- b: prey death rate (not being eaten)\n",
    "- c: consumption rate\n",
    "- d: predator death rate\n",
    "- e: predator birth rate"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "### Setup Solvers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "def euler(x,y,tau):\n",
    "    for step in range(maxstep-1):\n",
    "        xc=x[step]\n",
    "        yc=y[step]\n",
    "        t=step*tau\n",
    "        times[step]=t\n",
    "        \n",
    "        #Compute accel. vector at each time step\n",
    "        fx=(a-b*xc-c*yc)*xc\n",
    "        fy=(-d+e*xc)*yc\n",
    "        \n",
    "        #Compute r and v for next time step\n",
    "        x[step+1]=xc+tau*fx\n",
    "        y[step+1]=yc+tau*fy\n",
    "    return x,y\n",
    "\n",
    "def leapfrog(x,y,tau):\n",
    "    for step in range(maxstep-1):\n",
    "        xc=x[step]\n",
    "        yc=y[step]\n",
    "        t=step*tau\n",
    "        times[step]=t\n",
    "        \n",
    "        # Compute RHS at each time step\n",
    "        fy=(-d+e*xc)*yc\n",
    "\n",
    "        #Compute y and x for next time step    \n",
    "        y[step+1]=yc+tau*fy\n",
    "        \n",
    "        # Use updated y to compute new x (ala Euler-Cromer)\n",
    "        fx=(a-b*xc-c*y[step+1])*xc\n",
    "        x[step+1]=xc+tau*fx\n",
    "\n",
    "    return x,y\n",
    "    \n",
    "    \n",
    "plt.close('all')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "#### Run the solvers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "#parameters\n",
    "\n",
    "a=10        #prey birth rate\n",
    "b=0.01    #prey death rate (other than being eaten)\n",
    "c=0.4       #rate of consumption\n",
    "d=12       #pred. death rate\n",
    "e=0.4       #pred. growth rate\n",
    "\n",
    "#Initial conditions\n",
    "#Y is predator population and x is prey population\n",
    "\n",
    "xsteady=d/e\n",
    "ysteady=a/c - b*d/(c*e)\n",
    "\n",
    "#yi=ysteady*1.8\n",
    "#xi=xsteady\n",
    "xi=20\n",
    "yi=5\n",
    "#Simulation parameters\n",
    "tau=0.001\n",
    "Tmax = 25 \n",
    "maxstep=int(Tmax/tau)\n",
    "\n",
    "#Create container 2-D arrays for position, velocity and theory position\n",
    "x=np.zeros(maxstep,dtype=float)\n",
    "y=np.zeros(maxstep,dtype=float)\n",
    "\n",
    "times=np.zeros(maxstep,dtype=float)\n",
    "x[0]=xi\n",
    "y[0]=yi\n",
    "\n",
    "x,y=euler(x,y,tau)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "### Visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "times[-1]=maxstep*tau\n",
    "plt.plot(times,x,'b-')\n",
    "plt.plot(times,y,'g-')\n",
    "plt.hlines(xsteady,0,max(times)*1.1,linestyles='dashed',color='b')\n",
    "plt.hlines(ysteady,0,max(times)*1.1,linestyles='dashed',color='g')\n",
    "\n",
    "plt.xlabel(\"Time\")\n",
    "plt.ylabel(\"Populations\")\n",
    "plt.legend((\"Rabbits\", \"Wolves\"))\n",
    "plt.title(\"Lotka-Volterra Predator-Prey Model\")\n",
    "\n",
    "plt.figure(2)\n",
    "plt.plot(x,y,'r-')\n",
    "plt.plot([xi],[yi],'go')\n",
    "plt.plot([xsteady],[ysteady],'bo')\n",
    "plt.legend((\"Trajectory\",\"Starting Point\",\"Fixed Point\"))\n",
    "plt.title(\"Predator vs Prey\")\n",
    "plt.xlabel(\"Prey\")\n",
    "plt.ylabel(\"Predator\")\n",
    "plt.show()\n"
   ]
  },
  {
   "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
}