text_html_output_1.png application_vnd.jupyter.stderr_output_1.png
import csv

import numpy as np

import numpy as np # linear algebra

import pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)

import scipy.optimize as optim

train_data = pd.read_csv('../input/covid19-global-forecasting-week-4/train.csv')
train_data = train_data.replace(np.nan, '', regex=True)
test_data = pd.read_csv('../input/covid19-global-forecasting-week-4/test.csv')
test_data = test_data.replace(np.nan, '', regex=True)

def filter_train_data(country, region):
    selector = train_data['Country_Region'] == country
    onlyonecountry = train_data[selector]
    selector2 = onlyonecountry['Province_State'] == region
    onlyoneregion = onlyonecountry[selector2]
    onlyoneregion = onlyoneregion.reset_index(drop=False)
    del onlyoneregion['index']
    del onlyoneregion['Id']
    onlyoneregion['Timestep'] = onlyoneregion.index
    return onlyoneregion

def infection_start(country, region):
    try:
        infection_start_df = pd.read_csv('../input/covid19-global-forecasting-week-4/train.csv')
        infection_start_df = infection_start_df.replace(np.nan, '', regex=True)
        selector = infection_start_df['Country_Region'] == country
        onlyonecountry = infection_start_df[selector]
        selector2 = onlyonecountry['Province_State'] == region
        onlyoneregion = onlyonecountry[selector2]
        return onlyoneregion['ConfirmedCases'].iloc[0]
    except:
        return 2

population_reader = csv.reader(open('../input/population/population.csv', 'r'))
population_dict = {}
next(population_reader)
for row in population_reader:
    k, v = row
    population_dict[k] = int(v)

pop_by_region_reader = csv.reader(open('../input/populationbycity/populationbycity.csv', 'r'))
populationbyregion_dict = {}
next(pop_by_region_reader)
for row in pop_by_region_reader:
    k, v = row
    populationbyregion_dict[k] = int(v)

def get_population(country, region):
    if region != '':
        if region in populationbyregion_dict:
            return populationbyregion_dict[region]
        else:
            return 1000000
    elif country != '':
        if country in population_dict:
            return population_dict[country]
        else:
            return 1000000
    else:
        return 1000000

def filter_test_data(country, region, date):
    selector = test_data['Country_Region'] == country
    onlyonecountry = test_data[selector]
    selector2 = onlyonecountry['Province_State'] == region
    onlyoneregion = onlyonecountry[selector2]
    onlyoneregion = onlyoneregion.reset_index(drop=False)
    del onlyoneregion['index']
    onlyoneregion['Timestep'] = onlyoneregion.index + 71
    if date != '':
        dateselect = test_data['Date'] == date
        onlyoneregion = onlyoneregion[dateselect]
    return onlyoneregion

def my_logistic(t, a, b, c):
    return c / (1 + a * np.exp(-b * t))

def calculate_infection(country_entry, region_entry, date_entry):
    train_set = filter_train_data(country_entry, region_entry)
    local_x = np.array(train_set['Timestep']) + 1
    local_y = np.array(train_set['ConfirmedCases'])
    p0 = np.random.exponential(size=3)
    bounds = (0, [100000.0, 3.0, 1000000000.0])
    try:
        (a, b, c), cov = optim.curve_fit(my_logistic, local_x, local_y, p0=p0, bounds=bounds)
    except:
        initial_infected = infection_start(country_entry, region_entry)
        c = get_population(country_entry, region_entry)
        b = 2
        a = c - 1
    test_set = filter_test_data(country_entry, region_entry, date_entry)
    test_set['Infected'] = round(my_logistic(test_set['Timestep'], a, b, c))
    return test_set.iloc[0]['Infected']

uruguay_forecast = calculate_infection('Uruguay', '', '2020-05-11')
print(uruguay_forecast)