Epidemic simulation COVID-19

[pirogone]

Hi all, I'd like to share with you a bit of code, for those of you who start to look deeper into epidemic simulations.

It is not the best piece of code ever written (as I initially learned coding on a TI-83+), anyway it works just fine and could probably easily be improved and optimized.

Feel free to modify, improve or just play with it.

I can answer questions about the code if needed.

Don't forget to "pip install matplotlib" if not done yet..

Best regards, Pierre.

#from random import *
import matplotlib.pyplot as plt
from tkinter import *
from tkinter import messagebox
from matplotlib.lines import Line2D



def start_simulation(_): #Function that will run the simulation after the start button or enter is pressed.
    font = {'size': 8}
    plt.rc('font', **font)
    plt.figure(1, figsize=(10, 6))
    plt.xlabel('Time (in days)')
    plt.ylabel('Population')

    #Setting variables :
    simulation_length = int(simulation_length_field.get())
    increment_1=1
    infected_total = 0
    immunized = 0
    desimmunized = 0
    death_total = 0
    icu_in_the_same_time = 0
    icu_total = 0
    contagiousness = float(contagiousness_field.get()) / 100
    number_of_contact_per_day = float(contact_per_day_field.get())
    healthy = int(population_size_field.get())
    infected = int(initially_infected_field.get())
    time_before_immunization = int(disease_duration_field.get())
    death_rate = float(death_rate_field.get()) / 100
    icu_rate = float(icu_rate_field.get()) / 100
    mean_icu_stay_length = int(icu_duration_field.get())
    table_immunization = simulation_length * [0]
    table_icu = time_before_immunization * [0]
    death_multiplier = int(death_multiplier_v.get())
    icu_multiplier = int(icu_multiplier_v.get())
    display_healthy = int(display_healthy_v.get())
    immunization_length = int(immunisation_duration_field.get())
    display_death = int(display_death_v.get())
    display_icu = int(diplay_icu_v.get())
    dynamic_graph = int(dynamic_display_v.get())
    isolate_infected = int(infected_detection_v.get())
    time_before_detection = int(sick_days_before_detection_v.get())
    table_isolated = time_before_detection * [0]
    isolation_rate = int(proportion_isolated_v.get()) / 100
    isolated_total = 0
    healthy = healthy - infected
    infected_initially=infected
    healthy_initially = healthy
    table_infected = [infected]
    growth_rate = 0
    death_before_reducing_contact=int(death_before_quarantine_v.get())
    reduce_contact = int(quarantine_v.get())

    # find and fix errors before starting graph :
    if time_before_detection > time_before_immunization and isolate_infected == 1:
        time_before_detection = time_before_immunization
        messagebox.showinfo('Error', 'Delay before infected detection must be lower than the disease duration. Detection time has been set to ' + str(time_before_immunization) + '.')

    # enable static view if ticked :
    if dynamic_graph==0 :
        plt.axis([0, simulation_length, 0, healthy])

    # enable/disable death and/or ICU multiplier, to have a better overview when plotted :
    if death_multiplier == 1:
        death_multiply = 10
    else:
        death_multiply = 1

    if icu_multiplier ==1:
        icu_multiply = 10
    else:
        icu_multiply = 1

    # allow line 132 to run properly if isolation of infected is disabled :
    if isolate_infected == 0:
        table_isolated = simulation_length * [0]

    # Start of the main loop, which calculate data for each day :
    for actual_day in range(0,simulation_length):
        # infected_per_day=0  #reset the count of infected for each new day (only used if slow calculation is used)
        increment_1 = increment_1+1  # allows line 182 to 187 to run properly, avoid unwanted early epidemic termination.

        # reduces the number of contact per day to 1, when the total number of death is superior to the chosen limit:
        if death_total > death_before_reducing_contact and reduce_contact==1:
            number_of_contact_per_day = 1

        # This secondary loop is deactivated by default, it could be interesting if you plan on doing simulation of small populations (<50 000), to have a unique simulation at each run
        # It calculate each day and for each subject if they are infected or not by running a random function
        # But it is way too slow for big populations
        # If you decide to activate it, delete the triple quotes at the beginning and at the end of the paragraph and add some to the next paragraph, also remove the # before infected_per_day=0 (line 81) and before random import at the beginning.
        '''for subject in range(0,healthy):
            probability_of_being_contaminated = number_of_contact_per_day*contagiousness*infected/(healthy+infected+immunized)
            if random() <= probability_of_being_contaminated:
               infected_per_day = infected_per_day + 1
               healthy = healthy -1'''

        # Add triple quotes here if needed (cf line 91)
        # Calculation of the daily new infected :
        probability_of_being_contaminated = number_of_contact_per_day * contagiousness * infected / (
                healthy + infected + immunized)
        infected_per_day =  int((healthy * probability_of_being_contaminated))
        healthy = healthy - int((healthy * probability_of_being_contaminated)) # Add triple quotes here if needed (cf line 91)


        table_infected.append(infected_per_day) # Saving the infected of the day in a table in order to remove them after healing time has past

        # Increment the infected counters
        infected=infected + infected_per_day
        infected_total = infected_total + infected_per_day

        # Growth rate calculation & displaying title + growth rate on the graph :
        if table_infected[actual_day-1] != 0:
            growth_rate=table_infected[actual_day]/table_infected[actual_day-1]
        graph_title = 'Epidemic Simulation,  Growth factor = ' + str(round(growth_rate,2))
        plt.suptitle(graph_title)

        # Detection and isolation of the infected :
        if isolate_infected == 1 and actual_day >= time_before_detection:
            infected_isolated = int(isolation_rate * table_infected[actual_day - time_before_detection])
            table_infected[actual_day - time_before_detection] = table_infected[actual_day - time_before_detection] - infected_isolated
            infected = infected - infected_isolated
            table_isolated.append(infected_isolated)
            isolated_total = isolated_total + infected_isolated

        plt.pause(0.05) # For live display

        # Infected become immunized after the time_before_immunization, or go to ICU or die :
        if actual_day >= time_before_immunization:
            infected = infected - table_infected[(actual_day - time_before_immunization)]
            daily_death = death_rate * (table_infected[actual_day - time_before_immunization])
            daily_icu = icu_rate*(table_infected[actual_day - time_before_immunization])
            immunized = immunized + table_infected[(actual_day - time_before_immunization)] - daily_death - daily_icu

            #same job for the infected in quarantine (become immunized, go to ICU or die) :
            if isolate_infected == 1 and actual_day >= time_before_detection:
                isolated_icu = icu_rate * (table_isolated[actual_day - time_before_immunization])
                daily_icu = daily_icu +isolated_icu
                isolated_death = death_rate * (table_isolated[actual_day - time_before_immunization])
                daily_death = daily_death + isolated_death
                immunized = immunized + table_isolated[actual_day - time_before_immunization] - isolated_death - isolated_icu
                isolated_total = isolated_total - table_isolated[actual_day - time_before_immunization]

            # If the immunization is not permanent, immunized subject get back to the healthy state and can be infected again :
            if immunization_length !=0:
                table_immunization[actual_day] = (table_infected[(actual_day - time_before_immunization)] - daily_death - daily_icu) + (table_isolated[actual_day - time_before_immunization])
                desimmunized = round(table_immunization[actual_day - immunization_length])
            if actual_day >= immunization_length != 0:
                healthy = healthy + desimmunized
                immunized = immunized - desimmunized

            # ICU count, adding daily_ICU to ICU_table in order to be able to retrieve and switch ICU subject from ICU to Immunized, Death count :
            icu_in_the_same_time = icu_in_the_same_time + daily_icu
            icu_total = icu_total + daily_icu
            table_icu.append(daily_icu)
            death_total = death_total + daily_death

            # When the mean ICU stay length is gone, subject from ICU switch to Immunized subjects :
            if actual_day >= mean_icu_stay_length:
                icu_in_the_same_time = icu_in_the_same_time - table_icu[actual_day - mean_icu_stay_length]
                immunized = immunized + table_icu[actual_day - mean_icu_stay_length]

            #Display of the legend on the graph :
            infected_patch = Line2D([0], [0], marker='o', color='w', label='Actual infected : ' + str(round(infected)), markerfacecolor='red', markersize=7)
            infected_per_day_patch = Line2D([0], [0], marker='^', color='w', label='Infected per day : ' + str(round(infected_per_day)), markerfacecolor='red', markersize=7)
            infected_total_patch = Line2D([0], [0], marker='*', color='w',label='Inf total : ' + str(round(infected_total)) + ' (' + str(round(infected_total / (infected_initially + healthy_initially) * 100, 1)) + ' %)', markerfacecolor='r', markersize=11)
            immunized_patch = Line2D([0], [0], marker='o', color='w', label='Immunized : ' + str(round(immunized)) + ' (' + str(round(immunized / (infected_initially + healthy_initially) * 100, 1)) + ' %)',markerfacecolor='green', markersize=7)
            death_patch = Line2D([0], [0], marker='o', color='w', label='Actual deaths : ' + str(round(daily_death)), markerfacecolor='black', markersize=7)
            death_total_patch = Line2D([0], [0], marker='*', color='w', label='Deaths total : ' + str(round(death_total)) + ' (' + str(round(death_total / (infected_initially + healthy_initially) * 100, 1)) + ' %)', markerfacecolor='black', markersize=11)
            icu_patch = Line2D([0], [0], marker='o', color='w', label='Actual ICU : ' + str(round(daily_icu)), markerfacecolor='blue', markersize=7)
            icu_total_patch = Line2D([0], [0], marker='*', color='w', label='ICU total : ' + str(round(icu_total)) + ' (' + str(round(icu_total / (infected_initially + healthy_initially) * 100, 1)) + ' %)', markerfacecolor='blue', markersize=11)
            healthy_patch = Line2D([0], [0], marker='o', color='w', label='Healthy not immunized', markerfacecolor='purple', markersize=7)
            isolated_patch = Line2D([0], [0], marker='o', color='w', label='Actual isolated : ' + str(round(isolated_total)), markerfacecolor='orange', markersize=7)
            info_prct_patch = Line2D([0], [0], marker='o', color='w', label='(% are on total population)', markerfacecolor='white',markersize=7)

            #Display the healthy subject on the legend if ticked, else not :
            if display_healthy == 1:
                plt.legend(handles=[infected_patch, infected_per_day_patch, infected_total_patch, immunized_patch, icu_patch, icu_total_patch, death_patch, death_total_patch, healthy_patch, isolated_patch,info_prct_patch ])
            else:
                plt.legend(handles=[infected_patch, infected_per_day_patch, infected_total_patch, immunized_patch, icu_patch, icu_total_patch,death_patch, death_total_patch, isolated_patch,info_prct_patch])


        else: # If the time of immunization/death/ICU (time_before_immunization) is not yet passed, this legend is shown. It's because the calculations of immunized, death and ICU can start only after this delay.
            infected_patch = Line2D([0], [0], marker='o', color='w', label='Inf : ' + str(round(infected-infected_initially)), markerfacecolor='r', markersize=7)
            infected_per_day_patch = Line2D([0], [0], marker='^', color='w', label='Inf per day : ' + str(round(infected_per_day)), markerfacecolor='r', markersize=7)
            infected_total_patch = Line2D([0], [0], marker='*', color='w', label='Inf total : ' + str(round(infected_total)) +' ---> ' + str(round(infected_total / (infected_initially + healthy_initially)*100, 1)) + ' %', markerfacecolor='r', markersize=11)
            plt.legend(handles=[infected_patch, infected_per_day_patch, infected_total_patch])

        # If the epidemic is gone, there's no more cases, no more virus to spread, a pop-up window appears :
        if infected == 0 and increment_1 > 10:
            result_end_question = messagebox.askquestion('Epidemic is gone', 'Do you want to quit all?')
            if result_end_question == 'yes':
               quit()
            else:
               break

        # Plot the different curves on the graph, sometimes with conditions (if the checkbox is ticked or not..)
        plt.scatter(actual_day, infected, color='red', s=1)
        plt.scatter(actual_day, immunized, color='green', s=1)
        if display_death == 1:
            plt.scatter(actual_day, death_total * death_multiply, color='black', s=1)
        if display_icu == 1:
            plt.scatter(actual_day, icu_in_the_same_time * icu_multiply, color='blue', s=1)
        if isolate_infected == 1:
            plt.scatter(actual_day, isolated_total, color='orange', s=1)
        if display_healthy == 1:
            plt.scatter(actual_day, healthy, color='purple', s=1)

    plt.show()


