def validate_layer_pars(self):
'''
Handle layer parameters, since they need to be validated after the population
creation, rather than before.
'''
# First, try to figure out what the layer keys should be and perform basic type checking
layer_keys = self.layer_keys()
layer_pars = cvpar.layer_pars # The names of the parameters that are specified by layer
for lp in layer_pars:
val = self[lp]
if sc.isnumber(
val
): # It's a scalar instead of a dict, assume it's all contacts
self[lp] = {k: val for k in layer_keys}
# Handle key mismaches
for lp in layer_pars:
lp_keys = set(self.pars[lp].keys())
if not lp_keys == set(layer_keys):
errormsg = f'Layer parameters have inconsistent keys with the layer keys {layer_keys}:'
for lp2 in layer_pars: # Fail on first error, but re-loop to list all of them
errormsg += f'\n{lp2} = ' + ', '.join(self.pars[lp].keys())
raise sc.KeyNotFoundError(errormsg)
# Handle mismatches with the population
if self.people is not None:
pop_keys = set(self.people.contacts.keys())
if pop_keys != set(layer_keys):
errormsg = f'Please update your parameter keys {layer_keys} to match population keys {pop_keys}. You may find sim.reset_layer_pars() helpful.'
raise sc.KeyNotFoundError(errormsg)
return
def validate_pars(self):
''' Some parameters can take multiple types; this makes them consistent '''
# Handle types
for key in ['pop_size', 'pop_infected', 'pop_size', 'n_days']:
self[key] = int(self[key])
# Handle start day
start_day = self['start_day'] # Shorten
if start_day in [None, 0]: # Use default start day
start_day = '2020-03-01'
self['start_day'] = cvm.date(start_day)
# Handle contacts
contacts = self['contacts']
if sc.isnumber(
contacts
): # It's a scalar instead of a dict, assume it's all contacts
self['contacts'] = {'a': contacts}
# Handle key mismaches
beta_layer_keys = set(self.pars['beta_layer'].keys())
contacts_keys = set(self.pars['contacts'].keys())
quar_eff_keys = set(self.pars['quar_eff'].keys())
if not (beta_layer_keys == contacts_keys == quar_eff_keys):
errormsg = f'Layer parameters beta={beta_layer_keys}, contacts={contacts_keys}, quar_eff={quar_eff_keys} have inconsistent keys'
raise sc.KeyNotFoundError(errormsg)
if self.people is not None:
pop_keys = set(self.people.contacts.keys())
if pop_keys != beta_layer_keys:
errormsg = f'Please update your parameter keys {beta_layer_keys} to match population keys {pop_keys}. You may find sim.reset_layer_pars() helpful.'
raise sc.KeyNotFoundError(errormsg)
# Handle population data
popdata_choices = ['random', 'hybrid', 'clustered', 'synthpops']
choice = self['pop_type']
if choice not in popdata_choices:
choicestr = ', '.join(popdata_choices)
errormsg = f'Population type "{choice}" not available; choices are: {choicestr}'
raise sc.KeyNotFoundError(errormsg)
# Handle interventions
self['interventions'] = sc.promotetolist(self['interventions'],
keepnone=False)
for i, interv in enumerate(self['interventions']):
if isinstance(
interv, dict
): # It's a dictionary representation of an intervention
self['interventions'][i] = cvi.InterventionDict(**interv)
return
def handle_args(fig_args=None, plot_args=None, scatter_args=None, axis_args=None, fill_args=None,
legend_args=None, date_args=None, show_args=None, mpl_args=None, **kwargs):
''' Handle input arguments -- merge user input with defaults; see sim.plot for documentation '''
# Set defaults
defaults = sc.objdict()
defaults.fig = sc.objdict(figsize=(10, 8))
defaults.plot = sc.objdict(lw=1.5, alpha= 0.7)
defaults.scatter = sc.objdict(s=20, marker='s', alpha=0.7, zorder=0)
defaults.axis = sc.objdict(left=0.10, bottom=0.08, right=0.95, top=0.95, wspace=0.30, hspace=0.30)
defaults.fill = sc.objdict(alpha=0.2)
defaults.legend = sc.objdict(loc='best', frameon=False)
defaults.date = sc.objdict(as_dates=True, dateformat=None, interval=None, rotation=None, start_day=None, end_day=None)
defaults.show = sc.objdict(data=True, ticks=True, interventions=True, legend=True)
defaults.mpl = sc.objdict(dpi=None, fontsize=None, fontfamily=None) # Use Covasim global defaults
# Handle directly supplied kwargs
for dkey,default in defaults.items():
keys = list(kwargs.keys())
for kw in keys:
if kw in default.keys():
default[kw] = kwargs.pop(kw)
# Merge arguments together
args = sc.objdict()
args.fig = sc.mergedicts(defaults.fig, fig_args)
args.plot = sc.mergedicts(defaults.plot, plot_args)
args.scatter = sc.mergedicts(defaults.scatter, scatter_args)
args.axis = sc.mergedicts(defaults.axis, axis_args)
args.fill = sc.mergedicts(defaults.fill, fill_args)
args.legend = sc.mergedicts(defaults.legend, legend_args)
args.date = sc.mergedicts(defaults.date, fill_args)
args.show = sc.mergedicts(defaults.show, show_args)
args.mpl = sc.mergedicts(defaults.mpl, mpl_args)
# If unused keyword arguments remain, raise an error
if len(kwargs):
notfound = sc.strjoin(kwargs.keys())
valid = sc.strjoin(sorted(set([k for d in defaults.values() for k in d.keys()]))) # Remove duplicates and order
errormsg = f'The following keywords could not be processed:\n{notfound}\n\n'
errormsg += f'Valid keywords are:\n{valid}\n\n'
errormsg += 'For more precise plotting control, use fig_args, plot_args, etc.'
raise sc.KeyNotFoundError(errormsg)
# Handle what to show
show_keys = defaults.show.keys()
args.show = {k:True for k in show_keys}
if show_args in [True, False]: # Handle all on or all off
args.show = {k:show_args for k in show_keys}
else:
args.show = sc.mergedicts(args.show, show_args)
# Handle global Matplotlib arguments
args.mpl_orig = sc.objdict()
for key,value in args.mpl.items():
if value is not None:
args.mpl_orig[key] = cvset.options.get(key)
cvset.options.set(key, value)
return args
def InterventionDict(which, pars):
'''
Generate an intervention from a dictionary. Although a function, it acts
like a class, since it returns a class instance.
**Example**::
interv = cv.InterventionDict(which='change_beta', pars={'days': 30, 'changes': 0.5, 'layers': None})
'''
mapping = dict(
dynamic_pars=dynamic_pars,
sequence=sequence,
change_beta=change_beta,
clip_edges=clip_edges,
test_num=test_num,
test_prob=test_prob,
contact_tracing=contact_tracing,
)
try:
IntervClass = mapping[which]
except:
available = ', '.join(mapping.keys())
errormsg = f'Only interventions "{available}" are available in dictionary representation, not "{which}"'
raise sc.KeyNotFoundError(errormsg)
intervention = IntervClass(**pars)
return intervention
def update_pars(self, pars=None, create=False):
'''
Update internal dict with new pars.
Args:
pars (dict): the parameters to update (if None, do nothing)
create (bool): if create is False, then raise a KeyNotFoundError if the key does not already exist
'''
if pars is not None:
if not isinstance(pars, dict):
raise TypeError(
f'The pars object must be a dict; you supplied a {type(pars)}'
)
if not hasattr(self, 'pars'):
self.pars = pars
if not create:
available_keys = list(self.pars.keys())
mismatches = [
key for key in pars.keys() if key not in available_keys
]
if len(mismatches):
errormsg = f'Key(s) {mismatches} not found; available keys are {available_keys}'
raise sc.KeyNotFoundError(errormsg)
self.pars.update(pars)
return
def set_prognoses(self):
'''
Set the prognoses for each person based on age during initialization. Need
to reset the seed because viral loads are drawn stochastically.
'''
pars = self.pars # Shorten
if 'prognoses' not in pars or 'rand_seed' not in pars:
errormsg = 'This people object does not have the required parameters ("prognoses" and "rand_seed"). Create a sim (or parameters), then do e.g. people.set_pars(sim.pars).'
raise sc.KeyNotFoundError(errormsg)
def find_cutoff(age_cutoffs, age):
'''
Find which age bin each person belongs to -- e.g. with standard
age bins 0, 10, 20, etc., ages [5, 12, 4, 58] would be mapped to
indices [0, 1, 0, 5]. Age bins are not guaranteed to be uniform
width, which is why this can't be done as an array operation.
'''
return np.nonzero(age_cutoffs <= age)[0][-1] # Index of the age bin to use
cvu.set_seed(pars['rand_seed'])
progs = pars['prognoses'] # Shorten the name
inds = np.fromiter((find_cutoff(progs['age_cutoffs'], this_age) for this_age in self.age), dtype=cvd.default_int, count=len(self)) # Convert ages to indices
self.symp_prob[:] = progs['symp_probs'][inds] # Probability of developing symptoms
self.severe_prob[:] = progs['severe_probs'][inds]*progs['comorbidities'][inds] # Severe disease probability is modified by comorbidities
self.crit_prob[:] = progs['crit_probs'][inds] # Probability of developing critical disease
self.death_prob[:] = progs['death_probs'][inds] # Probability of death
self.rel_sus[:] = progs['sus_ORs'][inds] # Default susceptibilities
self.rel_trans[:] = progs['trans_ORs'][inds] * cvu.sample(**self.pars['beta_dist'], size=len(inds)) # Default transmissibilities, with viral load drawn from a distribution
return
def make_edgelist(self, contacts):
'''
Parse a list of people with a list of contacts per person and turn it
into an edge list.
'''
# Parse the list
lkeys = self.layer_keys()
new_contacts = Contacts(layer_keys=lkeys)
for lkey in lkeys:
new_contacts[lkey]['p1'] = [] # Person 1 of the contact pair
new_contacts[lkey]['p2'] = [] # Person 2 of the contact pair
try:
for p, cdict in enumerate(contacts):
for lkey, p_contacts in cdict.items():
n = len(p_contacts) # Number of contacts
new_contacts[lkey]['p1'].extend([p] *
n) # e.g. [4, 4, 4, 4]
new_contacts[lkey]['p2'].extend(
p_contacts) # e.g. [243, 4538, 7,19]
except KeyError:
lkeystr = ', '.join(lkeys)
errormsg = f'Layer "{lkey}" could not be loaded since it was not among parameter keys "{lkeystr}". Please update manually or via sim.reset_layer_pars().'
raise sc.KeyNotFoundError(errormsg)
# Turn into a dataframe
for lkey in lkeys:
new_layer = Layer()
for ckey, value in new_contacts[lkey].items():
new_layer[ckey] = np.array(value, dtype=new_layer.meta[ckey])
new_contacts[lkey] = new_layer
return new_contacts
def __getitem__(self, key):
''' Allow sim['par_name'] instead of sim.pars['par_name'] '''
try:
return self.pars[key]
except:
all_keys = '\n'.join(list(self.pars.keys()))
errormsg = f'Key "{key}" not found; available keys:\n{all_keys}'
raise sc.KeyNotFoundError(errormsg)
def set_option(key=None, value=None, set_global=True, **kwargs):
'''
Set a parameter or parameters. Use ``cv.options.set('defaults')`` to reset all
values to default, or ``cv.options.set(dpi='default')`` to reset one parameter
to default. See ``cv.options.help()`` for more information.
Args:
key (str): the parameter to modify, or 'defaults' to reset eerything to default values
value (varies): the value to specify; use None or 'default' to reset to default
set_global (bool): if true (default), sets plotting options globally (rather than just for Covasim)
kwargs (dict): if supplied, set multiple key-value pairs
Options are (see also ``cv.options.help()``):
- verbose: default verbosity for simulations to use
- font_size: the font size used for the plots
- font_family: the font family/face used for the plots
- dpi: the overall DPI for the figure
- show: whether to show figures
- close: whether to close the figures
- backend: which Matplotlib backend to use
- precision: the arithmetic to use in calculations
- numba_parallel: whether to parallelize Numba
**Examples**::
cv.options.set('font_size', 18)
cv.options.set(font_size=18, show=False, backend='agg', precision=64)
cv.options.set('defaults') # Reset to default options
'''
if key is not None:
kwargs = sc.mergedicts(kwargs, {key: value})
reload_required = False
# Reset to defaults
if key in ['default', 'defaults']:
kwargs = orig_options # Reset everything to default
# Reset options
for key, value in kwargs.items():
if key not in options:
keylist = orig_options.keys()
keys = '\n'.join(keylist)
errormsg = f'Option "{key}" not recognized; options are "defaults" or:\n{keys}\n\nSee help(cv.options.set) for more information.'
raise sc.KeyNotFoundError(errormsg)
else:
if value in [None, 'default']:
value = orig_options[key]
options[key] = value
if key in numba_keys:
reload_required = True
if key in matplotlib_keys and set_global:
set_matplotlib_global(key, value)
if reload_required:
reload_numba()
return
def __setitem__(self, key, value):
''' Ditto '''
if key in self.pars:
self.pars[key] = value
else:
all_keys = '\n'.join(list(self.pars.keys()))
errormsg = f'Key "{key}" not found; available keys:\n{all_keys}'
raise sc.KeyNotFoundError(errormsg)
return
def parse(self, variant=None, label=None):
''' Unpack variant information, which may be given as either a string or a dict '''
# Option 1: variants can be chosen from a list of pre-defined variants
if isinstance(variant, str):
choices, mapping = cvpar.get_variant_choices()
known_variant_pars = cvpar.get_variant_pars()
label = variant.lower()
for txt in ['.', ' ', 'variant', 'variant', 'voc']:
label = label.replace(txt, '')
if label in mapping:
label = mapping[label]
variant_pars = known_variant_pars[label]
else:
errormsg = f'The selected variant "{variant}" is not implemented; choices are:\n{sc.pp(choices, doprint=False)}'
raise NotImplementedError(errormsg)
# Option 2: variants can be specified as a dict of pars
elif isinstance(variant, dict):
default_variant_pars = cvpar.get_variant_pars(default=True)
default_keys = list(default_variant_pars.keys())
# Parse label
variant_pars = variant
label = variant_pars.pop(
'label', label) # Allow including the label in the parameters
if label is None:
label = 'custom'
# Check that valid keys have been supplied...
invalid = []
for key in variant_pars.keys():
if key not in default_keys:
invalid.append(key)
if len(invalid):
errormsg = f'Could not parse variant keys "{sc.strjoin(invalid)}"; valid keys are: "{sc.strjoin(cvd.variant_pars)}"'
raise sc.KeyNotFoundError(errormsg)
# ...and populate any that are missing
for key in default_keys:
if key not in variant_pars:
variant_pars[key] = default_variant_pars[key]
else:
errormsg = f'Could not understand {type(variant)}, please specify as a dict or a predefined variant:\n{sc.pp(choices, doprint=False)}'
raise ValueError(errormsg)
# Set label and parameters
self.label = label
self.p = sc.objdict(variant_pars)
return
def set_matplotlib_global(key, value):
''' Set a global option for Matplotlib -- not for users '''
import pylab as pl
if value: # Don't try to reset any of these to a None value
if key == 'font_size': pl.rc('font', size=value)
elif key == 'font_family': pl.rc('font', family=value)
elif key == 'dpi': pl.rc('figure', dpi=value)
elif key == 'backend': pl.switch_backend(value)
else: raise sc.KeyNotFoundError(f'Key {key} not found')
return
def __getitem__(self, key):
''' Lightweight odict -- allow indexing by number, with low performance '''
try:
return super().__getitem__(key)
except KeyError as KE:
try: # Assume it's an integer
dictkey = self.keys()[key]
return self[dictkey]
except:
raise sc.KeyNotFoundError(KE) # Raise the original error
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 get(self, key=None):
''' Retrieve a snapshot from the given key (int, str, or date) '''
if key is None:
key = self.days[0]
day = cvm.day(key, start_day=self.start_day)
date = cvm.date(day, start_date=self.start_day, as_date=False)
if date in self.snapshots:
snapshot = self.snapshots[date]
else:
dates = ', '.join(list(self.snapshots.keys()))
errormsg = f'Could not find snapshot date {date} (day {day}): choices are {dates}'
raise sc.KeyNotFoundError(errormsg)
return snapshot
def checkmem(unit='mb', fmt='0.2f', start=0, to_string=True):
''' For use with logger, check current memory usage '''
process = psutil.Process(os.getpid())
mapping = {'b': 1, 'kb': 1e3, 'mb': 1e6, 'gb': 1e9}
try:
factor = mapping[unit.lower()]
except KeyError:
raise sc.KeyNotFoundError(f'Unit {unit} not found')
mem_use = process.memory_info().rss / factor - start
if to_string:
output = f'{mem_use:{fmt}} {unit.upper()}'
else:
output = mem_use
return output
def __init__(self, pars, do_plot=None):
super().__init__(do_plot=do_plot)
subkeys = ['days', 'vals']
for parkey in pars.keys():
for subkey in subkeys:
if subkey not in pars[parkey].keys():
errormsg = f'Parameter {parkey} is missing subkey {subkey}'
raise sc.KeyNotFoundError(errormsg)
if sc.isnumber(pars[parkey][subkey]): # Allow scalar values or dicts, but leave everything else unchanged
pars[parkey][subkey] = sc.promotetoarray(pars[parkey][subkey])
len_days = len(pars[parkey]['days'])
len_vals = len(pars[parkey]['vals'])
if len_days != len_vals:
raise ValueError(f'Length of days ({len_days}) does not match length of values ({len_vals}) for parameter {parkey}')
self.pars = pars
self._store_args()
return
def __init__(self, pars, **kwargs):
super().__init__(**kwargs) # Initialize the Intervention object
self._store_args() # Store the input arguments so the intervention can be recreated
subkeys = ['days', 'vals']
for parkey in pars.keys():
for subkey in subkeys:
if subkey not in pars[parkey].keys():
errormsg = f'Parameter {parkey} is missing subkey {subkey}'
raise sc.KeyNotFoundError(errormsg)
if sc.isnumber(pars[parkey][subkey]): # Allow scalar values or dicts, but leave everything else unchanged
pars[parkey][subkey] = sc.promotetoarray(pars[parkey][subkey])
len_days = len(pars[parkey]['days'])
len_vals = len(pars[parkey]['vals'])
if len_days != len_vals:
raise ValueError(f'Length of days ({len_days}) does not match length of values ({len_vals}) for parameter {parkey}')
self.pars = pars
return
def _get_from_pars(pars, default=False, key=None, defaultkey='default'):
''' Helper function to get the right output from vaccine and variant functions '''
# If a string was provided, interpret it as a key and swap
if isinstance(default, str):
key, default = default, key
# Handle output
if key is not None:
try:
return pars[key]
except Exception as E:
errormsg = f'Key "{key}" not found; choices are: {sc.strjoin(pars.keys())}'
raise sc.KeyNotFoundError(errormsg) from E
elif default:
return pars[defaultkey]
else:
return pars
def set_pars(self, pars=None):
'''
Re-link the parameters stored in the people object to the sim containing it,
and perform some basic validation.
'''
if pars is None:
pars = {}
elif sc.isnumber(pars): # Interpret as a population size
pars = {'pop_size': pars} # Ensure it's a dictionary
orig_pars = self.__dict__.get(
'pars') # Get the current parameters using dict's get method
pars = sc.mergedicts(orig_pars, pars)
if 'pop_size' not in pars:
errormsg = f'The parameter "pop_size" must be included in a population; keys supplied were:\n{sc.newlinejoin(pars.keys())}'
raise sc.KeyNotFoundError(errormsg)
pars['pop_size'] = int(pars['pop_size'])
pars.setdefault('location', None)
self.pars = pars # Actually store the pars
return
def validate_popdict(popdict, pars, verbose=True):
'''
Check that the popdict is the correct type, has the correct keys, and has
the correct length
'''
# Check it's the right type
try:
popdict.keys(
) # Although not used directly, this is used in the error message below, and is a good proxy for a dict-like object
except Exception as E:
errormsg = f'The popdict should be a dictionary or cv.People object, but instead is {type(popdict)}'
raise TypeError(errormsg) from E
# Check keys and lengths
required_keys = ['uid', 'age', 'sex']
popdict_keys = popdict.keys()
pop_size = pars['pop_size']
for key in required_keys:
if key not in popdict_keys:
errormsg = f'Could not find required key "{key}" in popdict; available keys are: {sc.strjoin(popdict.keys())}'
sc.KeyNotFoundError(errormsg)
actual_size = len(popdict[key])
if actual_size != pop_size:
errormsg = f'Could not use supplied popdict since key {key} has length {actual_size}, but all keys must have length {pop_size}'
raise ValueError(errormsg)
isnan = np.isnan(popdict[key]).sum()
if isnan:
errormsg = f'Population not fully created: {isnan:,} NaNs found in {key}. This can be caused by calling cv.People() instead of cv.make_people().'
raise ValueError(errormsg)
if ('contacts' not in popdict_keys) and (not hasattr(
popdict, 'contacts')) and verbose:
warnmsg = 'No contacts found. Please remember to add contacts before running the simulation.'
cvm.warn(warnmsg)
return
def validate_pars(self):
''' Some parameters can take multiple types; this makes them consistent '''
# Handle types
for key in ['pop_size', 'pop_infected', 'pop_size']:
try:
self[key] = int(self[key])
except Exception as E:
errormsg = f'Could not convert {key}={self[key]} of {type(self[key])} to integer'
raise ValueError(errormsg) from E
# Handle start day
start_day = self['start_day'] # Shorten
if start_day in [None, 0]: # Use default start day
start_day = '2020-03-01'
self['start_day'] = cvm.date(start_day)
# Handle end day and n_days
end_day = self['end_day']
n_days = self['n_days']
if end_day:
self['end_day'] = cvm.date(end_day)
n_days = cvm.daydiff(self['start_day'], self['end_day'])
if n_days <= 0:
errormsg = f"Number of days must be >0, but you supplied start={str(self['start_day'])} and end={str(self['end_day'])}, which gives n_days={n_days}"
raise ValueError(errormsg)
else:
self['n_days'] = int(n_days)
else:
if n_days:
self['n_days'] = int(n_days)
self['end_day'] = self.date(
n_days) # Convert from the number of days to the end day
else:
errormsg = f'You must supply one of n_days and end_day, not "{n_days}" and "{end_day}"'
raise ValueError(errormsg)
# Handle parameters specified by layer
# Try to figure out what the layer keys should be
layer_keys = None # e.g. household, school
layer_pars = ['beta_layer', 'contacts', 'iso_factor', 'quar_factor']
if self.people is not None:
layer_keys = set(self.people.contacts.keys())
elif isinstance(self['beta_layer'], dict):
layer_keys = list(
self['beta_layer'].keys()
) # Get keys from beta_layer since the "most required" layer parameter
else:
layer_keys = [
'a'
] # Assume this by default, corresponding to random/no layers
# Convert scalar layer parameters to dictionaries
for lp in layer_pars:
val = self[lp]
if sc.isnumber(
val
): # It's a scalar instead of a dict, assume it's all contacts
self[lp] = {k: val for k in layer_keys}
# Handle key mismaches
for lp in layer_pars:
lp_keys = set(self.pars[lp].keys())
if not lp_keys == set(layer_keys):
errormsg = f'Layer parameters have inconsistent keys:'
for lp2 in layer_pars: # Fail on first error, but re-loop to list all of them
errormsg += f'\n{lp2} = ' + ', '.join(self.pars[lp].keys())
raise sc.KeyNotFoundError(errormsg)
if self.people is not None:
pop_keys = set(self.people.contacts.keys())
if pop_keys != layer_keys:
errormsg = f'Please update your parameter keys {layer_keys} to match population keys {pop_keys}. You may find sim.reset_layer_pars() helpful.'
raise sc.KeyNotFoundError(errormsg)
# Handle population data
popdata_choices = ['random', 'hybrid', 'clustered', 'synthpops']
choice = self['pop_type']
if choice not in popdata_choices:
choicestr = ', '.join(popdata_choices)
errormsg = f'Population type "{choice}" not available; choices are: {choicestr}'
raise ValueError(errormsg)
# Handle interventions
self['interventions'] = sc.promotetolist(self['interventions'],
keepnone=False)
for i, interv in enumerate(self['interventions']):
if isinstance(
interv, dict
): # It's a dictionary representation of an intervention
self['interventions'][i] = cvi.InterventionDict(**interv)
return
def parse_synthpop(population, layer_mapping=None, community_contacts=0):
'''
Make a population using SynthPops, including contacts. Usually called automatically,
but can also be called manually. Either a simulation object or a population must
be supplied; if a population is supplied, transform it into the correct format;
otherwise, create the population and then transform it.
Args:
population (list): a pre-generated SynthPops population
layer_mapping (dict): a custom mapping from SynthPops layers to Covasim layers
community_contacts (int): create this many community contacts on average
New in version 1.10.0.
'''
# Handle layer mapping
default_layer_mapping = {
'H': 'h',
'S': 's',
'W': 'w',
'C': 'c',
'LTCF': 'l'
} # Remap keys from old names to new names
layer_mapping = sc.mergedicts(default_layer_mapping, layer_mapping)
# Create the basic lists
pop_size = len(population)
uids, ages, sexes, contacts = [], [], [], []
for uid, person in population.items():
uids.append(uid)
ages.append(person['age'])
sexes.append(person['sex'])
# Replace contact UIDs with ints
uid_mapping = {uid: u for u, uid in enumerate(uids)}
for uid in uids:
iid = uid_mapping[uid] # Integer UID
person = population.pop(uid)
uid_contacts = sc.dcp(person['contacts'])
int_contacts = {}
for spkey in uid_contacts.keys():
try:
lkey = layer_mapping[
spkey] # Map the SynthPops key into a Covasim layer key
except KeyError: # pragma: no cover
errormsg = f'Could not find key "{spkey}" in layer mapping "{layer_mapping}"'
raise sc.KeyNotFoundError(errormsg)
int_contacts[lkey] = []
for cid in uid_contacts[spkey]: # Contact ID
icid = uid_mapping[cid] # Integer contact ID
if icid > iid: # Don't add duplicate contacts
int_contacts[lkey].append(icid)
int_contacts[lkey] = np.array(int_contacts[lkey],
dtype=spu.default_int)
contacts.append(int_contacts)
# Add community contacts
c_contacts, _ = make_random_contacts(pop_size, {'c': community_contacts})
for i in range(int(pop_size)):
contacts[i]['c'] = c_contacts[i][
'c'] # Copy over community contacts -- present for everyone
# Finalize
popdict = {}
popdict['uid'] = np.array(list(uid_mapping.values()),
dtype=spu.default_int)
popdict['age'] = np.array(ages)
popdict['sex'] = np.array(sexes)
popdict['contacts'] = sc.dcp(contacts)
popdict['layer_keys'] = list(layer_mapping.values())
return popdict
def make_synthpop(sim=None,
population=None,
layer_mapping=None,
community_contacts=None,
**kwargs):
'''
Make a population using SynthPops, including contacts. Usually called automatically,
but can also be called manually. Either a simulation object or a population must
be supplied; if a population is supplied, transform it into the correct format;
otherise, create the population and then transform it.
Args:
sim (Sim): a Covasim simulation object
population (list): a pre-generated SynthPops population (otherwise, create a new one)
layer_mapping (dict): a custom mapping from SynthPops layers to Covasim layers
community_contacts (int): if a simulation is not supplied, create this many community contacts on average
kwargs (dict): passed to sp.make_population()
**Example**::
sim = cv.Sim(pop_type='synthpops')
sim.popdict = cv.make_synthpop(sim)
sim.run()
'''
try:
import synthpops as sp # Optional import
except ModuleNotFoundError as E:
errormsg = f'Please install the optional SynthPops module first, e.g. pip install synthpops' # Also caught in make_people()
raise ModuleNotFoundError(errormsg) from E
# Handle layer mapping
default_layer_mapping = {
'H': 'h',
'S': 's',
'W': 'w',
'C': 'c',
'LTCF': 'l'
} # Remap keys from old names to new names
layer_mapping = sc.mergedicts(default_layer_mapping, layer_mapping)
# Handle other input arguments
if population is None:
if sim is None:
errormsg = 'Either a simulation or a population must be supplied'
raise ValueError(errormsg)
pop_size = sim['pop_size']
population = sp.make_population(n=pop_size,
rand_seed=sim['rand_seed'],
**kwargs)
if community_contacts is None:
if sim is not None:
community_contacts = sim['contacts']['c']
else:
errormsg = 'If a simulation is not supplied, the number of community contacts must be specified'
raise ValueError(errormsg)
# Create the basic lists
pop_size = len(population)
uids, ages, sexes, contacts = [], [], [], []
for uid, person in population.items():
uids.append(uid)
ages.append(person['age'])
sexes.append(person['sex'])
# Replace contact UIDs with ints
uid_mapping = {uid: u for u, uid in enumerate(uids)}
for uid in uids:
iid = uid_mapping[uid] # Integer UID
person = population.pop(uid)
uid_contacts = sc.dcp(person['contacts'])
int_contacts = {}
for spkey in uid_contacts.keys():
try:
lkey = layer_mapping[
spkey] # Map the SynthPops key into a Covasim layer key
except KeyError:
errormsg = f'Could not find key "{spkey}" in layer mapping "{layer_mapping}"'
raise sc.KeyNotFoundError(errormsg)
int_contacts[lkey] = []
for cid in uid_contacts[spkey]: # Contact ID
icid = uid_mapping[cid] # Integer contact ID
if icid > iid: # Don't add duplicate contacts
int_contacts[lkey].append(icid)
int_contacts[lkey] = np.array(int_contacts[lkey],
dtype=cvd.default_int)
contacts.append(int_contacts)
# Add community contacts
c_contacts, _ = make_random_contacts(pop_size, {'c': community_contacts})
for i in range(int(pop_size)):
contacts[i]['c'] = c_contacts[i][
'c'] # Copy over community contacts -- present for everyone
# Finalize
popdict = {}
popdict['uid'] = np.array(list(uid_mapping.values()),
dtype=cvd.default_int)
popdict['age'] = np.array(ages)
popdict['sex'] = np.array(sexes)
popdict['contacts'] = sc.dcp(contacts)
popdict['layer_keys'] = list(layer_mapping.values())
return popdict
def set(self, key=None, value=None, **kwargs):
'''
Actually change the style. See ``cv.options.help()`` for more information.
Args:
key (str): the parameter to modify, or 'defaults' to reset everything to default values
value (varies): the value to specify; use None or 'default' to reset to default
kwargs (dict): if supplied, set multiple key-value pairs
**Example**::
cv.options.set(dpi=50) # Equivalent to cv.options(dpi=50)
'''
reload_required = False
# Reset to defaults
if key in ['default', 'defaults']:
kwargs = self.orig_options # Reset everything to default
# Handle other keys
elif key is not None:
kwargs = sc.mergedicts(kwargs, {key: value})
# Handle Jupyter
if 'jupyter' in kwargs.keys() and kwargs['jupyter']:
jupyter = kwargs['jupyter']
kwargs[
'returnfig'] = False # We almost never want to return figs from Jupyter, since then they appear twice
try: # This makes plots much nicer, but isn't available on all systems
if not os.environ.get(
'SPHINX_BUILD'
): # Custom check implemented in conf.py to skip this if we're inside Sphinx
try: # First try interactive
assert jupyter not in [
'default', 'retina'
] # Hack to intentionally go to the other part of the loop
from IPython import get_ipython
magic = get_ipython().magic
magic('%matplotlib widget')
except: # Then try retina
assert jupyter != 'default'
import matplotlib_inline
matplotlib_inline.backend_inline.set_matplotlib_formats(
'retina')
except:
pass
# Handle interactivity
if 'interactive' in kwargs.keys():
interactive = kwargs['interactive']
if interactive in [None, 'default']:
interactive = self.orig_options['interactive']
if interactive:
kwargs['show'] = True
kwargs['close'] = False
kwargs['backend'] = self.orig_options['backend']
else:
kwargs['show'] = False
kwargs['backend'] = 'agg'
# Reset options
for key, value in kwargs.items():
# Handle deprecations
rename = {'font_size': 'fontsize', 'font_family': 'font'}
if key in rename.keys():
from . import misc as cvm # Here to avoid circular import
oldkey = key
key = rename[oldkey]
warnmsg = f'Key "{oldkey}" is deprecated, please use "{key}" instead'
cvm.warn(warnmsg, FutureWarning)
if key not in self:
keylist = self.orig_options.keys()
keys = '\n'.join(keylist)
errormsg = f'Option "{key}" not recognized; options are "defaults" or:\n{keys}\n\nSee help(cv.options.set) for more information.'
raise sc.KeyNotFoundError(errormsg)
else:
if value in [None, 'default']:
value = self.orig_options[key]
self[key] = value
numba_keys = ['precision', 'numba_parallel', 'numba_cache'
] # Specify which keys require a reload
if key in numba_keys:
reload_required = True
if key in 'backend':
pl.switch_backend(value)
if reload_required:
reload_numba()
return
def with_style(self, style_args=None, use=False, **kwargs):
'''
Combine all Matplotlib style information, and either apply it directly
or create a style context.
To set globally, use ``cv.options.use_style()``. Otherwise, use ``cv.options.with_style()``
as part of a ``with`` block to set the style just for that block (using
this function outsde of a with block and with ``use=False`` has no effect, so
don't do that!). To set non-style options (e.g. warnings, verbosity) as
a context, see ``cv.options.context()``.
Args:
style_args (dict): a dictionary of style arguments
use (bool): whether to set as the global style; else, treat as context for use with "with" (default)
kwargs (dict): additional style arguments
Valid style arguments are:
- ``dpi``: the figure DPI
- ``font``: font (typeface)
- ``fontsize``: font size
- ``grid``: whether or not to plot gridlines
- ``facecolor``: color of the axes behind the plot
- any of the entries in ``pl.rParams``
**Examples**::
with cv.options.with_style(dpi=300): # Use default options, but higher DPI
pl.plot([1,3,6])
'''
# Handle inputs
rc = sc.dcp(
self.rc) # Make a local copy of the currently used settings
kwargs = sc.mergedicts(style_args, kwargs)
# Handle style, overwiting existing
style = kwargs.pop('style', None)
rc = self._handle_style(style, reset=False)
def pop_keywords(sourcekeys, rckey):
''' Helper function to handle input arguments '''
sourcekeys = sc.tolist(sourcekeys)
key = sourcekeys[0] # Main key
value = None
changed = self.changed(key)
if changed:
value = self[key]
for k in sourcekeys:
kwvalue = kwargs.pop(k, None)
if kwvalue is not None:
value = kwvalue
if value is not None:
rc[rckey] = value
return
# Handle special cases
pop_keywords('dpi', rckey='figure.dpi')
pop_keywords(['font', 'fontfamily', 'font_family'],
rckey='font.family')
pop_keywords(['fontsize', 'font_size'], rckey='font.size')
pop_keywords('grid', rckey='axes.grid')
pop_keywords('facecolor', rckey='axes.facecolor')
# Handle other keywords
for key, value in kwargs.items():
if key not in pl.rcParams:
errormsg = f'Key "{key}" does not match any value in Covasim options or pl.rcParams'
raise sc.KeyNotFoundError(errormsg)
elif value is not None:
rc[key] = value
# Tidy up
if use:
return pl.style.use(sc.dcp(rc))
else:
return pl.style.context(sc.dcp(rc))
def validate_pars(self):
''' Some parameters can take multiple types; this makes them consistent '''
# Handle types
for key in ['pop_size', 'pop_infected', 'pop_size']:
try:
self[key] = int(self[key])
except Exception as E:
errormsg = f'Could not convert {key}={self[key]} of {type(self[key])} to integer'
raise ValueError(errormsg) from E
# Handle start day
start_day = self['start_day'] # Shorten
if start_day in [None, 0]: # Use default start day
start_day = '2020-03-01'
self['start_day'] = cvm.date(start_day)
# Handle end day and n_days
end_day = self['end_day']
n_days = self['n_days']
if end_day:
self['end_day'] = cvm.date(end_day)
n_days = cvm.daydiff(self['start_day'], self['end_day'])
if n_days <= 0:
errormsg = f"Number of days must be >0, but you supplied start={str(self['start_day'])} and end={str(self['end_day'])}, which gives n_days={n_days}"
raise ValueError(errormsg)
else:
self['n_days'] = int(n_days)
else:
if n_days:
self['n_days'] = int(n_days)
self['end_day'] = self.date(
n_days) # Convert from the number of days to the end day
else:
errormsg = f'You must supply one of n_days and end_day, not "{n_days}" and "{end_day}"'
raise ValueError(errormsg)
# Handle contacts
contacts = self['contacts']
if sc.isnumber(
contacts
): # It's a scalar instead of a dict, assume it's all contacts
self['contacts'] = {'a': contacts}
# Handle key mismaches
beta_layer_keys = set(self.pars['beta_layer'].keys())
contacts_keys = set(self.pars['contacts'].keys())
quar_eff_keys = set(self.pars['quar_eff'].keys())
if not (beta_layer_keys == contacts_keys == quar_eff_keys):
errormsg = f'Layer parameters beta={beta_layer_keys}, contacts={contacts_keys}, quar_eff={quar_eff_keys} have inconsistent keys'
raise sc.KeyNotFoundError(errormsg)
if self.people is not None:
pop_keys = set(self.people.contacts.keys())
if pop_keys != beta_layer_keys:
errormsg = f'Please update your parameter keys {beta_layer_keys} to match population keys {pop_keys}. You may find sim.reset_layer_pars() helpful.'
raise sc.KeyNotFoundError(errormsg)
# Handle population data
popdata_choices = ['random', 'hybrid', 'clustered', 'synthpops']
choice = self['pop_type']
if choice not in popdata_choices:
choicestr = ', '.join(popdata_choices)
errormsg = f'Population type "{choice}" not available; choices are: {choicestr}'
raise ValueError(errormsg)
# Handle interventions
self['interventions'] = sc.promotetolist(self['interventions'],
keepnone=False)
for i, interv in enumerate(self['interventions']):
if isinstance(
interv, dict
): # It's a dictionary representation of an intervention
self['interventions'][i] = cvi.InterventionDict(**interv)
return
def make_synthpop(sim, generate=True, layer_mapping=None, **kwargs):
'''
Make a population using SynthPops, including contacts. Usually called automatically,
but can also be called manually.
Args:
sim (Sim): a Covasim simulation object
generate (bool): whether or not to generate a new population (otherwise, tries to load a pre-generated one)
layer_mapping (dict): a custom mapping from SynthPops layers to Covasim layers
kwars (dict): passed to sp.make_population()
**Example**::
sim = cv.Sim(pop_type='synthpops')
sim.popdict = cv.make_synthpop(sim)
sim.run()
'''
import synthpops as sp # Optional import
# Handle layer mapping
default_layer_mapping = {
'H': 'h',
'S': 's',
'W': 'w',
'C': 'c'
} # Remap keys from old names to new names
layer_mapping = sc.mergedicts(default_layer_mapping, layer_mapping)
# Handle other input arguments
pop_size = sim['pop_size']
population = sp.make_population(n=pop_size,
generate=generate,
rand_seed=sim['rand_seed'],
**kwargs)
uids, ages, sexes, contacts = [], [], [], []
for uid, person in population.items():
uids.append(uid)
ages.append(person['age'])
sexes.append(person['sex'])
# Replace contact UIDs with ints
uid_mapping = {uid: u for u, uid in enumerate(uids)}
for uid in uids:
iid = uid_mapping[uid] # Integer UID
person = population.pop(uid)
uid_contacts = sc.dcp(person['contacts'])
int_contacts = {}
for spkey in uid_contacts.keys():
try:
lkey = layer_mapping[
spkey] # Map the SynthPops key into a Covasim layer key
except KeyError:
errormsg = f'Could not find key "{spkey}" in layer mapping "{layer_mapping}"'
raise sc.KeyNotFoundError(errormsg)
int_contacts[lkey] = []
for cid in uid_contacts[spkey]: # Contact ID
icid = uid_mapping[cid] # Integer contact ID
if icid > iid: # Don't add duplicate contacts
int_contacts[lkey].append(icid)
int_contacts[lkey] = np.array(int_contacts[lkey],
dtype=cvd.default_int)
contacts.append(int_contacts)
# Add community contacts
c_contacts, _ = make_random_contacts(pop_size, {'c': sim['contacts']['c']})
for i in range(int(pop_size)):
contacts[i]['c'] = c_contacts[i][
'c'] # Copy over community contacts -- present for everyone
# Finalize
popdict = {}
popdict['uid'] = np.array(list(uid_mapping.values()),
dtype=cvd.default_int)
popdict['age'] = np.array(ages)
popdict['sex'] = np.array(sexes)
popdict['contacts'] = sc.dcp(contacts)
popdict['layer_keys'] = list(layer_mapping.values())
return popdict
def single_run(sim,
ind=0,
reseed=True,
noise=0.0,
noisepar=None,
verbose=None,
keep_people=False,
run_args=None,
sim_args=None,
**kwargs):
'''
Convenience function to perform a single simulation run. Mostly used for
parallelization, but can also be used directly.
Args:
sim (Sim): the sim instance to be run
ind (int): the index of this sim
reseed (bool): whether or not to generate a fresh seed for each run
noise (float): the amount of noise to add to each run
noisepar (string): the name of the parameter to add noise to
verbose (int): detail to print
run_args (dict): arguments passed to sim.run()
sim_args (dict): extra parameters to pass to the sim, e.g. 'n_infected'
kwargs (dict): also passed to the sim
Returns:
sim (Sim): a single sim object with results
**Example**::
import covasim as cv
sim = cv.Sim() # Create a default simulation
sim = cv.single_run(sim) # Run it, equivalent(ish) to sim.run()
'''
new_sim = sc.dcp(sim) # Copy the sim to avoid overwriting it
# Set sim and run arguments
sim_args = sc.mergedicts(sim_args, kwargs)
run_args = sc.mergedicts({'verbose': verbose}, run_args)
if verbose is None:
verbose = new_sim['verbose']
if not new_sim.label:
new_sim.label = f'Sim {ind:d}'
if reseed:
new_sim[
'rand_seed'] += ind # Reset the seed, otherwise no point of parallel runs
new_sim.set_seed()
# If the noise parameter is not found, guess what it should be
if noisepar is None:
noisepar = 'beta'
if noisepar not in sim.pars.keys():
raise sc.KeyNotFoundError(
f'Noise parameter {noisepar} was not found in sim parameters')
# Handle noise -- normally distributed fractional error
noiseval = noise * np.random.normal()
if noiseval > 0:
noisefactor = 1 + noiseval
else:
noisefactor = 1 / (1 - noiseval)
new_sim[noisepar] *= noisefactor
if verbose >= 1:
print(
f'Running a simulation using {new_sim["rand_seed"]} seed and {noisefactor} noise factor'
)
# Handle additional arguments
for key, val in sim_args.items():
print(f'Processing {key}:{val}')
if key in new_sim.pars.keys():
if verbose >= 1:
print(f'Setting key {key} from {new_sim[key]} to {val}')
new_sim[key] = val
else:
raise sc.KeyNotFoundError(
f'Could not set key {key}: not a valid parameter name')
# Run
new_sim.run(**run_args)
# Shrink the sim to save memory
if not keep_people:
new_sim.shrink()
return new_sim
def set_option(key=None, value=None, **kwargs):
'''
Set a parameter or parameters. Use ``cv.options.set('defaults')`` to reset all
values to default, or ``cv.options.set(dpi='default')`` to reset one parameter
to default. See ``cv.options.help()`` for more information.
Args:
key (str): the parameter to modify, or 'defaults' to reset everything to default values
value (varies): the value to specify; use None or 'default' to reset to default
kwargs (dict): if supplied, set multiple key-value pairs
Options are (see also ``cv.options.help()``):
- verbose: default verbosity for simulations to use
- font_size: the font size used for the plots
- font_family: the font family/face used for the plots
- dpi: the overall DPI for the figure
- show: whether to show figures
- close: whether to close the figures
- backend: which Matplotlib backend to use
- interactive: convenience method to set show, close, and backend
- precision: the arithmetic to use in calculations
- numba_parallel: whether to parallelize Numba
**Examples**::
cv.options.set('font_size', 18) # Larger font
cv.options.set(font_size=18, show=False, backend='agg', precision=64) # Larger font, non-interactive plots, higher precision
cv.options.set(interactive=False) # Turn off interactive plots
cv.options.set('defaults') # Reset to default options
'''
if key is not None:
kwargs = sc.mergedicts(kwargs, {key:value})
reload_required = False
# Reset to defaults
if key in ['default', 'defaults']:
kwargs = orig_options # Reset everything to default
# Handle interactivity
if 'interactive' in kwargs.keys():
interactive = kwargs['interactive']
if interactive in [None, 'default']:
interactive = orig_options['interactive']
if interactive:
kwargs['show'] = True
kwargs['close'] = False
kwargs['backend'] = orig_options['backend']
else:
kwargs['show'] = False
kwargs['close'] = True
kwargs['backend'] = 'agg'
# Reset options
for key,value in kwargs.items():
if key not in options:
keylist = orig_options.keys()
keys = '\n'.join(keylist)
errormsg = f'Option "{key}" not recognized; options are "defaults" or:\n{keys}\n\nSee help(cv.options.set) for more information.'
raise sc.KeyNotFoundError(errormsg)
else:
if value in [None, 'default']:
value = orig_options[key]
options[key] = value
if key in numba_keys:
reload_required = True
if key in matplotlib_keys:
set_matplotlib_global(key, value)
if reload_required:
reload_numba()
return
