text_plain_output_1.png
from sklearn.linear_model import LinearRegression

from sklearn.metrics import mean_squared_error, r2_score

from sklearn.model_selection import KFold

import numpy as np

import pandas as pd

FIGSIZE = (10, 6)
SAVE_PICKLE = True
FREE_MEMORY = True
OUTPUT_FILE = 'potential_energy_upd'
RANDOM_STATE = 123
N_SPLITS = 3
SHUFFLE = True
VERBOSE = False

DATA_PATH = '../input'
def csv_path(dataset='train', data_path=DATA_PATH):
    """
    """
    return '{}/{}.csv'.format(data_path, dataset)
def read_data(dataset='train', data_path=DATA_PATH):
    """
    """
    index_col = None
    index_type = ['train', 'test']
    if dataset in index_type:
        index_col = 'id'
    data_path = csv_path(dataset, data_path=data_path)
    return pd.read_csv(data_path, index_col=index_col)

train = read_data('train')
test = read_data('test')

molecule_train = pd.DataFrame({'molecule_name': train['molecule_name'].unique()})
molecule_test = pd.DataFrame({'molecule_name': test['molecule_name'].unique()})
structures = read_data('structures')
atom_list_df = structures.groupby('molecule_name')['atom'].apply(list)
atom_list_df = atom_list_df.to_frame()
if FREE_MEMORY:
    del train
    del test

molecule_train = pd.merge(molecule_train, atom_list_df, how='left', on='molecule_name')
molecule_test = pd.merge(molecule_test, atom_list_df, how='left', on='molecule_name')

potential_energy = read_data('potential_energy')
molecule_train = pd.merge(molecule_train, potential_energy)
if FREE_MEMORY:
    del potential_energy
    del structures

id_feature = 'molecule_name'
target_feature = (set(molecule_train) - set(molecule_test)).pop()
selected_features = list(molecule_test)
selected_features.remove(id_feature)
selected_features.remove('atom')

X = molecule_train[selected_features]
y = molecule_train[target_feature]

kfold = KFold(n_splits=N_SPLITS, random_state=RANDOM_STATE, shuffle=SHUFFLE)

fold = 0
r2_scores = []
mse_scores = []
lin_reg = LinearRegression()
for in_index, oof_index in kfold.split(X, y):
    fold += 1
    print('- Training Fold: ({}/{})'.format(fold, N_SPLITS))
    X_in, X_oof = (X.loc[in_index], X.loc[oof_index])
    y_in, y_oof = (y.loc[in_index], y.loc[oof_index])
    lin_reg.fit(X_in, y_in)
    y_pred = lin_reg.predict(X_oof)
    r2 = r2_score(y_oof, y_pred)
    r2_scores.append(r2)
    mse_score = mean_squared_error(y_oof, y_pred)
    mse_scores.append(mse_score)
print('\nkFold Validation Results:')
print(' * Average Variance Score (R2): \t{:.4f}'.format(np.mean(r2_scores)))
print(' * Average Mean squared error (MSE): \t{:.4f}'.format(np.mean(mse_score)))