def handle_to_plot(kind, to_plot, n_cols, sim, check_ready=True):
''' Handle which quantities to plot '''
# Check that results are ready
if check_ready and not sim.results_ready:
errormsg = 'Cannot plot since results are not ready yet -- did you run the sim?'
raise RuntimeError(errormsg)
# If not specified or specified as a string, load defaults
if to_plot is None or isinstance(to_plot, str):
to_plot = cvd.get_default_plots(to_plot, kind=kind, sim=sim)
# If a list of keys has been supplied
if isinstance(to_plot, list):
to_plot_list = to_plot # Store separately
to_plot = sc.odict() # Create the dict
reskeys = sim.result_keys()
for reskey in to_plot_list:
name = sim.results[reskey].name if reskey in reskeys else sim.results['strain'][reskey].name
to_plot[name] = [reskey] # Use the result name as the key and the reskey as the value
to_plot = sc.odict(sc.dcp(to_plot)) # In case it's supplied as a dict
# Handle rows and columns -- assume 5 is the most rows we would want
n_plots = len(to_plot)
if n_cols is None:
max_rows = 4 # Assumption -- if desired, the user can override this by setting n_cols manually
n_cols = int((n_plots-1)//max_rows + 1) # This gives 1 column for 1-4, 2 for 5-8, etc.
n_rows,n_cols = sc.get_rows_cols(n_plots, ncols=n_cols) # Inconsistent naming due to Covasim/Matplotlib conventions
return to_plot, n_cols, n_rows
def handle_to_plot(kind, to_plot, n_cols, sim, check_ready=True):
''' Handle which quantities to plot '''
# Allow default kind to be overwritten by to_plot -- used by msim.plot()
if isinstance(to_plot, tuple):
kind, to_plot = to_plot # Split the tuple
# Check that results are ready
if check_ready and not sim.results_ready:
errormsg = 'Cannot plot since results are not ready yet -- did you run the sim?'
raise RuntimeError(errormsg)
# If it matches a result key, convert to a list
reskeys = sim.result_keys('main')
varkeys = sim.result_keys('variant')
allkeys = reskeys + varkeys
if to_plot in allkeys:
to_plot = sc.tolist(to_plot)
# If not specified or specified as another string, load defaults
if to_plot is None or isinstance(to_plot, str):
to_plot = cvd.get_default_plots(to_plot, kind=kind, sim=sim)
# If a list of keys has been supplied or constructed
if isinstance(to_plot, list):
to_plot_list = to_plot # Store separately
to_plot = sc.odict() # Create the dict
invalid = sc.autolist()
for reskey in to_plot_list:
if reskey in allkeys:
name = sim.results[
reskey].name if reskey in reskeys else sim.results[
'variant'][reskey].name
to_plot[name] = [
reskey
] # Use the result name as the key and the reskey as the value
else:
invalid += reskey
if len(invalid):
errormsg = f'The following key(s) are invalid:\n{sc.strjoin(invalid)}\n\nValid main keys are:\n{sc.strjoin(reskeys)}\n\nValid variant keys are:\n{sc.strjoin(varkeys)}'
raise sc.KeyNotFoundError(errormsg)
to_plot = sc.odict(sc.dcp(to_plot)) # In case it's supplied as a dict
# Handle rows and columns -- assume 5 is the most rows we would want
n_plots = len(to_plot)
if n_cols is None:
max_rows = 5 # Assumption -- if desired, the user can override this by setting n_cols manually
n_cols = int((n_plots - 1) // max_rows +
1) # This gives 1 column for 1-4, 2 for 5-8, etc.
n_rows, n_cols = sc.get_rows_cols(
n_plots, ncols=n_cols
) # Inconsistent naming due to Covasim/Matplotlib conventions
return to_plot, n_cols, n_rows
def test_samples(do_plot=False, verbose=True):
sc.heading('Samples distribution')
n = 200_000
nbins = 100
# Warning, must match utils.py!
choices = [
'uniform', 'normal', 'lognormal', 'normal_pos', 'normal_int',
'lognormal_int', 'poisson', 'neg_binomial'
]
if do_plot:
pl.figure(figsize=(20, 14))
# Run the samples
nchoices = len(choices)
nsqr, _ = sc.get_rows_cols(nchoices)
results = sc.objdict()
mean = 11
std = 7
low = 3
high = 9
normal_dists = [
'normal', 'normal_pos', 'normal_int', 'lognormal', 'lognormal_int'
]
for c, choice in enumerate(choices):
kw = {}
if choice in normal_dists:
par1 = mean
par2 = std
elif choice == 'neg_binomial':
par1 = mean
par2 = 1.2
kw['step'] = 0.1
elif choice == 'poisson':
par1 = mean
par2 = 0
elif choice == 'uniform':
par1 = low
par2 = high
else:
errormsg = f'Choice "{choice}" not implemented'
raise NotImplementedError(errormsg)
# Compute
results[choice] = cv.sample(dist=choice,
par1=par1,
par2=par2,
size=n,
**kw)
# Optionally plot
if do_plot:
pl.subplot(nsqr, nsqr, c + 1)
plotbins = np.unique(results[choice]) if (
choice == 'poisson' or '_int' in choice) else nbins
pl.hist(x=results[choice], bins=plotbins, width=0.8)
pl.title(f'dist={choice}, par1={par1}, par2={par2}')
with pytest.raises(NotImplementedError):
cv.sample(dist='not_found')
# Do statistical tests
tol = 1 / np.sqrt(
n / 50 / len(choices)
) # Define acceptable tolerance -- broad to avoid false positives
def isclose(choice, tol=tol, **kwargs):
key = list(kwargs.keys())[0]
ref = list(kwargs.values())[0]
npfunc = getattr(np, key)
value = npfunc(results[choice])
msg = f'Test for {choice:14s}: expecting {key:4s} = {ref:8.4f} ± {tol*ref:8.4f} and got {value:8.4f}'
if verbose:
print(msg)
assert np.isclose(value, ref, rtol=tol), msg
return True
# Normal
for choice in normal_dists:
isclose(choice, mean=mean)
if all([k not in choice
for k in ['_pos', '_int']]): # These change the variance
isclose(choice, std=std)
# Negative binomial
isclose('neg_binomial', mean=mean)
# Poisson
isclose('poisson', mean=mean)
isclose('poisson', var=mean)
# Uniform
isclose('uniform', mean=(low + high) / 2)
return results