Corona virus simulation

0. Corona virus simulation.md

Corona virus simulation

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from matplotlib.patches import Patch

# Set simulation details
healthy, sick, recovered, dead = 0, 1, 2, 3
p_infect, p_recover, p_die = 0.3, 0.2, 0.1
p_nochange = 1 - p_recover - p_die
p_notinfect = 1 - p_infect

def update(population):
    # Sick people randomly recover or die
    for x in range(population.shape[1]):
        for y in range(population.shape[0]):
            if population[y, x] == sick:
                population[y, x] = np.random.choice([sick, recovered, dead], 1,
                    p=[p_nochange, p_recover, p_die])
    # Sick people randomly infect their neighbours
    new_population = population.copy()
    for x in range(population.shape[1]):
        for y in range(population.shape[0]):
            if population[y, x] == sick:
                neighbours = population[y-1:y+2,x-1:x+2]
                infections = np.random.choice([healthy, sick], neighbours.shape,
                    p=[p_notinfect, p_infect])
                new_population[y-1:y+2,x-1:x+2] = np.where(
                    neighbours == healthy, infections, neighbours)
    # Modify array in-place
    np.copyto(population, new_population)


def plot_frame(frame_num, population, quad_mesh):
    plt.title("Day {} of the Corona Virus".format(frame_num))
    update(population)
    quad_mesh.set_array(population.ravel())

if __name__ == "__main__":
    np.random.seed(0)

    ny, nx = 50, 50
    population = np.full([ny, nx], healthy)
    population[:2, :2] = sick

    # Initialise plot and format
    fig = plt.figure(figsize=[8, 6])
    plt.axis("equal")
    plt.xlim([0, 50])
    plt.ylim([0, 50])
    quad_mesh = plt.pcolormesh(population, cmap="hsv", vmin=0, vmax=4)

    colours = plt.get_cmap("hsv")(np.linspace(0, 1, 4, endpoint=False))
    plt.legend(handles=[
        Patch(color=colours[0], label="Healthy"),
        Patch(color=colours[1], label="Sick"),
        Patch(color=colours[2], label="Recovered"),
        Patch(color=colours[3], label="Dead")])
    n_days = 200
    fps = 5

    # Create animation object and save as gif
    anim = animation.FuncAnimation(fig, plot_frame,
        fargs=[population, quad_mesh], save_count=n_days)
    anim.save("corona_virus.gif", writer=animation.PillowWriter(fps=fps))

1. corona_virus.gif