benchmark.py

import numpy as np
import pandas as pd
import sys
import time

import preparer
import attr_map
from sklearn.neighbors import KNeighborsClassifier
from sklearn.tree import DecisionTreeClassifier
from sklearn.ensemble import RandomForestClassifier
from sklearn.ensemble import BaggingClassifier
from sklearn.ensemble import AdaBoostClassifier
from sklearn.linear_model import LogisticRegression
from lightgbm import LGBMClassifier

from sklearn.metrics import roc_auc_score
from sklearn.svm import SVC
from sklearn.model_selection import train_test_split
from sklearn.impute import SimpleImputer
from sklearn.preprocessing import StandardScaler
from category_encoders import TargetEncoder
from category_encoders import CatBoostEncoder
from category_encoders import CountEncoder
from category_encoders import JamesSteinEncoder
from sklearn.pipeline import Pipeline
from sklearn.model_selection import GridSearchCV

# Arg handling
if len(sys.argv) < 3:
    sys.exit(
        'Please provide arguments!\n'
        'Arguments: path/to/dataset algorithm encoding\n'
        'algorithm: cart logistic svm knn bagging rf adaboost lgbm\n'
        'encoding: target js catboost freq\n'
        f'Example: python3 {sys.argv[0]} path/to/dataset.csv cart target'
    )

dataset_path = sys.argv[1]
algorithm = sys.argv[2]  # cart, lgbm
encoding = sys.argv[3]  # freq, target, or loo

# KNN & SVM suffers if we have large data & dimension
if algorithm in ['knn', 'svm']:
    MAX_COLUMNS = 10
else:
    MAX_COLUMNS = 100

# Mapping
CLF_MAP = {
    'cart': DecisionTreeClassifier,
    'logistic': LogisticRegression,
    'svm': SVC,
    'knn': KNeighborsClassifier,
    'bagging': BaggingClassifier,
    'rf': RandomForestClassifier,
    'adaboost': AdaBoostClassifier,
    'lgbm': LGBMClassifier
}
ENCODER_MAP = {
    'target': TargetEncoder,
    'js': JamesSteinEncoder,
    'freq': CountEncoder,
    'catboost': CatBoostEncoder,
    # 'loo': LeaveOneOutEncoder # LOO is buggy on the pipeline
}
PARAM_GRID_MAP = {
    'cart': {
        'clf__criterion': ['gini', 'entropy'],
        'clf__min_samples_leaf': np.array(range(10, 21, 5)) / 10_000,
    },
    'logistic': {
        'clf__penalty': ['l2', 'none'],
        'clf__C': [0.5, 0.75, 1.0],
    },
    'svm': {
        'clf__kernel': ['linear', 'rbf'],
        'clf__C': [0.5, 1.0, 0.25]
    },
    'knn': {
        'clf__n_neighbors': [3, 5, 7],
        'clf__weights': ['uniform', 'distance']
    },
    'bagging': {
        'clf__n_estimators': [10, 15],
        'clf__max_samples': [0.5, 1.0],
        'clf__max_features': [0.5, 1.0],
    },
    'rf': {
        'clf__criterion': ['gini', 'entropy'],
        'clf__min_samples_leaf': np.array(range(10, 21, 5)) / 10_000,
    },
    'adaboost': {
        'clf__n_estimators': [50, 100],
        'clf__learning_rate': [0.5, 0.75, 1.0],
    },
    'lgbm': {
        'clf__reg_alpha': [0, 2.5, 5],
        'clf__min_child_samples': [50, 100, 150],
    }
}


# Attributes
attributes = attr_map.ATTR_MAP[encoding]['attrs'][:MAX_COLUMNS]
cat_attributes = np.intersect1d(
    attributes, attr_map.ATTR_MAP[encoding]['cat_attrs']
)
target = attr_map.TARGET

# Prepare & split data
df = preparer.get_data(dataset_path, encoding)
X_train, X_test, y_train, y_test = train_test_split(
    df[attributes], df[target], test_size=0.2, random_state=42
)

# Encoder & classifier
encoder = ENCODER_MAP[encoding](cols=cat_attributes)
if algorithm == 'svm':
    clf = CLF_MAP[algorithm](probability=True)   # so we can predict_proba
else:
    clf = CLF_MAP[algorithm]()

# Initialize pipeline
estimators = list()

# Encoder to pipeline
if encoding == 'freq':  # CountEncoder needn't be put in pipeline for CV
    encoder.fit(pd.concat([X_train, X_test]))   # Because needn't label
    X_train = encoder.transform(X_train)
    X_test = encoder.transform(X_test)
else:
    estimators.append(('encoder', encoder))

# Add scaler and classifier to the pipeline
estimators += [
    ('imputer', SimpleImputer()),
    ('scaler', StandardScaler()),
    ('clf', clf)
]
pipe = Pipeline(steps=estimators)

# Parameter grid
param_grid = PARAM_GRID_MAP[algorithm]

# Fit
start = time.time()
grid = GridSearchCV(pipe, param_grid=param_grid, scoring='roc_auc')
grid.fit(X_train, y_train)

# Print grid search results
print(f'{algorithm} {encoding} {X_train.shape[0] + X_test.shape[0]}\n')
print('Grid Search Results')
print('Best params:')
print(grid.best_params_)
print('Best score:')
print(grid.best_score_)
print('Scores:')
print(grid.cv_results_['mean_test_score'])

# Test
y_pred = grid.predict_proba(X_test)
print()
print('TEST SCORE:')
print(roc_auc_score(y_test, y_pred[:, 1]))
end = time.time()

print()
print('Running Time:', end - start, 'seconds')

# Feature importances
TREE_ALGOS = ['cart', 'rf', 'adaboost', 'lgbm']
if algorithm in TREE_ALGOS:
    fi_df = pd.DataFrame(index=X_test.columns, columns=['Importance'])
    fi_df['Importance'] = grid.best_estimator_.steps[-1][1].feature_importances_
    print()
    print('Feature Importances')
    print(fi_df.sort_values('Importance', ascending=False))

# # Fit
# pipe.fit(X_train, y_train)

# # Test
# y_pred = pipe.predict(X_test)
# print()
# print('TEST SCORE:')
# print(roc_auc_score(y_test, y_pred))