Machine_Learning_in_Python/base_models_regression.py at master · markwang1040/Machine_Learning_in_Python · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
import pandas as pd
import numpy as np

from sklearn.ensemble import RandomForestRegressor, GradientBoostingRegressor
from sklearn.linear_model import ElasticNet
from xgboost import XGBRegressor

from sklearn.metrics import mean_squared_error, mean_absolute_error

import datetime as dt


# Authors: Mark Wang <markswang@uchicago.edu>


def fit_regression_models(X_train, y_train, X_test, y_test, random_state):

    rfr = RandomForestRegressor(n_estimators=500, n_jobs=-1, random_state=random_state, oob_score=True)
    gbr = GradientBoostingRegressor(learning_rate=0.1, n_estimators=500, random_state=random_state)
    en = ElasticNet(l1_ratio=0.5, max_iter=100000, random_state=random_state)
    xgbr = XGBRegressor(learning_rate=0.1, n_estimators=500, random_state=random_state, n_jobs=-1)

    regressors = [rfr, gbr, en, xgbr]

    model = []
    training_error_mse = []
    testing_error_mse = []
    training_error_mae = []
    testing_error_mae = []
    time_taken = []

    for i in regressors:
        begin = dt.datetime.now()

        i.fit(X_train, y_train)
        y_pred_train = i.predict(X_train)
        y_pred_test = i.predict(X_test)
        mse_train = mean_squared_error(y_train, y_pred_train)
        mse_test = mean_squared_error(y_test, y_pred_test)
        mae_train = mean_absolute_error(y_train, y_pred_train)
        mae_test = mean_absolute_error(y_test, y_pred_test)


        run_time = dt.datetime.now() - begin

        model.append(str(i).split('(')[0])
        training_error_mse.append(mse_train)
        testing_error_mse.append(mse_test)
        training_error_mae.append(mae_train)
        testing_error_mae.append(mae_test)

        time_taken.append(run_time)

    results = pd.DataFrame(zip(model, training_error_mse, testing_error_mse, training_error_mae, testing_error_mae,
                               time_taken), columns = ['model', 'trainint_error_mse', 'testing_error_mse',
                                                       'trainint_error_mae', 'testing_error_mae','time_taken'])

    return(rfr, gbr, en, xgbr, results)