def tstep(self):
'''
takes a time step in the simulation
'''
start = time.time() # start clock
#======================================================================================#
#check destinations if active
#define motion vectors if destinations active and not everybody is at destination
active_dests = len(self.population[
self.population[:, 11] != 0]) # look op this only once
if active_dests > 0 and len(
self.population[self.population[:, 12] == 0]) > 0:
self.population = set_destination(self.population,
self.destinations)
self.population = check_at_destination(
self.population,
self.destinations,
wander_factor=self.Config.wander_factor_dest)
if active_dests > 0 and len(
self.population[self.population[:, 12] == 1]) > 0:
#keep them at destination
self.population = keep_at_destination(self.population,
self.Config.isolation_bounds)
#======================================================================================#
#gravity wells
if self.Config.gravity_strength > 0:
[self.population, self.last_step_change] = update_gravity_forces(
self.population, self.time, self.last_step_change,
self.Config.wander_step_size, self.Config.gravity_strength,
self.Config.wander_step_duration)
#======================================================================================#
#activate social distancing above a certain infection threshold
if not self.above_act_thresh and self.Config.social_distance_threshold_on > 0:
# If not previously above infection threshold activate when threshold reached
if self.Config.thresh_type == 'hospitalized':
self.above_act_thresh = sum(
self.population[:, 11] ==
1) >= self.Config.social_distance_threshold_on
elif self.Config.thresh_type == 'infected':
self.above_act_thresh = sum(
self.population[:, 6] ==
1) >= self.Config.social_distance_threshold_on
elif self.Config.social_distance_threshold_on == 0:
self.above_act_thresh = True
#deactivate social distancing after infection drops below threshold after using social distancing
if self.above_act_thresh and not self.above_deact_thresh and self.Config.social_distance_threshold_off > 0:
# If previously went above infection threshold deactivate when threshold reached
self.above_deact_thresh = sum(self.population[:,6][self.population[:,11] == 0] == 1) <= \
self.Config.social_distance_threshold_off
# activate social distancing at the onset of infection
if not self.Config.SD_act_onset:
act_social_distancing = self.above_act_thresh and not self.above_deact_thresh and sum(
self.population[:, 6] == 1) > 0
# activate social distancing from start of simulation
elif self.Config.SD_act_onset:
act_social_distancing = self.above_act_thresh and not self.above_deact_thresh
#activate social distancing only for compliant individuals
if self.Config.social_distance_factor > 0 and act_social_distancing:
self.population[(self.population[:,17] == 0) &\
(self.population[:,11] == 0)] = update_repulsive_forces(self.population[(self.population[:,17] == 0) &\
(self.population[:,11] == 0)], self.Config.social_distance_factor)
#======================================================================================#
#out of bounds
#define bounds arrays, excluding those who are marked as having a custom destination
if len(self.population[:, 11] == 0) > 0:
buffer = 0.0
_xbounds = np.array([[
self.Config.xbounds[0] + buffer,
self.Config.xbounds[1] - buffer
]] * len(self.population[self.population[:, 11] == 0]))
_ybounds = np.array([[
self.Config.ybounds[0] + buffer,
self.Config.ybounds[1] - buffer
]] * len(self.population[self.population[:, 11] == 0]))
self.population[self.population[:, 11] == 0] = update_wall_forces(
self.population[self.population[:, 11] == 0], _xbounds,
_ybounds)
#======================================================================================#
#update velocities
self.population[(self.population[:,11] == 0) |\
(self.population[:,12] == 1)] = update_velocities(self.population[(self.population[:,11] == 0) |\
(self.population[:,12] == 1)],
self.Config.max_speed,self.Config.dt)
#for dead ones: set velocity and social distancing to 0 for dead ones
self.population[:, 3:5][self.population[:, 6] == 3] = 0
self.population[:, 17][self.population[:, 6] == 3] = 1
#update positions
self.population = update_positions(self.population, self.Config.dt)
#======================================================================================#
#find new infections
if not self.above_test_thresh and self.Config.testing_threshold_on > 0:
# If not previously above infection threshold activate when threshold reached
self.above_test_thresh = sum(
self.population[:, 6] == 1) >= self.Config.testing_threshold_on
# self.above_test_thresh = sum(self.population[:,6] == 1) >= self.Config.social_distance_threshold_on
elif self.Config.testing_threshold_on == 0:
self.above_test_thresh = True
act_testing = self.above_test_thresh and sum(
self.population[:, 6] == 1) > 0
self.population, self.destinations = infect(
self.population,
self.Config,
self.frame,
send_to_location=self.Config.self_isolate,
location_bounds=self.Config.isolation_bounds,
destinations=self.destinations,
location_no=1,
location_odds=self.Config.self_isolate_proportion,
test_flag=act_testing)
#recover and die
self.population = recover_or_die(self.population, self.frame,
self.Config)
#======================================================================================#
#send cured back to population if self isolation active
#perhaps put in recover or die class
#send cured back to population
self.population[:, 11][self.population[:, 6] == 2] = 0
#======================================================================================#
#update population statistics
self.pop_tracker.update_counts(self.population, self.frame)
#======================================================================================#
#visualise
if self.Config.visualise and (
self.frame % self.Config.visualise_every_n_frame) == 0:
draw_tstep(self.Config, self.population, self.pop_tracker,
self.frame, self.fig, self.spec, self.ax1, self.ax2,
self.tight_bbox)
#report stuff to console
if (self.Config.verbose) and ((self.frame % self.Config.report_freq)
== 0):
end = time.time()
time_elapsed = end - start # elapsed time
sys.stdout.write('\r')
sys.stdout.write(
'%i: S: %i, I: %i, R: %i, in treatment: %i, F: %i, of total: %i, D: %.5f, GC: %.5f, time: %.5f'
% (self.frame, self.pop_tracker.susceptible[-1],
self.pop_tracker.infectious[-1],
self.pop_tracker.recovered[-1],
len(self.population[self.population[:, 10] == 1]),
self.pop_tracker.fatalities[-1], self.Config.pop_size,
self.pop_tracker.distance_travelled[-1],
self.pop_tracker.mean_perentage_covered[-1], time_elapsed))
#save popdata if required
if self.Config.save_pop and (self.frame %
self.Config.save_pop_freq) == 0:
save_population(self.population, self.frame,
self.Config.save_pop_folder)
#run callback
self.callback()
#======================================================================================#
#update frame
self.frame += 1
self.time += self.Config.dt
def tstep(self):
"""
takes a time step in the simulation
"""
# check destinations if active
# define motion vectors if destinations active and not everybody is at destination
active_dests = len(
self.population[self.population[:, 11] != 0]
) # look op this only once
if active_dests > 0 and len(self.population[self.population[:, 12] == 0]) > 0:
self.population = set_destination(self.population, self.destinations)
self.population = check_at_destination(
self.population,
self.destinations,
wander_factor=self.Config.wander_factor_dest,
speed=self.Config.speed,
)
if active_dests > 0 and len(self.population[self.population[:, 12] == 1]) > 0:
# keep them at destination
self.population = keep_at_destination(
self.population, self.destinations, self.Config.wander_factor
)
# out of bounds
# define bounds arrays, excluding those who are marked as having a custom destination
if len(self.population[:, 11] == 0) > 0:
_xbounds = np.array(
[[self.Config.xbounds[0] + 0.02, self.Config.xbounds[1] - 0.02]]
* len(self.population[self.population[:, 11] == 0])
)
_ybounds = np.array(
[[self.Config.ybounds[0] + 0.02, self.Config.ybounds[1] - 0.02]]
* len(self.population[self.population[:, 11] == 0])
)
self.population[self.population[:, 11] == 0] = out_of_bounds(
self.population[self.population[:, 11] == 0], _xbounds, _ybounds
)
# set randoms
if self.Config.lockdown:
if len(self.pop_tracker.infectious) == 0:
mx = 0
else:
mx = np.max(self.pop_tracker.infectious)
if len(self.population[self.population[:, 6] == 1]) >= len(
self.population
) * self.Config.lockdown_percentage or mx >= (
len(self.population) * self.Config.lockdown_percentage
):
# reduce speed of all members of society
self.population[:, 5] = np.clip(
self.population[:, 5], a_min=None, a_max=0.001
)
# set speeds of complying people to 0
self.population[:, 5][self.Config.lockdown_vector == 0] = 0
else:
# update randoms
self.population = update_randoms(
self.population, self.Config.pop_size, self.Config.speed
)
else:
# update randoms
self.population = update_randoms(
self.population, self.Config.pop_size, self.Config.speed
)
# for dead ones: set speed and heading to 0
self.population[:, 3:5][self.population[:, 6] == 3] = 0
# update positions
self.population = update_positions(self.population)
# find new infections
self.population, self.destinations = infect(
self.population,
self.Config,
self.frame,
send_to_location=self.Config.self_isolate,
location_bounds=self.Config.isolation_bounds,
destinations=self.destinations,
location_no=1,
location_odds=self.Config.self_isolate_proportion,
)
# recover and die
self.population = recover_or_die(self.population, self.frame, self.Config)
# send cured back to population if self isolation active
# perhaps put in recover or die class
# send cured back to population
self.population[:, 11][self.population[:, 6] == 2] = 0
# update population statistics
self.pop_tracker.update_counts(self.population)
# visualise
if self.Config.visualise:
draw_tstep(
self.Config,
self.population,
self.pop_tracker,
self.frame,
self.fig,
self.spec,
self.ax1,
self.ax2,
)
# report stuff to console
sys.stdout.write("\r")
sys.stdout.write(
"%i: healthy: %i, infected: %i, immune: %i, in treatment: %i, \
dead: %i, of total: %i"
% (
self.frame,
self.pop_tracker.susceptible[-1],
self.pop_tracker.infectious[-1],
self.pop_tracker.recovered[-1],
len(self.population[self.population[:, 10] == 1]),
self.pop_tracker.fatalities[-1],
self.Config.pop_size,
)
)
# save popdata if required
if self.Config.save_pop and (self.frame % self.Config.save_pop_freq) == 0:
save_population(self.population, self.frame, self.Config.save_pop_folder)
# run callback
self.callback()
# update frame
self.frame += 1
def update(frame, population, destinations, configuration):
# #add one infection to jumpstart
if frame == 10:
population[0,6] = 1
#update out of bounds
#define bounds arrays
_xbounds = np.array([[configuration.xbounds[0] + 0.02, configuration.xbounds[1] - 0.02]] * len(population))
_ybounds = np.array([[configuration.ybounds[0] + 0.02, configuration.ybounds[1] - 0.02]] * len(population))
population = out_of_bounds(population, _xbounds, _ybounds)
#update randoms
population = update_randoms(population, pop_size)
#for dead ones: set speed and heading to 0
population[:,3:5][population[:,6] == 3] = 0
#update positions
population = update_positions(population)
#find new infections
population = infect(population, configuration, frame)
infected_plot.append(len(population[population[:,6] == 1]))
#recover and die
population = recover_or_die(population, frame, configuration)
fatalities_plot.append(len(population[population[:,6] == 3]))
if configuration.visualise:
#construct plot and visualise
spec = fig.add_gridspec(ncols=1, nrows=2, height_ratios=[5,2])
ax1.clear()
ax2.clear()
ax1.set_xlim(configuration.xbounds[0], configuration.xbounds[1])
ax1.set_ylim(configuration.ybounds[0], configuration.ybounds[1])
healthy = population[population[:,6] == 0][:,1:3]
ax1.scatter(healthy[:,0], healthy[:,1], color='gray', s = 2, label='healthy')
infected = population[population[:,6] == 1][:,1:3]
ax1.scatter(infected[:,0], infected[:,1], color='red', s = 2, label='infected')
immune = population[population[:,6] == 2][:,1:3]
ax1.scatter(immune[:,0], immune[:,1], color='green', s = 2, label='immune')
fatalities = population[population[:,6] == 3][:,1:3]
ax1.scatter(fatalities[:,0], fatalities[:,1], color='black', s = 2, label='fatalities')
#add text descriptors
ax1.text(configuration.xbounds[0],
configuration.ybounds[1] + ((configuration.ybounds[1] - configuration.ybounds[0]) / 100),
'timestep: %i, total: %i, healthy: %i infected: %i immune: %i fatalities: %i' %(frame,
len(population),
len(healthy),
len(infected),
len(immune),
len(fatalities)),
fontsize=6)
ax2.set_title('number of infected')
ax2.text(0, configuration.pop_size * 0.05,
'https://github.com/paulvangentcom/python-corona-simulation',
fontsize=6, alpha=0.5)
ax2.set_xlim(0, simulation_steps)
ax2.set_ylim(0, configuration.pop_size + 100)
ax2.plot(infected_plot, color='gray')
ax2.plot(fatalities_plot, color='black', label='fatalities')
if treatment_dependent_risk:
#ax2.plot([healthcare_capacity for x in range(simulation_steps)], color='red',
# label='healthcare capacity')
infected_arr = np.asarray(infected_plot)
indices = np.argwhere(infected_arr >= healthcare_capacity)
ax2.plot(indices, infected_arr[infected_arr >= healthcare_capacity],
color='red')
#ax2.legend(loc = 1, fontsize = 6)
#plt.savefig('render/%i.png' %frame)
return population
def update(
frame,
population,
destinations,
pop_size,
infection_range=0.01,
infection_chance=0.03,
recovery_duration=(200, 500),
mortality_chance=0.02,
xbounds=[0.02, 0.98],
ybounds=[0.02, 0.98],
wander_range_x=0.05,
wander_range_y=0.05,
risk_age=55,
critical_age=75,
critical_mortality_chance=0.1,
risk_increase="quadratic",
no_treatment_factor=3,
treatment_factor=0.5,
healthcare_capacity=250,
age_dependent_risk=True,
treatment_dependent_risk=True,
visualise=True,
verbose=True,
):
# add one infection to jumpstart
if frame == 100:
# make C
# first leg
destinations[:, 0][0:100] = 0.05
destinations[:, 1][0:100] = 0.7
population[:, 13][0:100] = 0.01
population[:, 14][0:100] = 0.05
# Top
destinations[:, 0][100:200] = 0.1
destinations[:, 1][100:200] = 0.75
population[:, 13][100:200] = 0.05
population[:, 14][100:200] = 0.01
# Bottom
destinations[:, 0][200:300] = 0.1
destinations[:, 1][200:300] = 0.65
population[:, 13][200:300] = 0.05
population[:, 14][200:300] = 0.01
# make O
# first leg
destinations[:, 0][300:400] = 0.2
destinations[:, 1][300:400] = 0.7
population[:, 13][300:400] = 0.01
population[:, 14][300:400] = 0.05
# Top
destinations[:, 0][400:500] = 0.25
destinations[:, 1][400:500] = 0.75
population[:, 13][400:500] = 0.05
population[:, 14][400:500] = 0.01
# Bottom
destinations[:, 0][500:600] = 0.25
destinations[:, 1][500:600] = 0.65
population[:, 13][500:600] = 0.05
population[:, 14][500:600] = 0.01
# second leg
destinations[:, 0][600:700] = 0.3
destinations[:, 1][600:700] = 0.7
population[:, 13][600:700] = 0.01
population[:, 14][600:700] = 0.05
# make V
# First leg
destinations[:, 0][700:800] = 0.35
destinations[:, 1][700:800] = 0.7
population[:, 13][700:800] = 0.01
population[:, 14][700:800] = 0.05
# Bottom
destinations[:, 0][800:900] = 0.4
destinations[:, 1][800:900] = 0.65
population[:, 13][800:900] = 0.05
population[:, 14][800:900] = 0.01
# second leg
destinations[:, 0][900:1000] = 0.45
destinations[:, 1][900:1000] = 0.7
population[:, 13][900:1000] = 0.01
population[:, 14][900:1000] = 0.05
# Make I
# leg
destinations[:, 0][1000:1100] = 0.5
destinations[:, 1][1000:1100] = 0.7
population[:, 13][1000:1100] = 0.01
population[:, 14][1000:1100] = 0.05
# I dot
destinations[:, 0][1100:1200] = 0.5
destinations[:, 1][1100:1200] = 0.8
population[:, 13][1100:1200] = 0.01
population[:, 14][1100:1200] = 0.01
# make D
# first leg
destinations[:, 0][1200:1300] = 0.55
destinations[:, 1][1200:1300] = 0.67
population[:, 13][1200:1300] = 0.01
population[:, 14][1200:1300] = 0.03
# Top
destinations[:, 0][1300:1400] = 0.6
destinations[:, 1][1300:1400] = 0.75
population[:, 13][1300:1400] = 0.05
population[:, 14][1300:1400] = 0.01
# Bottom
destinations[:, 0][1400:1500] = 0.6
destinations[:, 1][1400:1500] = 0.65
population[:, 13][1400:1500] = 0.05
population[:, 14][1400:1500] = 0.01
# second leg
destinations[:, 0][1500:1600] = 0.65
destinations[:, 1][1500:1600] = 0.7
population[:, 13][1500:1600] = 0.01
population[:, 14][1500:1600] = 0.05
# dash
destinations[:, 0][1600:1700] = 0.725
destinations[:, 1][1600:1700] = 0.7
population[:, 13][1600:1700] = 0.03
population[:, 14][1600:1700] = 0.01
# Make 1
destinations[:, 0][1700:1800] = 0.8
destinations[:, 1][1700:1800] = 0.7
population[:, 13][1700:1800] = 0.01
population[:, 14][1700:1800] = 0.05
# Make 9
# right leg
destinations[:, 0][1800:1900] = 0.91
destinations[:, 1][1800:1900] = 0.675
population[:, 13][1800:1900] = 0.01
population[:, 14][1800:1900] = 0.08
# roof
destinations[:, 0][1900:2000] = 0.88
destinations[:, 1][1900:2000] = 0.75
population[:, 13][1900:2000] = 0.035
population[:, 14][1900:2000] = 0.01
# middle
destinations[:, 0][2000:2100] = 0.88
destinations[:, 1][2000:2100] = 0.7
population[:, 13][2000:2100] = 0.035
population[:, 14][2000:2100] = 0.01
# left vertical leg
destinations[:, 0][2100:2200] = 0.86
destinations[:, 1][2100:2200] = 0.72
population[:, 13][2100:2200] = 0.01
population[:, 14][2100:2200] = 0.01
###################
##### ROW TWO #####
###################
# S
# first leg
destinations[:, 0][2200:2300] = 0.115
destinations[:, 1][2200:2300] = 0.5
population[:, 13][2200:2300] = 0.01
population[:, 14][2200:2300] = 0.03
# Top
destinations[:, 0][2300:2400] = 0.15
destinations[:, 1][2300:2400] = 0.55
population[:, 13][2300:2400] = 0.05
population[:, 14][2300:2400] = 0.01
# second leg
destinations[:, 0][2400:2500] = 0.2
destinations[:, 1][2400:2500] = 0.45
population[:, 13][2400:2500] = 0.01
population[:, 14][2400:2500] = 0.03
# middle
destinations[:, 0][2500:2600] = 0.15
destinations[:, 1][2500:2600] = 0.48
population[:, 13][2500:2600] = 0.05
population[:, 14][2500:2600] = 0.01
# bottom
destinations[:, 0][2600:2700] = 0.15
destinations[:, 1][2600:2700] = 0.41
population[:, 13][2600:2700] = 0.05
population[:, 14][2600:2700] = 0.01
# Make I
# leg
destinations[:, 0][2700:2800] = 0.25
destinations[:, 1][2700:2800] = 0.45
population[:, 13][2700:2800] = 0.01
population[:, 14][2700:2800] = 0.05
# I dot
destinations[:, 0][2800:2900] = 0.25
destinations[:, 1][2800:2900] = 0.55
population[:, 13][2800:2900] = 0.01
population[:, 14][2800:2900] = 0.01
# M
# Top
destinations[:, 0][2900:3000] = 0.37
destinations[:, 1][2900:3000] = 0.5
population[:, 13][2900:3000] = 0.07
population[:, 14][2900:3000] = 0.01
# Left leg
destinations[:, 0][3000:3100] = 0.31
destinations[:, 1][3000:3100] = 0.45
population[:, 13][3000:3100] = 0.01
population[:, 14][3000:3100] = 0.05
# Middle leg
destinations[:, 0][3100:3200] = 0.37
destinations[:, 1][3100:3200] = 0.45
population[:, 13][3100:3200] = 0.01
population[:, 14][3100:3200] = 0.05
# Right leg
destinations[:, 0][3200:3300] = 0.43
destinations[:, 1][3200:3300] = 0.45
population[:, 13][3200:3300] = 0.01
population[:, 14][3200:3300] = 0.05
# set all destinations active
population[:, 11] = 1
elif frame == 400:
population[:, 11] = 0
population[:, 12] = 0
population = update_randoms(population, pop_size, 1, 1)
# define motion vectors if destinations active and not everybody is at destination
active_dests = len(
population[population[:, 11] != 0]) # look op this only once
if active_dests > 0 and len(population[population[:, 12] == 0]) > 0:
population = set_destination(population, destinations)
population = check_at_destination(population, destinations)
if active_dests > 0 and len(population[population[:, 12] == 1]) > 0:
# keep them at destination
population = keep_at_destination(population,
destinations,
wander_factor=1)
# update out of bounds
# define bounds arrays
_xbounds = np.array([[xbounds[0] + 0.02, xbounds[1] - 0.02]] *
len(population))
_ybounds = np.array([[ybounds[0] + 0.02, ybounds[1] - 0.02]] *
len(population))
population = out_of_bounds(population, _xbounds, _ybounds)
# update randoms
population = update_randoms(population, pop_size)
# for dead ones: set speed and heading to 0
population[:, 3:5][population[:, 6] == 3] = 0
# update positions
population = update_positions(population)
# find new infections
population = infect(
population,
pop_size,
infection_range,
infection_chance,
frame,
healthcare_capacity,
verbose,
)
infected_plot.append(len(population[population[:, 6] == 1]))
# recover and die
population = recover_or_die(
population,
frame,
recovery_duration,
mortality_chance,
risk_age,
critical_age,
critical_mortality_chance,
risk_increase,
no_treatment_factor,
age_dependent_risk,
treatment_dependent_risk,
treatment_factor,
verbose,
)
fatalities_plot.append(len(population[population[:, 6] == 3]))
if visualise:
# construct plot and visualise
spec = fig.add_gridspec(ncols=1, nrows=2, height_ratios=[5, 2])
ax1.clear()
ax2.clear()
ax1.set_xlim(xbounds[0], xbounds[1])
ax1.set_ylim(ybounds[0], ybounds[1])
healthy = population[population[:, 6] == 0][:, 1:3]
ax1.scatter(healthy[:, 0],
healthy[:, 1],
color="gray",
s=2,
label="healthy")
infected = population[population[:, 6] == 1][:, 1:3]
ax1.scatter(infected[:, 0],
infected[:, 1],
color="red",
s=2,
label="infected")
immune = population[population[:, 6] == 2][:, 1:3]
ax1.scatter(immune[:, 0],
immune[:, 1],
color="green",
s=2,
label="immune")
fatalities = population[population[:, 6] == 3][:, 1:3]
ax1.scatter(fatalities[:, 0],
fatalities[:, 1],
color="black",
s=2,
label="fatalities")
# add text descriptors
ax1.text(
xbounds[0],
ybounds[1] + ((ybounds[1] - ybounds[0]) / 100),
"timestep: %i, total: %i, healthy: %i infected: %i immune: %i fatalities: %i"
% (
frame,
len(population),
len(healthy),
len(infected),
len(immune),
len(fatalities),
),
fontsize=6,
)
ax2.set_title("number of infected")
ax2.text(
0,
pop_size * 0.05,
"https://github.com/paulvangentcom/python-corona-simulation",
fontsize=6,
alpha=0.5,
)
ax2.set_xlim(0, simulation_steps)
ax2.set_ylim(0, pop_size + 100)
ax2.plot(infected_plot, color="gray")
ax2.plot(fatalities_plot, color="black", label="fatalities")
if treatment_dependent_risk:
# ax2.plot([healthcare_capacity for x in range(simulation_steps)], color='red',
# label='healthcare capacity')
infected_arr = np.asarray(infected_plot)
indices = np.argwhere(infected_arr >= healthcare_capacity)
ax2.plot(indices,
infected_arr[infected_arr >= healthcare_capacity],
color="red")
# ax2.legend(loc = 1, fontsize = 6)
# plt.savefig('render/%i.png' %frame)
return population
def update(frame,
population,
destinations,
pop_size,
infection_range=0.01,
infection_chance=0.03,
speed=0.01,
recovery_duration=(200, 500),
mortality_chance=0.02,
xbounds=[0.02, 0.98],
ybounds=[0.02, 0.98],
x_plot=[0, 1],
y_plot=[0, 1],
wander_range=0.05,
risk_age=55,
critical_age=75,
critical_mortality_chance=0.1,
risk_increase='quadratic',
no_treatment_factor=3,
treatment_factor=0.5,
healthcare_capacity=250,
age_dependent_risk=False,
treatment_dependent_risk=False,
visualise=False,
verbose=False,
self_isolate=True,
self_isolate_proportion=0.6,
isolation_bounds=[0, 0, 0.1, 0.1],
traveling_infects=False,
lockdown=False,
lockdown_percentage=0.1,
lockdown_vector=[],
plot_style='default'):
#add one infection to jumpstart
if frame == 50:
population[0][6] = 1
population[0][8] = 75
population[0][10] = 1
#define motion vectors if destinations active and not everybody is at destination
active_dests = len(
population[population[:, 11] != 0]) # look op this only once
if active_dests > 0 and len(population[population[:, 12] == 0]) > 0:
population = set_destination(population, destinations)
population = check_at_destination(population,
destinations,
wander_factor=1.5)
if active_dests > 0 and len(population[population[:, 12] == 1]) > 0:
#keep them at destination
population = keep_at_destination(population,
destinations,
wander_factor=1)
#update out of bounds
#define bounds arrays, excluding those who are marked as having a custom destination
if len(population[:, 11] == 0) > 0:
_xbounds = np.array([[xbounds[0] + 0.02, xbounds[1] - 0.02]] *
len(population[population[:, 11] == 0]))
_ybounds = np.array([[ybounds[0] + 0.02, ybounds[1] - 0.02]] *
len(population[population[:, 11] == 0]))
population[population[:, 11] == 0] = out_of_bounds(
population[population[:, 11] == 0], _xbounds, _ybounds)
if lockdown:
if len(infected_plot) == 0:
mx = 0
else:
mx = np.max(infected_plot)
if len(population[population[:,6] == 1]) >= len(population) * lockdown_percentage or\
mx >= (len(population) * lockdown_percentage):
#reduce speed of all members of society
population[:, 5] = np.clip(population[:, 5],
a_min=None,
a_max=0.001)
#set speeds of complying people to 0
population[:, 5][lockdown_vector == 0] = 0
else:
#update randoms
population = update_randoms(population, pop_size, speed=speed)
else:
#update randoms
population = update_randoms(population, pop_size, speed=speed)
#for dead ones: set speed and heading to 0
population[:, 3:5][population[:, 6] == 3] = 0
#update positions
population = update_positions(population)
#find new infections
population, destinations = infect(population,
pop_size,
infection_range,
infection_chance,
frame,
healthcare_capacity,
verbose,
send_to_location=self_isolate,
location_bounds=isolation_bounds,
destinations=destinations,
location_no=1,
location_odds=self_isolate_proportion,
traveling_infects=traveling_infects)
infected_plot.append(len(population[population[:, 6] == 1]))
#recover and die
population = recover_or_die(population, frame, recovery_duration,
mortality_chance, risk_age, critical_age,
critical_mortality_chance, risk_increase,
no_treatment_factor, age_dependent_risk,
treatment_dependent_risk, treatment_factor,
verbose)
#send cured back to population
population[:, 11][population[:, 6] == 2] = 0
fatalities_plot.append(len(population[population[:, 6] == 3]))
if visualise:
#construct plot and visualise
spec = fig.add_gridspec(ncols=1, nrows=2, height_ratios=[5, 2])
ax1.clear()
ax2.clear()
ax1.set_xlim(x_plot[0], x_plot[1])
ax1.set_ylim(y_plot[0], y_plot[1])
if self_isolate and isolation_bounds != None:
build_hospital(isolation_bounds[0],
isolation_bounds[2],
isolation_bounds[1],
isolation_bounds[3],
ax1,
addcross=False)
#plot population segments
healthy = population[population[:, 6] == 0][:, 1:3]
ax1.scatter(healthy[:, 0],
healthy[:, 1],
color='gray',
s=2,
label='healthy')
infected = population[population[:, 6] == 1][:, 1:3]
ax1.scatter(infected[:, 0],
infected[:, 1],
color='red',
s=2,
label='infected')
immune = population[population[:, 6] == 2][:, 1:3]
ax1.scatter(immune[:, 0],
immune[:, 1],
color='green',
s=2,
label='immune')
fatalities = population[population[:, 6] == 3][:, 1:3]
ax1.scatter(fatalities[:, 0],
fatalities[:, 1],
color='black',
s=2,
label='dead')
#add text descriptors
ax1.text(
x_plot[0],
y_plot[1] + ((y_plot[1] - y_plot[0]) / 100),
'timestep: %i, total: %i, healthy: %i infected: %i immune: %i fatalities: %i'
% (frame, len(population), len(healthy), len(infected),
len(immune), len(fatalities)),
fontsize=6)
ax2.set_title('number of infected')
ax2.text(0,
pop_size * 0.05,
'https://github.com/paulvangentcom/python-corona-simulation',
fontsize=6,
alpha=0.5)
#ax2.set_xlim(0, simulation_steps)
ax2.set_ylim(0, pop_size + 200)
if treatment_dependent_risk:
infected_arr = np.asarray(infected_plot)
indices = np.argwhere(infected_arr >= healthcare_capacity)
ax2.plot([healthcare_capacity for x in range(len(infected_plot))],
color='red',
label='healthcare capacity')
ax2.plot(infected_plot, color='gray')
ax2.plot(fatalities_plot, color='black', label='fatalities')
if treatment_dependent_risk:
ax2.plot(indices,
infected_arr[infected_arr >= healthcare_capacity],
color='red')
ax2.legend(loc='best', fontsize=6)
#plt.savefig('render/%i.png' %frame)
return population
def tstep(self):
if self.frame == 0:
self.fig, self.spec, self.ax1, self.ax2 = build_fig(self.Config)
xbounds = np.array(
[[self.Config.xbounds[0] + 0.02, self.Config.xbounds[1] - 0.02]] *
self.Config.pop_size)
ybounds = np.array(
[[self.Config.ybounds[0] + 0.02, self.Config.ybounds[1] - 0.02]] *
self.Config.pop_size)
self.population = out_of_bounds(self.population, xbounds, ybounds)
if self.Config.is_lockdown and self.Config.lockdown == False:
if len(self.population[(self.population[:, 6] == 1)]
) >= self.Config.lockdown_percentage * self.Config.pop_size:
self.Config.lockdown = True
print("\nLockdown Started")
left_range = [
self.Config.xbounds[0] + 0.02, self.Config.xbounds[1] / 3 - 0.02
]
mid_range = [
self.Config.xbounds[1] / 3 + 0.02,
2 * self.Config.xbounds[1] / 3 - 0.02
]
right_range = [
2 * self.Config.xbounds[1] / 3 + 0.02,
self.Config.xbounds[1] - 0.02
]
bottom_range = [
self.Config.ybounds[0] + 0.02, self.Config.ybounds[1] / 2 - 0.02
]
top_range = [
self.Config.ybounds[1] / 2 + 0.02, self.Config.ybounds[1] - 0.02
]
left_condition = (self.population[:, 1] <= self.Config.xbounds[1] / 3)
mid_condition = (
self.population[:, 1] > self.Config.xbounds[1] / 3) & (
self.population[:, 1] <= 2 * self.Config.xbounds[1] / 3)
right_condition = (self.population[:, 1] >
2 * self.Config.xbounds[1] / 3)
bottom_condition = (self.population[:, 2] <=
self.Config.ybounds[1] / 2)
top_condition = (self.population[:, 2] > self.Config.ybounds[1] / 2)
if self.Config.lockdown:
x_left_bottom = np.array(
[left_range] *
len(self.population[left_condition & bottom_condition]))
y_left_bottom = np.array(
[bottom_range] *
len(self.population[left_condition & bottom_condition]))
self.population[left_condition & bottom_condition] = out_of_bounds(
self.population[left_condition & bottom_condition],
x_left_bottom, y_left_bottom)
x_left_top = np.array(
[left_range] *
len(self.population[left_condition & top_condition]))
y_left_top = np.array(
[top_range] *
len(self.population[left_condition & top_condition]))
self.population[left_condition & top_condition] = out_of_bounds(
self.population[left_condition & top_condition], x_left_top,
y_left_top)
x_mid_bottom = np.array(
[mid_range] *
len(self.population[mid_condition & bottom_condition]))
y_mid_bottom = np.array(
[bottom_range] *
len(self.population[mid_condition & bottom_condition]))
self.population[mid_condition & bottom_condition] = out_of_bounds(
self.population[mid_condition & bottom_condition],
x_mid_bottom, y_mid_bottom)
x_mid_top = np.array(
[mid_range] *
len(self.population[mid_condition & top_condition]))
y_mid_top = np.array(
[top_range] *
len(self.population[mid_condition & top_condition]))
self.population[mid_condition & top_condition] = out_of_bounds(
self.population[mid_condition & top_condition], x_mid_top,
y_mid_top)
x_right_bottom = np.array(
[right_range] *
len(self.population[right_condition & bottom_condition]))
y_right_bottom = np.array(
[bottom_range] *
len(self.population[right_condition & bottom_condition]))
self.population[right_condition
& bottom_condition] = out_of_bounds(
self.population[right_condition
& bottom_condition],
x_right_bottom, y_right_bottom)
x_right_top = np.array(
[right_range] *
len(self.population[right_condition & top_condition]))
y_right_top = np.array(
[top_range] *
len(self.population[right_condition & top_condition]))
self.population[right_condition & top_condition] = out_of_bounds(
self.population[right_condition & top_condition], x_right_top,
y_right_top)
self.population = update_randoms(self.population, self.Config.pop_size,
self.Config.speed)
self.population[:, 5][self.population[:, 6] == 3] = 0
self.population = update_positions(self.population)
self.population = infect(self.population, self.Config, self.frame)
self.population = recover_or_die(self.population, self.frame,
self.Config)
self.pop_tracker.update_counts(self.population)
draw_tstep(self.Config, self.population, self.pop_tracker, self.frame,
self.fig, self.spec, self.ax1, self.ax2)
self.peak_infections = max(
self.peak_infections,
len(self.population[self.population[:, 6] == 1]))
if self.frame == 50:
print('\ninfecting patient zero')
self.population[0][6] = 1
self.frame += 1
def update(frame,
population,
destinations,
pop_size,
infection_range=0.01,
infection_chance=0.03,
recovery_duration=(200, 500),
mortality_chance=0.02,
xbounds=[0.02, 0.98],
ybounds=[0.02, 0.98],
x_plot=[-0.1, 1],
y_plot=[-0.1, 1],
wander_range_x=0.05,
wander_range_y=0.05,
risk_age=55,
critical_age=75,
critical_mortality_chance=0.1,
risk_increase='quadratic',
no_treatment_factor=3,
treatment_factor=0.5,
healthcare_capacity=250,
age_dependent_risk=True,
treatment_dependent_risk=True,
visualise=True,
verbose=True,
healthcare_workers=50,
hospital_bounds=None,
healthcare_worker_risk=0):
#add one infection to jumpstart
if frame == 1:
population[healthcare_workers + 1][6] = 1
#define motion vectors if destinations active and not everybody is at destination
active_dests = len(
population[population[:, 11] != 0]) # look op this only once
if active_dests > 0 and len(population[population[:, 12] == 0]) > 0:
population = set_destination(population, destinations)
population = check_at_destination(population,
destinations,
wander_factor=1)
if active_dests > 0 and len(population[population[:, 12] == 1]) > 0:
#keep them at destination
population = keep_at_destination(population,
destinations,
wander_factor=1)
#update out of bounds
#define bounds arrays
if len(population[:, 11] == 0) > 0:
_xbounds = np.array([[xbounds[0] + 0.02, xbounds[1] - 0.02]] *
len(population[population[:, 11] == 0]))
_ybounds = np.array([[ybounds[0] + 0.02, ybounds[1] - 0.02]] *
len(population[population[:, 11] == 0]))
population[population[:, 11] == 0] = out_of_bounds(
population[population[:, 11] == 0], _xbounds, _ybounds)
#update randoms
population = update_randoms(population, pop_size)
#for dead ones: set speed and heading to 0
population[:, 3:5][population[:, 6] == 3] = 0
#update positions
population = update_positions(population)
#find new infections
population, destinations = infect(population,
pop_size,
infection_range,
infection_chance,
frame,
healthcare_capacity,
verbose,
send_to_location=True,
location_bounds=hospital_bounds,
destinations=destinations,
location_no=1)
#apply risk factor to healthcare worker pool
if healthcare_worker_risk != 0: #if risk is not zero, affect workers
workers = population[0:healthcare_workers]
workers = healthcare_infection_correction(workers,
healthcare_worker_risk)
population[0:healthcare_workers] = workers
infected_plot.append(len(population[population[:, 6] == 1]))
#recover and die
population = recover_or_die(population, frame, recovery_duration,
mortality_chance, risk_age, critical_age,
critical_mortality_chance, risk_increase,
no_treatment_factor, age_dependent_risk,
treatment_dependent_risk, treatment_factor,
verbose)
#send cured back to population
population[:, 11][population[:, 6] == 2] = 0
fatalities_plot.append(len(population[population[:, 6] == 3]))
if visualise:
#construct plot and visualise
spec = fig.add_gridspec(ncols=1, nrows=2, height_ratios=[5, 2])
ax1.clear()
ax2.clear()
ax1.set_xlim(x_plot[0], x_plot[1])
ax1.set_ylim(y_plot[0], y_plot[1])
if hospital_bounds != None:
build_hospital(hospital_bounds[0], hospital_bounds[2],
hospital_bounds[1], hospital_bounds[3], ax1)
healthy = population[population[:, 6] == 0][:, 1:3]
ax1.scatter(healthy[:healthcare_workers][:, 0],
healthy[:healthcare_workers][:, 1],
marker='P',
s=2,
color='green',
label='healthy')
ax1.scatter(healthy[healthcare_workers:][:, 0],
healthy[healthcare_workers:][:, 1],
color='gray',
s=2,
label='healthy')
infected = population[population[:, 6] == 1][:, 1:3]
ax1.scatter(infected[:, 0],
infected[:, 1],
color='red',
s=2,
label='infected')
immune = population[population[:, 6] == 2][:, 1:3]
ax1.scatter(immune[:, 0],
immune[:, 1],
color='green',
s=2,
label='immune')
fatalities = population[population[:, 6] == 3][:, 1:3]
ax1.scatter(fatalities[:, 0],
fatalities[:, 1],
color='black',
s=2,
label='dead')
#add text descriptors
ax1.text(
x_plot[0],
y_plot[1] + ((y_plot[1] - y_plot[0]) / 8),
'timestep: %i, total: %i, healthy: %i infected: %i immune: %i fatalities: %i'
% (frame, len(population), len(healthy), len(infected),
len(immune), len(fatalities)),
fontsize=6)
ax2.set_title('number of infected')
ax2.text(0,
pop_size * 0.05,
'https://github.com/paulvangentcom/python-corona-simulation',
fontsize=6,
alpha=0.5)
#ax2.set_xlim(0, simulation_steps)
ax2.set_ylim(0, pop_size + 100)
if treatment_dependent_risk:
infected_arr = np.asarray(infected_plot)
indices = np.argwhere(infected_arr >= healthcare_capacity)
ax2.plot([healthcare_capacity for x in range(len(infected_plot))],
color='red',
label='healthcare capacity')
ax2.plot(infected_plot, color='gray')
ax2.plot(fatalities_plot, color='black', label='fatalities')
ax2.legend(loc=1, fontsize=6)
if treatment_dependent_risk:
ax2.plot(indices,
infected_arr[infected_arr >= healthcare_capacity],
color='red')
#plt.savefig('render/%i.png' %frame)
return population