class MyModel(Model):
def __init__(
self, student_class, params
): #student_num, sick_num, bus_cols, bus_rows, highrisk_radius, highrisk, lowrisk_radius, lowrisk, masks):
# mesa required attributes
self.running = True # determines if model should keep on running
# should be specified to false when given conditions are met
self.grid = GeoSpace(
) # To learn more about other type of space grid, check mesa documentation
self.schedule = SimultaneousActivation(
self
) # scheduler dictates model level agent behavior, aka. step function
# Here we are using a BaseScheduler which computes step of the agents by order
# To learn more about other type of scheduler, check mesa documentation
student_num = params['student_num']
sick_num = params['sick_num']
bus_cols = params['bus_cols']
bus_rows = params['bus_rows']
breathe_rate = params["breathe_rate"],
windows_open = params["windows_open"],
seating_array = params["seating_array"],
highrisk_radius = params['highrisk_radius']
highrisk = params['highrisk']
lowrisk_radius = params['lowrisk_radius']
lowrisk = params['lowrisk']
masks = params['masks']
bus_stop_student_count = params['bus_stop_student_count']
bus_stop_minutes_count = params['bus_stop_minutes_count']
sick = 0
locs = []
# initializing the model by making the students and putting them into their seats
stop_counter = 0
for i in range(student_num):
bus_stop = bus_stop_minutes_count[bus_stop_student_count.index(
min(filter(lambda x: x > stop_counter,
bus_stop_student_count)))]
stop_counter += 1
# finding an empty seat for the next student
new = False
while (new == False):
loc = (np.random.randint(bus_cols),
np.random.randint(bus_rows))
if loc not in locs:
new = True
locs.append(loc)
pnt = Point(loc)
# adding sick and healthy students
if sick < sick_num:
sick += 1
a = Student(model=self,
shape=pnt,
unique_id="Passenger #" + str(i),
sick=True,
mask=False,
params=params,
spreads=True,
bus_stop=bus_stop,
breathe_rate=breathe_rate)
else:
a = Student(model=self,
shape=pnt,
unique_id="Passenger #" + str(i),
sick=False,
mask=False,
params=params,
spreads=False,
bus_stop=bus_stop,
breathe_rate=breathe_rate)
self.grid.add_agents(a)
self.schedule.add(a)
def step(self):
'''
step function of the model that would essentially call the step function of all agents
'''
self.schedule.step()
#self.schedule.advance()
self.grid._recreate_rtree(
) # this is some history remaining issue with the mesa-geo package
class School(Model):
def __init__(self,
map_path,
schedule_path,
grade_N,
KG_N,
preschool_N,
special_education_N,
faculty_N,
seat_dist,
init_patient=3,
attend_rate=1,
mask_prob=0.516,
inclass_lunch=False,
username="******"):
# zipcode etc, for access of more realistic population from KG perhaps
# model param init
self.__mask_prob = mask_prob
self.inclass_lunch = inclass_lunch
self.seat_dist = math.ceil(seat_dist / (attend_rate**(1 / 2)))
self.idle_teachers = [] # teachers to be assigned without a classroom
self.init_patient = init_patient
# mesa model init
self.running = True
self.grid = GeoSpace()
self.schedule = BaseScheduler(self)
#data collect init
model_reporters = {
"day": "day_count",
"cov_positive": "infected_count"
}
agent_reporters = {
"unique_id": "unique_id",
"health_status": "health_status",
"symptoms": "symptoms",
"x": "x",
"y": "y",
"viral_load": "viral_load"
}
self.datacollector = datacollection.DataCollector(
model_reporters=model_reporters, agent_reporters=agent_reporters)
school_gdf = load_map(map_path)
# room agent init
self.room_agents = school_gdf.apply(
lambda x: Classroom(unique_id=x["Id"],
model=self,
shape=x["geometry"],
room_type=x["room_type"]),
axis=1).tolist()
self.grid.add_agents(self.room_agents)
# stats tracking init
self.infected_count = 0
self.step_count = 0
self.day_count = 0
self.num_exposed = 0
# student activity init
self.schoolday_schedule = pd.read_csv(schedule_path)
self.activity = None
# id tracking init
self.__teacher_id = 0
self.__student_id = 0
self.__faculty_N = faculty_N
self.schedule_ids = self.schoolday_schedule.columns
self.recess_yards = find_room_type(self.room_agents, 'recess_yard')
def init_agents(room_type, N, partition=False):
'''
batch initialize human agents into input room type rooms with equal partition size
room_type: a valid string of room type: [None, 'restroom_grade_boys', 'lunch_room', 'classroom_grade',
'restroom_all', 'restroom_grade_girls', 'restroom_KG',
'classroom_KG', 'community_room', 'library',
'restroom_special_education', 'restroom_faculty',
'classroom_special_education', 'health_room', 'faculty_lounge',
'classroom_preschool', 'restroom_preschool']
'''
rooms = find_room_type(self.room_agents, room_type)
# if student group should be seperated to different day schedules
# assigning schedule_id to equally partitioned rooms
# currently only grade 1-5 "grade" students need to be partitioned,
partition_size = len(rooms)
if partition:
partition_size = math.ceil(partition_size /
len(self.schedule_ids))
class_size = N // len(rooms)
remaining_size = N % len(rooms)
for i, classroom in zip(range(len(rooms)), rooms):
classroom.generate_seats(class_size, self.seat_dist)
classroom.schedule_id = self.schedule_ids[i // partition_size]
for idx in range(class_size):
pnt = classroom.seats[idx]
mask_on = np.random.choice([True, False],
p=[mask_prob, 1 - mask_prob])
agent_point = Student(model=self,
shape=pnt,
unique_id="S" +
str(self.__student_id),
room=classroom,
mask_on=mask_on)
self.grid.add_agents(agent_point)
self.schedule.add(agent_point)
self.__student_id += 1
# spread remaining student into all classrooms
if remaining_size > 0:
pnt = classroom.seats[class_size]
mask_on = np.random.choice([True, False],
p=[mask_prob, 1 - mask_prob])
agent_point = Student(model=self,
shape=pnt,
unique_id="S" +
str(self.__student_id),
room=classroom,
mask_on=mask_on)
self.grid.add_agents(agent_point)
self.schedule.add(agent_point)
self.__student_id += 1
remaining_size -= 1
#add teacher to class
pnt = generate_random(classroom.shape)
agent_point = Teacher(model=self,
shape=pnt,
unique_id="T" + str(self.__teacher_id),
room=classroom)
self.grid.add_agents(agent_point)
self.schedule.add(agent_point)
self.idle_teachers.append(agent_point)
self.__teacher_id += 1
self.__faculty_N -= 1
# initialize all students and teachers in classrooms
init_agents("classroom_grade",
int(grade_N * attend_rate),
partition=True)
# keep track of student types
#self.grade_students = [a for a in list(self.schedule.agents) if isinstance(a, Student)]
init_agents("classroom_KG", int(KG_N * attend_rate))
init_agents("classroom_preschool", int(preschool_N * attend_rate))
#self.pkg_students = [a for a in list(set(self.schedule.agents).difference(self.grade_students)) if isinstance(a, Student)]
init_agents("classroom_special_education",
int(special_education_N * attend_rate))
# dump remaining teacher to faculty lounge
for f_lounge in find_room_type(self.room_agents, "faculty_lounge"):
f_lounge.schedule_id = self.schedule_ids[0]
for i in range(self.__faculty_N):
pnt = generate_random(f_lounge.shape)
agent_point = Teacher(model=self,
shape=pnt,
unique_id="T" + str(self.__teacher_id),
room=f_lounge)
self.grid.add_agents(agent_point)
self.schedule.add(agent_point)
self.__teacher_id += 1
#self.people = list(self.schedule.agents)
# add rooms to scheduler at last
for room in self.room_agents:
self.schedule.add(room)
self.lunchroom = find_room_type(self.room_agents, 'lunch_room')[0]
self.lunchroom.generate_seats_lunch(3, 12)
def small_step(self):
self.schedule.step()
self.grid._recreate_rtree()
def add_N_patient(self, N):
patients = random.sample(
[a for a in self.schedule.agents if isinstance(a, Student)], N)
for p in patients:
p.health_status = "exposed"
p.asymptomatic = True
p.infective = True
def show(self):
'''
plot current step visualization
deprecated since end of model visualization update
'''
# UPDATE 10/16: add deprecation warning
message = "this function is no longer used for performance issues, check output_image.py for end of model visualization"
warnings.warn(message, DeprecationWarning)
school_geometry = gpd.GeoSeries([a.shape for a in self.room_agents])
school_map = gpd.GeoDataFrame(
{"viral_load": [min(a.viral_load, 5) for a in self.room_agents]})
school_map.geometry = school_geometry
basemap = school_map.plot(column="viral_load",
cmap="Reds",
alpha=0.5,
vmin=0,
vmax=5)
school_map.boundary.plot(ax=basemap, color='k', linewidth=0.2)
list(
map(lambda a: a.plot(), [
a for a in self.schedule.agents if issubclass(type(a), Human)
]))
hour = 9 + self.step_count * 5 // 60 # assume plot start at 9am
minute = self.step_count * 5 % 60
plt.title("Iteration: Day {}, ".format(self.day_count + 1) +
"%d:%02d" % (hour, minute),
fontsize=30)
def __update_day(self):
'''
update incubation time, reset viral_load, remove symptomatic agents, etc for end of day
'''
for a in self.schedule.agents[:]:
if issubclass(type(a), Human):
if a.symptoms:
# remove agent if symptom onset
if isinstance(a, Teacher):
# assign a new teacher to position
new_teacher = self.idle_teachers.pop()
new_teacher.shape = a.shape
new_teacher.room = a.room
new_teacher.classroom = a.classroom
self.schedule.remove(a)
self.grid.remove_agent(a)
# UPDATE 10/16: infectious made obsolete, end of day update rework
elif a.health_status == "exposed":
# UPDATE 10/17: update infective delay if agent is not infective by end of day
a.infective = True
a.symptom_countdown -= 1
# calculate when symptoms begin to show using 0-15 density
if a.symptom_countdown <= 0:
if a.symptom_countdown == 0:
self.infected_count += 1
# update model stat for total infected
# negative countdown means this agent is asymptomatic
if not a.asymptomatic:
# this is a really small chance, however possible
# set symtoms to true
# next day this agent will be removed from the model
a.symptoms = True
else:
# reset viral_load of room agents
a.viral_load = 0
def step(self):
'''
simulate a day with school day schedule
'''
if not self.schedule.steps:
self.add_N_patient(self.init_patient)
for i, row in self.schoolday_schedule.iterrows():
self.activity = row
self.datacollector.collect(self)
self.schedule.step()
self.grid._recreate_rtree()
self.step_count += 1
self.__update_day()
self.grid._recreate_rtree()
self.day_count += 1
self.step_count = 0
class NaiveModel(Model):
'''
this class represents the environment that the agents are present in, and the steps are the steps in the
agent based model. There are a few parameters that are included in this class:
- agent_class: this is basically what the agent is, which in our case is the busAgent that we created earlier
- dim_bus: these are the dimensions of the bus. the format is a list: [num_left_columns, num_middle_columns,
num right_columns, num_rows]. There will not be passengers in the middle column; this is simply to provide
some area between passengers (which is realistic amongst the bus).
- distance_seats: this sets the distance between seats in feet. This will affect how close the passangers are
to each other and thus affect the rate of infection
- num_infected: this is the number of people initially infected
'''
def __init__(self, agent_class, num_col_left, num_col_mid, num_col_right,
num_row, dist_bw_seats, num_infected, breath_prob, cough_prob,
sneeze_prob, breath_dist, cough_dist, sneeze_dist,
prob_infected):
# mesa required attributes
self.running = True
self.grid = GeoSpace()
self.schedule = BaseScheduler(
self
) # scheduler dictates model level agent behavior, aka. step function
# variables used for later functions that need descriptions of the model
dim_bus = [num_col_left, num_col_mid, num_col_right, num_row]
self.max_columns = (dim_bus[0] + dim_bus[1] +
dim_bus[2]) * dist_bw_seats
self.max_rows = dim_bus[3] * dist_bw_seats
self.seat_dist = dist_bw_seats
i = 1
for x in range(0, dim_bus[0] * dist_bw_seats, dist_bw_seats):
for y in range(0, dim_bus[3] * dist_bw_seats, dist_bw_seats):
pnt = Point(x, y)
i += 1
a = agent_class(model=self,
shape=pnt,
unique_id="na" + str(i),
breath_prob=breath_prob,
cough_prob=cough_prob,
sneeze_prob=sneeze_prob,
breath_dist=breath_dist,
cough_dist=cough_dist,
sneeze_dist=sneeze_dist,
prob_infected=prob_infected)
self.grid.add_agents(a)
self.schedule.add(a)
for x in range((dim_bus[0] + dim_bus[1]) * dist_bw_seats,
(dim_bus[0] + dim_bus[1] + dim_bus[2]) * dist_bw_seats,
dist_bw_seats):
for y in range(0, dim_bus[3] * dist_bw_seats, dist_bw_seats):
pnt = Point(x, y)
i += 1
a = agent_class(model=self,
shape=pnt,
unique_id="na" + str(i),
breath_prob=breath_prob,
cough_prob=cough_prob,
sneeze_prob=sneeze_prob,
breath_dist=breath_dist,
cough_dist=cough_dist,
sneeze_dist=sneeze_dist,
prob_infected=prob_infected)
self.grid.add_agents(a)
self.schedule.add(a)
infected_agents = random.sample(self.grid.agents, num_infected)
for i in infected_agents:
i.infected = True
def step(self):
'''
step function of the model that would essentially call the step function of all agents
'''
self.schedule.step()
self.grid._recreate_rtree(
) # this is some history remaining issue with the mesa-geo package
class School(Model):
schedule_types = {
"Sequential": BaseScheduler,
"Random": RandomActivation,
"Simultaneous": SimultaneousActivation
}
def __init__(self,
map_path,
schedule_path,
grade_N,
KG_N,
preschool_N,
special_education_N,
faculty_N,
seat_dist,
init_patient=3,
attend_rate=1,
mask_prob=0.516,
inclass_lunch=False,
student_vaccine_prob=0,
student_testing_freq=14,
teacher_vaccine_prob=1,
teacher_testing_freq=7,
teacher_mask='N95',
schedule_type="Simultaneous"):
# zipcode etc, for access of more realistic population from KG perhaps
# model param init
self.__mask_prob = mask_prob
self.inclass_lunch = inclass_lunch
self.seat_dist = math.ceil(seat_dist / (attend_rate**(1 / 2)))
self.idle_teachers = [] # teachers to be assigned without a classroom
self.init_patient = init_patient
# testing param init
self.teacher_testing_freq = teacher_testing_freq
self.student_testing_freq = student_testing_freq
# mesa model init
self.running = True
self.grid = GeoSpace()
self.schedule_type = schedule_type
self.schedule = self.schedule_types[self.schedule_type](self)
#data collect init
model_reporters = {
"day": "day_count",
"cov_positive": "infected_count"
}
agent_reporters = {
"unique_id": "unique_id",
"health_status": "health_status",
"symptoms": "symptoms",
"x": "x",
"y": "y",
"viral_load": "viral_load"
}
self.datacollector = datacollection.DataCollector(
model_reporters=model_reporters, agent_reporters=agent_reporters)
school_gdf = gpd.read_file(map_path)
# minx miny maxx maxy
# use minx maxy
# gdf.dessolve
# minx miny maxx maxy = geometery.bounds
# for loop:
# bus = bus(shape(minx,maxy))
# minx = minx - width
# maxy = maxy + length
# room agent init
self.room_agents = school_gdf.apply(
lambda x: room_agent.Classroom(unique_id=x["Id"],
model=self,
shape=x["geometry"],
room_type=x["room_type"]),
axis=1).tolist()
self.grid.add_agents(self.room_agents)
# stats tracking init
self.infected_count = 0
self.step_count = 0
self.day_count = 1
self.num_exposed = 0
# student activity init
self.schoolday_schedule = pd.read_csv(schedule_path)
self.activity = None
# id tracking init
self.__teacher_id = 0
self.__student_id = 0
self.__cohort_id = 0
self.__faculty_N = faculty_N
self.schedule_ids = self.schoolday_schedule.columns
# geo-object tracking init
self.recess_yards = util.find_room_type(self.room_agents,
'recess_yard')
self.cohorts = []
# UPDATE Christmas cohort generation
def generate_cohorts(students, N):
'''
generate cohorts with within/out-of classroom probability, cohort size probablity
example: students have 80% chance to have a friend in same room, 20% chance to have a friend in different room
and 50% to have a cohort size of 5, 20% size 2, 15% size 4, 10% size 3, 5% size 1
students: a 2d list containing list of students in each room
students[k] is a list of student agents (in same classroom room)
'''
size_prob = eval(cohort_config['size_prob'])
same_room_prob = eval(cohort_config['same_room_prob'])
radius = eval(cohort_config['radius'])
size_val_list = list(size_prob.keys())
size_prob_list = list(size_prob.values())
same_room_val_list = list(same_room_prob.keys())
same_room_prob_list = list(same_room_prob.values())
# start at room 0
cur_room = 0
while N > 0:
cur_size = np.random.choice(size_val_list, p=size_prob_list)
if N <= max(size_val_list):
cur_size = N
# get same-room cohort size
cur_same = sum(
np.random.choice(same_room_val_list,
size=cur_size,
p=same_room_prob_list))
# add students from current room to current cohort
cur_same = min(cur_same, len(students[cur_room]))
cohort = students[cur_room][:cur_same]
students[cur_room] = students[cur_room][cur_same:]
room_idx = list(range(len(students)))
other_room = room_idx[:]
other_room.remove(cur_room)
# add students from other rooms to cohort
if not len(other_room):
rand_room = [cur_room] * (cur_size - cur_same)
else:
rand_room = np.random.choice(other_room,
size=(cur_size - cur_same))
for r in rand_room:
# update and remove r if r is an empty room
while True:
try:
cohort.append(students[r][0])
students[r] = students[r][1:]
break
except:
if r in other_room:
other_room.remove(r)
if not len(other_room):
r = cur_room
else:
r = np.random.choice(other_room)
# TODO: recess yard is current hard coded
recess_yard = self.recess_yards[0]
if cohort[0].grade != 'grade':
recess_yard = self.recess_yards[1]
# make cohort agent with dummy shape
cur_cohort = cohort_agent.Cohort(
"Cohort" + str(self.__cohort_id), self, Point(0, 0),
cohort, recess_yard, cur_size * radius)
self.grid.add_agents(cur_cohort)
self.schedule.add(cur_cohort)
self.cohorts.append(cur_cohort)
self.__cohort_id += 1
# remove empty rooms
students = [room for room in students if len(room) > 0]
# rolling update to minimize student pop edge cases
# fail safe break
if not len(students):
break
cur_room = (cur_room + 1) % len(students)
# update student population
N -= cur_size
def init_agents(room_type, N, partition=False):
'''
batch initialize human agents into input room type rooms with equal partition size
room_type: a valid string of room type: [None, 'restroom_grade_boys', 'lunch_room', 'classroom_grade',
'restroom_all', 'restroom_grade_girls', 'restroom_KG',
'classroom_KG', 'community_room', 'library',
'restroom_special_education', 'restroom_faculty',
'classroom_special_education', 'health_room', 'faculty_lounge',
'classroom_preschool', 'restroom_preschool']
'''
rooms = util.find_room_type(self.room_agents, room_type)
# if student group should be seperated to different day schedules
# assigning schedule_id to equally partitioned rooms
# currently only grade 1-5 "grade" students need to be partitioned,
partition_size = len(rooms)
if partition:
partition_size = math.ceil(partition_size /
len(self.schedule_ids))
class_size = N // len(rooms)
remaining_size = N % len(rooms)
#track all students of same grade type
all_students = []
for i, classroom in zip(range(len(rooms)), rooms):
# spread remaining student into all classrooms
c_size = class_size
if remaining_size > 0:
remaining_size -= 1
c_size += 1
#each classroom has its own possibility to have circular desks instead of normal grid seating
#TODO: strongly believe this is subject to change
prob_circular = eval(population_config['circular_desk_prob'])
if np.random.choice([True, False],
p=[prob_circular, 1 - prob_circular]):
classroom.generate_seats(c_size,
self.seat_dist,
style='circular')
else:
classroom.generate_seats(c_size, self.seat_dist)
classroom.schedule_id = self.schedule_ids[i // partition_size]
#track students within the same room
students = []
for idx in range(c_size):
pnt = classroom.seats[idx]
mask_on = np.random.choice([True, False],
p=[mask_prob, 1 - mask_prob])
agent_point = human_agent.Student(model=self,
shape=pnt,
unique_id="S" +
str(self.__student_id),
room=classroom,
mask_on=mask_on)
# vaccinate students accordingly
agent_point.vaccinated = np.random.choice(
[True, False],
p=[student_vaccine_prob, 1 - student_vaccine_prob])
if classroom.seating_pattern == 'circular':
desks = gpd.GeoSeries(classroom.desks)
agent_point.desk = desks[desks.distance(
agent_point.shape).sort_values().index[0]]
self.grid.add_agents(agent_point)
self.schedule.add(agent_point)
self.__student_id += 1
# add student to room temp list
students.append(agent_point)
#add teacher to class
pnt = util.generate_random(classroom.shape)
agent_point = human_agent.Teacher(model=self,
shape=pnt,
unique_id="T" +
str(self.__teacher_id),
room=classroom)
# teacher mask/vaccination protocol
agent_point.vaccinated = np.random.choice(
[True, False],
p=[teacher_vaccine_prob, 1 - teacher_vaccine_prob])
agent_point.mask_type = teacher_mask
agent_point.mask_passage_prob = trans_rate.return_mask_passage_prob(
teacher_mask)
self.grid.add_agents(agent_point)
self.schedule.add(agent_point)
self.__teacher_id += 1
self.__faculty_N -= 1
# add room students list to all students
# shuffle students for efficiency improvement
np.random.shuffle(students)
all_students.append(students)
#UPDATE Christmas
#generate cohort with temp student list
generate_cohorts(all_students, N)
# initialize all students and teachers in classrooms
init_agents("classroom_grade",
int(grade_N * attend_rate),
partition=True)
# keep track of student types
#self.grade_students = [a for a in list(self.schedule.agents) if isinstance(a, Student)]
init_agents("classroom_KG", int(KG_N * attend_rate))
init_agents("classroom_preschool", int(preschool_N * attend_rate))
#self.pkg_students = [a for a in list(set(self.schedule.agents).difference(self.grade_students)) if isinstance(a, Student)]
init_agents("classroom_special_education",
int(special_education_N * attend_rate))
# dump remaining teacher to faculty lounge
for f_lounge in util.find_room_type(self.room_agents,
"faculty_lounge"):
f_lounge.schedule_id = self.schedule_ids[0]
for i in range(self.__faculty_N):
pnt = util.generate_random(f_lounge.shape)
agent_point = human_agent.Teacher(model=self,
shape=pnt,
unique_id="T" +
str(self.__teacher_id),
room=f_lounge)
# teacher mask/vaccination protocol
agent_point.vaccinated = np.random.choice(
[True, False],
p=[teacher_vaccine_prob, 1 - teacher_vaccine_prob])
agent_point.mask_type = teacher_mask
agent_point.mask_passage_prob = trans_rate.return_mask_passage_prob(
teacher_mask)
self.grid.add_agents(agent_point)
self.schedule.add(agent_point)
#teacher from faculty lounge can be used later if on duty teachers test positive
self.idle_teachers.append(agent_point)
self.__teacher_id += 1
# add rooms to scheduler at last
for room in self.room_agents:
self.schedule.add(room)
self.lunchroom = util.find_room_type(self.room_agents, 'lunch_room')[0]
self.lunchroom.generate_seats_lunch(1, 4)
def small_step(self):
self.schedule.step()
self.grid._recreate_rtree()
def add_N_patient(self, N):
patients = random.sample([
a for a in self.schedule.agents
if isinstance(a, human_agent.Student)
], N)
for p in patients:
p.health_status = "exposed"
p.asymptomatic = True
p.infective = True
def show(self):
'''
plot current step visualization
deprecated since end of model visualization update
'''
# UPDATE 10/16: add deprecation warning
message = "this function is no longer used for performance issues, check output_image.py for end of model visualization"
warnings.warn(message, DeprecationWarning)
school_geometry = gpd.GeoSeries([a.shape for a in self.room_agents])
school_map = gpd.GeoDataFrame(
{"viral_load": [min(a.viral_load, 5) for a in self.room_agents]})
school_map.geometry = school_geometry
basemap = school_map.plot(column="viral_load",
cmap="Reds",
alpha=0.5,
vmin=0,
vmax=5)
school_map.boundary.plot(ax=basemap, color='k', linewidth=0.2)
list(
map(lambda a: a.plot(), [
a for a in self.schedule.agents
if issubclass(type(a), human_agent.Human)
]))
hour = 9 + self.step_count * 5 // 60 # assume plot start at 9am
minute = self.step_count * 5 % 60
plt.title("Iteration: Day {}, ".format(self.day_count) + "%d:%02d" %
(hour, minute),
fontsize=30)
def __update_day(self):
'''
update incubation time, reset viral_load, remove symptomatic agents, aerosol transmission etc for end of day
'''
for a in self.schedule.agents[:]:
# update human agent disease stats
if issubclass(type(a), human_agent.Human):
if a.symptoms:
# remove agent if symptom onset
if isinstance(a, human_agent.Teacher) and len(
self.idle_teachers) > 0:
# assign a new teacher to position
new_teacher = self.idle_teachers.pop()
new_teacher.shape = a.shape
new_teacher.room = a.room
new_teacher.classroom = a.classroom
self.schedule.remove(a)
self.grid.remove_agent(a)
# UPDATE 10/16: infectious made obsolete, end of day update rework
elif a.health_status == "exposed":
# UPDATE 2/28: merge testing implementation
# test student and teacher accordingly
# Q: why testing is here under exposed case?:
# testing only matters if infect the result is a hit
# therefore the agnent gets removed only if two conditions are met
# 1.testing is arranged; 2. testing result the agent is indeed exposed
if isinstance(a, human_agent.Teacher) and (
self.day_count % self.teacher_testing_freq == 0):
# if hit teacher, try to assign a new teacher to position
if len(self.idle_teachers) > 0:
new_teacher = self.idle_teachers.pop()
new_teacher.shape = a.shape
new_teacher.room = a.room
new_teacher.classroom = a.classroom
#remove teacher if testing conditions are met
self.schedule.remove(a)
self.grid.remove_agent(a)
elif isinstance(a, human_agent.Student) and (
self.day_count % self.student_testing_freq == 0):
#remove student if testing conditions are met
self.schedule.remove(a)
self.grid.remove_agent(a)
# UPDATE 10/17: update infective delay if agent is not infective by end of day
a.infective = True
a.symptom_countdown -= 1
# calculate when symptoms begin to show using 0-15 density
if a.symptom_countdown <= 0:
if a.symptom_countdown == 0:
self.infected_count += 1
# update model stat for total infected
# negative countdown means this agent is asymptomatic
if not a.asymptomatic:
# this is a really small chance, however possible
# set symtoms to true
# next day this agent will be removed from the model
a.symptoms = True
# update room agent aerosal stats
elif issubclass(type(a), room_agent.Classroom):
room = a
mean_aerosol_transmissions = sum(
room.aerosol_transmission_rate)
if np.isnan(mean_aerosol_transmissions):
mean_aerosol_transmissions = 0
occupants = [
a for a in list(self.grid.get_intersecting_agents(room))
if issubclass(type(a), human_agent.Human)
]
healthy_occupants = [
a for a in occupants if a.health_status == 'healthy'
]
# failsafe for rare case where this can exceed one
mean_aerosol_transmissions = min(mean_aerosol_transmissions, 1)
# treating aerosal transmissions as a probability for each healthy occupant in this room to get sick
for healthy_occupant in healthy_occupants:
if np.random.choice([True, False],
p=[
mean_aerosol_transmissions,
1 - mean_aerosol_transmissions
]):
if not healthy_occupant.vaccinated:
healthy_occupant.health_status = 'exposed'
def step(self):
'''
simulate a day with school day schedule
'''
if not self.schedule.steps:
self.add_N_patient(self.init_patient)
for i, row in self.schoolday_schedule.iterrows():
self.activity = row
self.datacollector.collect(self)
self.schedule.step()
self.grid._recreate_rtree()
self.step_count += 1
self.__update_day()
self.grid._recreate_rtree()
self.day_count += 1
self.step_count = 0
class InfectedModel(Model):
"""Model class for a simplistic infection model."""
# Geographical parameters for desired map
MAP_COORDS = COORDS[CITY]
unique_id = "NAME"
def __init__(self, monitored_statistic, show_schools, show_restaurants):
"""
Create a new InfectedModel
"""
self.schedule = BaseScheduler(self)
self.grid = GeoSpace()
self.steps = 0
self.counts = None
self.reset_counts()
self.monitored_statistic = 'infected-per-home-series'
self.show_schools = show_schools
self.show_restaurants = show_restaurants
self.maximum = {self.monitored_statistic: 0}
self.minimum = {self.monitored_statistic: sys.maxsize}
with open(LOG_FILE) as json_file:
self.simulation_data = json.load(json_file)
#for key in self.simulation_data[self.monitored_statistic]:
# for v in self.simulation_data[self.monitored_statistic][key]:
# if v > self.maximum[self.monitored_statistic]: self.maximum[self.monitored_statistic] = v
# if v < self.minimum[self.monitored_statistic]: self.minimum[self.monitored_statistic] = v
self.minimum['infected-per-home-series'] = 0
self.maximum['infected-per-home-series'] = 500
self.running = True
self.datacollector = DataCollector({
"infected": get_infected_count,
"susceptible": get_susceptible_count,
"recovered": get_recovered_count,
"dead": get_dead_count,
})
# Neighboorhoods
AC = AgentCreator(NeighbourhoodAgent, {"model": self})
neighbourhood_agents = AC.from_file(geojson_neighborhoods,
unique_id=self.unique_id)
for agent in neighbourhood_agents:
for neighborhood in NEIGHBORHOODS['features']:
if agent.unique_id == neighborhood['properties']['NAME']:
agent2feature[agent.unique_id] = neighborhood
break
self.grid.add_agents(neighbourhood_agents)
# Schools
AC = AgentCreator(SchoolAgent, {"model": self})
school_agents = AC.from_file(geojson_schools, unique_id=self.unique_id)
for agent in school_agents:
for school in SCHOOLS['features']:
if agent.unique_id == school['properties']['NAME']:
agent2feature[agent.unique_id] = school
break
self.grid.add_agents(school_agents)
# Restaurants
AC = AgentCreator(RestaurantAgent, {"model": self})
restaurant_agents = AC.from_file(geojson_restaurants,
unique_id=self.unique_id)
for agent in restaurant_agents:
for restaurant in RESTAURANTS['features']:
if agent.unique_id == restaurant['properties']['NAME']:
agent2feature[agent.unique_id] = restaurant
break
self.grid.add_agents(restaurant_agents)
for agent in neighbourhood_agents + school_agents:
self.schedule.add(agent)
self.datacollector.collect(self)
def reset_counts(self):
self.counts = {
"susceptible": 0,
"infected": 0,
"recovered": 0,
"dead": 0,
"safe": 0,
"hotspot": 0,
}
def step(self):
#if self.steps >= len(self.simulation_data['infected-series']) - 2:
if self.steps > 50:
self.running = False
return
self.steps += 1
self.reset_counts()
self.schedule.step()
self.grid._recreate_rtree(
) # Recalculate spatial tree, because agents are moving
self.datacollector.collect(self)
return True
class InfectedModel(Model):
"""Model class for a simplistic infection model."""
# Geographical parameters for desired map
MAP_COORDS = [43.741667, -79.373333] # Toronto
geojson_regions = "TorontoNeighbourhoods.geojson"
unique_id = "HOODNUM"
def __init__(self,
pop_size,
init_infected,
exposure_distance,
infection_risk=0.2):
"""
Create a new InfectedModel
:param pop_size: Size of population
:param init_infected: Probability of a person agent to start as infected
:param exposure_distance: Proximity distance between agents to be exposed to each other
:param infection_risk: Probability of agent to become infected, if it has been exposed to another infected
"""
self.schedule = BaseScheduler(self)
self.grid = GeoSpace()
self.steps = 0
self.counts = None
self.reset_counts()
# SIR model parameters
self.pop_size = pop_size
self.counts["susceptible"] = pop_size
self.exposure_distance = exposure_distance
self.infection_risk = infection_risk
self.running = True
self.datacollector = DataCollector({
"infected": get_infected_count,
"susceptible": get_susceptible_count,
"recovered": get_recovered_count,
"dead": get_dead_count,
})
# Set up the Neighbourhood patches for every region in file (add to schedule later)
AC = AgentCreator(NeighbourhoodAgent, {"model": self})
neighbourhood_agents = AC.from_file(self.geojson_regions,
unique_id=self.unique_id)
self.grid.add_agents(neighbourhood_agents)
# Generate PersonAgent population
ac_population = AgentCreator(PersonAgent, {
"model": self,
"init_infected": init_infected
})
# Generate random location, add agent to grid and scheduler
for i in range(pop_size):
this_neighbourhood = self.random.randint(
0,
len(neighbourhood_agents) - 1) # Region where agent starts
center_x, center_y = neighbourhood_agents[
this_neighbourhood].shape.centroid.coords.xy
this_bounds = neighbourhood_agents[this_neighbourhood].shape.bounds
spread_x = int(
this_bounds[2] -
this_bounds[0]) # Heuristic for agent spread in region
spread_y = int(this_bounds[3] - this_bounds[1])
this_x = center_x[0] + self.random.randint(0,
spread_x) - spread_x / 2
this_y = center_y[0] + self.random.randint(0,
spread_y) - spread_y / 2
this_person = ac_population.create_agent(Point(this_x, this_y),
"P" + str(i))
self.grid.add_agents(this_person)
self.schedule.add(this_person)
# Add the neighbourhood agents to schedule AFTER person agents,
# to allow them to update their color by using BaseScheduler
for agent in neighbourhood_agents:
self.schedule.add(agent)
self.datacollector.collect(self)
def reset_counts(self):
self.counts = {
"susceptible": 0,
"infected": 0,
"recovered": 0,
"dead": 0,
"safe": 0,
"hotspot": 0,
}
def step(self):
"""Run one step of the model."""
self.steps += 1
self.reset_counts()
self.schedule.step()
self.grid._recreate_rtree(
) # Recalculate spatial tree, because agents are moving
self.datacollector.collect(self)
# Run until no one is infected
if self.counts["infected"] == 0:
self.running = False