# This is the first window of the program :
window = Tk()
window.title('EPIDEMIC SIMULATION')

label_1 = Label()
label_1.grid(column=0, rowspan=4) #this is just a space...

#Every left sliders are declared here :
contagiousness_v = StringVar()
contagiousness_v.set(5)
contagiousness_field = Scale(window, variable = contagiousness_v, orient='horizontal', from_=0, to=20,
                             resolution=0.1, tickinterval=0, length=350,
                             label='% Contagiousness')
contagiousness_field.grid(column=0, rowspan=4, padx=20)

contact_per_day_v = StringVar()
contact_per_day_v.set(4)
contact_per_day_field = Scale(window, variable = contact_per_day_v, orient='horizontal', from_=0, to=30,
                              resolution=1, tickinterval=0, length=350,
                              label='N contact per day')
contact_per_day_field.grid(column=0, rowspan=4)

population_size_v = StringVar()
population_size_v.set(1000000)
population_size_field = Scale(window, variable = population_size_v, orient='horizontal', from_=100000, to=10000000,
                              resolution=100000, tickinterval=0, length=800,
                              label='Population size')
population_size_field.grid(column=0, rowspan=4, columnspan=4, sticky='W', padx=20)

initially_infected_v = StringVar()
initially_infected_v.set(500)
initially_infected_field = Scale(window, variable = initially_infected_v, orient='horizontal', from_=1, to=10000,
                                 resolution=1, tickinterval=0, length=800, sliderlength=15,
                                 label='N initially infected')
