From cd38312b55dc3d4c13a95b46b8fa0107b72fc307 Mon Sep 17 00:00:00 2001 From: Mikhail Martin Date: Sun, 28 Sep 2025 17:12:08 +0300 Subject: [PATCH 1/3] memoryview --- build-extension.py | 8 ++++---- pyproject.toml | 2 +- smarttree/_builder.py | 5 ++--- smarttree/_criterion.pxd | 8 ++++---- smarttree/_criterion.pyi | 8 ++++---- smarttree/_criterion.pyx | 11 +++++++---- smarttree/_cy_column_splitter.pyx | 12 ++++++------ tests/test__criterion.py | 4 ++-- 8 files changed, 30 insertions(+), 28 deletions(-) diff --git a/build-extension.py b/build-extension.py index 992f5a5..ede3fd9 100644 --- a/build-extension.py +++ b/build-extension.py @@ -11,17 +11,17 @@ from setuptools import Extension from setuptools.command.build_ext import build_ext +import numpy as np + if sys.platform == "win32": COMPILE_ARGS = ["/O2", "/fp:fast"] - LINK_ARGS = [] - INCLUDE_DIRS = [] LIBRARIES = [] else: COMPILE_ARGS = ["-march=native", "-O3", "-msse", "-msse2", "-mfma", "-mfpmath=sse"] - LINK_ARGS = [] - INCLUDE_DIRS = [] LIBRARIES = ["m"] +LINK_ARGS = [] +INCLUDE_DIRS = [np.get_include()] def build() -> None: diff --git a/pyproject.toml b/pyproject.toml index ab9cc82..bec4ad7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,5 +1,5 @@ [build-system] -requires = ["poetry-core", "cython", "setuptools"] +requires = ["poetry-core", "setuptools", "cython", "numpy"] build-backend = "poetry.core.masonry.api" [tool.poetry] diff --git a/smarttree/_builder.py b/smarttree/_builder.py index ca37236..b11f8e6 100644 --- a/smarttree/_builder.py +++ b/smarttree/_builder.py @@ -1,6 +1,5 @@ import bisect -import numpy as np import pandas as pd from ._criterion import ClassificationCriterion, Entropy, Gini @@ -45,7 +44,7 @@ def build(self, tree: Tree) -> None: self.available_features.remove(value) mask = self.y.apply(lambda x: True) - mask_np = mask.to_numpy(dtype=np.int8) + mask_np = mask.to_numpy() root = tree.create_node( mask=mask, hierarchy=self.hierarchy, @@ -76,7 +75,7 @@ def build(self, tree: Tree) -> None: else: # str node.available_features.append(value) - child_mask_np = child_mask.to_numpy(dtype=np.int8) + child_mask_np = child_mask.to_numpy() child_node = tree.create_node( mask=child_mask, hierarchy=node.hierarchy, diff --git a/smarttree/_criterion.pxd b/smarttree/_criterion.pxd index 138e718..6012cf5 100644 --- a/smarttree/_criterion.pxd +++ b/smarttree/_criterion.pxd @@ -1,4 +1,4 @@ -from libc.stdint cimport int8_t +cimport numpy as cnp cdef class ClassificationCriterion: @@ -7,11 +7,11 @@ cdef class ClassificationCriterion: cdef Py_ssize_t n_classes cdef Py_ssize_t n_samples - cpdef long[:] distribution(self, int8_t[:] mask) + cpdef long[:] distribution(self, cnp.npy_bool[:] mask) cdef class Gini(ClassificationCriterion): - cpdef double impurity(self, int8_t[:] mask) + cpdef double impurity(self, cnp.npy_bool[:] mask) cdef class Entropy(ClassificationCriterion): - cpdef double impurity(self, int8_t[:] mask) + cpdef double impurity(self, cnp.npy_bool[:] mask) diff --git a/smarttree/_criterion.pyi b/smarttree/_criterion.pyi index 326571d..509bf01 100644 --- a/smarttree/_criterion.pyi +++ b/smarttree/_criterion.pyi @@ -11,16 +11,16 @@ class ClassificationCriterion(ABC): ... @abstractmethod - def impurity(self, mask: NDArray[np.int8]) -> float: + def impurity(self, mask: NDArray[np.bool_]) -> float: raise NotImplementedError - def distribution(self, mask: NDArray[np.int8]) -> NDArray[np.int32]: + def distribution(self, mask: NDArray[np.bool_]) -> NDArray[np.int32]: ... class Gini(ClassificationCriterion): - def impurity(self, mask: NDArray[np.int8]) -> float: + def impurity(self, mask: NDArray[np.bool_]) -> float: r""" Calculates Gini index in a tree node. @@ -35,7 +35,7 @@ class Gini(ClassificationCriterion): class Entropy(ClassificationCriterion): - def impurity(self, mask: NDArray[np.int8]) -> float: + def impurity(self, mask: NDArray[np.bool_]) -> float: r""" Calculates entropy in a tree node. diff --git a/smarttree/_criterion.pyx b/smarttree/_criterion.pyx index 6d99a58..7410cd6 100644 --- a/smarttree/_criterion.pyx +++ b/smarttree/_criterion.pyx @@ -1,12 +1,15 @@ cimport cython from libc.math cimport log2 -from libc.stdint cimport int8_t import numpy as np +cimport numpy as cnp from ._dataset import Dataset +cnp.import_array() + + cdef class ClassificationCriterion: def __cinit__(self, dataset: Dataset) -> None: @@ -16,7 +19,7 @@ cdef class ClassificationCriterion: @cython.boundscheck(False) @cython.wraparound(False) - cpdef long[:] distribution(self, int8_t[:] mask): + cpdef long[:] distribution(self, cnp.npy_bool[:] mask): cdef Py_ssize_t i cdef long[:] result @@ -34,7 +37,7 @@ cdef class Gini(ClassificationCriterion): @cython.boundscheck(False) @cython.cdivision(True) @cython.wraparound(False) - cpdef double impurity(self, int8_t[:] mask): + cpdef double impurity(self, cnp.npy_bool[:] mask): cdef Py_ssize_t i cdef long[:] distribution @@ -60,7 +63,7 @@ cdef class Entropy(ClassificationCriterion): @cython.boundscheck(False) @cython.cdivision(True) @cython.wraparound(False) - cpdef double impurity(self, int8_t[:] mask): + cpdef double impurity(self, cnp.npy_bool[:] mask): cdef Py_ssize_t i cdef long[:] distribution diff --git a/smarttree/_cy_column_splitter.pyx b/smarttree/_cy_column_splitter.pyx index 6f7d724..79e3dd0 100644 --- a/smarttree/_cy_column_splitter.pyx +++ b/smarttree/_cy_column_splitter.pyx @@ -1,7 +1,7 @@ cimport cython -from libc.stdint cimport int8_t import numpy as np +cimport numpy as cnp import pandas as pd from ._dataset import Dataset @@ -9,6 +9,8 @@ from ._types import Criterion from ._criterion cimport Entropy, Gini +cnp.import_array() + cdef int CRITERION_GINI = 1 @@ -30,15 +32,13 @@ cdef class CyBaseColumnSplitter: normalize: bool = False, ) -> float: - cdef int8_t[:] parent_mask_arr, child_mask_arr + cdef cnp.npy_bool[:] parent_mask_arr, child_mask_arr 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 - parent_mask_arr = parent_mask.to_numpy(dtype=np.int8) - child_mask_arrs = [ - child_mask.to_numpy(dtype=np.int8) for child_mask in child_masks - ] + parent_mask_arr = parent_mask.to_numpy() + child_mask_arrs = [child_mask.to_numpy() for child_mask in child_masks] N = 0 N_parent = 0 diff --git a/tests/test__criterion.py b/tests/test__criterion.py index 4a690ea..6213605 100644 --- a/tests/test__criterion.py +++ b/tests/test__criterion.py @@ -5,13 +5,13 @@ def test__gini(dataset): gini_criterion = Gini(dataset) - mask = np.ones(dataset.y.shape, dtype=np.int8) + mask = np.ones(dataset.y.shape, dtype=np.bool_) gini_index = gini_criterion.impurity(mask) assert gini_index == 0.6666591342419322 def test__entropy(dataset): entropy_criterion = Entropy(dataset) - mask = np.ones(dataset.y.shape, dtype=np.int8) + mask = np.ones(dataset.y.shape, dtype=np.bool_) entropy = entropy_criterion.impurity(mask) assert entropy == 1.584946181877191 From 43cdc0ea5c5eb7f4d928a4dbed289cd88017def7 Mon Sep 17 00:00:00 2001 From: Mikhail Martin Date: Sun, 28 Sep 2025 18:05:00 +0300 Subject: [PATCH 2/3] typing --- smarttree/_criterion.pxd | 4 ++-- smarttree/_criterion.pyi | 2 +- smarttree/_criterion.pyx | 38 +++++++++++++++++++------------ smarttree/_cy_column_splitter.pxd | 4 ++-- smarttree/_dataset.py | 2 +- 5 files changed, 29 insertions(+), 21 deletions(-) diff --git a/smarttree/_criterion.pxd b/smarttree/_criterion.pxd index 6012cf5..cb17bdc 100644 --- a/smarttree/_criterion.pxd +++ b/smarttree/_criterion.pxd @@ -3,11 +3,11 @@ cimport numpy as cnp cdef class ClassificationCriterion: - cdef int[:] y + cdef cnp.int64_t[:] y cdef Py_ssize_t n_classes cdef Py_ssize_t n_samples - cpdef long[:] distribution(self, cnp.npy_bool[:] mask) + cpdef cnp.int64_t[:] distribution(self, cnp.npy_bool[:] mask) cdef class Gini(ClassificationCriterion): diff --git a/smarttree/_criterion.pyi b/smarttree/_criterion.pyi index 509bf01..bdd250e 100644 --- a/smarttree/_criterion.pyi +++ b/smarttree/_criterion.pyi @@ -14,7 +14,7 @@ class ClassificationCriterion(ABC): def impurity(self, mask: NDArray[np.bool_]) -> float: raise NotImplementedError - def distribution(self, mask: NDArray[np.bool_]) -> NDArray[np.int32]: + def distribution(self, mask: NDArray[np.bool_]) -> NDArray[np.int64]: ... diff --git a/smarttree/_criterion.pyx b/smarttree/_criterion.pyx index 7410cd6..b42c5bc 100644 --- a/smarttree/_criterion.pyx +++ b/smarttree/_criterion.pyx @@ -19,15 +19,19 @@ cdef class ClassificationCriterion: @cython.boundscheck(False) @cython.wraparound(False) - cpdef long[:] distribution(self, cnp.npy_bool[:] mask): + cpdef cnp.int64_t[:] distribution(self, cnp.npy_bool[:] mask): cdef Py_ssize_t i - cdef long[:] result + cdef cnp.int64_t[:] result + cdef cnp.npy_bool mask_value + cdef cnp.int64_t class_index - result = np.zeros(self.n_classes, dtype=np.int32) + result = np.zeros(self.n_classes, dtype=np.int64) for i in range(self.n_samples): - if mask[i]: - result[self.y[i]] += 1 + mask_value = mask[i] + if mask_value: + class_index = self.y[i] + result[class_index] += 1 return result @@ -40,19 +44,21 @@ cdef class Gini(ClassificationCriterion): cpdef double impurity(self, cnp.npy_bool[:] mask): cdef Py_ssize_t i - cdef long[:] distribution - cdef long N + cdef cnp.int64_t[:] distribution + cdef cnp.int64_t N, count cdef double p_i, gini distribution = self.distribution(mask) N = 0 for i in range(self.n_classes): - N += distribution[i] + count = distribution[i] + N += count gini = 1.0 for i in range(self.n_classes): - if distribution[i] > 0: - p_i = distribution[i] / N + count = distribution[i] + if count > 0: + p_i = count / N gini -= p_i * p_i return gini @@ -66,19 +72,21 @@ cdef class Entropy(ClassificationCriterion): cpdef double impurity(self, cnp.npy_bool[:] mask): cdef Py_ssize_t i - cdef long[:] distribution - cdef long N + cdef cnp.int64_t[:] distribution + cdef cnp.int64_t N, count cdef double p_i, gini distribution = self.distribution(mask) N = 0 for i in range(self.n_classes): - N += distribution[i] + count = distribution[i] + N += count entropy = 0.0 for i in range(self.n_classes): - if distribution[i] > 0: - p_i = distribution[i] / N + count = distribution[i] + if count > 0: + p_i = count / N entropy -= p_i * log2(p_i) return entropy diff --git a/smarttree/_cy_column_splitter.pxd b/smarttree/_cy_column_splitter.pxd index e297fb4..01dce08 100644 --- a/smarttree/_cy_column_splitter.pxd +++ b/smarttree/_cy_column_splitter.pxd @@ -1,4 +1,4 @@ -from libc.stdint cimport int8_t +cimport numpy as cnp from ._criterion cimport ClassificationCriterion @@ -6,6 +6,6 @@ from ._criterion cimport ClassificationCriterion cdef class CyBaseColumnSplitter: cdef ClassificationCriterion criterion - cdef int[:] y + cdef cnp.int64_t[:] y cdef Py_ssize_t n_classes cdef Py_ssize_t n_samples diff --git a/smarttree/_dataset.py b/smarttree/_dataset.py index e3a2811..5fcaf5b 100644 --- a/smarttree/_dataset.py +++ b/smarttree/_dataset.py @@ -7,7 +7,7 @@ class Dataset: def __init__(self, X: pd.DataFrame, y: pd.Series) -> None: self.X = X self.classes = np.sort(y.unique()) - self.y = np.searchsorted(self.classes, y.to_numpy()).astype(np.int32) + self.y = np.searchsorted(self.classes, y.to_numpy()).astype(np.int64) self.has_na: dict[str, bool] = dict() self.mask_na: dict[str, pd.Series] = dict() for column in self.X.columns: From 2ba86786572e68f567557f04b6010587a3b2f92d Mon Sep 17 00:00:00 2001 From: Mikhail Martin Date: Sun, 28 Sep 2025 18:57:59 +0300 Subject: [PATCH 3/3] mask.to_numpy() --- smarttree/_column_splitter.py | 22 +++++++++++++---- smarttree/_cy_column_splitter.pyi | 7 +++--- smarttree/_cy_column_splitter.pyx | 24 ++++++++----------- .../test__base_column_splitter.py | 4 ++++ 4 files changed, 35 insertions(+), 22 deletions(-) diff --git a/smarttree/_column_splitter.py b/smarttree/_column_splitter.py index 1a58ad1..01fb30a 100644 --- a/smarttree/_column_splitter.py +++ b/smarttree/_column_splitter.py @@ -61,6 +61,15 @@ def __init__( def split(self, *args, **kwargs) -> ColumnSplitResult: raise NotImplementedError + def pre_information_gain( + self, + parent_mask: pd.Series, + child_masks: list[pd.Series], + ) -> tuple[NDArray[np.bool_], list[NDArray[np.bool_]]]: + parent_mask_np = parent_mask.to_numpy() + child_masks_np = [child_mask.to_numpy() for child_mask in child_masks] + return parent_mask_np, child_masks_np + def foo( self, parent_mask: pd.Series, @@ -78,7 +87,8 @@ def foo( else: assert False else: - information_gain = self.information_gain(parent_mask, child_masks) + parent_mask_np, child_masks_np = self.pre_information_gain(parent_mask, child_masks) + information_gain = self.information_gain(parent_mask_np, child_masks_np) return information_gain, child_masks, -1 def include_all_split( @@ -93,7 +103,8 @@ def include_all_split( if child_masks[i].sum() < self.min_samples_leaf: return NO_INFORMATION_GAIN, [], -1 - information_gain = self.information_gain(parent_mask, child_masks, normalize=True) + parent_mask_np, child_masks_np = self.pre_information_gain(parent_mask, child_masks) + information_gain = self.information_gain(parent_mask_np, child_masks_np, normalize=True) return information_gain, child_masks, -1 @@ -119,7 +130,8 @@ def include_best_split( best_child_masks = [] best_child_na_index = -1 for child_na_index, child_masks in enumerate(candidates): - information_gain = self.information_gain(parent_mask, child_masks) + parent_mask_np, child_masks_np = self.pre_information_gain(parent_mask, child_masks) + information_gain = self.information_gain(parent_mask_np, child_masks_np) if best_information_gain < information_gain: best_information_gain = information_gain best_child_masks = child_masks @@ -129,8 +141,8 @@ def include_best_split( def information_gain( self, - parent_mask: pd.Series, - child_masks: list[pd.Series], + parent_mask: NDArray[np.bool_], + child_masks: list[NDArray[np.bool_]], normalize: bool = False, ) -> float: r""" diff --git a/smarttree/_cy_column_splitter.pyi b/smarttree/_cy_column_splitter.pyi index efa445f..e543199 100644 --- a/smarttree/_cy_column_splitter.pyi +++ b/smarttree/_cy_column_splitter.pyi @@ -1,4 +1,5 @@ -import pandas as pd +import numpy as np +from numpy.typing import NDArray from ._dataset import Dataset from ._types import Criterion @@ -10,8 +11,8 @@ class CyBaseColumnSplitter: def information_gain( self, - parent_mask: pd.Series, - child_masks: list[pd.Series], + parent_mask: NDArray[np.bool_], + child_masks: list[NDArray[np.bool_]], normalize: bool = False, ) -> float: r""" diff --git a/smarttree/_cy_column_splitter.pyx b/smarttree/_cy_column_splitter.pyx index 79e3dd0..71a883b 100644 --- a/smarttree/_cy_column_splitter.pyx +++ b/smarttree/_cy_column_splitter.pyx @@ -27,39 +27,35 @@ cdef class CyBaseColumnSplitter: def information_gain( self, - parent_mask: pd.Series, - child_masks: list[pd.Series], - normalize: bool = False, - ) -> float: + cnp.npy_bool[:] parent_mask, + list[cnp.npy_bool[:]] child_masks, + bint normalize, + ): - cdef cnp.npy_bool[:] parent_mask_arr, child_mask_arr 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 - parent_mask_arr = parent_mask.to_numpy() - child_mask_arrs = [child_mask.to_numpy() for child_mask in child_masks] - N = 0 N_parent = 0 for i in range(self.n_samples): N += 1 - if parent_mask_arr[i]: + if parent_mask[i]: N_parent += 1 - impurity_parent = self.criterion.impurity(parent_mask_arr) + impurity_parent = self.criterion.impurity(parent_mask) N_childs = 0 - n_childs = len(child_mask_arrs) + n_childs = len(child_masks) weighted_impurity_childs = 0.0 for j in range(n_childs): N_child_j = 0 - child_mask_arr = child_mask_arrs[j] + child_mask = child_masks[j] for i in range(self.n_samples): - if child_mask_arr[i]: + if child_mask[i]: N_child_j += 1 N_childs += N_child_j - impurity_child_i = self.criterion.impurity(child_mask_arr) + impurity_child_i = self.criterion.impurity(child_mask) weighted_impurity_childs += (N_child_j / N_parent) * impurity_child_i cdef double norm_coef diff --git a/tests/column_splitter/test__base_column_splitter.py b/tests/column_splitter/test__base_column_splitter.py index 0e4adf0..0a053b4 100644 --- a/tests/column_splitter/test__base_column_splitter.py +++ b/tests/column_splitter/test__base_column_splitter.py @@ -35,6 +35,10 @@ def test__information_gain(concrete_column_splitter, y): right_child_mask = ~left_child_mask child_masks = [left_child_mask, right_child_mask] + + parent_mask = parent_mask.to_numpy() + child_masks = [child_mask.to_numpy() for child_mask in child_masks] + inf_gain = concrete_column_splitter.information_gain(parent_mask, child_masks) assert inf_gain == 0.0016794443115909496