Skip to content

Commit de16d93

Browse files
authored
MAINT Refactors __reduce__ to parent splitter (scikit-learn#25199)
1 parent 8256a48 commit de16d93

2 files changed

Lines changed: 27 additions & 31 deletions

File tree

‎sklearn/tree/_splitter.pyx‎

Lines changed: 7 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -91,6 +91,13 @@ cdef class Splitter:
9191
def __setstate__(self, d):
9292
pass
9393

94+
def __reduce__(self):
95+
return (type(self), (self.criterion,
96+
self.max_features,
97+
self.min_samples_leaf,
98+
self.min_weight_leaf,
99+
self.random_state), self.__getstate__())
100+
94101
cdef int init(
95102
self,
96103
object X,
@@ -240,13 +247,6 @@ cdef class BaseDenseSplitter(Splitter):
240247

241248
cdef class BestSplitter(BaseDenseSplitter):
242249
"""Splitter for finding the best split."""
243-
def __reduce__(self):
244-
return (BestSplitter, (self.criterion,
245-
self.max_features,
246-
self.min_samples_leaf,
247-
self.min_weight_leaf,
248-
self.random_state), self.__getstate__())
249-
250250
cdef int node_split(self, double impurity, SplitRecord* split,
251251
SIZE_t* n_constant_features) nogil except -1:
252252
"""Find the best split on node samples[start:end]
@@ -553,13 +553,6 @@ cdef void heapsort(DTYPE_t* Xf, SIZE_t* samples, SIZE_t n) nogil:
553553

554554
cdef class RandomSplitter(BaseDenseSplitter):
555555
"""Splitter for finding the best random split."""
556-
def __reduce__(self):
557-
return (RandomSplitter, (self.criterion,
558-
self.max_features,
559-
self.min_samples_leaf,
560-
self.min_weight_leaf,
561-
self.random_state), self.__getstate__())
562-
563556
cdef int node_split(self, double impurity, SplitRecord* split,
564557
SIZE_t* n_constant_features) nogil except -1:
565558
"""Find the best random split on node samples[start:end]
@@ -1057,14 +1050,6 @@ cdef inline void sparse_swap(SIZE_t[::1] index_to_samples, SIZE_t[::1] samples,
10571050

10581051
cdef class BestSparseSplitter(BaseSparseSplitter):
10591052
"""Splitter for finding the best split, using the sparse data."""
1060-
1061-
def __reduce__(self):
1062-
return (BestSparseSplitter, (self.criterion,
1063-
self.max_features,
1064-
self.min_samples_leaf,
1065-
self.min_weight_leaf,
1066-
self.random_state), self.__getstate__())
1067-
10681053
cdef int node_split(self, double impurity, SplitRecord* split,
10691054
SIZE_t* n_constant_features) nogil except -1:
10701055
"""Find the best split on node samples[start:end], using sparse features
@@ -1283,14 +1268,6 @@ cdef class BestSparseSplitter(BaseSparseSplitter):
12831268

12841269
cdef class RandomSparseSplitter(BaseSparseSplitter):
12851270
"""Splitter for finding a random split, using the sparse data."""
1286-
1287-
def __reduce__(self):
1288-
return (RandomSparseSplitter, (self.criterion,
1289-
self.max_features,
1290-
self.min_samples_leaf,
1291-
self.min_weight_leaf,
1292-
self.random_state), self.__getstate__())
1293-
12941271
cdef int node_split(self, double impurity, SplitRecord* split,
12951272
SIZE_t* n_constant_features) nogil except -1:
12961273
"""Find a random split on node samples[start:end], using sparse features

‎sklearn/tree/tests/test_tree.py‎

Lines changed: 20 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
"""
44
import copy
55
import pickle
6-
from itertools import product
6+
from itertools import product, chain
77
import struct
88
import io
99
import copyreg
@@ -59,6 +59,7 @@
5959
from sklearn import datasets
6060

6161
from sklearn.utils import compute_sample_weight
62+
from sklearn.tree._classes import DENSE_SPLITTERS, SPARSE_SPLITTERS
6263

6364

6465
CLF_CRITERIONS = ("gini", "log_loss")
@@ -2371,3 +2372,21 @@ def test_max_features_auto_deprecated():
23712372
)
23722373
with pytest.warns(FutureWarning, match=msg):
23732374
tree.fit(X, y)
2375+
2376+
2377+
@pytest.mark.parametrize(
2378+
"Splitter", chain(DENSE_SPLITTERS.values(), SPARSE_SPLITTERS.values())
2379+
)
2380+
def test_splitter_serializable(Splitter):
2381+
"""Check that splitters are serializable."""
2382+
rng = np.random.RandomState(42)
2383+
max_features = 10
2384+
n_outputs, n_classes = 2, np.array([3, 2], dtype=np.intp)
2385+
2386+
criterion = CRITERIA_CLF["gini"](n_outputs, n_classes)
2387+
splitter = Splitter(criterion, max_features, 5, 0.5, rng)
2388+
splitter_serialize = pickle.dumps(splitter)
2389+
2390+
splitter_back = pickle.loads(splitter_serialize)
2391+
assert splitter_back.max_features == max_features
2392+
assert isinstance(splitter_back, Splitter)

0 commit comments

Comments
 (0)