# Simulation of virus growth

The gist of my program is to simulate the growth of some virus. I used (attempted) OOP concepts to break down the problem into chunks and have the chunks talk to each other.

Not sure if my implementation is effective but it seems to work pretty good. Currently, the bottleneck seems to be in the plotting. I'm still learning about matplotlib so I'm not surprised.

The program has five classes. The first class just keeps track of simulation details, nothing too fancy.

class Details():
def __init__(self,num_people_1d,fig_size):
self.num_people_1d = num_people_1d
self.total_people = self.num_people_1d**2

self.fig_size = fig_size

self.x_length = self.fig_size[0]/num_people_1d
self.y_length = self.fig_size[0]/num_people_1d


The second class is the display. At each iteration, the 2D grid is updated with details about individuals that are either infected or have perished. The display class is updated with the information

import matplotlib.pyplot as plt
import matplotlib.patches as patches

class Display():
def __init__(self,details_instance):
self.size_x,self.size_y = details_instance.fig_size
self.length_1d = details_instance.num_people_1d

self.x_length = details_instance.x_length
self.y_length = details_instance.y_length

def create_plot(self,plot_size = (5,5)):
self.fig = plt.figure(figsize = plot_size)
self.ax = self.fig.subplots()

canvas = patches.Rectangle((0,0),1,1,fill=True,
edgecolor='none',facecolor='g')

def update_plot(self,infected_table=None,kill_table = None):
#Transposing tables
infected_table = list(map(list, zip(*infected_table)))
kill_table = list(map(list, zip(*kill_table)))

for i,row in enumerate(infected_table):
for j,col in enumerate(row):
infected_person = col
coord = i*self.x_length,j*self.y_length
square = patches.Rectangle(coord,
self.x_length,self.y_length,
fill=True,
edgecolor = 'none',
facecolor = 'r')
coord = i*self.x_length,j*self.y_length
square = patches.Rectangle(coord,
self.x_length,self.y_length,
fill=True,
edgecolor = 'none',
facecolor = 'y')
plt.show()
plt.pause(0.1)


The next class is the virus class. Not too much going on, just the infect and mortality rate.

class Virus():

def __init__(self,infectionRate = 0.1,mortalityRate = 0.01):
self.IR = infectionRate
self.MR = mortalityRate


Then is the person class. This class just keeps some basic information. If the person is infected or dead, and some simple methods.

import random

class Person():

def __init__(self,id = None,discrete_location = None,infected = False):
self.id = id
if discrete_location:
self.dl_x,self.dl_y = discrete_location
else:
raise Exception()
self.infected = infected
self.neighbors = []

def become_infected(self,virus):
self.infected = True
self.virus = virus

def do_i_live(self):
return random.random()>self.virus.MR

def kill(self):
self.infected = False


The last class is the Population class. This class holds most of the code that actually does stuff because it is what updates all of the individuals.

from person import Person
import random

class Population():
def __init__(self,persons=[],details_instance =None,virus_strain=None):
if len(persons)<1:
print('There is no population! Adding a member')
self.persons = persons
self.count = 0

else:
self.persons = persons
self.details_instance = details_instance
self.virus_strain = virus_strain

if len(self.persons)<1:
self.persons.append(Person(id=self.count,
discrete_location = (0,0),
infected = False)
)
self.count +=1
else:
loc_x = self.details_instance.x_length*(self.count%self.details_instance.num_people_1d)
loc_y = self.details_instance.y_length*((self.count - self.count%self.details_instance.num_people_1d)/self.details_instance.num_people_1d)
person = Person(id = self.count,
discrete_location = (loc_x,loc_y),
infected = False)
self.count +=1
self.persons.append(person)

def get_infected_table(self):
truth_table = []
current_list = []
i = 0
while i < self.count:
current_list.append(self.persons[i].infected)
i+=1
if (i)%(self.details_instance.num_people_1d) ==0:
truth_table.append(current_list)
current_list = []
if self.count%(self.details_instance.num_people_1d) !=0:
truth_table.append(current_list)
return truth_table

truth_table = []
current_list = []
i = 0
while i < self.count:
i+=1
if (i)%(self.details_instance.num_people_1d) ==0:
truth_table.append(current_list)
current_list = []
if self.count%(self.details_instance.num_people_1d) !=0:
truth_table.append(current_list)

return truth_table

def kill_infected(self,infected_table):
linear_indices = self.get_infected_indices(infected_table)
for index in linear_indices:
still_living = self.persons[index].do_i_live()
if not still_living:
self.persons[index].kill()

