Implementations

Let's try implementing the SIR model in Python. Specifically, we will code

deterministic model with homogeneous mixing assumption
stochastic model with homogeneous mixing assumption

and observe that the two models provide approximately same result in average.

Notebook setup

import numpy as np
import pandas as pd
from matplotlib import pyplot as plt
import networkx as nx
from tqdm import tqdm, trange
import cmocean

%config InlineBackend.figure_format = 'retina'
plt.rcParams['text.usetex'] = False 
plt.rcParams['font.family'] = 'sans-serif'
plt.rcParams['font.sans-serif'] = 'Helvetica'

def show_fig_svg(fig):
    import io
    from IPython.display import SVG, display

    # Save the figure into a string buffer as SVG
    buf = io.StringIO()
    fig.savefig(buf, format='svg')
    buf.seek(0)

    # Display the SVG image
    display(SVG(buf.getvalue()))

Paremeters

Here are the parameters that we use throughout the simulations.

params = {
    "N": 50000, # Total population size
    "n_steps": 1000, # Number of time steps
    "beta": 0.2, # Infection rate
    "gamma": 0.1, # Recovery rate
    "n_stoc_sim": 20, # Number of stochastic simulations
}
params["I_initial"] = 20 # Initial number of infected individuals
params["R_initial"] = 0 # Initial number of recovered individuals
params["S_initial"] = params["N"] - (params["I_initial"] + params["R_initial"])
params

{'N': 50000,
 'n_steps': 1000,
 'beta': 0.2,
 'gamma': 0.1,
 'n_stoc_sim': 20,
 'I_initial': 20,
 'R_initial': 0,
 'S_initial': 49980}

Deterministic model

def SIR_model(S_initial=1000, I_initial=1, R_initial=0, n_steps=1000, beta=0.8, gamma=0.4):
    """
    Simulate SIR model.

    Parameters
    ----------
    S_initial : int
        Initial number of Susceptible individuals.
    I_initial : int
        Initial number of Infected individuals.
    R_initial : int
        Initial number of Recovered individuals.
    n_steps : int
        Number of time steps to simulate.
    beta : float
        Infection rate (probability of infection on contact)
    gamma : float
        Recovery rate (probability of recovery per unit time)
    
    Returns
    -------
    S_dict : dict
        Dictionary with number of susceptible individuals at each time step.
        {t: S_t, ... }
    I_dict : dict
        Dictionary with number of infected individuals at each time step.
        {t: I_t, ... }
    R_dict : dict
        Dictionary with number of recovered individuals at each time step.
        {t: R_t, ... }
    """

    # Set S,I,R to initial values
    S = S_initial
    I = I_initial
    R = R_initial

    # N is always constant
    N = S+I+R

    # Initialize dictionaries to store values
    # {t: S_t, ... } etc.
    # t=0 is the initial state
    S_dict = {0: S}
    I_dict = {0: I}
    R_dict = {0: R}

    # Start simulation from t=1 to t=n_steps
    for t in range(1, n_steps+1):

        # Calculate number of individuals for each compartment at next time step
        S_new = S - S*beta*(I/N)
        I_new = I + S*beta*(I/N) - gamma*I
        R_new = R + gamma*I

        # Store values
        S_dict[t] = S_new
        I_dict[t] = I_new
        R_dict[t] = R_new

        # Update values for next time step
        S = S_new
        I = I_new
        R = R_new

    return S_dict, I_dict, R_dict

S_determ, I_determ, R_determ = SIR_model(
    S_initial=params["S_initial"],
    I_initial=params["I_initial"],
    R_initial=params["R_initial"],
    n_steps=params["n_steps"],
    beta=params["beta"],
    gamma=params["gamma"]
)

fig, ax = plt.subplots(1,1, figsize=(4,3))

# Color mapping
cmap = cmocean.cm.haline
colors = [cmap(0.3), cmap(0.5), cmap(0.75)]

