sir-with-variant-heterogeneous-new.sage

## Consider a population infected by a virus which has two forms, an
## ancestral one and a variant (with perfect cross-immunity) having a
## different basic reproduction number and possibly a different
## recovery rate.  We simulate the evolution of the epidemic possibly
## with heterogeneous susceptibility, different for both variants and
## possibly correlated and/or epidemic quenching (viz., mitigation
## measures which depend on number of infected).

## s  = susceptible
## i1 = infected by ancestral form
## i2 = infected by variant
## r  = recovered (immune to both)
## f1 = cumulated infection events

var('s i1 i2 r1 r2 f1 f2 t')

# Reproduction number of the ancestral form
rep1 = 1.0
# Reproduction number of the variant
rep2 = 1.4

# Recovery rate for the ancestral form
gamma1 = 1.0
# Recovery rate for the variant
gamma2 = 1.0

# Initial proportion infected
eps = 0.0075
# Proportion OF infected initially by variant
vprop = 1e-2

# Quenching parameter (use 0 for standard SIR)
quench = 0.0

# Relative variance of susceptibility for ancestral form
hetvar1 = 0.0
# Relative variance of susceptibility for variant
hetvar2 = 0.0
# Relative covariance of susceptibility between variantsn
hetcovar = 0.0

beta1 = rep1*gamma1
beta2 = rep2*gamma2
f1dot = beta1*i1*(1-quench*(i1+i2))
f2dot = beta2*i2*(1-quench*(i1+i2))
meansusc1 = (1+hetvar1*f1+hetvar2*f2-hetcovar*f2)/(1+hetvar1*f1)/(1+hetvar1*f1+hetvar2*f2)
meansusc2 = (1+hetvar1*f1+hetvar2*f2-hetcovar*f1)/(1+hetvar2*f2)/(1+hetvar1*f1+hetvar2*f2)
infect1 = f1dot*s*meansusc1
infect2 = f2dot*s*meansusc2
sdot = -infect1-infect2
i1dot = infect1-gamma1*i1
i2dot = infect2-gamma2*i2
r1dot = gamma1*i1
r2dot = gamma2*i2
p = desolve_system_rk4([sdot, i1dot, i2dot, r1dot, r2dot, f1dot, f2dot], [s,i1,i2,r1,r2,f1,f2], ics=[0, 1-eps, eps*(1-vprop), eps*vprop, 0, 0, 0, 0], ivar=t, end_points=[0, 50], step=0.1)

g1 = line([(tmp[0], tmp[1]) for tmp in p], color="green") + line([(tmp[0], tmp[2]+tmp[3]) for tmp in p], color="red") + line([(tmp[0], tmp[4]+tmp[5]) for tmp in p], color="blue") + line([(tmp[0], tmp[3]/(tmp[2]+tmp[3])) for tmp in p], color="purple", linestyle="dashed")
g2 = line([(tmp[0], tmp[2]+tmp[3]) for tmp in p], color="red", ymin=0, ymax=0.025) + line([(tmp[0], tmp[2]) for tmp in p], color="orange", ymin=0, ymax=0.025) + line([(tmp[0], tmp[3]) for tmp in p], color="purple", ymin=0, ymax=0.025)
g3 = line([(tmp[0], tmp[1]) for tmp in p], color="green", ymin=1e-4, ymax=1, scale="semilogy") + line([(tmp[0], tmp[2]+tmp[3]) for tmp in p], color="red", ymin=1e-4, ymax=1, scale="semilogy") + line([(tmp[0], tmp[4]+tmp[5]) for tmp in p], color="blue", ymin=1e-4, ymax=1, scale="semilogy") + line([(tmp[0], tmp[3]/(tmp[2]+tmp[3])) for tmp in p], color="purple", linestyle="dashed", ymin=1e-4, ymax=1, scale="semilogy")
g4 = line([(tmp[0], 1) for tmp in p], color="gray", linestyle="dashdot", ymin=0.5, ymax=1.5) + line([(tmp[0], (1+i1dot/r1dot).subs({s:tmp[1], i1:tmp[2], i2:tmp[3], f1:tmp[6], f2:tmp[7]})) for tmp in p], color=Color("orange").darker(), thickness=0.5, ymin=0.5, ymax=1.5) + line([(tmp[0], (1+i2dot/r2dot).subs({s:tmp[1], i1:tmp[2], i2:tmp[3], f1:tmp[6], f2:tmp[7]})) for tmp in p], color=Color("purple").darker(), thickness=0.5, ymin=0.5, ymax=1.5) + line([(tmp[0], (1+(i1dot+i2dot)/(i1+i2)).subs({s:tmp[1], i1:tmp[2], i2:tmp[3], f1:tmp[6], f2:tmp[7]})) for tmp in p], color="black", linestyle="dotted", ymin=0.5, ymax=1.5) + line([(tmp[0], (1+(i1dot+i2dot)/(r1dot+r2dot)).subs({s:tmp[1], i1:tmp[2], i2:tmp[3], f1:tmp[6], f2:tmp[7]})) for tmp in p], color="black", ymin=0.5, ymax=1.5)
graphics_array([[g1, g2], [g3, g4]]).show(dpi=192)