Search code examples
pythonrunge-kutta

Fourth Order Runge Kutta for SIR model


I am implementing a fourth order Runge Kutta method to solve a system of three ODEs that describe the SIR model. Based on other simulations, I would expect a different solution (mine diverges very quickly). Here is my code for the runge kutta method:

def rk4(f, s0, x0, xf, n):
  x = np.linspace(x0, xf, n+1) #x grid
  s = np.array((n+1)*[s0]) #array of the state of the sytem for each x
  h = x[1] - x[0] #stepsize
  for i in range(n):       #Fourth Order Runge Kutta Method
    k0 = h * f(s[i]) 
    k1= h * f(s[i] + 0.5 * k0)
    k2 = h * f(s[i] + 0.5 * k1)
    k3 = h * f(s[i] + k2)
    s[i+1] = s[i] + (k0 + 2*(k1 + k2) + k3) / 6.
  return x, s

And here is the implementation for the SIR model:

def f(u):
  dS = -a*u[0]*u[1] #dY/dt = -aS(t)I(t)
  dI = a*u[0]*u[1] - b*u[1]  #dI/dt = aS(t)I(t)-bI(t)
  dR = b*u[1]  #dR/dt = bI(t)
  return np.array([dS, dI, dR])

for:

a = 0.2
b = 0.1
x, s = rk4(f, np.array([235.0, 14.0, 0.0]), 0., 109., 109000)

the solution I get is this although I expected to approach these data. What am I doing wrong?

Thank you in advance!


Solution

  • With the given parameters a=1/5 and b=1/10 and the general context of infection models, it is reasonable to guess that the intention is that the model has one transmission per infected person every 5 days, and a resolution of every infection in 10 days on average.

    The model to implement this via population numbers, whatever unit they may have, is

    def f(u):
      S,I,R = u
      N = sum(u)
      dS = -a*S*I/N 
      dI = a*S*I/N - b*I  
      dR = b*I
      return np.array([dS, dI, dR])
    

    Only if the variables are the population densities the sum N will always be one and can be omitted.

    With this change the code as given produces the image

    enter image description here

    which looks very similar to the reference image