From bdf85fcadf326ee2cf60752c7517a81d7cf85b57 Mon Sep 17 00:00:00 2001 From: Mikhail Martin Date: Tue, 30 Sep 2025 23:39:48 +0300 Subject: [PATCH] CyColumnSplitter.information_gain() -> Criterion.impurity_decrease() --- smarttree/_column_splitter.py | 19 ++++----- smarttree/_criterion.pxd | 6 +++ smarttree/_criterion.pyi | 42 +++++++++++++++++++ smarttree/_criterion.pyx | 43 +++++++++++++++++++ smarttree/_cy_column_splitter.pxd | 11 ----- smarttree/_cy_column_splitter.pyi | 52 ----------------------- smarttree/_cy_column_splitter.pyx | 70 ------------------------------- smarttree/_types.py | 7 ---- 8 files changed, 99 insertions(+), 151 deletions(-) delete mode 100644 smarttree/_cy_column_splitter.pxd delete mode 100644 smarttree/_cy_column_splitter.pyi delete mode 100644 smarttree/_cy_column_splitter.pyx diff --git a/smarttree/_column_splitter.py b/smarttree/_column_splitter.py index e44eed9..6b963ba 100644 --- a/smarttree/_column_splitter.py +++ b/smarttree/_column_splitter.py @@ -8,10 +8,10 @@ import numpy as np from numpy.typing import NDArray -from ._cy_column_splitter import CyBaseColumnSplitter +from ._criterion import ClassificationCriterion, Entropy, Gini from ._dataset import Dataset from ._tree import TreeNode -from ._types import ClassificationCriterionType, Criterion, NaModeType +from ._types import ClassificationCriterionType, NaModeType NO_INFORMATION_GAIN = float("-inf") @@ -35,12 +35,6 @@ def no_split(cls) -> ColumnSplitResult: class BaseColumnSplitter(ABC): - mapping: dict[ClassificationCriterionType, Criterion] = { - "gini": Criterion.GINI, - "entropy": Criterion.ENTROPY, - "log_loss": Criterion.LOG_LOSS, - } - def __init__( self, dataset: Dataset, @@ -51,7 +45,11 @@ def __init__( ) -> None: self.dataset = dataset - self.criterion = self.mapping[criterion] + self.criterion: ClassificationCriterion + if criterion == "gini": + self.criterion = Gini(dataset) + else: # "entropy" | "log_loss" + self.criterion = Entropy(dataset) self.min_samples_split = min_samples_split self.min_samples_leaf = min_samples_leaf self.feature_na_mode = feature_na_mode @@ -166,8 +164,7 @@ def information_gain( \item $\text{impurity}_{\text{child}_i}$ — child node impurity. \end{itemize} """ - cs = CyBaseColumnSplitter(self.dataset, self.criterion) - return cs.information_gain(parent_mask, child_masks, normalize) + return self.criterion.impurity_decrease(parent_mask, child_masks, normalize) class NumColumnSplitter(BaseColumnSplitter): diff --git a/smarttree/_criterion.pxd b/smarttree/_criterion.pxd index cb17bdc..bc3984a 100644 --- a/smarttree/_criterion.pxd +++ b/smarttree/_criterion.pxd @@ -7,6 +7,12 @@ cdef class ClassificationCriterion: cdef Py_ssize_t n_classes cdef Py_ssize_t n_samples + cpdef double impurity_decrease( + self, + cnp.npy_bool[:] parent_mask, + list[cnp.npy_bool[:]] child_masks, + bint normalize, + ) cpdef cnp.int64_t[:] distribution(self, cnp.npy_bool[:] mask) diff --git a/smarttree/_criterion.pyi b/smarttree/_criterion.pyi index bdd250e..23b6910 100644 --- a/smarttree/_criterion.pyi +++ b/smarttree/_criterion.pyi @@ -10,6 +10,48 @@ class ClassificationCriterion(ABC): def __init__(self, dataset: Dataset) -> None: ... + def impurity_decrease( + self, + parent_mask: NDArray[np.bool_], + child_masks: list[NDArray[np.bool_]], + normalize: bool = False, + ) -> float: + r""" + Calculates information gain of the split. + + Parameters: + parent_mask: pd.Series + boolean mask of parent node. + child_masks: pd.Series + list of boolean masks of child nodes. + normalize: bool, default=False + if True, normalizes information gain by split factor to handle + unbalanced splits. Uses child node counts for normalization. + + Returns: + float: information gain. + + Formula in LaTeX: + \begin{align*} + \text{Information Gain} = + \frac{N_{\text{parent}}}{N} \cdot + \Biggl( & \text{impurity}_{\text{parent}} - \\ + & \sum^C_{i=1} \frac{N_{\text{child}_i}}{N_{\text{parent}}} + \cdot \text{impurity}_{\text{child}_i} \Biggr) + \end{align*} + where: + \begin{itemize} + \item $\text{Information Gain}$ — information gain; + \item $N$ — number of samples in entire training set; + \item $N_{\text{parent}}$ — number of samples in parent node; + \item $\text{impurity}_{\text{parent}}$ — parent node impurity; + \item $C$ — number of child nodes; + \item $N_{\text{child}_i}$ — number of samples in child node; + \item $\text{impurity}_{\text{child}_i}$ — child node impurity. + \end{itemize} + """ + ... + @abstractmethod def impurity(self, mask: NDArray[np.bool_]) -> float: raise NotImplementedError diff --git a/smarttree/_criterion.pyx b/smarttree/_criterion.pyx index b42c5bc..7cf6e4b 100644 --- a/smarttree/_criterion.pyx +++ b/smarttree/_criterion.pyx @@ -17,6 +17,49 @@ cdef class ClassificationCriterion: self.n_classes = len(dataset.classes) self.n_samples = len(dataset.y) + cpdef double impurity_decrease( + self, + cnp.npy_bool[:] parent_mask, + list[cnp.npy_bool[:]] child_masks, + bint normalize, + ): + + cdef Py_ssize_t i, j, n_childs + cdef long N, N_parent, N_childs, N_child_j + cdef double impurity_parent, weighted_impurity_childs, impurity_child_i, norm_coef, local_information_gain, information_gain + + N = 0 + N_parent = 0 + for i in range(self.n_samples): + N += 1 + if parent_mask[i]: + N_parent += 1 + + impurity_parent = self.impurity(parent_mask) + + N_childs = 0 + n_childs = len(child_masks) + weighted_impurity_childs = 0.0 + for j in range(n_childs): + N_child_j = 0 + child_mask = child_masks[j] + for i in range(self.n_samples): + if child_mask[i]: + N_child_j += 1 + N_childs += N_child_j + impurity_child_i = self.impurity(child_mask) + weighted_impurity_childs += (N_child_j / N_parent) * impurity_child_i + + if normalize: + norm_coef = N_parent / N_childs + weighted_impurity_childs *= norm_coef + + local_information_gain = impurity_parent - weighted_impurity_childs + + information_gain = (N_parent / N) * local_information_gain + + return information_gain + @cython.boundscheck(False) @cython.wraparound(False) cpdef cnp.int64_t[:] distribution(self, cnp.npy_bool[:] mask): diff --git a/smarttree/_cy_column_splitter.pxd b/smarttree/_cy_column_splitter.pxd deleted file mode 100644 index 01dce08..0000000 --- a/smarttree/_cy_column_splitter.pxd +++ /dev/null @@ -1,11 +0,0 @@ -cimport numpy as cnp - -from ._criterion cimport ClassificationCriterion - - -cdef class CyBaseColumnSplitter: - - cdef ClassificationCriterion criterion - cdef cnp.int64_t[:] y - cdef Py_ssize_t n_classes - cdef Py_ssize_t n_samples diff --git a/smarttree/_cy_column_splitter.pyi b/smarttree/_cy_column_splitter.pyi deleted file mode 100644 index e543199..0000000 --- a/smarttree/_cy_column_splitter.pyi +++ /dev/null @@ -1,52 +0,0 @@ -import numpy as np -from numpy.typing import NDArray - -from ._dataset import Dataset -from ._types import Criterion - -class CyBaseColumnSplitter: - - def __init__(self, dataset: Dataset, criterion: Criterion) -> None: - ... - - def information_gain( - self, - parent_mask: NDArray[np.bool_], - child_masks: list[NDArray[np.bool_]], - normalize: bool = False, - ) -> float: - r""" - Calculates information gain of the split. - - Parameters: - parent_mask: pd.Series - boolean mask of parent node. - child_masks: pd.Series - list of boolean masks of child nodes. - normalize: bool, default=False - if True, normalizes information gain by split factor to handle - unbalanced splits. Uses child node counts for normalization. - - Returns: - float: information gain. - - Formula in LaTeX: - \begin{align*} - \text{Information Gain} = - \frac{N_{\text{parent}}}{N} \cdot - \Biggl( & \text{impurity}_{\text{parent}} - \\ - & \sum^C_{i=1} \frac{N_{\text{child}_i}}{N_{\text{parent}}} - \cdot \text{impurity}_{\text{child}_i} \Biggr) - \end{align*} - where: - \begin{itemize} - \item $\text{Information Gain}$ — information gain; - \item $N$ — number of samples in entire training set; - \item $N_{\text{parent}}$ — number of samples in parent node; - \item $\text{impurity}_{\text{parent}}$ — parent node impurity; - \item $C$ — number of child nodes; - \item $N_{\text{child}_i}$ — number of samples in child node; - \item $\text{impurity}_{\text{child}_i}$ — child node impurity. - \end{itemize} - """ - ... diff --git a/smarttree/_cy_column_splitter.pyx b/smarttree/_cy_column_splitter.pyx deleted file mode 100644 index 71a883b..0000000 --- a/smarttree/_cy_column_splitter.pyx +++ /dev/null @@ -1,70 +0,0 @@ -cimport cython - -import numpy as np -cimport numpy as cnp -import pandas as pd - -from ._dataset import Dataset -from ._types import Criterion -from ._criterion cimport Entropy, Gini - - -cnp.import_array() - -cdef int CRITERION_GINI = 1 - - -cdef class CyBaseColumnSplitter: - - def __cinit__(self, dataset: Dataset, criterion: Criterion) -> None: - self.y = dataset.y - self.n_classes = len(dataset.classes) - self.n_samples = len(dataset.y) - if criterion.value == CRITERION_GINI: - self.criterion = Gini(dataset) - else: - self.criterion = Entropy(dataset) - - def information_gain( - self, - cnp.npy_bool[:] parent_mask, - list[cnp.npy_bool[:]] child_masks, - bint normalize, - ): - - cdef Py_ssize_t i, j, n_childs - cdef long N, N_parent, N_childs, N_child_j - cdef double impurity_parent, weighted_impurity_childs, impurity_child_i - - N = 0 - N_parent = 0 - for i in range(self.n_samples): - N += 1 - if parent_mask[i]: - N_parent += 1 - - impurity_parent = self.criterion.impurity(parent_mask) - - N_childs = 0 - n_childs = len(child_masks) - weighted_impurity_childs = 0.0 - for j in range(n_childs): - N_child_j = 0 - child_mask = child_masks[j] - for i in range(self.n_samples): - if child_mask[i]: - N_child_j += 1 - N_childs += N_child_j - impurity_child_i = self.criterion.impurity(child_mask) - weighted_impurity_childs += (N_child_j / N_parent) * impurity_child_i - - cdef double norm_coef - if normalize: - norm_coef = N_parent / N_childs - weighted_impurity_childs *= norm_coef - - cdef double local_information_gain = impurity_parent - weighted_impurity_childs - - cdef double information_gain = (N_parent / N) * local_information_gain - - return information_gain diff --git a/smarttree/_types.py b/smarttree/_types.py index 9ffa7b8..927e113 100644 --- a/smarttree/_types.py +++ b/smarttree/_types.py @@ -1,4 +1,3 @@ -from enum import Enum from typing import Literal @@ -12,9 +11,3 @@ VerboseType = Literal["CRITICAL", "ERROR", "WARNING", "INFO", "DEBUG"] | int SplitType = Literal["numerical", "categorical", "rank"] - - -class Criterion(Enum): - GINI = 1 - ENTROPY = 2 - LOG_LOSS = 2