Correct implementation of SI, SIS, SIR models (python)

Question

I have created some very basic implementations of the mentioned models. However, although graphs seem to look right, the numbers don't add up to a constant. That is for the sum of susceptible/infected/recovered people in each compartment should add up to N (which is total number of people), but it doesn't, for some reason it adds up to some bizarre decimal numbers, and I really don't know how to fix it, after looking at it for 3 days now.

The SI Model:

import matplotlib.pyplot as plt

N = 1000000
S = N - 1
I = 1
beta = 0.6

sus = [] # infected compartment
inf = [] # susceptible compartment
prob = [] # probability of infection at time t

def infection(S, I, N):
    t = 0
    while (t < 100):
        S = S - beta * ((S * I / N))
        I = I + beta * ((S * I) / N)
        p = beta * (I / N)

        sus.append(S)
        inf.append(I)
        prob.append(p)
        t = t + 1

infection(S, I, N)
figure = plt.figure()
figure.canvas.set_window_title('SI model')

figure.add_subplot(211)
inf_line, =plt.plot(inf, label='I(t)')

sus_line, = plt.plot(sus, label='S(t)')
plt.legend(handles=[inf_line, sus_line])

plt.ticklabel_format(style='sci', axis='y', scilimits=(0,0)) # use scientific notation

ax = figure.add_subplot(212)
prob_line = plt.plot(prob, label='p(t)')
plt.legend(handles=prob_line)

type(ax)  # matplotlib.axes._subplots.AxesSubplot

# manipulate
vals = ax.get_yticks()
ax.set_yticklabels(['{:3.2f}%'.format(x*100) for x in vals])

plt.xlabel('T')
plt.ylabel('p')

plt.show()

SIS Model:

import matplotlib.pylab as plt

N = 1000000
S = N - 1
I = 1
beta = 0.3
gamma = 0.1

sus = \[\]
inf = \[\]

def infection(S, I, N):
    for t in range (0, 1000):
        S = S - (beta*S*I/N) + gamma * I
        I = I + (beta*S*I/N) - gamma * I

        sus.append(S)
        inf.append(I)


infection(S, I, N)

figure = plt.figure()
figure.canvas.set_window_title('SIS model')

inf_line, =plt.plot(inf, label='I(t)')

sus_line, = plt.plot(sus, label='S(t)')
plt.legend(handles=\[inf_line, sus_line\])

plt.ticklabel_format(style='sci', axis='y', scilimits=(0,0))

plt.xlabel('T')
plt.ylabel('N')

plt.show()

SIR Model:

import matplotlib.pylab as plt

N = 1000000
S = N - 1
I = 1
R = 0
beta = 0.5
mu = 0.1

sus = []
inf = []
rec = []

def infection(S, I, R, N):
    for t in range (1, 100):
        S = S -(beta * S * I)/N
        I = I + ((beta * S * I)/N) - R
        R = mu * I

        sus.append(S)
        inf.append(I)
        rec.append(R)

infection(S, I, R, N)

figure = plt.figure()
figure.canvas.set_window_title('SIR model')

inf_line, =plt.plot(inf, label='I(t)')

sus_line, = plt.plot(sus, label='S(t)')

rec_line, = plt.plot(rec, label='R(t)')
plt.legend(handles=[inf_line, sus_line, rec_line])

plt.ticklabel_format(style='sci', axis='y', scilimits=(0,0))

plt.xlabel('T')
plt.ylabel('N')


plt.show()

Answer 1

I'll look only at the SI model.

Your two key variables are S and I. (You may have reversed the meanings of these two variables, though that does not affect what I write here.) You initialize them so their sum is N which is the constant 1000000.

You update your two key variables in the lines

S = S - beta * ((S * I / N))
I = I + beta * ((S * I) / N)

You apparently intend to add to I and subtract from S the same value, so the sum of S and I is unchanged. However, you actually first change S then use that new value to change I, so the values added and subtracted are not actually the same, and the sum of the variables has not remained constant.

You can fix this by using Python's ability to update multiple variables in one line. Replace those two lines with

S, I = S - beta * ((S * I / N)), I + beta * ((S * I) / N)

This calculates both of the new values before updating the variables, so the same value actually added and subtracted from the two variables. (There are other ways to get the same effect, such as temporary variables for the updated values, or one temporary variable to store the amount to add and subtract, but since you use Python you may as well use its capabilities.)

When I now run the program, I get these graphs:

which I think is what you want.

Answer 2

So the solution above worked for the SIS model as well.

As for the SIR model I had to solve differential equations using odeint, here is a simple solution to the SIR model:

import matplotlib.pylab as plt
from scipy.integrate import odeint
import numpy as np

N = 1000
S = N - 1
I = 1
R = 0
beta = 0.6 # infection rate
gamma = 0.2 # recovery rate

# differential equatinons
def diff(sir, t):
    # sir[0] - S, sir[1] - I, sir[2] - R
    dsdt = - (beta * sir[0] * sir[1])/N
    didt = (beta * sir[0] * sir[1])/N - gamma * sir[1]
    drdt = gamma * sir[1]
    print (dsdt + didt + drdt)
    dsirdt = [dsdt, didt, drdt]
    return dsirdt


# initial conditions
sir0 = (S, I, R)

# time points
t = np.linspace(0, 100)

# solve ODE
# the parameters are, the equations, initial conditions, 
# and time steps (between 0 and 100)
sir = odeint(diff, sir0, t)

plt.plot(t, sir[:, 0], label='S(t)')
plt.plot(t, sir[:, 1], label='I(t)')
plt.plot(t, sir[:, 2], label='R(t)')

plt.legend()

plt.xlabel('T')
plt.ylabel('N')

# use scientific notation
plt.ticklabel_format(style='sci', axis='y', scilimits=(0,0))

plt.show()