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