The python code simulates an environment of an X number of persons with an Y number of places they can go to, puts the persons randomly in the places and calculates how many persons get infected, die or recover. It mostly work with lists, if you can optimize it in any way, thanks in advance!

import random
import functools
from time import perf_counter

with open("results.txt", "w") as results:
    time1 = perf_counter()

def simulation():
    infected, non_infected = 10, 15999990
    infectation_chance_c, death_chance, recovery_chance, reinfectation_chance, incubation_time = 1.4, 1 - 0.03, 1 - 0.97, 1 - 1 / 150, 2
    death_chance, recovery_chance = death_chance / incubation_time, recovery_chance / incubation_time
    population_total, population_list = infected + non_infected, non_infected * [0] + infected * [1]
    place = 120000
    day = 1
    simulation_duration = 100000000
    with open("results.txt", "a") as results:
        print("Starting... \nPlease wait for results, this can take lots of time!")
        while infected > 0  and simulation_duration > 0:
            population = population_list.count(0) + population_list.count(-1) + population_list.count(1)
            healthy = population_list.count(0) + population_list.count(-1)
            recovered = population_list.count(-1)
            infected = population_list.count(1)
            died = population_total - len(population_list)
            p = {i: [] for i in range(1,place + 1)}
            results.write(f"Day {day}: Infected: {infected} Healthy: {healthy} p-Imune: {recovered} Alive: {population} Died: {died} \n")
            print(f"Day {day}: Infected: {infected} Healthy: {healthy} p-Imune: {recovered} Alive: {population} Died: {died}")
            for person in population_list:
                p[random.randint(1, place)].append(person)
            i = 0
            while i < place:
                i += 1
                p_infected = p[i].count(1)
                    infectation_chance =  1 - float(p_infected) / (float(len(p[i])) / infectation_chance_c)
                for j, crowd in enumerate(p[i]):
                    if crowd == -1:
                        if random.random() > reinfectation_chance:
                            p[i][j] = 1
                        elif random.random() > reinfectation_chance:
                            p[i][j] = 0
                    elif crowd:
                        if random.random() > death_chance:
                        elif random.random() > recovery_chance:
                            if random.random() > 0.4:
                                p[i][j] = -1
                                p[i][j] = 0
                    elif not crowd:
                        if random.random()>infectation_chance:
                            p[i][j] = 1

            i = 0
            population_list = []
            while i < place:
                i += 1
            simulation_duration -= 1
            day += 1
        return time1
print(f"Simulation finishsed... \nProcessing time: {perf_counter()-time1}")
  • 1
    \$\begingroup\$ persons get infected, die, recover, die - so like... zombies? \$\endgroup\$
    – Reinderien
    Sep 25, 2021 at 14:40
  • \$\begingroup\$ @Reinderien Sorry accidentally added one more "die", no it's about a virus, people can recover and be sort of immune for a period of time after recovered or die. We are talking about humans \$\endgroup\$
    – R0Best
    Sep 25, 2021 at 15:09

1 Answer 1


Readability ,Data Types, Potential Bugs

  • The code is quite hard to follow. We should break it up into functions and let's add a if __main__==..

  • I'm renaming a couple of variable like place to rooms

  • Multiple assignments with such long variables makes it quite hard to read the value for each variable. Furthermore if they're constant for the simulations we can declare them as global constants.

  • I'm not going to be using it later but for future reference : p = {i: [] for i in range(1,place + 1)} can be replaced with a collections.defaultdict

  • Since you're only ever need the number of infected/healthy/recovered individuals and not the actual objects, in your case integers. We could easily replace the representation with something much shorter like a dict: population = dict(healthy=n_healthy,infected=n_infected etc..)

  • Express the probabilities correctly: deach_chance = 0.97 and if random.random > death_chance : he dies means the actual death chance is 0.03

  • BUG Popping an item by index during a list iteration changes the index for all following iterations don't do it. I'm quite certain this doesn't do what you intended to do for the following iterations: if random.random() > death_chance: p[i].pop(j)

The algorithm

The infected and the recovered

What happens to the infected and the recovered is independent of the room their in, we can treat all of them in one go. And you don't need to draw random variables for each of them. You can use binomial variables. If you have 100 infected and the death chance is 10% you can simply draw the number of deaths from np.random.binomial(100,0.1). For both the infected and the recovered we therefore draw from succession of binomial variables and return their new states.

