def test__infectivity_profile():
hi = HealthIndexGenerator.from_file()
iss = InfectionSelectorSetter(infectivity_profile="xnexp")
infection_selector = iss.make_infection_selector(hi)
assert infection_selector.health_index_generator == hi
assert isinstance(infection_selector, InfectionSelector)
person = Person.from_attributes()
infection_selector.infect_person_at_time(person, 0)
assert person.infection
assert isinstance(person.infection.transmission, TransmissionXNExp)
iss = InfectionSelectorSetter(infectivity_profile="nature")
infection_selector = iss.make_infection_selector(hi)
person = Person.from_attributes()
infection_selector.infect_person_at_time(person, 0)
assert isinstance(person.infection.transmission, TransmissionGamma)
iss = InfectionSelectorSetter(infectivity_profile="correction_nature")
infection_selector = iss.make_infection_selector(hi)
person = Person.from_attributes()
infection_selector.infect_person_at_time(person, 0)
assert isinstance(person.infection.transmission, TransmissionGamma)
iss = InfectionSelectorSetter(infectivity_profile="constant")
infection_selector = iss.make_infection_selector(hi)
person = Person.from_attributes()
infection_selector.infect_person_at_time(person, 0)
assert isinstance(person.infection.transmission, TransmissionConstant)
def __init__(self, module_config):
self.care_home = None
self.people = []
# residents
for age in range(50, 101):
for _ in range(0, 2):
man = Person.from_attributes(sex='m', age=age)
self.people.append(man)
woman = Person.from_attributes(sex='f', age=age)
self.people.append(woman)
# workers/carers
self.super_area = MockSuperArea(module_config)
def test__carehome_for_geography(world, carehome_distributor):
# add two workers atificially
world.care_homes = CareHomes.for_areas(world.areas)
p1 = Person.from_attributes()
p1.sector = "Q"
p2 = Person.from_attributes()
p2.sector = "Q"
world.super_areas[0].workers = [p1, p2]
carehome_distributor.populate_care_home_in_areas(world.areas)
care_home = world.care_homes[0]
assert len(care_home.residents) == 24
assert len(care_home.workers) == 2
def test__smaller_than_one():
index_list = HealthIndexGenerator.from_file()
increasing_count = 0
for i in range(len(index_list.prob_lists[0])):
index_m = index_list(Person.from_attributes(age=i, sex="m"))
index_w = index_list(Person.from_attributes(age=i, sex="f"))
bool_m = np.sum(np.round(index_m, 7) <= 1)
bool_w = np.sum(np.round(index_w, 7) <= 1)
if bool_m + bool_w == 14:
increasing_count += 1
else:
increasing_count == increasing_count
assert increasing_count == 121
def test__household_mates():
house = Household()
person1 = Person.from_attributes()
house.add(person1, subgroup_type=house.SubgroupType.kids)
assert house.residents[0] == person1
person2 = Person.from_attributes()
person3 = Person.from_attributes()
house.add(person2)
house.add(person3)
assert person1 in person1.housemates
assert person2 in person1.housemates
assert person3 in person1.housemates
def test__growing_index():
index_list = HealthIndexGenerator.from_file()
increasing_count = 0
for i in range(len(index_list.prob_lists[0])):
index_m = index_list(Person.from_attributes(age=i, sex="m"))
index_w = index_list(Person.from_attributes(age=i, sex="f"))
if sum(np.sort(index_w) == index_w) != len(index_w):
increasing_count += 0
if sum(np.sort(index_m) == index_m) != len(index_m):
increasing_count += 0
assert increasing_count == 0
def test__simple_age_profile_test(selector, ):
n_people = 10000
ages = np.random.randint(low=0, high=100, size=n_people)
people = [Person.from_attributes(age=ages[n]) for n in range(n_people)]
seed = InfectionSeed(super_areas=None,
selector=selector,
age_profile={
'0-9': 0.3,
'10-39': 0.5,
'40-100': 0.2
})
choice = seed.select_from_susceptible(people,
1000,
age_profile=seed.age_profile)
ages_infected = np.array([person.age for person in people])[choice]
count = Counter(ages_infected)
count_0_9 = sum([
count_value for count_key, count_value in count.items()
if count_key < 10
])
assert count_0_9 / len(ages_infected) == pytest.approx(0.3, 0.05)
count_10_39 = sum([
count_value for count_key, count_value in count.items()
if count_key >= 10 and count_key < 40
])
assert count_10_39 / len(ages_infected) == pytest.approx(0.5, 0.05)
count_40_100 = sum([
count_value for count_key, count_value in count.items()
if count_key > 40
])
assert count_40_100 / len(ages_infected) == pytest.approx(0.2, 0.05)
def test__generate_leisure_from_world():
geography = Geography.from_file({"super_area": ["E02002135"]})
world = generate_world_from_geography(geography,
include_households=True,
include_commute=False)
world.pubs = Pubs.for_geography(geography)
world.cinemas = Cinemas.for_geography(geography)
world.groceries = Groceries.for_geography(geography)
person = Person.from_attributes(sex="m", age=27)
household = Household()
household.area = world.areas[0]
household.add(person)
person.area = geography.areas[0]
leisure = generate_leisure_for_world(
list_of_leisure_groups=["pubs", "cinemas", "groceries"], world=world)
leisure.distribute_social_venues_to_households([household])
leisure.generate_leisure_probabilities_for_timestep(0.1, False)
n_pubs = 0
n_cinemas = 0
n_groceries = 0
for _ in range(0, 1000):
subgroup = leisure.get_subgroup_for_person_and_housemates(person)
if subgroup is not None:
if subgroup.group.spec == "pub":
n_pubs += 1
elif subgroup.group.spec == "cinema":
n_cinemas += 1
elif subgroup.group.spec == "grocery":
n_groceries += 1
assert 0 < n_pubs < 100
assert 0 < n_cinemas < 100
assert 0 < n_groceries < 107
def test__isolation_compliance(selector):
isolation = Isolation(
testing_mean_time=3, testing_std_time=1, n_quarantine_days=7, compliance=0.5
)
go_isolation = set()
for _ in range(1000):
person = Person.from_attributes()
infect_person(person, selector, symptom_tag="mild")
isolation_units = IsolationUnits([IsolationUnit(area=None)])
for day in range(0, 100):
isolation.apply(
person, medical_facilities=[isolation_units], days_from_start=day
)
if 0 < day < person.infection.time_of_testing:
assert person not in isolation_units[0].people
elif (
person.infection.time_of_testing
< day
< person.infection.time_of_testing + isolation.n_quarantine_days
):
if person in isolation_units[0].people:
go_isolation.add(person.id)
else:
assert person not in isolation_units[0].people
isolation_units[0].clear()
assert np.isclose(len(go_isolation), 500, rtol=0.1)
def test__time_of_testing(selector, isolation):
person = Person.from_attributes(sex="m", age=27)
infect_person(person, selector, symptom_tag="mild")
testing_times = []
for _ in range(1000):
testing_times.append(isolation._generate_time_of_testing(person))
assert np.isclose(np.mean(testing_times), 3 + 5.3, atol=0.1)
assert np.isclose(np.std(testing_times), 1, atol=0.1)
def test__comoposition_play_groups():
kid_young = Person.from_attributes(age=3)
kid_middle = Person.from_attributes(age=8)
kid_old = Person.from_attributes(age=13)
kid_very_old = Person.from_attributes(age=16)
play_group = PlayGroup()
subgroup = play_group.get_leisure_subgroup(person=kid_young)
assert subgroup.subgroup_type == 0
subgroup = play_group.get_leisure_subgroup(person=kid_middle)
assert subgroup.subgroup_type == 1
subgroup = play_group.get_leisure_subgroup(person=kid_old)
assert subgroup.subgroup_type == 2
subgroup = play_group.get_leisure_subgroup(person=kid_very_old)
assert subgroup.subgroup_type == 2
not_kid = Person.from_attributes(age=50)
subgroup = play_group.get_leisure_subgroup(person=not_kid)
assert subgroup is None
def make_super_area():
super_area = SuperArea()
for i in range(3):
super_area.companies.append(Company(sector=i, n_workers_max=i))
person = Person.from_attributes()
person.sector = i
super_area.workers.append(person)
return super_area
def get_oc():
person = Person.from_attributes(age=90, sex="m")
area = Area(name='E00003255', super_area='E02000134', coordinates=(0, 0))
area.add(person)
super_area = SuperArea(areas=[area], coordinates=(1, 1), name='E02000134')
super_areas = SuperAreas([super_area])
health_index = HealthIndexGenerator.from_file()
return Observed2Cases.from_file(super_areas=super_areas,
health_index=health_index)
def __init__(self, module_config):
self.workers = []
n_workers = 5
# workers/carers
for _ in range(n_workers):
carer = Person.from_attributes()
carer.sector = list(module_config["sector"].keys())[0]
carer.sub_sector = None
self.workers.append(carer)
def make_dummy_world():
teacher = Person.from_attributes(age=100, sex="f")
pupil_shift_1 = Person.from_attributes(age=12, sex="f")
pupil_shift_2 = Person.from_attributes(age=5, sex="m")
pupil_shift_3 = Person.from_attributes(age=11, sex="f")
learning_center = LearningCenter(coordinates=None, n_pupils_max=None)
household = Household()
household.add(person=teacher)
household.add(person=pupil_shift_1)
household.add(person=pupil_shift_2)
household.add(person=pupil_shift_3)
learning_center.add(person=teacher,
shift=0,
subgroup_type=learning_center.SubgroupType.teachers)
learning_center.add(person=teacher,
shift=1,
subgroup_type=learning_center.SubgroupType.teachers)
learning_center.add(person=teacher,
shift=2,
subgroup_type=learning_center.SubgroupType.teachers)
learning_center.add(person=pupil_shift_1, shift=0)
learning_center.add(person=pupil_shift_2, shift=1)
learning_center.add(person=pupil_shift_3, shift=2)
world = World()
world.learning_centers = LearningCenters([learning_center],
learning_centers_tree=False,
n_shifts=3)
world.households = Households([household])
world.people = Population(
[teacher, pupil_shift_1, pupil_shift_2, pupil_shift_3])
for person in world.people.members:
person.busy = False
learning_center.clear()
household.clear()
return (
teacher,
pupil_shift_1,
pupil_shift_2,
pupil_shift_3,
learning_center,
household,
world,
)
def test__person_drags_household(leisure):
person1 = Person.from_attributes(sex="m", age=26)
person2 = Person.from_attributes(sex="f", age=26)
person3 = Person.from_attributes(sex="m", age=27)
household = Household()
household.add(person1)
household.add(person2)
household.add(person3)
person2.busy = False
person3.busy = False
social_venue = leisure.leisure_distributors["cinemas"].social_venues[0]
social_venue.add(person1)
leisure.send_household_with_person_if_necessary(
person1,
person1.leisure,
1.0,
)
for person in [person1, person2, person3]:
assert person.subgroups.leisure == social_venue.subgroups[0]
def test__mean_multiplier_reference():
comorbidity_multipliers = {"guapo": 0.8, "feo": 1.2, "no_condition": 1.0}
prevalence_reference_population = {
"feo": {
"f": {
"0-10": 0.2,
"10-100": 0.4
},
"m": {
"0-10": 0.6,
"10-100": 0.5
},
},
"guapo": {
"f": {
"0-10": 0.1,
"10-100": 0.1
},
"m": {
"0-10": 0.05,
"10-100": 0.2
},
},
"no_condition": {
"f": {
"0-10": 0.7,
"10-100": 0.5
},
"m": {
"0-10": 0.35,
"10-100": 0.3
},
},
}
health_index = HealthIndexGenerator.from_file(
comorbidity_multipliers=comorbidity_multipliers,
prevalence_reference_population=prevalence_reference_population,
)
dummy = Person.from_attributes(
sex="f",
age=40,
)
mean_multiplier_uk = (
prevalence_reference_population["feo"]["f"]["10-100"] *
comorbidity_multipliers["feo"] +
prevalence_reference_population["guapo"]["f"]["10-100"] *
comorbidity_multipliers["guapo"] +
prevalence_reference_population["no_condition"]["f"]["10-100"] *
comorbidity_multipliers["no_condition"])
assert (health_index.get_multiplier_from_reference_prevalence(
dummy.age, dummy.sex) == mean_multiplier_uk)
def test__apply_hospitalisation_correction():
health_index = HealthIndexGenerator.from_file(
adjust_hospitalisation_adults=False)
adjusted_health_index = HealthIndexGenerator.from_file(
adjust_hospitalisation_adults=True)
dummy = Person.from_attributes(
sex="f",
age=65,
)
hi = health_index(dummy)
adjusted_hi = adjusted_health_index(dummy)
np.testing.assert_allclose(adjusted_hi, hi)
dummy = Person.from_attributes(
sex="f",
age=18,
)
hi = np.diff(health_index(dummy), prepend=0., append=1.)
adjusted_hi = np.diff(adjusted_health_index(dummy), prepend=0., append=1.)
assert sum(adjusted_hi) == 1.
assert adjusted_hi[3] == pytest.approx(hi[3] / 3., rel=0.01)
assert adjusted_hi[4] == pytest.approx(hi[4] / 3., rel=0.01)
assert adjusted_hi[5] == hi[5]
assert adjusted_hi[6] == pytest.approx(hi[6], rel=0.01)
assert adjusted_hi[7] == pytest.approx(hi[7], rel=0.01)
dummy = Person.from_attributes(
sex="f",
age=40,
)
hi = np.diff(health_index(dummy), prepend=0., append=1.)
adjusted_hi = np.diff(adjusted_health_index(dummy), prepend=0., append=1.)
assert sum(adjusted_hi) == 1.
assert adjusted_hi[3] == pytest.approx(hi[3] * 0.65, rel=0.01)
assert adjusted_hi[4] == pytest.approx(hi[4] * 0.65, rel=0.01)
assert adjusted_hi[5] == hi[5]
assert adjusted_hi[6] == pytest.approx(hi[6], rel=0.01)
assert adjusted_hi[7] == pytest.approx(hi[7], rel=0.01)
def create_school(n_students, n_teachers):
school = School(
n_pupils_max=n_students,
age_min=6,
age_max=6,
coordinates=(1.0, 1.0),
sector="primary_secondary",
)
people = []
# create students
for _ in range(n_students):
person = Person.from_attributes(sex="f", age=6)
school.add(person)
people.append(person)
for _ in range(n_teachers):
person = Person.from_attributes(sex="m", age=40)
school.add(person, subgroup_type=school.SubgroupType.teachers)
people.append(person)
assert len(people) == n_students + n_teachers
assert len(school.people) == n_students + n_teachers
assert len(school.subgroups[1].people) == n_students
assert len(school.subgroups[0].people) == n_teachers
return people, school
def test__play_group_per_area(n_people):
people = [
Person.from_attributes(age=age)
for age in np.random.randint(low=3, high=16, size=n_people)
]
dummy_area = CampArea(name="dummy",
super_area=None,
coordinates=(12.0, 15.0))
areas = [dummy_area]
dummy_area.people = people
play_groups = PlayGroups.for_areas(areas=areas, venues_per_capita=1. / 20.)
assert len(play_groups) == int(np.ceil(1. / 20. * n_people))
def test__send_to_isolation(selector, isolation):
person = Person.from_attributes(sex="m", age=27)
infect_person(person, selector, symptom_tag="mild")
person.infection.time_of_testing = isolation._generate_time_of_testing(person)
isolation_units = IsolationUnits([IsolationUnit(area=None)])
for day in range(0, 100):
isolation.apply(
person, medical_facilities=[isolation_units], days_from_start=day
)
if 0 < day < person.infection.time_of_testing:
assert person not in isolation_units[0].people
elif (
person.infection.time_of_testing
< day
< person.infection.time_of_testing + isolation.n_quarantine_days
):
assert person in isolation_units[0].people
else:
assert person not in isolation_units[0].people
isolation_units[0].clear()
def populate(self,
area_name: str,
ethnicity=True,
socioecon_index=True,
comorbidity=True) -> Population:
"""
Generate a population for a given area. Age, sex and number of residents
are all based on census data for that area.
Parameters
----------
area_name
The name of an area a population should be generated for
Returns
-------
A population of people
"""
people = []
age_and_sex_generator = self.age_sex_generators[area_name]
for _ in range(age_and_sex_generator.n_residents):
if ethnicity:
ethnicity_value = age_and_sex_generator.ethnicity()
else:
ethnicity_value = None
if socioecon_index:
socioecon_index_value = age_and_sex_generator.socioecon_index()
else:
socioecon_index_value = None
person = Person.from_attributes(
age=age_and_sex_generator.age(),
sex=age_and_sex_generator.sex(),
ethnicity=ethnicity_value,
socioecon_index=socioecon_index_value,
)
if comorbidity:
person.comorbidity = generate_comorbidity(
person, self.comorbidity_data)
people.append(person) # add person to population
return Population(people=people)
def test__probability_of_leisure(leisure):
person = Person.from_attributes(sex="m", age=26)
household = Household(type="student")
household.add(person)
person.residence.group.social_venues = {
"cinemas": [leisure.leisure_distributors["cinemas"].social_venues[0]],
"pubs": [leisure.leisure_distributors["pubs"].social_venues[0]],
}
estimated_time_for_activity = 1 / (0.5 + 0.2)
times = []
times_goes_pub = 0
times_goes_cinema = 0
for _ in range(0, 300):
counter = 0
while True:
counter += 0.01
subgroup = leisure.get_subgroup_for_person_and_housemates(person)
if subgroup is None:
continue
if subgroup.group.spec == "pub":
times_goes_pub += 1
elif subgroup.group.spec == "cinema":
times_goes_cinema += 1
else:
raise ValueError
times.append(counter)
break
assert np.isclose(np.mean(times),
estimated_time_for_activity,
atol=0.2,
rtol=0)
assert np.isclose(times_goes_pub / times_goes_cinema,
0.5 / 0.2,
atol=0,
rtol=0.25)
def test__comorbidities_effect():
comorbidity_multipliers = {"guapo": 0.8, "feo": 1.2, "no_condition": 1.0}
prevalence_reference_population = {
"feo": {
"f": {
"0-10": 0.2,
"10-100": 0.4
},
"m": {
"0-10": 0.6,
"10-100": 0.5
},
},
"guapo": {
"f": {
"0-10": 0.1,
"10-100": 0.1
},
"m": {
"0-10": 0.05,
"10-100": 0.2
},
},
"no_condition": {
"f": {
"0-10": 0.7,
"10-100": 0.5
},
"m": {
"0-10": 0.35,
"10-100": 0.3
},
},
}
health_index = HealthIndexGenerator.from_file(
comorbidity_multipliers=comorbidity_multipliers,
prevalence_reference_population=prevalence_reference_population,
)
dummy = Person.from_attributes(
sex="f",
age=60,
)
feo = Person.from_attributes(sex="f", age=60, comorbidity="feo")
guapo = Person.from_attributes(sex="f", age=60, comorbidity="guapo")
dummy_health = health_index(dummy)
feo_health = health_index(feo)
guapo_health = health_index(guapo)
mean_multiplier_uk = health_index.get_multiplier_from_reference_prevalence(
dummy.age, dummy.sex)
dummy_probabilities = np.diff(dummy_health, prepend=0., append=1.)
feo_probabilities = np.diff(feo_health, prepend=0., append=1.)
guapo_probabilities = np.diff(guapo_health, prepend=0., append=1.)
np.testing.assert_allclose(
feo_probabilities[:2].sum(),
1 - comorbidity_multipliers['feo'] / mean_multiplier_uk *
dummy_probabilities[2:].sum(),
)
np.testing.assert_allclose(
feo_probabilities[2:].sum(),
comorbidity_multipliers['feo'] / mean_multiplier_uk *
dummy_probabilities[2:].sum(),
)
np.testing.assert_allclose(
guapo_probabilities[:2].sum(), 1 - comorbidity_multipliers['guapo'] /
mean_multiplier_uk * dummy_probabilities[2:].sum())
np.testing.assert_allclose(
guapo_probabilities[2:].sum(), comorbidity_multipliers['guapo'] /
mean_multiplier_uk * dummy_probabilities[2:].sum())
np.testing.assert_allclose(
guapo_probabilities[:2].sum(), 1 - comorbidity_multipliers['guapo'] /
mean_multiplier_uk * dummy_probabilities[2:].sum())
np.testing.assert_allclose(
guapo_probabilities[2:].sum(), comorbidity_multipliers['guapo'] /
mean_multiplier_uk * dummy_probabilities[2:].sum())
def test__reduce_household_visits(self, setup_policy_world):
world, pupil, student, worker, sim = setup_policy_world
super_area = world.super_areas[0]
leisure_instance = leisure.generate_leisure_for_config(
world=world, config_filename=test_config)
reduce_leisure_probabilities = ChangeLeisureProbability(
start_time="2020-03-02",
end_time="2020-03-05",
leisure_activities_probabilities={
"pubs": {
"men": {
"0-50": 0.2,
"50-100": 0.0
},
"women": {
"0-100": 0.2
},
},
},
)
policies = Policies([reduce_leisure_probabilities])
sim.activity_manager.policies = policies
sim.activity_manager.leisure = leisure_instance
sim.clear_world()
leisure_policies = LeisurePolicies.get_active_policies(
policies=policies, date=sim.timer.date)
leisure_policies.apply(date=sim.timer.date,
leisure=sim.activity_manager.leisure)
sim.activity_manager.leisure.generate_leisure_probabilities_for_timestep(
0.1, False)
original_male_pub_probabilities = sim.activity_manager.leisure.leisure_distributors[
"pubs"].male_probabilities
original_female_pub_probabilities = sim.activity_manager.leisure.leisure_distributors[
"pubs"].female_probabilities
assert str(sim.timer.date.date()) == "2020-03-01"
household = Household()
household.area = super_area.areas[0]
leisure_instance.distribute_social_venues_to_households([household])
person1 = Person.from_attributes(age=60, sex="m")
person1.area = super_area.areas[0]
household.add(person1)
person2 = Person.from_attributes(age=80, sex="f")
person2.area = super_area.areas[0]
sim.activity_manager.leisure.distribute_social_venues_to_households(
[household])
household.add(person2)
pubs1_visits_before = 0
pubs2_visits_before = 0
for _ in range(5000):
subgroup = sim.activity_manager.leisure.get_subgroup_for_person_and_housemates(
person1)
if subgroup is not None and subgroup.group.spec == "pub":
pubs1_visits_before += 1
person1.subgroups.leisure = None
subgroup = sim.activity_manager.leisure.get_subgroup_for_person_and_housemates(
person2)
if subgroup is not None and subgroup.group.spec == "pub":
pubs2_visits_before += 1
person2.subgroups.leisure = None
assert pubs1_visits_before > 0
assert pubs2_visits_before > 0
# next day leisure policies are
while str(sim.timer.date.date()) != "2020-03-02":
next(sim.timer)
leisure_policies = LeisurePolicies.get_active_policies(
policies=policies, date=sim.timer.date)
leisure_policies.apply(date=sim.timer.date,
leisure=sim.activity_manager.leisure)
sim.activity_manager.leisure.generate_leisure_probabilities_for_timestep(
0.1, False)
assert (sim.activity_manager.leisure.leisure_distributors["pubs"].
male_probabilities[60] == 0.0)
assert (sim.activity_manager.leisure.leisure_distributors["pubs"].
female_probabilities[60] == 0.2)
pubs1_visits_after = 0
pubs2_visits_after = 0
for _ in range(5000):
subgroup = sim.activity_manager.leisure.get_subgroup_for_person_and_housemates(
person1)
if subgroup is not None and subgroup.group.spec == "pub":
pubs1_visits_after += 1
person1.subgroups.leisure = None
subgroup = sim.activity_manager.leisure.get_subgroup_for_person_and_housemates(
person2)
if subgroup is not None and subgroup.group.spec == "pub":
pubs2_visits_after += 1
person2.subgroups.leisure = None
assert pubs1_visits_after == 0
assert 0 < pubs2_visits_after < pubs2_visits_before
# end of policy
while str(sim.timer.date.date()) != "2020-03-05":
next(sim.timer)
leisure_policies = LeisurePolicies.get_active_policies(
policies=policies, date=sim.timer.date)
leisure_policies.apply(date=sim.timer.date,
leisure=sim.activity_manager.leisure)
sim.activity_manager.leisure.generate_leisure_probabilities_for_timestep(
0.1, False)
pubs1_visits_restored = 0
pubs2_visits_restored = 0
for _ in range(5000):
subgroup = sim.activity_manager.leisure.get_subgroup_for_person_and_housemates(
person1)
if subgroup is not None and subgroup.group.spec == "pub":
pubs1_visits_restored += 1
person1.subgroups.leisure = None
subgroup = sim.activity_manager.leisure.get_subgroup_for_person_and_housemates(
person2)
if subgroup is not None and subgroup.group.spec == "pub":
pubs2_visits_restored += 1
person2.subgroups.leisure = None
assert np.isclose(pubs1_visits_restored, pubs1_visits_before, rtol=0.2)
assert np.isclose(pubs2_visits_restored, pubs2_visits_before, rtol=0.2)
assert (sim.activity_manager.leisure.leisure_distributors["pubs"].
male_probabilities == original_male_pub_probabilities)
assert (sim.activity_manager.leisure.leisure_distributors["pubs"].
female_probabilities == original_female_pub_probabilities)
def load_population_from_hdf5(file_path: str, chunk_size=100000):
"""
Loads the population from an hdf5 file located at ``file_path``.
Note that this object will not be ready to use, as the links to
object instances of other classes need to be restored first.
This function should be rarely be called oustide world.py
"""
print("loading population from hdf5 ", end="")
with h5py.File(file_path, "r", libver="latest", swmr=True) as f:
people = []
population = f["population"]
# read in chunks of 100k people
n_people = population.attrs["n_people"]
n_chunks = int(np.ceil(n_people / chunk_size))
for chunk in range(n_chunks):
print(".", end="")
idx1 = chunk * chunk_size
idx2 = min((chunk + 1) * chunk_size, n_people)
ids = population["id"][idx1:idx2]
ages = population["age"][idx1:idx2]
sexes = population["sex"][idx1:idx2]
ethns = population["ethnicity"][idx1:idx2]
socioecon_indices = population["socioecon_index"][idx1:idx2]
home_city = population["home_city"][idx1:idx2]
group_ids = population["group_ids"][idx1:idx2]
group_specs = population["group_specs"][idx1:idx2]
subgroup_types = population["subgroup_types"][idx1:idx2]
sectors = population["sector"][idx1:idx2]
sub_sectors = population["sub_sector"][idx1:idx2]
lockdown_status = population["lockdown_status"][idx1:idx2]
mode_of_transport_is_public_list = population[
"mode_of_transport_is_public"][idx1:idx2]
mode_of_transport_description_list = population[
"mode_of_transport_description"][idx1:idx2]
areas = population["area"][idx1:idx2]
for k in range(idx2 - idx1):
if ethns[k].decode() == " ":
ethnicity = None
else:
ethnicity = ethns[k].decode()
if socioecon_indices[k] == nan_integer:
socioecon_index = None
else:
socioecon_index = socioecon_indices[k]
person = Person.from_attributes(
id=ids[k],
age=ages[k],
sex=sexes[k].decode(),
ethnicity=ethnicity,
socioecon_index=socioecon_index,
)
mode_of_transport_description = mode_of_transport_description_list[
k]
mode_of_transport_is_public = mode_of_transport_is_public_list[
k]
# mode of transport
if mode_of_transport_description.decode() == " ":
person.mode_of_transport = None
else:
person.mode_of_transport = ModeOfTransport(
description=mode_of_transport_description.decode(),
is_public=mode_of_transport_is_public,
)
hc = home_city[k]
if hc == nan_integer:
person.home_city = None
else:
person.home_city = hc
subgroups = []
for group_id, subgroup_type, group_spec in zip(
group_ids[k], subgroup_types[k], group_specs[k]):
if group_id == nan_integer:
group_id = None
subgroup_type = None
group_spec = None
else:
group_spec = group_spec.decode()
subgroups.append([group_spec, group_id, subgroup_type])
person.subgroups = subgroups
person.area = areas[k]
if sectors[k].decode() == " ":
person.sector = None
else:
person.sector = sectors[k].decode()
if sub_sectors[k].decode() == " ":
person.sub_sector = None
else:
person.sub_sector = sub_sectors[k].decode()
if lockdown_status[k].decode() == " ":
person.lockdown_status = None
else:
person.lockdown_status = lockdown_status[k].decode()
people.append(person)
print("\n", end="")
return Population(people)
def make_dummy_world():
g = Geography.from_file(filter_key={"super_area": ["E02002559"]})
super_area = g.super_areas.members[0]
company = Company(super_area=super_area, n_workers_max=100, sector="Q")
school = School(
coordinates=super_area.coordinates,
n_pupils_max=100,
age_min=4,
age_max=10,
sector="primary",
)
household = Household()
household.area = super_area.areas[0]
hospital = Hospital(
n_beds=40,
n_icu_beds=5,
super_area=super_area.name,
coordinates=super_area.coordinates,
)
worker = Person.from_attributes(age=40)
worker.area = super_area
household.add(worker, subgroup_type=household.SubgroupType.adults)
worker.sector = "Q"
company.add(worker)
pupil = Person.from_attributes(age=6)
pupil.area = super_area
household.add(pupil, subgroup_type=household.SubgroupType.kids)
household.area = super_area
school.add(pupil)
student = Person.from_attributes(age=21)
student.area = super_area
household.add(student, subgroup_type=household.SubgroupType.adults)
university = University(
coordinates=super_area.coordinates,
n_students_max=100,
)
university.add(student)
commuter = Person.from_attributes(sex="m", age=30)
commuter.mode_of_transport = ModeOfTransport(description="bus",
is_public=True)
commuter.mode_of_transport = "public"
household.add(commuter)
world = World()
world.schools = Schools([school])
world.hospitals = Hospitals([hospital])
world.households = Households([household])
world.universities = Universities([])
world.companies = Companies([company])
world.universities = Universities([university])
world.care_homes = CareHomes([CareHome()])
world.people = Population([worker, pupil, student, commuter])
world.areas = Areas([super_area.areas[0]])
world.areas[0].people = world.people
world.super_areas = SuperAreas([super_area])
cinema = Cinema()
cinema.coordinates = super_area.coordinates
world.cinemas = Cinemas([cinema])
pub = Pub()
pub.coordinates = super_area.coordinates
world.pubs = Pubs([pub])
grocery = Grocery()
grocery.coordinates = super_area.coordinates
world.groceries = Groceries([grocery])
# commute
city = CommuteCity()
hub = CommuteHub(None, None)
city.commutehubs = [hub]
world.commutehubs = CommuteHubs([city])
world.commutehubs.members = [hub]
world.commutecities = CommuteCities()
world.commutecities.members = [city]
world.commutehubs[0].add(commuter)
world.commuteunits = CommuteUnits(world.commutehubs.members)
world.commuteunits.init_units()
world.commutecityunits = CommuteCityUnits(world.commutecities)
world.cemeteries = Cemeteries()
return world
def make_pop():
people = Population([])
for i in range(100):
people.add(Person.from_attributes(age=i))
return people
def test__distribute_teachers():
dummy_area = Area(name="dummy", super_area=None, coordinates=(12.0, 15.0))
dummy_areas = Areas(areas=[dummy_area], )
people = [Person.from_attributes(sex="f", age=age) for age in range(100)]
for person in people:
person.area = dummy_area
dummy_area.people = people
male_enrollment_rates = {
"dummy_region": {
"0-4": 0.0,
"4-6": 0.5,
"6-12": 0.3,
"12-16": 0.2,
"16-21": 0.1,
"21-100": 0.0,
}
}
female_enrollment_rates = {
"dummy_region": {
"0-4": 0.0,
"4-6": 0.0,
"6-12": 1.0,
"12-16": 0.0,
"16-21": 0.0,
"21-100": 0.0,
}
}
area_region_df = pd.DataFrame({
'area': ['dummy'],
'region': ['dummy_region']
})
coordinates_1 = (12.3, 15.6)
learning_center_1 = LearningCenter(coordinates=coordinates_1,
n_pupils_max=20)
coordinates_2 = (120.3, 150.6)
learning_center_2 = LearningCenter(coordinates=coordinates_2,
n_pupils_max=20)
coordinates = np.vstack(
(np.array(coordinates_1), np.array(coordinates_2))).T
learning_centers_tree = BallTree(np.deg2rad(coordinates),
metric="haversine")
learning_centers = LearningCenters(
learning_centers=[learning_center_1, learning_center_2],
learning_centers_tree=learning_centers_tree,
n_shifts=3)
learning_center_distributor = LearningCenterDistributor(
learning_centers=learning_centers,
female_enrollment_rates=female_enrollment_rates,
male_enrollment_rates=male_enrollment_rates,
area_region_df=area_region_df)
learning_center_distributor.distribute_teachers_to_learning_centers(
areas=dummy_areas)
for learning_center in learning_centers:
assert len(learning_center.teachers) == 3
assert len(learning_center.ids_per_shift[0]) == 1
assert learning_center.ids_per_shift[
0] == learning_center.ids_per_shift[1]
assert learning_center.ids_per_shift[
1] == learning_center.ids_per_shift[2]
assert learning_center.teachers[0].age >= 21
def test__age_distribution():
dummy_area = Area(name="dummy", super_area=None, coordinates=(12.0, 15.0))
dummy_areas = Areas(areas=[dummy_area], )
people = [Person.from_attributes(sex="f", age=age) for age in range(100)]
for person in people:
person.area = dummy_area
dummy_area.people = people
male_enrollment_rates = {
"dummy_region": {
"0-4": 0.0,
"4-6": 0.5,
"6-12": 0.3,
"12-16": 0.2,
"16-21": 0.1,
"21-100": 0.0,
}
}
female_enrollment_rates = {
"dummy_region": {
"0-4": 0.0,
"4-6": 0.0,
"6-12": 1.0,
"12-16": 0.0,
"16-21": 0.0,
"21-100": 0.0,
}
}
area_region_df = pd.DataFrame({
'area': ['dummy'],
'region': ['dummy_region']
})
coordinates_1 = (12.3, 15.6)
learning_center_1 = LearningCenter(coordinates=coordinates_1,
n_pupils_max=20)
coordinates_2 = (120.3, 150.6)
learning_center_2 = LearningCenter(coordinates=coordinates_2,
n_pupils_max=20)
coordinates = np.vstack(
(np.array(coordinates_1), np.array(coordinates_2))).T
learning_centers_tree = BallTree(np.deg2rad(coordinates),
metric="haversine")
learning_centers = LearningCenters(
learning_centers=[learning_center_1, learning_center_2],
learning_centers_tree=learning_centers_tree,
)
learning_center_distributor = LearningCenterDistributor(
learning_centers=learning_centers,
female_enrollment_rates=female_enrollment_rates,
male_enrollment_rates=male_enrollment_rates,
area_region_df=area_region_df)
learning_center_distributor.distribute_kids_to_learning_centers(
areas=dummy_areas)
for kid in people:
if kid.age < 6 or kid.age >= 12:
assert kid.primary_activity is None
else:
assert kid.primary_activity.group.spec == "learning_center"