def test_age_brackets_used_with_contact_matrix():
"""
Test that the age brackets used in sp.Pop.generate() matches the contact matrices used.
Note:
This is a test to ensure that within sp.Pop.generate() uses the right age brackets. By default, without specifying nbrackets in sp.get_census_age_brackets(), the number of age brackets will not match the granularity of the contact matrix.
"""
sp.logger.info(
"Test that the age brackets used in sp.Pop.generate() with the contact matrices have the same number of bins as the contact matrices."
)
pop_obj = sp.Pop(**pars)
sheet_name = pop_obj.sheet_name
pop = pop_obj.to_dict(
) # this is basically what sp.make_population does...
contact_matrix_dic = sp.get_contact_matrix_dic(sp.datadir,
sheet_name=sheet_name)
contact_matrix_nbrackets = contact_matrix_dic[list(
contact_matrix_dic.keys())[0]].shape[0]
cm_age_brackets = sp.get_census_age_brackets(
sp.datadir,
country_location=pop_obj.country_location,
state_location=pop_obj.state_location,
location=pop_obj.location,
nbrackets=contact_matrix_nbrackets)
assert contact_matrix_nbrackets == len(
cm_age_brackets
), f'Check failed, len(contact_matrix_nbrackets): {contact_matrix_nbrackets} does not match len(cm_age_brackets): {len(cm_age_brackets)}.'
print(
f'Check passed. The age brackets loaded match the number of age brackets for the contact matrices used for the location.'
)
def test_generate_larger_households(location='seattle_metro',
state_location='Washington',
country_location='usa'):
Nhomes_to_sample_smooth = 1000
household_size_distr = sp.get_household_size_distr(datadir, location,
state_location,
country_location)
hh_sizes = sp.generate_household_sizes(Nhomes_to_sample_smooth,
household_size_distr)
hha_brackets = sp.get_head_age_brackets(datadir,
country_location=country_location)
hha_by_size_counts = sp.get_head_age_by_size_distr(
datadir, country_location=country_location)
age_brackets_filepath = sp.get_census_age_brackets_path(
datadir, state_location, country_location)
age_brackets = sp.get_age_brackets_from_df(age_brackets_filepath)
age_by_brackets_dic = sp.get_age_by_brackets_dic(age_brackets)
contact_matrix_dic = sp.get_contact_matrix_dic(
datadir, sheet_name='United States of America')
single_year_age_distr = {}
for n in range(101):
single_year_age_distr[n] = float(1.0 / 101.0)
# generate households of size 3
size = 3
# first variable is the household size to be created, so here this means we want to create all households of size 3 and the hh_sizes variable tells us how many of size 3 will be created at index 3-1 (since hh_sizes is an array rather than a dictionary)
larger_households = sp.generate_larger_households(
size, hh_sizes, hha_by_size_counts, hha_brackets, age_brackets,
age_by_brackets_dic, contact_matrix_dic, single_year_age_distr)
assert larger_households is not None
print(larger_households)
def test_age_brackets_used_with_contact_matrix():
"""
Test that the age brackets used in sp.Pop.generate() matches the contact matrices used.
Note:
This is a test to ensure that within sp.Pop.generate() uses the right age brackets. By default, without specifying nbrackets in sp.get_census_age_brackets(), the number of age brackets will not match the granularity of the contact matrix.
"""
sp.logger.info(
"Test that the age brackets used in sp.Pop.generate() with the contact matrices have the same number of bins as the contact matrices."
)
pop = sp.Pop(**pars)
sheet_name = pop.sheet_name
loc_pars = pop.loc_pars
contact_matrix_dic = sp.get_contact_matrix_dic(sp.datadir,
sheet_name=sheet_name)
contact_matrix_nbrackets = contact_matrix_dic[list(
contact_matrix_dic.keys())[0]].shape[0]
cm_age_brackets = sp.get_census_age_brackets(
**sc.mergedicts(loc_pars, {'nbrackets': contact_matrix_nbrackets}))
assert contact_matrix_nbrackets == len(
cm_age_brackets
), f'Check failed, len(contact_matrix_nbrackets): {contact_matrix_nbrackets} does not match len(cm_age_brackets): {len(cm_age_brackets)}.'
print(
f'Check passed. The age brackets loaded match the number of age brackets for the contact matrices used for the location.'
)
def test_generate_all_households(location='seattle_metro',
state_location='Washington',
country_location='usa'):
N = 1000
household_size_distr = sp.get_household_size_distr(datadir, location,
state_location,
country_location)
hh_sizes = sp.generate_household_sizes_from_fixed_pop_size(
N, household_size_distr)
hha_brackets = sp.get_head_age_brackets(datadir,
country_location=country_location)
hha_by_size_counts = sp.get_head_age_by_size_distr(
datadir, country_location=country_location)
age_brackets_filepath = sp.get_census_age_brackets_path(
datadir, state_location, country_location)
age_brackets = sp.get_age_brackets_from_df(age_brackets_filepath)
age_by_brackets_dic = sp.get_age_by_brackets_dic(age_brackets)
contact_matrix_dic = sp.get_contact_matrix_dic(
datadir, sheet_name='United States of America')
single_year_age_distr = {}
for n in range(101):
single_year_age_distr[n] = float(1.0 / 101.0)
homes_dic, homes = sp.generate_all_households(
N, hh_sizes, hha_by_size_counts, hha_brackets, age_brackets,
age_by_brackets_dic, contact_matrix_dic, single_year_age_distr)
assert homes_dic, homes is not None
def test_webapp_synthpops_calls(n=default_n,
location='seattle_metro',
state_location='Washington',
country_location='usa',
sheet_name='United States of America'):
datadir = sp.datadir
sp.read_age_bracket_distr(datadir,
location=location,
state_location=state_location,
country_location=country_location)
sp.get_census_age_brackets(datadir,
state_location=state_location,
country_location=country_location)
num_agebrackets = 16
n_contacts_dic = {'H': 4.11, 'S': 11.41, 'W': 8.07, 'C': 7}
contact_matrix_dic = sp.get_contact_matrix_dic(datadir,
sheet_name=sheet_name)
contact_matrix_dic['M'] = sp.combine_matrices(contact_matrix_dic,
n_contacts_dic,
num_agebrackets)
for k in contact_matrix_dic:
print(contact_matrix_dic[k].shape)
n = int(n)
sp.get_age_n(datadir,
n=default_n,
location=location,
state_location=state_location,
country_location=country_location)
return
def test_send_students_to_school(n=10000,
location='seattle_metro',
state_location='Washington',
country_location='usa',
folder_name='contact_networks'):
homes = sprw.read_setting_groups(datadir,
location,
state_location,
country_location,
folder_name,
'households',
n,
with_ages=True)
homes_by_uids, age_by_uid_dic = sp.assign_uids_by_homes(homes)
uids_in_school, uids_in_school_by_age, ages_in_school_count = sp.get_uids_in_school(
datadir,
n,
location,
state_location,
country_location,
age_by_uid_dic,
homes_by_uids,
use_default=False)
school_size_distr_by_bracket = sp.get_school_size_distr_by_brackets(
datadir, location, state_location, country_location)
school_size_brackets = sp.get_school_size_brackets(datadir, location,
state_location,
country_location)
school_sizes = sp.generate_school_sizes(school_size_distr_by_bracket,
school_size_brackets,
uids_in_school)
age_brackets_filepath = sp.get_census_age_brackets_path(
datadir, state_location, country_location)
age_brackets = sp.get_age_brackets_from_df(age_brackets_filepath)
age_by_brackets_dic = sp.get_age_by_brackets_dic(age_brackets)
contact_matrix_dic = sp.get_contact_matrix_dic(
datadir, sheet_name='United States of America')
syn_schools, syn_school_uids, syn_school_types = sp.send_students_to_school(
school_sizes,
uids_in_school,
uids_in_school_by_age,
ages_in_school_count,
age_brackets,
age_by_brackets_dic,
contact_matrix_dic,
verbose=False)
assert syn_schools, syn_school_uids is not None
return syn_schools, syn_school_uids
def test_all(location='seattle_metro', state_location='Washington', country_location='usa', sheet_name='United States of America'):
''' Run all tests '''
sc.heading('Running all tests')
sp.validate() # Validate that data files can be found
# dropbox_path = sp.datadir
dropbox_path = sp.settings.datadir
age_bracket_distr = spdd.read_age_bracket_distr(dropbox_path, location, state_location, country_location)
gender_fraction_by_age = sp.read_gender_fraction_by_age_bracket(dropbox_path, location, state_location, country_location)
age_brackets_file, age_brackets_filepath = sp.get_census_age_brackets_path(dropbox_path, state_location, country_location)
print(age_brackets_filepath)
age_brackets = sp.get_age_brackets_from_df(age_brackets_filepath)
age_by_brackets_dic = sp.get_age_by_brackets_dic(age_brackets)
# ## Test selecting an age and sex for an individual ###
a, s = sp.get_age_sex(gender_fraction_by_age, age_bracket_distr, age_brackets)
print(a, s)
# ## Test age mixing matrix ###
# num_agebrackets = 18
# flu-like weights. calibrated to empirical diary survey data.
weights_dic = {'H': 4.11, 'S': 11.41, 'W': 8.07, 'C': 2.79}
age_mixing_matrix_dic = sp.get_contact_matrix_dic(dropbox_path, sheet_name)
# ## Test sampling contacts based on age ###
age, sex = sp.get_age_sex(gender_fraction_by_age, age_bracket_distr, age_brackets) # sample an age (and sex) from the seattle metro distribution
n_contacts = 30
contact_ages = sp.sample_n_contact_ages(n_contacts, age, age_brackets, age_by_brackets_dic, age_mixing_matrix_dic, weights_dic)
print(contact_ages)
# shut down schools
no_schools_weights = sc.dcp(weights_dic)
no_schools_weights['S'] = 0.1 # research shows that even with school closure, kids still have some contact with their friends from school.
f_reduced_contacts_students = 0.5
f_reduced_contacts_nonstudents = 0.2
if age < 20:
n_reduced_contacts = int(n_contacts * (1 - f_reduced_contacts_students))
else:
n_reduced_contacts = int(n_contacts * (1 - f_reduced_contacts_nonstudents))
contact_ages = sp.sample_n_contact_ages(n_reduced_contacts, age, age_brackets, age_by_brackets_dic, age_mixing_matrix_dic, no_schools_weights)
print(contact_ages)
return
def test_generate_workplace_sizes(location='seattle_metro', state_location='Washington',
country_location='usa', folder_name='contact_networks'):
Npeople = 10000
uids_in_school, uids_in_school_by_age, ages_in_school_count = sp.get_uids_in_school(datadir, Npeople, location,
state_location,
country_location,
folder_name=folder_name,
use_default=True)
school_size_distr_by_bracket = sp.get_school_size_distr_by_brackets(datadir, location, state_location,
country_location)
school_size_brackets = sp.get_school_size_brackets(datadir, location, state_location, country_location)
school_sizes = sp.generate_school_sizes(school_size_distr_by_bracket, school_size_brackets, uids_in_school)
age_brackets_filepath = sp.get_census_age_brackets_path(datadir, state_location, country_location)
age_brackets = sp.get_age_brackets_from_df(age_brackets_filepath)
age_by_brackets_dic = sp.get_age_by_brackets_dic(age_brackets)
contact_matrix_dic = sp.get_contact_matrix_dic(datadir, sheet_name='United States of America')
# Need to instead get syn_schools now
syn_schools, syn_school_uids = sp.send_students_to_school(school_sizes, uids_in_school, uids_in_school_by_age,
ages_in_school_count, age_brackets, age_by_brackets_dic,
contact_matrix_dic)
employment_rates = sp.get_employment_rates(datadir, location=location, state_location=state_location,
country_location=country_location, use_default=True)
age_by_uid_dic = sp.read_in_age_by_uid(datadir, location, state_location, country_location, folder_name, Npeople)
potential_worker_uids, potential_worker_uids_by_age, potential_worker_ages_left_count = sp.get_uids_potential_workers(
syn_school_uids, employment_rates, age_by_uid_dic)
workers_by_age_to_assign_count = sp.get_workers_by_age_to_assign(employment_rates, potential_worker_ages_left_count,
age_by_uid_dic)
workplace_size_brackets = sp.get_workplace_size_brackets(datadir, location, state_location, country_location,
use_default=True)
workplace_size_distr_by_brackets = sp.get_workplace_size_distr_by_brackets(datadir,
state_location=state_location,
country_location=country_location,
use_default=True)
workplace_sizes = sp.generate_workplace_sizes(workplace_size_distr_by_brackets, workplace_size_brackets,
workers_by_age_to_assign_count)
return workers_by_age_to_assign_count, workplace_size_brackets, workplace_size_distr_by_brackets, workplace_sizes
def test_send_students_to_school(n=10000, location='seattle_metro', state_location='Washington',
country_location='usa', folder_name='contact_networks'):
homes = sp.read_setting_groups(datadir, location, state_location, country_location, 'households', folder_name, n, with_ages=True)
homes_by_uids, age_by_uid_dic = sp.assign_uids_by_homes(homes)
uids_in_school, uids_in_school_by_age, ages_in_school_count = sp.get_uids_in_school(datadir, n, location,
state_location,
country_location,
age_by_uid_dic,
homes_by_uids,
use_default=False)
# assert uids_in_school is not None
# def test_send_students_to_school(n=1000, location='seattle_metro', state_location='Washington',
# country_location='usa'):
# homes = sp.get_head_age_by_size_distr(datadir, state_location, country_location, file_path=None,
# household_size_1_included=False, use_default=True)
# homes_by_uids, age_by_uid_dic = sp.assign_uids_by_homes(homes, id_len=16)
# uids_in_school, uids_in_school_by_age, ages_in_school_count = sp.get_uids_in_school(datadir, n, location,
# state_location,
# country_location,
# age_by_uid_dic,
# homes_by_uids,
# use_default=False)
# >>>>>>> origin/mf/update-saved-pop-fixes
school_size_distr_by_bracket = sp.get_school_size_distr_by_brackets(datadir, location, state_location,
country_location)
school_size_brackets = sp.get_school_size_brackets(datadir, location, state_location, country_location)
school_sizes = sp.generate_school_sizes(school_size_distr_by_bracket, school_size_brackets, uids_in_school)
age_brackets_filepath = sp.get_census_age_brackets_path(datadir, state_location, country_location)
age_brackets = sp.get_age_brackets_from_df(age_brackets_filepath)
age_by_brackets_dic = sp.get_age_by_brackets_dic(age_brackets)
contact_matrix_dic = sp.get_contact_matrix_dic(datadir, sheet_name='United States of America')
syn_schools, syn_school_uids = sp.send_students_to_school(school_sizes, uids_in_school, uids_in_school_by_age,
ages_in_school_count, age_brackets, age_by_brackets_dic,
contact_matrix_dic, verbose=False)
assert syn_schools, syn_school_uids is not None
return syn_schools, syn_school_uids
def test_older_ages_have_household_contacts():
"""
Test that older age groups (85+) have at least some household contacts with
other older individuals if expected. Together, if sp.Pop.generate() uses the
incorrect number of age brackets with the contact matrices, older age groups
will not be generated as household contacts for each other (when we look at
the generated contact matrix for households, the blocks between 85+ year
olds would then be 0 for relatively large populations, even though the
household contact matrix would have us expect otherwise.)
"""
test_pars = sc.dcp(pars)
test_pars['n'] = 20e3
pop = sp.Pop(**test_pars)
pop_dict = pop.to_dict()
contact_matrix_dic = sp.get_contact_matrix_dic(sp.datadir,
sheet_name=pop.sheet_name)
contact_matrix_nbrackets = contact_matrix_dic[list(
contact_matrix_dic.keys())[0]].shape[0]
cm_age_brackets = sp.get_census_age_brackets(
sp.datadir,
country_location=pop.country_location,
state_location=pop.state_location,
location=pop.location,
nbrackets=contact_matrix_nbrackets)
cm_age_by_brackets_dic = sp.get_age_by_brackets_dic(cm_age_brackets)
age_threshold = 85
age_threshold_bracket = cm_age_by_brackets_dic[age_threshold]
expected_older_contact = np.sum(
contact_matrix_dic['H'][age_threshold_bracket:,
age_threshold_bracket:])
matrix = sp.calculate_contact_matrix(pop_dict, setting_code='H')
gen_older_age_contacts = np.sum(matrix[age_threshold:, age_threshold:])
if expected_older_contact != 0:
assert gen_older_age_contacts != 0, f'Check failed, individuals over {age_threshold} years old have no contacts with each other in households even though the household contact matrix expects them to.'
else:
assert gen_older_age_contacts == 0, f'Check failed, individuals over {age_threshold} years old have {gen_older_age_contacts} household contacts with each other even though the household contact matrix expects them to have none.'
print('Check passed.')
def test_assign_rest_of_workers(state_location='Washington',
country_location='usa'):
workers_by_age_to_assign_count, workplace_size_brackets, workplace_size_distr_by_brackets, \
workplace_sizes = test_generate_workplace_sizes()
potential_worker_uids, potential_worker_uids_by_age, employment_rates, age_by_uid_dic = test_get_uids_potential_workers(
)
contact_matrix_dic = sp.get_contact_matrix_dic(
datadir, sheet_name='United States of America')
age_brackets_16 = sp.get_census_age_brackets(datadir, state_location,
country_location)
age_by_brackets_dic_16 = sp.get_age_by_brackets_dic(age_brackets_16)
syn_workplaces, syn_workplace_uids, potential_worker_uids, potential_worker_uids_by_age, workers_by_age_to_assign_count = sp.assign_rest_of_workers(
workplace_sizes, potential_worker_uids,
potential_worker_uids_by_age, workers_by_age_to_assign_count,
dict(age_by_uid_dic), age_brackets_16, age_by_brackets_dic_16,
contact_matrix_dic)
# TODO: Issue #116 assign_rest_of_workers returns empty syn_workplaces and syn_workplace_uids
# syn_workplaces should return a list of lists where each sublist is a workplace with the ages of workers, not empty
# for workplace in syn_workplaces:
# assert workplace is not None
# assert syn_workplaces != []
# syn_worplace_uids should be a list of workers ids, not empty
# assert syn_workplace_uids != []
# potential_worker_uids should return a list of potential worker ids
for worker_id in potential_worker_uids:
assert worker_id is not None
# potential_worker_uids_by_age should return a list of potential worker ids mapped by age
for worker_by_age in potential_worker_uids_by_age:
assert int(worker_by_age)
# workers_by_age_to_assign_count should be a dictionary mapping age to the count of workers left to assign
for worker in workers_by_age_to_assign_count.items():
assert tuple(worker)
def test_send_students_to_school(location='seattle_metro',
state_location='Washington',
country_location='usa'):
NPeople = 10000
uids_in_school, uids_in_school_by_age, ages_in_school_count = sp.get_uids_in_school(
datadir,
NPeople,
location,
state_location,
country_location,
use_default=True)
school_size_distr_by_bracket = sp.get_school_size_distr_by_brackets(
datadir, location, state_location, country_location)
school_size_brackets = sp.get_school_size_brackets(datadir, location,
state_location,
country_location)
school_sizes = sp.generate_school_sizes(school_size_distr_by_bracket,
school_size_brackets,
uids_in_school)
age_brackets_filepath = sp.get_census_age_brackets_path(
datadir, state_location, country_location)
age_brackets = sp.get_age_brackets_from_df(age_brackets_filepath)
age_by_brackets_dic = sp.get_age_by_brackets_dic(age_brackets)
contact_matrix_dic = sp.get_contact_matrix_dic(
datadir, sheet_name='United States of America')
syn_schools, syn_school_uids = sp.send_students_to_school(
school_sizes,
uids_in_school,
uids_in_school_by_age,
ages_in_school_count,
age_brackets,
age_by_brackets_dic,
contact_matrix_dic,
verbose=False)
assert syn_schools, syn_school_uids is not None
cmap2 = mplt.cm.get_cmap(cmocean.cm.curl_r)
cmap3 = mplt.cm.get_cmap(cmocean.cm.matter)
datadir = sp.datadir
state_location = 'Washington'
location = 'seattle_metro'
country_location = 'usa'
use_bayesian = False
age_brackets = sp.get_census_age_brackets(datadir, country_location,
use_bayesian)
age_by_brackets_dic = sp.get_age_by_brackets_dic(age_brackets)
num_agebrackets = 18
contact_matrix_dic = sp.get_contact_matrix_dic(datadir, state_location,
num_agebrackets)
household_size_distr = sp.get_household_size_distr(datadir, location,
state_location,
country_location,
use_bayesian)
print(household_size_distr)
# Nhomes = 20000
Nhomes = 10000
# create_homes = True
create_homes = False
if create_homes:
household_size_distr = sp.get_household_size_distr(datadir, location,
state_location,
def test_contact_matrix_has_all_layers(self):
contact_matrix = sp.get_contact_matrix_dic(
datadir=sp.datadir, sheet_name="United States of America")
for layer in ['H', 'S', 'W', 'C']:
self.assertIn(layer, contact_matrix)
pass
cmap2 = mplt.cm.get_cmap(cmocean.cm.curl_r)
cmap3 = mplt.cm.get_cmap(cmocean.cm.matter)
datadir = sp.datadir
state_location = 'Washington'
location = 'seattle_metro'
country_location = 'usa'
sheet_name = 'United States of America'
age_brackets = sp.get_census_age_brackets(datadir, state_location,
country_location)
age_by_brackets_dic = sp.get_age_by_brackets_dic(age_brackets)
num_agebrackets = 18
contact_matrix_dic = sp.get_contact_matrix_dic(datadir, sheet_name)
household_size_distr = sp.get_household_size_distr(datadir, location,
state_location,
country_location)
print(household_size_distr)
# Nhomes = 100000
# Nhomes = 10000
Nhomes = 6666
create_homes = True
# create_homes = False
if create_homes:
household_size_distr = sp.get_household_size_distr(datadir, location,
state_location,
