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