def test_rand():
default_seed = 1
@nb.njit
def numba_rand():
return np.random.rand()
# Check that two consecutive numbers don't match
cova.set_seed(default_seed)
a = np.random.rand()
b = np.random.rand()
v = numba_rand()
w = numba_rand()
assert a != b
assert v != w
# Check that after resetting the seed, they do
cova.set_seed(default_seed)
c = np.random.rand()
x = numba_rand()
assert a == c
assert v == x
# Check that resetting with no argument, they don't again
cova.set_seed()
d = np.random.rand()
y = numba_rand()
cova.set_seed()
e = np.random.rand()
z = numba_rand()
assert d != e
assert y != z
return a
def run_scenario(index, v=None, which=None, doshrink=True, verbose=True):
if verbose:
print(f'Now running {index}:')
cv.set_seed(index + start_seed)
if which == 'a':
if v is None:
sweeps = construct1dsweeps()
v = sweeps[index]
f = get_f_a(v)
scenpars = scenconvert_a(f)
elif which == 'b':
f = get_f_b(v, index=index)
scenpars = scenconvert_b(f)
if resample:
eind = (index + entry_offset) % n_top_entries
else:
eind = index + entry_offset
labelvals = [f'{v:0.2f}' for v in f.values()]
labelstr = f'i={index} e={eind} f=' + ','.join(labelvals)
runinfo = sc.objdict()
runinfo.f = f
runinfo.scenpars = scenpars
runinfo.ind = index
runinfo.eind = eind
sim = cs.run_sim(index=eind,
scenpars=scenpars,
label=labelstr,
runinfo=runinfo,
verbose=verbose)
# Try to save progress, but don't worry about it
try:
sc.savetext(f'progress/{index}.ind', str(f))
except Exception as E:
print(f'Could not save file {index}: {str(E)}')
return sim
def modify_sim(sim, scenpars, label=None, runinfo=None):
''' Modify the simulation for the scenarios '''
print(
f' Note: modifying simulation {label} at day={sim.t}, date={sim.date(sim.t)}, scenpars:\n{scenpars}'
)
sim.sceninfo.scen_start = '2020-06-01'
sim.sceninfo.scale_down = '2020-07-01'
last_calib_day = sim.day(sim.sceninfo.calib_end)
first_scen_day = sim.day(sim.sceninfo.scen_start)
scale_down_day = sim.day(sim.sceninfo.scale_down)
n_days = 5
scale_range = np.arange(n_days)
ramp_up = np.linspace(0, 1, n_days)
ramp_down = 1 - ramp_up
for ilabel in ['clip_w', 'clip_c']:
interv = sim.get_interventions(ilabel)
valid_days = sc.findinds(interv.days <= last_calib_day)
# Do reopening: modify clip_edges
v0 = interv.changes[valid_days[-1]]
v1 = scenpars['reopen'][0]
v2 = scenpars['reopen'][1]
scale_up = ramp_down * v0 + ramp_up * v1
scale_down = ramp_down * v1 + ramp_up * v2
scale_days = np.append(scale_range + first_scen_day,
scale_range + scale_down_day)
scale_vals = np.append(scale_up, scale_down)
interv.days = np.append(interv.days[valid_days], scale_days)
interv.changes = np.append(interv.changes[valid_days], scale_vals)
# Change iso_factor and quar_factor
sim.pars['iso_factor'] = sc.dcp(scenpars['i_factor'])
sim.pars['quar_factor'] = sc.dcp(scenpars['q_factor'])
# Implement testing & tracing interventions
ctpars = {k: scenpars[k] for k in ['trace_probs', 'trace_time']}
tn = sim.get_interventions('tn')
if scenpars['which'] == 'actual':
tn.test_delay = scenpars.test_delay
sim['interventions'] += [
cv.contact_tracing(start_day=first_scen_day,
label='contact_tracing',
**ctpars)
]
elif scenpars['which'] == 'low':
offset = 7 # So it's constant
tn.daily_tests[first_scen_day - tn.start_day -
offset:] = scenpars['n_tests']
sim['interventions'] += [
cv.contact_tracing(start_day=first_scen_day,
label='contact_tracing',
**ctpars)
]
elif scenpars['which'] == 'high':
cv.set_seed(
sim['rand_seed']) # Just in case since we use a random number
tn.end_day = last_calib_day # End the test_num intervention
n_weeks = 4 # Ramp up the test_prob and contact_tracing interventions
days = first_scen_day + np.arange(n_weeks) * 7
tpramp = np.maximum(
0,
np.linspace(0.2, 0.7, n_weeks + 1)[1:] +
np.random.randn(n_weeks) * 0.1
) # Don't start from 0 and don't make it a perfectly smooth ramp
ctramp = np.linspace(0.0, 0.7, n_weeks + 1)[1:] # Don't start from 0
tppars = {
k: scenpars[k]
for k in [
'symp_prob', 'asymp_prob', 'symp_quar_prob', 'asymp_quar_prob',
'test_delay'
]
}
tplist = []
ctlist = []
for w in range(n_weeks):
tp = cv.test_prob(quar_policy=[0], **tppars, label=f'tp_{w}')
ct = cv.contact_tracing(label=f'ct_{w}', **ctpars)
tp.symp_prob *= tpramp[w]
tp.asymp_prob *= tpramp[w]
ct.trace_probs = sc.dcp(ct.trace_probs)
for key in ct.trace_probs.keys():
ct.trace_probs[key] *= ctramp[w]
tplist.append(tp)
ctlist.append(ct)
print(f'I AM {w} {ct.trace_probs}')
tpseq = cv.sequence(days=days, interventions=tplist)
ctseq = cv.sequence(days=days, interventions=ctlist)
tpseq.initialize(sim)
ctseq.initialize(sim)
sim['interventions'] += [tpseq, ctseq]
# Final tidying
sim.label = label
sim.runinfo = runinfo
sim.sceninfo.scenpars = scenpars
return sim
def plot():
fig = pl.figure(num='Fig. 4: Suppression scenarios', figsize=(figw, figh))
rx = 0.07
r1y = 0.74
rdx = 0.26
r1dy = 0.20
rδ = 0.30
r1δy = 0.29
r1ax = {}
for i in range(6):
xi = i % 3
yi = i // 3
r1ax[i] = pl.axes([rx + rδ * xi, r1y - r1δy * yi, rdx, r1dy])
r2y = 0.05
r2dy = rdx * figw / figh # To ensure square
r2ax = {}
for i in range(3):
r2ax[i] = pl.axes([rx + rδ * i, r2y, rdx, r2dy])
cax = pl.axes([0.96, r2y, 0.01, r2dy])
# Labels
lx = 0.015
pl.figtext(lx, r1y + r1dy + 0.02, 'a', fontsize=40)
pl.figtext(lx, r2y + r2dy + 0.02, 'b', fontsize=40)
slopes = sc.objdict().make(keys=df1map.keys(), vals=[])
slopes2 = sc.objdict().make(keys=df1map.keys(), vals=[])
for plotnum, key, label in df1map.enumitems():
cv.set_seed(plotnum)
ax = r1ax[plotnum]
for ei in eis:
theseinds = sc.findinds(~np.isnan(df1[key].values))
if sepinds:
ei_ok = sc.findinds(df1['eind'].values == ei)
theseinds = np.intersect1d(theseinds, ei_ok)
x = xvals[key][theseinds]
rawy = df1[ykey].values[theseinds]
y = rawy / kcpop * 100 if logy else rawy
xm = x.max()
if key in ['testdelay', 'trtime']:
xnoise = 0.01
ynoise = 0.05
else:
xnoise = 0
ynoise = 0
rndx = (np.random.randn(len(x))) * xm * xnoise
rndy = (np.random.randn(len(y))) * ynoise
ax.scatter(x + rndx,
y * (1 + rndy),
alpha=0.2,
c=[cols[ei]],
edgecolor='w')
# Calculate slopes
slopey = np.log(rawy) if logy else rawy
slopex = xvals[key].values[theseinds]
tmp, res = np.polyfit(slopex, slopey, 1, cov=True)
fitm, fitb = tmp
factor = np.exp(fitm * slopepoint[key] +
fitb) / slopedenom[key] * slopex.max()
slope = fitm * factor
slopes[key].append(slope)
# Calculate slopes, method 2 -- used for std calculation
X = sm.add_constant(slopex)
mod = sm.OLS(slopey, X)
res = mod.fit()
conf = res.conf_int(alpha=0.05, cols=None)
best = res.params[1] * factor
high = conf[1, 1] * factor
slopes2[key].append(res)
# Plot fit line
bflx = np.array([x.min(), x.max()])
ploty = np.log10(y) if logy else y
plotm, plotb = np.polyfit(x, ploty, 1)
bfly = plotm * bflx + plotb
if logy:
ax.semilogy(bflx, 10**(bfly), lw=3, c=c2)
else:
ax.plot(bflx, bfly, lw=3, c=c2)
plot_baseinds = False
if plot_baseinds:
default_x = default_xvals[key][baseinds]
default_rawy = df1[ykey].values[baseinds]
default_y = default_rawy / kcpop * 100 if logy else default_rawy
ax.scatter(default_x,
default_y,
marker='x',
alpha=1.0,
c=[cols[i]])
if verbose:
print(
f'Slope for {key:10s}: {np.mean(slopes[key]):0.3f} ± {high-best:0.3f}'
)
sc.boxoff(ax=ax)
ax.yaxis.set_major_formatter(mpl.ticker.ScalarFormatter())
if plotnum in [0, 3]:
if logy:
ax.set_ylabel('Attack rate (%)')
ax.set_yticks((0.01, 0.1, 1.0, 10, 100))
ax.set_yticklabels(('0.01', '0.1', '1.0', '10', '100'))
else:
ax.set_ylabel(r'$R_{e}$')
ax.set_xlabel(xlabelmap[key])
if key in ['iqfactor']:
ax.set_xticks(np.linspace(0, 100, 6))
elif key in ['trprob']:
ax.set_xticks(np.linspace(0, 10, 6))
elif key in ['testprob']:
ax.set_xticks(np.arange(7))
elif key in ['testqprob']:
ax.set_xticks(np.arange(7))
else:
ax.set_xticks(np.arange(8))
if logy:
ax.set_ylim([0.1, 100])
else:
ax.set_ylim([0.7, 1.7])
xl = ax.get_xlim()
if logy:
ypos1 = 150
ypos2 = 40
else:
ypos1 = 1.9
ypos2 = 1.7
xlpos = dict(
iqfactor=0.86,
testprob=0.13,
trprob=0.83,
testqprob=0.86,
testdelay=0.13,
trtime=0.00,
)
if key in ['iqfactor', 'testqprob', 'trprob']:
align = 'right'
else:
align = 'left'
ax.text((xl[0] + xl[1]) / 2,
ypos1,
label,
fontsize=26,
horizontalalignment='center')
ax.text(xlpos[key] * xl[1],
ypos2,
f'{abs(best):0.2f} ± {high-best:0.2f} {slopelabels[key]}',
color=pointcolor,
horizontalalignment=align)
ax.axvline(slopepoint[key],
ymax=0.83,
linestyle='--',
c=pointcolor,
alpha=0.5,
lw=2)
reop = [0.6, 0.8, 1.0]
for ri, r in enumerate(reop):
dfr = df2[df2['reopen'] == r]
im = plot_surface(r2ax[ri], dfr, col=ri, colval=r)
bbox = dict(facecolor='w', alpha=0.0, edgecolor='none')
pointsize = 150
if ri == 0:
dotx = 1900 * 1000 / kcpop
doty = 0.06
r2ax[ri].scatter([dotx], [doty],
c='k',
s=pointsize,
zorder=10,
marker='d')
r2ax[ri].text(dotx * 1.20,
doty * 1.50,
'Estimated\nconditions\non June 1',
bbox=bbox)
if ri == 1:
dotx = 3000 * 1000 / kcpop
doty = 0.20
r2ax[ri].scatter([dotx], [doty],
c=[pointcolor2],
s=pointsize,
zorder=10,
marker='d')
r2ax[ri].text(dotx * 1.10,
doty * 0.20,
'Estimated\nconditions\non July 15',
color=pointcolor2,
bbox=bbox)
if ri == 2:
dotx = 2.80 # 7200*1000/kcpop
doty = 0.70 # 0.66
r2ax[ri].scatter([dotx], [doty],
c=[pointcolor],
s=pointsize,
zorder=10,
marker='d')
r2ax[ri].text(dotx * 1.05,
doty * 1.05,
'High mobility,\nhigh test + trace\nscenario',
color=pointcolor,
bbox=bbox)
cbar = pl.colorbar(im, ticks=np.linspace(0.4, 1.6, 7), cax=cax)
cbar.ax.set_title('$R_{e}$', rotation=0, pad=20, fontsize=24)
return fig
s0.run(until=today)
s0.finalize() # Finalize the results
if dosave: s0.save(f'{resultsfolder}/nsw_{whattorun}_single.obj')
if doplot:
s0.plot(to_plot=to_plot,
do_save=True,
do_show=False,
fig_path=f'{figsfolder}/nsw_{whattorun}.png',
legend_args={'loc': 'upper left'},
axis_args={'hspace': 0.4},
interval=35)
# Full parameter/seed search
elif whattorun == 'fullfit':
fitsummary = sc.objdict()
cv.set_seed(1)
betas = cvu.sample(dist='normal', par1=0.025, par2=0.002, size=100)
fitsummary.betas = betas
fitsummary.mismatches = []
for bn, beta in enumerate(betas):
sc.blank()
print('---------------\n')
print(f'BN: {bn}, Beta: {beta},... ')
print('---------------\n')
s0 = make_sim(beta,
do_make_ints=True,
mask_uptake=0.3,
venue_trace_prob=0.5,
future_test_prob=0.9,
mask_eff=0.3,
load_pop=True,