def infect(population, Config, frame, send_to_location=False, location_bounds=[], destinations=[], location_no=1, location_odds=1.0): '''finds new infections. Function that finds new infections in an area around infected persens defined by infection_range, and infects others with chance infection_chance Keyword arguments ----------------- population : ndarray array containing all data on the population pop_size : int the number if individuals in the population infection_range : float the radius around each infected person where transmission of virus can take place infection_chance : float the odds that the virus infects someone within range (range 0 to 1) frame : int the current timestep of the simulation healthcare_capacity : int the number of places available in the healthcare system verbose : bool whether to report illness events send_to_location : bool whether to give infected people a destination location_bounds : list the location bounds where the infected person is sent to and can roam within (xmin, ymin, xmax, ymax) destinations : list or ndarray the destinations vector containing destinations for each individual in the population. Needs to be of same length as population location_no : int the location number, used as index for destinations array if multiple possible destinations are defined location_odds: float the odds that someone goes to a location or not. Can be used to simulate non-compliance to for example self-isolation. traveling_infects : bool whether infected people heading to a destination can still infect others on the way there ''' #mark those already infected first infected_previous_step = population[population[:, 6] == 1] healthy_previous_step = population[population[:, 6] == 0] new_infections = [] #if less than half are infected, slice based on infected (to speed up computation) if len(infected_previous_step) < (Config.pop_size // 2): for patient in infected_previous_step: #define infection zone for patient infection_zone = [ patient[1] - Config.infection_range, patient[2] - Config.infection_range, patient[1] + Config.infection_range, patient[2] + Config.infection_range ] #find healthy people surrounding infected patient if Config.traveling_infects or patient[11] == 0: indices = find_nearby(population, infection_zone, kind='healthy') else: indices = [] for idx in indices: #roll die to see if healthy person will be infected if np.random.random() < Config.infection_chance: population[idx][6] = 1 population[idx][8] = frame if len(population[population[:, 10] == 1]) <= Config.healthcare_capacity: population[idx][10] = 1 if send_to_location: #send to location if die roll is positive if np.random.uniform() <= location_odds: population[idx],\ destinations[idx] = go_to_location(population[idx], destinations[idx], location_bounds, dest_no=location_no) else: pass new_infections.append(idx) else: #if more than half are infected slice based in healthy people (to speed up computation) for person in healthy_previous_step: #define infecftion range around healthy person infection_zone = [ person[1] - Config.infection_range, person[2] - Config.infection_range, person[1] + Config.infection_range, person[2] + Config.infection_range ] if person[ 6] == 0: #if person is not already infected, find if infected are nearby #find infected nearby healthy person if Config.traveling_infects: poplen = find_nearby(population, infection_zone, traveling_infects=True, kind='infected') else: poplen = find_nearby( population, infection_zone, traveling_infects=True, kind='infected', infected_previous_step=infected_previous_step) if poplen > 0: if np.random.random() < (Config.infection_chance * poplen): #roll die to see if healthy person will be infected population[np.int32(person[0])][6] = 1 population[np.int32(person[0])][8] = frame if len(population[population[:, 10] == 1]) <= Config.healthcare_capacity: population[np.int32(person[0])][10] = 1 if send_to_location: #send to location and add to treatment if die roll is positive if np.random.uniform() < location_odds: population[np.int32(person[0])],\ destinations[np.int32(person[0])] = go_to_location(population[np.int32(person[0])], destinations[np.int32(person[0])], location_bounds, dest_no=location_no) new_infections.append(np.int32(person[0])) if len(new_infections) > 0 and Config.verbose: print('\nat timestep %i these people got sick: %s' % (frame, new_infections)) if len(destinations) == 0: return population else: return population, destinations
def test_population(population, Config, frame, positive, send_to_location=False, location_bounds=[], destinations=[], location_no=1, location_odds=1.0): if frame % Config.min_ticks_between_tests != 0: # testing the individuals periodically return None, None positive_rolling_avg = 0 if len(positive) == 0 else np.mean(positive[-3:]) number_of_tests = 0 positive = 0 tests_pool_max = int(Config.max_tests_daily_proportion * Config.pop_size) tests_pool_min = int(Config.min_tests_daily_proportion * Config.pop_size) desired_test_count = int(1 / Config.desired_positive_proportion * positive_rolling_avg) tests_pool = int( np.clip(desired_test_count, tests_pool_min, tests_pool_max)) #sev2 cases sev2_mask = (population[:, 16] != 1) & (population[:, 15] == 2) sev2_idxs = np.asarray(population[sev2_mask, 0], dtype=np.int32) sev2_count = sum(sev2_mask) sev2_tests_count = min(sev2_count, tests_pool_max) sev2_selected_idxs = np.random.choice(sev2_idxs, sev2_tests_count, replace=False) population[sev2_selected_idxs, 16] = 1 number_of_tests += sev2_tests_count positive += sev2_tests_count #sev1 cases tests_left = tests_pool - number_of_tests if tests_left > 0: sev1_mask = (population[:, 16] != 1) & (population[:, 15] == 1) sev1_idxs = np.asarray(population[sev1_mask, 0], dtype=np.int32) sev1_count = sum(sev1_mask) sev1_tests_count = min(sev1_count, tests_left) sev1_selected_idxs = np.random.choice(sev1_idxs, sev1_tests_count, replace=False) population[sev1_selected_idxs, 16] = 1 number_of_tests += sev1_tests_count positive += sev1_tests_count tests_left = tests_pool - number_of_tests if tests_left > 0: sev0_mask = (population[:, 16] != 1) & (population[:, 15] < 1) sev0_idxs = np.asarray(population[sev0_mask, 0], dtype=np.int32) sev0_count = sum(sev0_mask) sev0_tests_count = min(sev0_count, tests_left) sev0_selected_idxs = np.random.choice(sev0_idxs, sev0_tests_count, replace=False) population[sev0_selected_idxs, 16] = population[sev0_selected_idxs, 6] == 1 number_of_tests += sev0_tests_count positive += sum(population[sev0_selected_idxs, 6] == 1) for idx in range(len(population)): is_sick = population[idx, 6] severity = int(population[idx, 15]) if severity == 2 and len( population[population[:, 10] == 1]) <= Config.healthcare_capacity: population[idx, 10] = 1 if is_sick and severity != -1 and send_to_location and np.random.uniform( ) <= max( population[idx, 16] == 1 + Config.self_isolate_severity_proportion[severity], 1) * location_odds: population[idx], destinations[idx] = go_to_location( population[idx, :], destinations[idx], location_bounds, dest_no=location_no) if Config.verbose: print('testing at timestep %i: %i/%i' % (frame, positive, number_of_tests)) return positive, number_of_tests