Animated SIR Model for Coronavirus Spread

A python 3.x program that animates the spread of a virus using a SIR model. Runs on any computer with numpy and matplotlib.

IntermediateFull instructions provided2 hours4,627

Things used in this project

Hardware components

Raspberry Pi 3 Model B

Software apps and online services

Python 3.x

Matplotlib

Story

During the coronavirus outbreak I was at home and I decided to review some of the things about object-oriented programming that I had learned some time ago. I didn't find anything worth coding untiI watched this 3blue1brown video (highly recommended) modelling the spread of a disease, evaluating the effect of different experiment setups and the effectiveness of different containment measures on the behavior of the cases over time.

I decided to code a far more simple animation with Python and Matplotlib also using a SIR model, which divides a population in Susceptible, Infected and Recovered individuals in each particular time. The simulation consists of an array of dots moving randomly through a 100x100px box. A susceptible individual has a given probability of getting the disease from an infected individual if he gets within a certain radius (this probability can be determined in the real world by how often do we wash our hand or touch our face), and after a certain amount time being sick he recovers and can no longer be infected.

As you can see in the picture, the model behavior reassembles the real behavior of the virus. At the beginning there is a phase of exponential growth, which is gradually substituted by a phase of linear growth until the peak is reached. Then, as more people start to heal and less people is left to infect, the number of infected starts to drop until it finally gets to zero. When this happens, the recovered curve stops growing, as there is no one left to heal.

Results of my SIR animation

The model includes parameters that can be modified to play around with the simulation, such as the number of individuals, the distance of transmission, the probability of transmission, etc. It also features a "quarantine mode" in which is possible to fix some of the points in its initial position through the whole simulation (as if they stayed home). As you will see at the end, this actually has the effect of "flattening the curve".

You can download the code and try it yourself!

Step 1: Creating a Python Class representing an individual

As the main objective of my project was to review some concepts of object-oriented programming, I started by creating a class to represent an individual taking part in the simulation. Each individual has the following attributes:

ID and name: a unique number and name used for identification purposes.
Status: three Boolean variables representing if the individual is Susceptible, Infected or Recovered.
Position: two coordinates representing the position of the individual in each time.
Speed: the speed with which the individual will move through the simulation area.
Recover time: time that the individual takes to recover.
Infection time: the moment of the simulation in which the individual got infected.
Fixed?: a Boolean that tells if the individual is quarantined.

I also defined some methods to be used throughout the simulation. They allow to change the status of the individual, get and set its position, animate the movement frame by frame, calculate the Euclidean distance to another point, etc. They can be found commented in the code.

Finally, I saved this class on a separated file to be imported in the code that runs the simulation.

Step 2: Coding the simulation

After creating the class, I coded the simulation in a different file. The first step is, as usual, importing the necessary packages including the class created in step 1 (both files must be saved in the same folder so the code can work).

Then the following simulation parameters are defined (the ones you can play with):

Number of individuals. if it's too high (>500 in my PC) simulation will run very slowly.
Percentage of infected individuals at the beginning of the simulation.
Radius of transmission in pixels.
Probability of transmission in percentage.
Recover time in number of frames
Percentage of the population not able to move (in quarantine).

Later all the individuals are created. A scatter plot is created and configured to host the dots moving and a 2D line plot is created and configured to display the curves showing amount of infected and the recovered at each time.

Finally, I coded the function that is continuously executed to generate each frame of the animation. It checks how much time has each individual been sick, and if it's currently infected it search for susceptible people within the transmission radius to infect them given the transmission probability. It also calculates the new position of each dot in the box, counts the number of infected and recovered people, and updates the information on the line plot.

You can find the code to run this simulation in your computer at the end of the page, if something doesn't work out of the box write me a message. The variable names are written in Spanish but the code is commented in English, however if you would like to get deeper into how the code works you can also reach me.

Step 3: Analyzing the results

I anticipated most of the results at the beginning, which can be summarized saying that the simulation does represent partly the real behavior of the Coronavirus, although this is still a very limited model.

As I said, in the model the quarantine does have the effect of flattening the curve, reducing its peak magnitude. The two images below were generated by the simulation running with the exact same parameters, but in the second one 70% of the individuals remained fixed in their initial position, and as you can see the difference is significant.

1 / 2

I hope this encourages you stay home and download the code to try it yourself, maybe changing parameters and running the simulation over and over you can discover new things that I missed.

PD: I'm a spanish speaker, so if you find any mistakes in the text please let me know!

import math
import numpy as np

