In [22]:

```
# Libraries
# The code was initially inspired by the following sources:
# https://www.kaggle.com/pstarszyk/seir-model-with-extensions/comments#Define-Extended-Model
# https://www.kaggle.com/anjum48/seir-hcd-model/data#SEIR-HCD-Model
# http://gabgoh.github.io/COVID/index.html
# https://neherlab.org/covid19/about
import pandas as pd
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
import os
from scipy.integrate import solve_ivp
from scipy.optimize import minimize
from sklearn.metrics import mean_squared_log_error, mean_squared_error
from collections import namedtuple
from collections import defaultdict as ddict
from copy import deepcopy
import warnings
from datetime import datetime, timedelta
warnings.filterwarnings("ignore")
import seaborn as sns
import matplotlib.pylab as plt
from seaborn import heatmap
from pyDOE import *
```

In [23]:

```
# This is a function used to control the font size in generating the figures
def update_font_size(font_size = 30,rcParams = mpl.rcParams):
rcParams.update({'font.size': font_size})
plt.rc('font', size=font_size) # controls default text sizes
plt.rc('axes', titlesize=font_size) # fontsize of the axes title
plt.rc('axes', labelsize=font_size*1.1) # fontsize of the x and y labels
plt.rc('xtick', labelsize=font_size) # fontsize of the tick labels
plt.rc('ytick', labelsize=font_size) # fontsize of the tick labels
plt.rc('legend', fontsize=font_size) # legend fontsize
plt.rc('figure', titlesize=font_size)
```

In [24]:

```
# Make model labels a global so we can access anywhere
model_labels = {'Cashier' :'Public Facing',
'USPS' :'Non-Public Facing',
'Healthcare' :'Healthcare',
'No Structure':'No EWs',
'Healthcare_Sigmoid' :'Healthcare'}
```

In [25]:

```
# Susceptible equation
def dS_dt(y_i, p, N_i, I, Beta_t, H, Beta_hosp, groupID):
if groupID == 1 and p.model == 'Healthcare_Sigmoid':
return - y_i.S * (sum(Beta_t * I) * (1./N_i) + Beta_hosp)
else:
return -(1./N_i) * y_i.S * ( sum(Beta_t * I) + Beta_hosp * H)
# Exposed equation
def dE_dt(y_i, p, N_i, I, Beta_t, H, Beta_hosp, groupID):
if groupID == 1 and p.model == 'Healthcare_Sigmoid':
incoming = y_i.S * (sum(Beta_t * I) * (1./N_i) + Beta_hosp)
else:
incoming = (1./N_i) * y_i.S * ( sum(Beta_t * I) + Beta_hosp * H )
return incoming - (y_i.E / p.t_inc)
# Infected, asymptomatic, and will recover equation
def dIA_dt(y_i, p,groupID):
if groupID == 1:
p_IA = p.p_IA * p.change_in_p_IA
else:
p_IA = p.p_IA
return (y_i.E / p.t_inc) * p_IA - y_i.IA / p.t_AR
# Infected, symptomatic, and will recover equation
def dIR_dt(y_i, p,groupID):
if groupID == 1:
p_IA = p.p_IA * p.change_in_p_IA
p_IH = p.p_IH * p.change_in_p_IH
else:
p_IA = p.p_IA
p_IH = p.p_IH
return (y_i.E / p.t_inc)*(1-p_IA)*(1-p_IH) - y_i.IR / p.t_IR
# Infected and will require hospitalization equation
def dIH_dt(y_i, p,groupID):
if groupID == 1:
p_IA = p.p_IA*p.change_in_p_IA
p_IH = p.p_IH*p.change_in_p_IH
else:
p_IA = p.p_IA
p_IH = p.p_IH
return (y_i.E / p.t_inc)*(1-p_IA)*p_IH - y_i.IH / p.t_IH
# Recovered equation
def dR_dt(y_i, p,groupID):
return y_i.IR / p.t_IR + y_i.IA / p.t_AR + y_i.HR / p.t_HR + y_i.L / p.t_LR
# Hospitalized and will recover equation
def dHR_dt(y_i, p,groupID):
return y_i.IH / p.t_IH * (1 - p.p_HC) - y_i.HR / p.t_HR
# Hospitalized and will require critical care equation
def dHC_dt(y_i, p,groupID):
return y_i.IH / p.t_IH * p.p_HC - y_i.HC / p.t_HC
# Critical care and will recover equation
def dCR_dt(y_i, p,groupID):
return y_i.HC / p.t_HC * (1 - p.p_CD) - (y_i.CR / p.t_CL)
# Critical and will die equation
def dCD_dt(y_i, p, groupID):
return y_i.HC / p.t_HC * p.p_CD - y_i.CD / p.t_CD
# Post-Critical care but still hospitalized equation
def dL_dt(y_i, p,groupID):
return y_i.CR / p.t_CL - y_i.L / p.t_LR
# Dead equation
def dD_dt(y_i, p,groupID):
return y_i.CD / p.t_CD
```

In [26]:

```
# Parameters: Using a namedtuple allows you to access parameters by name instead of index (e.g. args.N)
parameters = namedtuple('parameters',["R_0","N","C",
"social_distancing_time","theta","model","proportion_essential","rho",
"max_days","num_groups","compartment_names",
"t_inc","t_AR","t_IR","t_IH","t_HR","t_HC","t_CD","t_CL","t_LR",
"p_IA","p_IH","p_HC","p_CD",
"change_in_p_IH","change_in_p_IA","adjust_beta_hosp"])
# Define the compartments and link names with formulas
compartment_names = ["S","E","IA","IR","IH","HR","HC","CR","CD","L","R","D"]
compartment_formulas = [dS_dt,dE_dt,dIA_dt,dIR_dt,dIH_dt,dHR_dt,dHC_dt,dCR_dt,dCD_dt,dL_dt,dR_dt,dD_dt]
compartments = namedtuple("compartments",compartment_names)
```

