def cum_subrecruits(ax, recruits, first_plot_day, color, linestyle):
"""Plot the cumulative sum of recruits to compare across many populations.
Args:
ax: A series of axis instances to plot on.
recruits: An xr.DataArray that representing the expected or
observed recruits. Must have a time dimension.
first_plot_day: An int representing the first date to plot.
color: A mpl color.
linestyle: A mpl linestyle.
"""
sel_recruits = plot_utils.unpack_participant_labels(recruits)
labels_to_plot = plot_utils.get_labels_to_plot(recruits)
num_plots = len(ax)
for i, label in enumerate(labels_to_plot):
a = ax[i]
participants = sel_recruits.sel(participant_label=label, drop=True)
a.set_title(label)
time_dim = plot_utils.find_time_dim(participants)
array_over_time(a, participants.cumsum(time_dim), first_plot_day,
{'color': color, 'ls': linestyle})
if i in [num_plots-2, num_plots-1]:
format_time_axis(a, 3, date_format='%b-%d')
else:
format_time_axis(a, 3, include_labels=False)
def array_over_time(ax, array_to_plot, first_plot_day=None, plot_kwargs={'color':'b', 'ls':'-'}):
"""Plot array_to_plot as a function of time.
If array has a `sample` or `scenario` dimension, then all samples will be plotted with a
low opacity (alpha) value.
Args:
ax: An axes instance to plot our data on.
array_to_plot: A xr.DataArray with a time dimension, and optionally a sample OR
scenario dimension
first_plot_day: Optional, a time coordinate indicating the first date to plot.
plot_kwargs: Optional, a dictionary with keyword arguments to pass to matplotlib.plot
"""
time_dim = plot_utils.find_time_dim(array_to_plot)
shaped_data = array_to_plot.transpose(time_dim, ...)
if first_plot_day in array_to_plot.coords[time_dim].values:
data = shaped_data.sel({time_dim:slice(first_plot_day, None)})
else:
data = shaped_data
if any(item in data.dims for item in ['sample', 'scenario', 'sample_flattened']):
alpha = 0.1
else:
alpha = 1.0
ax.plot(data[time_dim], data.values, **plot_kwargs, alpha=alpha)
def cum_recruits(ax, recruits, first_plot_day, color, linestyle):
"""Plot cumulative recruits over a time dimension.
Args:
ax: The axis instance we want to plot on.
recruits: The xr.DataArray that we want to plot. Must have either
'time' OR 'historical_time' dimension.
first_plot_day: An int representing the first date to plot.
color: A mpl color.
linestyle: A mpl linestyle.
"""
time_dim = plot_utils.find_time_dim(recruits)
cum_recruits = recruits.cumsum(time_dim)
array_over_time(ax, cum_recruits, first_plot_day, {'color': color, 'ls': linestyle})
ax.set_ylabel('Cumulative recruits')
def loc_plots(ds, box, loc_to_plot=None):
"""Make plots summarizing incidence and recruitment in one location.
Show data specific to one location. Plot the incidence, cumulative
recruits, and recuits in each subgroup as functions of time.
Args:
ds: An xr.dataset containing the vill' to visualize.
box: An ipywidgets.Box containing three outputs.
loc_to_plot: A ds.location.coord specifiying the location to plot
"""
# setup
if loc_to_plot is None:
loc_to_plot = ds.location.values[0]
fpd = ville_config.FIRST_PLOT_DAY
pc = colors_config.ville_styles['gray_ville_3']['color']
oc = colors_config.ville_styles['highlight_ville_2']['color']
ls = '-'
# select just what we want to look at
incidence = ds.incidence_flattened.sel(location=loc_to_plot)
historical_incidence = ds.historical_incidence.sel(location=loc_to_plot)
proposed_part = plot_utils.join_in_time(
ds, 'participants').sel(location=loc_to_plot)
original_part = plot_utils.join_in_time(
ds, 'original_participants').sel(location=loc_to_plot)
# Make plots
# Incidence
# No original as the user cannot interactively control the incidence.
fig, axis = plot_utils.new_figure()
plot.incidence(axis, incidence, fpd, 'k', '-')
plot.incidence(axis, historical_incidence, fpd, 'k', '-')
plot.format_time_axis(axis, date_format='%b-%d')
# TODO find a better way to align plots
axis.text(-0.25,
1.1,
f'Individual Trial Site \n {loc_to_plot}',
ha='left',
va='bottom',
transform=axis.transAxes,
fontsize=16.)
int_utils.update_disp(box.children[0], fig)
# Total recruits over time
fig, axis = plot_utils.new_figure()
# TODO figure out grids
axis.text(0.5,
1.1,
' ',
ha='left',
va='bottom',
transform=axis.transAxes,
fontsize=16.)
# Assume everything but time is a participant dimension.
# The arrays aren't unpacked, so this is not equivalent to
# get_labels_to_plot.
p_dims = list(proposed_part.dims)
remove_dims = [plot_utils.find_time_dim(proposed_part)]
for rd in remove_dims:
p_dims.remove(rd)
plot.cum_recruits(axis, original_part.sum(p_dims), fpd, oc, ls)
plot.cum_recruits(axis, proposed_part.sum(p_dims), fpd, pc, ls)
plot.format_time_axis(axis, date_format='%b-%d')
axis.set_title('Cumulative Recruits \n All Participants')
int_utils.update_disp(box.children[1], fig)
# Subrecruits over time
num_labels = len(plot_utils.get_labels_to_plot(proposed_part))
num_cols = 2
num_rows = num_labels // num_cols + (num_labels % num_cols != 0)
fig, a = plot_utils.make_subplots(num_rows, num_cols)
fig.suptitle('Cumulative Recruits')
plot.cum_subrecruits(a, original_part, fpd, oc, ls)
plot.cum_subrecruits(a, proposed_part, fpd, pc, ls)
int_utils.update_disp(box.children[2], fig)