<jupyter_start><jupyter_text>A toy model for vaccinations, derived from SEIR<jupyter_code>%matplotlib inline
import numpy as np
import pyross
import copy
import matplotlib.pyplot as plt
#from matplotlib import rc; rc('text', usetex=True)<jupyter_output><empty_output><jupyter_text>The basic SEIR modelBelow you will find the model-specification dictionary for the SEIR model, without vaccinations yet<jupyter_code>model_spec = {
    "classes" : ["S", "E", "I", "R"],

    "S" : {
        "infection" : [ ["I", "S", "-beta"] ]
    },

    "E" : {
        "linear"    : [ ["E", "-gammaE"] ],
        "infection" : [ ["I", "S", "beta"] ]
    },
    
    "I" : {
        "linear"    : [ ["E", "gammaE"],
                        ["I", "-gammaI"]]  
    },
    
    "R" : {
        "linear"    : [ ["I", "gammaI"] ],     
    }
}<jupyter_output><empty_output><jupyter_text>Next, we define the parameter values:<jupyter_code>parameters = {
    'beta' : 0.05,
    'gammaE' : 0.5,
    'gammaI' : 0.1,
    'seed' : 12345  ## fix seed to get same result after each simulations
}<jupyter_output><empty_output><jupyter_text>For simplicity, we just use 3 equally sized age groups<jupyter_code>M = 3                
Ni = 50000*np.ones(M)
N = np.sum(Ni) 



# initial conditions as a dictionary 
E0 = [30, 30, 30]
S0 = [n-30 for n in Ni]
I0 = [0, 0, 0]
R0 = [0, 0, 0]

x0 = {
    'S' : S0,
    'E' : E0,
    'I' : I0, 
    'R' : R0 
}

CM = np.array([
    [3,   0.8, 0.1],  # youngest cohort spreading most
    [0.8, 2,   0.5],
    [0.1, 0.5, 1  ]
], dtype=float)

# the contact matrix is time-dependent
def contactMatrix(t):
    if t<40:
        xx = CM
    elif 40<=t<200:
        xx = 0.5*CM
    elif 200<=t<280:
        xx = 0.9*CM
    elif 280<=t<370:
        xx = 0.5*CM
    else:
        xx = CM
    return xx

# duration of simulation and data file
Tf = 1000;  Nf=Tf+1; 

model = pyross.stochastic.Model(model_spec, parameters, M, Ni)
# simulate model 
data = model.simulate(x0, contactMatrix, Tf, Nf)

I = np.sum(model.model_class_data('I', data), axis=1)
print(np.sum(I), 'previous simulation: 797679')
# plot the data and obtain the epidemic curve
S = np.sum(model.model_class_data('S', data), axis=1)
E = np.sum(model.model_class_data('E', data), axis=1)
I = np.sum(model.model_class_data('I', data), axis=1)
R = np.sum(model.model_class_data('R', data), axis=1)
t = data['t']

fig = plt.figure(num=None, figsize=(10, 8), dpi=80, facecolor='w', edgecolor='k')
plt.rcParams.update({'font.size': 22})


plt.semilogy(t, S, '-',  label='$S$', lw=4)
plt.semilogy(t, E, '-',  label='$E$', lw=4)
plt.semilogy(t, I, '-', label='$I$', lw=4)
plt.semilogy(t, R, '-',  label='$R$', lw=4)

plt.legend(fontsize=26); plt.grid() 
plt.autoscale(enable=True, axis='x', tight=True)
plt.ylabel('Compartment value')
plt.xlabel('Days');
# age structure
Ii = model.model_class_data('I', data)

fig = plt.figure(num=None, figsize=(10, 8), dpi=80, facecolor='w', edgecolor='k')
plt.rcParams.update({'font.size': 22})
plt.stackplot(t,np.transpose(Ii),labels=('cohort 1','cohort 2','cohort 3'))
plt.legend(fontsize=26); plt.grid()
plt.ylabel('Infections')
plt.xlabel('Days');<jupyter_output><empty_output><jupyter_text>Without vaccinations, herd immunity is reached after the 3rd wave (no further lockdown). Total number of people that had been infected, by age group:<jupyter_code>np.round(np.sum(Ii,axis=0)*parameters['gammaI'])<jupyter_output><empty_output><jupyter_text>The SEIR model with vaccinationsThis code takes the `Model` specification of SEIR, extends the number of classes by the vaccinated version of classes and sets up all the transitions. There are additional infection terms, the linear terms are the same in the vaccinated compartments (though they might be changed later, to account for a reduction in severity), and finite resource terms move individuals from the unvaccinated to the vaccinated classes. This avoids typing the whole model specification by hand, and should also work for our more detailed, calibrated models.The vaccinated version of class X is labeled XV, The extension to different types/doses of vaccinations should be straightforward (XV1, XV2, et.c).<jupyter_code>model_spec_vaccinations = copy.deepcopy(model_spec)