In [27]:

```
# This function returns an array of differential equations that is used for the ODE solver in the following section
def SEIR_HCD_model(t, y_vals, parameters):
# Convert to matrix
# Each row is a group
# Each column corresponds to a given class within each group
# Columns indexed according to position in name array: compartment_names = ["S", "E", "I", "R", "H", "C", "D"]
y_matrix = np.reshape(y_vals,[parameters.num_groups,len(compartment_names)])
output = np.zeros(np.shape(y_matrix))
# Iterate over each group
# Nonessential = 0
# Essential = 1
for group_idx in range(parameters.num_groups):
Beta_t_i = get_betas(t,parameters,group_idx)
# Components for group i
compartments_group_i = compartments(*y_matrix[group_idx,:])
# Total population within group i
N_i = sum(y_matrix[group_idx,:])
# Get the infected population for each group for all 3 compartments
I = sum(np.transpose(y_matrix[:,2:5]))
# Get the hospitalized population for each group
H = sum(sum(np.transpose(y_matrix[:,5:10])))
# If using a healthcare model and considering the healthcare group,
# we need to consider infections stemming from hospitals
if 'Healthcare' in parameters.model and group_idx == 1:
# Determine the beta or fraction of workers infected
beta_hosp = get_hosp_beta(t,parameters,H)
else:
beta_hosp = 0
# Susceptibles and Exposed DiffEq use special parameters I and Beta_t_i)
S_out_i = dS_dt(y_i = compartments_group_i,
p = parameters,
N_i = N_i,
I = I,
Beta_t = Beta_t_i,
H = H,
Beta_hosp = beta_hosp,
groupID = group_idx)
output[group_idx,0] = S_out_i
E_out_i = dE_dt(compartments_group_i, parameters, N_i, I, Beta_t_i, H, beta_hosp, group_idx)
output[group_idx,1] = E_out_i
# The remainder of the differentiail equations
for index in np.arange(2,len(compartment_names)):
output[group_idx,index] = compartment_formulas[index](compartments_group_i, parameters,group_idx)
# Convert the matrix back into a linear array for the ODE solver
output_array = np.reshape(output,[1,parameters.num_groups*len(compartment_names)])[0]
return output_array
```

In [28]:

```
# Return the matrix of thetas that determine how to adjust betas given lockdown
def get_effect_of_social_distancing(p,t=0):
if p.model == 'No Structure' or (p.rho > 0 and (p.model == 'USPS' or p.model == 'Healthcare' or p.model == 'Healthcare_Sigmoid')):
social_distancing_effect = np.array([p.theta]*4)
elif p.model == 'Cashier':
social_distancing_effect = np.array([p.theta]+[1]*3)
elif p.model == 'USPS' or p.model == 'Healthcare' or p.model == 'Healthcare_Sigmoid':
social_distancing_effect = np.array([p.theta]*3 + [1])
else:
raise ValueError('Model should be either Cashier,USPS,or No Structure. '
+ p.model + ' is not an implemented model')
social_distancing_effect = np.reshape(social_distancing_effect,[p.num_groups,p.num_groups])
return social_distancing_effect
# Return the beta for within hospital interactions
def get_hosp_beta(t,p,H):
if p.model == 'Healthcare_Sigmoid':
return get_sigmoid_beta(t,p,H)
else:
# Default: assume the within beta hospital interaction is the same as the beta between I compartments and S
return get_infectious_beta(p)*p.adjust_beta_hosp
def get_sigmoid_beta(t,p,H):
exponent_parameter = 1
halfway_point = 1/4
max_infection_rate = get_infectious_beta(p)*p.adjust_beta_hosp*1.5*H_max/p.N
halfway_exponent = halfway_point**exponent_parameter
h_exponent = (H/H_max)**exponent_parameter
return ( max_infection_rate * ( h_exponent ) / ( h_exponent+halfway_exponent ) )
def get_infectious_beta(p,model = ''):
# llow user to set model, but otherwise use parameter model
if not model:
model = p.model
# Fraction of essential workers
f = p.proportion_essential
# Average time spent in I compartments
avg_time_infectious = (p.p_IA * p.t_AR + (1 - p.p_IA) * ((1-p.p_IH) * p.t_IR + p.p_IH * p.t_IH))
# Average time spent in the hospital
if model == 'Healthcare':
avg_time_critical = (1 - p.p_CD) * (p.t_CL + p.t_LR) + p.p_CD * (p.t_CD)
avg_time_hospital = (1-p.p_HC) * p.t_HR + p.p_HC * (p.t_HC +avg_time_critical)
avg_time_healthcare = (1 - p.p_IA) * p.p_IH * avg_time_hospital
else:
# ignore if not a healthcare model
avg_time_healthcare = 0
# Get the mixing matrix or set to proportional mixing if not provided
C = p.C
if C == []:
C = np.array([[1-f,1-f],[f,f]])
# Vector to describe the distribution of people
people_array = np.array([1-f,f])
# The weighted sum across mixings
beta_coefficient = avg_time_infectious*(sum(np.matmul(C,people_array)) + p.rho * f)
hosp_coefficient = avg_time_healthcare
beta = p.R_0/(beta_coefficient + hosp_coefficient * p.adjust_beta_hosp * keep_R0_constant)
return beta
# This allows for greater within-group mixing than under proportional mixing
def get_extra_rho_effect(t,p,groupID):
if groupID != 1:
return 0
if t > p.social_distancing_time and p.model == 'No Structure':
return p.rho * (p.theta)
return p.rho
def get_betas(t,p,groupID,model = ''):
if not model:
model = p.model
social_distancing_effect = get_effect_of_social_distancing(p,t)
if t>p.social_distancing_time:
social_distancing_effect = social_distancing_effect[:,groupID]
else:
social_distancing_effect = social_distancing_effect[:,groupID] * 0 + 1
# beta is R_0 * mean infectious time
# This is the same beta across all infectious compartments and groups
nointervention_beta = get_infectious_beta(p,model)
# This gives you the betas including the mixing between groups (specific to group i)
betas = list(nointervention_beta*social_distancing_effect*p.C[groupID,:])
extra_rho = get_extra_rho_effect(t,p,groupID)
betas[1] += extra_rho * nointervention_beta
return betas
```

