Random walk simulation of disease spread

covidPeeps.py

# Module for simulating infectious-disease spread with a random-walk model
# with widgets to allow interactive control in a Jupyter notebook
# David Goldenberg, August 2020

# Revision to add vaccinated peeps!
# Sept 2021


import matplotlib.pyplot as plt
import numpy as np
import random as rand

from matplotlib import animation, rc

from matplotlib import cm
from matplotlib.colors import ListedColormap, LinearSegmentedColormap


from ipywidgets import interact, interactive, fixed
import ipywidgets as widgets
from IPython.display import display, clear_output, HTML, FileLink

display(HTML("<style>div.output_scroll { height: 500ex; }</style>"))


class Peep():
    '''Represents an individual who can move around on a two-dimensional grid and catch and transmit disease.
    Functions in conjunction with a two-dimensional grid represented by a two-d numpy array of integers.'''
    
    inf_state_names = ['susceptible','vaxxed','infected','recovered','dead']
    def __init__(self,x=0,y=0,inf_state=0,inf_clock=0, inf_per = 10, inf_prob = 0.1, mort = 0.01,vax_eff=0.8):
        self.x = x # x-coordinate on grid.
        self.y = y # y-coordinate on grid.
        self.inf_state=inf_state # infection state of peep
        self.inf_clock = inf_clock # integer counter to keep track of infection period
        self.inf_per = inf_per # Period of time (in move steps) before infected peep becomes immune.
        self.inf_prob = inf_prob # probability of transmission per time unit per contact.
        self.mort = mort # mortality probablility over period of infection
        
    def __str__(self):
        retStr = 'x='+str(self.x)+', y=' + str(self.y) + ', '+ self.inf_state_names[self.inf_state]
        return retStr
    
    def grid_place(self,grid):
        '''Changes value of grid position to represent infection state.'''
        grid[self.x,self.y] = self.inf_state + 1
        
    def rand_move(self,grid):
        '''Moves position of peep on the grid to an open adjacent space, choosing randomly from the open spaces.
        Can move to any of the eight closest spaces.'''
        # first locate adjacent open spaces
        # look at the 8 spaces surrounding the current location
        # array these from upper left clockwise
        spaces = [(self.x-1,self.y+1),(self.x,self.y+1),(self.x+1,self.y+1), (self.x+1,self.y),
                 (self.x+1,self.y-1),(self.x,self.y-1), (self.x-1,self.y-1),(self.x-1,self.y)]
        # Make a list of open spaces.
        open_spaces = []
        for space in spaces:
            # first check that we aren't going out of the grid
            if 0 <= space[0] < len(grid) and 0 <= space[1] < len(grid[0]):
                # then check to see if space is empty
                if grid[space]==0:
                    open_spaces.append(space)
        if len(open_spaces) > 0:
            # choose one of the open spaces
            space = open_spaces[rand.randrange(len(open_spaces))]
            # remove existing mark from grid
            grid[self.x,self.y] = 0
            # change coordinates to new value
            self.x, self.y = space
            # place new position on grid
            grid[self.x,self.y] = self.inf_state + 1
            
    def allow_inf(self,grid):
        '''Allows peep to become infected if one or more of the adjacent positions is occupied by another
        peep that is infected. Probability of infection is determined by the value of inf_prob, whether the
        peep is vaccinated, the effectiveness of the vaccine and the number of infected neighbors.'''
        # check to see if the peep is already infected, immune or dead.
        if self.inf_state > 1:
            return
        # check to see if inf_prob has a reasonable value.
        if self.inf_prob <=0 or self.inf_prob > 1:
            return
        # count infected neighbors
        spaces = [(self.x-1,self.y+1),(self.x,self.y+1),(self.x+1,self.y+1), (self.x+1,self.y),
                 (self.x+1,self.y-1),(self.x,self.y-1), (self.x-1,self.y-1),(self.x-1,self.y)]
        inf_neighbors = 0
        # check surrounding spaces
        for space in spaces:
            if 0 <= space[0] < len(grid) and 0 <= space[1] < len(grid[0]):
                if grid[space]==3:
                    inf_neighbors += 1
        if self.inf_state == 0: # unvaccinated
            total_prob = self.inf_prob*inf_neighbors
        elif self.inf_state == 1: # vaccinated
            total_prob = self.inf_prob*inf_neighbors*(1-vax_eff.value)
        rand_test = rand.random() # random value between 0 and 1 (not including 1)
        if total_prob > rand_test:
            self.inf_state = 2
            grid[self.x,self.y] = 3
                    
    def move(self,grid):
        '''Basic step in simulation. Makes random move and allows infection or death.
        If infection period is over, changes infection state to immune.'''
        # check to see if peep is dead. Dead peeps don't move!
        if self.inf_state==4:
            return
        # make random move
        self.rand_move(grid)
        # allow infection.
        self.allow_inf(grid)
        # if infected step the infection clock
        if self.inf_state == 2:
            self.inf_clock += 1
            # check to see if it is time to recover.
            if self.inf_clock > self.inf_per:
                self.inf_state = 3
                grid[self.x,self.y]=4
            # allow peep to die with defined probability
            else:
                mort_prob = self.mort/self.inf_per
                rand_test = rand.random()
                if mort_prob > rand_test:
                    self.inf_state = 4
                    grid[self.x,self.y]=5
    