for cl in model_spec["classes"]:
    clV = cl+'V'
    model_spec_vaccinations["classes"].append(clV)
    
    model_spec_vaccinations[clV] = {}
    if "infection" in model_spec[cl]:
        model_spec_vaccinations[clV]["infection"] = []
        for term in model_spec[cl]["infection"]:
            
            termI = term.copy() # vaccinated infecting unvaccinated
            termI[0] += 'V'
            termI[2] += "*thetaI"
            model_spec_vaccinations[cl]["infection"].append(termI)
            
            termS = term.copy() # unvaccinated infecting vaccinated
            termS[1] += 'V'
            termS[2] += "*thetaS"
            model_spec_vaccinations[clV]["infection"].append(termS)
            
            termSI = term.copy() # vaccinated infecting vaccinated
            termSI[0] += 'V'
            termSI[1] += 'V'
            termSI[2] += "*thetaS*thetaI"
            model_spec_vaccinations[clV]["infection"].append(termSI)
            
    if "linear" in model_spec[cl]:
        model_spec_vaccinations[clV]["linear"] = []
        for term in model_spec[cl]["linear"]:
            termV = term.copy()
            termV[0] += 'V'
            model_spec_vaccinations[clV]["linear"].append(termV)
    
    if "finite-resource" in model_spec[cl]:
        model_spec_vaccinations[clV]["finite-resource"] = []
        for term in model_spec[cl]["finite-resource"]:
            termV = term.copy()
            termV[0] += 'V'
            model_spec_vaccinations[clV]["finite-resource"].append(termV)
    else:
        model_spec_vaccinations[cl]["finite-resource"] = []
        model_spec_vaccinations[clV]["finite-resource"] = []
    term = [cl, "vaccination_rate", "vaccination_priority", "one"]   # one: probability that the vaccination is successful. Set to 1, because we already 
                                                                     # account for the imperfectness of vaccinations through the theta parameters
    model_spec_vaccinations[clV]["finite-resource"].append(term.copy())
    term[1] = "-"+term[1]
    model_spec_vaccinations[cl]["finite-resource"].append(term)
    
model_spec_vaccinations<jupyter_output><empty_output><jupyter_text>The set of parameters gets extended by the parameters $\theta_S$ and $\theta_I$, representing the reduction in susceptibility and infectiousness in vaccinated individuals. See https://www.medrxiv.org/content/10.1101/2021.03.14.21253544v2<jupyter_code>parameters = {
    'beta' : 0.05,
    'gammaE' : 0.5,
    'gammaI' : 0.1,
    'thetaS' : 0.6,  
    'thetaI' : 0.6,
}<jupyter_output><empty_output><jupyter_text>The age-dependence of vaccinations is introduced explicitly through an age-dependent vaccination rate. Alternatively, one could have specified a scalar vaccination rate, and implicilty realised the age-dependence though age-dependent priorities.<jupyter_code>def vaccination_rate(t):
    rate=500*np.array([0.5*np.tanh((t-420)/10)-0.5*np.tanh((t-480)/10),
                       0.5*np.tanh((t-360)/10)-0.5*np.tanh((t-420)/10),
                       0.5*np.tanh((t-300)/20)-0.5*np.tanh((t-360)/10)  ])
    rate*=(rate>1)
    return rate

def parameter_mapping(input_parameters, t):
    output_parameters = {
        'beta' : input_parameters['beta'],
        'gammaE' : input_parameters['gammaE'],
        'gammaI' : input_parameters['gammaI'],
        'beta*thetaS' : input_parameters['beta']*input_parameters['thetaS'],
        'beta*thetaI' : input_parameters['beta']*input_parameters['thetaI'],
        'beta*thetaS*thetaI' : input_parameters['beta']*input_parameters['thetaS']*input_parameters['thetaI'],
        'one' : 1,
        'vaccination_priority' : 1,
        'vaccination_rate' : vaccination_rate(t)     
    }
        
    return output_parameters
