2
\$\begingroup\$

I'm trying to implement the dynamics of Dengue transmission between mosquitoes and humans. I decided create two classes, one for the human population and another for the mosquitoes. Both are based on an abstract class. I tried to optimize the different methods of each class by vectorizing them with numpy matrices. It works okay. I'm want it to be as readable, maintainable and extendable as possible.

"""
File: population.py
Project: dengueSim
File Created: Tuesday, 3rd October 2023 5:43:09 pm
Author: Athansya
-----
License: MIT License
-----
Description: Population abstract class and its subclasses to simulate
human and mosquito populations dynamics during dengue infections.
"""

from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from icecream import ic
import numpy as np
import pygame

# TODO Define scene to simulate (e.g. room, street, park)
# TODO Define time scale (e.g. minutes, hours, days, etc)
# TODO Define valid states given previous decisions
# TODO Define movement in scene for each agent (e.g. commuting, walking, flying)
# TODO Save states and counters in a file for plotting
# TODO Change Pygame framework to Manim for better visualization.

@dataclass(slots=True)
class Population(ABC):
    """
    Abstract base population class. Handles positions, velocities and states of
    the population. Includes abstract methods to be overriden by subclasses.
    
    Attributes
    ----------
        size : int
            Population size.
        max_position_x : float
            Maximum x coordinate.
        max_position_y : float
            Maximum y coordinate.
        max_velocity : float
            Maximum velocity.
        states : dict[str, int]
            Dictionary of states.
        states_color_map : dict[int, str]
            Dictionary of colors.
        _positions_matrix : np.ndarray
            Matrix of positions.
        _velocities_matrix np.ndarray
            Matrix of velocities.
        _states_matrix : np.ndarray
            Matrix of states.
    
    """
    size: int = 10
    max_position_x: float = 500
    max_position_y: float = 500
    max_velocity: float = 10
    states: dict[str, int] = field(
        default_factory=lambda: {"susceptible": 0, "infected": 1, "recovered": 2}
    )
    states_color_map: dict[int, str] = field(
        default_factory=lambda: {0: "blue", 1: "red", 2: "green"}
    )

    _positions_matrix: np.ndarray = field(init=False)
    _velocities_matrix: np.ndarray = field(init=False)
    _states_matrix: np.ndarray = field(init=False)

    def __post_init__(self) -> None:
        """
        Initialize positions, velocities and states matrices.

        Returns
        -------
            None

        """ 
        self._init_positions()
        self._init_velocities()
        self._init_states()

    def _init_positions(self) -> None:
        """
        PRIVATE METHOD
        Initialize a matrix array of size (N, 2) where N is the 
        population size and 2 the x and y coordinates. Fills it
        with random values from a uniform distribution with boundaries
        between 0 and max_positions.

        Returns
        -------
            None

        """
        self._positions_matrix = np.zeros(shape=(self.size, 2))

        self._positions_matrix[:, 0] = np.random.uniform(
            0, self.max_position_x, size=self.size
        )
        self._positions_matrix[:, 1] = np.random.uniform(
            0, self.max_position_y, size=self.size
        )

    def _init_velocities(self) -> None:
        """
        PRIVATE METHOD
        Initialize a matrix array of size (N, 2) where N is the 
        population size and 2 the vx and vy velocity components. 
        Fills it with random values from a uniform distribution
        based on the max_velocity.

        Returns
        -------
            None

        """
        self._velocities_matrix = np.random.uniform(
            -self.max_velocity, self.max_velocity, size=(self.size, 2)
        )

    def _init_states(self) -> None:
        """
        PRIVATE METHOD
        Initialize a matrix array of size (N, 1) where N is the 
        population size and 1 the state. Assigns a random state,
        excluding the "recovered" state.

        Returns
        -------
            None
            
        """
        self._states_matrix = np.random.choice(
            list(self.states.values())[:-1], size=(self.size, 1)
        )

    def _handle_borders(self) -> None:
        """
        PRIVATE METHOD
        Check if positions are out of bounds and replace them.

        Returns
        -------
            None

        """
        # X-coordinate
        self._positions_matrix[:, 0] = np.where(
            self._positions_matrix[:, 0] > self.max_position_x,  # Exceeded width
            0,  # Value to replace it with
            self._positions_matrix[:, 0],  # else keep the same
        )

        self._positions_matrix[:, 0] = np.where(
            self._positions_matrix[:, 0] < 0,  # Less than 0
            self.max_position_x,  # Value to replace it with
            self._positions_matrix[:, 0],  # else keep the same
        )
        # Y-coordinate
        self._positions_matrix[:, 1] = np.where(
            self._positions_matrix[:, 1] > self.max_position_y,  # Exceeded height
            0,  # Value to replace it with
            self._positions_matrix[:, 1],  # else keep the same
        )

        self._positions_matrix[:, 1] = np.where(
            self._positions_matrix[:, 1] < 0,  # Less than 0
            self.max_position_y,  # Value to replace it with
            self._positions_matrix[:, 1],  # else keep the same
        )

    def draw(self, screen: pygame.Surface, radius: int | float) -> None:
        """
        Draws the population as circles on the screen.

        Parameters
        ----------
            screen : pygame.Surface
                Window to draw on. 
            radius : int | float
                Radius of the circles.

        Returns
        -------
            None

        """
        for index, agent_pos in enumerate(self._positions_matrix):
            pygame_vector_pos = pygame.Vector2(agent_pos[0], agent_pos[1])
            pygame.draw.circle(
                screen,
                color=self.states_color_map[self._states_matrix[index][0]],
                center=pygame_vector_pos,
                radius=radius,
            )

    @abstractmethod
    def move(self) -> None:
        """
        ABSTRACT METHOD
        Move the agents. MUST be overriden by subclasses.

        Returns
        -------
            None
        """
        pass

    @abstractmethod
    def update_velocity(self) -> None:
        """
        ABSTRACT METHOD
        Update agents velocities. MUST be overriden by subclasses.

        Returns
        -------
            None

        """
        # TODO CHECK IF IT IS NECESSARY OR DELETE IT!
        pass