#Currently returns the linear index! Compatible with persons!!
if len(self.persons)<=1:
return

#One method:
#Use self.count and modulos to identify neighbors
#Possibly a better method that I do not follow:
#Using discrete location to identify neighbors

#Using first method
for i in range(self.count):
#at left boundary
if i%self.details_instance.num_people_1d==0:
left = -1
else:
left = i-1
#at right boundary
if (i+1)%self.details_instance.num_people_1d==0:
right = -1
else:
right = i+1

up = i+self.details_instance.num_people_1d
down = i - self.details_instance.num_people_1d

#First build potential neighbors
potential_neighbors = [left,right,up,down]

#Second identify if any potential neighbors don't exist
neighbor_list = []
for j in potential_neighbors:
if (j >= 0) and (j<self.count):
neighbor_list.append(j)

#Third update the person with neighbors
self.persons[i].neighbors = neighbor_list

linear_indices = self.get_infected_indices(infected_table)
for index in linear_indices:
current_infected_person = self.persons[index]
neighbors = current_infected_person.neighbors
for neighbor in neighbors:
if random.random()<current_infected_person.virus.IR:
self.persons[neighbor].become_infected(self.virus_strain)

def get_infected_count(self):
infected_people = 0
for person in self.persons:
if person.infected:
infected_people+=1
return infected_people

for person in self.persons:

def get_infected_indices(self,infected_table):
#returns the linear indices of those infected
linear_indices=[]
for i,row in enumerate(infected_table):
for j,col in enumerate(row):
if col:
linear_indices.append(j+i*self.details_instance.num_people_1d)
return linear_indices


To run all of this code I wrote the following script:

from person import Person
from virus import Virus
from display import Display
from details import Details
from population import Population

import random
num_people_1d = 10

simul_details = Details(num_people_1d = num_people_1d,fig_size = (1,1))
virus_strain1 = Virus()
pop = Population(details_instance = simul_details,virus_strain=virus_strain1)

number_people = num_people_1d**2-1
for i in range(number_people):
starting_person = random.randint(0,number_people-1)
print('The starting person is %d' % starting_person)

pop.persons[starting_person].become_infected(virus_strain1)
current_infected = pop.get_infected_table()

simul_display = Display(details_instance=simul_details)
simul_display.create_plot()

total = 100
for iter in range(total):
infected_people = pop.get_infected_count()
print('The iteration we are on is %d with %d infected' %(iter,infected_people))

current_infected=pop.get_infected_table()
pop.kill_infected(current_infected)

print('All individuals are infected or dead!')
break


This is all of the code. Any comments would be gratefully appreciated.

# Files

Python is not Java, not every class needs it's own module. You can keep Details, Virus and Person in 1 file, Display can go in its own file, since that serves another purpose

# Programming tools

You can let the IDE help you a lot by using a few tools. I myself use black as code formatter, isort to sort the imports, pylama with the linters mccabe,pep8,pycodestyle,pyflakes to check the code quality, mypy for the static type analysis and py.test for unit tests. All these tools nicely integrate in most common python IDEs

This is my setup.cfg

[pylama]
linters = mccabe,pep8,pycodestyle,pyflakes,mypy,isort

[pylama:*/__init__.py]
ignore=W0611

[pylama:pydocstyle]

[pylama:mccabe]
max-complexity = 2

[pydocstyle]

[isort]
multi_line_output=3
include_trailing_comma=True
force_grid_wrap=0
use_parentheses=True
line_length=79

[mypy]
check_untyped_defs = true
disallow_any_generics = true
disallow_untyped_defs = true
ignore_missing_imports = true
no_implicit_optional = true
warn_redundant_casts = true
warn_return_any = true
warn_unused_ignores = true

[mypy-tests.*]
disallow_untyped_defs = false
check_untyped_defs = false

[tool:pytest]
junit_family = xunit2
testpaths = tests


My project directory looks like this

project_name/
- data/
- raw/
- processed/
- docs/
- build/
- source/
- notebooks/
- 20200402 analysis interference.ipynb
- ...
- output/
- analysis1/
- ...
- src/
- package_name/
- sub_module/
- __init__.py
- module1.py
- module2.py
- __init__.py
- module1.py
- module2.py
- ...
- tests/
- data/
- conftest.py
- test_feature1.py
- ...
-.gitignore
- requirements_dev.txt
- requirements.txt
- setup.cfg
- setup.py


# Documentation

Some docstring (PEP-257) would help users of your package to know what's going on. That includes you if you revisit this project a few months later.

# Type analysis

I use type annotations for 2 reasons. It serves as additional documentation, and helps the IDE spot bugs for you, especially combined with a strict mypy configuration.

