def test_plot_generated_trimmed_contact_matrix(setting_code='H', n=5000, aggregate_flag=True, logcolors_flag=True,
density_or_frequency='density', with_facilities=False, cmap='cmr.freeze_r', fontsize=16, rotation=50):
"""
Plot the age mixing matrix for a specific setting where the edges are trimmed.
Args:
setting_code (str) : name of the physial contact setting: H for households, S for schools, W for workplaces, C for community or other
n (int) : number of people in the population
aggregate_flag (book) : If True, plot the contact matrix for aggregate age brackets, else single year age contact matrix.
logcolors_flag (bool) : If True, plot heatmap in logscale
density_or_frequency (str) : If 'density', then each contact counts for 1/(group size -1) of a person's contact in a group, elif 'frequency' then count each contact. This means that more people in a group leads to higher rates of contact/exposure.
with_facilities (bool) : If True, create long term care facilities
cmap(str or matplotlib colormap) : colormap
fontsize (int) : base font size
rotation (int) : rotation for x axis labels
Returns:
A fig object.
"""
datadir = sp.datadir
state_location = 'Washington'
location = 'seattle_metro'
country_location = 'usa'
# popdict = {}
options_args = {'use_microstructure': True}
network_distr_args = {'Npop': int(n)}
# contacts = sp.make_contacts(popdict, state_location=state_location, location=location, options_args=options_args,
# network_distr_args=network_distr_args)
# contacts = sp.trim_contacts(contacts, trimmed_size_dic=None, use_clusters=False)
population = sp.make_population(n, generate=True, with_facilities=with_facilities)
age_brackets = sp.get_census_age_brackets(datadir, state_location=state_location, country_location=country_location)
age_by_brackets_dic = sp.get_age_by_brackets_dic(age_brackets)
ages = []
for uid in population:
ages.append(population[uid]['age'])
age_count = Counter(ages)
aggregate_age_count = sp.get_aggregate_ages(age_count, age_by_brackets_dic)
matrix = sp.calculate_contact_matrix(population, density_or_frequency, setting_code)
fig = sp.plot_contact_matrix(matrix, age_count, aggregate_age_count, age_brackets, age_by_brackets_dic,
setting_code, density_or_frequency, logcolors_flag, aggregate_flag, cmap, fontsize, rotation)
return fig
def plot_contact_matrix_after_intervention(n,
n_days,
interventions,
intervention_name,
location='seattle_metro',
state_location='Washington',
country_location='usa',
aggregate_flag=True,
logcolors_flag=True,
density_or_frequency='density',
setting_code='H',
cmap='cmr.freeze_r',
fontsize=16,
rotation=50):
"""
Args:
intervention (cv.intervention): a single intervention
"""
pars = sc.objdict(pop_size=n, n_days=n_days, pop_type='synthpops')
# sim = sc.objdict()
sim = cv.Sim(pars=pars, interventions=interventions)
sim.run()
age_brackets = sp.get_census_age_brackets(
sp.datadir,
state_location=state_location,
country_location=country_location)
age_by_brackets_dic = sp.get_age_by_brackets_dic(age_brackets)
ages = sim.people.age
ages = np.round(ages, 1)
ages = ages.astype(int)
max_age = max(ages)
age_count = Counter(ages)
age_count = dict(age_count)
for i in range(max_age + 1):
if i not in age_count:
age_count[i] = 0
aggregate_age_count = sp.get_aggregate_ages(age_count, age_by_brackets_dic)
matrix = calculate_contact_matrix(sim, density_or_frequency, setting_code)
fig = sp.plot_contact_matrix(matrix, age_count, aggregate_age_count,
age_brackets, age_by_brackets_dic,
setting_code, density_or_frequency,
logcolors_flag, aggregate_flag, cmap,
fontsize, rotation)
return fig