@dataclass
class HumanPopulation(Population):
    """
    Human agents population class. Inherits from Population class. Adds additional
    properties, methods and overrides the move and update_velocity abstract methods.

    Attributes 
    ----------
    time_to_recover : float = 50
        Defines how much time must pass before a human can recover from infection.
    time_to_susceptible : float = 50
        Defines how much time must pass before a human can become susceptible.
        In other words, it how much time a human remains immune to infection.
    _infection_timer_matrix : np.ndarray = field(init=False)
        Timer that keeps track of how long each human has been infected.
    _recover_timer_matrix : np.ndarray = field(init=False)
        Timer that keeps track of how long each human has been recovered from infection.

    """
    # TODO CHECK IF FOLLOWING ATTRIBUTES ARE NECESSARY OR DELETE THEM
    # Infected rate per infected vector
    # B_i: float = 0
    # Infection rate per infected host
    # Beta_i: float = 0
    # Susceptible birth rate
    # mu: float = 1 / 65
    # Recovery rate
    # gamma: float = 365 / 7
    # Temporary cross-immunity rate
    # alpha: float = 2

    time_to_recover: float = 50
    time_to_susceptible: float = 50
    _infection_timer_matrix: np.ndarray = field(init=False)
    _recover_timer_matrix: np.ndarray = field(init=False)
    

    def _init_states(self) -> None:
        """
        Initializes the states matrix and its related timers.

        Returns
        -------
            None

        """
        # Every human is susceptible at the start
        self._states_matrix = np.zeros(shape=(self.size, 1))
        self._infection_timer_matrix = np.zeros(shape=(self.size, 1))
        self._recover_timer_matrix = np.zeros(shape=(self.size, 1))

    def move(self, random: bool = False) -> None:
        """
        Moves the humans, updating their positions.

        Parameters
        ----------
        random : bool, optional
            Whether to move randomly or not, by default False.

        Returns
        -------
            None

        """
        if random:
            self._init_velocities()
        self._positions_matrix += self._velocities_matrix
        self._handle_borders()

    def update_velocity(self) -> None:
        pass

    def _time_infected(self) -> None:
        """
        PRIVATE METHOD
        Updates the timers that keep track of how long each human has been infected.

        Returns
        -------
            None

        """
        infected_humans = self._states_matrix == self.states["infected"]
        if np.any(infected_humans):
            self._infection_timer_matrix[infected_humans] += 1

    def recover(self) -> None:
        """
        Checks which humans are ready to recover from infection and changes their state.

        Returns
        -------
            None

        """
        self._time_infected()
        humans_ready_to_recover = self._infection_timer_matrix >= self.time_to_recover
        if np.any(humans_ready_to_recover):
            self._states_matrix[humans_ready_to_recover] = self.states["recovered"]
            self._infection_timer_matrix[humans_ready_to_recover] = 0

    def _time_recovered(self) -> None:
        """
        PRIVATE METHOD
        Updates the timers that keep track of how long each human has been recovered (immune).

        Returns
        -------
            None

        """
        infected_humans = self._states_matrix == self.states["recovered"]
        if np.any(infected_humans):
            self._recover_timer_matrix[infected_humans] += 1

    def make_susceptible(self) -> None:
        """
        Checks which humans are ready to be suscpetible to infection and changes their state.

        Returns
        -------
            None

        """
        self._time_recovered()
        humans_ready_to_susceptible = self._recover_timer_matrix >= self.time_to_susceptible
        if np.any(humans_ready_to_susceptible):
            self._states_matrix[humans_ready_to_susceptible] = self.states["susceptible"]
            self._recover_timer_matrix[humans_ready_to_susceptible] = 0


