def test_plot_school_sizes(do_show, do_save, artifact_dir): """ Test that the school size distribution by type plotting method in sp.Pop class works for a large population. This will be a longer test that is run as part of our end-to-end testing suite. Visually show how the school size distribution generated compares to the data for the location being simulated. Notes: The larger the population size, the better the generated school size distributions by school type can match the expected data. If generated populations are too small, larger schools will be missed and in general there won't be enough schools generated to apply statistical tests. """ sp.logger.info("Test that the school size distribution by type plotting method in sp.Pop class works. Note: For small population sizes, the expected and generated size distributions may not match very well given that the model is stochastic and demographics are based on much larger populations.") pop = sp.Pop(**pars) kwargs = sc.objdict(sc.mergedicts(pars, pop.loc_pars)) kwargs.figname = f"test_school_size_distributions_{kwargs.location}_pop" kwargs.do_show = do_show kwargs.do_save = do_save if artifact_dir: kwargs.figdir = artifact_dir kwargs.screen_height_factor = 0.20 kwargs.hspace = 0.8 kwargs.bottom = 0.09 kwargs.keys_to_exclude = ['uv'] kwargs.cmap = cmr.get_sub_cmap('cmo.curl', 0.08, 1) fig, ax = pop.plot_school_sizes(**kwargs) assert isinstance(fig, mplt.figure.Figure), 'End-to-end school sizes check failed.' print('Check passed. Figure made.') return fig, ax, pop
def test_plot_school_sizes(do_show=False, do_save=False): """ Test that the school size distribution by type plotting method in sp.Pop class works. Visually show how the school size distribution generated compares to the data for the location being simulated. Notes: The larger the population size, the better the generated school size distributions by school type can match the expected data. If generated populations are too small, larger schools will be missed and in general there won't be enough schools generated to apply statistical tests. """ sp.logger.info( "Test that the school size distribution by type plotting method in sp.Pop class works. Note: For small population sizes, the expected and generated size distributions may not match very well given that the model is stochastic and demographics are based on much larger populations." ) pop = sp.Pop(**pars) kwargs = sc.objdict(sc.mergedicts(pars, pop.loc_pars)) kwargs.figname = f"test_school_size_distributions_{kwargs.location}_pop" kwargs.do_show = do_show kwargs.do_save = do_save kwargs.screen_height_factor = 0.20 kwargs.hspace = 0.8 kwargs.bottom = 0.09 kwargs.keys_to_exclude = ['uv'] kwargs.cmap = cmr.get_sub_cmap('cmo.curl', 0.08, 1) fig, ax = pop.plot_school_sizes(**kwargs) assert isinstance(fig, mplt.figure.Figure), 'Check 1 failed.' print('Check passed. Figure 1 made.') # works on popdict sp.logger.info("Test school size distribution plotting method on popdict.") popdict = pop.popdict kwargs.datadir = sp.settings.datadir kwargs.do_show = False kwargs.figname = f"test_school_size_distributions_{kwargs.location}_popdict" fig2, ax2 = sp.plot_school_sizes(popdict, **kwargs) if not kwargs.do_show: plt.close() assert isinstance(fig2, mplt.figure.Figure), 'Check 2 failed.' print('Check passed. Figure 2 made.') sp.logger.info( "Test school size distribution plotting method with keys_to_exclude as a string and without comparison." ) kwargs.keys_to_exclude = 'uv' kwargs.comparison = False fig3, ax3 = pop.plot_school_sizes(**kwargs) assert isinstance(fig3, mplt.figure.Figure), 'Check 3 failed.' print( 'Check passed. Figure 3 made with keys_to_exclude as a string and without comparison.' ) return fig, ax, pop
def smooth_binned_age_distribution(pars, do_show=False): sp.logger.info(f"Smoothing out age distributions with moving averages.") s = dict() # raw_age_bracket_distr = sp.read_age_bracket_distr(sp.datadir, location=pars['location'], state_location=pars['state_location'], country_location=pars['country_location']) raw_age_distr = sp.get_smoothed_single_year_age_distr( sp.datadir, location=pars['location'], state_location=pars['state_location'], country_location=pars['country_location'], window_length=1) age_brackets = sp.get_census_age_brackets( sp.datadir, country_location=pars['country_location'], state_location=pars['state_location'], location=pars['location']) max_age = age_brackets[max(age_brackets.keys())][-1] age_range = np.arange(max_age + 1) for si in np.arange(3, 8, 2): smoothed_age_distr = sp.get_smoothed_single_year_age_distr( sp.datadir, location=pars['location'], state_location=pars['state_location'], country_location=pars['country_location'], window_length=si) s[si] = np.array([smoothed_age_distr[a] for a in age_range]) fig, ax = plt.subplots(1, 1, figsize=(10, 6)) cmap = mplt.cm.get_cmap('cmr.ember') if len(s) > 3: cmap1 = cmr.get_sub_cmap('cmr.rainforest', 0.13, 0.85) cmap2 = cmr.get_sub_cmap('cmr.rainforest', 0.20, 0.92) else: cmap1 = cmr.get_sub_cmap('cmr.rainforest', 0.18, 0.68) cmap2 = cmr.get_sub_cmap('cmr.rainforest', 0.25, 0.75) delta = 1 / (len(s) - 1) age_range = np.array(sorted(smoothed_age_distr.keys())) r = np.array([raw_age_distr[a] for a in age_range]) ax.plot(age_range, r, color=cmap(0.55), marker='o', markerfacecolor=cmap(0.65), markersize=3, markeredgewidth=1, alpha=0.65, label='Raw') for ns, si in enumerate(sorted(s.keys())): ax.plot(age_range, s[si], color=cmap1(ns * delta), marker='o', markerfacecolor=cmap2(ns * delta), markeredgewidth=1, markersize=3, alpha=.75, label=f'Smoothing window = {si}') leg = ax.legend(loc=3) leg.draw_frame(False) ax.set_xlim(age_range[0], age_range[-1]) ax.set_ylim(bottom=0.) ax.set_xlabel('Age') ax.set_ylabel('Distribution (%)') ax.set_title( f"Smoothing Binned Age Distribution: {pars['location'].replace('_', ' ').replace('-', ' ')}" ) if do_show: plt.show() return fig, ax
def plot_school_sizes_by_type(pop, pars, do_show=False): """ Plot the school size distribution by type compared with the expected data. """ sp.logger.info( f"Plotting to show that school sizes are generated by school type when the parameter 'with_school_types' is set to True." ) if pars['with_school_types']: expected_school_size_distr = sp.get_school_size_distr_by_type( sp.datadir, location=pars['location'], state_location=pars['state_location'], country_location=pars['country_location'], use_default=pars['use_default']) school_size_brackets = sp.get_school_size_brackets( sp.datadir, location=pars['location'], state_location=pars['state_location'], country_location=pars['country_location'] ) # for right now the size distribution for all school types will use the same brackets or bins else: expected_school_size_distr = { None: sp.get_school_size_distr_by_brackets( sp.datadir, location=pars['location'], state_location=pars['state_location'], country_location=pars['country_location'], use_default=pars['use_default']) } school_size_brackets = sp.get_school_size_brackets( sp.datadir, location=pars['location'], state_location=pars['state_location'], country_location=pars['country_location']) bins = [school_size_brackets[0][0]] + [ school_size_brackets[b][-1] + 1 for b in school_size_brackets ] schools = dict() enrollment_by_school_type = dict() gen_school_size_distr = dict() for i, person in pop.items(): if person['scid'] is not None and person['sc_student']: schools.setdefault(person['scid'], dict()) schools[person['scid']]['sc_type'] = person['sc_type'] schools[person['scid']].setdefault('enrolled', 0) schools[person['scid']]['enrolled'] += 1 for i, school in schools.items(): enrollment_by_school_type.setdefault(school['sc_type'], []) enrollment_by_school_type[school['sc_type']].append(school['enrolled']) for sc_type in enrollment_by_school_type: sizes = enrollment_by_school_type[sc_type] hist, bins = np.histogram(sizes, bins=bins, density=0) gen_school_size_distr[sc_type] = { i: hist[i] / sum(hist) for i in school_size_brackets } gen_school_size_distr = sc.objdict(gen_school_size_distr) width = 6 height = 3 * len(gen_school_size_distr) hspace = 0.4 cmap = cmr.get_sub_cmap('cmo.curl', 0.12, 1) fig, ax = plt.subplots(len(gen_school_size_distr), 1, figsize=(width, height), tight_layout=True) plt.subplots_adjust(hspace=hspace) if len(gen_school_size_distr) == 1: ax = [ax] bin_labels = [ f"{school_size_brackets[b][0]}-{school_size_brackets[b][-1]}" for b in school_size_brackets ] sorted_school_types = sorted(gen_school_size_distr.keys()) for ns, school_type in enumerate(sorted_school_types): x = np.arange(len(school_size_brackets)) c = ns / len(gen_school_size_distr) c2 = min(c + 0.1, 1) sorted_bins = sorted(expected_school_size_distr[school_type].keys()) ax[ns].bar( x, [expected_school_size_distr[school_type][b] for b in sorted_bins], color=cmap(c), edgecolor='white', label='Expected', zorder=0) ax[ns].plot( x, [gen_school_size_distr[school_type][b] for b in sorted_bins], color=cmap(c2), ls='--', marker='o', markerfacecolor=cmap(c2), markeredgecolor='white', markeredgewidth=.5, markersize=5, label='Simulated', zorder=1) leg = ax[ns].legend(loc=1) leg.draw_frame(False) ax[ns].set_xticks(x) ax[ns].set_xticklabels(bin_labels, rotation=25) ax[ns].set_xlim(0, x[-1]) ax[ns].set_ylim(0, 1) if school_type is None: title = "without school types defined" else: title = f"{school_type}" if ns == 0: if pars['location'] is not None: location_text = f"{pars['location'].replace('_', ' ').title()}" else: location_text = f"{sp.config.default_location.replace('_', ' ').title()} Default Sizes" ax[ns].text(0., 1.1, location_text, horizontalalignment='left', fontsize=10) ax[ns].set_title(title, fontsize=10) ax[ns].set_xlabel('School size') if do_show: plt.show() return fig, ax, sorted_school_types
import synthpops as sp import matplotlib.pyplot as plt import cmasher as cmr import cmocean pars = dict( n = 40e3, rand_seed = 123, location = 'Spokane_County', state_location = 'Washington', country_location = 'usa', smooth_ages = 1, household_method = 'fixed_ages', with_facilities = 1, with_non_teaching_staff = 1, # also include non teaching staff with_school_types = 1, school_mixing_type = {'pk': 'random', 'es': 'age_and_class_clustered', 'ms': 'age_and_class_clustered', 'hs': 'age_clustered', } ) pop = sp.Pop(**pars) kwargs = sc.dcp(pars) kwargs['cmap'] = cmr.get_sub_cmap('cmo.curl', 0.05, 1) # let's change the colormap used a little fig, ax = pop.plot_school_sizes(**kwargs) # plot school sizes by school type plt.show()
def show_observable(obs, realization=0, title=None, cartesian_axes='yz', **kwargs): """ Displays the contents of a single realisation of an Observable object. Parameters ---------- obs : imagine.observables.observable.Observable Observable object whose contents one wants to plot realization : int Index of the ensemble realization to be plotted cartesian_axes : str If plotting a tabular observable using cartesian coordinates, this allows selecting which two axes should be used for the plot. E.g. 'xy', 'zy', 'xz'. Default: 'yz'. **kwargs Parameters to be passed to the apropriate plotting routine (either `healpy.visufunc.mollview` or `matplotlib.pyplot.imshow`). """ if obs.otype == 'HEALPix': default_cmap = _choose_cmap(title=title) mollview_args = {'norm': 'hist', 'cmap': copy(default_cmap), 'unit': obs.unit._repr_latex_()} mollview_args.update(kwargs) return hp.mollview(obs.global_data[realization], title=title, **mollview_args) elif obs.otype == 'tabular': ax = None if 'sub' in kwargs: ax = plt.subplot(*kwargs['sub']) if obs.coords['type'] == 'galactic': x, y = obs.coords['lon'], obs.coords['lat'] plt.xlabel('Gal. lon. [{}]'.format(x.unit._repr_latex_())) plt.ylabel('Gal. lat. [{}]'.format(y.unit._repr_latex_())) elif obs.coords['type'] == 'cartesian': x, y = obs.coords[cartesian_axes[0]], obs.coords[cartesian_axes[1]] plt.xlabel('{} [{}]'.format(cartesian_axes[0], x.unit._repr_latex_())) plt.ylabel('{} [{}]'.format(cartesian_axes[1], y.unit._repr_latex_())) else: raise ValueError('Unsupported coordinates type', obs.coords['type']) values = obs.global_data[realization] #values -= values.min() #values /= values.max() if 'cmap' in kwargs: cmap = kwargs['cmap'] else: cmap = 'cmr.chroma' cmap = cmr.get_sub_cmap(cmap, 0.15, 0.85) #colors = cmap(values) plt.title(title) im = plt.scatter(x, y, c=values, cmap=cmap) plt.grid(alpha=0.2) plt.colorbar(im, ax=ax, label=obs.unit._repr_latex_()) else: # Includes the title in the corresponding subplot, even in # the unsupported case (so that the user remembers it) print("Plotting observable type '{}' not yet implemented".format(obs.otype))