PGF simulation

Setup

import numpy as np
import matplotlib.pyplot as plt
from tqdm import trange
import pandas as pd
from collections import Counter

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()))

Parameters

# SIR model setup
beta = 0.2
mu = 0.1

# Stochastic simulation
S_initial=1e4
I_initial=20
n_steps=200
n_sim=10000

N_pgf = 3000 # Support for PGF calculation (approximation)

beta/mu

2.0

Probability Generating Function for SIR simulation

From Guillaume's slide (p110):

For SIR in a well-mixed population, we can use the following approximations in the early phase:

$$\textrm{Bin}(S_t, \lambda_t) \simeq \textrm{Poisson}(S_t \lambda_t) \simeq \textrm{Poisson}(\beta I_t) = \sum_{i=1}^{I_t}\textrm{Poisson}(\beta)$$ $$\textrm{Bin}(I_t, \mu) = \sum_{i=1}^{I_t}\textrm{Bern}(\mu)$$

This approximation allows us to treat the SIR model as a branching process.

Offspring Distribution:
$Q(x)$ is the probability generating function (PGF) for the offspring distribution, i.e., the number of new infections produced by a single infected individual.

$$Q(x) = \big[ \mu + (1-\mu) x \big] \exp \big[ \beta (x-1) \big]$$

Total Infections $I_t$:
The total number of infections at time $t$, $I_t$, is obtained by the cumulative effect of these offspring over successive generations. This can be solved iteratively.

If we set the initial PGF as

$$G_0(x) = x^{I_{\text{initial}}},$$

then the PGF for $I_t$ is given by the $t$-fold composition:

$$G_t(x) = Q^{(t)}(x) \equiv Q\big(Q(\cdots Q(x) \cdots )\big)$$

# PGF
Q = lambda x: ( mu + ((1-mu) *x) ) * np.exp(beta*(x-1))
G0 = lambda x: x**I_initial

# We use numerical inversion to get probability distribution from PGF
# https://cosmo-notes.github.io/pgfunk/chapters/numerical_inversion.html

n = np.arange(N_pgf)
c = np.exp(2*np.pi*1j*n/N_pgf)

# Iteratively solve for I_t using the offspring distribution
# G_{t+1}(x) = Q(G_t(x))
tset = range(0,n_steps+1)
x = c.copy()
pn_dict_pgf = {}
for t in range(max(tset)+1):
    if t in tset:
        pn = abs(np.fft.fft(G0(x))/N_pgf)
        pn_dict_pgf[t] = dict(zip(n, pn))
    x = Q(x)

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

for t in [1,3,5,10]:
    ax.plot(
        list(pn_dict_pgf[t].keys())[:100],
        list(pn_dict_pgf[t].values())[:100],
        "o-",
        markersize=3,
        label=fr"t={t}"
    )

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

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

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

show_fig_svg(fig)

# Calculate mean I_t for each timestep from probability distribution
# We'll compare with mean from stochastic simulation
It_mean_pgf = {}
for t, pn_dict in pn_dict_pgf.items():
    It_mean_pgf[t] = np.sum(
        np.prod([list(pn_dict.keys()), list(pn_dict.values())], axis=0)
    )