# Color map for image plots
# see https://matplotlib.org/3.1.1/tutorials/colors/colormap-manipulation.html

cdict = {'red':   [[0.0,  0.8, 0.8],
                   [0.16, 0.8, 0.0],
                   [0.33, 0.0, 0.0],
                   [0.49, 0.0, 1.0],
                   [0.66, 1.0, 1.0],
                   [0.83, 1.0, 1.0],
                   [1.0,  1.0, 1.0]],
         'green': [[0.0,  0.8, 0.8],
                   [0.16, 0.8, 0.0],
                   [0.33, 0.0, 1.0],
                   [0.49, 1.0, 0.0],
                   [0.66, 0.0, 0.0],
                   [0.83, 0.0, 1.0],
                   [1.0,  1.0, 1.0]],
         'blue':  [[0.0,  0.8, 0.8],
                   [0.16, 0.8, 1.0],
                   [0.33, 0.0, 1.0],
                   [0.49, 1.0, 0.0],
                   [0.66, 0.0, 1.0],
                   [0.83, 1.0, 0.0],
                   [1.0,  0.0, 0.0]]}

peepcmp = LinearSegmentedColormap('peepCmap', segmentdata=cdict, N=256)

def grid_cts(grid):
    '''Counts frequencies of different values in np array used to
    represent grids. Input is an nxn numpy array. Output is an list of 
    four integer values, representing the number of grid cells with values of
    1,2,3,4 and 5, corresponding to the number of peeps in the uninfected, vaxxed,
    infected, recovered and dead states. (grid==0 is empty.)'''
    cts = []
    cts.append((grid==1).sum()) # uninfected (and unvaxxed)
    cts.append((grid==2).sum()) # vaxxed (and uninfected)
    cts.append((grid==3).sum()) # infected
    cts.append((grid==4).sum()) # recovered
    cts.append((grid==5).sum()) # dead
    return cts

def plot_counts(counts):
    '''Function to plot counts of peeps in different infection states.
    Accepts two arguments:
        4xn array with data to plot, as returned by sim_run().
        optional file name for graph to be saved. If None, no file is saved.'''
    
    fig = plt.figure(figsize=(10,5))
    ax= fig.add_subplot(1,1,1)

    ax.plot(counts[2]+counts[3]+counts[4],linestyle='--',color='k',label='Cases')
    ax.plot(counts[0],color='b',label='Susceptable-unvaccinated')
    ax.plot(counts[1],color='c',label='Susceptible-vaccinated')
    ax.plot(counts[2],color='r',label='Infected')
    ax.plot(counts[3],color='m',label='Recovered')
    ax.plot(counts[4],color='y',label='Dead')
    
    ax.set_xlabel('Time',fontsize=14)
    ax.set_ylabel('Count',fontsize=14)
    ax.tick_params(axis='both', direction='out', top=False, right=False, length=5, width=1.0, 
                    pad = 6, labelsize=12)   

    ax.legend()
    plt.close()
    return fig
    
