binary_classification_rf.py

# -*- coding: utf-8 -*-
"""
classification using random forest
"""

import pandas as pd
import matplotlib.pyplot as plt

from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score, ConfusionMatrixDisplay

from data_module.alignment_data import AlignmentData
from data_module.feature_extractor import FeatureExtractor
from data_module.features import Mean



rm_filenames = ['misaligned_0.03_1500rpm_10min.txt',
                'misaligned_0.06_1500rpm_10min.txt',
                'misaligned_0.09_1500rpm_10min.txt',
                'misaligned_0.12_1500rpm_10min.txt',
                'misaligned_0.15_1500rpm_10min.txt']

sampling_method = {'name': 'segment', 'interval': 10000}

features = [Mean()]

def main():
    
    # prep data
    data = AlignmentData()
    data.load_all_filenames()
    data.remove_filenames(rm_filenames)
    data.load_raw_data()
    
    # create dataset
    feature_extractor = FeatureExtractor(features)
    df = data.create_dataset(feature_extractor, sampling_method)
    print('Number of Observations: ', df.shape[0])
    
    # train and test sets 
    X_train, X_test, y_train, y_test = train_test_split(df.iloc[:,:14], df['Label'], test_size=0.33)
    
    # train model
    rf = RandomForestClassifier().fit(X_train, y_train)
    
    # evaluate
    y_pred = rf.predict(X_test)
    accuracy = accuracy_score(y_test, y_pred)
    print('Accuracy: ', accuracy)
    
    ConfusionMatrixDisplay.from_predictions(y_test, y_pred)
    
    forest_importances = pd.Series(rf.feature_importances_, index=X_train.columns).sort_values()
    fig, ax = plt.subplots()
    forest_importances.plot.barh(ax=ax)
    ax.set_title("Feature importances using MDI")
    ax.set_ylabel("Mean decrease in impurity")
    fig.tight_layout()
    
    
if __name__ == "__main__":
    main()