compartments_to_plot = ["S", "I", "R"]

data_determ = {
    "S": S_determ,
    "I": I_determ,
    "R": R_determ
}

for idx, compartment in enumerate(compartments_to_plot):
    ax.plot(
        data_determ[compartment].keys(),
        data_determ[compartment].values(),
        color=colors[idx],
        label=f'{compartment}'
    )
ax.set_title("SIR model (deterministic)", loc='left')
ax.set_xlabel("Timestep")
ax.set_ylabel("Number of individuals")

ax.set_xlim(0,200)

ax.legend()

fig.tight_layout()
# plt.savefig('fig/fig_name.pdf', bbox_inches='tight')
plt.close()

Stochastic model

def SIR_model_stochastic(S_initial=1000, I_initial=1, R_initial=0, n_steps=1000, beta=0.8, gamma=0.4):
    """
    Simulate SIR model.

    Parameters
    ----------
    S_initial : int
        Initial number of Susceptible individuals.
    I_initial : int
        Initial number of Infected individuals.
    R_initial : int
        Initial number of Recovered individuals.
    n_steps : int
        Number of time steps to simulate.
    beta : float
        Infection rate (probability of infection on contact)
    gamma : float
        Recovery rate (probability of recovery per unit time)
    
    Returns
    -------
    S_dict : dict
        Dictionary with number of susceptible individuals at each time step.
        {t: S_t, ... }
    I_dict : dict
        Dictionary with number of infected individuals at each time step.
        {t: I_t, ... }
    R_dict : dict
        Dictionary with number of recovered individuals at each time step.
        {t: R_t, ... }
    """

    # Set S,I,R to initial values
    S = S_initial
    I = I_initial
    R = R_initial

    # N is always constant
    N = S+I+R

    # Initialize dictionaries to store values
    # {t: S_t, ... } etc.
    # t=0 is the initial state
    S_dict = {0: S}
    I_dict = {0: I}
    R_dict = {0: R}

    # Start simulation from t=1 to t=n_steps
    for t in range(1, n_steps+1):

        # Calculate number of individuals for each compartment at next time step
        lambda_t = beta*(I/N)

        S_diff = np.random.binomial(S, lambda_t)
        I_diff = np.random.binomial(I, gamma)

        S_new = S - S_diff
        I_new = I + S_diff - I_diff
        R_new = R + I_diff

        # Store values
        S_dict[t] = S_new
        I_dict[t] = I_new
        R_dict[t] = R_new

        # Update values for next time step
        S = S_new
        I = I_new
        R = R_new

    return S_dict, I_dict, R_dict

# This is a single determination
S_stoc, I_stoc, R_stoc = SIR_model_stochastic(
    S_initial=params["N"],
    I_initial=params["I_initial"],
    R_initial=params["R_initial"],
    n_steps=params["n_steps"],
    beta=params["beta"],
    gamma=params["gamma"]
)

# We'll repeat the stochastic simulation for multiple times.

# Each determination will be added to the array
S_stoc_dict, I_stoc_dict, R_stoc_dict = [], [], []

for _ in trange(params["n_stoc_sim"]):
    S_stoc, I_stoc, R_stoc = SIR_model_stochastic(
        S_initial=params["S_initial"],
        I_initial=params["I_initial"],
        R_initial=params["R_initial"],
        n_steps=params["n_steps"],
        beta=params["beta"],
        gamma=params["gamma"]
    )

    S_stoc_dict.append(S_stoc)
    I_stoc_dict.append(I_stoc)
    R_stoc_dict.append(R_stoc)

100%|██████████| 20/20 [00:00<00:00, 1020.72it/s]

# Reorganize data into a list of dictionaries, so that it can be converted into a dataframe.
data_stoc_sim = []
for id, _ in enumerate(S_stoc_dict):
    for t, _ in S_stoc_dict[id].items():
        data_stoc_sim.append(
            {
                'iter': id,
                'time': t,
                'S': S_stoc_dict[id][t],
                'I': I_stoc_dict[id][t],
                'R': R_stoc_dict[id][t]
            }
        )

