epidemicSoluce.py

# -------------------------------------------------------------------------
#
# PYTHON for DUMMIES 
#
# Solution detaillée
#  Vincent Legat
#
# -------------------------------------------------------------------------

from numpy import *

# -------------------------------------------------------------------------
#
# -1- Equations SIR 
#
#           dSdt(t) = - beta*S(i)*I(t)
#           dIdt(t) =   beta*S(i)*I(t) - gamma*I(t)
#           dRdt(t) =                  + gamma*I(t)
# 
#

def f(u,beta,gamma):
  dSdt = - beta * u[0] * u[1]
  dIdt =   beta * u[0] * u[1] - gamma * u[1]
  dRdt =                        gamma * u[1]  
  return array([dSdt,dIdt,dRdt])
 
 
# -------------------------------------------------------------------------    
#
# -1- Schema de d'Euler explicite d'ordre 1
# 
  
def epidemicEuler(Xstart,Xend,Ustart,n,beta,gamma):
  X = linspace(Xstart,Xend,n+1)
  U = zeros((n+1,3)); U[0,:] = Ustart
  h = (Xend - Xstart)/n
  for i in range(n):  
    U[i+1,:] = U[i,:] + h*f(U[i,:],beta,gamma)    
  return X,U
 

# -------------------------------------------------------------------------    
#
# -2- Schema de Runge-Kutta classique d'ordre 4
# 
  
def epidemicRungeKutta(Xstart,Xend,Ustart,n,beta,gamma):
  X = linspace(Xstart,Xend,n+1)
  U = zeros((n+1,3)); U[0,:] = Ustart
  h = (Xend - Xstart)/n
  for i in range(n):  
    K1 = f(U[i,:]       ,beta,gamma)
    K2 = f(U[i,:]+K1*h/2,beta,gamma)
    K3 = f(U[i,:]+K2*h/2,beta,gamma)
    K4 = f(U[i,:]+K3*h  ,beta,gamma)
    U[i+1,:] = U[i,:] + h*(K1+2*K2+2*K3+K4)/6     
  return X,U

# -------------------------------------------------------------------------    
#
# -3- Schema de Taylor d'ordre 4
# 
  
def epidemicTaylor(Xstart,Xend,Ustart,n,beta,gamma):
  X = linspace(Xstart,Xend,n+1)
  U = zeros((n+1,3)); U[0,:] = Ustart
  h = (Xend - Xstart)/n
  for i in range(n): 
    u = U[i,:]
    F    = array([-beta * u[0]*u[1],
                   beta * u[0]*u[1] - gamma * u[1],
                   gamma * u[1] ])
    dF   = array([-beta * (F[0]*u[1] + u[0]*F[1]),
                   beta * (F[0]*u[1] + u[0]*F[1]) - gamma * F[1],
                   gamma * F[1] ])
    ddF  = array([-beta * (dF[0]*u[1] + 2*F[1]*F[0] + u[0]*dF[1]),
                   beta * (dF[0]*u[1] + 2*F[1]*F[0] + u[0]*dF[1]) - gamma * dF[1],
                   gamma * dF[1] ])
    dddF = array([-beta * (ddF[0]*u[1] + 3*dF[1]*F[0] + 3*F[1]*dF[0] + u[0]*ddF[1]),
                   beta * (ddF[0]*u[1] + 3*dF[1]*F[0] + 3*F[1]*dF[0] + u[0]*ddF[1]) - gamma * ddF[1],
                   gamma * ddF[1] ])
    U[i+1,:] = U[i,:] + h*F + h**2*dF/2 + h**3*ddF/6 + h**4*dddF/24
  return X,U

# -------------------------------------------------------------------------