# Argument defaults

You allow creating a person with a empty id, but do it nowhere. Why allow this feature? Then you would also not need to do this:

    if discrete_location:
self.dl_x,self.dl_y = discrete_location
else:
raise Exception()


That can be prevented as well by removing the default argument for discrete_location, and either removing the default for id as well or moving id to the back, or instructing python to get all arguments as keyword arguments.

# builtin sum

You can use sum and the fact that a boolean counts as 1:

def get_infected_count(self):
infected_people = 0
for person in self.persons:
if person.infected:
infected_people+=1
return infected_people


can become

def get_infected_count(self):
return sum(person.infected for person in self.persons)


# generators

def get_infected_indices(self,infected_table):
#returns the linear indices of those infected
linear_indices=[]
for i,row in enumerate(infected_table):
for j,col in enumerate(row):
if col:
linear_indices.append(j+i*self.details_instance.num_people_1d)
return linear_indices


can become a lot clearer as a generator

def get_infected_indices(self,infected_table):
"""The linear indices of those infected."""
for i, row in enumerate(infected_table):
for j, col in enumerate(row):
if col:
yield j+i*self.details_instance.num_people_1d


# more functional programming

Each iteration you follow these steps:

for iter in range(total):
infected_people = pop.get_infected_count()
print('The iteration we are on is %d with %d infected' %(iter,infected_people))

current_infected=pop.get_infected_table()
pop.kill_infected(current_infected)


You are mutating you population in-place, and to do so, you need to follow a complicated series of steps. A simpler option would be to have a Population.advance method that returns a new Population instance representing the state of the population. That way you can keep track of what happened, who died, ...

# dataclasses

These kind of classes lend themselves very good to be implemented using dataclasses

## Person

@dataclasses.dataclass(frozen=True)
class Person:
"""A Person."""

alive: bool = True
virus: typing.Optional[Virus] = None

@property
"""A person can spread the virus when he's alive and infected."""
return self.alive and self.virus is not None

def infect(self, virus: Virus) -> Person:
"""Returns a new, infected Person."""
return dataclasses.replace(self, virus=virus)

def die(self) -> Person:
return dataclasses.replace(self, alive=False)


Using an @property to check whether someone can spread the disease, and returning a new person when dying or getting infected instead of changing in-place.

In a later stage, allowing persons with multiple infections can be as simple as changing the virus to a set[Virus] and a small tweak to the infect method

## Virus

@dataclasses.dataclass(frozen=True)
class Virus:
"""A Virus."""

infection_rate: float
mortality_rate: float

def spread(self, subject: Person) -> Person:
"""Possibly infects the subject.

In this simple algorithm, it just picks a random number
in the range [0.0, 1.0)]

i this number is lower than the virus's infection rate,
the person gets inected
"""
dice_roll = random.random()
if dice_roll < self.infection_rate:
return subject.infect(self)
return subject

def advance_infection(self, subject: Person) -> Person:
"""Advance the virus infection in the subject.

If not infected, does nothing.
I infected, checks whether the subject dies.

In this simple algorithm, it just picks a random number
in the range [0.0, 1.0)]

i this number is lower than the virus's mortality rate,
the person dies
"""
dice_roll = random.random()
if dice_roll < self.mortality_rate:
return subject.die()
return subject


This is rather self-explanatory. Doing it like this lets you easily implement more sophisticated viruses with incubation_periods, ...

## Population

People = typing.List[typing.List["Person"]]  # for typing purposes

@dataclasses.dataclass(frozen=True)
class Population:
"""A Population."""

people: People
virus: Virus

@property
def infected_count(self) -> int:
"""Returns the number of alive people who have been infected."""
return sum(person.can_spread for person in self)

@property
"""Returns the number of dead people."""
return sum(not person.alive for person in self)

def __iter__(self) -> typing.Iterator[Person]:
"""Yield all the people in the population."""
return itertools.chain.from_iterable(self.people)

@property
def grid_size(self) -> typing.Tuple[int, int]:
"""The gridsize of the population."""
return len(self.people), len(self.people[0])


Defines a simple population. Instead of keeping one list of all the people in the population, we use a grid. This makes searching for the neighbours a lot simpler later on. Doing it like this also allows us to calculate the number of dead and infected on the fly, instead of having to keep track of that separately.

As a convenience, we supply a method to generate a pristine population:

    @classmethod
def new(cls, gridsize: int, virus: Virus) -> Population:
"""Generates a new Population of healthy people."""
return cls(
people=[
[Person() for _ in range(gridsize)] for _ in range(gridsize)
],
virus=virus,
)