t=np.linspace(200,600,1024)
v=[vaccination_rate(tt) for tt in t]
plt.plot(t,v)
plt.xlabel('days')
plt.ylabel('vaccination rate (by age)')
plt.show()<jupyter_output><empty_output><jupyter_text>We have chosen the vaccination rates such that 60% of each cohort will get vaccinated<jupyter_code># initial conditions as a dictionary 
E0 = [30, 30, 30]
S0 = [n-30 for n in Ni]
I0 = [0, 0, 0]
R0 = [0, 0, 0]

x0 = {
    'S' : S0,
    'E' : E0,
    'I' : I0, 
    'R' : R0 
}

for cl in model_spec["classes"]:
    x0[cl+'V'] = np.zeros(M)


# duration of simulation and data file
Tf = 1000;  Nf=Tf+1; 

model = pyross.stochastic.Model(model_spec_vaccinations, parameters, M, Ni, time_dep_param_mapping=parameter_mapping)
# simulate model 
data = model.simulate(x0, contactMatrix, Tf, Nf)



# plot the data and obtain the epidemic curve
S = np.sum(model.model_class_data('S', data), axis=1)
E = np.sum(model.model_class_data('E', data), axis=1)
I = np.sum(model.model_class_data('I', data), axis=1)
R = np.sum(model.model_class_data('R', data), axis=1)
SV = np.sum(model.model_class_data('SV', data), axis=1)
EV = np.sum(model.model_class_data('EV', data), axis=1)
IV = np.sum(model.model_class_data('IV', data), axis=1)
RV = np.sum(model.model_class_data('RV', data), axis=1)
V=SV+EV+IV+RV
t = data['t']

fig = plt.figure(num=None, figsize=(10, 8), dpi=80, facecolor='w', edgecolor='k')
plt.rcParams.update({'font.size': 22})


plt.semilogy(t, S+SV, '-',  label='$S$', lw=4)
plt.semilogy(t, E+EV, '-',  label='$E$', lw=4)
plt.semilogy(t, I+IV, '-', label='$I$', lw=4)
plt.semilogy(t, R+RV, '-',  label='$R$', lw=4)
plt.semilogy(t, V, '-',  label='$V$', lw=4)

plt.legend(fontsize=26); plt.grid() 
plt.autoscale(enable=True, axis='x', tight=True)
plt.ylabel('Compartment value')
plt.xlabel('Days');
# age structure
Ii = model.model_class_data('I', data)
IVi = model.model_class_data('IV', data)


fig = plt.figure(num=None, figsize=(10, 8), dpi=80, facecolor='w', edgecolor='k')
plt.rcParams.update({'font.size': 22})
plt.stackplot(t,np.transpose(Ii+IVi),labels=('cohort 1','cohort 2','cohort 3'))
plt.legend(fontsize=26); plt.grid()
plt.ylabel('Infections')
plt.xlabel('Days');<jupyter_output><empty_output><jupyter_text>With vaccinations, the 3rd wave is reduced, despite the lifting of lockdown. Total number of people that had been infected, by age group:<jupyter_code># unvaccinated
np.round(np.sum(Ii,axis=0)*parameters['gammaI'])
# vaccinated
np.round(np.sum(IVi,axis=0)*parameters['gammaI'])<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.15879052961871257, 'fraction_numerical': 0.028620328991157174, 'mean_word_length': 3.7293165467625897, 'pattern_counts': {'":': 0, '<': 48, '<?xml version=': 0, '>': 42, 'https://': 1, 'lorem ipsum': 0, 'www.': 1, 'xml': 0}, 'pii_count': 1, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 3, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '10', 'dir': 'jupyter-structured-clean-dedup', 'id': '115758', 'license': 'MIT', 'n_tokens_mistral': 4032, 'n_tokens_neox': 3725, 'n_words': 852, 'path': 'examples/stochastic/ex01b_Model.ipynb', 'repo_name': 'vishalbelsare/pyross', 'size': '10517'}