indonesia_infection_simulation.py

# run this with python 3
# no external dependencies needed

import random

# initial number of case
initial_case = 2
# number of carrier
carrier = 2
# number of infected
infected = 2
# incubation period (middle number between 7-14 days)
period = 14
# number of people infected by each carrier
infect_power = 10
# number of people infected but with good immunity
infected_drop_rate = 6
carrier_drop_rate = 6
# Days until carrier no longer able to infect
max_carrier_infection_period = 15
# end of simulation date
max_days = 30
current_day = 1
carrier_data = []

def get_drop_rate(max_range):
    rate = random.randint(1,max_range)
    return float(rate/10)
def initial_carrier_node(carrier_amount):
    return dict(amount=carrier_amount, current_day=1)

def increment_carrier_node(carrier_node):
    if carrier_node['current_day'] < max_carrier_infection_period:
        carrier_node['current_day'] += 1
    return carrier_node

def calculate_additional_infected(drop_rate):
    new_infected = 0
    if len(carrier_data) > 0:
        for idx, data in enumerate(carrier_data):
            if data.get('current_day') > max_carrier_infection_period:
                del carrier_data[idx]
            if data.get('current_day') == period:
                new_case = data.get('amount') - int(data.get('amount') * drop_rate)
                data['amount'] -= new_case
                data['current_day'] += 1
                new_infected += new_case
                carrier_data[idx] = data
            elif data.get('current_day') < period:
                updated_node = increment_carrier_node(data)
                carrier_data[idx] = updated_node
    return new_infected

def get_infecting_carrier(drop_rate):
    total_carrier = 1
    if len(carrier_data) > 0:
        for idx, data in enumerate(carrier_data):
            if data.get('current_day') == max_carrier_infection_period:
                del carrier_data[idx]
            elif data.get('current_day') < max_carrier_infection_period:
                total_carrier += data.get('amount') - int(data.get('amount') * drop_rate)
    return total_carrier

while(current_day <= max_days):
    current_carrier = get_infecting_carrier(get_drop_rate(carrier_drop_rate))
    new_carrier = current_carrier * infect_power
    carrier += new_carrier
    data = initial_carrier_node(new_carrier)
    new_infected = calculate_additional_infected(get_drop_rate(infected_drop_rate))
    infected += new_infected
    carrier_data.append(data)
    print("Day {} -  Total Infected: {:,.2f} New Infection: {:,.2f}".format(current_day, infected, new_infected))
    current_day += 1