application_vnd.jupyter.stderr_output_2.png text_plain_output_1.png
from sklearn.metrics import r2_score

from sklearn.model_selection import train_test_split, cross_val_score, KFold

from sklearn.preprocessing import MinMaxScaler, LabelEncoder

import matplotlib.pyplot as plt

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

import xgboost as xgb

verbose = False
loop_logic = True
scale_data = True
use_base_model = True
one_hot_encode = False
estimators = 5000
public_leaderboard_end_date = None

def get_test_train_for_country_state(one_hot_encode_flag, df_train, df_test, country, state):
    if one_hot_encode_flag == True:
        cs_train = df_train[(df_train['Country_Region_' + country] == 1) & (df_train['Province_State_' + state] == 1)]
        cs_test = df_test[(df_test['Country_Region_' + country] == 1) & (df_test['Province_State_' + state] == 1)]
    else:
        cs_train = df_train[(df_train['Country_Region'] == country) & (df_train['Province_State'] == state)]
        cs_test = df_test[(df_test['Country_Region'] == country) & (df_test['Province_State'] == state)]
    return (cs_train, cs_test)

def transform_dates(df):
    dates = pd.to_datetime(df['Date'])
    min_dates = dates.min()
    df['Date_Year'] = dates.dt.year
    df['Date_Month'] = dates.dt.month
    df['Date_Day'] = dates.dt.day
    df.drop(['Date'], axis=1, inplace=True)

def setup_df_encode_and_dates(df, encode_flag, dummy_cols, target_cols=[]):
    enc_df = df.copy()
    enc_df = enc_df[[enc_df.columns[0], enc_df.columns[2], enc_df.columns[1], enc_df.columns[3]]]
    if encode_flag == True:
        enc_df = pd.get_dummies(enc_df, columns=dummy_cols)
    else:
        le = LabelEncoder()
        for dum_col in dummy_cols:
            enc_df[dum_col] = le.fit_transform(enc_df[dum_col])
    transform_dates(enc_df)
    for col in target_cols:
        enc_df[col] = df[col]
    return enc_df

def prepare_train_set(df_train):
    train_x, train_target1, train_target2 = (df_train.iloc[:, :-2], df_train.iloc[:, -2], df_train.iloc[:, -1])
    return (train_x, train_target1, train_target2)

def prepare_submission(preds):
    preds['ForecastId'] = preds['ForecastId'].fillna(0.0).astype('int32')
    preds['Fatalities'] = preds['Fatalities'].fillna(0.0).astype('int32')
    preds['ConfirmedCases'] = preds['ConfirmedCases'].fillna(0.0).astype('int32')
    preds.clip(lower=0, inplace=True)
    preds.to_csv('submission.csv', index=False)

def model_and_predict(model, X, y, test, estimators=5000):
    if model != None:
        run_model = model
    else:
        run_model = xgb.XGBRegressor(objective='reg:squarederror', n_estimators=estimators)
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=12345)
    run_model.fit(X_train, y_train)
    y_train_pred = run_model.predict(X_train)
    y_test_pred = run_model.predict(X_test)
    y_pred = run_model.predict(test)
    y_pred[y_pred < 0] = 0
    r2 = r2_score(y_train_pred, y_train, multioutput='variance_weighted')
    return (y_pred, r2)

def show_results(model):
    # Code based on "Selecting Optimal Parameters for XGBoost Model Training" by Andrej Baranovskij (Medium)
    results = model.evals_result()
    epochs = len(results['validation_0']['error'])
    x_axis = range(0, epochs)
    # plot log loss
    fig, ax = plt.subplots()
    ax.plot(x_axis, results['validation_0']['logloss'], label='Train')
    ax.plot(x_axis, results['validation_1']['logloss'], label='Test')
    ax.legend()
    plt.ylabel('Log Loss')
    plt.title('XGBoost Log Loss')
    plt.show()
    # plot classification error
    fig, ax = pyplot.subplots()
    ax.plot(x_axis, results['validation_0']['error'], label='Train')
    ax.plot(x_axis, results['validation_1']['error'], label='Test')
    ax.legend()
    plt.ylabel('Classification Error')
    plt.title('XGBoost Classification Error')
    plt.show()

