def initialize(self, sim):
''' Fix days and store beta '''
if sc.isstring(self.days) or not sc.isiterable(self.days):
self.days = sc.promotetolist(self.days)
if isinstance(self.days, list):
for d, day in enumerate(self.days):
self.days[d] = sim.day(
day
) # Ensure it's an integer and not a string or something
self.days = sc.promotetoarray(self.days)
self.changes = sc.promotetoarray(self.changes)
if len(self.days) != len(self.changes):
errormsg = f'Number of days supplied ({len(self.days)}) does not match number of changes in beta ({len(self.changes)})'
raise ValueError(errormsg)
self.orig_betas = {}
self.layers = sc.promotetolist(self.layers, keepnone=True)
for lkey in self.layers:
if lkey is None:
self.orig_betas['overall'] = sim['beta']
else:
self.orig_betas[lkey] = sim['beta_layer'][lkey]
self.initialized = True
return
def validate_pars(self):
''' Some parameters can take multiple types; this makes them consistent '''
# Handle start day
start_day = self['start_day'] # Shorten
if start_day in [None, 0]: # Use default start day
start_day = dt.datetime(2020, 1, 1)
if not isinstance(start_day, dt.datetime):
start_day = sc.readdate(start_day)
self['start_day'] = start_day # Convert back
# Handle population data
popdata_choices = ['random', 'bayesian', 'data']
if sc.isnumber(self['usepopdata']) or isinstance(
self['usepopdata'],
bool): # Convert e.g. usepopdata=1 to 'bayesian'
self['usepopdata'] = popdata_choices[int(
self['usepopdata'])] # Choose one of these
if self['usepopdata'] not in popdata_choices:
choice = self['usepopdata']
choicestr = ', '.join(popdata_choices)
errormsg = f'Population data option "{choice}" not available; choices are: {choicestr}'
raise ValueError(errormsg)
# Handle interventions
self['interventions'] = sc.promotetolist(self['interventions'],
keepnone=False)
return
def story(self, uid, *args):
'''
Print out a short history of events in the life of the specified individual.
Args:
uid (int/list): the person or people whose story is being regaled
args (list): these people will tell their stories too
**Example**::
sim = cv.Sim(pop_type='hybrid', verbose=0)
sim.run()
sim.people.story(12)
sim.people.story(795)
'''
def label_lkey(lkey):
''' Friendly name for common layer keys '''
if lkey.lower() == 'a':
llabel = 'default contact'
if lkey.lower() == 'h':
llabel = 'household'
elif lkey.lower() == 's':
llabel = 'school'
elif lkey.lower() == 'w':
llabel = 'workplace'
elif lkey.lower() == 'c':
llabel = 'community'
else:
llabel = f'"{lkey}"'
return llabel
uids = sc.promotetolist(uid)
uids.extend(args)
for uid in uids:
p = self[uid]
sex = 'female' if p.sex == 0 else 'male'
intro = f'\nThis is the story of {uid}, a {p.age:.0f} year old {sex}'
print(intro)
total_contacts = 0
no_contacts = []
for lkey in p.contacts.keys():
llabel = label_lkey(lkey)
n_contacts = len(p.contacts[lkey])
total_contacts += n_contacts
if n_contacts:
print(
f'{uid} is connected to {n_contacts} people in the {llabel} layer'
)
else:
no_contacts.append(llabel)
if len(no_contacts):
nc_string = ', '.join(no_contacts)
print(f'{uid} has no contacts in the {nc_string} layer(s)')
print(f'{uid} has {total_contacts} contacts in total')
return
def remove_dynamic_contacts(self, dynamic_keys='c'):
''' Remove all contacts labeled as dynamic '''
dynamic_keys = sc.promotetolist(dynamic_keys)
for key in dynamic_keys:
if key in self.pars['contacts']:
self.contacts.pop(key)
return
def validate_pars(self):
''' Some parameters can take multiple types; this makes them consistent '''
# Handle start day
start_day = self['start_day'] # Shorten
if start_day in [None, 0]: # Use default start day
start_day = dt.date(2020, 1, 1)
elif sc.isstring(start_day):
start_day = sc.readdate(start_day)
if isinstance(start_day,dt.datetime):
start_day = start_day.date()
self['start_day'] = 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}
self['beta_layers'] = {'a':1.0}
# Handle population data
popdata_choices = ['random', 'microstructure', 'synthpops']
if sc.isnumber(self['pop_type']) or isinstance(self['pop_type'], bool): # Convert e.g. pop_type=1 to 'microstructure'
self['pop_type'] = popdata_choices[int(self['pop_type'])] # Choose one of these
if self['pop_type'] not in popdata_choices:
choice = self['pop_type']
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)
return
def update_contacts(self, dynamic_keys='c'):
''' Set dynamic contacts, by default, community ('c') '''
# Remove existing dynamic contacts
self.remove_dynamic_contacts()
# Figure out if anything needs to be done
dynamic_keys = sc.promotetolist(dynamic_keys)
for dynamic_key in dynamic_keys:
if dynamic_key in self.layer_keys():
pop_size = len(self)
n_contacts = self.pars['contacts'][dynamic_key]
beta = self.pars['beta_layer'][dynamic_key]
# Create new contacts
n_new = n_contacts * pop_size
new_contacts = {} # Initialize
new_contacts['p1'] = np.array(cvu.choose_r(max_n=pop_size,
n=n_new),
dtype=np.int32)
new_contacts['p2'] = np.array(cvu.choose_r(max_n=pop_size,
n=n_new),
dtype=np.int32)
# Set the things for the entire list
new_contacts['layer'] = np.array([dynamic_key] * n_new)
new_contacts['beta'] = np.array([beta] * n_new,
dtype=np.float32)
# Add to contacts
self.add_contacts(new_contacts, lkey=dynamic_key)
self.contacts[dynamic_key].validate()
return self.contacts
def date(self, ind, *args, dateformat=None):
'''
Convert an integer or list of integer simulation days to a date/list of dates.
Args:
ind (int, list, or array): the day(s) in simulation time
Returns:
dates (str or list): the date relative to the simulation start day, as an integer
**Example**::
sim.date(35) # Returns '2020-04-05'
'''
if sc.isnumber(ind): # If it's a number, convert it to a list
ind = sc.promotetolist(ind)
ind.extend(args)
if dateformat is None:
dateformat = '%Y-%m-%d'
dates = []
for i in ind:
tmp = self['start_day'] + dt.timedelta(days=int(i))
dates.append(tmp.strftime(dateformat))
# Return a string rather than a list if only one provided
if len(ind) == 1:
dates = dates[0]
return dates
def __init__(self, days, changes, layers=None, do_plot=None):
super().__init__(do_plot=do_plot)
self.days = days
self.changes = sc.promotetoarray(changes)
self.layers = sc.promotetolist(layers, keepnone=True)
self.orig_betas = None
self._store_args()
return
def day(obj, *args, start_day=None):
'''
Convert a string, date/datetime object, or int to a day (int), the number of
days since the start day. See also date() and daydiff(). Used primarily via
sim.day() rather than directly.
Args:
obj (str, date, int, or list): convert any of these objects to a day relative to the start day
args (list): additional days
start_day (str or date): the start day; if none is supplied, return days since 2020-01-01.
Returns:
days (int or str): the day(s) in simulation time
**Example**::
sim.day('2020-04-05') # Returns 35
'''
# Do not process a day if it's not supplied
if obj is None:
return None
if start_day is None:
start_day = '2020-01-01'
# Convert to list
if sc.isstring(obj) or sc.isnumber(obj) or isinstance(
obj, (dt.date, dt.datetime)):
obj = sc.promotetolist(obj) # Ensure it's iterable
elif isinstance(obj, np.ndarray):
obj = obj.tolist() # Convert to list if it's an array
obj.extend(args)
days = []
for d in obj:
if d is None:
days.append(d)
elif sc.isnumber(d):
days.append(int(d)) # Just convert to an integer
else:
try:
if sc.isstring(d):
d = sc.readdate(d).date()
elif isinstance(d, dt.datetime):
d = d.date()
d_day = (d - date(start_day)
).days # Heavy lifting -- actually compute the day
days.append(d_day)
except Exception as E:
errormsg = f'Could not interpret "{d}" as a date: {str(E)}'
raise ValueError(errormsg)
# Return an integer rather than a list if only one provided
if len(days) == 1:
days = days[0]
return days
def __init__(self, days, *args, **kwargs):
super().__init__(**kwargs) # Initialize the Intervention object
days = sc.promotetolist(days) # Combine multiple days
days.extend(args) # Include additional arguments, if present
self.days = days # Converted to integer representations
self.dates = None # String representations
self.start_day = None # Store the start date of the simulation
self.snapshots = sc.odict() # Store the actual snapshots
return
def initialize(self, sim):
self.days = process_days(sim, self.days)
self.changes = process_changes(sim, self.changes, self.days)
if self.layers is None:
self.layers = sim.layer_keys()
else:
self.layers = sc.promotetolist(self.layers)
self.contacts = cvb.Contacts(layer_keys=self.layers)
self.initialized = True
return
def __init__(self, days, changes, layers=None):
super().__init__()
self.days = sc.promotetoarray(days)
self.changes = sc.promotetoarray(changes)
self.layer_keys = sc.promotetolist(layers, keepnone=True)
if len(self.days) != len(self.changes):
errormsg = f'Number of days supplied ({len(self.days)}) does not match number of changes in beta ({len(self.changes)})'
raise ValueError(errormsg)
self.orig_betas = None
return
def process_graphs(figs):
jsons = []
for fig in sc.promotetolist(figs):
fig.update_layout(paper_bgcolor=bgcolor, plot_bgcolor=plotbg)
output = {'json': fig.to_json(), 'id': str(sc.uuid())}
d = json.loads(output['json'])
d['config'] = {'responsive': True}
output['json'] = json.dumps(d)
jsons.append(output)
return jsons
def add_student(self, pid):
"""
Add student / students to class.
Args:
pid (int): a student's id or a list of students' id
Returns:
None
"""
self.students.update(sc.promotetolist(pid))
def add_teacher(self, pid):
"""
Add teacher / teachers to class.
Args:
pid (int): a teacher's id or a list of teachers' id
Returns:
None
"""
self.teachers.update(sc.promotetolist(pid))
def resize(self, pop_size=None, keys=None):
''' Resize arrays if any mismatches are found '''
if pop_size is None:
pop_size = len(self)
self.pop_size = pop_size
if keys is None:
keys = self.keys()
keys = sc.promotetolist(keys)
for key in keys:
self[key].resize(pop_size, refcheck=False)
return
def to_json(self, filename=None, keys=None, tostring=False, indent=2, verbose=False, *args, **kwargs):
'''
Export results as JSON.
Args:
filename (str): if None, return string; else, write to file
keys (str or list): attributes to write to json (default: results, parameters, and summary)
tostring (bool): if not writing to file, whether to write to string (alternative is sanitized dictionary)
indent (int): if writing to file, how many indents to use per nested level
verbose (bool): detail to print
args (list): passed to savejson()
kwargs (dict): passed to savejson()
Returns:
A unicode string containing a JSON representation of the results,
or writes the JSON file to disk
**Examples**::
json = sim.to_json()
sim.to_json('results.json')
sim.to_json('summary.json', keys='summary')
'''
# Handle keys
if keys is None:
keys = ['results', 'pars', 'summary']
keys = sc.promotetolist(keys)
# Convert to JSON-compatible format
d = {}
for key in keys:
if key == 'results':
resdict = self.export_results(for_json=True)
d['results'] = resdict
elif key in ['pars', 'parameters']:
pardict = self.export_pars()
d['parameters'] = pardict
elif key == 'summary':
d['summary'] = dict(sc.dcp(self.summary))
else:
try:
d[key] = sc.sanitizejson(getattr(self, key))
except Exception as E:
errormsg = f'Could not convert "{key}" to JSON: {str(E)}; continuing...'
print(errormsg)
if filename is None:
output = sc.jsonify(d, tostring=tostring, indent=indent, verbose=verbose, *args, **kwargs)
else:
output = sc.savejson(filename=filename, obj=d, indent=indent, *args, **kwargs)
return output
def date(self, ind, *args, dateformat=None, as_date=False):
'''
Convert one or more integer days of simulation time to a date/list of dates --
by default returns a string, or returns a datetime Date object if as_date is True.
See also cv.date(), which provides a partly overlapping set of date conversion
features.
Args:
ind (int, list, or array): the index day(s) in simulation time (NB: strings and date objects are accepted, and will be passed unchanged)
args (list): additional day(s)
dateformat (str): the format to return the date in
as_date (bool): whether to return as a datetime date instead of a string
Returns:
dates (str, Date, or list): the date(s) corresponding to the simulation day(s)
**Examples**::
sim = cv.Sim()
sim.date(34) # Returns '2020-04-04'
sim.date([34, 54]) # Returns ['2020-04-04', '2020-04-24']
sim.date([34, '2020-04-24']) # Returns ['2020-04-04', '2020-04-24']
sim.date(34, 54, as_date=True) # Returns [datetime.date(2020, 4, 4), datetime.date(2020, 4, 24)]
'''
# Handle inputs
if not isinstance(
ind, list
): # If it's a number, string, or dateobj, convert it to a list
ind = sc.promotetolist(ind)
ind.extend(args)
if dateformat is None:
dateformat = '%Y-%m-%d'
# Do the conversion
dates = []
for raw in ind:
if sc.isnumber(raw):
date_obj = cvm.date(self['start_day'],
as_date=True) + dt.timedelta(days=int(raw))
else:
date_obj = cvm.date(raw, as_date=True)
if as_date:
dates.append(date_obj)
else:
dates.append(date_obj.strftime(dateformat))
# Return a string rather than a list if only one provided
if len(ind) == 1:
dates = dates[0]
return dates
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 inds2dates(self, inds, dateformat=None):
''' Convert a set of indices to a set of dates '''
if sc.isnumber(inds): # If it's a number, convert it to a list
inds = sc.promotetolist(inds)
if dateformat is None:
dateformat = '%b-%d'
dates = []
for ind in inds:
tmp = self['start_day'] + dt.timedelta(days=int(ind))
dates.append(tmp.strftime(dateformat))
return dates
def process_days(sim, days):
'''
Ensure lists of days are in consistent format. Used by change_beta, clip_edges,
and some analyzers. If day is 'end' or -1, use the final day of the simulation.
'''
if sc.isstring(days) or not sc.isiterable(days):
days = sc.promotetolist(days)
for d, day in enumerate(days):
if day in ['end', -1]:
day = sim['end_day']
days[d] = sim.day(
day) # Ensure it's an integer and not a string or something
days = sc.promotetoarray(days)
return days
def initialize(self, sim):
''' Fix days and store beta '''
self.days = process_days(sim, self.days)
self.changes = process_changes(sim, self.changes, self.days)
self.layers = sc.promotetolist(self.layers, keepnone=True)
self.orig_betas = {}
for lkey in self.layers:
if lkey is None:
self.orig_betas['overall'] = sim['beta']
else:
self.orig_betas[lkey] = sim['beta_layer'][lkey]
self.initialized = True
return
def plot_scens(to_plot=None, scens=None, do_save=None, fig_path=None, 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, n_cols=None, grid=False, commaticks=True, setylim=True,
log_scale=False, colors=None, labels=None, do_show=None, sep_figs=False, fig=None, ax=None, **kwargs):
''' Plot the results of a scenario -- see Scenarios.plot() for documentation '''
# Handle inputs
args = handle_args(fig_args=fig_args, plot_args=plot_args, scatter_args=scatter_args, axis_args=axis_args, fill_args=fill_args,
legend_args=legend_args, show_args=show_args, date_args=date_args, mpl_args=mpl_args, **kwargs)
to_plot, n_cols, n_rows = handle_to_plot('scens', to_plot, n_cols, sim=scens.base_sim, check_ready=False) # Since this sim isn't run
fig, figs = create_figs(args, sep_figs, fig, ax)
# Do the plotting
default_colors = sc.gridcolors(ncolors=len(scens.sims))
for pnum,title,reskeys in to_plot.enumitems():
ax = create_subplots(figs, fig, ax, n_rows, n_cols, pnum, args.fig, sep_figs, log_scale, title)
reskeys = sc.promotetolist(reskeys) # In case it's a string
for reskey in reskeys:
resdata = scens.results[reskey]
for snum,scenkey,scendata in resdata.enumitems():
sim = scens.sims[scenkey][0] # Pull out the first sim in the list for this scenario
strain_keys = sim.result_keys('strain')
if reskey in strain_keys:
ns = sim['n_strains']
strain_colors = sc.gridcolors(ns)
for strain in range(ns):
res_y = scendata.best[strain,:]
color = strain_colors[strain] # Choose the color
label = 'wild type' if strain == 0 else sim['strains'][strain - 1].label
ax.fill_between(scens.tvec, scendata.low[strain,:], scendata.high[strain,:], color=color, **args.fill) # Create the uncertainty bound
ax.plot(scens.tvec, res_y, label=label, c=color, **args.plot) # Plot the actual line
if args.show['data']:
plot_data(sim, ax, reskey, args.scatter, color=color) # Plot the data
else:
res_y = scendata.best
color = set_line_options(colors, scenkey, snum, default_colors[snum]) # Choose the color
label = set_line_options(labels, scenkey, snum, scendata.name) # Choose the label
ax.fill_between(scens.tvec, scendata.low, scendata.high, color=color, **args.fill) # Create the uncertainty bound
ax.plot(scens.tvec, res_y, label=label, c=color, **args.plot) # Plot the actual line
if args.show['data']:
plot_data(sim, ax, reskey, args.scatter, color=color) # Plot the data
if args.show['interventions']:
plot_interventions(sim, ax) # Plot the interventions
if args.show['ticks']:
reset_ticks(ax, sim, args.date) # Optionally reset tick marks (useful for e.g. plotting weeks/months)
if args.show['legend']:
title_grid_legend(ax, title, grid, commaticks, setylim, args.legend, pnum==0) # Configure the title, grid, and legend -- only show legend for first
return tidy_up(fig, figs, sep_figs, do_save, fig_path, do_show, args)
def map_entries(json, location, which):
'''
Find a match between the JSON file and the provided location(s).
Args:
json (list or dict): the data being loaded
location (list or str): the list of locations to pull from
which (str): either 'age' for age data or 'household' for household size distributions
'''
# The data have slightly different formats: list of dicts or just a dict
if which == 'age':
countries = [entry["country"].lower()
for entry in json] # Pull out available countries
else:
countries = [key.lower() for key in json.keys()]
# Set parameters
if location is None:
location = countries
else:
location = sc.promotetolist(location)
# Define a mapping for common mistakes
mapping = get_country_aliases()
mapping = {key.lower(): val.lower() for key, val in mapping.items()}
entries = {}
for loc in location:
lloc = loc.lower()
if lloc not in countries and lloc in mapping:
lloc = mapping[lloc]
try:
ind = countries.index(lloc)
if which == 'age':
entry = json[ind]
else:
entry = list(json.values())[ind]
entries[loc] = entry
except ValueError as E:
suggestions = sc.suggest(loc, countries, n=4)
if suggestions:
errormsg = f'Location "{loc}" not recognized, did you mean {suggestions}? ({str(E)})'
else:
errormsg = f'Location "{loc}" not recognized ({str(E)})'
raise ValueError(errormsg)
return entries
def validate_pars(self):
''' Some parameters can take multiple types; this makes them consistent '''
# Handle start day
start_day = self['start_day'] # Shorten
if start_day in [None, 0]: # Use default start day
start_day = dt.date(2020, 1, 1)
elif sc.isstring(start_day):
start_day = sc.readdate(start_day)
if isinstance(start_day, dt.datetime):
start_day = start_day.date()
self['start_day'] = 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 cvm.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 cvm.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 cvm.KeyNotFoundError(errormsg)
# Handle interventions
self['interventions'] = sc.promotetolist(self['interventions'],
keepnone=False)
return
def date(obj, *args, **kwargs):
'''
Convert a string or a datetime object to a date object. To convert to an integer
from the start day, use sim.date() instead.
Args:
obj (str, date, datetime): the object to convert
args (str, date, datetime): additional objects to convert
Returns:
dates (date or list): either a single date object, or a list of them
**Examples**::
cv.date('2020-04-05') # Returns datetime.date(2020, 4, 5)
'''
# Convert to list
if sc.isstring(obj) or sc.isnumber(obj) or isinstance(
obj, (dt.date, dt.datetime)):
obj = sc.promotetolist(obj) # Ensure it's iterable
obj.extend(args)
dates = []
for d in obj:
try:
if type(
d
) == dt.date: # Do not use isinstance, since must be the exact type
pass
elif sc.isstring(d):
d = sc.readdate(d).date()
elif isinstance(d, dt.datetime):
d = d.date()
else:
errormsg = f'Could not interpret "{d}" of type {type(d)} as a date'
raise TypeError(errormsg)
dates.append(d)
except Exception as E:
errormsg = f'Conversion of "{d}" to a date failed: {str(E)}'
raise ValueError(errormsg)
# Return an integer rather than a list if only one provided
if len(dates) == 1:
dates = dates[0]
return dates
def day(self, day, *args):
'''
Convert a string, date/datetime object, or int to a day (int).
Args:
day (str, date, int, or list): convert any of these objects to a day relative to the simulation's start day
Returns:
days (int or str): the day(s) in simulation time
**Example**::
sim.day('2020-04-05') # Returns 35
'''
# Do not process a day if it's not supplied
if day is None:
return None
# Convert to list
if sc.isstring(day) or sc.isnumber(day) or isinstance(
day, (dt.date, dt.datetime)):
day = sc.promotetolist(day) # Ensure it's iterable
day.extend(args)
days = []
for d in day:
if sc.isnumber(d):
days.append(int(d)) # Just convert to an integer
else:
try:
if sc.isstring(d):
d = sc.readdate(d).date()
elif isinstance(d, dt.datetime):
d = d.date()
d_day = (d - self['start_day']).days
days.append(d_day)
except Exception as E:
errormsg = f'Could not interpret "{d}" as a date: {str(E)}'
raise ValueError(errormsg)
# Return an integer rather than a list if only one provided
if len(days) == 1:
days = days[0]
return days
def process_days(sim, days, return_dates=False):
'''
Ensure lists of days are in consistent format. Used by change_beta, clip_edges,
and some analyzers. If day is 'end' or -1, use the final day of the simulation.
Optionally return dates as well as days.
'''
if sc.isstring(days) or not sc.isiterable(days):
days = sc.promotetolist(days)
for d,day in enumerate(days):
if day in ['end', -1]:
day = sim['end_day']
days[d] = sim.day(day) # Ensure it's an integer and not a string or something
days = np.sort(sc.promotetoarray(days)) # Ensure they're an array and in order
if return_dates:
dates = [sim.date(day) for day in days] # Store as date strings
return days, dates
else:
return days
def plot_scens(scens, data_check=None, dday=None, test_data=None, to_plot=None, do_save=None, fig_path=None, fig_args=None, plot_args=None,
scatter_args=None, axis_args=None, fill_args=None, legend_args=None, show_args=None,
as_dates=True, dateformat=None, interval=None, n_cols=None, font_size=18, font_family=None,
grid=False, commaticks=True, setylim=True, log_scale=False, colors=None, labels=None,
do_show=True, sep_figs=False, fig=None):
''' Plot the results of a scenario -- see Scenarios.plot() for documentation '''
# Handle inputs
args = handle_args(fig_args, plot_args, scatter_args, axis_args, fill_args, legend_args)
to_plot, n_cols, n_rows = handle_to_plot('scens', to_plot, n_cols, sim=scens.base_sim)
fig, figs, ax = create_figs(args, font_size, font_family, sep_figs, fig)
was_plot=False
# Do the plotting
default_colors = sc.gridcolors(ncolors=len(scens.sims))
for pnum,title,reskeys in to_plot.enumitems():
ax = create_subplots(figs, fig, ax, n_rows, n_cols, pnum, args.fig, sep_figs, log_scale, title)
reskeys = sc.promotetolist(reskeys) # In case it's a string
for reskey in reskeys:
resdata = scens.results[reskey]
for snum,scenkey,scendata in resdata.enumitems():
sim = scens.sims[scenkey][0] # Pull out the first sim in the list for this scenario
res_y = scendata.best
color = set_line_options(colors, scenkey, snum, default_colors[snum]) # Choose the color
label = set_line_options(labels, scenkey, snum, scendata.name) # Choose the label
if dday is not None:
ax.fill_between(scens.tvec[-dday:], smooth(scendata.low)[-dday:], smooth(scendata.high)[-dday:], color=color, **args.fill) # Create the uncertainty bound
ax.plot(scens.tvec[-dday:], smooth(res_y)[-dday:], label=label, c=color, **args.plot) # Plot the actual line
else:
ax.fill_between(scens.tvec, smooth(scendata.low), smooth(scendata.high), color=color, **args.fill) # Create the uncertainty bound
ax.plot(scens.tvec, smooth(res_y), label=label, c=color, **args.plot)
if data_check is not None and not was_plot:
ax.scatter(range(data_check.index.size), data_check[reskey], label='Data')
was_plot=True
elif data_check is None:
if args.show['data']:
plot_data(sim, ax, reskey, args.scatter, dday=dday, test_data=test_data, color=color) # Plot the data
if args.show['interventions']:
plot_interventions(sim, ax) # Plot the interventions
if args.show['ticks']:
reset_ticks(ax, sim, interval, as_dates, dateformat) # Optionally reset tick marks (useful for e.g. plotting weeks/months)
if args.show['legend']:
title_grid_legend(ax, title, grid, commaticks, setylim, args.legend, pnum==0) # Configure the title, grid, and legend -- only show legend for first
return tidy_up(fig, figs, sep_figs, do_save, fig_path, do_show)
def _resize_arrays(self, new_size=None, keys=None):
''' Resize arrays if any mismatches are found '''
# Handle None or tuple input
if new_size is None:
new_size = len(self)
pop_size = new_size if not isinstance(new_size, tuple) else new_size[1]
self.pars['pop_size'] = pop_size
# Reset sizes
if keys is None:
keys = self.keys()
keys = sc.promotetolist(keys)
for key in keys:
self[key].resize(
new_size, refcheck=False
) # Don't worry about cross-references to the arrays
return
