def run(days: int, stage: int, days_per_interval: int) -> List[int]:
# init globals
ps.init_globals(seed=100)
# setup simulator config and options
sim_config = ps.sh.small_town_config
sim_opts = ps.env.PandemicSimOpts(use_contact_tracer=True) # use defaults
# make sim
sim = ps.env.PandemicSim.from_config(
sim_config,
sim_opts,
person_routine_assignment=ps.sh.DefaultPersonRoutineAssignment())
# setup viz
viz = ps.viz.GraphViz(sim,
num_stages=len(ps.sh.austin_regulations),
days_per_interval=days_per_interval)
# execute the given regulation
sim.impose_regulation(
regulation=ps.sh.austin_regulations[stage]) # stage 0
print(f'Stage {stage}:')
# run regulation steps in the simulator
for _ in trange(days, desc='Simulating day'):
sim.step_day()
viz.record(sim.state)
return viz.num_components_per_interval
def run_pandemic_sim() -> None:
"""Here we execute the simulator using austin regulations, a small town config and default person routines."""
print(
'\nA tutorial that runs the simulator using austin regulations and default person routines',
flush=True)
# init globals
ps.init_globals(seed=0)
# select a simulator config
sim_config = ps.sh.small_town_config
# make sim
sim = ps.env.PandemicSim.from_config(
sim_config,
person_routine_assignment=ps.sh.DefaultPersonRoutineAssignment())
# setup viz to show plots
viz = ps.viz.MatplotLibViz(
num_persons=sim_config.num_persons,
max_hospital_capacity=sim_config.max_hospital_capacity,
num_stages=len(ps.sh.austin_regulations),
show_stages=False)
# impose a regulation
sim.impose_regulation(regulation=ps.sh.austin_regulations[0]) # stage 0
# run regulation steps in the simulator
for _ in trange(100, desc='Simulating day'):
sim.step_day()
viz.record(sim.state)
# generate plots
viz.plot()
def run_pandemic_sim() -> None:
"""Here we execute the simulator using austin regulations, a small town config and default person routines."""
print('\nA tutorial that runs the simulator using austin regulations and default person routines', flush=True)
# init globals
ps.init_globals(seed=1)
# select a simulator config
sim_config = ps.sh.small_town_config
# make sim
sim = ps.env.PandemicSim.from_config(sim_config)
# setup viz to show plots
viz = ps.viz.SimViz.from_config(sim_config)
# impose a regulation
sim.impose_regulation(regulation=ps.sh.austin_regulations[0]) # stage 0
# run regulation steps in the simulator
for _ in trange(100, desc='Simulating day'):
sim.step_day()
viz.record(sim.state)
# generate plots
viz.plot()
def run_pandemic_gym_env() -> None:
"""Here we execute the gym envrionment wrapped simulator using austin regulations,
a small town config and default person routines."""
print('\nA tutorial that runs the OpenAI Gym environment wrapped simulator', flush=True)
# init globals
ps.init_globals(seed=0)
# select a simulator config
sim_config = ps.sh.small_town_config
# make env
env = ps.env.PandemicGymEnv.from_config(sim_config,
pandemic_regulations=ps.sh.austin_regulations,
person_routine_assignment=ps.sh.DefaultPersonRoutineAssignment())
# setup viz
viz = ps.viz.MatplotLibViz(num_persons=sim_config.num_persons,
max_hospital_capacity=sim_config.max_hospital_capacity,
num_stages=len(ps.sh.austin_regulations),
show_stages=False)
# run stage-0 action steps in the environment
env.reset()
for _ in trange(100, desc='Simulating day'):
obs, reward, done, aux = env.step(action=0) # here the action is the discrete regulation stage identifier
viz.record(obs, reward=reward)
# generate plots
viz.plot()
def test_job_counselor() -> None:
ps.init_globals()
cr = ps.env.globals.registry
assert cr
config = ps.env.PandemicSimConfig(
num_persons=10,
location_configs=[
ps.env.LocationConfig(ps.env.GroceryStore, 1, 3),
ps.env.LocationConfig(ps.env.Office, 2, 5)
])
ps.env.make_locations(config)
job_counselor = ps.env.JobCounselor(config.location_configs)
all_work_ids = cr.location_ids_of_type(ps.env.BusinessBaseLocation)
total_jobs = sum([
config.num * config.num_assignees for config in config.location_configs
])
for _ in range(total_jobs):
work_package = job_counselor.next_available_work()
assert work_package
assert work_package.work in all_work_ids
# should return None when asked for next available work id since all are taken
assert job_counselor.next_available_work() is None
def run_pandemic_gym_env() -> None:
"""Here we execute the gym envrionment wrapped simulator using austin regulations,
a small town config and default person routines."""
print(
'\nA tutorial that runs the OpenAI Gym environment wrapped simulator',
flush=True)
# init globals
ps.init_globals(seed=0)
# select a simulator config
sim_config = ps.sh.small_town_config
# make env
env = ps.env.PandemicGymEnv.from_config(
sim_config, pandemic_regulations=ps.sh.austin_regulations)
# setup viz
viz = ps.viz.GymViz.from_config(sim_config=sim_config)
# run stage-0 action steps in the environment
env.reset()
for _ in trange(100, desc='Simulating day'):
obs, reward, done, aux = env.step(
action=0
) # here the action is the discrete regulation stage identifier
viz.record((obs, reward))
# generate plots
viz.plot()
def test_single_parent_households() -> None:
ps.init_globals()
config = ps.sh.small_town_config
ps.env.make_locations(config)
ps.env.make_population(config)
cr = ps.env.globals.registry
assert cr
minor_homes_list = []
homes = cr.location_ids_of_type(ps.env.Home)
for home in homes:
household = cr.get_persons_in_location(home)
minors = 0
adults = 0
retirees = 0
for member in household:
if member.age <= 18:
minors += 1
elif 18 < member.age <= 65:
adults += 1
else:
retirees += 1
if minors > 0:
minor_homes_list.append([minors, adults, retirees])
minor_homes = np.asarray(minor_homes_list)
nm = minor_homes[:, 1] + minor_homes[:, 2]
single_parent_homes = sum(nm == 1)
assert (single_parent_homes / len(minor_homes)) > 0.22
def test_retiree_households() -> None:
ps.init_globals()
config = ps.sh.small_town_config
ps.env.make_locations(config)
ps.env.make_population(config)
cr = ps.env.globals.registry
assert cr
home_ids = cr.location_ids_of_type(ps.env.Home)
retirees_in_nursing_home = 0
all_retirees = 0
for home in home_ids:
household = cr.get_persons_in_location(home)
is_nursing_home = True
for member in household:
if member.age <= 65:
is_nursing_home = False
else:
all_retirees += 1
if is_nursing_home:
retirees_in_nursing_home += len(household)
assert (retirees_in_nursing_home / all_retirees) > 0.065
def using_sim_config() -> None:
"""In simple_worker tutorials, we created each person and location manually. However, this can become tedious if
we want to generate a large population. To remedy this, the simulator provides a config interface and this
tutorial shows how to use it."""
print('\nA tutorial to use PandemicSimConfig', flush=True)
# the first thing to do at the start of any experiment is to initialize a few global parameters
# these parameters are shared across the entire repo
ps.init_globals(seed=0)
# generate a simulator config (see `python/pandemic_simulator/script_helpers/sim_configs.py` for more configs)
sim_config = ps.env.PandemicSimConfig(
num_persons=10,
location_configs=[
ps.env.LocationConfig(location_type=ps.env.Home, num=3),
ps.env.LocationConfig(location_type=ps.env.GroceryStore, num=1),
ps.env.LocationConfig(location_type=ps.env.Office, num=1),
ps.env.LocationConfig(location_type=ps.env.School, num=1)
])
# Init simulator
sim = ps.env.PandemicSim.from_config(sim_config)
# Iterate by advancing in days by calling step_day in the simulator
for _ in trange(10, desc='Simulating day'):
sim.step_day()
def impose_regulations() -> None:
"""In all the tutorial so far, we ran the simulator under no movement restrictions (regulations). This tutorial
shows how to impose regulations."""
print('\nA tutorial to impose pandemic regulations in the environment', flush=True)
# the first thing to do at the start of any experiment is to initialize a few global parameters
# these parameters are shared across the entire repo
ps.init_globals(seed=0)
# generate a simulator config (see `python/pandemic_simulator/script_helpers/sim_configs.py` for more configs)
sim_config = ps.env.PandemicSimConfig(
num_persons=10,
location_configs=[
ps.env.LocationConfig(location_type=ps.env.Home, num=4),
ps.env.LocationConfig(location_type=ps.env.GroceryStore, num=1),
ps.env.LocationConfig(location_type=ps.env.Office, num=1),
ps.env.LocationConfig(location_type=ps.env.School, num=1)
])
# init simulator
sim = ps.env.PandemicSim.from_config(sim_config)
# define two custom pandemic regulations (see ps.sh.austin_regulations for realistic regulations)
regulation_1 = ps.env.PandemicRegulation( # moderate restriction
stay_home_if_sick=True, # stay home if sick
location_type_to_rule_kwargs={
ps.env.Office: {'lock': False}, # unlock office (if locked)
},
stage=1 # a discrete identifier for this regulation
)
regulation_2 = ps.env.PandemicRegulation( # restricted movement
stay_home_if_sick=True, # stay home if sick
location_type_to_rule_kwargs={
ps.env.Office: {'lock': True}, # also lock office
},
stage=2 # a discrete identifier for this regulation
)
# Iterate with no restrictions
for _ in trange(3, desc='Simulating day (no restrictions)'):
sim.step_day()
# Iterate after imposing stage 1 restrictions
sim.impose_regulation(regulation_1)
for _ in trange(3, desc='Simulating day (stage 1)'):
sim.step_day()
# Iterate after imposing stage 2 restrictions
sim.impose_regulation(regulation_2)
for _ in trange(3, desc='Simulating day (stage 2)'):
sim.step_day()
def simple_worker_loop() -> None:
"""A simple worker loop tutorial, where a person goes to an assigned office during work time and goes back home
after work."""
print('\nSimple worker loop tutorial', flush=True)
# the first thing to do at the start of any experiment is to initialize a few global parameters
# these parameters are shared across the entire repo
ps.init_globals(
seed=
0, # if None, the experiment is not seeded and would initialized differently each time
registry=None, # if None, a registry is created and used
# a registry does bookkeeping of all people and locations used in the experiment
)
# init locations
home = ps.env.Home()
work = ps.env.Office(
) # any subclass of BusinessLocation can be a workplace, e.g. Bar, Restaurant, Hospital, etc.
# init a worker
person = ps.env.Worker(
person_id=ps.env.PersonID(
'worker', age=35), # person_id is a unique id for this person
home=home.id, # specify the home_id that person is assigned to
work=work.id, # specify the id of the person's workplace
)
# Init simulator
sim = ps.env.PandemicSim(
locations=[work, home], # a list of all locations
persons=[person] # a list of all persons
)
# PandemicSim by default creates and uses randomized testing and an SEIR infection model
# Iterate through steps in the simulator, where each step advances an hour
for _ in trange(24, desc='Simulating hour'):
sim.step()
# Or iterate by advancing in days by calling step_day in the simulator
for _ in trange(10, desc='Simulating day'):
sim.step_day()
# The above loop iterates the simulator with no movement restrictions
# To impose restrictions, for example, Stage-2 of austin_regulations
sim.impose_regulation(ps.sh.austin_regulations[2])
# Calling step_day now will run the simulator under Stage-2 regulation
for _ in trange(10, desc='Simulating day (Under Stage-2)'):
sim.step_day()
def test_person_routine_with_status_state_trigger() -> None:
ps.init_globals()
home = ps.env.Home()
store = ps.env.GroceryStore()
r = ps.env.PersonRoutine(
start_loc=home.id,
end_loc=store.id,
start_trigger=StateRoutineTrigger(),
reset_when_done_trigger=ps.env.SimTimeRoutineTrigger(day=1))
rws = ps.env.PersonRoutineWithStatus(r)
# check flags for non-trigger state
rws.sync(ps.env.SimTime(hour=5),
person_state=PersonState(home.id, risk=ps.env.Risk.LOW))
assert not rws.started
assert not rws.due
assert not rws.done
# due should be true at state trigger
rws.sync(ps.env.SimTime(hour=5),
person_state=PersonState(home.id, risk=ps.env.Risk.HIGH))
assert rws.due
# due should remain true even after trigger
rws.sync(ps.env.SimTime(hour=11),
person_state=PersonState(home.id, risk=ps.env.Risk.LOW))
assert rws.due
# due should switch to false when routine has started
rws.started = True
rws.sync(ps.env.SimTime(hour=12),
person_state=PersonState(home.id, risk=ps.env.Risk.HIGH))
assert not rws.due
# same when done is True
rws.done = True
rws.sync(ps.env.SimTime(hour=13),
person_state=PersonState(home.id, risk=ps.env.Risk.HIGH))
assert not rws.due
# flags reset at reset trigger and due is again set (because of state trigger)
rws.sync(ps.env.SimTime(day=1, hour=0),
person_state=PersonState(home.id, risk=ps.env.Risk.HIGH))
assert not rws.started
assert rws.due
assert not rws.done
def test_person_routine_with_status_sim_time_trigger() -> None:
ps.init_globals()
home = ps.env.Home()
store = ps.env.GroceryStore()
r = ps.env.PersonRoutine(
start_loc=home.id,
end_loc=store.id,
start_trigger=ps.env.SimTimeRoutineTrigger(day=1, offset_hour=10),
reset_when_done_trigger=ps.env.SimTimeRoutineTrigger(day=1,
offset_hour=14))
rws = ps.env.PersonRoutineWithStatus(r)
# at start all flags are false
rws.sync(ps.env.SimTime(hour=0))
assert not rws.started
assert not rws.due
assert not rws.done
# due should be true at trigger
rws.sync(ps.env.SimTime(hour=10))
assert rws.due
# due should remain true even after trigger
rws.sync(ps.env.SimTime(hour=11))
assert rws.due
# due should switch to false when routine has started
rws.started = True
rws.sync(ps.env.SimTime(hour=12))
assert not rws.due
# same when done is True
rws.done = True
rws.sync(ps.env.SimTime(hour=13))
assert not rws.due
# flags reset at reset trigger
rws.sync(ps.env.SimTime(hour=14))
assert not rws.started
assert not rws.due
assert not rws.done
def test_register_persons() -> None:
ps.init_globals()
cr = ps.env.globals.registry
assert cr
config = ps.env.PandemicSimConfig(
num_persons=10,
location_configs=[
ps.env.LocationConfig(ps.env.Home, 5),
ps.env.LocationConfig(ps.env.Office, 2, 5),
ps.env.LocationConfig(ps.env.School, 1)
])
ps.env.make_locations(config)
home_id = cr.location_ids_of_type(ps.env.Home)[0]
school_id = cr.location_ids_of_type(ps.env.School)[0]
office_id = cr.location_ids_of_type(ps.env.Office)[0]
m = ps.env.Minor(ps.env.PersonID('minor_0', 3), home_id, school=school_id)
a = ps.env.Worker(ps.env.PersonID('adult_0', 36), home_id, work=office_id)
assert (m.id in cr.get_persons_in_location(home_id))
assert (a.id in cr.get_persons_in_location(home_id))
def test_location_and_person_reset() -> None:
ps.init_globals()
config = ps.sh.tiny_town_config
locations = ps.env.make_locations(config)
persons = ps.env.make_population(config)
loc_states = [copy.deepcopy(loc.state) for loc in locations]
per_states = [copy.deepcopy(per.state) for per in persons]
for loc in locations:
loc.reset()
for per in persons:
per.reset()
new_loc_states = [copy.deepcopy(loc.state) for loc in locations]
new_per_states = [copy.deepcopy(per.state) for per in persons]
for st1, st2 in zip(loc_states, new_loc_states):
assert st1 == st2
for st1, st2 in zip(per_states, new_per_states):
assert st1 == st2
def simple_worker_loop_with_routines() -> None:
"""A simple worker loop tutorial with extra routines to other locations."""
print('\nSimple worker loop with routines tutorial', flush=True)
# the first thing to do at the start of any experiment is to initialize a few global parameters
# these parameters are shared across the entire repo
ps.init_globals(
seed=0, # if None, the experiment is not seeded and would initialized differently each time
registry=None, # if None, a registry is created and used
# a registry does bookkeeping of all people and locations used in the experiment
)
# init locations
home = ps.env.Home()
work = ps.env.Office() # any subclass of BusinessLocation can be a workplace, e.g. Bar, Restaurant, Hospital, etc.
restaurant = ps.env.Restaurant()
store = ps.env.GroceryStore()
# init a worker
person = ps.env.Worker(
person_id=ps.env.PersonID('worker', age=35), # person_id is a unique id for this person
home=home.id, # specify the home_id that person is assigned to
work=work.id, # specify the id of the person's workplace
work_time=work.get_worker_work_time(), # assign work time for the person
)
# setup person routines
# routines are abstractions that define when, where and how-often a person should transition
# between locations. A routine is specified using a PersonRoutine dataclass. Here we'll show an
# example of a during-work routine to visit a restaurant for lunch.
during_work_routines = [
ps.env.PersonRoutine(
start_loc=work.id,
# the location where this routine can start. Here, we specify it to be the person's workplace id
end_loc=restaurant.id,
# the location where the person needs to go. Here, we specify it to be the restaurant
valid_time=ps.env.SimTimeTuple(hours=tuple(range(11, 14)), week_days=tuple(range(0, 5))),
# the time when this routine is due to be executed. Here, we specify it to be around noon during weekdays
)
]
# more routines can be added to the list and the routines will be executed (if and when due) in the order
# presented in the list.
person.set_during_work_routines(during_work_routines)
# similarly, let's add an outside-work routine to visit a grocery store once every week
outside_work_routines = [
ps.env.PersonRoutine(
start_loc=None, # we specify it as None, so a person can go to the store directly after work if due
end_loc=store.id, # store id
start_trigger=ps.env.SimTimeRoutineTrigger(day=7),
# notice that we set here start_trigger argument instead of valid_time. start_trigger specifies
# a trigger to enable the routine. In this case, the routine is triggered every
# seventh day. Once triggered, it is queued to be executed until it gets triggered again.
# Advanced tip: SimTimeRoutineTrigger triggers based on sim_time only. If you want to create state
# based triggers, you can implement it similar to SimTimeRoutineTrigger and use person_state to return
# a boolean (see test/environment/test_person_routines.py for an example).
)
]
person.set_outside_work_routines(outside_work_routines)
# Init simulator
sim = ps.env.PandemicSim(locations=[work, home, restaurant, store], persons=[person])
# setup viz to show plots
viz = ps.viz.SimViz(num_persons=1)
# Iterate by advancing in days by calling step_day in the simulator
for _ in trange(10, desc='Simulating day'):
sim.step_day()
viz.record(sim.state)
# show plot
viz.plot([ps.viz.PlotType.location_assignee_visits, ps.viz.PlotType.location_visitor_visits])
def test_family_households() -> None:
ps.init_globals()
config = ps.sh.small_town_config
cr = ps.env.globals.registry
assert cr
ps.env.make_locations(config)
ps.env.make_population(config)
retiree_homes_list = []
minor_homes_list = []
adult_homes_list = []
homes = cr.location_ids_of_type(ps.env.Home)
tot_persons = 0
for home in homes:
household = cr.get_persons_in_location(home)
adults = 0
minors = 0
retirees = 0
for member in household:
if member.age <= 18:
minors += 1
elif 18 < member.age <= 65:
adults += 1
else:
retirees += 1
if minors > 0:
minor_homes_list.append([minors, adults, retirees])
elif adults > 0:
adult_homes_list.append([minors, adults, retirees])
elif retirees > 0:
retiree_homes_list.append([minors, adults, retirees])
tot_persons += len(household)
assert len(minor_homes_list)
assert len(adult_homes_list)
assert len(retiree_homes_list)
assert tot_persons == config.num_persons
minor_homes = np.asarray(minor_homes_list)
adult_homes = np.asarray(adult_homes_list)
retiree_homes = np.asarray(retiree_homes_list)
# there should be non-zero homes with 1, 2, and 3 children
for i in range(1, 4):
assert len(minor_homes[minor_homes[:, 0] == i]) > 0
# each minor home must contain an adult
assert all(minor_homes[:, 1] > 0)
# minor homes in general must also have retirees for small town config
assert np.sum(minor_homes, axis=0)[2] > 0
# there should be non-zeros homes with 1 and >1 adults
assert len(adult_homes[adult_homes[:, 1] == 1]) > 0
assert len(adult_homes[adult_homes[:, 1] > 1]) > 0
# no minors in adult homes
assert np.sum(adult_homes, axis=0)[0] == 0
# adult homes in general must also have retirees for small town config
assert np.sum(adult_homes, axis=0)[2] > 0
# there should be non-zeros with only retirees and no adults and minors
assert np.sum(retiree_homes, axis=0)[2] > 0
assert np.sum(retiree_homes, axis=0)[1] == 0
assert np.sum(retiree_homes, axis=0)[0] == 0
def eval_params(params: np.ndarray,
max_episode_length: int,
trial_cnt: int = 0) -> CalibrationData:
"""Evaluate the params and return the result
:param params: spread rate and social distancing rate
:param max_episode_length: length of simulation run in days
:param trial_cnt: evaluation trial count
:returns: CalibrationData instance
"""
if trial_cnt >= MAX_EVAL_TRIALS_TO_VALID:
raise Exception(
f'Could not find a valid evaluation for the params: {params} within the specified number'
f'of trials: {MAX_EVAL_TRIALS_TO_VALID}.')
spread_rate = np.round(params[:, 0][0], decimals=3)
social_distancing = np.round(params[:, 1][0], decimals=3)
deaths = []
hospitalizations = []
seed = SEED + trial_cnt
if trial_cnt == 0:
print(
f'Running with spread rate: {spread_rate} and social distancing: {social_distancing}'
)
else:
print(f'Re-Running with a different seed: {seed}')
ps.init_globals(seed=seed)
sim_config = ps.sh.small_town_config
reset_patient_capacity(sim_config)
sim_opts = ps.env.PandemicSimOpts(infection_spread_rate_mean=spread_rate)
sim = ps.env.PandemicSim.from_config(sim_config, sim_opts)
# using swedish stage 1 regulation with the given social distancing to calibrate
covid_regulation = dataclasses.replace(ps.sh.swedish_regulations[1],
social_distancing=social_distancing)
sim.impose_regulation(regulation=covid_regulation)
hospital_ids = sim.registry.location_ids_of_type(ps.env.Hospital)
hospital_weekly = 0
for i in trange(max_episode_length, desc='Simulating day'):
sim.step_day()
state = sim.state
num_deaths = state.global_infection_summary[
ps.env.InfectionSummary.DEAD]
deaths.append(num_deaths)
num_hospitalizations = sum([
cast(ps.env.HospitalState,
state.id_to_location_state[loc_id]).num_admitted_patients
for loc_id in hospital_ids
])
hospital_weekly += num_hospitalizations
if i % 7 == 0:
hospitalizations.append(hospital_weekly)
hospital_weekly = 0
deaths_arr = np.asarray(deaths)
deaths_arr = deaths_arr[1:] - deaths_arr[:-1]
hosp_arr = np.asarray(hospitalizations)
hosp_arr = hosp_arr[1:] - hosp_arr[:-1]
eval_result = CalibrationData(deaths=deaths_arr, hospitalizations=hosp_arr)
return eval_result if eval_result.is_valid() else eval_params(
params, max_episode_length, trial_cnt=trial_cnt + 1)
def plot_household_distribution() -> None:
ps.init_globals()
config = ps.sh.small_town_config
cr = ps.env.globals.registry
assert cr
ps.env.make_locations(config)
ps.env.make_population(config)
retiree_homes_list = []
minor_homes_list = []
adult_homes_list = []
homes = cr.location_ids_of_type(ps.env.Home)
tot_persons = 0
for home in homes:
household = cr.get_persons_in_location(home)
adults = 0
minors = 0
retirees = 0
for member in household:
if member.age <= 18:
minors += 1
elif 18 < member.age <= 65:
adults += 1
else:
retirees += 1
if minors > 0:
minor_homes_list.append([minors, adults, retirees])
elif adults > 0:
adult_homes_list.append([minors, adults, retirees])
elif retirees > 0:
retiree_homes_list.append([minors, adults, retirees])
tot_persons += len(household)
minor_homes = np.asarray(minor_homes_list)
adult_homes = np.asarray(adult_homes_list)
retiree_homes = np.asarray(retiree_homes_list)
n_rows = 2
n_cols = 3
plt.figure(figsize=(3 * n_cols, 3 * n_rows))
plt_i = 0
x = np.arange(3)
colors = ['g', 'r', 'b']
ylims = [
0,
max(np.max(minor_homes.sum(axis=0)), np.max(adult_homes.sum(axis=0)),
np.max(retiree_homes.sum(axis=0))) * 1.2
]
def plot_percent(n_homes: np.ndarray) -> None:
for i, n in enumerate(n_homes):
if n > 0:
plt.text(i,
n + 5,
f'{n / config.num_persons * 100: 0.2f}%',
ha="center",
color=colors[i])
plt_i += 1
plt.subplot(n_rows, n_cols, plt_i)
plt.bar(x, minor_homes.sum(axis=0), color=colors, alpha=0.5, width=0.3)
plot_percent(minor_homes.sum(axis=0))
plt.xticks(x, ['minors', 'adults', 'retirees'], fontsize=8)
plt.title(f'{len(minor_homes)} Homes with minors')
plt.ylim(ylims)
plt_i += 1
plt.subplot(n_rows, n_cols, plt_i)
plt.bar(x, adult_homes.sum(axis=0), color=colors, alpha=0.5, width=0.3)
plot_percent(adult_homes.sum(axis=0))
plt.xticks(x, ['minors', 'adults', 'retirees'], fontsize=8)
plt.title(f'{len(adult_homes)} Homes with adults\n(no minors)')
plt.ylim(ylims)
plt_i += 1
plt.subplot(n_rows, n_cols, plt_i)
plt.bar(x, retiree_homes.sum(axis=0), color=colors, alpha=0.5, width=0.3)
retirees_in_nursing = sum(retiree_homes[:, 2])
all_retirees = (sum(minor_homes[:, 2]) + sum(adult_homes[:, 2]) +
sum(retiree_homes[:, 2]))
plt.text(
2,
retirees_in_nursing + 7,
f'{retirees_in_nursing / config.num_persons * 100: 0.2f}% (total-person)\n'
f'{retirees_in_nursing / all_retirees * 100: 0.2f}% (total-retirees)',
ha="right",
color=colors[2])
plt.xticks(x, ['minors', 'adults', 'retirees'], fontsize=8)
plt.title(f'{len(retiree_homes)} Homes with only retirees')
plt.ylim(ylims)
plt_i += 1
plt.subplot(n_rows, n_cols, plt_i)
plt.hist(minor_homes[:, 0], [0, 1, 2, 3, 4],
alpha=0.5,
facecolor=colors[0],
align='left',
rwidth=0.3)
plt.title(
f'minor occ. (minor homes) \n(avg minors/home: {minor_homes[:, 0].mean(): 0.2f})'
)
plt.ylabel('num homes')
plt.xlabel('minors')
plt_i += 1
plt.subplot(n_rows, n_cols, plt_i)
nm = minor_homes[:, 1] + minor_homes[:, 2]
plt.hist(nm,
list(range(6)),
alpha=0.5,
facecolor=colors[1],
align='left',
rwidth=0.3)
plt.title(
f'non-minor occ. (minor homes) \n(avg per home: {nm.mean(): 0.2f})')
single_parent_homes = sum(nm == 1) / len(minor_homes)
plt.text(1,
sum(nm == 1) + 5,
f'{single_parent_homes * 100: 0.2f}%\n(single-parent)',
ha="center",
color=colors[1])
plt.ylabel('num homes')
plt.xlabel('all adults (parents/relatives, retirees)')
plt.tight_layout()
plt.show()
def using_person_routine_assignment() -> None:
"""In the tutorial 2_simple_worker_loop_with_routines.py, we showed how to add routines manually for each person.
However, this can become tedious if we want to add for an entire population. To this end, the simulator provides
a person routine assignment interface and this tutorial shows how to use it."""
print('\nA tutorial to use PersonRoutineAssignment', flush=True)
# the first thing to do at the start of any experiment is to initialize a few global parameters
# these parameters are shared across the entire repo
ps.init_globals(seed=0)
# generate a simulator config (see `python/pandemic_simulator/script_helpers/sim_configs.py` for more configs)
sim_config = ps.env.PandemicSimConfig(
num_persons=10,
location_configs=[
ps.env.LocationConfig(location_type=ps.env.Home, num=3),
ps.env.LocationConfig(location_type=ps.env.GroceryStore, num=1),
ps.env.LocationConfig(location_type=ps.env.Office, num=1),
ps.env.LocationConfig(location_type=ps.env.School, num=1)
])
# define an implementation of the PersonRoutineAssignment interface (see also
# ps.sh.DefaultPersonRoutineAssignment() for a more realistic example)
class MyAssignment(ps.env.PersonRoutineAssignment):
"""This is a callable class that gets used by the simulator to assign a routine for each person."""
@property
def required_location_types(self) -> Sequence[Type[Location]]:
"""Specify the a tuple of location types that are required for this routine assignment"""
return ps.env.GroceryStore,
def __call__(self, persons: Sequence[Person]) -> None:
"""
Here, we implement a person routine for each person in the simulator.
:param persons: A sequence of all person in the simulator
:return: None
"""
# iterate through each person
for p in persons:
# we check the category of the person
if isinstance(p, ps.env.Retired):
# routines for retirees
routines = [
# we use a helper that sets up a triggered routine once every 7 days to visit the grocery store
ps.env.triggered_routine(None, ps.env.GroceryStore, 7),
]
p.set_routines(routines) # set the routine
elif isinstance(p, ps.env.Minor):
# routines for minors, not implemented in this example
pass
elif isinstance(p, ps.env.Worker):
# routines for workers
routines = [
# we use a helper that sets up a triggered routine once every 7 days to visit the grocery store
ps.env.triggered_routine(None, ps.env.GroceryStore, 7),
]
p.set_outside_work_routines(
routines) # set as a outside work routine
# Init simulator by passing the person routine assignment instance
sim = ps.env.PandemicSim.from_config(
sim_config, person_routine_assignment=MyAssignment())
# Iterate by advancing in days by calling step_day in the simulator
for _ in trange(10, desc='Simulating day'):
sim.step_day()