initially_infected_field.grid(column=0, rowspan=4, columnspan=4, sticky='W', padx=20)

death_rate_v = StringVar()
death_rate_v.set(1)
death_rate_field = Scale(window, variable = death_rate_v, orient='horizontal', from_=0, to=100,
                         resolution=0.5, tickinterval=0, length=350,
                         label='Death rate in %, when infected')
death_rate_field.grid(column=0, rowspan=4)

icu_rate_v = StringVar()
icu_rate_v.set(4)
icu_rate_field = Scale(window, variable = icu_rate_v, orient='horizontal', from_=0, to=100,
                       resolution=0.5, tickinterval=0, length=350,
                       label='ICU rate in %, when infected')
icu_rate_field.grid(column=0, rowspan=4)

icu_duration_v = StringVar()
icu_duration_v.set(12)
icu_duration_field = Scale(window, variable = icu_duration_v, orient='horizontal', from_=1, to=20,
                           resolution=1, tickinterval=0, length=350,
                           label='ICU mean duration of stay, in days')
icu_duration_field.grid(column=0, rowspan=4)

simulation_length_v = StringVar()
simulation_length_v.set(100)
simulation_length_field = Scale(window, variable = simulation_length_v, orient='horizontal', from_=50, to=1000,
                                resolution=50, tickinterval=0, length=350,
                                label='Simulation length, in days')
