How to Simulate a Pandemic in Python | by Terence Shin | Towards Data Science

What’s a better time to simulate the spread of a disease than during a global pandemic? I don’t have much more to say — let’s jump right into programming a simple disease simulation.
In real life, there are hundreds of factors that affect how fast a contagion spreads, both from person to person and on a broader population-wide scale. I’m no epidemiologist but I’ve done my best to set up a fairly basic simulation that can mimic how a virus can infect people and spread throughout a population.
In my program, I will be using object-based programming. With this method, we could theoretically customize individual people and add in more events and factors — such as more complicated social dynamics.
Keep in mind that this is an introduction and serves as the most basic model that can be built on top of.
Fundamentally, our program will function around a single concept: any given person who is infected by our simulation’s disease has the potential to spread it to whoever they meet. Each person in our “peopleDictionary” will have a set number of friends (gaussian randomization for accuracy) and they may meet any one or more of these friends on a day to day basis.
For our starting round of simulations, we won’t implement face masks or lockdowns — we’ll just let the virus spread when people meet their friends and see if we can get that iconic pandemic “curve” which the news always talks about flattening.
So, we’ll use a Person() class and add a few characteristics. Firstly, we’ll assume that some very tiny percentage of characters simulated will already have immunity to our disease from the get-go, for whatever reason. I’m setting that at 1% (in reality, it’d be far lower but because our simulation runs so fast, a large portion like this makes a bit more sense). At the start of the simulation, the user will be prompted to enter this percentage.
Next, we have contagiousness, the all-important factor. When a person is not infected, this remains at 0. It also returns to 0 once a person ceases to be contagious and gains immunity. However, when a person is infected, this contagious value is somewhere between 0 and 100%, and it massively changes their chance of infecting a friend.
Before we implement this factor, we need to understand Gaussian Distribution. This mathematical function allows us to more accurately calculate random values between 1 and 100. Rather than the values being distributed purely randomly across the spectrum, most of them cluster around the median average point, making for a more realistic output:
As you can see, this bell-shaped function will be a lot better for our random characteristic variables because most people will have an average level of contagiousness, rather than a purely random percentage. I’ll show you how to implement this later.
We then have the variables “mask” and “lockdown” which are both boolean variables. These will be used to add a little bit of variety to our simulation after it is running.
Lastly, we have the “friends” variable for any given person. Just like contagiousness, this is a Gaussian Distribution that ends up with most people having about 5 friends that they regularly see. In our simulation, everyone lives in a super social society where on average a person meets with 2 people face to face every day. In real life, this is probably not as realistic but we’re using it because we don’t want a super slow simulation. Of course, you can make any modifications to the code that you like.
There are also a couple of other variables that will be used actively in the simulation and I’ll get to those as we go on!
So let’s get coding this simulation! First, there are three imports we have to do:
from scipy.stats import normimport randomimport time
SciPy will allow us to calculate values within the Gaussian Distribution we talked about. The random library will be for any variables we need that should be purely random, and the time library is just for convenience if we want to run the simulation slowly and watch the spread of the disease.
Next, we create our Person() class:
# simulation of a single personclass Person():    def __init__(self, startingImmunity):        if random.randint(0,100)<startingImmunity:            self.immunity = True        else:            self.immunity = False        self.contagiousness = 0        self.mask = False        self.contagiousDays = 0        #use gaussian distribution for number of friends; average is 5 friends        self.friends = int((norm.rvs(size=1,loc=0.5,scale=0.15)[0]*10).round(0))    def wearMask(self):        self.contagiousness /= 2
Why are we passing the variable startingImmunity to this class exactly? Remember how we could enter what percentage of the population would have natural immunity from day 1? When the user gives this percentage, for every person “spawned” into our simulation we’ll use random to find out if they’re one of those lucky few to already be immune — in which case the self.immunity boolean is set to True, protecting them from all infection down the line.
The remaining class variables are self-explanatory, except self.friends, which is the Gaussian Distribution we talked about. It’s definitely worth reading the documentation to get a better idea of how this works!
def initiateSim():    numPeople = int(input("Population: "))    startingImmunity = int(input("Percentage of people with natural immunity: "))    startingInfecters = int(input("How many people will be infectious at t=0: "))    for x in range(0,numPeople):        peopleDictionary.append(Person(startingImmunity))    for x in range(0,startingInfecters):        peopleDictionary[random.randint(0,len(peopleDictionary)-1)].contagiousness = int((norm.rvs(size=1,loc=0.5,scale=0.15)[0]*10).round(0)*10)    daysContagious = int(input("How many days contagious: "))    lockdownDay = int(input("Day for lockdown to be enforced: "))    maskDay = int(input("Day for masks to be used: "))    return daysContagious, lockdownDay, maskDay
After setting up our class, we need a function to initiate the simulation. I’m calling this initiateSim() and it’ll prompt the user for four inputs — population, natural immunity population, contagious people at day 0, and how many days a person will stay contagious for. This daysContagious variable should actually be random — or even better, dependent on any number of personal health conditions, such as immune compromisation — but let’s keep it like this for a basic simulation. I found from testing that it is most interesting to run the simulation with a 4–9 day contagious period.
We spawn the inputted number of people into the simulation. To start the disease, we pick people at random to be our “startingInfecters”. As you can see, we’re assigning a Gaussian variable to each one for their level of contagiousness! (Any time a person is made contagious in the simulation we’ll repeat this process.)
We return the number of days someone will stay contagious for, like mentioned.
Now, this simulation will be done day by day, so let’s set up a function:
def runDay(daysContagious, lockdown):    #this section simulates the spread, so it only operates on contagious people, thus:    for person in [person for person in peopleDictionary if person.contagiousness>0 and person.friends>0]:        peopleCouldMeetToday = int(person.friends/2)        if peopleCouldMeetToday > 0:            peopleMetToday = random.randint(0,peopleCouldMeetToday)        else:            peopleMetToday = 0                    if lockdown == True:            peopleMetToday= 0                    for x in range(0,peopleMetToday):            friendInQuestion = peopleDictionary[random.randint(0,len(peopleDictionary)-1)]            if random.randint(0,100)<person.contagiousness and friendInQuestion.contagiousness == 0 and friendInQuestion.immunity==False:                friendInQuestion.contagiousness = int((norm.rvs(size=1,loc=0.5,scale=0.15)[0]*10).round(0)*10)                print(peopleDictionary.index(person), " >>> ", peopleDictionary.index(friendInQuestion))
The runDay function takes daysContagious for reasons explained later. In our first for loop, we’re using a list comprehension to find the people who are capable of spreading the disease — that is, they are contagious and have friends. We’re then calculating the number of people they could meet on that day. The maximum is 50% of their friends, and then we’re using a standard random.randint() to generate how many they actually do meet on that day.
Then we use another embedded for loop to randomly select each friend that was met from the peopleDictionary[]. For the friend to have a chance of being infected, they can’t be immune to the disease. They also have to have a contagiousness of 0 — if they’re already infected, the encounter won’t influence them. We then use the infecter’s contagiousness percentage in a random function to find out if the friendInQuestion will be infected. Finally, if they do get infected, we go ahead and assign them a Gaussian Distribution variable for their contagiousness!
I added in a simple print statement as a marker which will allow us to follow the simulation in the console as it is running. At the end of our program, we’ll add functionality to save the results to a text file anyway, but it’s cool to see little tags that tell you who is infecting who.
Next part of our runDay() function:
for person in [person for person in peopleDictionary if person.contagiousness>0]:        person.contagiousDays += 1        if person.contagiousDays > daysContagious:            person.immunity = True            person.contagiousness = 0            print("|||", peopleDictionary.index(person), " |||")
Basically, all we’re doing here is finding all the people who are contagious and incrementing their contagiousDays variable by 1. If they’ve been contagious for more days than the daysContagious time the user selected, they will become immune and hence their contagiousness drops to 0. (Again, another print marker to show that the given person has gained immunity.)
I know I could have put this in the previous for loop but not to make my programming too dense, I separated it. Sue me.
Finally, to tie it all together, we need to do a bit of admin:
lockdown = FalsedaysContagious, lockdownDay, maskDay = initiateSim()saveFile = open("pandemicsave3.txt", "a")for x in range(0,100):    if x==lockdownDay:        lockdown = True            if x == maskDay:        for person in peopleDictionary:            person.wearMask()                print("DAY ", x)    runDay(daysContagious,lockdown)    write = str(len([person for person in peopleDictionary if person.contagiousness>0])) + "\n"    saveFile.write(write)    print(len([person for person in peopleDictionary if person.contagiousness>0]), " people are contagious on this day.")saveFile.close()
This is pretty self-explanatory. We get the daysContagious value by initiating the simulation, we open our save file, then cycle through the days up to day 100. Each day we use a list comprehension to get the number of people contagious and write it to our save file. I also added one final print statement so we can track the disease’s progression in the console.
And that’s it! I only explained the basics of the code, but let’s talk about the extra variables that you may have noticed...
Adding a lockdown variable is quite simple. First, add this in before the section where we cycle through each of the friends a person meets (see code above):
if lockdown == True:   peopleMetToday = 0for x in range(0, peopleMetToday):
Now, you want to select when the lockdown is enforced? No problem. Add a user prompt tight inside your initiateSim() function.
lockdownDay = int(input("Day for lockdown to be enforced: "))return daysContagiousreturn lockdownDay
Return it, and update the function call. Then, we need to define our lockdown boolean, and set it to true when we reach the correct date:
lockdown = FalsedaysContagious, lockdownDay = initiateSim()saveFile = open("pandemicsave2.txt", "a")for x in range(0,100):   if x == lockdownDay:      lockdown = True   print("DAY ", x)
You can see that I just added 3 more lines into where we manage the simulation. Simple and easy, then you will want to pass the lockdown boolean to your runDay() function and make sure the runDay() function can accept it:
runDay(daysContagious, lockdown)
And:
def runDay(daysContagious, lockdown):
That’s the lockdown added. See the results section to find out how the implementation of a lockdown affected the spread of the disease!
Finally, we want to add facemasks. I could add all sorts of ways that this changes how a disease spreads, but for us, we’ll just use it to decrease each person’s contagiousness. All we have to do is give the Person() class a function that tells them to wear a face mask:
def wearMask(self):   self.contagiousness /= 2
Yep, just halving their contagiousness is they wear a mask. Update initiateSim() so we can ask the user for what date the masks should come into use:
maskDay = int(input("Day for masks to be used: "))return daysContagious, lockdownDay, maskDay
And update our call:
daysContagious, lockdownDay, maskDay = initiateSim()
Finally, we’ll edit the section where we cycle through the days so that if the day reaches maskDay, then we tell every person to run their wearMask() function:
if x == maskDay:   for person in peopleDictionary:      person.wearMask()
If only it was this easy in real life, right?
Well what do you know, we’ve created a simple pandemic simulation with the ability to simulate each individual person, change attributes of the virus, enforce lockdowns, and make people wear face masks. Let’s look at our results:
I’m putting all the data gathered from my text save files into Excel.
5000 people, 1 starting infecter, 1% starting immunity, 7 days contagious, no lockdown or masks:
As expected, a nice smooth curve — almost mathematically perfect. By the end of the simulation, every has gained immunity and the cases drop to a 0, which continues until all the days have completed.
Now let’s see what happens to the previous result when you implement some countermeasures:
Now what we have here is really interesting. Take the blue line. This is the simulation without any countermeasures, just like our previous result. However, when we implement a lockdown on day 15, it has a huge effect on the orange line; the spread of the disease is curbed before it can really take off, and look at that gradual curve back down again — that’s where there are no new cases and people are gradually becoming immune!
We can then compare that to the gray line, where we implement lockdown just 5 days later than orange. It has a drastically lower effect because that five-day delay really made a difference to the number of cases.
Finally, take a look at the yellow line. This is where we implement face masks, and it’s probably the most interesting simulation of all. You can see at day 15, there is a sudden change in the gradient of the line which affects how fast the disease spreads. It probably would have increased much more rapidly without the face masks! About day 21, there is a peak, and thanks to the masks, it is substantially less than the blue line, where there were no countermeasures! There is also a tiny secondary peak, and the overall summit of the curve lasts longer than any other simulation. Can you figure out why?
Just to clarify, this was supposed to be a simple simulation. It is, of course, very basic with very limited parameters and functionality. However, it is incredible to see how much we can learn from a simulation that takes up barely a hundred lines of code. It really puts into perspective the impact lockdowns and face masks had.
I encourage anyone reading this with a programming mindset to go out and improve my code. I’d recommend the following features:
Face masks randomly (Gaussian?) affect contagiousness
Not everyone obeys lockdown, and even for those who do, there is a chance of an infection happening, say, during a grocery shopping trip
A certain percentage of people wear face masks, and this varies on a day to day basis
More social dynamics, or parameters in general.
The idea of communities.
If anyone does take on the challenge of upgrading this code, I’d love to see what results you get from playing around with the factors. Thanks for reading!
Full code:
from scipy.stats import normimport randomimport timepeopleDictionary = []#simulation of a single personclass Person():    def __init__(self, startingImmunity):        if random.randint(0,100)<startingImmunity:            self.immunity = True        else:            self.immunity = False        self.contagiousness = 0        self.mask = False        self.contagiousDays = 0        #use gaussian distribution for number of friends; average is 5 friends        self.friends = int((norm.rvs(size=1,loc=0.5,scale=0.15)[0]*10).round(0))    def wearMask(self):        self.contagiousness /= 2        def initiateSim():    numPeople = int(input("Population: "))    startingImmunity = int(input("Percentage of people with natural immunity: "))    startingInfecters = int(input("How many people will be infectious at t=0: "))    for x in range(0,numPeople):        peopleDictionary.append(Person(startingImmunity))    for x in range(0,startingInfecters):        peopleDictionary[random.randint(0,len(peopleDictionary)-1)].contagiousness = int((norm.rvs(size=1,loc=0.5,scale=0.15)[0]*10).round(0)*10)    daysContagious = int(input("How many days contagious: "))    lockdownDay = int(input("Day for lockdown to be enforced: "))    maskDay = int(input("Day for masks to be used: "))    return daysContagious, lockdownDay, maskDaydef runDay(daysContagious, lockdown):    #this section simulates the spread, so it only operates on contagious people, thus:    for person in [person for person in peopleDictionary if person.contagiousness>0 and person.friends>0]:        peopleCouldMeetToday = int(person.friends/2)        if peopleCouldMeetToday > 0:            peopleMetToday = random.randint(0,peopleCouldMeetToday)        else:            peopleMetToday = 0                    if lockdown == True:            peopleMetToday= 0                    for x in range(0,peopleMetToday):            friendInQuestion = peopleDictionary[random.randint(0,len(peopleDictionary)-1)]            if random.randint(0,100)<person.contagiousness and friendInQuestion.contagiousness == 0 and friendInQuestion.immunity==False:                friendInQuestion.contagiousness = int((norm.rvs(size=1,loc=0.5,scale=0.15)[0]*10).round(0)*10)                print(peopleDictionary.index(person), " >>> ", peopleDictionary.index(friendInQuestion))                for person in [person for person in peopleDictionary if person.contagiousness>0]:        person.contagiousDays += 1        if person.contagiousDays > daysContagious:            person.immunity = True            person.contagiousness = 0            print("|||", peopleDictionary.index(person), " |||")            lockdown = FalsedaysContagious, lockdownDay, maskDay = initiateSim()saveFile = open("pandemicsave3.txt", "a")for x in range(0,100):    if x==lockdownDay:        lockdown = True            if x == maskDay:        for person in peopleDictionary:            person.wearMask()                print("DAY ", x)    runDay(daysContagious,lockdown)    write = str(len([person for person in peopleDictionary if person.contagiousness>0])) + "\n"    saveFile.write(write)    print(len([person for person in peopleDictionary if person.contagiousness>0]), " people are contagious on this day.")saveFile.close()
I hope you found this entertaining and possibly inspiring! There are so many ways that you can improve this model, so I encourage you to see what you can build and see if you can simulate real-life even closer.
As always, I wish you the best in your endeavors!
Not sure what to read next? I’ve picked another article for you:
towardsdatascience.com
If you enjoyed this, follow me on Medium for more
Sign up for my email list here!
Let’s connect on LinkedIn
Interested in collaborating? Check out my website.
Note from the editors: Towards Data Science is a Medium publication primarily based on the study of data science and machine learning. We are not health professionals or epidemiologists, and the opinions of this article should not be interpreted as professional advice. See our Reader Terms for details. To learn about the coronavirus pandemic, you can click here.