PyDLLib 1.0__py3-none-any.whl

DLL/MachineLearning/SupervisedLearning/SupportVectorMachines/_SVR.py

@@ -0,0 +1,100 @@
+import torch
+import cvxopt
+import numpy as np
+
+from ..Kernels import RBF, _Base
+from ....Exceptions import NotFittedError
+
+
+class SVR:
+    """
+    The support vector machine regressor with a quadratic programming solver.
+    
+    Args:
+        kernel (:ref:`kernel_section_label`, optional): The non-linearity function for fitting the model. Defaults to RBF.
+        C (float or int, optional): A regularization parameter. Defaults to 1. Must be positive real number.
+        epsilon (float or int, optional): The width of the tube of no penalty in epsilon-SVR. Must be a positive real number.
+    Attributes:
+        n_features (int): The number of features. Available after fitting.
+        alpha (torch.Tensor of shape (n_samples,)): The optimized dual coefficients. Available after fitting.
+        alpha_star (torch.Tensor of shape (n_samples,)): The optimized dual coefficients. Available after fitting.
+    """
+    def __init__(self, kernel=RBF(), C=1, epsilon=0.1):
+        if not isinstance(kernel, _Base):
+            raise ValueError("kernel must be from DLL.MachineLearning.SupervisedLearning.Kernels")
+        if not isinstance(C, float | int) or C <= 0:
+            raise ValueError("C must be must be positive real number.")
+        if not isinstance(epsilon, float | int) or epsilon <= 0:
+            raise ValueError("epsilon must be must be positive real number.")
+
+        self.kernel = kernel
+        self.C = C
+        self.epsilon = epsilon
+    
+    def _kernel_matrix(self, X1, X2):
+        return self.kernel(X1, X2).to(X1.dtype)
+    
+    def fit(self, X, y):
+        """
+        Fits the SVR model to the input data by finding the hypertube that contains the data with minimum loss.
+
+        Args:
+            X (torch.Tensor of shape (n_samples, n_features)): The input data, where each row is a sample and each column is a feature.
+            y (torch.Tensor of shape (n_samples,)): The target values corresponding to each sample.
+        Returns:
+            None
+        Raises:
+            TypeError: If the input matrix or the target matrix is not a PyTorch tensor.
+            ValueError: If the input matrix or the target matrix is not the correct shape.
+        """
+        if not isinstance(X, torch.Tensor) or not isinstance(y, torch.Tensor):
+            raise TypeError("The input matrix and the target matrix must be a PyTorch tensor.")
+        if X.ndim != 2:
+            raise ValueError("The input matrix must be a 2 dimensional tensor.")
+        if y.ndim != 1 or y.shape[0] != X.shape[0]:
+            raise ValueError("The targets must be 1 dimensional with the same number of samples as the input data")
+
+        self.y = y.reshape((-1, 1)).to(X.dtype)
+        self.X = X
+        self.n_features = X.shape[1]
+        n = X.shape[0]
+        K = self._kernel_matrix(X, X).numpy()
+
+        P = cvxopt.matrix(np.block([[K, -K], [-K, K]]).tolist())  # [[K, -K], [-K, K]]
+        q = cvxopt.matrix(np.hstack([self.epsilon - y, self.epsilon + y]).tolist())
+        A = cvxopt.matrix(np.vstack([np.ones((n, 1)), -np.ones((n, 1))]).tolist())
+        b = cvxopt.matrix(0.0)
+        G = cvxopt.matrix(np.hstack([-np.eye(2 * n), np.eye(2 * n)]).tolist())
+        h = cvxopt.matrix(np.hstack([np.zeros((2 * n,)), np.ones((2 * n,)) * self.C]).tolist())
+
+        cvxopt.solvers.options['show_progress'] = False
+        sol = cvxopt.solvers.qp(P, q, G, h, A, b)
+        alpha = torch.tensor(np.array(sol["x"]), dtype=torch.float64).squeeze(dim=1)
+        self.alpha = alpha[:n]
+        self.alpha_star = alpha[n:]
+
+    def predict(self, X):
+        """
+        Applies the fitted SVR model to the input data, predicting the correct values.
+
+        Args:
+            X (torch.Tensor of shape (n_samples, n_features)): The input data to be regressed.
+        Returns:
+            target values (torch.Tensor of shape (n_samples,)): The predicted values corresponding to each sample.
+        Raises:
+            NotFittedError: If the SVR model has not been fitted before predicting.
+            TypeError: If the input matrix is not a PyTorch tensor.
+            ValueError: If the input matrix is not the correct shape.
+        """
+        if not hasattr(self, "n_features"):
+            raise NotFittedError("SVR.fit() must be called before predicting.")
+        if not isinstance(X, torch.Tensor):
+            raise TypeError("The input matrix must be a PyTorch tensor.")
+        if X.ndim != 2 or X.shape[1] != self.n_features:
+            raise ValueError("The input matrix must be a 2 dimensional tensor with the same number of features as the fitted tensor.")
+        
+        # Use every datapoint as a support vector, since otherwise the results seem to be bad.
+        # This makes no sense to me, but yields the best results.
+        bias = (self.y - torch.sum((self.alpha - self.alpha_star) * self._kernel_matrix(self.X, self.X), dim=1)).mean()
+        prediction = (self.alpha - self.alpha_star) @ self._kernel_matrix(self.X, X).to(dtype=self.alpha.dtype) + bias
+        return prediction