In [29]:

```
#These are the parameters used in the models. These are listed in Table 1 of the main text.
#By default, R0 is 2.5, however this is changed to 3.2 for most of the subsequent plots
keep_R0_constant = False
args_default = parameters(R_0 = 2.5, # R_0: used to calculate betas
N = 8e6, # total population size
C = [], # Mixing matrix
social_distancing_time = 50, # When does social distancing start
theta = 1, # How effective is social distancing
model = 'Cashier', # Who social distances?
proportion_essential = 0.1, # What fraction of the population is essential workers?
rho = 1,
max_days = 1000, # bookkeeping of how many days simulated
num_groups = 2, # bookkeeping of how many groups
compartment_names = compartment_names, # bookkeeping of how many compartments
t_inc = 4.6, # !time from E to I: exposed to infected
t_AR = 5, # !time from IA to R: infected and asymptomatic to recovered
t_IR = 5, # !time from IR to R: infected, symptomatic and will recover to recovered
t_IH = 5, # !time from IH to H: infected and will need hospitalization to hospitalization
t_HR = 8, # !time from HR to R: hospitalized and will recover to recovered
t_HC = 6, # !time from HC to C: hospitalization to critical care
t_CD = 10,# time from CD to D: critical care and will die to death
t_CL = 7, # time from CR to L: critical care and will recover to post-critical care hospitalization
t_LR = 3, # time from L to R: post-critical care hospitalization to recovered
p_IA = 1/3, # Proportions of infections that are asymptomatic
p_IH = 0.066, # Proportion of symptomatic infections that require hospitalizations (this corresponds to 4.4% of all infections that require hospitalization)
p_HC = 0.3, # Proportion of hospitalizations that require critical care
p_CD = 0.5, # Proportion of critical cases that are fatal
change_in_p_IH = 1, # How much is the probability of hopsitalization increased in OLD
change_in_p_IA = 1, # How much is the probability of asymptomatic increased in OLD
adjust_beta_hosp = 1) # change the beta for within hospital interactions relative to beta between I and S
```

In [30]:

```
# Returns number of cases, hospitalizations, and deaths given model solution, number of groups, and compartment names
def get_compartment_counts(solution, num_groups,
compartment_names = ["S","E","IA","IR","IH","HR","HC","CR","CD","L","R","D"]):
# convert solution into a matrix
max_days = len(solution.y[0])
solution_matrix = np.reshape(solution.y,[num_groups,len(compartment_names),max_days])
# sum across groups
total_solution = np.zeros([len(compartment_names),max_days])
for i in range(num_groups):
total_solution += solution_matrix[i]
# sum across relevant compartments
cases = sum(total_solution[2:])
deaths = total_solution[-1]
hospitilized = sum(total_solution[5:10])
return cases, hospitilized, deaths
# get R_0 from the doubling time
def R0_from_doubling_time(p,doubling_time,model= None,adjust_beta_hosp = np.nan,verbose = False):
# allow user to specify the model and override the default model
if not model:
model = p.model
if np.isnan(adjust_beta_hosp):
adjust_beta_hosp = p.adjust_beta_hosp
# Get average infectious time
avg_time_infectious = (p.p_IA * p.t_AR + (1 - p.p_IA) * ((1-p.p_IH) * p.t_IR + p.p_IH * p.t_IH))
# if healthcare model, add average healthcare time to the infectious time
if 'Healthcare' in model:
avg_time_critical = (1 - p.p_CD) * (p.t_CL + p.t_LR) + p.p_CD * (p.t_CD)
avg_time_hospital = (1 - p.p_HC) * p.t_HR + p.p_HC * (p.t_HC + avg_time_critical)
# only affects essential workers and include the relative deviation in the hospital beta
avg_time_infectious += keep_R0_constant * (1 - p.p_IA) * p.p_IH * avg_time_hospital * adjust_beta_hosp
if verbose:
print(avg_time_infectious,(p.p_IA * p.t_AR + (1 - p.p_IA) * ((1-p.p_IH) * p.t_IR + p.p_IH * p.t_IH)),avg_time_critical,avg_time_hospital)
R0 = 1 + (p.t_inc + avg_time_infectious) * np.log(2) / doubling_time
return R0
```

In [31]:

```
#This function will run the model above, and return the solution and parameters used in a tuple
def run_model(theta,R_0,t_lockdown,model,n_infected=50,n_exposed=50,prop_essential=0.1,max_days = 1000,
rho = 1, change_in_p_IH = 1, change_in_p_IA = 1,num_groups = 2,doubling_time = None,
adjust_beta_hosp = 1):
# Define the initial conditions & population size
N = 8e6
# Number of groups
num_groups = 2
# Social interaction effects: C
# C[i][j] = proportion of group I's contacts that are with group
# rho allows deviations from proportionate mixing
C = np.array([ (1 - prop_essential) + rho * prop_essential,
(1 - rho) * (1 - prop_essential),
(1 - rho) * (prop_essential),
1 - (1 - rho) * (1 - prop_essential) - rho])
C = np.reshape(C,[num_groups,num_groups])
if doubling_time:
args = parameters(R_0 = R_0, # R_0: used to calculate betas
N = N, # total population size
C = C, # Mixing matrix
social_distancing_time = t_lockdown, # when does social distancing start
theta = theta, # How effective is social distancing
model = model, # Who social distances?
proportion_essential = prop_essential, #
rho = rho, # record changes from proportionate mixing
max_days = max_days, # bookkeeping of how many days simulated
num_groups = num_groups, # bookkeeping of how many groups
compartment_names = compartment_names, # bookkeeping of how many compartments
t_inc = 4.6, # !time from E to I: exposed to infected
t_AR = 5, # !time from IA to R: infected and asymptomatic to recovered
t_IR = 5, # !time from IR to R: infected, symptomatic and will recover to recovered
t_IH = 5, # !time from IH to H: infected and will need hospitalization to hospitalization
t_HR = 8, # !time from HR to R: hospitalized and will recover to recovered
t_HC = 6, # !time from HC to C: hospitalization to critical care
t_CD = 10,# time from CD to D: critical care and will die to death
t_CL = 7, # time from CR to L: critical care and will recover to post-critical care hospitalization
t_LR = 3, # time from L to R: post-critical care hospitalization to recovered
p_IA = 1/3, # Proportions of infections that are asymptomatic
p_IH = 0.066, # Proportion of symptomatic infections that require hospitalizations, conditional on not being asymptomatic
p_HC = 0.3, # Proportion of hospitalizations that require critical care
p_CD = 0.5, # Proportion of critical cases that are fatal
change_in_p_IH = 1, # How much is the probability of hopsitalization increased in OLD
change_in_p_IA = 1, # How much is the probability of asymptomatic increased in OLD
adjust_beta_hosp = adjust_beta_hosp) # change the beta for within hospital interactions relative to beta between I and S
if doubling_time: R_0 = R0_from_doubling_time(args,doubling_time)
args = parameters(R_0 = R_0, # R_0: used to calculate betas
N = N, # total population size
C = C, # Mixing matrix
social_distancing_time = t_lockdown, # when does social distancing start
theta = theta, # How effective is social distancing
model = model, # Who social distances?
proportion_essential = prop_essential, #
rho = rho, # record changes from proportionate mixing
max_days = max_days, # bookkeeping of how many days simulated
num_groups = num_groups, # bookkeeping of how many groups
compartment_names = compartment_names, # bookkeeping of how many compartments
t_inc = 4.6, # !time from E to I: exposed to infected
t_AR = 5, # !time from IA to R: infected and asymptomatic to recovered
t_IR = 5, # !time from IR to R: infected, symptomatic and will recover to recovered
t_IH = 5, # !time from IH to H: infected and will need hospitalization to hospitalization
t_HR = 8, # !time from HR to R: hospitalized and will recover to recovered
t_HC = 6, # !time from HC to C: hospitalization to critical care
t_CD = 10,# time from CD to D: critical care and will die to death
t_CL = 7, # time from CR to L: critical care and will recover to post-critical care hospitalization
t_LR = 3, # time from L to R: post-critical care hospitalization to recovered
p_IA = 1/3, # Proportions of infections that are asymptomatic
p_IH = 0.066, # Proportion of symptomatic infections that require hospitalizations
p_HC = 0.3, # Proportion of hospitalizations that require critical care
p_CD = 0.5, # Proportion of critical cases that are fatal
change_in_p_IH = 1, # How much is the probability of hopsitalization increased in OLD
change_in_p_IA = 1, # How much is the probability of asymptomatic increased in OLD
adjust_beta_hosp = adjust_beta_hosp) # change the beta for within hospital interactions relative to beta between I and S
single_population = np.zeros(len(compartment_names))
single_population[1] = n_exposed
single_population[2] = n_infected*args.p_IA
single_population[3] = n_infected*(1-args.p_IA)*(1-args.p_IH)
single_population[4] = n_infected*(1-args.p_IA)*(args.p_IH)
single_population[0] = N - sum(single_population)
initial_state = [(1 - prop_essential) * single_population,
prop_essential * single_population]
# Convert initial state into 1-dimensional array
initial_state = [params for group in initial_state for params in group]
# Evaluate the model
# Here, a stiff solver (Radau) not RK45 (default) is used
sol = solve_ivp(lambda t,y: SEIR_HCD_model(t,y,args), [0, max_days], initial_state, t_eval=np.arange(max_days),method = 'Radau')
return (sol,args)
```

In [32]:

```
# Get the cumulative infection rate at some time
def get_infection_rate(solution, p, time = 365):
# Turn the solution array into a matrix
solution_matrix = np.reshape(solution.y,[p.num_groups,len(p.compartment_names),p.max_days])
# If no time specified, get cumulative infection at end of time
if not time or time == -1:
time = p.max_days-1
# Get the number of infections, return it as a fraction of the populations
infections = sum(sum(solution_matrix[:,1:,time]))
return infections/p.N
# Return the metrics to show the added disease risk within the essential workers
def added_disease_risk(solution, p,method = 'g(t)',groupID = 1):
solution_matrix = np.reshape(solution.y,[p.num_groups,len(p.compartment_names),p.max_days])
f = p.proportion_essential
if method == 'Fraction Infected':
if groupID == 1:
N = (p.N*(p.proportion_essential))
else:
N = (p.N*(1-p.proportion_essential))
return sum(solution_matrix[groupID,1:,-1])/N
elif method == 'Nonessential Comparison':
L = sum(solution_matrix[1,1:,-1])/sum(solution_matrix[0,1:,-1])*((1-f)/f)
return L
elif method == 'g(t) additive':
return sum(solution_matrix[1,1:,-1])/(sum(solution_matrix[1,1:,-1])+sum(solution_matrix[0,1:,-1]))
elif method == 'g(t) multiplicative':
return sum(solution_matrix[1,1:,-1])/(sum(solution_matrix[1,1:,-1])+sum(solution_matrix[0,1:,-1]))/f
elif method == 'EisenA':
return sum(solution_matrix[1,1:,-1])/(p.N*f)-sum(solution_matrix[0,1:,-1])/(p.N*(1-f))
elif method == 'No Structure Comparison':
NS_sol,NS_p = run_model(theta=p.theta,
R_0=p.R_0,
t_lockdown=p.social_distancing_time,
model='No Structure',
n_infected=20,
n_exposed=0,
prop_essential=p.proportion_essential,
max_days = p.max_days,
rho = 0,
change_in_p_IH = p.change_in_p_IH,
change_in_p_IA = p.change_in_p_IA,
doubling_time = 3,
adjust_beta_hosp = 1)
NS_solution_matrix = np.reshape(NS_sol.y,[p.num_groups,len(p.compartment_names),p.max_days])
# Turn the solution array into a matrix
return sum(solution_matrix[1,1:,-1])/sum(NS_solution_matrix[1,1:,-1])
else:
raise ValueError("Method not understood")
# Given a solution to the model and the corresponding parameters, get several different metrics of model outcomes
def get_metrics(result_dict):
R0 = R0_from_doubling_time(args_default,3)
# Create a nested dict, where the structure is {model:{theta:{f:([],[])}}}
reversed_solved_dict = ddict(dict)
for model,model_dict in result_dict.items():
reversed_solved_dict[model] = ddict(dict)
for f,f_dict in model_dict.items():
for theta,(sol,args) in f_dict.items():
reversed_solved_dict[model][theta][f] = (get_infection_rate(solution=sol,p=args, time = 365),
added_disease_risk(solution=sol,p = args,method = 'Fraction Infected',groupID = 0),
added_disease_risk(solution=sol,p = args,method = 'Fraction Infected',groupID = 1))
return reversed_solved_dict
# Assign row and column labels for the plots of the effective R over time
def assign_row_label_grid(ax,f, pad = 5, xy = (0,0.5),font_size = 30):
ax.annotate('f = {}'.format(f)+'\n', xy=xy, xytext=(-ax.yaxis.labelpad - pad, 0),
xycoords=ax.yaxis.label, textcoords='offset points',
size=font_size, ha='right', va='center',rotation=90)
# Label is the value of R_effective
def assign_column_label_grid(ax,value, pad = 25, xy = (0.5,1),font_size = 30):
text = r'{}'.format(round(value,2))
ax.annotate(text, xy=xy, xytext=(0, pad),
xycoords='axes fraction', textcoords='offset points',
size=font_size, ha='center', va='baseline')
def get_R_effective(solution_matrix,p):
# Set up an empty array for both effective R and the betas over time
R_eff = np.zeros(p.max_days+1)
betas_eff = np.zeros([2,2,p.max_days+1])
# Get the time values
t = np.arange(0,p.max_days)
# Fraction of population an essential worker
f = p.proportion_essential
# Get average time an infectious person is infectious
avg_time_infectious = (p.p_IA * p.t_AR + (1 - p.p_IA) * ((1-p.p_IH) * p.t_IR + p.p_IH * p.t_IH))
# Find the effects of social distancing. Second variable is from a different implementation
social_distancing_effect = get_effect_of_social_distancing(p,None)
# beta is R_0 * mean infectious time
# This is the same beta across all infectious compartments and groups
nointervention_beta = get_infectious_beta(p)
# Get fraction of infected population that is essential over time
I_t_0 = sum(solution_matrix[0,2:5,:])
I_t_1 = sum(solution_matrix[1,2:5,:])
S_t_0 = solution_matrix[0,0,:]/(p.N*(1-f))
S_t_1 = solution_matrix[1,0,:]/(p.N*f)
f_t_all = (I_t_1*f)/(I_t_0*(1-f)+I_t_1*f)
for j in t:
# Get fraction of infected that are essential at this time, j
f_t = f_t_all[j]
# Calculate betas. This is elementwise multiplication not matrix multiplication
betas = nointervention_beta*(p.C*(1-social_distancing_effect*(j>p.social_distancing_time)))
betas_eff[:,:,j] = betas
# Include extra effect from within group mixing
betas_eff[1,1,j] += get_extra_rho_effect(j,p,1) * nointervention_beta
# Calculate R effective by weighting betas by f(t) and the average infectious time
R_eff[j] = sum(np.matmul(betas_eff[:,:,j]*np.array([S_t_0[j],S_t_1[j]]),np.array([1-f_t,f_t])))
R_eff[j] = R_eff[j]*avg_time_infectious
# If a healthcare model, include the within hospital infections
if 'Healthcare' in p.model:
H = sum(np.transpose(solution_matrix[:,5:10,j]))
f_th = (H[1]*f)/(H[0]*(1-f)+H[1]*f)
hosp_beta = get_hosp_beta(j,p,sum(H))
avg_time_healthcare = (1 - p.p_IA) * p.p_IH * ((1 - p.p_HC) * p.t_HR + p.p_HC * (p.t_HC + (1 - p.p_CD) * (p.t_CL + p.t_LR) + p.p_CD * p.t_CD))
R_eff[j] += avg_time_healthcare*hosp_beta
return R_eff,betas_eff,f_t_all
def create_legend(fig,models: list,ls_styles: dict,linewidth = 2,legend_entries = {},legend_keywords = {}):
legend_entries['EWs'] = mpl.lines.Line2D([],[],color = 'k',ls = ':', linewidth = linewidth)
legend_entries['nEWs'] = mpl.lines.Line2D([],[],color = 'k',ls = '-', linewidth = linewidth)
for model in models:
ls,color = ls_styles[model]
legend_entries[model_labels[model]] = mpl.lines.Line2D([],[],color = color,ls = ls, linewidth = linewidth)
legend = fig.legend(legend_entries.values(),legend_entries.keys(),**legend_keywords)
return legend
def plot_grid(theta_values, f_values, metric, models, t_lockdown, max_days = 500,R_0=2.5,I0=20,adjust_beta_hosp = 1,
figsize = (30,10),model_dependent_rho = None, linewidth = 2,font_size = 30):
update_font_size(font_size)
fig, axes = plt.subplots(len(f_values), len(theta_values),figsize=(14*len(theta_values),12*len(f_values)))
# Plotting LS and color for each model
ls_styles = {'No Structure':('-',[0.5,0.5,0.5,0.5]),
'Cashier':('-',[0,0,1,1]),
'Healthcare':('-','r'),
'USPS':('-','g'),
'Healthcare_Sigmoid':('-','purple')}
for i,f in enumerate(f_values):
for j,theta in enumerate(theta_values):
for model_index,model in enumerate(models):
rho = model_dependent_rho[model]
# run model and reshape output
solution,p = run_model(theta = theta,
R_0 = R_0,
t_lockdown = t_lockdown,
model = model,
n_infected = I0,
n_exposed = I0,
prop_essential = f,
max_days = max_days,
rho = rho,
adjust_beta_hosp = adjust_beta_hosp*f)
solution_matrix = np.reshape(solution.y,[p.num_groups,len(p.compartment_names),p.max_days])
# Get average cumulative infection rate across population
avg_infection_rate = sum(sum(solution_matrix[:,1:,:]))/8e6
# Get proportion of population in each compartment relative to group size
# Note, here N is hardcoded as 8e6 (approx. NYC)
solution_matrix[0,:,:] = solution_matrix[0,:,:]/(8e6*(1-f))
solution_matrix[1,:,:] = solution_matrix[1,:,:]/(8e6*f)
# Get the effective R, betas, and g(t)
R_effective, Betas, G = get_R_effective(solution_matrix,p)
R_effective = R_effective[:-1]
R_effective[:t_lockdown+1] = np.nan
# Get color for each model but don't use LS
ls, color = ls_styles[model]
ls = '-'
if len(f_values) > 1:
if len(theta_values) > 1:
ax = axes[i][j]
else:
ax = axes[i]
else:
if len(theta_values) > 1:
ax = axes[j]
else:
ax = axes
plt.sca(ax)
if metric == 'Average_Infection':
# Plot cumulative infection rate in population
plt.plot(np.arange(0,max_days),avg_infection_rate,ls = '-',color=color,label=model,linewidth = linewidth)
plt.ylabel("Cumulative Infections")
elif metric == 'Essential_Infection':
plt.plot(np.arange(0,max_days),sum(solution_matrix[1,1:,:]),ls = ':',color=color,linewidth = linewidth,alpha = 1)
plt.ylabel("Cumulative Infections\nin Essential Workers")
elif metric == 'Nonessential_Infection':
plt.plot(np.arange(0,max_days),sum(solution_matrix[0,1:,:]),ls = ':',color=color,linewidth = linewidth,alpha = 1)
plt.ylabel("Cumulative Infections\nin Non-Essential Workers")
elif metric == 'All_Infection':
plt.plot(np.arange(0,max_days),sum(solution_matrix[1,1:,:]),ls = ':',color=color,linewidth = linewidth,alpha = 1)
plt.plot(np.arange(0,max_days),sum(solution_matrix[0,1:,:]),ls = '-',color=color,linewidth = linewidth,alpha = 1)
plt.ylabel("Cumulative Infections")
elif metric == 'R(t)':
plt.plot(np.arange(0,max_days),R_effective,color = color,linewidth = linewidth,ls=ls,label=model,rho = 1)
plt.ylabel("R(t)")
else: raise ValueError('Metric "{}" not understood'.format(metric))
plt.xlabel("Days")
plt.xlim([0,200])
if metric != 'R(t)':
plt.fill_betweenx(y=[1e-7,1],x1=0,x2=p.social_distancing_time,color=[0.8,0.8,0.8,0.5])
# Assign a column label
if i == 0 and len(theta_values) > 1:
assign_column_label_grid(ax,R_0*(theta),font_size = font_size)
# Assign a row column
if j == 0 and len(f_values) > 1:
assign_row_label_grid(ax,f,font_size = font_size)
legend = create_legend(fig = fig,
ls_styles = ls_styles,
models = models,
linewidth = linewidth,
legend_keywords = {'loc':'upper left',
'bbox_to_anchor':(0.82, 0.55),
'framealpha': 1,
'edgecolor':'k'})
for ax in axes.flat:
ax.tick_params(width = 5, length = 10)
ax.tick_params(width = 2.5, length = 5,which='minor')
plt.subplots_adjust(wspace=0.3)
plt.subplots_adjust(hspace=0.3)
plt.savefig('Figure2.png',bbox_extra_artists = legend)
plt.show()
```

