Intrinsic periodicity of the SIR system¶

import math

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from laser.core.propertyset import PropertySet

import laser.core
import laser.generic
import laser.core.distributions as dists
from laser.core.demographics import AliasedDistribution
from laser.core.demographics import KaplanMeierEstimator
from laser.generic import SIR
from laser.generic import Model
from laser.generic.utils import ValuesMap
from laser.generic.vitaldynamics import BirthsByCBR, MortalityByEstimator
from laser.core.utils import grid

from scipy.ndimage import gaussian_filter1d
from scipy.signal import find_peaks


print(f"{np.__version__=}")
print(f"{laser.core.__version__=}")
print(f"{laser.generic.__version__=}")

model = Model(scenario, parameters, birthrates=birthrate_map)

model.components = [
    SIR.Susceptible(model),
    SIR.Recovered(model),
    SIR.Infectious(model, infdurdist),
    Importation(model, infdurdist),
    SIR.Transmission(model, infdurdist),
    BirthsByCBR(model, birthrate_map, pyramid),
    MortalityByEstimator(model, estimator=survival),
]

model.run()

plt.plot(model.nodes.S, color="blue")
plt.plot(model.nodes.I, color="red")
plt.plot(model.nodes.R, color="green")
plt.plot(model.nodes.S+model.nodes.I+model.nodes.R, color="black")
plt.legend(["S", "I", "R", "N"])
plt.show()

mu = (1 + model.params.cbr / 1000) ** (1 / 365) - 1
R0 = model.params.beta / (1 / model.params.inf_mean + mu)
A = 1 / ((R0 - 1) * mu) / 365
G = 1 / (mu + 1 / model.params.inf_mean) / 365
T_exp = 2 * np.pi * np.sqrt(A * G)


def ID_freq_peakfinder(y0, T_exp, cutoff=18250, plot=False):
    y = y0[cutoff:]
    y = y - np.mean(y)
    y = gaussian_filter1d(y, sigma=50)
    peaks, _ = find_peaks(y, distance=T_exp * 365 / 2)
    peaks2, _ = find_peaks(-1*y, distance=T_exp * 365 / 2)
    if plot:
        plt.figure()
        plt.plot(y, 'b-', lw=1)
        plt.plot(peaks, y[peaks], "rx")
        plt.plot(peaks2, y[peaks2], "kx")
    return np.median(np.diff(peaks)) / 365, np.median(np.diff(peaks2)) / 365


T_obs_pf, T_obs_pf2 = ID_freq_peakfinder(np.squeeze(model.nodes.I), T_exp, plot=True)


def ID_freq_autocorr(y0, cutoff=18250):
    # Compute the FFT
    Y1 = np.fft.fft(y0[cutoff:] - np.mean(y0[cutoff:]))

    # Compute the circular autocorrelation using the inverse FFT
    circular_autocorr = np.fft.ifft(Y1 * np.conj(Y1)).real
    # Plot only the positive frequency spectrum
    peaks, _ = find_peaks(circular_autocorr, distance=300)
    return peaks[0] / 365


T_obs_fft = ID_freq_autocorr(np.squeeze(model.nodes.I))

plt.text(0.5, 0.99, f"T expected: {T_exp:.2f} y", transform=plt.gca().transAxes)
plt.text(0.5, 0.95, f"T observed, peakfinding: {T_obs_pf:.2f} y", transform=plt.gca().transAxes)
plt.text(0.5, 0.91, f"T observed, valley-finding : {T_obs_pf:.2f} y", transform=plt.gca().transAxes)
plt.text(0.5, 0.87, f"T observed, FFT: {T_obs_fft:.2f} y", transform=plt.gca().transAxes)
plt.gca().spines['top'].set_visible(False)
plt.show()

params_df

for i in range(nsims):
    row = params_df.iloc[i]
    scenario = grid(M=1, N=1, population_fn=lambda x,y: pop, origin_x=0, origin_y=0)
    parameters = PropertySet(
        {"seed": 42+i, 
         "nticks": 365*100, 
         "verbose": True, 
         "beta": row["R0"] * (row["mu"] + 1 / row["inf_mean"]), 
         "inf_mean": row["inf_mean"], 
         "cbr": row["cbr"], 
         "importation_period": 180, 
         "importation_count": 3}
    )

    scenario["I"] = initial_infected
    scenario["R"] = np.round(pop * (1 / row["R0"])).astype(np.int32)
    scenario["S"] = scenario.population - scenario["I"] - scenario["R"]

    birthrate_map = ValuesMap.from_scalar(parameters.cbr, nticks=parameters.nticks, nnodes=len(scenario))
    model = Model(scenario, parameters, birthrates=birthrate_map)
    infdurdist = dists.exponential(scale=parameters.inf_mean)
    rate_const = 365 * ((1 + parameters.cbr / 1000) ** (1/365) - 1) #to get stable population, small correction factor for the annual age-dist/KM estimators.
    pyramid = AliasedDistribution(stable_age_dist := np.array(1000*np.exp(-rate_const*np.arange(89))))
    survival = KaplanMeierEstimator(stable_age_dist.cumsum())
    
    model.components = [
        SIR.Susceptible(model),
        SIR.Recovered(model),
        SIR.Infectious(model, infdurdist),
        Importation(model, infdurdist),
        SIR.Transmission(model, infdurdist),
        BirthsByCBR(model, birthrate_map, pyramid),
        MortalityByEstimator(model, estimator=survival),
    ]

    model.run(f"{i+1:2}/{nsims} simulations")
    mycases[i, :] = np.squeeze(model.nodes.I)

params_df["T_obs_peakfinder1"] = np.nan
params_df["T_obs_peakfinder2"] = np.nan
params_df["T_obs_autocorr"] = np.nan

for i in range(mycases.shape[0]):
    #plt.plot(mycases[i, :])
    pf1, pf2 = ID_freq_peakfinder(np.squeeze(mycases[i, :]), params_df.loc[i, "T_exp"], plot=True)
    params_df.loc[i, "T_obs_peakfinder1"] = pf1
    params_df.loc[i, "T_obs_peakfinder2"] = pf2
    params_df.loc[i, "T_obs_autocorr"] = ID_freq_autocorr(np.squeeze(mycases[i, :]))
    plt.text(0.5, 0.99, f"T expected: {params_df.loc[i, 'T_exp']:.2f} y", transform=plt.gca().transAxes)
    plt.text(0.5, 0.95, f"T observed, peakfinding: {params_df.loc[i, 'T_obs_peakfinder1']:.2f} y", transform=plt.gca().transAxes)
    plt.text(0.5, 0.91, f"T observed, valleyfinding: {params_df.loc[i, 'T_obs_peakfinder2']:.2f} y", transform=plt.gca().transAxes)
    plt.text(0.5, 0.87, f"T observed, FFT: {params_df.loc[i, 'T_obs_autocorr']:.2f} y", transform=plt.gca().transAxes)
    plt.gca().spines['top'].set_visible(False)

    plt.show()

params_df

plt.plot(params_df["T_exp"], params_df["T_obs_peakfinder1"], "o")
plt.plot(params_df["T_exp"], params_df["T_obs_peakfinder2"], "o")
plt.plot(params_df["T_exp"], params_df["T_obs_autocorr"], "rx")
plt.plot([1, 4], [1, 4], "k--")
plt.show()