Pepis
Published © GPL3+

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 hours1,993

Things used in this project

Hardware components

Raspberry Pi 3 Model B
Raspberry Pi 3 Model B
×1

Software apps and online services

Python 3.x
Matplotlib

Story

Read more

Code

persona

Python
Person Class
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)

simulacion

Python
Simulation
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

Pepis

0 projects • 0 followers

Comments