Source code for fpsim.analyzers

Specify the core analyzers available in FPsim. Other analyzers can be
defined by the user by inheriting from these classes.

import numpy as np
import sciris as sc
import pylab as pl
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable

#%% Generic intervention classes

__all__ = ['Analyzer', 'snapshot', 'age_pyramids', 'empowerment_recorder', 'education_recorder']

[docs] class Analyzer(sc.prettyobj): ''' Base class for analyzers. Based on the Intervention class. Analyzers are used to provide more detailed information about a simulation than is available by default -- for example, pulling states out of sim.people on a particular timestep before it gets updated in the next timestep. To retrieve a particular analyzer from a sim, use sim.get_analyzer(). Args: label (str): a label for the Analyzer (used for ease of identification) ''' def __init__(self, label=None): if label is None: label = self.__class__.__name__ # Use the class name if no label is supplied self.label = label # e.g. "Record ages" self.initialized = False self.finalized = False return
[docs] def initialize(self, sim=None): ''' Initialize the analyzer, e.g. convert date strings to integers. ''' self.initialized = True self.finalized = False return
[docs] def finalize(self, sim=None): ''' Finalize analyzer This method is run once as part of `sim.finalize()` enabling the analyzer to perform any final operations after the simulation is complete (e.g. rescaling) ''' if self.finalized: raise RuntimeError('Analyzer already finalized') # Raise an error because finalizing multiple times has a high probability of producing incorrect results e.g. applying rescale factors twice self.finalized = True return
[docs] def apply(self, sim): ''' Apply analyzer at each time point. The analyzer has full access to the sim object, and typically stores data/results in itself. This is the core method which each analyzer object needs to implement. Args: sim: the Sim instance ''' pass
[docs] def to_json(self): ''' Return JSON-compatible representation Custom classes can't be directly represented in JSON. This method is a one-way export to produce a JSON-compatible representation of the intervention. This method will attempt to JSONify each attribute of the intervention, skipping any that fail. Returns: JSON-serializable representation ''' # Set the name json = {} json['analyzer_name'] = self.label if hasattr(self, 'label') else None json['analyzer_class'] = self.__class__.__name__ # Loop over the attributes and try to process attrs = self.__dict__.keys() for attr in attrs: try: data = getattr(self, attr) try: attjson = sc.jsonify(data) json[attr] = attjson except Exception as E: json[attr] = f'Could not jsonify "{attr}" ({type(data)}): "{str(E)}"' except Exception as E2: json[attr] = f'Could not jsonify "{attr}": "{str(E2)}"' return json
[docs] class snapshot(Analyzer): ''' Analyzer that takes a "snapshot" of the sim.people array at specified points in time, and saves them to itself. Args: timesteps (list): list of timesteps on which to take the snapshot args (list): additional timestep(s) die (bool): whether or not to raise an exception if a date is not found (default true) kwargs (dict): passed to Analyzer() **Example**:: sim = fp.Sim(analyzers=fps.snapshot('2020-04-04', '2020-04-14')) snapshot =['analyzers'][0] people = snapshot.snapshots[0] ''' def __init__(self, timesteps, *args, die=True, **kwargs): super().__init__(**kwargs) # Initialize the Analyzer object timesteps = sc.promotetolist(timesteps) # Combine multiple days timesteps.extend(args) # Include additional arguments, if present self.die = die # Whether or not to raise an exception self.timesteps = timesteps # String representations self.snapshots = sc.odict() # Store the actual snapshots return
[docs] def apply(self, sim): """ Apply snapshot at each timestep listed in timesteps and save result at snapshot[str(timestep)] """ for t in self.timesteps: if np.isclose(sim.i, t): self.snapshots[str(sim.i)] = sc.dcp(sim.people) # Take snapshot! return
[docs] class education_recorder(Analyzer): ''' Analyzer records all education attributes of females + pregnancy + living status for all timesteps. Made for debugging purposes. Args: args (list): additional timestep(s) kwargs (dict): passed to Analyzer() ''' def __init__(self, **kwargs): super().__init__(**kwargs) # Initialize the Analyzer object self.snapshots = sc.odict() # Store the actual snapshots self.keys = ['edu_objective', 'edu_attainment', 'edu_completed', 'edu_dropout', 'edu_interrupted', 'pregnant', 'alive', 'age'] self.max_agents = 0 # maximum number of agents this analyzer tracks self.time = [] self.trajectories = {} # Store education trajectories return
[docs] def apply(self, sim): """ Apply snapshot at each timestep listed in timesteps and save result at snapshot[str(timestep)] """ females = sim.people.filter(sim.people.is_female) self.snapshots[str(sim.i)] = {} for key in self.keys: self.snapshots[str(sim.i)][key] = sc.dcp(females[key]) # Take snapshot! self.max_agents = max(self.max_agents, len(females)) return
[docs] def finalize(self, sim=None): """ Process data in snapshots so we can plot it easily """ if self.finalized: raise RuntimeError('Analyzer already finalized') self.finalized = True # Process data so we can plot it easily self.time = np.array([key for key in self.snapshots.keys()], dtype=int) for state in self.keys: self.trajectories[state] = np.full((len(self.time), self.max_agents), np.nan) for ti, t in enumerate(self.time): stop_idx = len(self.snapshots[t][state]) self.trajectories[state][ti, 0:stop_idx] = self.snapshots[t][state] return
[docs] def plot(self, index=0, fig_args=None, pl_args=None): """ Plots time series of each state as a line graph Args: index: index of the female individual, must be less the analyzer's max_pop_size """ fig_args = sc.mergedicts(fig_args, {'figsize': (5, 7)}) pl_args = sc.mergedicts(pl_args) rows, cols = sc.get_rows_cols(2) fig = pl.figure(**fig_args) keys2 = ['edu_completed', 'edu_interrupted', 'edu_dropout'] keys3 = ['pregnant', 'alive'] k = 0 pl.subplot(rows, cols, k + 1) age_data = self.trajectories["age"] state = "edu_attainment" data = self.trajectories[state] pl.step(self.time, data[:, index], color="black", label=f"{state}", where='mid', **pl_args) state = "edu_objective" data = self.trajectories[state] pl.step(self.time, data[:, index], color="red", ls="--", label=f"{state}", where='mid', **pl_args) pl.ylim([0, 24]) pl.title('Education') pl.ylabel('Education (years)') pl.xlabel('Timesteps') pl.legend() k += 1 for state in sc.mergelists(keys2, keys3): pl.subplot(rows, cols, k + 1) data = self.trajectories[state] if state in keys2: if state == 'edu_interrupted': pl.step(self.time, 3*data[:, index], color=[0.7, 0.7, 0.7], label=f"{state}", ls=":", where='mid', **pl_args) elif state == "edu_dropout": pl.step(self.time, 3*data[:, index], color="black", label=f"{state}", ls=":", where='mid', **pl_args) else: pl.step(self.time, 3*data[:, index], color="#2ca25f", label=f"{state}", where='mid', **pl_args) elif state == 'pregnant': pl.step(self.time, data[:, index], color="#dd1c77", label=f"{state}", where='mid', **pl_args) elif state == 'alive': plt.step(self.time, 4*data[:, index], color="black", ls="--", label=f"{state}", where='mid', **pl_args) pl.title(f"Education trajectories - Start age: {int(age_data[0, index])}; final age {int(age_data[-1, index])}.") pl.ylabel('State') pl.xlabel('Timesteps') pl.legend() return fig
[docs] def plot_waterfall(self, max_timepoints=30, min_age=18, max_age=40, fig_args=None, pl_args=None): """ Plot a waterfall plot showing the evolution of education objective and attainment over time for a specified age group. Args: max_timepoints (int, optional): The maximum number of timepoints to plot, defaults to 30. min_age (int, optional): The minimum age for the age group, defaults to 18. max_age (int, optional): The maximum age for the age group, defaults to 20. Returns: figure handle The function generates uses kernel density estimation to visualize the data. If there's not data for the min max age specified, for a specific time step (ie, there are no agents in that age group), it adds a textbox. This is an edge case that can happen for a simulation with very few agents, and a very narrow age group. """ from scipy.stats import gaussian_kde data_att = self.trajectories["edu_attainment"] data_obj = self.trajectories["edu_objective"] data_age = self.trajectories["age"] mask = (data_age < min_age) | (data_age > max_age) | np.isnan(data_age) data_att =, mask=mask) data_obj =, mask=mask) n_tpts = data_att.shape[0] if n_tpts <= max_timepoints: tpts_to_plot = np.arange(n_tpts) else: tpts_to_plot = np.linspace(0, n_tpts - 1, max_timepoints, dtype=int) fig_args = sc.mergedicts(fig_args, {'figsize': (3, 10)}) pl_args = sc.mergedicts(pl_args, {'y_scaling': 0.9}) fig = plt.figure(**fig_args) ax = fig.add_subplot(111) edu_min, edu_max = 0, 25 edu_mid = (edu_max-edu_min)/2 + edu_min edu_years = np.linspace(edu_min, edu_max, 50) y_scaling = pl_args['y_scaling'] # Set the y-axis (time) labels ax.set_yticks(y_scaling*np.arange(len(tpts_to_plot))) ax.set_yticklabels(tpts_to_plot) # Initialize legend labels edu_att_label = None edu_obj_label = None # Loop through the selected time points and create kernel density estimates for idx, ti in enumerate(tpts_to_plot): data_att_ti = np.sort(data_att[ti, :][~data_att[ti, :].mask].data) data_obj_ti = np.sort(data_obj[ti, :][~data_obj[ti, :].mask].data) try: kde_att = gaussian_kde(data_att_ti) kde_obj = gaussian_kde(data_obj_ti) y_att = kde_att(edu_years) y_obj = kde_obj(edu_years) if idx == len(tpts_to_plot) - 1: edu_obj_label = 'Distribution of education objectives' edu_att_label = 'Current distribution of education attainment' ax.fill_between(edu_years, y_scaling*idx, y_obj / y_obj.max() + y_scaling*idx, color='#2f72de', alpha=0.3, label=edu_obj_label) ax.plot(edu_years, y_att / y_att.max() + y_scaling*idx, color='black', alpha=0.7, label=edu_att_label) except: # No data available for this age group or age range, ax.plot(edu_years, (y_scaling * idx) * np.ones_like(edu_years), color='black', alpha=0.2, label=edu_att_label) ax.annotate('No data available ', xy=(edu_mid, y_scaling*idx), xycoords='data', fontsize=8, ha='center', va='center', bbox=dict(boxstyle='round,pad=0.4', fc='none', ec="none")) # Labels and annotations ax.set_xlim([edu_min, edu_max]) ax.set_xlabel('Education years') ax.set_ylabel('Timesteps') ax.legend() ax.set_title(f"Evolution of education \n objective and attainment for age group:\n{min_age}-{max_age}.") # Show the plot return fig
[docs] class empowerment_recorder(Analyzer): ''' Records timeseries of empowerment attributes for different age groups. - For boolean attributes it computes the percentage returns percentage. - For float attributes it computes the median of the attribute from the population of interes Attributes: self.bins: A list of ages, default is a sequence from 0 to max_age + 1. self.keys: A list of people's empowerment attributes. A dictionary where[attribute] is a a matrix of shape (number of timesteps, number of bins - 1) containing age pyramid data. ''' def __init__(self, bins=None): """ Initializes self.i/t/y as empty lists and as empty dictionary """ super().__init__() self.bins = bins = sc.objdict() self.keys = ['partnered', 'urban', 'paid_employment', 'decision_wages', 'decision_health', 'sexual_autonomy', 'age'] self.nbins = None return
[docs] def initialize(self, sim): """ Initializes self.keys from sim.people """ super().initialize() if self.bins is None: self.bins = np.arange(0,['max_age']+2) self.nbins = len(self.bins)-1 for key in self.keys:[key] = np.full((self.nbins, sim.npts), np.nan) return
[docs] def apply(self, sim): """ Records histogram of empowerment attribute of all **alive female** individuals """ # Alive and female living_females = sc.findinds(sim.people.alive, sim.people.is_female) ages = sim.people.age[living_females] age_group = np.digitize(ages, self.bins) - 1 for key in self.keys: data = sim.people[key][living_females] if key == 'age': # Count how many living females we have in this age group temp = np.histogram(ages, self.bins)[0] vals = temp / temp.sum() # Transform to density elif key in ['partnered', 'urban', 'paid_employment']: vals = [np.mean(data[age_group == group_idx]) for group_idx in range(1, len(self.bins))] else: # assume float vals = [np.median(data[age_group == group_idx]) for group_idx in range(1, len(self.bins))][key][:, sim.i] = vals
[docs] def plot(self, to_plot=None, fig_args=None, pl_args=None): """ Plot all keys in self.keys or in to_plot as a heatmaps """ fig_args = sc.mergedicts(fig_args) pl_args = sc.mergedicts(pl_args) fig = pl.figure(**fig_args) if to_plot is None: to_plot = self.keys nkeys = len(to_plot) rows, cols = sc.get_rows_cols(nkeys) axs = [] for k, key in enumerate(to_plot): axs.append(fig.add_subplot(rows, cols, k+1)) try: data = np.array([key], dtype=float) label = f'metric: {key}' if key in ['partnered', 'urban', 'paid_employment']: clabel = f"proportion of {key}" cmap = 'RdPu' vmin, vmax = 0, 1 if key in ['urban']: cmap = 'RdYlBu_r' elif key in ['age']: clabel = "proportion of agents" cmap = 'Blues' vmin, vmax = 0, np.nanmax(data[:]) else: clabel = "average (median)" cmap = 'coolwarm' vmin, vmax = 0, 1 pcm = axs[k].pcolormesh(data, label=label, cmap=cmap, vmin=vmin, vmax=vmax, **pl_args) # Add colorbar to the right of the subplot divider = make_axes_locatable(axs[k]) cax = divider.append_axes("right", size="2.5%", pad=0.05) # Add colorbar to the right of the subplot plt.colorbar(pcm, cax=cax, label=clabel) # Generate age group labels and tick positions ytick_labels = [f"{self.bins[i]:.0f}-{self.bins[i+1]-1:.0f}" for i in range(self.nbins)] ytick_positions = np.arange(0.5, self.nbins + 0.5) # Center positions for ticks # Reduce the number of labels if we have too many bins max_labels = 10 if len(ytick_labels) > max_labels: step_size = len(ytick_labels) // max_labels ytick_labels = ytick_labels[::step_size] ytick_positions = ytick_positions[::step_size] # Label plots axs[k].set_yticks(ytick_positions) axs[k].set_yticklabels(ytick_labels) axs[k].set_title(key) axs[k].set_xlabel('Timestep') axs[k].set_ylabel('Age (years)') except: pl.title(f'Could not plot {key}') return fig
[docs] class age_pyramids(Analyzer): ''' Records age pyramids for each timestep. Attributes: self.bins: A list of ages, default is a sequence from 0 to max_age + 1. A matrix of shape (number of timesteps, number of bins - 1) containing age pyramid data. ''' def __init__(self, bins=None): """ Initializes bins and data variables """ super().__init__() self.bins = bins = None return
[docs] def initialize(self, sim): """ Initializes bins and data with proper shapes """ super().initialize() if self.bins is None: self.bins = np.arange(0,['max_age']+2) nbins = len(self.bins)-1 = np.full((sim.npts, nbins), np.nan) self._raw = sc.dcp( return
[docs] def apply(self, sim): """ Records histogram of ages of all alive individuals at a timestep such that[timestep] = list of proportions where index signifies age """ ages = sim.people.age[sc.findinds(sim.people.alive)] self._raw[sim.i, :] = np.histogram(ages, self.bins)[0][sim.i, :] = self._raw[sim.i, :]/self._raw[sim.i, :].sum()
[docs] def plot(self): """ Plots as 2D pyramid plot """ fig = pl.figure() pl.pcolormesh( pl.xlabel('Timestep') pl.ylabel('Age (years)') return fig
[docs] def plot3d(self): """ Plots as 3D pyramid plot """ print('Warning, very slow...') fig = pl.figure() sc.bar3d( pl.xlabel('Timestep') pl.ylabel('Age (years)') return fig