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