Agents simulation COVID

sim.py

import numpy as np
from numpy.random import uniform, randint
from typing import Tuple, List, Set

class Agent:
    def __init__(self, x: float, y: float, contamined: bool):
        self.__x: float = x
        self.__y: float = y
        self.__contamined: bool = False
        self.__days_contamined: int = 0
        self.__dead: bool = False
        self.__recovered: bool = False

        if contamined:
            self.contamine()
    
    def contamine(self):
        self.__contamined = True
        self.__days_contamined = 1

    def is_contamined(self) -> bool:
        return self.__contamined

    def is_dead(self) -> bool:
        return self.__dead

    def __recover(self):
        self.__contamined = False
        # self.__days_contamined = 0
        self.__recovered = True

    def advance_sickness():
        self.__days_contamined += 1
        should_die: int = randint(0, 100)
        if should_die <= 1:  # Die with 1% chance
            self.__dead = True
            return
        if self.__days_contamined > 7:
            should_recover: int = randint(0, 100)
            if should_recover <= 20:  # Recover with 20% chance if sick since at least 7 days
                self.__recover()



class Grid:
    def __init__(self, shape: Tuple[int, int], n_agents: int = 1000):
        self.__plane: np.ndarray = np.zeros(shape)
        self.__n_agents: int = n_agents
        w, h = shape
        randoms = zip(uniform(0, w, n_agents), uniform(0, h, n_agents), randint(0, 100, n_agents))
        self.__agents: List[Agent] = [Agent(x, y, not contamined) for x, y, contamined in randoms]


    def run(self, days: int = 100):
        for i in range(days):
            infected_today: Set[Agent] = set()

            for agent in self.__agents:
                if agent not in infected_today and not agent.is_dead() and agent.is_contamined():
                    agent.advance_sickness()
                    # TODO: Contamine people around him with a probability N (use np.random)
                    # To determine people around, you can use a kd tree datastructure
                    pass