5
\$\begingroup\$

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')
        self.ax.add_patch(canvas)

    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
                dead_person = kill_table[i][j]
                if dead_person:
                    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')
                    self.ax.add_patch(square)
                if infected_person and not dead_person:
                    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')
                    self.ax.add_patch(square)
        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 = []
        self.dead = False

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

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

    def kill(self):
        self.dead = True
        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
            self.add_person()

        else:
            self.persons = persons
        self.details_instance = details_instance
        self.virus_strain = virus_strain
        self.dead_persons = [[]]

    def add_person(self):
        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

    def get_dead_table(self):
        truth_table = []
        current_list = []
        i = 0
        while i < self.count:
            current_list.append(self.persons[i].dead)
            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()
                self.dead_persons.append(index)

    def add_neighbors(self):
        #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

    def spread_infection(self,infected_table):
        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

    def get_dead_count(self):
        dead_people = 0
        for person in self.persons:
            if person.dead:
                dead_people+=1
        return dead_people

    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):
    pop.add_person()
pop.add_neighbors()
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()
current_dead = pop.get_dead_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()
    dead_people = pop.get_dead_count()
    print('The iteration we are on is %d with %d infected' %(iter,infected_people))
    simul_display.update_plot(current_infected,current_dead)

    pop.spread_infection(current_infected)
    current_infected=pop.get_infected_table()
    pop.kill_infected(current_infected)
    current_dead = pop.get_dead_table()

    if infected_people+dead_people > number_people:
        print('All individuals are infected or dead!')
        break

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

\$\endgroup\$
4
\$\begingroup\$

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]
convention = google

[pylama:mccabe]
max-complexity = 2

[pydocstyle]
convention = google

[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()
    dead_people = pop.get_dead_count()
    print('The iteration we are on is %d with %d infected' %(iter,infected_people))
    simul_display.update_plot(current_infected,current_dead)


    pop.spread_infection(current_infected)
    current_infected=pop.get_infected_table()
    pop.kill_infected(current_infected)
    current_dead = pop.get_dead_table()

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
    def can_spread(self) -> bool:
        """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:
        """Returns a new, dead 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
    def dead_count(self) -> int:
        """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)

spread the 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:
                    # dead or already infected
                    continue

                people_copy[x][y] = person.virus.spread(neighbour)
    return people_copy

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

Instead of

if person.alive:
    # spread the virus

I do:

if not person.alive:
    continue
# spread the virus

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
            people_copy[i][j] = virus.advance_infection(people_copy[i][j])
    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
            people_copy[i][j] = virus.advance_infection(people_copy[i][j])
    return people_copy

Population.advance

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

def advance(self) -> Population:
    """Advances the population 1 tick.

    1. Spread the virus
    2. Kill some of the infected

    This returns a new Population
    """
    people_post_spread = _spread(self.people)
    people_post_deaths = _kill(
        original_people=self.people, people_post_spread=people_post_spread
    )

    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)

    # print(population.dead_count, population.infected_count)
    assert population.dead_count == 0
    assert population.infected_count == 1

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

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

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
    2 = not infected, dead
    3 = infected, dead
    """
    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...

| improve this answer | |
\$\endgroup\$
  • \$\begingroup\$ 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? \$\endgroup\$ – Bob Jeans Apr 2 at 18:31
  • 2
    \$\begingroup\$ 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__ \$\endgroup\$ – Maarten Fabré Apr 3 at 7:07
  • \$\begingroup\$ I added the code for plotting and fixed a bug in the _spread \$\endgroup\$ – Maarten Fabré Apr 3 at 7:07
  • \$\begingroup\$ This is a very nice & complete answer! \$\endgroup\$ – Grajdeanu Alex. Apr 3 at 7:33

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.