@dataclass
class MosquitoPopulation(Population):
    """
    Mosquito agents population class. Inherits from Population class. Adds additional
    properties, methods and overrides the move and update_velocity abstract methods.

    Attributes 
    ----------
    bite_radius : float = 1
        Mosquito bite radius area.
    bite_probability : float = 0.5
        Probability of biting a human.
    transmission_probability : float = 0.5
        Probability of Dengue transmission from a mosquito to a human

    """
    # TODO CHECK IF FOLLOWING ATTRIBUTES ARE NECESSARY OR DELETE THEM
    # Susceptible birth rate
    # upsilon: float = 36.5
    # Infection rate per host
    # var_theta: float = 73
    # Magnitude of sinusoidal fluctuations
    # eta: float = 0  # or 0.35
    # Ratio of likelihood of transmission from hosts with
    # secondary and hosts primary infection to vectors
    # var_phi: float = 0  # or 12
    # Phase
    # phi: float = 0

    bite_radius: float = 1
    bite_probability: float = 0.5
    transmission_probability: float = 0.5

    def move(self, random: bool = False) -> None:
        """
        Moves the mosquitos

        Parameters
        ----------
        random : bool, optional
            _description_, by default False
        """
        if random:
            self._init_velocities()
        self._positions_matrix += self._velocities_matrix
        self._handle_borders()

    def update_velocity(self) -> None:
        pass

    def bite_humans(self, human_population: HumanPopulation) -> None:
        """
        For each mosquito in the population, check if it is close to a human.
        If it is, checks the probability of biting and infecting the human. 

        Parameters
        ----------
        human_population : HumanPopulation
            Human agents population.

        Returns
        -------
            None

        """
        # For each mosquito check if it is close to a human
        for mosquito_pos in self._positions_matrix:
            distances = np.linalg.norm(
                human_population._positions_matrix - mosquito_pos, axis=1
            )
            close_to_human = distances < self.bite_radius
            susceptible_humans = human_population._states_matrix == self.states["susceptible"]
            susceptible_humans = susceptible_humans.reshape(len(susceptible_humans))
            # If there are close humans, update their state
            if np.any(close_to_human):
                # Bite and infect with certain probability. Note that events are dependent and
                # the probability of an event A and B happening is the product of the probabilities
                probabilities_array = np.random.random(
                    size=close_to_human.shape[0]
                ) <= (self.bite_probability * self.transmission_probability)
                # Only changes humans state if they are close to the mosquito AND the transmission probability is met
                # AND the human is susceptible
                human_population._states_matrix[
                    close_to_human
                    & probabilities_array
                    & susceptible_humans
                ] = self.states["infected"]

