def test_get_uids_potential_workers(location='seattle_metro',
state_location='Washington',
country_location='usa'):
n = 10000
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)
employment_rates = sp.get_employment_rates(
datadir,
location=location,
state_location=state_location,
country_location=country_location,
use_default=True)
potential_worker_uids, potential_worker_uids_by_age, potential_worker_ages_left_count = sp.get_uids_potential_workers(
uids_in_school, employment_rates, age_by_uid_dic)
assert potential_worker_ages_left_count is not None
return potential_worker_uids, potential_worker_uids_by_age, employment_rates, age_by_uid_dic
def test_get_uids_in_school(location='seattle_metro',
state_location='Washington',
country_location='usa',
folder_name='contact_networks'):
Npeople = 10000
homes = sprw.read_setting_groups(datadir,
location,
state_location,
country_location,
folder_name,
'households',
Npeople,
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,
Npeople,
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=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_get_uids_in_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)
assert uids_in_school is not None
def test_assign_teachers_to_work(location='seattle_metro',
state_location='Washington',
country_location='usa',
folder_name='contact_networks',
n=10000):
# Assign students to school
gen_schools, gen_school_uids = test_send_students_to_school()
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, n)
uids_in_school = sp.get_uids_in_school(datadir,
n,
location,
state_location,
country_location,
folder_name=folder_name,
use_default=True)
potential_worker_uids, potential_worker_uids_by_age, \
potential_worker_ages_left_count = sp.get_uids_potential_workers(uids_in_school, 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)
# Assign teachers and update school lists
syn_schools, syn_school_uids, potential_worker_uids, potential_worker_uids_by_age, \
workers_by_age_to_assign_count = sp.assign_teachers_to_work(gen_schools, gen_school_uids, employment_rates,
workers_by_age_to_assign_count,
potential_worker_uids, potential_worker_uids_by_age,
potential_worker_ages_left_count,
student_teacher_ratio=30, teacher_age_min=25,
teacher_age_max=75, verbose=False)
for n in range(len(syn_schools)):
print(syn_schools[n])
assert syn_schools[n] is not None
assert syn_school_uids[n] is not None
assert syn_schools == gen_schools
assert syn_school_uids == gen_school_uids
assert potential_worker_uids == potential_worker_uids
assert potential_worker_uids_by_age == potential_worker_uids_by_age
assert workers_by_age_to_assign_count == workers_by_age_to_assign_count
def test_generate_school_sizes(location='seattle_metro', state_location='Washington',
country_location='usa', folder_name='contact_networks'):
Nhomes = 10000
uids_in_school = sp.get_uids_in_school(datadir, Nhomes, 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)
assert school_sizes is not None
def test_get_uids_potential_workers(location='seattle_metro', state_location='Washington',
country_location='usa', folder_name='contact_networks'):
Nhomes = 10000
uids_in_school = sp.get_uids_in_school(datadir, Nhomes, location,
state_location,
country_location,
folder_name=folder_name,
use_default=True)
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, Nhomes)
potential_worker_uids, potential_worker_uids_by_age, potential_worker_ages_left_count = sp.get_uids_potential_workers(
uids_in_school, employment_rates, age_by_uid_dic)
assert potential_worker_ages_left_count is not None
return potential_worker_uids, potential_worker_uids_by_age, employment_rates, age_by_uid_dic
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_generate_workplace_sizes(location='seattle_metro',
state_location='Washington',
country_location='usa'):
Npeople = 10000
uids_in_school, uids_in_school_by_age, uids_in_school_count = sp.get_uids_in_school(
datadir,
Npeople,
location,
state_location,
country_location,
use_default=True)
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, Npeople)
potential_worker_uids, potential_worker_uids_by_age, potential_worker_ages_left_count = sp.get_uids_potential_workers(
uids_in_school, 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)
print(workplace_sizes)
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
# age_by_uid_dic = sp.read_in_age_by_uid(datadir,location,state_location,country_location,Nhomes)
sc_df = sp.get_school_enrollment_rates_df(datadir, location, 'county')
rates_by_age = sp.get_school_enrollment_rates(datadir, location, 'county')
school_sizes_count = sp.get_school_sizes_by_bracket(datadir, location)
# print(school_sizes_count)
# uids_by_age_dic = sp.get_ids_by_age_dic(age_by_uid_dic)
# create_work_and_school = True
create_work_and_school = False
if create_work_and_school:
uids_in_school, uids_in_school_by_age, ages_in_school_count = sp.get_uids_in_school(
datadir, location, state_location, country_location, 'county', Nhomes)
gen_school_sizes = sp.generate_school_sizes(school_sizes_count,
uids_in_school)
# print(gen_school_sizes)
# print(age_brackets)
# print(age_by_brackets_dic[34])
gen_schools, gen_school_uids = sp.send_students_to_school(
gen_school_sizes, uids_in_school, uids_in_school_by_age,
ages_in_school_count, age_brackets, age_by_brackets_dic,
contact_matrix_dic)
# for s in range(5):
# print(Counter(gen_schools[s]))
