The python code simulates an environment of a disease and calculates how many people get infected

Question

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:
    results.seek(0)
    results.write("")
    results.close()
    time1 = perf_counter()

@functools.lru_cache(maxsize=128)
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)
                try:
                    infectation_chance =  1 - float(p_infected) / (float(len(p[i])) / infectation_chance_c)
                except:
                    pass
                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:
                            p[i].pop(j)
                        elif random.random() > recovery_chance:
                            if random.random() > 0.4:
                                p[i][j] = -1
                            else:
                                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
                population_list.extend(p[i])
            simulation_duration -= 1
            day += 1
        return time1
simulation()
print(f"Simulation finishsed... \nProcessing time: {perf_counter()-time1}")

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

INCUBATION_TIME = 2
INFECTION_CHANCE = 1.4
DEATH_CHANCE = 1 - (1 - 0.03) / INCUBATION_TIME
RECOVERY_CHANCE = 1 - (1 - 0.97) / INCUBATION_TIME
REINFECTION_CHANCE =  1 / 150


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(
        population["infected"]
    )  # independent from rooms
    next_days_populations += simulate_n_recovered_people(
        population["recovered"]
    )  # 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
        print(
            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()
    simulation()
    print(f"Simulation finishsed... \nProcessing time: {perf_counter()-time1}")

Speed-up:

• The original code took about 10 seconds per day, had to kill it.
• The updated code finished running at Day 507, Time : 95s