def test_workplace_contact_distribution(do_show, do_save, create_sample_pop_e2e, get_fig_dir_by_module): # calculate the workplace contacts count and plot sp.logger.info( "Test workplace contact distribution: workers in workplace such that worksize <= max_contacts must have" "contacts equal to worksize-1, for workers in workplace with size > max_contacts, the distribution " "should be closer to poisson distribution with mean = max_contacts ") plotting_kwargs = sc.objdict(do_show=do_show, do_save=do_save, figdir=get_fig_dir_by_module) contacts, contacts_by_id = cn.get_contact_counts_by_layer( create_sample_pop_e2e.popdict, layer="w", with_layer_ids=1) plotting_kwargs.append( "title_prefix", f"Total Workers = {len(contacts.get('wpid').get('all'))}") plotting_kwargs.append("figname", f"workers_contact_count") sp.plot_contact_counts(contacts, **plotting_kwargs) plotting_kwargs.remove("title_prefix") plotting_kwargs.remove("figname") # check workplace with worksize <= max_contacts max_contacts = create_sample_pop_e2e.max_contacts['W'] upperbound = st.poisson.interval(alpha=0.95, mu=max_contacts)[1] group_size_contacts = { f'all_worksize_contacts size > {max_contacts//2}': [], # capture size > max_contacts//2 f'large_worksize_contacts size > {upperbound}': [], # capture size > upperbound f'medium_large_worksize_contacts size between {max_contacts}, {upperbound}': [], # capture size between max_contacts and upperbound f'small_medium_worksize_contacts size between {max_contacts//2}, {max_contacts}': [], # capture size between max_contacts//2 and max_contacts } for k, v in contacts_by_id.items(): if len(v) <= max_contacts // 2: assert len([i for i in v if i != len(v) - 1]) == 0, \ "Failed, not all contacts in {len(k)} are equal to {len(v)} : {v}" else: if len(v) > upperbound: group_size_contacts[ f'large_worksize_contacts size > {upperbound}'] += v elif len(v) >= max_contacts: group_size_contacts[ f'medium_large_worksize_contacts size between {max_contacts}, {upperbound}'] += v else: group_size_contacts[ f'small_medium_worksize_contacts size between {max_contacts//2}, {max_contacts}'] += v group_size_contacts[ f'all_worksize_contacts size > {max_contacts//2}'] += v file_pattern = re.compile(r'([\s><=])') for i in group_size_contacts: plotting_kwargs["title_prefix"] = i plotting_kwargs["figname"] = file_pattern.sub("_", i) sp.check_truncated_poisson(testdata=group_size_contacts[i], mu=max_contacts, lowerbound=max_contacts // 2, skipcheck=True if "small" in i else True, **plotting_kwargs)
def test_plot_contact_counts(do_show=False, do_save=False): sp.logger.info( "Test plot_contact_counts method. Unit test --- for actual use with a sp.Pop object see e2etests/test_workplace_e2e.py." ) # multiple subplots contacts1 = { "people_type1": { "contact_type1": [2, 3, 4, 3, 4, 3, 4, 2, 1, 2], "contact_type2": [10, 22, 11, 12, 9, 9, 8, 10, 11, 10] }, "people_type2": { "contact_type1": [1, 3, 4, 3, 5, 3, 4, 4, 1, 2], "contact_type2": [10, 21, 11, 11, 10, 9, 8, 10, 11, 10] } } # single subplot contacts2 = { "people_type": { "contact_type": [2, 3, 4, 3, 4, 3, 4, 2, 1, 2] } } params = sc.objdict(do_save=do_save, do_show=do_show, title_prefix='test=true') fig1, ax1 = sp.plot_contact_counts(contact_counter=contacts1, **params) assert isinstance( fig1, mplt.figure.Figure), 'Check failed. Figure not generated.' fig2, ax2 = sp.plot_contact_counts(contact_counter=contacts2, varname="test", **params) assert isinstance( fig2, mplt.figure.Figure), 'Check failed. Figure not generated.' print('Check passed. Figures made.')
def test_average_class_size(average_class_size, do_show, do_save, get_fig_dir, quantiles=None): """ Test case to check average_class_size by taking average of student-student contacts Args: average_class_size: The average classroom size. Returns: None """ testpars = dict( average_class_size=average_class_size, # average_student_teacher_ratio = average_class_size, # DM: note that this parameter will overide the average class size parameter when school mixing types are something other than random or undefined (which defaults to random) --- method refactor work for schools will clarify these relationships ) pop = sp.Pop(**pars, **testpars) plotting_kwargs = sc.objdict(do_show=do_show, do_save=do_save, figdir=get_fig_dir) contacts = sp.get_contact_counts_by_layer(pop.popdict, with_layer_ids=1)[0] plotting_kwargs.append("title_prefix", f"Average Class Size = {average_class_size}") plotting_kwargs.append("figname", f"contact_average_class_size_{average_class_size}") sp.plot_contact_counts(contacts, **plotting_kwargs) counts = [] if not pop.school_pars.with_school_types: counts.extend(contacts['sc_student']['all']) counts.extend(contacts['sc_teacher']['all']) counts.extend(contacts['sc_staff']['all']) elif pop.school_pars.with_school_types and pop.school_pars.school_mixing_type == 'age_and_class_clustered': counts.extend(contacts['sc_student']['sc_student']) sp.check_poisson(actual=counts, expected=average_class_size, label='average_class_size', check='dist') # visual check with scipy.stats.probplot -- temporary, just to show that the null hypothesis should pass here for the distribution fig, ax = plt.subplots(1, 1) res = stats.probplot(counts, dist=stats.poisson, sparams=(average_class_size, ), plot=ax) if do_show: plt.show() plt.close() return
def test_average_teacher_teacher_degree(average_teacher_teacher_degree, do_show, do_save, get_fig_dir, threshold=2): """ Test case for average_teacher_teacher_degree by taking average of teachers' contacts per teacher Args: average_teacher_teacher_degree: The average number of contacts per teacher with other teachers Returns: None """ testpars = dict( average_teacher_teacher_degree=average_teacher_teacher_degree, with_school_types=1, school_mixing_type={ 'pk': 'age_and_class_clustered', # average_teacher_teacher_degree will not be used in school_mixing_type == 'random' scenario 'es': 'age_and_class_clustered', 'ms': 'age_and_class_clustered', 'hs': 'age_clustered', 'uv': 'age_clustered' }) pop = sp.Pop(**pars, **testpars) plotting_kwargs = sc.objdict(do_show=do_show, do_save=do_save, figdir=get_fig_dir) contacts = sp.get_contact_counts_by_layer(pop.popdict, with_layer_ids=1)[0] plotting_kwargs.append( "title_prefix", f"Average Teacher-Teacher Degree = {average_teacher_teacher_degree}") plotting_kwargs.append( "figname", f"contact_average_teacher_teacher_degree_{average_teacher_teacher_degree}" ) sp.plot_contact_counts(contacts, **plotting_kwargs) counts = contacts['sc_teacher']['sc_teacher'] sp.check_normal(actual=counts, expected=average_teacher_teacher_degree, label='teacher degree', check='mean') return
def test_average_additional_staff_degree(average_additional_staff_degree, do_show, do_save, get_fig_dir, threshold=2): """ Test case to check average_additional_staff_degree by taking average of all contacts per staff Args: average_additional_staff_degree: The average number of contacts per additional non teaching staff in schools Returns: None """ # note this must be greater than default average_student_all_staff_ratio (20) testpars = dict( average_additional_staff_degree=average_additional_staff_degree, with_school_types=1, ) pop = sp.Pop(**pars, **testpars) plotting_kwargs = sc.objdict(do_show=do_show, do_save=do_save, figdir=get_fig_dir) contacts = sp.get_contact_counts_by_layer(pop.popdict, with_layer_ids=1)[0] plotting_kwargs.append( "title_prefix", f"Average Additional Staff Degree = {average_additional_staff_degree}") plotting_kwargs.append( "figname", f"contact_average_additional_staff_degree_{average_additional_staff_degree}" ) sp.plot_contact_counts(contacts, **plotting_kwargs) counts = contacts['sc_staff']['all'] sp.check_normal(actual=counts, expected=average_additional_staff_degree, label='staff degree', check='mean') return