def draw_tstep(Config, population, pop_tracker, frame, fig, spec, ax1, ax2):
# construct plot and visualise
# set plot style
set_style(Config)
# get color palettes
palette = Config.get_palette()
spec = fig.add_gridspec(ncols=1, nrows=2, height_ratios=[5, 2])
ax1.clear()
ax2.clear()
ax1.set_xlim(Config.x_plot[0], Config.x_plot[1])
ax1.set_ylim(Config.y_plot[0], Config.y_plot[1])
if Config.self_isolate and Config.isolation_bounds != None:
build_hospital(
Config.isolation_bounds[0],
Config.isolation_bounds[2],
Config.isolation_bounds[1],
Config.isolation_bounds[3],
ax1,
addcross=False,
)
# plot population segments
healthy = population[population[:, 6] == 0][:, 1:3]
ax1.scatter(healthy[:, 0],
healthy[:, 1],
color=palette[0],
s=2,
label="healthy")
infected = population[population[:, 6] == 1][:, 1:3]
ax1.scatter(infected[:, 0],
infected[:, 1],
color=palette[1],
s=2,
label="infected")
immune = population[population[:, 6] == 2][:, 1:3]
ax1.scatter(immune[:, 0],
immune[:, 1],
color=palette[2],
s=2,
label="immune")
fatalities = population[population[:, 6] == 3][:, 1:3]
ax1.scatter(fatalities[:, 0],
fatalities[:, 1],
color=palette[3],
s=2,
label="dead")
# add text descriptors
ax1.text(
Config.x_plot[0],
Config.y_plot[1] + ((Config.y_plot[1] - Config.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,
Config.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, Config.pop_size + 200)
if Config.treatment_dependent_risk:
infected_arr = np.asarray(pop_tracker.infectious)
indices = np.argwhere(infected_arr >= Config.healthcare_capacity)
ax2.plot(
[
Config.healthcare_capacity
for x in range(len(pop_tracker.infectious))
],
"r:",
label="healthcare capacity",
)
if Config.plot_mode.lower() == "default":
ax2.plot(pop_tracker.infectious, color=palette[1])
ax2.plot(pop_tracker.fatalities, color=palette[3], label="fatalities")
elif Config.plot_mode.lower() == "sir":
ax2.plot(pop_tracker.infectious, color=palette[1], label="infectious")
ax2.plot(pop_tracker.fatalities, color=palette[3], label="fatalities")
ax2.plot(pop_tracker.susceptible,
color=palette[0],
label="susceptible")
ax2.plot(pop_tracker.recovered, color=palette[2], label="recovered")
else:
raise ValueError(
"incorrect plot_style specified, use 'sir' or 'default'")
ax2.legend(loc="best", fontsize=6)
plt.draw()
plt.pause(0.0001)
if Config.save_plot:
try:
plt.savefig("%s/%i.png" % (Config.plot_path, frame))
except:
check_folder(Config.plot_path)
plt.savefig("%s/%i.png" % (Config.plot_path, frame))
def draw_tstep(Config, population, pop_tracker, frame, fig, spec, ax1, ax2,
ax3):
#construct plot and visualise
#set plot style
set_style(Config)
#get color palettes
palette = Config.get_palette()
spec = fig.add_gridspec(ncols=1, nrows=3, height_ratios=[5, 2, 2])
ax1.clear()
ax2.clear()
ax3.clear()
ax1.set_xlim(Config.x_plot[0], Config.x_plot[1])
ax1.set_ylim(Config.y_plot[0], Config.y_plot[1])
if Config.self_isolate and Config.isolation_bounds != None:
build_hospital(Config.isolation_bounds[0],
Config.isolation_bounds[2],
Config.isolation_bounds[1],
Config.isolation_bounds[3],
ax1,
addcross=False)
#plot population segments
healthy = population[population[:, 6] == 0][:, 1:3]
ax1.scatter(healthy[:, 0],
healthy[:, 1],
color=palette[0],
s=2,
label='healthy')
infected = population[population[:, 6] == 1][:, 1:3]
ax1.scatter(infected[:, 0],
infected[:, 1],
color=palette[1],
s=2,
label='infected')
immune = population[population[:, 6] == 2][:, 1:3]
ax1.scatter(immune[:, 0],
immune[:, 1],
color=palette[2],
s=2,
label='immune')
fatalities = population[population[:, 6] == 3][:, 1:3]
ax1.scatter(fatalities[:, 0],
fatalities[:, 1],
color=palette[3],
s=2,
label='dead')
#add text descriptors
ax1.text(
Config.x_plot[0],
Config.y_plot[1] + ((Config.y_plot[1] - Config.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.set_xlim(0, simulation_steps)
ax2.set_ylim(0, Config.pop_size + 200)
if Config.treatment_dependent_risk:
infected_arr = np.asarray(pop_tracker.infectious)
indices = np.argwhere(infected_arr >= Config.healthcare_capacity)
ax2.plot([
Config.healthcare_capacity
for x in range(len(pop_tracker.infectious))
],
'r:',
label='healthcare capacity')
if Config.plot_mode.lower() == 'default':
ax2.plot(pop_tracker.infectious, color=palette[1])
ax2.plot(pop_tracker.fatalities, color=palette[3], label='fatalities')
elif Config.plot_mode.lower() == 'sir':
ax2.plot(pop_tracker.susceptible,
color=palette[0],
label='susceptible')
ax2.plot(pop_tracker.infectious, color=palette[1], label='infectious')
ax2.plot(pop_tracker.recovered, color=palette[2], label='recovered')
ax2.plot(pop_tracker.fatalities, color=palette[3], label='fatalities')
else:
raise ValueError(
'incorrect plot_style specified, use \'sir\' or \'default\'')
ax2.legend(loc='best', fontsize=6)
ax3.set_title('GDP impacts')
# ax2.set_xlim(0, simulation_steps)
ax3.set_ylim(0, 50000000)
ax3.plot(pop_tracker.peoplegdp, color=palette[0], label='people')
ax3.plot(pop_tracker.businessgdp, color=palette[1], label='business')
ax3.plot(pop_tracker.governgdp, color=palette[2], label='govern')
ax3.plot(pop_tracker.totalgdp, color=palette[3], label='total')
ax3.legend(loc='best', fontsize=6)
plt.draw()
plt.pause(0.0000001)
if Config.save_plot:
try:
plt.savefig('%s/%i.png' % (Config.plot_path, frame))
except:
check_folder(Config.plot_path)
plt.savefig('%s/%i.png' % (Config.plot_path, frame))
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 build_fig(Config, figsize=(10,5)):
set_style(Config)
if not Config.self_isolate:
fig = plt.figure(figsize=(10,5))
spec = fig.add_gridspec(ncols=2, nrows=1, width_ratios=[5,5])
elif Config.self_isolate:
fig = plt.figure(figsize=(12,5))
spec = fig.add_gridspec(ncols=2, nrows=1, width_ratios=[7,5])
ax1 = fig.add_subplot(spec[0,0])
# plt.title('infection simulation')
plt.xlim(Config.xbounds[0], Config.xbounds[1])
plt.ylim(Config.ybounds[0], Config.ybounds[1])
lower_corner = (Config.xbounds[0],Config.ybounds[0])
width = Config.xbounds[1] - Config.xbounds[0]
height = Config.ybounds[1] - Config.ybounds[0]
# Draw boundary of world
if Config.plot_style.lower() == 'dark':
bound_color = 'w'
elif Config.plot_style.lower() == 'default':
bound_color = 'k'
rect = patches.Rectangle(lower_corner, width, height, linewidth=1, edgecolor=bound_color, facecolor='none', fill='None', hatch=None)
# Add the patch to the Axes
ax1.add_patch(rect)
if Config.self_isolate and Config.isolation_bounds != None:
build_hospital(Config.isolation_bounds[0], Config.isolation_bounds[2],
Config.isolation_bounds[1], Config.isolation_bounds[3], ax1,
bound_color, addcross = False)
ax1.axis('off')
# SIR graph
ax2 = fig.add_subplot(spec[0,1])
# ax2.set_title('number of infected')
#ax2.set_xlim(0, simulation_steps)
ax2.set_ylim(0, Config.pop_size)
ax2.set_xlabel('Simulation Steps', fontsize = 14)
ax2.set_ylabel('Number sof people', fontsize = 14)
#get color palettes
palette = Config.get_palette()
# Legend actors
# a1 = mlines.Line2D([], [], color=palette[1], marker='', markersize=5, linestyle=':')
# a2 = mlines.Line2D([], [], color=palette[1], marker='', markersize=5, linestyle='-')
# a3 = mlines.Line2D([], [], color=palette[3], marker='', markersize=5, linestyle='-')
# a4 = mlines.Line2D([], [], color=palette[0], marker='', markersize=5, linestyle='-')
# a5 = mlines.Line2D([], [], color=palette[2], marker='', markersize=5, linestyle='-')
# Legend actors type 2
a1 = mlines.Line2D([], [], color=palette[1], marker='', markersize=5, linestyle=':')
a2 = patches.Rectangle((20,20), 20, 20, linewidth=1, edgecolor='none', facecolor=palette[1], fill='None', hatch=None)
a3 = patches.Rectangle((20,20), 20, 20, linewidth=1, edgecolor='none', facecolor=palette[0], fill='None', hatch=None)
a4 = patches.Rectangle((20,20), 20, 20, linewidth=1, edgecolor='none', facecolor=palette[2], fill='None', hatch=None)
a5 = patches.Rectangle((20,20), 20, 20, linewidth=1, edgecolor='none', facecolor=palette[3], fill='None', hatch=None)
handles, labels = [[a1,a2,a3,a4,a5], ['healthcare capacity','infectious','healthy','recovered','dead']]
fig.legend(handles, labels, loc='upper center', ncol=5, fontsize = 10)
# Get tight figure bbox
tight_bbox_raw = fig.get_tightbbox(fig.canvas.get_renderer())
# tight_bbox = TransformedBbox(tight_bbox_raw, Affine2D().scale(1./fig.dpi))
#if
return fig, spec, ax1, ax2, tight_bbox_raw
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