In [33]:

```
# Make plot for Figure 2
desired_Reff = np.array([0.5,0.9,1.5])
R_0 = R0_from_doubling_time(args_default,3)
print('t',(1-desired_Reff/R_0))
plt.rc('legend', fontsize=30)
plot_grid(theta_values = desired_Reff/R_0,
f_values = [0.05],
max_days = 400,
metric = 'All_Infection',
R_0 = R_0,
I0 = 20,
t_lockdown = 59,
models = ['Cashier','USPS','Healthcare','No Structure'],
model_dependent_rho = {'No Structure': 0,
'Cashier': 0,
'Healthcare': 0.5,
'USPS': 0.5,
'Healthcare_Sigmoid': 0},
linewidth = 4,
adjust_beta_hosp = 9.5649565,
font_size=35)
```

In [34]:

```
# Make a similar plot to Figure 2, but change f
plt.rc('legend', fontsize=30)
plot_grid(theta_values = desired_Reff/R_0,
f_values = [0.10],
max_days = 400,
metric = 'All_Infection',
R_0 = R_0,
I0 = 20,
t_lockdown = 59,
models = ['Cashier','USPS','Healthcare','No Structure'],
model_dependent_rho = {'No Structure': 0,
'Cashier': 0,
'Healthcare': 0.5,
'USPS': 0.5,
'Healthcare_Sigmoid': 0},
linewidth = 4,
adjust_beta_hosp = 9.5649565,
font_size=35)
```

