def split_elementary_school_types_and_save(do_save=False):
""""""
school_size_distr_by_type = sp.get_school_size_distr_by_type(**pars)
school_types = sorted(school_size_distr_by_type.keys())
# splitting up pk and es school size distributions if they're combined
if 'pk-es' in school_types:
new_distr = school_size_distr_by_type.copy()
new_distr['pk'] = school_size_distr_by_type['pk-es'].copy()
new_distr['es'] = school_size_distr_by_type['pk-es'].copy()
new_distr.pop('pk-es', None)
if do_save:
# save a copy of the old distribution to another file
location_alias_original = f"{pars['location']}_original"
write_school_size_distr_by_type(
pars['datadir'],
location_alias=location_alias_original,
location=pars['location'],
state_location=pars['state_location'],
country_location=pars['country_location'],
school_size_distr_by_type=school_size_distr_by_type)
# save a copy of the new distribution with pk-es split
location_alias_new = f"{pars['location']}_split"
write_school_size_distr_by_type(
pars['datadir'],
location_alias=location_alias_new,
location=pars['location'],
state_location=pars['state_location'],
country_location=pars['country_location'],
school_size_distr_by_type=new_distr)
def test_school_types_created():
"""
Test that unique school types are created.
Returns:
A list of the school types expected for the specified location.
"""
sp.logger.info(
f"Test that unique school types are created for each school.\nRun this first to see what school types you are working with."
)
test_pars = sc.dcp(pars)
test_pars['n'] = 20e3
pop = sp.make_population(**test_pars)
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'])
expected_school_types = sorted(expected_school_size_distr.keys())
else:
expected_school_types = [None]
schools_by_type = dict()
for i, person in pop.items():
if person['scid'] is not None:
schools_by_type.setdefault(person['scid'], set())
schools_by_type[person['scid']].add(person['sc_type'])
for s, school_type in schools_by_type.items():
assert len(
school_type
) == 1, f'Check failed. School {s} is listed as having more than one type.'
schools_by_type[s] = list(school_type)[0]
gen_school_types = sorted(set(schools_by_type.values()))
assert gen_school_types == expected_school_types, f"Check failed. generated types: {gen_school_types}, expected: {expected_school_types}"
print(
f"School types generated for {test_pars['location']}: {set(schools_by_type.values())}"
)
return list(set(schools_by_type.values()))
def test_school_types_created():
"""
Test that unique school types are created.
Returns:
A list of the school types expected for the specified location.
"""
sp.logger.info(
f"Test that unique school types are created for each school.\nRun this first to see what school types you are working with."
)
test_pars = sc.dcp(pars)
test_pars.n = settings.pop_sizes.small
pop = sp.Pop(**pars)
popdict = pop.to_dict()
loc_pars = pop.loc_pars
if pars['with_school_types']:
expected_school_size_dist = sp.get_school_size_distr_by_type(
**loc_pars)
expected_school_types = sorted(expected_school_size_dist.keys())
else:
expected_school_types = [None]
schools_by_type = dict()
for i, person in popdict.items():
if person['scid'] is not None:
schools_by_type.setdefault(person['scid'], set())
schools_by_type[person['scid']].add(person['sc_type'])
for s, school_type in schools_by_type.items():
assert len(
school_type
) == 1, f'Check failed. School {s} is listed as having more than one type.'
schools_by_type[s] = list(school_type)[0]
gen_school_types = sorted(set(schools_by_type.values()))
assert gen_school_types == expected_school_types, f"Check failed. generated types: {gen_school_types}, expected: {expected_school_types}"
print(
f"School types generated for {pars['location']}: {set(schools_by_type.values())}"
)
return list(set(schools_by_type.values()))
age_brackets = sp.get_census_age_brackets(datadir=datadir,
state_location=state_location,
country_location=country_location)
age_brackets_labels = [
str(age_brackets[b][0]) + '-' + str(age_brackets[b][-1])
for b in sorted(age_brackets.keys())
]
with_school_types = True
school_mixing_type = 'random'
school_type = None
if with_school_types:
# use synthpops to find the school types available for the location
expected_school_size_distr = sp.get_school_size_distr_by_type(
sp.datadir,
location=location,
state_location=state_location,
country_location=country_location)
school_type = sorted(expected_school_size_distr.keys())
for seed in range(1, 100, 100):
test_prefix = f"{n}_seed{seed}"
print("seed:", seed) # Random seed
params = dict(n=n,
location=location,
state_location=state_location,
country_location=country_location,
generate=True,
rand_seed=seed,
with_school_types=with_school_types,
school_mixing_type=school_mixing_type)
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