If you want to test it here is a simple code to run the simulation:

from population import HumanPopulation, MosquitoPopulation
import pygame

WIDTH = 500 
HEIGHT = 500
FPS = 60

MOSQUITO_RADIUS = 1
HUMAN_RADIUS = 5
# Host population
N = 100
# Vector population
M = 1000


def main():
    # States dict
    states = {"susceptible": 0, "infected": 1, "recovered": 2}
    color_map = {
        states["susceptible"]: "blue",
        states["infected"]: "red",
        states["recovered"]: "green",
    }

    mosquitoes = MosquitoPopulation(
        size=M,
        max_position_x=WIDTH,
        max_position_y=HEIGHT,
        max_velocity=50,
        states=states,
        states_color_map=color_map
    )

    humans = HumanPopulation(
        size=N,
        max_position_x=WIDTH,
        max_position_y=HEIGHT,
        max_velocity=10,
        states=states,
        states_color_map=color_map
    )

    pygame.init()
    screen = pygame.display.set_mode((WIDTH, HEIGHT))
    clock = pygame.time.Clock()
    running = True

    # Time to move or not
    timer_to_move = 0

    while running:
        screen.fill((255, 255, 255))
        # poll for events
        # pygame.QUIT event means the user clicked X to close your window
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False

        # Mosquito population
        mosquitoes.draw(screen, MOSQUITO_RADIUS)
        # Move mosquitos every certain time
        if timer_to_move % 10 == 0:
            mosquitoes.move(random=True)
        mosquitoes.bite_humans(humans)

        # Human population
        humans.draw(screen, HUMAN_RADIUS)
        if timer_to_move % 10 == 0:
            humans.move(random=True)
        else:
            humans.move()
        humans.recover()
        humans.make_susceptible()

        pygame.display.flip()
        clock.tick(FPS)
        timer_to_move += 1


if __name__ == "__main__":
    main()
\$\endgroup\$

1 Answer 1

2
\$\begingroup\$

use Enum

This is a bit odd:

    states = {"susceptible": 0, "infected": 1, "recovered": 2}

Normally we'd expect to see an Enum in that situation. Having a separate color_map seems tedious, when color could just be an Enum attribute.

In SIR modeling you might prefer to substitute "removed" for "recovered". That clarifies we're describing individuals who are deceased or, through activity of the immune system, are now removed from the "susceptible" population.


use type hints

class Population(ABC):

Thank you for the lovely docstring. The English sentences were helpful.

But then you go into a long "Attributes" section. It's nice enough. But I encourage you to delete it. All the identifiers are wonderful and informative. The docstring does not further illuminate them. You mention their types, but you do it in a way that mypy cannot access them. Much better to delete them and add corresponding type annotations. Then humans would read them just the same, but significantly, humans would believe them, secure in the knowledge that mypy had verified them.

Oh, wait, now I see down in the code that you gave the type of everything twice. I imagine they're all in sync? But I'm not going to scan back and forth to check on that.

Comments are nice. Machine verified comments are much nicer.

If I see same thing in two ways, I will just assume that one of them is out of sync, because sooner or later that's sure to be the case.

Several methods are annotated def ...() -> None:, in which case "Returns NONE" in a docstring just isn't helpful. Couldn't sphinx infer that from annotations?