simulation_length_field.grid(column=0, rowspan=4)

disease_duration_v = StringVar()
disease_duration_v.set(10)
disease_duration_field = Scale(window, variable = disease_duration_v, orient='horizontal', from_=1, to=30,
                               resolution=1, tickinterval=0, length=350,
                               label='Disease duration, in days (before dying or being immunized)')
disease_duration_field.grid(column=0, rowspan=4)

immunization_duration_v = StringVar()
immunization_duration_v.set(0)
immunisation_duration_field = Scale(window, variable = immunization_duration_v, orient='horizontal', from_=0, to=200,
                                    resolution=5, tickinterval=0, length=350,
                                    label='Immunization duration, in days (0 = forever)')
immunisation_duration_field.grid (column=0, rowspan=4)

label_2 = Label()
label_2.grid(column=0, rowspan=4) # Just a space..

# Every checkbox are declared here + Start button :
infected_detection_v = IntVar()
infected_detection_v.set(0)
infected_detection_checkbox = Checkbutton(window, text='Detect INF after ', variable=infected_detection_v)
infected_detection_checkbox.grid(column=2, row=7, sticky='W')

sick_days_before_detection_v = StringVar()
sick_days_before_detection_v.set(5)
sick_days_before_detection_field = Scale(window, variable = sick_days_before_detection_v, orient='horizontal',
                                         from_=1, to=30, resolution=1, tickinterval=0, length=35,
                                         sliderlength=5, width=3)