class Persona:
    def __init__(self,i, posx, posy, objx, objy, v, t_contagiado, fijo):
        # movement speed
        self.v=v
        # target position
        self.objx=objx
        self.objy=objy
        #ID and name
        self.indice=i
        self.nombre="Persona "+str(i)
        #State: Susceptible, Infected or Retired
        self.infectado = False
        self.suceptible = True
        self.retirado = False
        #Current position
        self.posx = posx
        self.posy=posy
        #is it fixed (in quarantine)?
        self.fijo = fijo

        # displacement per iteration
        if self.fijo:

            self.deltax = 0
            self.deltay = 0
        else:
            self.deltax = (self.objx - self.posx) / self.v
            self.deltay = (self.objy - self.posy) / self.v
        #time in which the person was infected
        self.i_contagio=-1
        #time that the infection lasts, recover time
        self.t_contagiado = t_contagiado


    def __str__(self):
        return self.nombre+" en posicin "+str(self.posx)+", "+str(self.posy)

    def infectar(self,i):
        #infect
        self.infectado=True
        self.suceptible=False
        self.retirado = False
        self.i_contagio=i

    def retirar(self):
        #heal
        self.retirado=True
        self.suceptible=False
        self.infectado=False

    def set_objetivo(self,objx,objy):
        #this function is used to create a new target position
        self.objx=objx
        self.objy=objy
        if self.fijo:
            self.deltax = 0
            self.deltay=0
        else:
            self.deltax = (self.objx - self.posx) / self.v
            self.deltay = (self.objy - self.posy) / self.v
        print("Nuevo OBJ   ", self.objx,self.objy,"  ",self.indice)

    def check_contagio(self,i):
        #this function is used to heal the person if the established infection time has passed
        if self.i_contagio>-1:
            if i-self.i_contagio>self.t_contagiado:
                self.retirar()


    def update_pos(self, n_posx, n_posy):
        #this funcion animates the movement
        if(n_posx==0 and n_posy==0):
            self.posx=self.posx+self.deltax
            self.posy=self.posy+self.deltay
        else:
            self.posx=n_posx
            self.posy=n_posy

        if abs(self.posx-self.objx)<3 and abs(self.posy-self.objy)<3:
            self.set_objetivo(np.random.random()*100, np.random.random()*100)
        if self.posx>100:
            self.posx=100
        if self.posy>100:
            self.posy=100
        if self.posx<0:
            self.posx=0
        if self.posy<0:
            self.posy=0

    def get_color(self):
        if self.infectado:
            return 'red'
        if self.suceptible:
            return 'blue'
        if self.retirado:
            return 'gray'

    def get_pos(self):
        return (self.posx,self.posy)

    def get_dist(self,x,y):
        #this funcion calculates the distance between this person an another.
        return math.sqrt((self.posx-x)**2+(self.posy-y)**2)

from persona import Persona
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

#SIMULATION PARAMETERS
n=350  #number of individuals
p_infectadas = 1  #percentage of infected people at the beginning of the simulation (0-100%)
r_contagio=2  #radius of transmission in pixels (0-100)
p_contagio=4  #probability of transmission in percentage (0-100%)
p_aislamiento = 70  #percentage of the people in quarantine (0-100%)
t_contagiado=100   #time taken to recover in number of frames (0-infinity)


contagiados=0
personas=[]

#creating all the individuals in random positions. Infecting some of them
for i in range(n):
    p = Persona(i,np.random.random()*100, np.random.random()*100,
                np.random.random() * 100, np.random.random() * 100,
                (np.random.random()+0.5)*100,t_contagiado, False)

    if np.random.random()<p_infectadas/100:
        p.infectar(0)
        contagiados=contagiados+1
    if np.random.random()<p_aislamiento/100:
        p.fijo=True

    personas.append(p)


#this creates all the graphics
fig = plt.figure(figsize=(18,9))
ax = fig.add_subplot(1,2,1)
cx = fig.add_subplot(1,2,2)
ax.axis('off')
cx.axis([0,1000,0,n])
scatt=ax.scatter([p.posx for p in personas],
           [p.posy for p in personas],c='blue',s=8)
caja = plt.Rectangle((0,0),100,100,fill=False)
ax.add_patch(caja)
cvst,=cx.plot(contagiados,color="red",label="Currently infected")
rvst,=cx.plot(contagiados,color="gray",label="Recovered")
cx.legend(handles=[rvst,cvst])
cx.set_xlabel("Time")
cx.set_ylabel("People")


ct=[contagiados]
rt=[0]
t=[0]



#function excecuted frame by frame
def update(frame,rt,ct,t):
    contciclo = 0
    recuciclo = 0
    colores = []
    sizes = [8 for p in personas]
    for p in personas:
        #check how much time the person has been sick
        p.check_contagio(frame)
        #animate the movement of each person
        p.update_pos(0,0)
        if p.retirado:
            recuciclo+=1 #count the amount of recovered
        if p.infectado:
            contciclo=contciclo+1 #count the amount of infected
            #check for people around the sick individual and infect the ones within the
            # transmission radius given the probability
            for per in personas:
                if per.indice==p.indice or per.infectado or per.retirado:
                    pass
                else:
                    d=p.get_dist(per.posx,per.posy)
                    if d<r_contagio:
                        if np.random.random() < p_contagio / 100:
                            per.infectar(frame)
                            sizes[per.indice]=80


        colores.append(p.get_color()) #change dot color according to the person's status

    #update the plotting data
    ct.append(contciclo)
    rt.append(recuciclo)
    t.append(frame)


    #tramsfer de data to the matplotlib graphics
    offsets=np.array([[p.posx for p in personas],
                     [p.posy for p in personas]])
    scatt.set_offsets(np.ndarray.transpose(offsets))
    scatt.set_color(colores)
    scatt.set_sizes(sizes)
    cvst.set_data(t,ct)
    rvst.set_data(t,rt)
    return scatt,cvst,rvst

#run the animation indefinitely
animation = FuncAnimation(fig, update, interval=25,fargs=(rt,ct,t),blit=True)
plt.show()

Credits

Pepis

1 project • 3 followers

Thanks to Grant Sanderson (3blue1brown).

Animated SIR Model for Coronavirus Spread