def sim_setup():
    '''Function to set up a simulation by creating Peep objects and a grid.
    Uses parameters set with widgets.
    Returns a list of peep objects and a nxn numpy array representing the grid.
    peep[0] is placed at the middle of the grid.
    The rest of the peeps are placed randomly, avoiding overlaps with existing peeps. For each peep, 
    random placement is attempted 1,000 times before giving up. This should only be a problem
    sith very hign densities'''
    # create grid
    grid_len = grid_dim()
    grid = np.zeros((grid_len,grid_len),dtype=int)
    #create peeps
    peeps = []
    for i in range(peep_number.value):
        peeps.append(Peep(inf_prob=inf_prob.value,inf_per=inf_per.value,mort=mort_prob.value, vax_eff = vax_eff.value))    


    # place the first peep at the middle of the grid
    # to make it easy to start the infection there
    peeps[0].x = int(len(grid)/2)
    peeps[0].y = int(len(grid[0])/2)
    peeps[0].grid_place(grid)

    # place rest of peeps on grid randomly, avoiding overlaps
    for peep in peeps[1:]:
        placed = False
        tries = 0
        while placed == False:
            x = rand.randrange(len(grid))
            y = rand.randrange(len(grid[0]))
            if grid[x,y] == 0:
                peep.x = x
                peep.y = y
                peep.grid_place(grid)
                placed = True

            tries += 1
            if tries >= 1000:
                placed = True
                
    # infect init_infect peeps
    if init_infect.value == 1:
        # if only a single peep is to start out infected
        # it is the one at the center at the grid
        peeps[0].inf_state = 2
        peeps[0].grid_place(grid)
    elif init_infect.value > 1:
        # if more than one peep is to start out infected
        # they are placed randomly
        i,j = 0,0
        while j < peep_number.value and i < init_infect.value:
            if peeps[j].inf_state == 0:
                peeps[j].inf_state = 2
                peeps[j].grid_place(grid)
                i += 1
            j += 1
            
    # vaccinate vaxxed peeps
    n_vaxxed = int(peep_number.value*frac_vaxxed.value)
    uninf = peep_number.value - init_infect.value
    i,j = 0,0
    while j < uninf and i < n_vaxxed:
        if peeps[j].inf_state == 0:
            peeps[j].inf_state = 1
            peeps[j].grid_place(grid)
            i += 1
        j += 1

    return peeps, grid
    
def grid_dim():
    if grid_size.value == '100 x 100 (10,000)':
        return 100
    elif grid_size.value == '200 x 200 (40,000)':
        return 200
    elif grid_size.value == '400 x 400 (160,000)':
        return 400
        
def run_sim():
        
    time = tot_time.value
    peeps,grid = sim_setup()
    
    fig1 = plt.figure(figsize=(8,8))
    #creating a subplot 
    ax1 = fig1.add_subplot(1,1,1)
    counts = []
    prog_bar.value=0
    
    def animate(i):
        for peep in peeps:
            peep.move(grid)
        counts.append(grid_cts(grid))
            #print(peep)
        ax1.clear()
        ax1.imshow(grid.T*0.2,origin='lower',cmap = peepcmp,vmin=0,vmax=1)
        prog_bar.value += 1/time
        return 
        
    ani = animation.FuncAnimation(fig1, animate, frames=time,interval=100);
    ani_html = HTML(ani.to_html5_video())

    
    counts = np.array(counts)
    counts = counts.transpose()
    plt.close()
    
    return ani_html, counts
   