In [35]:

```
# Make a similar plot to Figure 2, but change t_lockdown
plt.rc('legend', fontsize=30)
plot_grid(theta_values = desired_Reff/R_0,
f_values = [0.05],
max_days = 400,
metric = 'All_Infection',
R_0 = R_0,
I0 = 20,
t_lockdown = 59 - 10,
models = ['Cashier','USPS','Healthcare','No Structure'],
model_dependent_rho = {'No Structure': 0,
'Cashier': 0,
'Healthcare': 0.5,
'USPS': 0.5,
'Healthcare_Sigmoid': 0},
linewidth = 4,
adjust_beta_hosp = 9.5649565,
font_size=35)
```

In [36]:

```
# Run several models over different values of theta and f
# Here, we look at theta between 0 and 1, and f between 0 and 0.25
def run_iterations(model_names, theta_values = np.linspace(0,1,21),f_values = np.linspace(0,0.25,21),
R_0 = 2.5, t_lockdown = 59, n_infected = 50, n_exposed = 50, max_days = 1000,
change_in_p_IH = 1, change_in_p_IA = 1, doubling_time = None, rho = 9.5649565):
results_dict = {}
# For each of the specified models
for model in model_names:
print(model)
results_dict[model] = {}
if model == 'USPS' or model == 'Healthcare':
rho = 0.5
else:
rho = 0
# For each of the given f values (f = fraction essential)
for f in f_values:
# f can't be 0
if f == 0: continue
#print('\t',f)
results_dict[model][f] = {}
# for each of the given theta values (theta = effectiveness of social distancing)
for theta in theta_values:
# run the model and record the results
results_dict[model][f][theta] = run_model(theta = theta,
R_0 = R_0,
t_lockdown=t_lockdown,
model=model,
n_infected=n_infected,
n_exposed=n_exposed,
prop_essential=f,
max_days = max_days,
rho = rho,
change_in_p_IH = change_in_p_IH,
change_in_p_IA = change_in_p_IA,
adjust_beta_hosp = f*rho)
return results_dict
def assign_row_label_figure3(ax,model, pad = 25, xy = (0,0.5), model_labels = model_labels,font_size = 12):
# Adds the text to left of the subplot
ax.annotate(model_labels[model], xy=xy, xytext=(-ax.yaxis.labelpad - pad, 0),
xycoords=ax.yaxis.label, textcoords='offset points',
size=font_size, ha='right', va='center',rotation=90)
def assign_column_label_figure3(ax,text,index, pad = 5, xy = (0.5,1),font_size=12):
letter_maps = {0: 'A',
1: 'B',
2: 'C'}
label = letter_maps[index] + ': ' + text
pad = 5 # in points
ax.annotate(label+'\n', xy=xy, xytext=(0, pad),
xycoords='axes fraction', textcoords='offset points',
size=font_size, ha='center', va='baseline')
def plot_heatmap(sol_dict,model_names = ['No Structure','Cashier','USPS'],
Xname = '',Yname = '',font_size = 30,center_point = None):
update_font_size(font_size)
# The models we are plotting heatmaps for
model_keys = list(sol_dict.keys())
#
variable_1_keys = np.sort([i for i in sol_dict[model_keys[0]].keys()])
variable_2_keys = np.sort([i for i in sol_dict[model_keys[0]][variable_1_keys[0]].keys()])
fig, axes_all = plt.subplots(len(model_names),3,figsize=[45,10*len(model_names)])
cbar_ax = fig.add_axes([0.85, 0.15, 0.02, 0.7])
# Only considering 2 metrics
for metric in range(3):
# Name of the metrics and levels to plot contour curves for
metric_names = {0:('Total infections after one year',[center_point]),
1:('Infections among non-essentials',[]),
2:('Infections among essential workers',[])}
metric_name,levels = metric_names[metric]
# Get the axis for this model
axes = [row_axes[metric] for row_axes in axes_all]
# Create a matrix to build the heatmaps from
sol_array = np.zeros([len(model_keys),len(variable_1_keys),len(variable_2_keys)])
for i,model in enumerate(model_names):
# Reformat keys to look nice
axis_1_keys = []
axis_2_keys = []
# Get R0 to convert from theta to r0(theta)
R0 = R0_from_doubling_time(args_default,3)
for index,key1 in enumerate(variable_1_keys):
axis_1_keys.append(round(R0*(key1),2))
for index,key2 in enumerate(variable_2_keys):
axis_2_keys.append(key2)
# Fill in the matrix
model_dict = sol_dict[model]
for j,(var1,var1_dict) in enumerate(model_dict.items()):
for k,(var2,solutions) in enumerate(var1_dict.items()):
sol_array[i,j,k] = solutions[metric]
plt.sca(axes[i])
# Get the relevant part of the matrix and orientate the matrix correctly (tranpose then flip vertically)
data = np.flip(np.transpose(sol_array[i,:,:]))
data = pd.DataFrame(data)
data.index = np.flip(axis_2_keys)
data.columns = axis_1_keys
# Make a heatmap from it
HM = heatmap(data=data,
yticklabels=5,xticklabels=5,
cmap = 'Spectral',
center = center_point,
vmin = 0,
vmax = 1,
cbar_ax=cbar_ax)
CS = plt.contour(data,levels=levels,colors='r',linewidths=6)
plt.yticks(rotation=0)
axes[i].tick_params(width = 5,length = 10)
# Label heatmaps with model names
model_labels = {'Cashier' :'Public Facing',
'USPS' :'Non-Public Facing',
'Healthcare' :'Healthcare',
'No Structure' :'No essential workers',
'Healthcare_Sigmoid' :'Healthcare'}
if i == 0:
assign_column_label_figure3(ax=axes[i],text = metric_name, index = metric,font_size=font_size)
if metric == 0:
assign_row_label_figure3(ax=axes[i],model = model,font_size=font_size)
# Set variable names on the axis
if Yname:
axes[i].set_ylabel(Yname)
if Xname:
axes[i].set_xlabel(Xname)
cbar_ax.tick_params(width = 5,length = 10)
fig.subplots_adjust(hspace=0.3,wspace=0.3,right = 0.825)
plt.show()
```

