Stochastic simulation

Stochastic SIR

def SIR_model_stochastic(S_initial=1000, I_initial=1, R_initial=0, n_steps=1000, beta=0.4, 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 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, ... }
    df_data : list
        List of dictionary recording number of S, I, R at each timestep. This can be converted into a Pandas dataframe by pd.DataFrame(df_data).
        [ {'t': t, 'S': S_new, 'I': I_new, 'R': R_new}, ... ]
    """

    # 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}
    df_data = [{'t': 0, 'S': S, 'I': I, 'R': R, 'lambda': beta*(I/N)}]

    # 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)
        # lambda_t = beta*(I)

        # S_diff = np.random.poisson(lambda_t)
        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
        
        # Data for pandas
        df_data.append({'t': t, 'S': S_new, 'I': I_new, 'R': R_new, 'lambda': lambda_t})
        
        # Update values for next time step
        S = S_new
        I = I_new
        R = R_new

    return S_dict, I_dict, R_dict, df_data

It_2darray_sim = []
sim_data = []

for it in trange(n_sim):

    # Run simulation
    _, I, _, df_list = SIR_model_stochastic(
        S_initial=S_initial,
        I_initial=I_initial,
        R_initial=0,
        n_steps=n_steps,
        beta=beta,
        gamma=mu
    )
    It_2darray_sim.append(list(I.values()))

    # For dataframe
    for item in df_list:
        item['iter'] = it
    sim_data.extend(df_list)

It_2darray_sim = np.array(It_2darray_sim).transpose()

100%|██████████| 10000/10000 [00:02<00:00, 3371.06it/s]

# Convert data to dataframe for summarization
df_sim = pd.DataFrame(sim_data)

def q25(x):
    return x.quantile(0.25)

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

It_summary = df_sim.groupby('t')['I'].agg(['median', 'mean', q25, q75])
It_summary

	median	mean	q25	q75
t
0	20.0	20.0000	20.0	20.0
1	22.0	22.0017	20.0	24.0
2	24.0	24.2170	22.0	27.0
3	26.0	26.6029	23.0	30.0
4	29.0	29.2637	25.0	33.0
...	...	...	...	...
196	1.0	1.1972	0.0	2.0
197	0.0	1.1243	0.0	2.0
198	0.0	1.0563	0.0	2.0
199	0.0	0.9973	0.0	2.0
200	0.0	0.9395	0.0	1.0

201 rows × 4 columns

lambda_summary = df_sim.groupby('t')['lambda'].agg(['median', 'mean', q25, q75])

# S_summary = df_sim.groupby('t')['S'].agg(['median', 'mean', q25, q75])

St_summary = df_sim.groupby('t')['S'].agg(['median', 'mean', q25, q75])

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

# ax.plot(St_summary.index, St_summary['median'], label="S (Median)")
# ax.fill_between(St_summary.index, St_summary['q25'], St_summary['q75'], alpha=0.2)

ax.plot(It_summary.index, It_summary['median'], label="Median")
ax.fill_between(It_summary.index, It_summary['q25'], It_summary['q75'], alpha=0.2, label="IQR")

# ax.plot(Rt_summary.index, Rt_summary['median'], label="R (Median)")
# ax.fill_between(Rt_summary.index, Rt_summary['q25'], Rt_summary['q75'], alpha=0.2)


ax.set_title("", loc='left')
ax.set_xlabel("Timestep")
ax.set_ylabel("$I_t$")

ax.legend(loc='upper right', frameon=False)

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

show_fig_svg(fig)

# Covert 2d array to dictionary by time
It_dict_sim = {t: It_2darray_sim[t] for t in range(len(It_2darray_sim))}

# Calculate probability for each t
pn_dict_sim = {}
for t, data in It_dict_sim.items():
    p_n = {int(It): It_count/len(data) for It, It_count in dict(Counter(It_dict_sim[t])).items() } # Count each I and calculate probability
    p_n = {k: p_n[k] for k in sorted(p_n)} # Sort by key
    pn_dict_sim[t] = p_n

Stochastic SIR​

Stochastic SIR