freeEDU/Log-Decoder

1"""2The implementation of the decision tree model for anomaly detection.3 4Authors: 5    LogPAI Team6 7Reference: 8    [1] Mike Chen, Alice X. Zheng, Jim Lloyd, Michael I. Jordan, Eric Brewer. 9        Failure Diagnosis Using Decision Trees. IEEE International Conference 10        on Autonomic Computing (ICAC), 2004.11 12"""13 14import numpy as np15from sklearn import tree16from ..utils import metrics17 18class DecisionTree(object):19 20    def __init__(self, criterion='gini', max_depth=None, max_features=None, class_weight=None):21        """ The Invariants Mining model for anomaly detection22        Arguments23        ---------24        See DecisionTreeClassifier API: https://scikit-learn.org/stable/modules/generated/sklearn.svm.LinearSVC.html25 26        Attributes27        ----------28            classifier: object, the classifier for anomaly detection29 30        """31        self.classifier = tree.DecisionTreeClassifier(criterion=criterion, max_depth=max_depth,32                          max_features=max_features, class_weight=class_weight)33 34    def fit(self, X, y):35        """36        Arguments37        ---------38            X: ndarray, the event count matrix of shape num_instances-by-num_events39        """40        print('====== Model summary ======')41        self.classifier.fit(X, y)42 43    def predict(self, X):44        """ Predict anomalies with mined invariants45 46        Arguments47        ---------48            X: the input event count matrix49 50        Returns51        -------52            y_pred: ndarray, the predicted label vector of shape (num_instances,)53        """54        55        y_pred = self.classifier.predict(X)56        return y_pred57 58    def predict_proba(self, X):59        """ Predict anomalies with mined invariants60 61        Arguments62        ---------63            X: the input event count matrix64 65        Returns66        -------67            y_pred: ndarray, the predicted label vector of shape (num_instances,)68        """69        70        y_pred = self.classifier.predict_proba(X)71        return y_pred72 73    def evaluate(self, X, y_true):74        print('====== Evaluation summary ======')75        y_pred = self.predict(X)76        precision, recall, f1 = metrics(y_pred, y_true)77        print('Precision: {:.3f}, recall: {:.3f}, F1-measure: {:.3f}\n'.format(precision, recall, f1))78        return precision, recall, f179