def fit_models_and_train(country, state, model, train, test):
    X, y_cases, y_fatal = prepare_train_set(train)
    X = X.drop(['Id'], axis=1)
    forecast_IDs = test.iloc[:, 0]
    test_no_id = test.iloc[:, 1:]
    if scale_data == True:
        scaler = MinMaxScaler()
        X = scaler.fit_transform(X.values)
        test_no_id = scaler.transform(test_no_id.values)
    y_cases_pred, cases_r2 = model_and_predict(model, X, y_cases, test_no_id)
    y_fatal_pred, fatal_r2 = model_and_predict(model, X, y_fatal, test_no_id)
    preds = pd.DataFrame(forecast_IDs)
    preds['ConfirmedCases'] = y_cases_pred
    preds['Fatalities'] = y_fatal_pred
    return preds

def cv_model(country, state, train, test):
    X, y_cases, y_fatal = prepare_train_set(train)
    X = X.drop(['Id'], axis=1)
    X_test = test.iloc[:, 1:]
    data_train_cases_matrix = xgb.DMatrix(data=X, label=y_cases)
    data_train_fatal_matrix = xgb.DMatrix(data=X, label=y_fatal)
    cv_results_cases = xgb.cv(dtrain=data_train_cases_matrix, params=parms, nfold=3, num_boost_round=50, early_stopping_rounds=50, metrics='rmse', as_pandas=True, seed=12345)
    cv_results_fatal = xgb.cv(dtrain=data_train_fatal_matrix, params=parms, nfold=3, num_boost_round=50, early_stopping_rounds=50, metrics='rmse', as_pandas=True, seed=12345)

df_train = pd.read_csv('../input/covid19-global-forecasting-week-4/train.csv')

df_test = pd.read_csv('../input/covid19-global-forecasting-week-4/test.csv')

df_train.shape

df_test.shape

df_train_original = df_train.copy()
df_test_original = df_test.copy()
df_train_original['Datetime'] = pd.to_datetime(df_train_original['Date'])
df_test_original['Datetime'] = pd.to_datetime(df_test_original['Date'])

date_filter = df_train[df_train.Date > '4/1/2020'].index
df_train.drop(date_filter, inplace=True)

df_train[df_train.Date > '2020/04/01']

base_model = xgb.XGBRegressor(n_estimators=estimators, random_state=12345, max_depth=15)

if one_hot_encode == True:
    country_groups = df_train_original.groupby(['Country_Region', 'Province_State']).groups
    df_country_list = pd.DataFrame.from_dict(list(country_groups))
    train_country_list = df_country_list[0].unique()
df_train_dd = setup_df_encode_and_dates(df_train, one_hot_encode, ['Country_Region', 'Province_State'], ['ConfirmedCases', 'Fatalities'])
df_test_dd = setup_df_encode_and_dates(df_test, one_hot_encode, ['Country_Region', 'Province_State'])
if one_hot_encode == False:
    country_groups = df_train_dd.groupby(['Country_Region', 'Province_State']).groups
    df_country_list = pd.DataFrame.from_dict(list(country_groups))
    train_country_list = df_country_list[0].unique()
df_preds = pd.DataFrame({'ForecastId': [], 'ConfirmedCases': [], 'Fatalities': []})
if loop_logic == True:
    print('Starting forecasting for {0} countries.'.format(len(train_country_list)))
    for country in train_country_list:
        print('Starting country {0}.'.format(country))
        country_states = df_country_list[df_country_list[0] == country][1].values
        for state in country_states:
            curr_cs_train, curr_cs_test = get_test_train_for_country_state(one_hot_encode, df_train_dd, df_test_dd, country, state)
            preds = fit_models_and_train(country, state, base_model if use_base_model == True else None, curr_cs_train, curr_cs_test)
            preds = preds.round(5)
            df_preds = pd.concat([df_preds, preds], axis=0)
        print('Country {0} complete.'.format(country))
else:
    print('Starting forecasting for all {0} countries.'.format(len(train_country_list)))
    preds = fit_models_and_train('All', 'All', base_model if use_base_model == True else None, df_train_dd, df_test_dd)
    df_preds = pd.concat([df_preds, preds], axis=0)
print('All countries complete.')
if ~(public_leaderboard_end_date == None):
    df_preds.loc[df_test_original.Datetime > pd.to_datetime(public_leaderboard_end_date), 'ConfirmedCases'] = 1
    df_preds.loc[df_test_original.Datetime > pd.to_datetime(public_leaderboard_end_date), 'Fatalities'] = 1
df_preds[df_test_original.Datetime > pd.to_datetime(public_leaderboard_end_date)].head()