To infect our initial person, we add an infect_person method:

    def infect_person(self, x: int, y: int) -> Population:
"""Infects the person a location x, y.

Returns a new Population.
"""
people_copy: People = [row[:] for row in self.people]
people_copy[x][y] = people_copy[x][y].infect(self.virus)
return Population(people=people_copy, virus=self.virus)


To spread the virus, I would use a helper method that operates on a grid of people. We iterate over the grid, looking for people who are alive and have the virus. Then looking in the cells around that person looking for people who can be infected.

def _spread(people: People) -> People:
"""Spread the disease in a population.

returns a new people
"""
rows = len(people)
columns = len(people[0])
people_copy: People = [row[:] for row in people]

person: Person
for i, row in enumerate(people):
for j, person in enumerate(row):
if not person.alive:
continue
if person.virus is None:
continue
for di, dj in [
(-1, 0),
(1, 0),
(0, -1),
(0, 1),
]:
# iterate over the neighbours
x, y = i + di, j + dj
if (not 0 <= x < rows) or not (0 <= y < columns):
# out of bounds
continue

neighbour = people[x][y]
if not neighbour.alive or neighbour.virus is person.virus:
continue

return people_copy


I use the technique of the negative check a few time.

if person.alive:


I do:

if not person.alive:
continue


This saves on a few levels of indentation, and makes the code easier to read.

# kill

the _kill helper method works much in the same way:

def _kill(people: People) -> People:
"""Kills a portion of the infected.

returns a new people
"""
people_copy: People = [row[:] for row in people]
person: Person
for i, row in enumerate(people):
for j, person in enumerate(row):
if not person.alive:
continue
if person.virus is None:
continue
virus = person.virus
return people_copy


If you want to give the people who have been infected just that tick a tick of repriece,you need to do something like this:

def _kill(original_people: People, people_post_spread: People) -> People:
"""Kills a portion of the infected of the previous tick.

returns a new people
"""
people_copy: People = [row[:] for row in people_post_spread]
person: Person
for i, row in enumerate(original_people):
for j, person in enumerate(row):
if not person.alive:
continue
if person.virus is None:
continue
virus = person.virus
return people_copy


# Population.advance

And now to the method why we did all this work has become very simple:

def advance(self) -> Population:

2. Kill some of the infected

This returns a new Population
"""
people_post_deaths = _kill(
)

return Population(people=people_post_deaths, virus=self.virus)


# using the simulation:

This simulation can be used very easily:

if __name__ == "__main__":
virus = Virus(infection_rate=.1, mortality_rate=.1)
population = Population.new(gridsize=10, virus=virus).infect_person(4, 4)

assert population.infected_count == 1

populations: typing.List[Population] = [population]
for i in range(1, 30):
populations.append(population)

print(
f"after {i} iterations: {population.infected_count} infected and "
)


And now you can use those populations to do analyses later, plotting, ...

## plotting

Now you have the population as a grid, you can convert this grid to a numpy array

def matrix(self) -> np.array:
"""Creates a numpy array of the grid.

A kind of bitmap
0 = fine
1 = infected, alive
"""
return np.array(
[
[
(person.virus is not None) + 2 * (not person.alive)
for person in row
]
for row in self.people
],
dtype="int8",
)


Then the plotting is as simple as

fig, ax = plt.subplots()
im = ax.imshow(population.matrix())
plt.show()


You can choose the colormap...

• This is exactly what I was looking for. Thank you! If you have the time, I do have some follow up questions. First, why use dataclasses over standard classes? Is it because the amount of detail in the class is minimal and the classes primarily store information (such as alive, dead, infected, etc.)? Second, I'm not sure I follow why the @property decorator is utilized. Third, why do you recommend creating new population instances, and creating new Persons, instead of an update-in place strategy? Fourth, do you thoughts about following an update-in place structure? – Bob Jeans Apr 2 at 18:31
• 1: dataclasses are just more convenient so you can skip some of the boiler plate. The frozen=True is a nice side effect if you want immutability; 2:the @property decorator is just a convenience as well.That way they act as variables. 3: the immutability allows you to prevent unwanted side effects. If you want in-place you need to make sure that what you do in 1 place doesn't affect another place. The only downside is a little bit more memory use, but even that is limited. I that really bothers you, you can use __slots__ – Maarten Fabré Apr 3 at 7:07
• I added the code for plotting and fixed a bug in the _spread – Maarten Fabré Apr 3 at 7:07
• This is a very nice & complete answer! – Grajdeanu Alex. Apr 3 at 7:33