text_plain_output_1.png
from sklearn.tree import DecisionTreeClassifier

import numpy as np

import seaborn as sns

NUM_FEATURES = 10
NUM_CLASSES = 3
NUM_POINTS = 100
MAX_DEPTH = 5
METRIC = 'gini'

def plotDists(y):
    return
def acc(true, pred):
    assert len(true) == len(pred), 'Truth and Pred Lengths not same'
    true = np.array(true)
    pred = np.array(pred).astype(np.int32)
    return np.sum(true == pred) / len(true)

class Node:
    def __init__(self):
        self.left = None
        self.right = None
        self.feature_id = None
        self.feature_thresh = None
        self.is_leaf = False
        self.result = None
        self.metric = 0
        self.level = -1
    def set_leaf(self, result, level):
        self.is_leaf = True
        self.result = result
        self.level = level
    def set_branch(self, best_feature, best_feature_thresh, best_metric, left, right, level):
        self.feature_id = best_feature
        self.feature_thresh = best_feature_thresh
        self.metric = best_metric
        self.left = left
        self.right = right
        self.level = level
    def traverse(self, inp):
        if self.is_leaf:
            return (self.result, True)
        if inp[self.feature_id] <= self.feature_thresh:
            return (self.left, False)
        else:
            return (self.right, False)
    def __repr__(self):
        s = ''
        s += f'Level   : {self.level}\n'
        s += f'Leaf    : {self.is_leaf}\n'
        if self.is_leaf:
            s += f'Result  : {self.result}\n'
        else:
            s += f'Feature : {self.feature_id}\n'
            s += f'Thresh  : {self.feature_thresh}\n'
        return s

class DecisionTree:
    def __init__(self, x, y, MAX_LEVELS=5, metric='gini'):
        self.x = np.array(x)
        self.y = np.array(y)
        self.data = np.concatenate((self.x, self.y.reshape(-1, 1)), axis=1)
        self.n_features = self.x.shape[1]
        self.classes = set(y)
        self.n_classes = len(self.classes)
        self.MAX_LEVELS = MAX_LEVELS
        self.root = None
        assert metric in ['gini'], 'Invalid metric'
        self.metric = metric
    def fit(self):
        self.root = self.buildTree(0, self.data)
    def get_metric(self, y):
        _, counts = np.unique(y, return_counts=True)
        total = len(y)
        if self.metric == 'gini':
            return 1 - np.sum((counts / total) ** 2)
        elif self.metric == 'entropy':
            return 0
    def split_data(self, data, feature_idx, thresh):
        return (data[data[:, feature_idx] <= thresh], data[data[:, feature_idx] > thresh])
    def buildTree(self, level, data):
        n = Node()
        if level == self.MAX_LEVELS or len(set(data[:, -1])) == 1:
            n.set_leaf(self.max_reps_class(data[:, -1]), level)
        elif self.metric == 'gini':
            best_metric = float('inf')
            left_partition = None
            right_partition = None
            best_feature = None
            best_feature_thresh = None
            for feature_idx in range(self.n_features):
                for thresh in set(self.data[:, feature_idx]):
                    left, right = self.split_data(data, feature_idx, thresh)
                    m = len(left) * self.get_metric(left[:, -1]) + len(right) * self.get_metric(right[:, -1])
                    if m < best_metric:
                        best_metric = m
                        left_partition = left
                        right_partition = right
                        best_feature = feature_idx
                        best_feature_thresh = thresh
            n.set_branch(best_feature, best_feature_thresh, best_metric, self.buildTree(level + 1, left_partition), self.buildTree(level + 1, right_partition), level)
        return n
    def max_reps_class(self, y):
        classes, counts = np.unique(y, return_counts=True)
        return classes[np.argmax(counts)]
    def predict(self, inputs, log_track=False):
        if not self.root:
            raise Exception('Tree Not Fit Yet!')
        results = []
        for inp in inputs:
            res = self.root
            fin = False
            while not fin:
                res, fin = res.traverse(inp)
            results.append(int(res))
        return results

d = DecisionTree(x, y, MAX_DEPTH, METRIC)
d.fit()

m = DecisionTreeClassifier(max_depth=MAX_DEPTH, criterion=METRIC)
m.fit(x, y)

acc(y, d.predict(x))

acc(y, m.predict(x))

d.predict([x[1]], True)