text_plain_output_1.png
from sklearn.preprocessing import LabelEncoder

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

covid_train = pd.read_csv('../input/covid19-global-forecasting-week-4/train.csv', index_col='Id', parse_dates=['Date'])
covid_test = pd.read_csv('../input/covid19-global-forecasting-week-4/test.csv', index_col='ForecastId', parse_dates=['Date'])
last_register = pd.to_datetime(covid_train['Date'].iloc[covid_train.shape[0] - 1])

def adjustState(row):
    if pd.isna(row['Province_State']):
        row['Province_State'] = row['Country_Region']
    return row
covid_train = covid_train.apply(adjustState, axis=1)
covid_test = covid_test.apply(adjustState, axis=1)
covid_train.fillna('NA', inplace=True)
covid_test.fillna('NA', inplace=True)
n_cases_cc = 50
n_cases_ft = 50
data_mark_date = pd.DataFrame(columns=['Country_Region', 'Province_State', 'Date_cc', 'Date_ft'])
data_mark_date.set_index(['Country_Region', 'Province_State'])
for country in covid_train['Country_Region'].unique():
    for state in covid_train[covid_train['Country_Region'] == country]['Province_State'].unique():
        data_df = covid_train[(covid_train['Country_Region'] == country) & (covid_train['Province_State'] == state)]
        if data_df[data_df['ConfirmedCases'] >= n_cases_cc].shape[0] > 0:
            date_cc = data_df[data_df['ConfirmedCases'] >= n_cases_cc].iloc[0]['Date']
        else:
            date_cc = last_register
        if data_df[data_df['Fatalities'] >= n_cases_ft].shape[0] > 0:
            date_ft = data_df[data_df['Fatalities'] >= n_cases_ft].iloc[0]['Date']
        else:
            date_ft = last_register
        data_state = pd.DataFrame({'Country_Region': [country], 'Province_State': [state], 'Date_cc': [date_cc], 'Date_ft': [date_ft]})
        data_state.set_index(['Country_Region', 'Province_State'])
        data_mark_date = data_mark_date.append(data_state.iloc[0])
def mark_date(row):
    data_df = data_mark_date[(data_mark_date['Country_Region'] == row['Country_Region']) & (data_mark_date['Province_State'] == row['Province_State'])].iloc[0]
    if not pd.isna(data_df['Date_cc']):
        row['Date_cc'] = (row['Date'] - data_df['Date_cc']).days
    if not pd.isna(data_df['Date_ft']):
        row['Date_ft'] = (row['Date'] - data_df['Date_ft']).days
    return row
covid_train = covid_train[(covid_train['ConfirmedCases'] > 0) | (covid_train['Fatalities'] > 0)]
covid_train['Date_cc'] = [0 for i in range(covid_train.shape[0])]
covid_train['Date_ft'] = [0 for i in range(covid_train.shape[0])]
covid_train = covid_train.apply(mark_date, axis=1)
covid_test['Date_cc'] = [0 for i in range(covid_test.shape[0])]
covid_test['Date_ft'] = [0 for i in range(covid_test.shape[0])]
covid_test = covid_test.apply(mark_date, axis=1)
covid_train['Date_st'] = covid_train['Date'].map(lambda x: x.timestamp())
covid_test['Date_st'] = covid_test['Date'].map(lambda x: x.timestamp())
enc = LabelEncoder()
covid_train['Province_State_enc'] = enc.fit_transform(covid_train['Province_State'])
covid_test['Province_State_enc'] = enc.transform(covid_test['Province_State'])
enc = LabelEncoder()
covid_train['Country_Region_enc'] = enc.fit_transform(covid_train['Country_Region'])
covid_test['Country_Region_enc'] = enc.transform(covid_test['Country_Region'])
X_features = ['Province_State', 'Country_Region', 'Date_st', 'Date_cc', 'Date_ft']
X = covid_train[X_features]
y_cc = covid_train['ConfirmedCases']
y_ft = covid_train['Fatalities']
print('Adjust data complete')