"""
Description: The goal is to model a population of people who are infected with a virus.
"""
# TODO: Make a curve for each case so we can see in real time how much the virus is spreading.
# -------------------- IMPORTS --------------------
from random import randint, random
from matplotlib import pyplot as plt, animation as anim
import math
# -------------------- GLOBAL VARIABLES --------------------
number_of_dots = 100 # number of dots to generate
shape = "o" # tip: use '.' instead if you put a value < 3 in minimal_distance
HEIGHT_WIDTH = 100 # Window height and width (yes, window shape must be a square)
BORDER_MIN = 1 # minimum distance from the border
BORDER_MAX = HEIGHT_WIDTH - 1 # maximum distance from the border
minimal_distance = 3 # Minimal distance at initialization and for contamination
time = 0 # Time to initialize
time_step = 0.1 # Time step for the simulation
transmission_rate = 0.7 # Chance of a dot to be infected
time_to_cure = 40 # Time to cure a dot
time_before_being_contagious_again = 40 # Time before being contagious again
virus_mortality = 0.0005 # Chance of a dot to die per tick
# -------------------- CLASSES & METHODS --------------------
class Dot:
def __init__(self, x: int, y: int):
"""Constructor for the Dot class
Args:
x (int): abscissa of the dot
y (int): ordinate of the dot
"""
self.x = x
self.y = y
self.velx = (random() - 0.5) / 5
self.vely = (random() - 0.5) / 5
self.is_infected = False
self.infected_at = -1
self.has_been_infected = False
self.cured_at = -1
def init_checker(x: int, y: int, already_used_coords: list):
"""Checks if the dot is in a distance of a minimal_distance from another dot
Args:
x (int): absissa of the dot
y (int): ordinate of the dot
already_used_coords (list): list of already occupated coordinates (by initialized dots)
Returns:
boolean: Whether the Dot should be initialized or not
"""
for coord in already_used_coords:
if Dot.get_distance(coord[0], x, coord[1], y) < minimal_distance:
return False
return True
def get_distance(x1: float, y1: float, x2: float, y2: float):
"""Gets the distance between two Dot objects
Args:
x1 (float): abscissa of the first dot
y1 (float): ordinate of the first dot
x2 (float): abscissa of the second dot
y2 (float): ordinate of the second dot
Returns:
float: distance between the two dots
"""
return math.sqrt((x2 - x1) ** 2 (y2 - y1) ** 2)
def initalize_multiple_dots():
"""Generates a list of Dots
Returns:
list: initialized dots
"""
dots = []
already_used_coords = []
while len(dots) != number_of_dots:
randx = randint(BORDER_MIN, BORDER_MAX)
randy = randint(BORDER_MIN, BORDER_MAX)
# So the dots keep distances between each other
if Dot.init_checker(randx, randy, already_used_coords):
dot = Dot(randx, randy)
already_used_coords.append((randx, randy))
else:
continue
dots.append(dot)
print("There are", len(dots), "dots in the area.")
return dots
def move(self):
"""Moves the dot and makes sure they don't go out of the area or touch each other. They've 4% chance to change direction."""
global dots, dead_dots
if random() < 0.96:
self.x = self.x self.velx
self.y = self.y self.vely
else:
self.x = self.x self.velx
self.y = self.y self.vely
# Change 2 to lower value to make the dots go faster
self.velx = (random() - 0.5) / (2 / (time_step 1))
# Change 2 to lower value to make the dots go faster
self.vely = (random() - 0.5) / (2 / (time_step 1))
if self.x >= BORDER_MAX:
self.x = BORDER_MAX
self.velx = -1 * self.velx
if self.x <= BORDER_MIN:
self.x = BORDER_MIN
self.velx = -1 * self.velx
if self.y >= BORDER_MAX:
self.y = BORDER_MAX
self.vely = -1 * self.vely
if self.y <= BORDER_MIN:
self.y = BORDER_MIN
self.vely = -1 * self.vely
if (
random() < transmission_rate
and not self.has_been_infected
and not self.is_infected
):
for dot in dots:
if (
dot.is_infected
and Dot.get_distance(self.x, self.y, dot.x, dot.y)
< minimal_distance
):
self.is_infected = True
self.infected_at = time
break
if self.is_infected and random() < virus_mortality:
dead_dots.append(self)
dots.remove(self)
def move_all(dots: list):
"""Moves a given list of dots. Make sur that infected dots goes in the correct axes
Args:
dots (list): List of Dot objects
"""
global healthy_dots, infected_dots, cured_dots, time
for dot in dots:
dot.move()
if (
dot.is_infected
and dot.infected_at != -1
and dot.infected_at time_to_cure < time
):
dot.is_infected = False
dot.has_been_infected = True
dot.cured_at = time
if (
dot.has_been_infected
and dot.cured_at != -1
and dot.cured_at time_before_being_contagious_again < time
):
dot.has_been_infected = False
dot.infected_at = -1
dot.cured_at = -1
healthy_dots.set_data(
[
dot.x
for dot in dots
if not dot.is_infected and not dot.has_been_infected
],
[
dot.y
for dot in dots
if not dot.is_infected and not dot.has_been_infected
],
)
infected_dots.set_data(
[dot.x for dot in dots if dot.is_infected],
[dot.y for dot in dots if dot.is_infected],
)
cured_dots.set_data(
[dot.x for dot in dots if dot.has_been_infected],
[dot.y for dot in dots if dot.has_been_infected],
)
time = time_step
plt.title(
f"Healthy: {len([dot for dot in dots if not dot.is_infected and not dot.has_been_infected])}"
f" | Infected: {len([dot for dot in dots if dot.is_infected])}"
f" | Cured: {len([dot for dot in dots if dot.has_been_infected])}"
f" | Dead: {len([dot for dot in dead_dots])}",
color="black",
)
def updateAxes():
"""Updates axes of the second window"""
# -------------------- MAIN FUNCTION --------------------
def main():
global dots, dead_dots, axes
dots = Dot.initalize_multiple_dots()
random_infected = randint(0, len(dots) - 1)
dots[random_infected].is_infected = True
dots[random_infected].infected_at = time
figureDots = plt.figure(facecolor="white")
axes = plt.axes(xlim=(0, HEIGHT_WIDTH), ylim=(0, HEIGHT_WIDTH))
global healthy_dots
healthy_dots = axes.plot(
[dot.x for dot in dots if not dot.is_infected],
[dot.y for dot in dots if not dot.is_infected],
f"g{shape}",
)[0]
global infected_dots
infected_dots = axes.plot(
[dot.x for dot in dots if dot.is_infected],
[dot.y for dot in dots if dot.is_infected],
f"r{shape}",
)[0]
global cured_dots
cured_dots = axes.plot(
[dot.x for dot in dots if dot.has_been_infected],
[dot.y for dot in dots if dot.has_been_infected],
f"b{shape}",
)[0]
global dead_dots
dead_dots = []
# We need to keep this in an unsed variable, otherwise the function won't work
a = anim.FuncAnimation(figureDots, lambda z: Dot.move_all(dots), frames=60, interval=5)
plt.axis('off')
plt.show()
# -------------------- MAIN CALL --------------------
if __name__ == "__main__":
main()
I'm trying to model a population with a virus and I would like to add another window where we can see in real time a graph of each type of the population (infected, healthy, dead, cured).
CodePudding user response：
Correct me if I misunderstood your question but I think you're looking for subfigures: https://matplotlib.org/devdocs/gallery/subplots_axes_and_figures/subfigures.html
You can create subfigures like so:
import matplotlib.pyplot as plt
fig = plt.figure()
ax1 = fig.add_subplot(211)
ax2 = fig.add_subplot(212)
There are multiple ways of adding subfigures, this is how I usually do it, the linked example above has another method.
The subfigures are added in the format [row column subfig_number] where subfig_number goes from the top left subfigure row-wise from left to right to the bottom right figure.