DLL/MachineLearning/SupervisedLearning/SupportVectorMachines/__init__.py

@@ -0,0 +1,2 @@
+from ._SVC import SVC
+from ._SVR import SVR

DLL/MachineLearning/SupervisedLearning/Trees/_AdaBoostClassifier.py

@@ -0,0 +1,179 @@
+import torch
+from math import log
+from warnings import warn
+
+from ._DecisionTree import DecisionTree
+from ....Exceptions import NotFittedError
+
+
+class AdaBoostClassifier:
+    """
+    AdaBoostClassifier implements a classification algorithm fitting many consecutive :class:`DecisionTrees <DLL.MachineLearning.SupervisedLearning.Trees.DecisionTree>` to previously missclassified samples.
+
+    Args:
+        n_trees (int, optional): The number of trees used for predicting. Defaults to 10. Must be a positive integer.
+        max_depth (int, optional): The maximum depth of the tree. Defaults to 25. Must be a positive integer.
+        min_samples_split (int, optional): The minimum required samples in a leaf to make a split. Defaults to 2. Must be a positive integer.
+        criterion (str, optional): The information criterion used to select optimal splits. Must be one of "entropy" or "gini". Defaults to "gini".
+    Attributes:
+        n_features (int): The number of features. Available after fitting.
+        n_classes (int): The number of classes. 2 for binary classification. Available after fitting.
+        confidences (torch.tensor of shape (n_trees,)): The confidence on each tree.
+    """
+    def __init__(self, n_trees=10, max_depth=25, min_samples_split=2, criterion="gini"):
+        if not isinstance(n_trees, int) or n_trees < 1:
+            raise ValueError("n_trees must be a positive integer.")
+        if not isinstance(max_depth, int) or max_depth < 1:
+            raise ValueError("max_depth must be a positive integer.")
+        if not isinstance(min_samples_split, int) or min_samples_split < 1:
+            raise ValueError("min_samples_split must be a positive integer.")
+        if criterion not in ["entropy", "gini"]:
+            raise ValueError('The chosen criterion must be one of "entropy" or "gini".')
+        
+        self.n_trees = n_trees
+        self.max_depth = max_depth
+        self.min_samples_split = min_samples_split
+        self.criterion = criterion
+        self.trees = None
+
+    def fit(self, X, y, verbose=True):
+        """
+        Fits the AdaBoostClassifier model to the input data by fitting trees to the errors made by previous trees.
+
+        Args:
+            X (torch.Tensor of shape (n_samples, n_features)): The input data, where each row is a sample and each column is a feature.
+            y (torch.Tensor of shape (n_samples,)): The labels corresponding to each sample. Every element must be in [0, ..., n_classes - 1].
+            verbose (bool, optional): Determines if warnings are given if the training ends due to a weak learner being worse than random guessing. Defaults to True.
+        Returns:
+            The average errors after each tree.
+        Raises:
+            TypeError: If the input matrix or the label vector is not a PyTorch tensor or if the problem is binary and metrics is not a list or a tuple.
+            ValueError: If the input matrix or the label vector is not the correct shape or the label vector contains wrong values.
+        """
+        if not isinstance(X, torch.Tensor) or not isinstance(y, torch.Tensor):
+            raise TypeError("The input matrix and the label matrix must be a PyTorch tensor.")
+        if X.ndim != 2:
+            raise ValueError("The input matrix must be a 2 dimensional tensor.")
+        if y.ndim != 1 or y.shape[0] != X.shape[0]:
+            raise ValueError("The labels must be 1 dimensional with the same number of samples as the input data")
+        vals = torch.unique(y).numpy()
+        if set(vals) != {*range(len(vals))}:
+            raise ValueError("y must only contain the values in [0, ..., n_classes - 1].")
+        if not isinstance(verbose, bool):
+            raise TypeError("verbose must be a boolean.")
+        
+        self.n_classes = len(vals)
+        self.classes = vals
+        y = y.to(X.dtype)
+        self.n_features = X.shape[1]
+
+        trees = []
+        weights = torch.full_like(y, 1 / len(y))
+        self.confidences = []
+
+        errors = []
+
+        for i in range(self.n_trees):
+            indices = torch.multinomial(weights, len(y), replacement=True)
+            X_sample = X[indices]
+            y_sample = y[indices]
+
+            #  Make sure atleast 1 datapoint is present from each class
+            sample_classes = torch.unique(y_sample)
+            if len(sample_classes) < self.n_classes:
+                unseen_classes = list(set(self.classes) - set(sample_classes))
+                X_new = []
+                y_new = []
+                for class_ in unseen_classes:
+                    index = torch.multinomial((y == class_).float(), 1)
+                    X_new.append(X[index].squeeze(0))
+                    y_new.append(y[index].squeeze(0))
+                X_sample = torch.cat([X_sample, torch.stack(X_new)])
+                y_sample = torch.cat([y_sample, torch.stack(y_new)])
+            
+            tree = DecisionTree(max_depth=self.max_depth, min_samples_split=self.min_samples_split, criterion=self.criterion)
+            tree.fit(X_sample, y_sample)
+            prediction = tree.predict(X)
+            incorrect = prediction != y
+            eps = torch.sum(incorrect * weights)
+            errors.append(eps)
+
+            #  If better than random quessing, continue
+            if eps < 1 - 1 / self.n_classes:
+                alpha = 0.5 * (torch.log((1 - eps) / (eps + 1e-8)) + log(self.n_classes - 1))
+                weights = weights * torch.exp(alpha * incorrect)
+                weights /= weights.sum()  # keep weights as a probability distribution
+                self.confidences.append(alpha)
+            else:
+                if verbose: warn(f"The latest weak learner is worse than random guessing. The training is stopped to reduce over fitting. Only {i} learners are used.")
+                self.n_trees = i
+                break
+
+            trees.append(tree)
+        
+        self.trees = trees
+        return errors
+
+    def predict(self, X):
+        """
+        Applies the fitted AdaBoostClassifier model to the input data, predicting the correct classes.
+
+        Args:
+            X (torch.Tensor of shape (n_samples, n_features)): The input data to be classified.
+
+        Returns:
+            labels (torch.Tensor of shape (n_samples,)): The predicted labels corresponding to each sample.
+
+        Raises:
+            NotFittedError: If the AdaBoostClassifier model has not been fitted before predicting.
+            TypeError: If the input matrix is not a PyTorch tensor.
+            ValueError: If the input matrix is not the correct shape.
+        """
+        if self.trees is None:
+            raise NotFittedError("AdaBoostClassifier.fit() must be called before predicting.")
+        if not isinstance(X, torch.Tensor):
+            raise TypeError("The input matrix must be a PyTorch tensor.")
+        if X.ndim != 2 or X.shape[1] != self.n_features:
+            raise ValueError("The input matrix must be a 2 dimensional tensor with the same number of features as the fitted tensor.")
+
+        preds = torch.zeros((len(X), self.n_classes))
+        for alpha, tree in zip(self.confidences, self.trees):
+            pred = tree.predict(X).int()
+            one_hot_pred = torch.eye(self.n_classes)[pred]
+            preds += alpha * one_hot_pred
+        
+        return preds.argmax(dim=1)
+    
+    def predict_proba(self, X):
+        """
+        Applies the fitted AdaBoostClassifier model to the input data, predicting the probabilities of each class.
+
+        Args:
+            X (torch.Tensor of shape (n_samples, n_features)): The input data for which to predict probabilities.
+
+        Returns:
+            torch.Tensor of shape (n_samples, n_classes): The predicted probabilities for each class.
+        
+        Raises:
+            NotFittedError: If the AdaBoostClassifier model has not been fitted before predicting.
+            TypeError: If the input matrix is not a PyTorch tensor.
+            ValueError: If the input matrix is not the correct shape.
+        """
+        if self.trees is None:
+            raise NotFittedError("AdaBoostClassifier.fit() must be called before predicting.")
+        if not isinstance(X, torch.Tensor):
+            raise TypeError("The input matrix must be a PyTorch tensor.")
+        if X.ndim != 2 or X.shape[1] != self.n_features:
+            raise ValueError("The input matrix must be a 2 dimensional tensor with the same number of features as the fitted tensor.")
+
+        preds = torch.zeros((len(X), self.n_classes))
+        for alpha, tree in zip(self.confidences, self.trees):
+            pred = tree.predict(X).int()
+            one_hot_pred = torch.eye(self.n_classes)[pred]
+            preds += alpha * one_hot_pred
+
+        probs = torch.exp(preds - torch.max(preds, dim=1, keepdim=True).values)
+        probs = probs / torch.sum(probs, dim=1, keepdim=True)
+        if self.n_classes == 2:
+            return probs[:, 1]
+        return probs