def peepWidgets():
    '''function to set up and display widgets'''
    display(intro)
       
    row1 = widgets.HBox([peep_number,grid_size])
    row2 = widgets.HBox([init_infect,inf_prob])
    row3 =widgets.HBox([frac_vaxxed,vax_eff])
    row4 =widgets.HBox([mort_prob])
    row5 = widgets.HBox([tot_time,inf_per])
    row6 = widgets.HBox([run_button,prog_bar])
    display(widgets.VBox([row1, row2, row3, row4,row5,row6]))
    
    display(fig_output)
    
    return


def run_click(b):
    global fig
    run_button.description = 'Running'

    fig_output.clear_output()
    display(HTML("<style>div.output_scroll { height: 1800px; }</style>"))
    
    
    ani_html, counts =run_sim()
 
    
    
    fig = plot_counts(counts)
    with fig_output:
        clear_output(wait=True)
        display(fig)
        display(ani_html)
        
        
    run_button.description = 'Run simulation'
    

    
##### Widgets  #####


intro_html = """
<h2> A Random-walk Model of Disease Spread </h2>
<p style="font-size:16px">David P. Goldenberg, September 2021
<br>
School of Biological Sciences, University of Utah 
<br>
Salt Lake City, Utah 84105
</p>
<hr border-width:6px, color:black>
<p style="font-size:16px">
This web page demonstrates a simple random-walk model of disease spread in a population, as might be applied to the COVID-19 pandemic. The model and simulation were developed for use in Biology 3550, Physical Prinicples in Biology, for the fall semester of 2020. More information about the course can be found at: <a href="https://goldenberg.biology.utah.edu/courses/biol3550/index.shtml" target="new">https://goldenberg.biology.utah.edu/courses/biol3550.</a>
<br>
The basic assumptions of the model are that: </p>
<ol style="font-size:16px">
	<li>individuals, called peeps, live on a two-dimensional grid.</li>
	<li>Each peep can be in one of four infection states:</li>
	<ul style="font-size:16px">
	    <li> Susceptible and unvaccinated</li>
	    <li> Susceptible and vaccinated</li>
	    <li> Infected and contagious</li>
	    <li> Recoverd and non-contagious</li>
	    <li> Dead</li>
	</ul> 
	<li> Peeps move about on the grid by taking random single steps to one of the eight adjacent spots on the grid, provided that the spot is unoccupied.</li>
	<li> Once infected, a peep remains infected and contagious for a specified period, unless it dies first. </li>
	<li> Infected peeps die with a specified probability that remains constant over the infection period. </li>	
</ol>
<p style="font-size:16px">
The parameters for the simulation can be adjusted using the controls below. After setting the parameters, click the "Run Simulation" button.
<br>
The simulation will take a minute or few to run, depending on the total number of peeps and the total time.
After the simulation is complete, a animation will be displayed, along with a graph of the number of peeps in the different states as a function of time.
<br>
To download an image file of the graph, control-click (MacOS) or right-cick (Windows) on the image and choose the "save image" item in the drop-down menu. To download a movie file of the animation, click on the three dots on the right-hand end of the movie control bar and choose "Download".
</p>
<hr border-width: 6px,color:black>

"""

fig_output = widgets.Output()


intro= widgets.HTML(
    value=intro_html)

peep_number=widgets.BoundedIntText(
    value=2500,
    min=0,
    max=5000,
    step=500,
    description='No. of Peeps',
    disabled=False,
    layout=widgets.Layout(height='40px',margin='0px 20px 0px 0px',width='200px')
)
grid_size=widgets.Dropdown(
options=['100 x 100 (10,000)', '200 x 200 (40,000)', '400 x 400 (160,000)'],
    value='100 x 100 (10,000)',
    description='Grid size:',
    disabled=False,
    layout=widgets.Layout(width='250px',margin='0px 20px 0px 5px')
)

label_style = {'description_width': 'initial'}
inf_prob=widgets.BoundedFloatText(
    value=0.2,
    min=0, 
    max=1,
    step = 0.05,
    description='Infection prob.',
    style = label_style,
    layout=widgets.Layout(height='40px',margin='0px 20px 0px 75px',width='200px')
)
mort_prob=widgets.BoundedFloatText(
    value=0.01,
    min=0, 
    max=1,
    step = 0.01,
    description='Mortality rate',
    style = label_style,
    layout=widgets.Layout(height='40px',width='200px',margin='0px 20px 0px 0px')
)

