def convert_df_to_json_array(df, cols, int_cols=None):
"""
Convert desired data from a pandas dataframe into a json array.
Args:
df (pandas dataframe) : the dataframe with data
cols (list) : list of the columns to convert to the json array format
int_cols (str or list) : a str or list of columns to convert to integer values
Returns:
array: An array version of the pandas dataframe to be added to synthpops
json data objects.
"""
df = df[cols]
# make into a list to iterate over
int_cols = sc.tolist(int_cols)
# some columns as ints
df = df.astype({k: int for k in int_cols})
# make an array of arrays --- dtype=object to preserve each columns type
arr = df.to_numpy(dtype=object).tolist()
return arr
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 member_ages(self, age_by_uid, subgroup_member_uids=None):
"""Return the ages of members in the layer group given the pop object."""
if len(age_by_uid) == 0:
print(
"age_by_uid is empty. Returning an empty array for member_ages."
)
return np.array([])
if subgroup_member_uids is None:
return np.array(age_by_uid[self['member_uids']])
else:
subgroup_member_uids = sc.tolist(subgroup_member_uids)
return np.array(age_by_uid[subgroup_member_uids])
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
def member_ages(self, age_by_uid, subgroup_member_uids=None):
"""
Return the ages of members in the layer group given the pop object.
Args:
age_by_uid (np.ndarray) : mapping of age to uid
subgroup_member_uids (np.ndarray, list) : subgroup of uids to return ages for
Returns:
nd.ndarray : ages of members in group or subgroup
"""
if len(age_by_uid) == 0:
print(
"age_by_uid is empty. Returning an empty array for member_ages."
)
return np.array([])
if subgroup_member_uids is None:
return np.array(age_by_uid[self['member_uids']])
else:
subgroup_member_uids = sc.tolist(subgroup_member_uids)
return np.array(age_by_uid[subgroup_member_uids])
def check_valid_probability_distributions(property_name,
valid_properties=None):
"""
Check that the property_name is a valid probability distribution.
Args:
property_name (str) : the property name
valid_properties (str or list) : a list of the valid probability distributions
Returns:
None.
"""
# check the property_name is in the list of valid_probability_distributions()
if valid_properties is None:
valid_properties = defaults.valid_probability_distributions
# if a single str, make into a list so next check will work
valid_properties = sc.tolist(valid_properties)
if property_name not in valid_properties: # pragma: no cover
raise NotImplementedError(
f"{property_name} is not one of the expected probability distributions. The list of expected probability distributions is {valid_properties}. If you wish to use this method on the attribute {property_name}, you can supply it as the parameter valid_properties={property_name}."
)
def savefig(filename=None, comments=None, fig=None, **kwargs):
'''
Wrapper for Matplotlib's ``pl.savefig()`` function which automatically stores
Covasim metadata in the figure.
By default, saves (git) information from both the Covasim version and the calling
function. Additional comments can be added to the saved file as well. These can
be retrieved via ``cv.get_png_metadata()`` (or ``sc.loadmetadata``). Metadata can
also be stored for PDF, but cannot be automatically retrieved.
Args:
filename (str/list): name of the file to save to (default, timestamp); can also be a list of names
comments (str/dict): additional metadata to save to the figure
fig (fig/list): figure to save (by default, current one); can also be a list of figures
kwargs (dict): passed to ``fig.savefig()``
**Example**::
cv.Sim().run().plot()
cv.savefig()
'''
# Handle inputs
dpi = kwargs.pop('dpi', 150)
metadata = kwargs.pop('metadata', {})
if fig is None:
fig = pl.gcf()
figlist = sc.tolist(fig)
if filename is None: # pragma: no cover
now = sc.getdate(dateformat='%Y-%b-%d_%H.%M.%S')
filename = f'covasim_{now}.png'
filenamelist = sc.tolist(filename)
if len(figlist) != len(filenamelist):
errormsg = f'You have supplied {len(figlist)} figures and {len(filenamelist)} filenames: these must be the same length'
raise ValueError(errormsg)
metadata = {}
metadata['Covasim version'] = cvv.__version__
gitinfo = git_info()
for key,value in gitinfo['covasim'].items():
metadata[f'Covasim {key}'] = value
for key,value in gitinfo['called_by'].items():
metadata[f'Covasim caller {key}'] = value
metadata['Covasim current time'] = sc.getdate()
metadata['Covasim calling file'] = sc.getcaller()
if comments:
metadata['Covasim comments'] = comments
# Loop over the figures (usually just one)
for thisfig, thisfilename in zip(figlist, filenamelist):
# Handle different formats
lcfn = thisfilename.lower() # Lowercase filename
if lcfn.endswith('pdf') or lcfn.endswith('svg'):
metadata = {'Keywords':str(metadata)} # PDF and SVG doesn't support storing a dict
# Save the figure
thisfig.savefig(thisfilename, dpi=dpi, metadata=metadata, **kwargs)
return filename