conventional naming

    def _init_positions(self) -> None:
        """
        PRIVATE METHOD

Yeah, I got that loud and clear from the function name, it's _private. You don't have to tell me twice.

        self._positions_matrix = np.zeros(shape=(self.size, 2))

        self._positions_matrix[:, 0] = np.random.uniform(
            0, self.max_position_x, size=self.size
        )
        self._positions_matrix[:, 1] = np.random.uniform(
            0, self.max_position_y, size=self.size
        )

Sooo, this is lovely, very clear, thank you.

It calls into question the "max_position_{x,y}" decision. You might have been happier with a "max_position" vector of 2-tuples, as the assignments here would be more compact.

Ooohh, lookitdat, it appears _init_velocities adopted that very approach. Good.


clip

The _handle_borders method is lovely, it does exactly what you'd expect.

We import numpy, and it turns out that is rather a large library. If there's a math concept you can imagine finding on wikipedia, there's a fair chance you can find it in the numpy documentation, as well. Here, you might prefer to use the clip function.


float annotations

    def draw(self, screen: pygame.Surface, radius: int | float) -> None:

Thank you for the optional type annotations, they are helpful.

nit: It would be enough to declare radius: float.

Why?

https://peps.python.org/pep-0484

this PEP proposes a straightforward shortcut ... : when an argument is annotated as having type float, an argument of type int is acceptable

Linters such as mypy definitely follow this advice.

            pygame_vector_pos = ...

I understand where the name for that local variable came from, but it's a trifle verbose -- pos would suffice.


no pass

    @abstractmethod
    def move(self) -> None:
        """
        ABSTRACT METHOD
        Move the agents. MUST be overriden by subclasses.

        Returns
        -------
            None
        """
        pass

nit: We write pass where it is syntactically necessary. But here, a """docstring""" is present, so no need for pass.

Plus the usual remark about @abstractmethod and "ABSTRACT METHOD" saying the same thing. Similarly for ) -> None: and "Returns None".

Oh, wait, here's trouble!

signal error upon failure to override

... MUST be overriden by subclasses.

And then we permissively pass. No. Don't do that. It would be much more helpful to raise a fatal error explaining that subclass neglected to implement an important override.


DRY

class HumanPopulation(Population):
    """
    Human agents population class. Inherits from Population class.

Yup, we got that loud and clear from the ...(Population): part, no need to say it again in English.


discretize

The scaling on this looks ugly as both populations increase:

    def bite_humans(self, human_population: HumanPopulation) -> None:
        ...
        # For each mosquito check if it is close to a human
        for mosquito_pos in self._positions_matrix:
            distances = np.linalg.norm(
                human_population._positions_matrix - mosquito_pos, axis=1
            )
            close_to_human = distances < self.bite_radius

Consider discretizing to some grid size, and narrow the checks based on whether insect + human are within the same grid cell.


This codebase would benefit from the addition of a robust test suite plus code coverage measurements.

The author was clearly trying to produce maintainable code. As features are added in the coming months, I am confident that the codebase will continue to be of high quality and easily approached by new maintenance engineers that are recruited to the team.

This code appears to achieve its design goals.

I would be willing to delegate or accept maintenance tasks on this codebase.

\$\endgroup\$
2
  • \$\begingroup\$ Thank you for the review! I've implemented most of your suggestions. I'm just wrapping my head around the idea of using Enum. I haven't use it before... I forgot to mention that I want a Torus-like behavior, so np.clip won't work I think. Regarding the discretization of my grid, I already have a couple ideas and questions, but I think those would be better discussed in StackOverflow if I'm not mistaken. I haven't added tests yet, but plan to do. I'm not versed in code coverage measurements, so I'm a bit lost there. Once again, thank you for your help. \$\endgroup\$
    – Athansya
    Oct 16, 2023 at 4:25
  • 1
    \$\begingroup\$ Cool. Implementing "torus" is much easier than worrying about four different boundaries. Just use % modulo and you're done! \$\endgroup\$
    – J_H
    Oct 16, 2023 at 5:34

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

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