The healthy

There probably is some mathematical trickery to be done. I'm certain a probability distribution can be found for the number of heathy people who become infected as a function of all the other variables. But I'm too tired to solve the integrals so I simply kept the exact same logic except in numpy. The performance is acceptable IMO.

Final code

import random
from collections import Counter
import numpy as np
from time import perf_counter


def simulate_n_infected_people(n):
    dead = np.random.binomial(n, DEATH_CHANCE)
    recovered_and_healty = np.random.binomial(n - dead, RECOVERY_CHANCE)
    healthy = np.random.binomial(recovered_and_healty, 0.4)
    recovered = recovered_and_healty - healthy
    infected = n - dead - recovered_and_healty
    return Counter(healthy=healthy, infected=infected, recovered=recovered)

def simulate_n_recovered_people(n):
    infected = np.random.binomial(n, REINFECTION_CHANCE)
    healthy = np.random.binomial(n - infected, REINFECTION_CHANCE)
    recovered = n - infected - healthy
    return Counter(healthy=healthy, infected=infected, recovered=recovered)

def simulate_n_healthy_people(population, n_rooms):
    rooms_array = np.zeros((n_rooms, 3), dtype=int)

    for i, key in enumerate(["infected", "recovered", "healthy"]):
        rooms_array[:, i] = np.bincount(
            np.random.randint(0, n_rooms, size=population[key]), minlength=n_rooms

    local_infection_chances = (
        rooms_array[:, 0] / np.sum(rooms_array, axis=1) / INFECTION_CHANCE
    infected = np.sum(np.random.binomial(rooms_array[:, 2], local_infection_chances))
    return Counter(
        healthy=population["healthy"] - infected, infected=infected, recovered=0

def single_day_simulation(population, n_rooms) -> dict:
    next_days_populations = Counter(healthy=0, infected=0, recovered=0)

    next_days_populations += simulate_n_infected_people(
    )  # independent from rooms
    next_days_populations += simulate_n_recovered_people(
    )  # independent from rooms
    next_days_populations += simulate_n_healthy_people(population, n_rooms)

    return next_days_populations

def simulation():
    infected, non_infected = 10, 15999990
    population_total = infected + non_infected
    population = Counter(healthy=15999990, infected=10, recovered=0)
    n_rooms = 120000
    day = 1
    simulation_duration = 100000000
    print("Starting... \nPlease wait for results, this can take lots of time!")
    while infected > 0 and day <= simulation_duration:
        alive_population = sum(population.values())
        healthy = population["healthy"]
        recovered = population["recovered"]
        infected = population["infected"]
        total_dead = population_total - alive_population
            f"Day {day}: Infected: {infected} Healthy: {healthy} p-Imune: {recovered} Alive: {alive_population} Died: {total_dead}",
        population = single_day_simulation(population, n_rooms)
        day += 1
    return time1

if __name__ == "__main__":
    time1 = perf_counter()
    print(f"Simulation finishsed... \nProcessing time: {perf_counter()-time1}")


  • The original code took about 10 seconds per day, had to kill it.
  • The updated code finished running at Day 507, Time : 95s
  • \$\begingroup\$ There seems to be a problem, I sometimes get the following error: RuntimeWarning: invalid value encountered in true_divide Traceback (most recent call last): simulation() population = single_day_simulation(population, n_rooms) next_days_populations += simulate_n_healthy_people(population, n_rooms) infected = np.sum(np.random.binomial(rooms_array[:, 2], local_infection_chances)) File "mtrand.pyx", line 3368, in numpy.random.mtrand.RandomState.binomial File "_common.pyx", line 338, in numpy.random._common.check_array_constraint ValueError: p < 0, p > 1 or p contains NaNs \$\endgroup\$
    – R0Best
    Oct 24, 2021 at 14:23
  • \$\begingroup\$ The warning is not a problem. As for the error, it's weird that i don't get it, might be a different version you have. You might have to fill nan values in the local_infection_chances array to zeros. \$\endgroup\$
    – kubatucka
    Oct 26, 2021 at 15:13

Your Answer

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

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