def check_at_destination(population,
destinations,
wander_factor=1.5,
speed=0.01):
"""check who is at their destination already
Takes subset of population with active destination and
tests who is at the required coordinates. Updates at destination
column for people at destination.
Keyword arguments
-----------------
population : ndarray
the array containing all the population information
destinations : ndarray
the array containing all destinations information
wander_factor : int or float
defines how far outside of 'wander range' the destination reached
is triggered
"""
# how many destinations are active
active_dests = np.unique(population[:, 11][(population[:, 11] != 0)])
# see who is at destination
for d in active_dests:
dest_x = destinations[:, int((d - 1) * 2)]
dest_y = destinations[:, int(((d - 1) * 2) + 1)]
# see who arrived at destination and filter out who already was there
at_dest = population[(np.abs(population[:, 1] - dest_x) <
(population[:, 13] * wander_factor))
& (np.abs(population[:, 2] - dest_y) <
(population[:, 14] * wander_factor))
& (population[:, 12] == 0)]
if len(at_dest) > 0:
# mark those as arrived
at_dest[:, 12] = 1
# insert random headings and speeds for those at destination
at_dest = update_randoms(
at_dest,
pop_size=len(at_dest),
speed=speed,
heading_update_chance=1,
speed_update_chance=1,
)
# at_dest[:,5] = 0.001
# reinsert into population
population[(np.abs(population[:, 1] - dest_x) <
(population[:, 13] * wander_factor))
& (np.abs(population[:, 2] - dest_y) <
(population[:, 14] * wander_factor))
& (population[:, 12] == 0)] = at_dest
return population
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 tstep(config, vir, pop, pop_tracker, soc, fig, spec, ax1, ax2):
'''
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(
pop.population[pop.population[:, 11] != 0]) # look op this only once
if active_dests > 0 and len(pop.population[pop.population[:,
12] == 0]) > 0:
pop.population = pop.set_destination()
pop.population = pop.check_at_destination()
if active_dests > 0 and len(pop.population[pop.population[:,
12] == 1]) > 0:
#keep them at destination
pop.population = pop.keep_at_destination()
#out of bounds
#define bounds arrays, excluding those who are marked as having a custom destination
if len(pop.population[:, 11] == 0) > 0:
_xbounds = np.array(
[[config.xbounds[0] + 0.02, config.xbounds[1] - 0.02]] *
len(pop.population[pop.population[:, 11] == 0]))
_ybounds = np.array(
[[config.ybounds[0] + 0.02, config.ybounds[1] - 0.02]] *
len(pop.population[pop.population[:, 11] == 0]))
pop.population[pop.population[:, 11] == 0] = out_of_bounds(
pop.population[pop.population[:, 11] == 0], _xbounds, _ybounds)
#set randoms
if soc.lockdown:
if len(pop_tracker.infectious) == 0:
mx = 0
else:
#mx = np.max(pop_tracker.infectious)
mx = pop_tracker.infectious[-1]
if len(pop.population[pop.population[:,6] == 1]) >= len(pop.population) * soc.lockdown_percentage or\
mx >= (len(pop.population) * soc.lockdown_percentage) or soc.lockdown_act:
soc.lockdown_act = True
#reduce speed of all members of society
pop.population[:, 5] = np.clip(pop.population[:, 5],
a_min=None,
a_max=0.00001)
#set speeds of complying people to 0
pop.population[:, 5][soc.lockdown_vector == 0] = 0
if len(pop.population[pop.population[:, 6] == 1]) <= len(
pop.population) * soc.lockdown_percentage / 2:
soc.lockdown_act = False
else:
#update randoms
pop.population = update_randoms(pop.population, pop.pop_size,
pop.speed)
else:
#update randoms
pop.population = update_randoms(pop.population, pop.pop_size,
pop.speed)
#for dead ones: set speed and heading to 0
pop.population[:, 3:5][pop.population[:, 6] == 3] = 0
#update positions
pop.population = update_positions(pop.population)
#find new infections
pop.population, pop.destinations = vir.infect(
pop,
soc,
config,
send_to_location=soc.self_isolate,
location_bounds=soc.isolation_bounds,
destinations=pop.destinations,
location_no=1,
location_odds=soc.self_isolate_proportion)
#recover and die
pop.population = vir.recover_or_die(pop, soc, config)
#send cured back to population if self isolation active
#perhaps put in recover or die class
#send cured back to population
pop.population[:, 11][pop.population[:, 6] == 2] = 0
#update population statistics
pop_tracker.update_counts(pop.population)
#visualise
if config.visualise:
draw_tstep(config, soc, pop.pop_size, pop.population, pop_tracker,
config.frame, fig, spec, ax1, 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' %
(config.frame, pop_tracker.susceptible[-1], pop_tracker.infectious[-1],
pop_tracker.recovered[-1],
len(pop.population[pop.population[:, 10] == 1]),
pop_tracker.fatalities[-1], pop.pop_size))
#save popdata if required
if config.save_pop and (config.frame % config.save_pop_freq) == 0:
pop.save_population(pop.population, config.frame,
config.save_pop_folder)
#run callback
callback(pop, config)
#update frame
config.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],
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
