text_plain_output_1.png
from collections import OrderedDict

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

N_FEATURES = 5

countries = pd.read_csv('/kaggle/input/countries-of-the-world/countries of the world.csv', decimal=',')
countries['Country'] = countries.Country.str.lower()
countries = countries.set_index('Country')

train_data = pd.read_csv('/kaggle/input/covid19-global-forecasting-week-4/train.csv').set_index('Id')

def transform_area(data):
    data['Area'] = ''
    data.loc[~data.Province_State.isna(), 'Area'] = data.loc[~data.Province_State.isna()].Country_Region.str.lower() + '/' + data.Province_State.str.lower()
    data.loc[data.Province_State.isna(), 'Area'] = data.loc[data.Province_State.isna()].Country_Region.str.lower()
    data = data.drop('Province_State', axis='columns')
    return data
def transform_countries(data):
    data['Density'] = countries['Pop. Density (per sq. mi.)'].mean()
    data['InfantMortality'] = countries['Infant mortality (per 1000 births)'].mean()
    for i in data.index:
        country = data.loc[i, 'Country_Region'].lower()
        if country in countries.index:
            data.loc[i, 'Density'] = countries.loc[country, 'Pop. Density (per sq. mi.)']
    return data

train_data = transform_area(train_data)
train_data = transform_countries(train_data)

X_train = {}
X_test = {}
Y_train = {}
Y_test = {}
for area in list(set(train_data.Area)):
    print(area)
    area_data = train_data.loc[train_data.Area == area].set_index('Date').sort_index()
    Y = area_data[['ConfirmedCases', 'Fatalities']] - area_data.shift(1)[['ConfirmedCases', 'Fatalities']]
    dic = OrderedDict()
    for i in range(1, 1 + N_FEATURES):
        dic['CC_{}'.format(i)] = area_data.shift(i)['ConfirmedCases']
        dic['F_{}'.format(i)] = area_data.shift(i)['Fatalities']
    dic['Density'] = area_data['Density']
    dic['InfantMortality'] = area_data['InfantMortality']
    X = pd.DataFrame(dic, index=area_data.index)
    X = X.dropna(axis='index')
    Y = Y.loc[X.index]
    X_train[area] = X.iloc[:int(len(X) * 0.8)]
    X_test[area] = X.iloc[int(len(X) * 0.8):]
    Y_train[area] = Y.iloc[:int(len(X) * 0.8)]
    Y_test[area] = Y.iloc[int(len(X) * 0.8):]