tot_time=widgets.BoundedIntText(
    value=150,
    min=0, 
    max=500,
    step = 50,
    description='Total time',
    style = label_style,
    layout=widgets.Layout(height='40px',width='175px',margin='0px 20px 0px 25px')
)

inf_per=widgets.BoundedIntText(
    value=14,
    min=0, 
    max=101,
    step = 5,
    description='Infection period',
    style = label_style,
    layout=widgets.Layout(height='40px',width='200px',margin='0px 20px 0px 55px')
)


init_infect=widgets.BoundedIntText(
    value=4,
    min=0, 
    max=5000,
    step = 5,
    description='Initially infected',
    style = label_style,
    layout=widgets.Layout(height='40px',width='200px',margin='0px 0px 0px 0px')
)

frac_vaxxed=widgets.BoundedFloatText(
    value=0,
    min=0, 
    max=1,
    step = 0.1,
    description='Fraction vaccinated',
    style = label_style,
    layout=widgets.Layout(height='40px',width='200px',margin='0px 20px 0px 0px')
)

vax_eff = widgets.BoundedFloatText(
    value=0.8,
    min=0, 
    max=1,
    step = 0.1,
    description='Vaccine effectiveness',
    style = label_style,
    layout=widgets.Layout(height='40px',width='200px',margin='0px 20px 0px 55px')
)


run_button=widgets.Button(
    description='Run Simulation')
    
prog_bar= widgets.FloatProgress(
    description='',
    value=0,
    min=0,
    max=1,
    step=0.05,
    orientation='horizontal')
    
save_plot_button=widgets.Button(
    description='Save plot')

run_button.on_click(run_click)


def on_grid_change(change):
    grid_spots = grid_dim()**2
    peep_number.max=int(0.8*grid_spots)
    
def on_pop_change(change):
    n_vaxxed = int(frac_vaxxed.value*peep_number.value)
    test_numb = init_infect.value + n_vaxxed
    if test_numb > peep_number.value:        
        n_vaxxed = peep_number.value - init_infect.value
        frac_vaxxed.value = n_vaxxed/peep_number.value
        
    init_infect.max=peep_number.value
    
    
def on_frac_vaxxed_change(change):
    n_vaxxed = int(peep_number.value*frac_vaxxed.value)
    test_numb = init_infect.value + n_vaxxed
    if test_numb > peep_number.value:
        if test_numb <= peep_number.max:
            peep_number.value = test_numb
        else:
            peep_number.value = peep_number.max
            frac_vaxxed.value = (peep_number.value - init_infect.value)/peep_number.value
            
    
def on_inf_change(change):
    n_vaxxed = int(peep_number.value*frac_vaxxed.value)
    test_numb = init_infect.value + n_vaxxed
    if test_numb > peep_number.value:
        if test_numb <= peep_number.max:
            peep_number.value = init_infect.value + n_vaxxed
        else:
            peep_number.value = peep_number.max
            n_vaxxed = peep_number.value - init_infect.value
            frac_vaxxed.value = n_vaxxed/peep_number.value
            
def on_time_change(change):
    inf_per.max = tot_time.value +1
    
grid_size.observe(on_grid_change,names='value')
peep_number.observe(on_pop_change,names='value')
init_infect.observe(on_inf_change,names='value')
frac_vaxxed.observe(on_frac_vaxxed_change,names='value')
tot_time.observe(on_time_change,names='value')

peepWidgets()



    


    


    

covidSimulation.ipynb

environment.yml

name: covid_peeps
channels:
  - conda-forge
  - defaults
  - conda-forge/label/broken
dependencies:
  - python=3.8.3
  - numpy=1.19.2
  - scipy = 1.5.2
  - matplotlib=3.3.2
  - ipython=7.19.0
  - ipywidgets=7.5.1
  - ffmpeg