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
Example #2
0
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
Example #3
0
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