In [37]:

```
keep_R0_constant = False
# Get R_0 assuming a doubling time of 3 days
R_0 = R0_from_doubling_time(args_default,3)
R_0_values = np.linspace(2,0,21)
# Get the values of theta and f to iterate over
# Only use values of theta such that R0(theta) < 2
theta_values=np.sort(R_0_values/R_0)
f_values = np.linspace(0,0.25,26)
all_results_phaseDiagram2 = run_iterations(theta_values=theta_values,
f_values = f_values,
R_0 = R_0,
n_exposed = 0,
n_infected = 20,
model_names = ['Cashier','USPS','Healthcare'])
# Make the plot for Figure 3
# This takes the resulting model solutions and converts them to the metrics we want
solved_results = get_metrics(all_results_phaseDiagram2)
# This plots the metrics we want as a heatmap.
sol, pp = run_model(t_lockdown=59,theta = (1/R_0),R_0 = R_0, n_exposed = 0, n_infected = 20, rho = 0,model = 'No Structure',prop_essential=0.05)
center_point = get_infection_rate(solution=sol,p=pp, time = 365)
plot_heatmap(solved_results,Yname=r'',Xname=r'',model_names = ['Cashier','USPS','Healthcare'],font_size = 42,center_point = center_point)
```

In [38]:

```
# Get the values of theta and f to iterate over
# Only use values of theta such that R0(theta) < 2
theta_values=np.sort(R_0_values/R_0)
f_values = np.linspace(0,0.25,26)
all_results_phaseDiagram2 = run_iterations(theta_values=theta_values,
f_values = f_values,
R_0 = R_0,
n_exposed = 0,
n_infected = 20,
t_lockdown=49, # Changing the lockdown date
model_names = ['Cashier','USPS','Healthcare'])
# Make the plot for Figure 3, but with the t_lockdown moved 10 days earlier
# This takes the resulting model solutions and converts them to the metrics we want
solved_results = get_metrics(all_results_phaseDiagram2)
# This plots the metrics we want as a heatmap.
sol, pp = run_model(t_lockdown=49,theta = (1/R_0),R_0 = R_0, n_exposed = 0, n_infected = 20, rho = 0,model = 'No Structure',prop_essential=0.05)
center_point = get_infection_rate(solution=sol,p=pp, time = 365)
plot_heatmap(solved_results,Yname=r'',Xname=r'',model_names = ['Cashier','USPS','Healthcare'],font_size = 42,center_point = center_point)
```

In [39]:

```
# Assign row and column labels for the plots of the effective R over time
def assign_row_label_figure4(ax,model,letter,pad = 5, xy = (0,0.5),letter_pad = 0):
if not letter_pad:
letter_pad = pad*20
ax.annotate(letter, xy=xy, xytext=(-ax.yaxis.labelpad - letter_pad, 0),
xycoords=ax.yaxis.label, textcoords='offset points',
size='large', ha='center', va='center',rotation=0)
ax.annotate(model+'\n', xy=xy, xytext=(-ax.yaxis.labelpad - pad, 0),
xycoords=ax.yaxis.label, textcoords='offset points',
size='large', ha='center', va='center',rotation=90)
# Label is the value of R_effective
def assign_column_label_figure4(ax,model, pad = 5, xy = (0.5,1)):
text = model_labels[model]
pad = 5 # in points
ax.annotate(text+'\n', xy=xy, xytext=(0, pad),
xycoords='axes fraction', textcoords='offset points',
size='large', ha='center', va=
```