def test_Dakar(): """Test that a Dakar population can be created with the basic SynthPops API.""" sp.logger.info("Test that a Dakar population can be created with the basic SynthPops API.") sp.set_nbrackets(18) # Dakar age distributions available are up to 18 age brackets pop = sp.make_population(**pars) assert len(pop) == pars['n'], 'Check failed.' print('Check passed') sp.set_location_defaults('defaults') # Reset default values after this test is complete. sp.logger.info("Test that the default country was reset.") assert sp.default_country == 'usa', f'Check failed: default_country is {sp.default_country}' print('2nd Check passed') return pop
def test_nbrackets(test_cleanup_set_default): sp.logger.info( "Testing that nbrackets can be set outside of the recommended range and warning message returned." ) nbrackets = max(min(sp.settings.valid_nbracket_ranges), 2) # make sure new nbrackets is at least 2 sp.set_nbrackets( n=nbrackets - 1) # testing a valid outside the range currently supported. assert nbrackets - 1 == sp.settings.nbrackets, f'Check failed. sp.settings.nbrackets not set to {nbrackets-1} outside of the official supported range.' print( f'Check passed. synthpops.settings.nbrackets updated to {nbrackets-1} outside of the official supported range.' ) sp.set_nbrackets(n=nbrackets) # resetting to the default value assert nbrackets == sp.settings.nbrackets, f'Check failed. sp.settings.nbrackets not reset to {nbrackets}.' print(f'Check passed. Reset default synthpops.settings.nbrackets.') nbrackets = max(sp.settings.valid_nbracket_ranges) sp.set_nbrackets(n=nbrackets) # set to different value assert nbrackets == sp.settings.nbrackets, f'Check failed. sp.default_config.nbrackets not reset to {nbrackets}.' print(f'Check passed. Reset default synthpops.default_config.nbrackets.') test_cleanup_set_default
def make_population(pop_size, rand_seed=1, max_pop_seeds=None, do_save=True, popfile=None, cohorting=True, n_brackets=20, community_contacts=20, **kwargs): ''' Generate the synthpops population. Args: pop_size (int): number of people in the model rand_seed (int): random seed to use for generating the population max_pop_seeds (int): if supplied, take the random seed as modulus of this to limit number of populations generated do_save (bool): whether to save the population popfile (str): if so, where to save it to cohorting (bool): whether to use cohorting n_brackets (int): whether to use 16- or 20-bracket age bins community_contacts (int): how many community contacts there are kwargs (dict): passed to sp.make_population() ''' sp.set_nbrackets( n_brackets) # Essential for getting the age distribution right pars = sc.objdict( n=pop_size, rand_seed=rand_seed, with_facilities=True, use_two_group_reduction=True, average_LTCF_degree=20, ltcf_staff_age_min=20, ltcf_staff_age_max=60, with_school_types=True, average_class_size=20, inter_grade_mixing=0.1, average_student_teacher_ratio=20, average_teacher_teacher_degree=3, teacher_age_min=25, teacher_age_max=75, with_non_teaching_staff=True, # if with_non_teaching_staff is False, but generate is True, then ,average_all_staff_ratio should be average_student_teacher_ratio or 0 average_student_all_staff_ratio=11, average_additional_staff_degree=20, staff_age_min=20, staff_age_max=75, ) pars.update(kwargs) # Update any parameters # For reference re: school_types # school_mixing_type = 'random' means that students in the school have edges randomly chosen from other students, teachers, and non teaching staff across the school. Students, teachers, and non teaching staff are treated the same in terms of edge generation. # school_mixing_type = 'age_clustered' means that students in the school have edges mostly within their own age/grade, with teachers, and non teaching staff. Strict classrooms are not generated. Teachers have some additional edges with other teachers. # school_mixing_type = 'age_and_class_clustered' means that students are cohorted into classes of students of the same age/grade with at least 1 teacher, and then some have contact with non teaching staff. Teachers have some additional edges with other teachers. if cohorting: strategy = 'clustered' # students in pre-k, elementary, and middle school are cohorted into strict classrooms pars.school_mixing_type = { 'pk': 'age_and_class_clustered', 'es': 'age_and_class_clustered', 'ms': 'age_and_class_clustered', 'hs': 'random', 'uv': 'random' } else: strategy = 'normal' pars.school_mixing_type = { 'pk': 'age_clustered', 'es': 'age_clustered', 'ms': 'age_clustered', 'hs': 'random', 'uv': 'random' } if popfile is None: popfile = os.path.join( 'inputs', f'kc_{strategy}_{int(pars.n)}_seed{pars.rand_seed}.ppl') T = sc.tic() print('Making population...') # Make the population population = sp.make_population(**pars) # Convert to a popdict popdict = cv.make_synthpop(population=sc.dcp(population), community_contacts=community_contacts) school_ids = [None] * int(pop_size) teacher_flag = [False] * int(pop_size) staff_flag = [False] * int(pop_size) student_flag = [False] * int(pop_size) school_types = {'pk': [], 'es': [], 'ms': [], 'hs': [], 'uv': []} school_type_by_person = [None] * int(pop_size) schools = dict() for uid, person in population.items(): if person['scid'] is not None: school_ids[uid] = person['scid'] school_type_by_person[uid] = person['sc_type'] if person['scid'] not in school_types[person['sc_type']]: school_types[person['sc_type']].append(person['scid']) if person['scid'] in schools: schools[person['scid']].append(uid) else: schools[person['scid']] = [uid] if person['sc_teacher'] is not None: teacher_flag[uid] = True elif person['sc_student'] is not None: student_flag[uid] = True elif person['sc_staff'] is not None: staff_flag[uid] = True popdict['school_id'] = np.array(school_ids) popdict['schools'] = schools popdict['teacher_flag'] = teacher_flag popdict['student_flag'] = student_flag popdict['staff_flag'] = staff_flag popdict['school_types'] = school_types popdict['school_type_by_person'] = school_type_by_person assert sum( popdict['teacher_flag'] ), 'Uh-oh, no teachers were found: as a school analysis this is treated as an error' assert sum( popdict['student_flag'] ), 'Uh-oh, no students were found: as a school analysis this is treated as an error' # Actually create the people people_pars = dict( pop_size=pars.n, beta_layer={k: 1.0 for k in 'hswcl' }, # Since this is used to define hat layers exist beta= 1.0, # TODO: this is required for plotting (people.plot()), but shouldn't be ) people = cv.People(people_pars, strict=False, uid=popdict['uid'], age=popdict['age'], sex=popdict['sex'], contacts=popdict['contacts'], school_id=popdict['school_id'], schools=popdict['schools'], school_types=popdict['school_types'], student_flag=popdict['student_flag'], teacher_flag=popdict['teacher_flag'], staff_flag=popdict['staff_flag'], school_type_by_person=popdict['school_type_by_person']) if do_save: print(f'Saving to "{popfile}"...') sc.saveobj(popfile, people) sc.toc(T) print('Done') return people
''' Pre-generate the synthpops population including school types. Takes ~102s per seed ''' import psutil import sciris as sc import covasim as cv import synthpops as sp sp.set_nbrackets(20) # Essential for getting the age distribution right def cache_populations(seed=0, popfile=None): ''' Pre-generate the synthpops population ''' pars = sc.objdict( # pop_size = 2.25e6, pop_size=225e3, pop_type='synthpops', rand_seed=seed, ) use_two_group_reduction = True average_LTCF_degree = 20 ltcf_staff_age_min = 20 ltcf_staff_age_max = 60 with_school_types = True average_class_size = 20 inter_grade_mixing = 0.1 average_student_teacher_ratio = 20 average_teacher_teacher_degree = 3