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,