12223737_proj_2_2.py

# -*- coding: utf-8 -*-
"""12223737_Proj#2-2.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1tueSiF7huFFWtFr-AUKwEU3rWT34PZUD
"""

import pandas as pd
import numpy as np
from sklearn.tree import DecisionTreeRegressor
from sklearn.ensemble import RandomForestRegressor
from sklearn.svm import SVR
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import mean_squared_error

def sort_dataset(dataset_df):
  sorted_df = dataset_df.sort_values('year', ascending=True)
  return sorted_df

def split_dataset(dataset_df):
  dataset_df['salary'] = dataset_df['salary'] * 0.001
  dataset_train = dataset_df.iloc[:1718]
  dataset_test = dataset_df.iloc[1718:]

  X_train = dataset_train.drop('salary', axis=1)
  Y_train = dataset_train['salary']

  X_test = dataset_test.drop('salary', axis=1)
  Y_test = dataset_test['salary']

  return X_train, X_test, Y_train, Y_test

def extract_numerical_cols(dataset_df):
  numerical_cols = {'age', 'G', 'PA', 'AB', 'R', 'H', '2B', '3B', 'HR', 'RBI',
                    'SB', 'CS', 'BB', 'HBP', 'SO', 'GDP', 'fly', 'war'}
  numerical_df = dataset_df[numerical_cols]
  return numerical_df

def train_predict_decision_tree(X_train, Y_train, X_test):
  dt_reg = DecisionTreeRegressor()
  dt_reg.fit(X_train, Y_train)

  dt_predicted = dt_reg.predict(X_test)
  return dt_predicted

def train_predict_random_forest(X_train, Y_train, X_test):
  rf_reg = RandomForestRegressor()
  rf_reg.fit(X_train, Y_train)

  rf_predicted = rf_reg.predict(X_test)
  return rf_predicted

def train_predict_svm(X_train, Y_train, X_test):
  svm_pipe = make_pipeline(
      StandardScaler(),
      SVR()
  )

  svm_pipe.fit(X_train, Y_train)
  svm_predicted = svm_pipe.predict(X_test)
  return svm_predicted

def calculate_RMSE(labels, predictions):
  RMSE = np.sqrt(np.mean((predictions - labels)**2))
  return RMSE

if __name__=='__main__':
  data_df = pd.read_csv('/content/drive/MyDrive/OSS_Project2/2019_kbo_for_kaggle_v2.csv')

  sorted_df = sort_dataset(data_df)
  X_train, X_test, Y_train, Y_test = split_dataset(sorted_df)

  X_train = extract_numerical_cols(X_train)
  X_test = extract_numerical_cols(X_test)

  dt_predictions = train_predict_decision_tree(X_train, Y_train, X_test)
  rf_predictions = train_predict_random_forest(X_train, Y_train, X_test)
  svm_predictions = train_predict_svm(X_train, Y_train, X_test)

  print ("Decision Tree Test RMSE: ", calculate_RMSE(Y_test, dt_predictions))
  print ("Random Forest Test RMSE: ", calculate_RMSE(Y_test, rf_predictions))
  print ("SVM Test RMSE: ", calculate_RMSE(Y_test, svm_predictions))