df_stochastic = pd.DataFrame(data_stoc_sim)
df_stochastic

	iter	time	S	I	R
0	0	0	49980	20	0
1	0	1	49979	19	2
2	0	2	49970	27	3
3	0	3	49967	28	5
4	0	4	49963	31	6
...	...	...	...	...	...
20015	19	996	9878	0	40122
20016	19	997	9878	0	40122
20017	19	998	9878	0	40122
20018	19	999	9878	0	40122
20019	19	1000	9878	0	40122

20020 rows × 5 columns

# Because we are dealing with multiple realizations, we'll calculate the interquantile range (25-percentile and 75-percentile)
def q25(x):
    return x.quantile(0.25)

def q75(x):
    return x.quantile(0.75)

# Groupby "time" and get the median and IQR
df_stochastic.groupby('time')['S'].agg(["median", q25, q75])

	median	q25	q75
time
0	49980.0	49980.00	49980.0
1	49977.0	49975.00	49978.0
2	49973.0	49970.00	49975.0
3	49968.0	49964.75	49971.0
4	49963.0	49958.75	49966.0
...	...	...	...
996	9851.0	9688.75	9998.5
997	9851.0	9688.75	9998.5
998	9851.0	9688.75	9998.5
999	9851.0	9688.75	9998.5
1000	9851.0	9688.75	9998.5

1001 rows × 3 columns

fig, ax = plt.subplots(1,1, figsize=(4,3))

compartments_to_plot = ["S", "I", "R"]

for idx, compartment in enumerate(compartments_to_plot):

    quantiles = df_stochastic.groupby('time')[compartment].agg(["median", q25, q75])

    # Median
    ax.plot(
        quantiles.index,
        quantiles['median'],
        color=colors[idx],
        label=f'{compartment}'
    )

    # IQR
    ax.fill_between(
        quantiles.index,
        quantiles['q25'],
        quantiles['q75'],
        color=colors[idx],
        alpha=0.2,
    )

ax.set_title("SIR model (stochastic)", loc='left')
ax.set_xlabel("Timestep")
ax.set_ylabel("Number of individuals")

ax.set_xlim(0,200)

ax.legend()


fig.tight_layout()
# plt.savefig('fig/fig_name.pdf', bbox_inches='tight')
plt.close()

Comparison

fig, ax = plt.subplots(1,1, figsize=(4,3))

compartments_to_plot = ["S", "I", "R"]

# === Stochastic simulation ===
for idx, compartment in enumerate(compartments_to_plot):

    quantiles = df_stochastic.groupby('time')[compartment].agg(["median", q25, q75])

    # Median
    ax.plot(
        quantiles.index,
        quantiles['median'],
        color=colors[idx],
        label=f'{compartment} (stochastic)'
    )

    # IQR
    ax.fill_between(
        quantiles.index,
        quantiles['q25'],
        quantiles['q75'],
        color=colors[idx],
        alpha=0.2,
    )

# === Deterministic result ===
for idx, compartment in enumerate(compartments_to_plot):
    
    ax.plot(
        data_determ[compartment].keys(),
        data_determ[compartment].values(),
        linestyle=(0,(1,2)),
        color="#2e2e2e",
    )

ax.set_title("SIR model (comparison)", loc='left')
ax.set_xlabel("Timestep")
ax.set_ylabel("Number of individuals")

ax.set_xlim(0,200)


ax.legend()

fig.tight_layout()
# plt.savefig('fig/fig_name.pdf', bbox_inches='tight')
plt.close()

Notebook setup​

Paremeters​

Deterministic model​

Stochastic model​

Comparison​

Notebook setup

Paremeters

Deterministic model

Stochastic model

Comparison