DLL/MachineLearning/SupervisedLearning/Trees/_AdaBoostRegressor.py

@@ -0,0 +1,141 @@
+import torch
+from warnings import warn
+
+from ._RegressionTree import RegressionTree
+from ....Exceptions import NotFittedError
+
+
+class AdaBoostRegressor:
+    """
+    AdaBoostRegressor implements a regression algorithm fitting many consecutive :class:`RegressionTrees <DLL.MachineLearning.SupervisedLearning.Trees.RegressionTree>` to previously incorrectly predicted samples.
+
+    Args:
+        n_trees (int, optional): The number of trees used for predicting. Defaults to 10. Must be a positive integer.
+        max_depth (int, optional): The maximum depth of the tree. Defaults to 25. Must be a positive integer.
+        min_samples_split (int, optional): The minimum required samples in a leaf to make a split. Defaults to 2. Must be a positive integer.
+        loss (str, optional): The loss function used. Must be in ["linear", "square", "exponential"]. Defaults to "square".
+    Attributes:
+        n_features (int): The number of features. Available after fitting.
+    """
+    def __init__(self, n_trees=10, max_depth=25, min_samples_split=2, loss="square"):
+        if not isinstance(n_trees, int) or n_trees < 1:
+            raise ValueError("n_trees must be a positive integer.")
+        if not isinstance(max_depth, int) or max_depth < 1:
+            raise ValueError("max_depth must be a positive integer.")
+        if not isinstance(min_samples_split, int) or min_samples_split < 1:
+            raise ValueError("min_samples_split must be a positive integer.")
+        if loss not in ["linear", "square", "exponential"]:
+            raise ValueError('loss must be in ["linear", "square", "exponential"].')
+        
+        self.n_trees = n_trees
+        self.max_depth = max_depth
+        self.min_samples_split = min_samples_split
+        self.loss = loss
+        self.trees = None
+
+    def fit(self, X, y, verbose=True):
+        """
+        Fits the AdaBoostRegressor model to the input data by fitting trees to the errors made by previous trees.
+
+        Args:
+            X (torch.Tensor of shape (n_samples, n_features)): The input data, where each row is a sample and each column is a feature.
+            y (torch.Tensor of shape (n_samples,)): The target values corresponding to each sample.
+            verbose (bool, optional): Determines if warnings are given if the training ends due to a weak learner having over 0.5 weighted loss. Defaults to True.
+        Returns:
+            The average errors after each tree.
+        Raises:
+            TypeError: If the input matrix or the label vector is not a PyTorch tensor or if the problem is binary and metrics is not a list or a tuple.
+            ValueError: If the input matrix or the label vector is not the correct shape or the label vector contains wrong values.
+        """
+        if not isinstance(X, torch.Tensor) or not isinstance(y, torch.Tensor):
+            raise TypeError("The input matrix and the label matrix must be a PyTorch tensor.")
+        if X.ndim != 2:
+            raise ValueError("The input matrix must be a 2 dimensional tensor.")
+        if y.ndim != 1 or y.shape[0] != X.shape[0]:
+            raise ValueError("The labels must be 1 dimensional with the same number of samples as the input data")
+        if not isinstance(verbose, bool):
+            raise TypeError("verbose must be a boolean.")
+        
+        y = y.to(X.dtype)
+        self.n_features = X.shape[1]
+
+        trees = []
+        weights = torch.full_like(y, 1 / len(y))
+        self.confidences = []
+        errors = []
+
+        for i in range(self.n_trees):
+            indices = torch.multinomial(weights, len(y), replacement=True)
+            X_sample = X[indices]
+            y_sample = y[indices]
+            
+            tree = RegressionTree(max_depth=self.max_depth, min_samples_split=self.min_samples_split)
+            tree.fit(X_sample, y_sample)
+            prediction = tree.predict(X)
+            max_absolute_error = torch.abs(y - prediction).max()
+            if self.loss == "linear":
+                error = torch.abs(y - prediction) / max_absolute_error
+            elif self.loss == "square":
+                error = (y - prediction) ** 2 / max_absolute_error ** 2
+            elif self.loss == "exponential":
+                error = 1 - torch.exp(-torch.abs(y - prediction) / max_absolute_error)
+            
+            average_loss = torch.sum(weights * error)
+            if average_loss > 0.5:
+                if verbose: warn(f"The average error exceeds 0.5. The training is stopped to reduce over fitting. Only {i} trees are used.")
+                self.n_trees = i
+                break
+
+            beta = average_loss / (1 - average_loss)
+            weights = weights * beta ** (1 - error)
+            weights /= weights.sum()
+
+            self.confidences.append(torch.log(1 / beta))
+            trees.append(tree)
+            errors.append(average_loss)
+        
+        self.confidences = torch.stack(self.confidences)
+        self.confidences /= self.confidences.sum()
+        self.trees = trees
+        return errors
+
+    def predict(self, X, method="average"):
+        """
+        Applies the fitted AdaBoostRegressor model to the input data, predicting the correct values.
+
+        Args:
+            X (torch.Tensor of shape (n_samples, n_features)): The input data to be regressed.
+            method (str, optional): The method for computing the prediction. Must be one of "average" or "weighted_median". Defaults to average.
+
+        Returns:
+            values (torch.Tensor of shape (n_samples,)): The predicted values corresponding to each sample.
+
+        Raises:
+            NotFittedError: If the AdaBoostRegressor model has not been fitted before predicting.
+            TypeError: If the input matrix is not a PyTorch tensor.
+            ValueError: If the input matrix is not the correct shape.
+        """
+        if self.trees is None:
+            raise NotFittedError("AdaBoostRegressor.fit() must be called before predicting.")
+        if not isinstance(X, torch.Tensor):
+            raise TypeError("The input matrix must be a PyTorch tensor.")
+        if X.ndim != 2 or X.shape[1] != self.n_features:
+            raise ValueError("The input matrix must be a 2 dimensional tensor with the same number of features as the fitted tensor.")
+        if method not in ["average", "weighted_median"]:
+            raise ValueError('method must be one of "average" or "weighted_median".')
+
+        if method == "average":
+            preds = torch.zeros(len(X))
+            for log_beta, tree in zip(self.confidences, self.trees):
+                pred = tree.predict(X)
+                preds += log_beta * pred
+        elif method == "weighted_median":
+            tree_preds = torch.stack([tree.predict(X) for tree in self.trees]).T  # n_samples, n_trees
+            indicies = torch.argsort(tree_preds, dim=1)
+            accumulated = torch.cumsum(self.confidences[indicies], dim=1)
+            after_middle = accumulated >= 0.5 * accumulated[:, -1]
+            middle_indicies = torch.argmax(after_middle, dim=1)
+            original_middle_indicies = indicies[:, middle_indicies]
+            preds = tree_preds[:, original_middle_indicies]
+        
+        return preds

