Simulation of a spreading virus

vsp.py

"""Simple simulation of a virus spreading,
using numpy matrices for internal board representation,
and pillow to generate a colored gif visualization.

    pip install numpy Pillow
    python vsp.py

Outputs: vsp.gif

A cell of the 2D world is either None or an integer >= -2.
See is_* functions for the meaning of each value.
In short, contaminating a human is made by making it equal to 2 or more.
Each step it will decrease by 1, 1 meaning death.

"""

import random
import itertools
import numpy as np
from PIL import Image

SIZE = 200
HUMAN_DENSITY = 0.65         # initial ratio of human/empty space
IMMUNE_LIKELIHOOD = 0.01     # ratio of humans discovered naturally immune
CARRIER_LIKELIHOOD = 0.01    # ratio of humans discovered healthy carrier
VACCINE_COVER = 0.15         # initial vaccine cover in population
VIRUS_CONTAGIOUSNESS = 0.4   # probability to be sick when next to a carrier (cumulative)
VIRUS_INCUBATION_TIME = 2    # time during which a human is sick before dying or recovering
VIRUS_MORTALITY_RATE = 0.99  # probability that the virus kills at the end of incubation
VIRUS_REMISSION_RATE = 0.01  # probability that a human recover at each step of incubation
VIRUS_SPREAD_FROM_DEAD = True  # virus is transmitted by dead humans too


STATE_IMMUNE = -1
STATE_HEALTHY_CARRIER = -2
STATE_DEAD = 1
STATE_HEALTHY = 0
def is_human(cell:int) -> bool:
    "True if given cell represents an human"
    return cell is not None
def is_healthy(cell:int) -> bool:
    "True if given cell represents an healthy (non-carrier) human"
    return cell == STATE_HEALTHY
def is_immune(cell:int) -> bool:
    "True if given cell represents an immune human"
    return cell == STATE_IMMUNE
def is_healthy_carrier(cell:int) -> bool:
    "True if given cell represents a healthy carrier human"
    return cell == STATE_HEALTHY_CARRIER
def is_dead(cell:int) -> bool:
    "True if given cell represents a dead human"
    return cell == STATE_DEAD
def is_sick(cell:int) -> bool:
    "True if given cell represents a sick human (not dead)"
    return cell is not None and cell > 1


def next_board(board):
    """Return a new matrix, the next state of given board"""
    return np.matrix([
        [next_state(board, row, col) for col in range(board[row].size)]
        for row in range(board[0].size)
    ])

def next_state(board, row, col) -> int or None:
    """Return the next state to be found at given place"""
    cell_state = board[row, col]
    if cell_state is None:
        return None  # nothing to do
    elif is_healthy(cell_state):  # maybe it will be infected
        nb_sick_neighbors = sum(
            1 for nei in neighbors(board, row, col)
            if is_sick(nei) or is_healthy_carrier(nei) or (VIRUS_SPREAD_FROM_DEAD and is_dead(nei))
        )
        infection_likelihood = nb_sick_neighbors * VIRUS_CONTAGIOUSNESS
        if random.random() < infection_likelihood:  # that's bad news
            if random.random() < CARRIER_LIKELIHOOD:
                return STATE_HEALTHY_CARRIER  # sorry everyone
            if random.random() < IMMUNE_LIKELIHOOD:
                return STATE_IMMUNE  # hold the do… virus !
            return 1 + VIRUS_INCUBATION_TIME  # nope, it's just sick
        return STATE_HEALTHY  # no problem !
    elif is_sick(cell_state):
        cell_state -= 1
        if cell_state == 1:  # just reach the final stage
            return STATE_DEAD if random.random() < VIRUS_MORTALITY_RATE else STATE_IMMUNE
        # their is a chance that the human just recover
        return STATE_IMMUNE if random.random() < VIRUS_REMISSION_RATE else cell_state
    else:
        assert is_immune(cell_state) or is_healthy_carrier(cell_state) or is_dead(cell_state), cell_state
        return cell_state

def neighbors(board, row, col) -> tuple:
    """Yield neighbors of given row,col position in board"""
    dimrow = board.size // len(board)
    dimcol = lambda row: board[row].size
    for drow, dcol in itertools.product((-1, 0, 1), repeat=2):
        if drow or dcol:  # discard the (0, 0) case
            x = (row + drow) % dimrow
            y = (col + dcol) % dimcol(x)
            yield board[x, y]


def image_of(board):
    """Build image from numpy array"""
    board = board_as_rgb(board)
    board = board.astype('uint8')
    return Image.fromarray(board, mode='RGB')

def board_as_rgb(board:np.array) -> np.array:
    """Return given numpy.array with RGB values instead of integer ones"""
    return np.array([
        [cell_as_rgb(board[row, col]) for col in range(board[row].size)]
        for row in range(board[0].size)
    ])
def cell_as_rgb(cell_state:int) -> (int, int, int):
    if cell_state is None:  # no human here
        return (0, 0, 0)
    elif is_healthy(cell_state):  # white
        return (205, 205, 205)
    elif is_immune(cell_state):
        return (100, 255, 100)  # white
    elif is_sick(cell_state):  # the near to death, the yellower
        # level = VIRUS_INCUBATION_TIMEcell_state * (255 / VIRUS_INCUBATION_TIME)
        level = 255 / (2 + VIRUS_INCUBATION_TIME - cell_state)
        return (level, level, 0)
    elif is_healthy_carrier(cell_state):
        return (255, 165, 0)  # orange
    elif is_dead(cell_state):
        return (255, 0, 0)  # red
    else:
        assert False, "can't be"

def make_gif(files, target:str):
    """Build gif from given grayscales images"""
    first, *lasts = files
    first.save(target, save_all=True, append_images=lasts, duration=10)


def vsp(board, steps:int=200):
    """Run the game of life on given board"""
    images = []
    try:
        for step in range(steps):
            print(f'\r    \r{step}', end='', flush=True)
            write_stats(board)
            new_board = next_board(board)
            if (board == new_board).all(): break  # we obtained the same configuration, nothing changed
            board = new_board
            images.append(image_of(board))
    except KeyboardInterrupt:
        pass
    make_gif(images, 'vsp.gif')

def write_stats(board):
    nb_dead = (board == STATE_DEAD).sum()
    nb_carrier = (board == STATE_HEALTHY_CARRIER).sum()
    nb_immune = (board == STATE_IMMUNE).sum()
    nb_healthy = (board == STATE_HEALTHY).sum()
    nb_alive = nb_healthy + nb_carrier + nb_immune
    print(f'\tTOTAL ALIVE: {nb_alive}\t\tDEATH: {nb_dead}\t\tIMMUNE/CARRIER/HEALTHY: {nb_immune}/{nb_carrier}/{nb_healthy}')

# build the world
board = np.random.choice([STATE_HEALTHY, STATE_IMMUNE, None], size=(SIZE, SIZE), p=[(HUMAN_DENSITY * (1-VACCINE_COVER)), (HUMAN_DENSITY * VACCINE_COVER), 1-HUMAN_DENSITY])
# inoculate the virus at the middle
board[SIZE//2][SIZE//2] = 1+VIRUS_INCUBATION_TIME
# run the simulation
vsp(board)