sick_days_before_detection_field.grid(column=2, row=7)

text_2 = Label(window, text ='days sick and isolate')
text_2.grid(column=2, row=7, sticky='E')

proportion_isolated_v=StringVar()
proportion_isolated_v.set(50)
proportion_isolated_field = Scale(window, variable = proportion_isolated_v, orient='horizontal', from_=1, to=100,
                                  resolution=5, tickinterval=0, length=60, sliderlength=10, width=3)
proportion_isolated_field.grid(column=3, row=7, sticky='W')

text_3 = Label(window, text ='% of them')
text_3.grid(column=3, row=7)

quarantine_v = IntVar()
quarantine_v.set(0)
quarantine_checkbox = Checkbutton(window, text='Reduce to 1 contact per day when reaching ', variable=quarantine_v)
quarantine_checkbox.grid(column=2, row=11, sticky='W')

death_before_quarantine_v=StringVar()
death_before_quarantine_v.set(100)
death_before_quarantine_field = Scale(window, variable = death_before_quarantine_v, orient='horizontal',
                                      from_=1, to=5000, resolution=50, tickinterval=0, length=120,
                                      sliderlength=10, width=3)
death_before_quarantine_field.grid(column=3, row=11, sticky='W')

text_1 = Label(window, text ='deaths')
text_1.grid(column=3, row=11, sticky='E', padx=20)

display_death_v = IntVar()
display_death_v.set(1)
display_death_checkbox = Checkbutton(window, text='Display Deaths', variable=display_death_v)
display_death_checkbox.grid(column=2, row=23, sticky='W')

death_multiplier_v = IntVar()
death_multiplier_v.set(1)
death_multiplier_checkbox = Checkbutton(window, text='Display Deaths x10', variable=death_multiplier_v)
death_multiplier_checkbox.grid(column=2, row=23, sticky='E')

diplay_icu_v = IntVar()
diplay_icu_v.set(1)
display_icu_checkbox = Checkbutton(window, text='Display ICU', variable=diplay_icu_v)
display_icu_checkbox.grid(column=2, row=27, sticky='W')

icu_multiplier_v = IntVar()
icu_multiplier_v.set(1)
icu_multiplier_checkbox = Checkbutton(window, text='Display ICU x10', variable=icu_multiplier_v)
icu_multiplier_checkbox.grid(column=2, row=27, sticky='E')

display_healthy_v = IntVar()
display_healthy_v.set(1)
display_healthy_checkbox = Checkbutton(window, text='Display Healthy persons', variable=display_healthy_v)
display_healthy_checkbox.grid(column=3, row = 39, sticky='W', padx=20)

dynamic_display_v = IntVar()
dynamic_display_v.set(0)
dynamic_display_checkbox = Checkbutton(window, text='Dynamic display', variable=dynamic_display_v)
dynamic_display_checkbox.grid(column=3, row=40, sticky='W', padx=20)

start_button = Button(window, text='Start simulation')
start_button.grid(column=3, row=43)
start_button.bind('<Button-1>', start_simulation)
window.bind('<Return>', start_simulation)

text_4 = Label(window, text ='@pirog, 2020')
text_4.grid(column=2, row=44)

window.mainloop()

[Anarab]

This is really cool good job man props I'll make sure to check it out.