DLL/MachineLearning/SupervisedLearning/Trees/_DecisionTree.py

@@ -0,0 +1,219 @@
+import torch
+
+from ....Exceptions import NotFittedError
+
+
+class Node:
+    def __init__(self, left=None, right=None, threshold=None, feature_index=None, value=None, probabilities=None, size=None, depth=None):
+        self.left = left
+        self.right = right
+        self.threshold = threshold
+        self.feature_index = feature_index
+        self.value = value
+        self.probabilities = probabilities
+        self.size = size
+        self.depth = depth
+
+    def is_leaf(self):
+        return self.value is not None
+
+
+class DecisionTree:
+    """
+    DecisionTree implements a classification algorithm splitting the data along features yielding the maximum entropy.
+
+    Args:
+        max_depth (int, optional): The maximum depth of the tree. Defaults to 10. Must be a positive integer.
+        min_samples_split (int, optional): The minimum required samples in a leaf to make a split. Defaults to 2. Must be a positive integer.
+        criterion (str, optional): The information criterion used to select optimal splits. Must be one of "entropy" or "gini". Defaults to "gini".
+        ccp_alpha (non-negative float, optional): Determines how easily subtrees are pruned in cost-complexity pruning. The larger the value, more subtrees are pruned. Defaults to 0.0.
+    Attributes:
+        n_classes (int): The number of classes. A positive integer available after calling DecisionTree.fit().
+        classes (torch.Tensor of shape (n_classes,)): The classes in an arbitrary order. Available after calling DecisionTree.fit().
+    """
+    def __init__(self, max_depth=10, min_samples_split=2, criterion="gini", ccp_alpha=0.0):
+        if not isinstance(max_depth, int) or max_depth < 1:
+            raise ValueError("max_depth must be a positive integer.")
+        if not isinstance(min_samples_split, int) or min_samples_split < 1:
+            raise ValueError("min_samples_split must be a positive integer.")
+        if criterion not in ["entropy", "gini"]:
+            raise ValueError('The chosen criterion must be one of "entropy" or "gini".')
+        if not isinstance(ccp_alpha, int | float) or ccp_alpha < 0:
+            raise ValueError("ccp_alpha must be non-negative.")
+
+        self.max_depth = max_depth
+        self.min_samples_split = min_samples_split
+        self.criterion = criterion
+        self.root = None
+        self.ccp_alpha = ccp_alpha
+    
+    def fit(self, X, y):
+        """
+        Fits the DecisionTree model to the input data by generating a tree, which splits the data appropriately.
+
+        Args:
+            X (torch.Tensor of shape (n_samples, n_features)): The input data, where each row is a sample and each column is a feature.
+            y (torch.Tensor of shape (n_samples,)): The labels corresponding to each sample.
+        Returns:
+            None
+        Raises:
+            TypeError: If the input matrix or the label matrix is not a PyTorch tensor.
+            ValueError: If the input matrix or the label matrix is not the correct shape.
+        """
+        if not isinstance(X, torch.Tensor) or not isinstance(y, torch.Tensor):
+            raise TypeError("The input matrix and the label matrix must be a PyTorch tensor.")
+        if X.ndim != 2:
+            raise ValueError("The input matrix must be a 2 dimensional tensor.")
+        if y.ndim != 1 or y.shape[0] != X.shape[0]:
+            raise ValueError("The labels must be 1 dimensional with the same number of samples as the input data")
+        self.n_samples, self.n_features = X.shape
+        self.classes = torch.unique(y)
+        self.n_classes = len(self.classes)
+        self.root = self._grow_tree(X, y, 0)
+        if self.ccp_alpha > 0:
+            self.root, _, _ = self._prune(self.root, X, y)
+
+    def _grow_tree(self, x, y, depth):
+        n_samples, n_features = x.size()
+        classes, counts = torch.unique(y, return_counts=True)
+
+        if depth >= self.max_depth or len(classes) == 1 or n_samples < self.min_samples_split:
+            largest_class = classes[torch.argmax(counts)]
+            probabilities = torch.tensor([(counts[torch.where(classes == _class)[0]] / len(y) if _class in classes else 0) for _class in self.classes])
+            return Node(value=largest_class, probabilities=probabilities)
+
+        feature_indicies = torch.randperm(n_features)[:self.n_features]
+        split_threshold, split_index = self._best_split(x, y, feature_indicies)
+        # if no split gains more information
+        if split_threshold is None:
+            largest_class = classes[torch.argmax(counts)]
+            probabilities = torch.tensor([(counts[torch.where(classes == _class)[0]] / len(y) if _class in classes else 0) for _class in self.classes])
+            return Node(value=largest_class, probabilities=probabilities)
+        
+        left_indicies, right_indicies = self._split(x[:, split_index], split_threshold)
+        left = self._grow_tree(x[left_indicies], y[left_indicies], depth + 1)
+        right = self._grow_tree(x[right_indicies], y[right_indicies], depth + 1)
+        return Node(left, right, split_threshold, split_index)
+    
+    def _best_split(self, x, y, feature_indicies):
+        max_entropy_gain = -1
+        split_index = None
+        split_threshold = None
+
+        for index in feature_indicies:
+            feature_values = x[:, index]
+            possible_thresholds = torch.unique(feature_values)
+            for threshold in possible_thresholds:
+                entropy_gain = self._information_gain(y, feature_values, threshold)
+                if entropy_gain > max_entropy_gain:
+                    max_entropy_gain = entropy_gain
+                    split_index = index
+                    split_threshold = threshold
+        if max_entropy_gain <= 0: return None, None
+        return split_threshold, split_index
+
+    def _split(self, feature_values, threshold):
+        left_indicies = torch.argwhere(feature_values <= threshold).flatten()
+        right_indicies = torch.argwhere(feature_values > threshold).flatten()
+        return left_indicies, right_indicies
+
+    def _information_gain(self, y, feature_values, threshold):
+        if self.criterion == "gini":
+            information_func = self._gini
+        elif self.criterion == "entropy":
+            information_func = self._entropy
+
+        left_indicies, right_indicies = self._split(feature_values, threshold)
+        if len(left_indicies) == 0 or len(right_indicies) == 0:
+            return 0
+        n = len(y)
+        return information_func(y) - len(left_indicies) / n * information_func(y[left_indicies]) - len(right_indicies) / n * information_func(y[right_indicies])
+    
+    def _entropy(self, values):
+        n = len(values)
+        data_type = values.dtype
+        p = torch.bincount(values.to(dtype=torch.int32)).to(dtype=data_type) / n
+        p = p[p != 0]  # remove non-zero classes from messing up the calculations
+        return -torch.sum(p * torch.log(p))
+    
+    def _gini(self, values):
+        n = len(values)
+        data_type = values.dtype
+        p = torch.bincount(values.to(dtype=torch.int32)).to(dtype=data_type) / n
+        return torch.sum(p * (1 - p))
+
+    def predict(self, X):
+        """
+        Applies the fitted DecisionTree model to the input data, predicting the correct classes.
+
+        Args:
+            X (torch.Tensor of shape (n_samples, n_features)): The input data to be classified.
+        Returns:
+            labels (torch.Tensor of shape (n_samples,)): The predicted labels corresponding to each sample.
+        Raises:
+            NotFittedError: If the DecisionTree model has not been fitted before predicting.
+            TypeError: If the input matrix is not a PyTorch tensor.
+            ValueError: If the input matrix is not the correct shape.
+        """
+        if self.root is None:
+            raise NotFittedError("DecisionTree.fit() must be called before predicting.")
+        if not isinstance(X, torch.Tensor):
+            raise TypeError("The input matrix must be a PyTorch tensor.")
+        if X.ndim != 2 or X.shape[1] != self.n_features:
+            raise ValueError("The input matrix must be a 2 dimensional tensor with the same number of features as the fitted tensor.")
+        return torch.tensor([self._predict_single(datapoint, self.root) for datapoint in X])
+
+    def _predict_single(self, x, current_node, proba=False):
+        if current_node.is_leaf():
+            if proba:
+                return current_node.probabilities
+            return current_node.value
+        
+        if x[current_node.feature_index] <= current_node.threshold:
+            return self._predict_single(x, current_node.left, proba=proba)
+        
+        return self._predict_single(x, current_node.right, proba=proba)
+    
+    def predict_proba(self, X):
+        """
+        Applies the fitted DecisionTree model to the input data, predicting the probabilities of each class.
+
+        Args:
+            X (torch.Tensor of shape (n_samples, n_features)): The input data to be classified.
+        Returns:
+            probabilities (torch.Tensor of shape (n_samples, n_classes)): The predicted probabilities corresponding to each sample.
+        Raises:
+            NotFittedError: If the DecisionTree model has not been fitted before predicting.
+            TypeError: If the input matrix is not a PyTorch tensor.
+            ValueError: If the input matrix is not the correct shape.
+        """
+        if self.root is None:
+            raise NotFittedError("DecisionTree.fit() must be called before predicting.")
+        if not isinstance(X, torch.Tensor):
+            raise TypeError("The input matrix must be a PyTorch tensor.")
+        if X.ndim != 2 or X.shape[1] != self.n_features:
+            raise ValueError("The input matrix must be a 2 dimensional tensor with the same number of features as the fitted tensor.")
+        return torch.stack([self._predict_single(datapoint, self.root, proba=True) for datapoint in X])
+
+    def _prune(self, node, X, y):
+        classes, counts = torch.unique(y, return_counts=True)
+        m = len(y) - torch.min(counts).item()
+        if node.is_leaf():
+            return node, m / self.n_samples, 1
+        
+        left_indicies, right_indicies = self._split(X[:, node.feature_index], node.threshold)
+        node.left, left_cost, left_leaf_nodes = self._prune(node.left, X[left_indicies], y[left_indicies])
+        node.right, right_cost, right_leaf_nodes = self._prune(node.right, X[right_indicies], y[right_indicies])
+        subtree_cost = left_cost + right_cost
+        total_subleaves = left_leaf_nodes + right_leaf_nodes
+
+        leaf_miss_classification_error_rate = m / self.n_samples
+        new_leaf_nodes = 1
+
+        subtree_cost_with_n_nodes_penalty = subtree_cost + self.ccp_alpha * total_subleaves
+        
+        if leaf_miss_classification_error_rate + self.ccp_alpha * new_leaf_nodes < subtree_cost_with_n_nodes_penalty:
+            largest_class = classes[torch.argmax(counts)]
+            probabilities = torch.tensor([(counts[torch.where(classes == _class)[0]] / len(y) if _class in classes else 0) for _class in self.classes])
+            return Node(value=largest_class, probabilities=probabilities), leaf_miss_classification_error_rate, 1
+        return node, subtree_cost, total_subleaves