From 1b9d231a47f5fe8d743151249dcd609ee1ecb814 Mon Sep 17 00:00:00 2001 From: Ron Wettenstein Date: Sun, 5 Apr 2026 11:19:42 +0300 Subject: [PATCH 01/11] Added the tree_accurate method --- sklearn/inspection/_partial_dependence.py | 106 +++++- .../tests/test_partial_dependence.py | 194 ++++++++++- sklearn/tree/_tree.pyx | 307 ++++++++++++++++++ 3 files changed, 603 insertions(+), 4 deletions(-) diff --git a/sklearn/inspection/_partial_dependence.py b/sklearn/inspection/_partial_dependence.py index 98786550425bd..95c4ea11af13f 100644 --- a/sklearn/inspection/_partial_dependence.py +++ b/sklearn/inspection/_partial_dependence.py @@ -10,13 +10,13 @@ from scipy.stats.mstats import mquantiles from sklearn.base import is_classifier, is_regressor -from sklearn.ensemble import RandomForestRegressor +from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor from sklearn.ensemble._gb import BaseGradientBoosting from sklearn.ensemble._hist_gradient_boosting.gradient_boosting import ( BaseHistGradientBoosting, ) from sklearn.inspection._pd_utils import _check_feature_names, _get_feature_index -from sklearn.tree import DecisionTreeRegressor +from sklearn.tree import DecisionTreeClassifier, DecisionTreeRegressor from sklearn.utils import Bunch, _safe_indexing, check_array from sklearn.utils._indexing import ( _determine_key_type, @@ -218,6 +218,68 @@ def _partial_dependence_recursion(est, grid, features): return averaged_predictions +def _partial_dependence_tree_accurate(est, X, all_features, all_grids): + """Calculate partial dependence using the tree_accurate method. + + Performs a single O(m * D²) background pass over ``X`` per tree to collect + per-leaf statistics, then combines them with a vectorised pass over each grid. + Only ``kind='average'`` is supported. + Returns one centred averaged-prediction array per feature, + shifted by the mean prediction so results are on the same scale as 'brute'. + + Parameters + ---------- + est : DecisionTreeRegressor or RandomForestRegressor + A fitted tree-based regressor (single- or multi-output). + X : ndarray of shape (n_samples, n_features), dtype=np.float32 + Background dataset used for marginalisation. + all_features : list of int + Global column indices of the target features, one per requested feature. + all_grids : list of ndarray of shape (n_grid_j,) + 1-D grid of values for each target feature. + + Returns + ------- + averaged_predictions : list of ndarray of shape (n_outputs, n_grid_j) + One array per feature in ``all_features``. + """ + m = X.shape[0] + + if is_classifier(est): + n_classes = est.n_classes_ + n_effective_outputs = 1 if n_classes == 2 else n_classes + else: + n_effective_outputs = est.n_outputs_ + + grid_sizes = [len(g) for g in all_grids] + n_grid_max = max(grid_sizes) + n_required = len(all_features) + + # Build padded grid: shape (n_grid_max, n_required). + grid_2d = np.zeros((n_grid_max, n_required), dtype=np.float32) + for i, grid_1d in enumerate(all_grids): + grid_2d[: grid_sizes[i], i] = grid_1d + + required_features_arr = np.array(all_features, dtype=np.intp) + out = np.zeros((n_effective_outputs, n_required, n_grid_max), dtype=np.float64) + + if isinstance(est, (DecisionTreeRegressor, DecisionTreeClassifier)): + est.tree_.compute_partial_dependence_tree_accurate( + X, grid_2d, required_features_arr, out + ) + out /= m + elif isinstance(est, (RandomForestRegressor, RandomForestClassifier)): + n_trees = len(est.estimators_) + for tree_est in est.estimators_: + tree_est.tree_.compute_partial_dependence_tree_accurate( + X, grid_2d, required_features_arr, out + ) + out /= m * n_trees + + # Extract per-feature results, dropping any padded grid entries. + return [out[:, i, : grid_sizes[i]] for i in range(n_required)] + + def _partial_dependence_brute( est, grid, features, X, response_method, sample_weight=None ): @@ -361,7 +423,7 @@ def _partial_dependence_brute( "response_method": [StrOptions({"auto", "predict_proba", "decision_function"})], "percentiles": [tuple], "grid_resolution": [Interval(Integral, 1, None, closed="left")], - "method": [StrOptions({"auto", "recursion", "brute"})], + "method": [StrOptions({"auto", "recursion", "brute", "tree_accurate"})], "kind": [StrOptions({"average", "individual", "both"})], "custom_values": [dict, None], }, @@ -601,6 +663,8 @@ def partial_dependence( raise ValueError( "The 'recursion' method only applies when 'kind' is set to 'average'" ) + if method == "tree_accurate": + raise ValueError("The 'tree_accurate' method only supports kind='average'.") method = "brute" if method == "recursion" and sample_weight is not None: @@ -608,6 +672,36 @@ def partial_dependence( "The 'recursion' method can only be applied when sample_weight is None." ) + if method == "tree_accurate": + if sample_weight is not None: + raise ValueError( + "The 'tree_accurate' method can only be applied when " + "sample_weight is None." + ) + if not isinstance( + estimator, + ( + DecisionTreeRegressor, + RandomForestRegressor, + DecisionTreeClassifier, + RandomForestClassifier, + ), + ): + raise ValueError( + "The 'tree_accurate' method only supports DecisionTreeRegressor, " + "RandomForestRegressor, DecisionTreeClassifier, and " + "RandomForestClassifier. Use method='brute' for other estimators." + ) + if is_classifier(estimator): + if response_method == "auto": + response_method = "decision_function" + if response_method != "decision_function": + raise ValueError( + "With the 'tree_accurate' method, the response_method for " + "classifiers must be 'decision_function'. " + "Got {}.".format(response_method) + ) + if method == "auto": if sample_weight is not None: method = "brute" @@ -750,6 +844,12 @@ def partial_dependence( predictions = predictions.reshape( -1, X.shape[0], *[val.shape[0] for val in values] ) + elif method == "tree_accurate": + X_bg = np.asarray(X, dtype=np.float32, order="C") + per_feature_preds = _partial_dependence_tree_accurate( + estimator, X_bg, list(features_indices), values + ) + averaged_predictions = np.concatenate(per_feature_preds, axis=1) else: averaged_predictions = _partial_dependence_recursion( estimator, grid, features_indices diff --git a/sklearn/inspection/tests/test_partial_dependence.py b/sklearn/inspection/tests/test_partial_dependence.py index 8de472c6e7114..57c4dfdb1f482 100644 --- a/sklearn/inspection/tests/test_partial_dependence.py +++ b/sklearn/inspection/tests/test_partial_dependence.py @@ -38,7 +38,7 @@ StandardScaler, scale, ) -from sklearn.tree import DecisionTreeRegressor +from sklearn.tree import DecisionTreeClassifier, DecisionTreeRegressor from sklearn.tree.tests.test_tree import assert_is_subtree from sklearn.utils._testing import assert_allclose, assert_array_equal from sklearn.utils.fixes import _IS_32BIT @@ -1253,3 +1253,195 @@ def test_partial_dependence_empty_categorical_features(): partial_dependence( estimator=clf, X=iris.data, features=[0], categorical_features=[] ) + + +# ============================================================================= +# tree_accurate method tests +# ============================================================================= + + +@pytest.mark.parametrize("seed", range(3)) +@pytest.mark.parametrize( + "Estimator", + [ + DecisionTreeRegressor, + RandomForestRegressor, + ], +) +def test_tree_accurate_matches_brute(Estimator, seed): + """tree_accurate must give the same averaged predictions as brute.""" + rng = np.random.RandomState(seed) + n_samples, n_features = 200, 5 + X = rng.randn(n_samples, n_features).astype(np.float64) + y = rng.randn(n_samples) + + kwargs = dict(max_depth=4, random_state=seed) + if Estimator is RandomForestRegressor: + kwargs["n_estimators"] = 3 + kwargs["max_features"] = "sqrt" + est = Estimator(**kwargs).fit(X, y) + + for feature in range(n_features): + pdp_brute = partial_dependence( + est, X, features=[feature], method="brute", grid_resolution=20 + ) + pdp_fast = partial_dependence( + est, X, features=[feature], method="tree_accurate", grid_resolution=20 + ) + np.testing.assert_allclose( + pdp_fast["average"], + pdp_brute["average"], + rtol=1e-4, + atol=1e-6, + err_msg=f"Mismatch on feature {feature} with seed {seed}", + ) + + +@pytest.mark.parametrize( + "Estimator", + [DecisionTreeRegressor, RandomForestRegressor], +) +def test_tree_accurate_output_shape(Estimator): + """Output shape should be (1, grid_resolution) for each feature.""" + rng = np.random.RandomState(0) + X = rng.randn(100, 4).astype(np.float64) + y = rng.randn(100) + + kwargs = dict(max_depth=3, random_state=0) + if Estimator is RandomForestRegressor: + kwargs["n_estimators"] = 2 + est = Estimator(**kwargs).fit(X, y) + + grid_resolution = 15 + for feature in range(4): + result = partial_dependence( + est, + X, + features=[feature], + method="tree_accurate", + grid_resolution=grid_resolution, + ) + assert result["average"].shape == (1, grid_resolution), ( + f"Expected (1, {grid_resolution}), got {result['average'].shape}" + ) + assert result["grid_values"][0].shape == (grid_resolution,) + + +def test_tree_accurate_repeated_feature_in_path(): + """tree_accurate must give correct values when a feature splits multiple + times on the same root-to-leaf path.""" + # Build a tree that is guaranteed to reuse feature 0 at multiple depths + # by restricting max_features so it can only choose feature 0. + rng = np.random.RandomState(42) + n_samples = 500 + # Single informative feature so the tree will reuse it + X = rng.randn(n_samples, 3).astype(np.float64) + y = X[:, 0] ** 2 # target depends on feature 0 only + + est = DecisionTreeRegressor(max_depth=6, random_state=0).fit(X, y) + + pdp_brute = partial_dependence( + est, X, features=[0], method="brute", grid_resolution=30 + ) + pdp_fast = partial_dependence( + est, X, features=[0], method="tree_accurate", grid_resolution=30 + ) + np.testing.assert_allclose( + pdp_fast["average"], + pdp_brute["average"], + rtol=1e-4, + atol=1e-6, + ) + + +def test_tree_accurate_kind_not_average_raises(): + """kind != 'average' must raise ValueError.""" + rng = np.random.RandomState(0) + X = rng.randn(50, 3).astype(np.float64) + y = rng.randn(50) + est = DecisionTreeRegressor(max_depth=3, random_state=0).fit(X, y) + + for kind in ("individual", "both"): + with pytest.raises(ValueError, match="'tree_accurate' method only supports"): + partial_dependence(est, X, features=[0], method="tree_accurate", kind=kind) + + +def test_tree_accurate_unsupported_estimator_raises(): + """Non-tree estimators must raise ValueError.""" + X, y = make_regression(n_samples=50, random_state=0) + est = LinearRegression().fit(X, y) + + with pytest.raises(ValueError, match="'tree_accurate' method only supports"): + partial_dependence(est, X, features=[0], method="tree_accurate") + + +def test_tree_accurate_sample_weight_raises(): + """sample_weight != None must raise ValueError.""" + rng = np.random.RandomState(0) + X = rng.randn(50, 3).astype(np.float64) + y = rng.randn(50) + est = DecisionTreeRegressor(max_depth=3, random_state=0).fit(X, y) + + with pytest.raises(ValueError, match="'tree_accurate' method can only be applied"): + partial_dependence( + est, X, features=[0], method="tree_accurate", sample_weight=np.ones(50) + ) + + +def test_tree_accurate_binary_classifier_matches_brute(): + """Binary DecisionTreeClassifier: tree_accurate PDP must match brute.""" + X, y = make_classification(n_samples=100, n_features=4, random_state=0) + est = DecisionTreeClassifier(max_depth=3, random_state=0).fit(X, y) + + result_ta = partial_dependence( + est, X, features=[0], method="tree_accurate", grid_resolution=10 + ) + result_br = partial_dependence( + est, X, features=[0], method="brute", grid_resolution=10 + ) + np.testing.assert_allclose(result_ta["average"], result_br["average"], atol=1e-6) + + +def test_tree_accurate_multiclass_classifier_matches_brute(): + """Multiclass DecisionTreeClassifier: tree_accurate PDP must match brute.""" + X, y = make_classification( + n_samples=100, + n_features=4, + n_classes=3, + n_informative=3, + n_redundant=0, + random_state=0, + ) + est = DecisionTreeClassifier(max_depth=3, random_state=0).fit(X, y) + + result_ta = partial_dependence( + est, X, features=[0], method="tree_accurate", grid_resolution=10 + ) + result_br = partial_dependence( + est, X, features=[0], method="brute", grid_resolution=10 + ) + np.testing.assert_allclose(result_ta["average"], result_br["average"], atol=1e-6) + + +def test_tree_accurate_multi_output_matches_brute(): + """Multi-output DecisionTreeRegressor: tree_accurate must match brute per output.""" + rng = np.random.RandomState(42) + X = rng.randn(120, 4).astype(np.float64) + # Two independent regression targets. + Y = np.column_stack([X[:, 0] + X[:, 1], X[:, 2] - X[:, 3]]) + + est = DecisionTreeRegressor(max_depth=4, random_state=0).fit(X, Y) + assert est.n_outputs_ == 2 + + result_accurate = partial_dependence( + est, X, features=[0], method="tree_accurate", grid_resolution=10 + ) + result_brute = partial_dependence( + est, X, features=[0], method="brute", grid_resolution=10 + ) + + # Both methods should return shape (n_outputs, n_grid) = (2, 10). + assert result_accurate["average"].shape == (2, 10) + np.testing.assert_allclose( + result_accurate["average"], result_brute["average"], atol=1e-6 + ) diff --git a/sklearn/tree/_tree.pyx b/sklearn/tree/_tree.pyx index d587400bba80a..0f4d9b2f92d3b 100644 --- a/sklearn/tree/_tree.pyx +++ b/sklearn/tree/_tree.pyx @@ -1532,6 +1532,313 @@ cdef class Tree: raise ValueError("Total weight should be 1.0 but was %.9f" % total_weight) + def compute_partial_dependence_tree_accurate( + self, + float32_t[:, ::1] X_bg, + float32_t[:, ::1] grid, + const intp_t[::1] required_features, + float64_t[:, :, ::1] out, + ): + """Partial dependence via background-data tree traversal (tree_accurate method). + + Uses a single DFS per background sample: for each sample the natural + prediction path is followed, and at every split on a required feature a shadow + path is spawned that diverges to the other branch. Shadows may re-diverge + on the same feature (handling repeated splits), so the per-sample work is + O(D + total shadow nodes). This is O(D^2) when each required feature + appears at most once per path, and degrades to O(L) when the tree splits + exclusively on a single required feature (L = number of leaves). + Per-leaf statistics (``reached`` and ``diverged_once``) are accumulated using O(1) + lookups via a depth-indexed path-state array maintained with LEAVE markers. + + Memory overhead is O(N * D) where N = node_count and D = max_depth. + + Accumulates raw sums into ``out``; the caller is responsible for dividing + by the number of background samples (and, for forests, by n_estimators). + + Parameters + ---------- + X_bg : float32 C-contiguous array of shape (n_background, n_features) + Background dataset used for marginalisation. + grid : float32 C-contiguous array of shape (n_grid, n_required_features) + Consumer grid; column j contains values for required_features[j]. + required_features : intp array of shape (n_required_features,) + Global feature indices of the features whose PDP is being computed. + out : float64 array of shape (n_outputs, n_required_features, n_grid) + Output array; sums are accumulated in-place. + """ + cdef: + intp_t n_required = required_features.shape[0] + intp_t n_grid = grid.shape[0] + intp_t n_background = X_bg.shape[0] + intp_t _TREE_LEAF_SENTINEL = TREE_LEAF + intp_t _LEAVE_MARKER = TREE_UNDEFINED + + if n_background == 0 or n_grid == 0 or self.node_count == 0: + return + + # ------------------------------------------------------------------ # + # feature_to_req_pos: global feature index -> position in # + # required_features (-1 if not a required feature) # + # ------------------------------------------------------------------ # + cdef intp_t[:] feature_to_req_pos = np.full(self.n_features, -1, dtype=np.intp) + cdef intp_t req_idx + for req_idx in range(n_required): + feature_to_req_pos[required_features[req_idx]] = req_idx + + # ----------------------------------------------------------------------- # + # Phase 0: single DFS over the tree structure, O(node_count). It builds: # + # - node_depth_arr: The depth of each node (needed by Phase 1) # + # - leaf_conditions: With the details about split conditions along its # + # root-to-leaf path (needed by Phase 2) # + # No large per-node arrays (no leaf_uniq, no leaf_n_uniq). # + # ----------------------------------------------------------------------- # + + children_left_arr = np.asarray(self.children_left) + children_right_arr = np.asarray(self.children_right) + features_arr = np.asarray(self.feature) + thresholds_arr = np.asarray(self.threshold) + + # Classifier detection: regressors have n_classes[0] == 1. + cdef intp_t _n_clf_classes = self.max_n_classes + cdef bint _is_clf = _n_clf_classes > 1 + + cdef intp_t[:] node_depth_arr = np.empty(self.node_count, dtype=np.intp) + + # leaf_conditions[node_id] = (path_conditions, path_req_slots, leaf_values) + # path_conditions : {req_pos: [(split_threshold, goes_left), ...]} + # path_req_slots : ordered list of (d, req_pos) for each required feature + # first seen on this path. d is the feature's position + # among ALL unique features on the path, matching Phase 1's + # feature_to_d assignment. + leaf_conditions = {} + + # Python DFS — stack entries: + # (node_id, depth, path_conditions, path_req_slots, seen_features, n_seen) + # seen_features : set of global feature indices already seen on this path + # n_seen : len(seen_features), tracked explicitly to assign d + dfs_stack = [(0, 0, {}, [], set(), 0)] + while dfs_stack: + node_id, depth, path_conditions, path_req_slots, seen_features, n_seen = dfs_stack.pop() + node_depth_arr[node_id] = depth + + if children_left_arr[node_id] == TREE_LEAF: + if _is_clf: + # Normalise raw class counts to probabilities, matching + # predict_proba: divide each class count by the total count. + # Sum the counts from the buffer directly (same as predict_proba). + total = sum( + self.value[node_id * self.value_stride + c] + for c in range(_n_clf_classes) + ) + if _n_clf_classes == 2: + # Binary: keep only positive-class probability. + leaf_values = np.array( + [self.value[node_id * self.value_stride + 1] / total], + dtype=np.float64, + ) + else: + leaf_values = np.array( + [self.value[node_id * self.value_stride + c] / total + for c in range(_n_clf_classes)], + dtype=np.float64, + ) + else: + leaf_values = np.array( + [self.value[node_id * self.value_stride + k] + for k in range(self.n_outputs)], + dtype=np.float64, + ) + leaf_conditions[node_id] = (path_conditions, path_req_slots, leaf_values) + else: + split_feature = int(features_arr[node_id]) + split_threshold = float(thresholds_arr[node_id]) + left_child_idx = int(children_left_arr[node_id]) + right_child_idx = int(children_right_arr[node_id]) + req_pos = int(feature_to_req_pos[split_feature]) + + # Compute d (position among ALL unique path features) for this feature. + # This mirrors Phase 1's uniform path-state update for all features. + is_new = split_feature not in seen_features + d = n_seen # used only when is_new + child_seen = seen_features | {split_feature} if is_new else seen_features + child_n_seen = n_seen + (1 if is_new else 0) + + for child_idx, goes_left in ((left_child_idx, True), (right_child_idx, False)): + child_conditions = {rp: list(conds) for rp, conds in path_conditions.items()} + child_req_slots = list(path_req_slots) + if req_pos >= 0: + if req_pos not in child_conditions: + # First occurrence of this required feature: record its d + # (position among ALL unique path features). + child_req_slots.append((d, req_pos)) + child_conditions[req_pos] = [] + child_conditions[req_pos].append((split_threshold, goes_left)) + dfs_stack.append((child_idx, depth + 1, child_conditions, child_req_slots, + child_seen, child_n_seen)) + + # ------------------------------------------------------------------ # + # Phase 1: background pass (Cython hot loop) # + # # + # Per-leaf accumulators (node_count sized, but only leaves used): # + # reached[leaf] — samples that naturally landed # + # here (main path) # + # diverged_once[leaf, d] — samples whose shadow (that # + # diverged on the required feature # + # at path position d) landed here # + # # + # Path state (O(D)) — tracks ALL features, not only required features: # + # n_unique_features_on_path — #unique features (all) # + # seen on the active path # + # feature_to_d[feature_idx] — d of feature in path # + # order (-1 if not yet # + # seen) # + # feature_first_added_at_depth[depth] — global feature index # + # first added at this depth # + # (-1 if none or already # + # on path) # + # # + # Stack encoding: # + # stack_node[stack_size] >= 0 → ENTER that node # + # stack_node[stack_size] == _LEAVE_MARKER # + # → LEAVE: undo path state at # + # depth stack_diverged_on_target[sz] # + # ------------------------------------------------------------------ # + + cdef intp_t[:] reached = np.zeros(self.node_count, dtype=np.intp) + cdef intp_t[:, :] diverged_once = np.zeros( + (self.node_count, max(self.max_depth, 1)), dtype=np.intp + ) + + # Path state — initialised once, self-cleaning via LEAVE markers + cdef intp_t n_unique_features_on_path = 0 + cdef intp_t[:] feature_to_d = np.full(self.n_features, -1, dtype=np.intp) + cdef intp_t[:] feature_first_added_at_depth = np.full(self.max_depth + 1, -1, dtype=np.intp) + + # Stack: LEAVE entries double the worst-case depth; +4 for safety + cdef intp_t stack_capacity = 2 * (n_required + 2) * (self.max_depth + 2) + 4 + cdef intp_t[:] stack_node = np.empty(stack_capacity, dtype=np.intp) + # When node_idx == _LEAVE_MARKER include the depth, + # when node_idx is of an inner node include the feature we diverged on. + cdef intp_t[:] stack_diverged_on_target = np.empty(stack_capacity, dtype=np.intp) + + cdef: + intp_t stack_size, sample_idx, node_idx, diverged_on_target + intp_t current_feature, current_depth, newly_added_feature + intp_t proceed_child, diverge_child + Node* node + bint sample_goes_left + + for sample_idx in range(n_background): + stack_size = 1 + stack_node[0] = 0 + stack_diverged_on_target[0] = -1 + + while stack_size > 0: + stack_size -= 1 + node_idx = stack_node[stack_size] + + if node_idx == _LEAVE_MARKER: + # ---- LEAVE marker: undo path state for this depth ---- # + current_depth = stack_diverged_on_target[stack_size] + newly_added_feature = feature_first_added_at_depth[current_depth] + if newly_added_feature >= 0: + feature_to_d[newly_added_feature] = -1 + n_unique_features_on_path -= 1 + feature_first_added_at_depth[current_depth] = -1 + continue + + # ---- ENTER node ---- # + diverged_on_target = stack_diverged_on_target[stack_size] + node = &self.nodes[node_idx] + + if node.left_child == _TREE_LEAF_SENTINEL: + # Leaf: O(1) update via feature_to_d + if diverged_on_target < 0: + reached[node_idx] += 1 + else: + diverged_once[node_idx, feature_to_d[diverged_on_target]] += 1 + else: + current_feature = node.feature + current_depth = node_depth_arr[node_idx] + sample_goes_left = X_bg[sample_idx, current_feature] <= node.threshold + + proceed_child = node.left_child if sample_goes_left else node.right_child + diverge_child = node.right_child if sample_goes_left else node.left_child + + stack_node[stack_size] = _LEAVE_MARKER + stack_diverged_on_target[stack_size] = current_depth + stack_size += 1 + + if feature_to_d[current_feature] < 0: + feature_to_d[current_feature] = n_unique_features_on_path + n_unique_features_on_path += 1 + feature_first_added_at_depth[current_depth] = current_feature + + # Push natural direction (main path or shadow continues) + stack_node[stack_size] = proceed_child + stack_diverged_on_target[stack_size] = diverged_on_target + stack_size += 1 + + # Spawn a shadow only at target-feature splits. + if feature_to_req_pos[current_feature] >= 0 and (diverged_on_target < 0 or diverged_on_target == current_feature): + stack_node[stack_size] = diverge_child + stack_diverged_on_target[stack_size] = current_feature + stack_size += 1 + + # ------------------------------------------------------------------ # + # Phase 2: consumer patterns + accumulation (numpy, O(L * n_grid)) # + # ------------------------------------------------------------------ # + + reached_arr = np.asarray(reached) + diverged_once_arr = np.asarray(diverged_once) + grid_arr = np.asarray(grid) + out_arr = np.asarray(out) + + # avg_pred_sum accumulates leaf_values * reached[leaf] over all leaves. + # Adding it to out_arr (then dividing by m in the caller) gives the absolute mean prediction + avg_pred_sum = np.zeros(out_arr.shape[0], dtype=np.float64) + + for leaf_node_idx, (path_conditions, path_req_slots, leaf_values) in leaf_conditions.items(): + avg_pred_sum += leaf_values * float(reached_arr[leaf_node_idx]) + + n_unique_req = len(path_req_slots) + if n_unique_req == 0: + # No required features on this path; only contributes to avg_pred. + continue + + grid_satisfies_conditions = np.ones((n_grid, n_unique_req), dtype=np.bool_) + for i, (d, req_pos) in enumerate(path_req_slots): + for split_threshold, goes_left in path_conditions.get(req_pos, []): + if goes_left: + grid_satisfies_conditions[:, i] &= grid_arr[:, req_pos] <= split_threshold + else: + grid_satisfies_conditions[:, i] &= grid_arr[:, req_pos] > split_threshold + + # d is the position of the feature among ALL unique path features + req_ds = [d for d, _ in path_req_slots] + # req_pos is the column index in the grid / out arrays + req_positions = [rp for _, rp in path_req_slots] + + diverged_once_for_leaf = diverged_once_arr[leaf_node_idx, req_ds].astype(np.float64) + multipliers = np.where( + grid_satisfies_conditions, # (n_grid, n_unique_req) + diverged_once_for_leaf, # broadcast over n_grid + -float(reached_arr[leaf_node_idx]), + ) + + # Accumulate into out_arr[n_outputs, n_required_features, n_grid]: + # leaf_values[:, None, None] : (n_outputs, 1, 1) + # multipliers.T[None, :, :] : (1, n_unique_req, n_grid) + # product : (n_outputs, n_unique_req, n_grid) + out_arr[:, req_positions, :] += ( + leaf_values[:, None, None] * multipliers.T[None, :, :] + ) + + # Fold avg_pred into out so the caller just divides by m (no external + # predict call needed). Broadcasts over all (n_required, n_grid) entries. + out_arr += avg_pred_sum[:, np.newaxis, np.newaxis] + def _check_n_classes(n_classes, expected_dtype): if n_classes.ndim != 1: From 9f8ed48564dc068bbe095e576df7b4e516c1582d Mon Sep 17 00:00:00 2001 From: Ron Wettenstein Date: Sun, 5 Apr 2026 18:13:13 +0300 Subject: [PATCH 02/11] added docs --- sklearn/inspection/_partial_dependence.py | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/sklearn/inspection/_partial_dependence.py b/sklearn/inspection/_partial_dependence.py index 95c4ea11af13f..8bdd592b1320e 100644 --- a/sklearn/inspection/_partial_dependence.py +++ b/sklearn/inspection/_partial_dependence.py @@ -549,7 +549,7 @@ def partial_dependence( .. versionadded:: 1.7 - method : {'auto', 'recursion', 'brute'}, default='auto' + method : {'auto', 'recursion', 'brute', 'tree_accurate'}, default='auto' The method used to calculate the averaged predictions: - `'recursion'` is only supported for some tree-based estimators @@ -572,6 +572,16 @@ def partial_dependence( - `'brute'` is supported for any estimator, but is more computationally intensive. + - `'tree_accurate'` is supported for + :class:`~sklearn.tree.DecisionTreeRegressor`, + :class:`~sklearn.ensemble.RandomForestRegressor`, + :class:`~sklearn.tree.DecisionTreeClassifier`, and + :class:`~sklearn.ensemble.RandomForestClassifier` + when `kind='average'` and `response_method='decision_function'`. + Joint PDP (tuples in `features`) is not supported. + This method is equivalent to the `'brute'` method, + but significantly faster for tree-based estimators. + - `'auto'`: the `'recursion'` is used for estimators that support it, and `'brute'` is used otherwise. If `sample_weight` is not `None`, then `'brute'` is used regardless of the estimator. @@ -673,6 +683,11 @@ def partial_dependence( ) if method == "tree_accurate": + if any(isinstance(f, (list, tuple)) for f in features): + raise ValueError( + "The 'tree_accurate' method does not support joint PDP " + "(tuples in features). Use method='brute' instead." + ) if sample_weight is not None: raise ValueError( "The 'tree_accurate' method can only be applied when " From cb14578bb7c36b6a0cab24892b63461dc5d04ab7 Mon Sep 17 00:00:00 2001 From: Ron Wettenstein Date: Fri, 12 Jun 2026 16:10:52 +0300 Subject: [PATCH 03/11] Use the fact we are only interested in 1 feature at a time --- sklearn/inspection/_partial_dependence.py | 60 +++--- sklearn/tree/_tree.pyx | 231 +++++++--------------- 2 files changed, 90 insertions(+), 201 deletions(-) diff --git a/sklearn/inspection/_partial_dependence.py b/sklearn/inspection/_partial_dependence.py index 8bdd592b1320e..6993f7686c8bb 100644 --- a/sklearn/inspection/_partial_dependence.py +++ b/sklearn/inspection/_partial_dependence.py @@ -218,30 +218,28 @@ def _partial_dependence_recursion(est, grid, features): return averaged_predictions -def _partial_dependence_tree_accurate(est, X, all_features, all_grids): - """Calculate partial dependence using the tree_accurate method. +def _partial_dependence_tree_accurate(est, X, feature, grid): + """Calculate partial dependence for a single feature using tree_accurate. Performs a single O(m * D²) background pass over ``X`` per tree to collect - per-leaf statistics, then combines them with a vectorised pass over each grid. + per-leaf statistics, then combines them with a vectorised pass over the grid. Only ``kind='average'`` is supported. - Returns one centred averaged-prediction array per feature, - shifted by the mean prediction so results are on the same scale as 'brute'. Parameters ---------- - est : DecisionTreeRegressor or RandomForestRegressor - A fitted tree-based regressor (single- or multi-output). + est : DecisionTreeRegressor, RandomForestRegressor, DecisionTreeClassifier, + or RandomForestClassifier + A fitted tree-based estimator. X : ndarray of shape (n_samples, n_features), dtype=np.float32 Background dataset used for marginalisation. - all_features : list of int - Global column indices of the target features, one per requested feature. - all_grids : list of ndarray of shape (n_grid_j,) - 1-D grid of values for each target feature. + feature : int + Global column index of the target feature. + grid : ndarray of shape (n_grid,), dtype=np.float32 + Grid of values for the target feature. Returns ------- - averaged_predictions : list of ndarray of shape (n_outputs, n_grid_j) - One array per feature in ``all_features``. + averaged_predictions : ndarray of shape (n_outputs, n_grid) """ m = X.shape[0] @@ -251,33 +249,20 @@ def _partial_dependence_tree_accurate(est, X, all_features, all_grids): else: n_effective_outputs = est.n_outputs_ - grid_sizes = [len(g) for g in all_grids] - n_grid_max = max(grid_sizes) - n_required = len(all_features) - - # Build padded grid: shape (n_grid_max, n_required). - grid_2d = np.zeros((n_grid_max, n_required), dtype=np.float32) - for i, grid_1d in enumerate(all_grids): - grid_2d[: grid_sizes[i], i] = grid_1d - - required_features_arr = np.array(all_features, dtype=np.intp) - out = np.zeros((n_effective_outputs, n_required, n_grid_max), dtype=np.float64) + out = np.zeros((n_effective_outputs, len(grid)), dtype=np.float64) if isinstance(est, (DecisionTreeRegressor, DecisionTreeClassifier)): - est.tree_.compute_partial_dependence_tree_accurate( - X, grid_2d, required_features_arr, out - ) + est.tree_.compute_partial_dependence_tree_accurate(X, grid, feature, out) out /= m elif isinstance(est, (RandomForestRegressor, RandomForestClassifier)): n_trees = len(est.estimators_) for tree_est in est.estimators_: tree_est.tree_.compute_partial_dependence_tree_accurate( - X, grid_2d, required_features_arr, out + X, grid, feature, out ) out /= m * n_trees - # Extract per-feature results, dropping any padded grid entries. - return [out[:, i, : grid_sizes[i]] for i in range(n_required)] + return out def _partial_dependence_brute( @@ -594,10 +579,10 @@ def partial_dependence( samples in the dataset or one value per sample or both. See Returns below. - Note that the fast `method='recursion'` option is only available for - `kind='average'` and `sample_weights=None`. Computing individual - dependencies and doing weighted averages requires using the slower - `method='brute'`. + Note that the fast `method='recursion'` and `method='tree_accurate'` + options are only available for `kind='average'`. + Computing individual dependencies and doing weighted averages requires + using the slower `method='brute'`. .. versionadded:: 0.24 @@ -861,10 +846,11 @@ def partial_dependence( ) elif method == "tree_accurate": X_bg = np.asarray(X, dtype=np.float32, order="C") - per_feature_preds = _partial_dependence_tree_accurate( - estimator, X_bg, list(features_indices), values + feature = int(features_indices[0]) + grid_1d = np.asarray(values[0], dtype=np.float32) + averaged_predictions = _partial_dependence_tree_accurate( + estimator, X_bg, feature, grid_1d ) - averaged_predictions = np.concatenate(per_feature_preds, axis=1) else: averaged_predictions = _partial_dependence_recursion( estimator, grid, features_indices diff --git a/sklearn/tree/_tree.pyx b/sklearn/tree/_tree.pyx index 0f4d9b2f92d3b..9e2c1a3320b46 100644 --- a/sklearn/tree/_tree.pyx +++ b/sklearn/tree/_tree.pyx @@ -1535,23 +1535,21 @@ cdef class Tree: def compute_partial_dependence_tree_accurate( self, float32_t[:, ::1] X_bg, - float32_t[:, ::1] grid, - const intp_t[::1] required_features, - float64_t[:, :, ::1] out, + float32_t[::1] grid, + intp_t required_feature, + float64_t[:, ::1] out, ): """Partial dependence via background-data tree traversal (tree_accurate method). Uses a single DFS per background sample: for each sample the natural - prediction path is followed, and at every split on a required feature a shadow - path is spawned that diverges to the other branch. Shadows may re-diverge - on the same feature (handling repeated splits), so the per-sample work is - O(D + total shadow nodes). This is O(D^2) when each required feature - appears at most once per path, and degrades to O(L) when the tree splits - exclusively on a single required feature (L = number of leaves). - Per-leaf statistics (``reached`` and ``diverged_once``) are accumulated using O(1) - lookups via a depth-indexed path-state array maintained with LEAVE markers. + prediction path is followed, and at every split on the required feature a + shadow path is spawned that diverges to the other branch. Shadows may + re-diverge on the same feature (handling repeated splits), so the + per-sample work is O(D + total shadow nodes). + Per-leaf statistics (``reached`` and ``diverged_on_required``) are + accumulated with O(1) per-leaf updates. - Memory overhead is O(N * D) where N = node_count and D = max_depth. + Memory overhead is O(N) where N = node_count. Accumulates raw sums into ``out``; the caller is responsible for dividing by the number of background samples (and, for forests, by n_estimators). @@ -1560,38 +1558,25 @@ cdef class Tree: ---------- X_bg : float32 C-contiguous array of shape (n_background, n_features) Background dataset used for marginalisation. - grid : float32 C-contiguous array of shape (n_grid, n_required_features) - Consumer grid; column j contains values for required_features[j]. - required_features : intp array of shape (n_required_features,) - Global feature indices of the features whose PDP is being computed. - out : float64 array of shape (n_outputs, n_required_features, n_grid) + grid : float32 C-contiguous array of shape (n_grid,) + Grid values for the required feature. + required_feature : intp + Global index of the feature whose PDP is being computed. + out : float64 array of shape (n_outputs, n_grid) Output array; sums are accumulated in-place. """ cdef: - intp_t n_required = required_features.shape[0] intp_t n_grid = grid.shape[0] intp_t n_background = X_bg.shape[0] intp_t _TREE_LEAF_SENTINEL = TREE_LEAF - intp_t _LEAVE_MARKER = TREE_UNDEFINED if n_background == 0 or n_grid == 0 or self.node_count == 0: return - # ------------------------------------------------------------------ # - # feature_to_req_pos: global feature index -> position in # - # required_features (-1 if not a required feature) # - # ------------------------------------------------------------------ # - cdef intp_t[:] feature_to_req_pos = np.full(self.n_features, -1, dtype=np.intp) - cdef intp_t req_idx - for req_idx in range(n_required): - feature_to_req_pos[required_features[req_idx]] = req_idx - # ----------------------------------------------------------------------- # - # Phase 0: single DFS over the tree structure, O(node_count). It builds: # - # - node_depth_arr: The depth of each node (needed by Phase 1) # - # - leaf_conditions: With the details about split conditions along its # - # root-to-leaf path (needed by Phase 2) # - # No large per-node arrays (no leaf_uniq, no leaf_n_uniq). # + # Phase 0: single DFS over the tree structure, O(node_count). It builds: # + # leaf_conditions: split conditions along each root-to-leaf path # + # (needed by Phase 2) # # ----------------------------------------------------------------------- # children_left_arr = np.asarray(self.children_left) @@ -1603,24 +1588,16 @@ cdef class Tree: cdef intp_t _n_clf_classes = self.max_n_classes cdef bint _is_clf = _n_clf_classes > 1 - cdef intp_t[:] node_depth_arr = np.empty(self.node_count, dtype=np.intp) - - # leaf_conditions[node_id] = (path_conditions, path_req_slots, leaf_values) - # path_conditions : {req_pos: [(split_threshold, goes_left), ...]} - # path_req_slots : ordered list of (d, req_pos) for each required feature - # first seen on this path. d is the feature's position - # among ALL unique features on the path, matching Phase 1's - # feature_to_d assignment. + # leaf_conditions[node_id] = (req_seen, conditions, leaf_values) + # req_seen : True if required_feature appeared on this path + # conditions : list of (split_threshold, goes_left) for required_feature + # leaf_values: ndarray of shape (n_outputs,) leaf_conditions = {} - # Python DFS — stack entries: - # (node_id, depth, path_conditions, path_req_slots, seen_features, n_seen) - # seen_features : set of global feature indices already seen on this path - # n_seen : len(seen_features), tracked explicitly to assign d - dfs_stack = [(0, 0, {}, [], set(), 0)] + # Python DFS — stack entries: (node_id, req_seen, req_conditions) + dfs_stack = [(0, False, [])] while dfs_stack: - node_id, depth, path_conditions, path_req_slots, seen_features, n_seen = dfs_stack.pop() - node_depth_arr[node_id] = depth + node_id, req_seen, req_conditions = dfs_stack.pop() if children_left_arr[node_id] == TREE_LEAF: if _is_clf: @@ -1649,82 +1626,45 @@ cdef class Tree: for k in range(self.n_outputs)], dtype=np.float64, ) - leaf_conditions[node_id] = (path_conditions, path_req_slots, leaf_values) + leaf_conditions[node_id] = (req_seen, req_conditions, leaf_values) else: split_feature = int(features_arr[node_id]) split_threshold = float(thresholds_arr[node_id]) left_child_idx = int(children_left_arr[node_id]) right_child_idx = int(children_right_arr[node_id]) - req_pos = int(feature_to_req_pos[split_feature]) - # Compute d (position among ALL unique path features) for this feature. - # This mirrors Phase 1's uniform path-state update for all features. - is_new = split_feature not in seen_features - d = n_seen # used only when is_new - child_seen = seen_features | {split_feature} if is_new else seen_features - child_n_seen = n_seen + (1 if is_new else 0) + child_req_seen = req_seen or (split_feature == required_feature) for child_idx, goes_left in ((left_child_idx, True), (right_child_idx, False)): - child_conditions = {rp: list(conds) for rp, conds in path_conditions.items()} - child_req_slots = list(path_req_slots) - if req_pos >= 0: - if req_pos not in child_conditions: - # First occurrence of this required feature: record its d - # (position among ALL unique path features). - child_req_slots.append((d, req_pos)) - child_conditions[req_pos] = [] - child_conditions[req_pos].append((split_threshold, goes_left)) - dfs_stack.append((child_idx, depth + 1, child_conditions, child_req_slots, - child_seen, child_n_seen)) + child_conditions = list(req_conditions) + if split_feature == required_feature: + child_conditions.append((split_threshold, goes_left)) + dfs_stack.append((child_idx, child_req_seen, child_conditions)) # ------------------------------------------------------------------ # # Phase 1: background pass (Cython hot loop) # # # # Per-leaf accumulators (node_count sized, but only leaves used): # - # reached[leaf] — samples that naturally landed # - # here (main path) # - # diverged_once[leaf, d] — samples whose shadow (that # - # diverged on the required feature # - # at path position d) landed here # - # # - # Path state (O(D)) — tracks ALL features, not only required features: # - # n_unique_features_on_path — #unique features (all) # - # seen on the active path # - # feature_to_d[feature_idx] — d of feature in path # - # order (-1 if not yet # - # seen) # - # feature_first_added_at_depth[depth] — global feature index # - # first added at this depth # - # (-1 if none or already # - # on path) # + # reached[leaf] — samples that naturally landed here # + # diverged_on_required[leaf] — samples whose shadow (spawned at a # + # required_feature split) landed here # # # - # Stack encoding: # - # stack_node[stack_size] >= 0 → ENTER that node # - # stack_node[stack_size] == _LEAVE_MARKER # - # → LEAVE: undo path state at # - # depth stack_diverged_on_target[sz] # + # Stack entries: (node_idx, diverged_on_target) # + # diverged_on_target >= 0 → this is a shadow traversal # + # diverged_on_target < 0 → this is the main traversal # # ------------------------------------------------------------------ # cdef intp_t[:] reached = np.zeros(self.node_count, dtype=np.intp) - cdef intp_t[:, :] diverged_once = np.zeros( - (self.node_count, max(self.max_depth, 1)), dtype=np.intp - ) - - # Path state — initialised once, self-cleaning via LEAVE markers - cdef intp_t n_unique_features_on_path = 0 - cdef intp_t[:] feature_to_d = np.full(self.n_features, -1, dtype=np.intp) - cdef intp_t[:] feature_first_added_at_depth = np.full(self.max_depth + 1, -1, dtype=np.intp) + cdef intp_t[:] diverged_on_required = np.zeros(self.node_count, dtype=np.intp) - # Stack: LEAVE entries double the worst-case depth; +4 for safety - cdef intp_t stack_capacity = 2 * (n_required + 2) * (self.max_depth + 2) + 4 + # Stack: at most 2 entries pushed per node (proceed + diverge); +4 for safety + cdef intp_t stack_capacity = 2 * (self.max_depth + 2) + 4 cdef intp_t[:] stack_node = np.empty(stack_capacity, dtype=np.intp) - # When node_idx == _LEAVE_MARKER include the depth, - # when node_idx is of an inner node include the feature we diverged on. cdef intp_t[:] stack_diverged_on_target = np.empty(stack_capacity, dtype=np.intp) cdef: intp_t stack_size, sample_idx, node_idx, diverged_on_target - intp_t current_feature, current_depth, newly_added_feature + intp_t current_feature intp_t proceed_child, diverge_child Node* node bint sample_goes_left @@ -1737,51 +1677,28 @@ cdef class Tree: while stack_size > 0: stack_size -= 1 node_idx = stack_node[stack_size] - - if node_idx == _LEAVE_MARKER: - # ---- LEAVE marker: undo path state for this depth ---- # - current_depth = stack_diverged_on_target[stack_size] - newly_added_feature = feature_first_added_at_depth[current_depth] - if newly_added_feature >= 0: - feature_to_d[newly_added_feature] = -1 - n_unique_features_on_path -= 1 - feature_first_added_at_depth[current_depth] = -1 - continue - - # ---- ENTER node ---- # diverged_on_target = stack_diverged_on_target[stack_size] node = &self.nodes[node_idx] if node.left_child == _TREE_LEAF_SENTINEL: - # Leaf: O(1) update via feature_to_d if diverged_on_target < 0: reached[node_idx] += 1 else: - diverged_once[node_idx, feature_to_d[diverged_on_target]] += 1 + diverged_on_required[node_idx] += 1 else: current_feature = node.feature - current_depth = node_depth_arr[node_idx] sample_goes_left = X_bg[sample_idx, current_feature] <= node.threshold proceed_child = node.left_child if sample_goes_left else node.right_child diverge_child = node.right_child if sample_goes_left else node.left_child - stack_node[stack_size] = _LEAVE_MARKER - stack_diverged_on_target[stack_size] = current_depth - stack_size += 1 - - if feature_to_d[current_feature] < 0: - feature_to_d[current_feature] = n_unique_features_on_path - n_unique_features_on_path += 1 - feature_first_added_at_depth[current_depth] = current_feature - # Push natural direction (main path or shadow continues) stack_node[stack_size] = proceed_child stack_diverged_on_target[stack_size] = diverged_on_target stack_size += 1 - # Spawn a shadow only at target-feature splits. - if feature_to_req_pos[current_feature] >= 0 and (diverged_on_target < 0 or diverged_on_target == current_feature): + # Spawn a shadow only at required_feature splits. + if current_feature == required_feature: stack_node[stack_size] = diverge_child stack_diverged_on_target[stack_size] = current_feature stack_size += 1 @@ -1791,53 +1708,39 @@ cdef class Tree: # ------------------------------------------------------------------ # reached_arr = np.asarray(reached) - diverged_once_arr = np.asarray(diverged_once) - grid_arr = np.asarray(grid) - out_arr = np.asarray(out) + diverged_on_required_arr = np.asarray(diverged_on_required) + grid_arr = np.asarray(grid) # 1D: (n_grid,) + out_arr = np.asarray(out) # 2D: (n_outputs, n_grid) # avg_pred_sum accumulates leaf_values * reached[leaf] over all leaves. - # Adding it to out_arr (then dividing by m in the caller) gives the absolute mean prediction + # Adding it to out_arr (then dividing by m in the caller) gives the + # absolute mean prediction. avg_pred_sum = np.zeros(out_arr.shape[0], dtype=np.float64) - for leaf_node_idx, (path_conditions, path_req_slots, leaf_values) in leaf_conditions.items(): + for leaf_node_idx, (req_seen, conditions, leaf_values) in leaf_conditions.items(): avg_pred_sum += leaf_values * float(reached_arr[leaf_node_idx]) - n_unique_req = len(path_req_slots) - if n_unique_req == 0: - # No required features on this path; only contributes to avg_pred. + if not req_seen: + # required_feature never appeared on this path; only contributes + # to avg_pred. continue - grid_satisfies_conditions = np.ones((n_grid, n_unique_req), dtype=np.bool_) - for i, (d, req_pos) in enumerate(path_req_slots): - for split_threshold, goes_left in path_conditions.get(req_pos, []): - if goes_left: - grid_satisfies_conditions[:, i] &= grid_arr[:, req_pos] <= split_threshold - else: - grid_satisfies_conditions[:, i] &= grid_arr[:, req_pos] > split_threshold - - # d is the position of the feature among ALL unique path features - req_ds = [d for d, _ in path_req_slots] - # req_pos is the column index in the grid / out arrays - req_positions = [rp for _, rp in path_req_slots] - - diverged_once_for_leaf = diverged_once_arr[leaf_node_idx, req_ds].astype(np.float64) - multipliers = np.where( - grid_satisfies_conditions, # (n_grid, n_unique_req) - diverged_once_for_leaf, # broadcast over n_grid - -float(reached_arr[leaf_node_idx]), - ) + # Which grid values satisfy all split conditions on required_feature? + mask = np.ones(n_grid, dtype=np.bool_) + for split_threshold, goes_left in conditions: + if goes_left: + mask &= grid_arr <= split_threshold + else: + mask &= grid_arr > split_threshold - # Accumulate into out_arr[n_outputs, n_required_features, n_grid]: - # leaf_values[:, None, None] : (n_outputs, 1, 1) - # multipliers.T[None, :, :] : (1, n_unique_req, n_grid) - # product : (n_outputs, n_unique_req, n_grid) - out_arr[:, req_positions, :] += ( - leaf_values[:, None, None] * multipliers.T[None, :, :] - ) + diverged = float(diverged_on_required_arr[leaf_node_idx]) + multipliers = np.where(mask, diverged, -float(reached_arr[leaf_node_idx])) + + # out_arr: (n_outputs, n_grid) + out_arr += leaf_values[:, None] * multipliers[None, :] - # Fold avg_pred into out so the caller just divides by m (no external - # predict call needed). Broadcasts over all (n_required, n_grid) entries. - out_arr += avg_pred_sum[:, np.newaxis, np.newaxis] + # Fold avg_pred into out so the caller just divides by m. + out_arr += avg_pred_sum[:, np.newaxis] def _check_n_classes(n_classes, expected_dtype): From 14326e6f91495bbd78a16e026f66c29dbb4977a8 Mon Sep 17 00:00:00 2001 From: Ron Wettenstein Date: Sat, 13 Jun 2026 13:54:44 +0300 Subject: [PATCH 04/11] simpler algorithm --- sklearn/tree/_tree.pyx | 243 ++++++++++++++--------------------------- 1 file changed, 80 insertions(+), 163 deletions(-) diff --git a/sklearn/tree/_tree.pyx b/sklearn/tree/_tree.pyx index 9e2c1a3320b46..6f726e4fa25ee 100644 --- a/sklearn/tree/_tree.pyx +++ b/sklearn/tree/_tree.pyx @@ -1539,20 +1539,31 @@ cdef class Tree: intp_t required_feature, float64_t[:, ::1] out, ): - """Partial dependence via background-data tree traversal (tree_accurate method). + """Partial dependence via background-data tree traversal (tree_accurate). - Uses a single DFS per background sample: for each sample the natural - prediction path is followed, and at every split on the required feature a - shadow path is spawned that diverges to the other branch. Shadows may - re-diverge on the same feature (handling repeated splits), so the - per-sample work is O(D + total shadow nodes). - Per-leaf statistics (``reached`` and ``diverged_on_required``) are - accumulated with O(1) per-leaf updates. + Two traversals are performed, both using the same split test as + prediction (so missing-value routing and any future split kind stay + consistent automatically): - Memory overhead is O(N) where N = node_count. + * Background pass -- the heavy, ``O(n_background * depth)`` part. Each + sample is routed through the tree, branching both ways at splits on + ``required_feature`` and following the sample everywhere else; the + number of arrivals per leaf is accumulated into ``count``. + ``count[leaf]`` is thus the number of background samples routed to + that leaf when ``required_feature`` is treated as a wildcard. - Accumulates raw sums into ``out``; the caller is responsible for dividing - by the number of background samples (and, for forests, by n_estimators). + * Grid pass -- ``O(n_grid * node_count)``. Each grid value is routed + through the tree, following it at ``required_feature`` splits and + branching both ways elsewhere; ``value * count`` is accumulated at + every leaf reached. The result for grid value ``g`` equals + ``n_background * PDP(g)``. + + ``required_feature`` of each sample is never read -- marginalising it is + exactly the both-ways branching above. Working memory is + ``O(node_count)``. + + Accumulates raw sums into ``out``; the caller divides by the number of + background samples (and, for forests, by n_estimators). Parameters ---------- @@ -1569,178 +1580,84 @@ cdef class Tree: intp_t n_grid = grid.shape[0] intp_t n_background = X_bg.shape[0] intp_t _TREE_LEAF_SENTINEL = TREE_LEAF + intp_t K = out.shape[0] + intp_t stack_size, sample_idx, node_idx, j, k, cnt + float64_t g + Node* node + float64_t[:, ::1] node_values if n_background == 0 or n_grid == 0 or self.node_count == 0: return - # ----------------------------------------------------------------------- # - # Phase 0: single DFS over the tree structure, O(node_count). It builds: # - # leaf_conditions: split conditions along each root-to-leaf path # - # (needed by Phase 2) # - # ----------------------------------------------------------------------- # - - children_left_arr = np.asarray(self.children_left) - children_right_arr = np.asarray(self.children_right) - features_arr = np.asarray(self.feature) - thresholds_arr = np.asarray(self.threshold) - - # Classifier detection: regressors have n_classes[0] == 1. - cdef intp_t _n_clf_classes = self.max_n_classes - cdef bint _is_clf = _n_clf_classes > 1 - - # leaf_conditions[node_id] = (req_seen, conditions, leaf_values) - # req_seen : True if required_feature appeared on this path - # conditions : list of (split_threshold, goes_left) for required_feature - # leaf_values: ndarray of shape (n_outputs,) - leaf_conditions = {} - - # Python DFS — stack entries: (node_id, req_seen, req_conditions) - dfs_stack = [(0, False, [])] - while dfs_stack: - node_id, req_seen, req_conditions = dfs_stack.pop() - - if children_left_arr[node_id] == TREE_LEAF: - if _is_clf: - # Normalise raw class counts to probabilities, matching - # predict_proba: divide each class count by the total count. - # Sum the counts from the buffer directly (same as predict_proba). - total = sum( - self.value[node_id * self.value_stride + c] - for c in range(_n_clf_classes) - ) - if _n_clf_classes == 2: - # Binary: keep only positive-class probability. - leaf_values = np.array( - [self.value[node_id * self.value_stride + 1] / total], - dtype=np.float64, - ) - else: - leaf_values = np.array( - [self.value[node_id * self.value_stride + c] / total - for c in range(_n_clf_classes)], - dtype=np.float64, - ) - else: - leaf_values = np.array( - [self.value[node_id * self.value_stride + k] - for k in range(self.n_outputs)], - dtype=np.float64, - ) - leaf_conditions[node_id] = (req_seen, req_conditions, leaf_values) - else: - split_feature = int(features_arr[node_id]) - split_threshold = float(thresholds_arr[node_id]) - left_child_idx = int(children_left_arr[node_id]) - right_child_idx = int(children_right_arr[node_id]) - - child_req_seen = req_seen or (split_feature == required_feature) - - for child_idx, goes_left in ((left_child_idx, True), (right_child_idx, False)): - child_conditions = list(req_conditions) - if split_feature == required_feature: - child_conditions.append((split_threshold, goes_left)) - dfs_stack.append((child_idx, child_req_seen, child_conditions)) - - # ------------------------------------------------------------------ # - # Phase 1: background pass (Cython hot loop) # - # # - # Per-leaf accumulators (node_count sized, but only leaves used): # - # reached[leaf] — samples that naturally landed here # - # diverged_on_required[leaf] — samples whose shadow (spawned at a # - # required_feature split) landed here # - # # - # Stack entries: (node_idx, diverged_on_target) # - # diverged_on_target >= 0 → this is a shadow traversal # - # diverged_on_target < 0 → this is the main traversal # - # ------------------------------------------------------------------ # - - cdef intp_t[:] reached = np.zeros(self.node_count, dtype=np.intp) - cdef intp_t[:] diverged_on_required = np.zeros(self.node_count, dtype=np.intp) - - # Stack: at most 2 entries pushed per node (proceed + diverge); +4 for safety - cdef intp_t stack_capacity = 2 * (self.max_depth + 2) + 4 - cdef intp_t[:] stack_node = np.empty(stack_capacity, dtype=np.intp) - cdef intp_t[:] stack_diverged_on_target = np.empty(stack_capacity, dtype=np.intp) + node_values = self._tree_accurate_node_values() - cdef: - intp_t stack_size, sample_idx, node_idx, diverged_on_target - intp_t current_feature - intp_t proceed_child, diverge_child - Node* node - bint sample_goes_left + # ---- Pass B: wildcard routing counts (heavy, m-dependent) ----------- + cdef intp_t[::1] count = np.zeros(self.node_count, dtype=np.intp) + cdef intp_t stack_capacity = 2 * (self.max_depth + 2) + 4 + cdef intp_t[::1] stack_node = np.empty(stack_capacity, dtype=np.intp) for sample_idx in range(n_background): stack_size = 1 stack_node[0] = 0 - stack_diverged_on_target[0] = -1 - while stack_size > 0: stack_size -= 1 node_idx = stack_node[stack_size] - diverged_on_target = stack_diverged_on_target[stack_size] node = &self.nodes[node_idx] - if node.left_child == _TREE_LEAF_SENTINEL: - if diverged_on_target < 0: - reached[node_idx] += 1 - else: - diverged_on_required[node_idx] += 1 + count[node_idx] += 1 + elif node.feature == required_feature: + stack_node[stack_size] = node.left_child + stack_size += 1 + stack_node[stack_size] = node.right_child + stack_size += 1 else: - current_feature = node.feature - sample_goes_left = X_bg[sample_idx, current_feature] <= node.threshold - - proceed_child = node.left_child if sample_goes_left else node.right_child - diverge_child = node.right_child if sample_goes_left else node.left_child - - # Push natural direction (main path or shadow continues) - stack_node[stack_size] = proceed_child - stack_diverged_on_target[stack_size] = diverged_on_target + if X_bg[sample_idx, node.feature] <= node.threshold: + stack_node[stack_size] = node.left_child + else: + stack_node[stack_size] = node.right_child stack_size += 1 - # Spawn a shadow only at required_feature splits. - if current_feature == required_feature: - stack_node[stack_size] = diverge_child - stack_diverged_on_target[stack_size] = current_feature - stack_size += 1 - - # ------------------------------------------------------------------ # - # Phase 2: consumer patterns + accumulation (numpy, O(L * n_grid)) # - # ------------------------------------------------------------------ # - - reached_arr = np.asarray(reached) - diverged_on_required_arr = np.asarray(diverged_on_required) - grid_arr = np.asarray(grid) # 1D: (n_grid,) - out_arr = np.asarray(out) # 2D: (n_outputs, n_grid) - - # avg_pred_sum accumulates leaf_values * reached[leaf] over all leaves. - # Adding it to out_arr (then dividing by m in the caller) gives the - # absolute mean prediction. - avg_pred_sum = np.zeros(out_arr.shape[0], dtype=np.float64) - - for leaf_node_idx, (req_seen, conditions, leaf_values) in leaf_conditions.items(): - avg_pred_sum += leaf_values * float(reached_arr[leaf_node_idx]) - - if not req_seen: - # required_feature never appeared on this path; only contributes - # to avg_pred. - continue - - # Which grid values satisfy all split conditions on required_feature? - mask = np.ones(n_grid, dtype=np.bool_) - for split_threshold, goes_left in conditions: - if goes_left: - mask &= grid_arr <= split_threshold + # ---- Pass C: per grid value, accumulate value * count --------------- + for j in range(n_grid): + g = grid[j] + stack_size = 1 + stack_node[0] = 0 + while stack_size > 0: + stack_size -= 1 + node_idx = stack_node[stack_size] + node = &self.nodes[node_idx] + if node.left_child == _TREE_LEAF_SENTINEL: + cnt = count[node_idx] + if cnt != 0: + for k in range(K): + out[k, j] += node_values[node_idx, k] * cnt + elif node.feature == required_feature: + if g <= node.threshold: + stack_node[stack_size] = node.left_child + else: + stack_node[stack_size] = node.right_child + stack_size += 1 else: - mask &= grid_arr > split_threshold - - diverged = float(diverged_on_required_arr[leaf_node_idx]) - multipliers = np.where(mask, diverged, -float(reached_arr[leaf_node_idx])) + stack_node[stack_size] = node.left_child + stack_size += 1 + stack_node[stack_size] = node.right_child + stack_size += 1 - # out_arr: (n_outputs, n_grid) - out_arr += leaf_values[:, None] * multipliers[None, :] + def _tree_accurate_node_values(self): + """Per-node output values matching predict / decision_function. - # Fold avg_pred into out so the caller just divides by m. - out_arr += avg_pred_sum[:, np.newaxis] + Returns a C-contiguous ``(node_count, K)`` float64 array: regression + values, all class probabilities for multiclass, or the positive-class + probability only for binary classification. + """ + value_3d = np.asarray(self._get_value_ndarray()) + if self.max_n_classes > 1: + class_counts = value_3d[:, 0, :] + probs = class_counts / class_counts.sum(axis=1, keepdims=True) + if self.max_n_classes == 2: + return np.ascontiguousarray(probs[:, 1:2]) + return np.ascontiguousarray(probs) + return np.ascontiguousarray(value_3d[:, :, 0]) def _check_n_classes(n_classes, expected_dtype): From 15c9bd8d8e55359c0c09397ba03773d737e545f9 Mon Sep 17 00:00:00 2001 From: Ron Wettenstein Date: Sat, 13 Jun 2026 13:58:54 +0300 Subject: [PATCH 05/11] bugfix --- sklearn/inspection/_partial_dependence.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/sklearn/inspection/_partial_dependence.py b/sklearn/inspection/_partial_dependence.py index 6993f7686c8bb..e2c1a09bddfc4 100644 --- a/sklearn/inspection/_partial_dependence.py +++ b/sklearn/inspection/_partial_dependence.py @@ -668,7 +668,8 @@ def partial_dependence( ) if method == "tree_accurate": - if any(isinstance(f, (list, tuple)) for f in features): + _features_iter = [features] if isinstance(features, (str, int)) else features + if any(isinstance(f, (list, tuple)) for f in _features_iter): raise ValueError( "The 'tree_accurate' method does not support joint PDP " "(tuples in features). Use method='brute' instead." From fb35923a46c47b129bd08bd03d0cb158dc931468 Mon Sep 17 00:00:00 2001 From: Ron Wettenstein Date: Sat, 13 Jun 2026 14:13:17 +0300 Subject: [PATCH 06/11] Pass A+B not B+C --- sklearn/tree/_tree.pyx | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/sklearn/tree/_tree.pyx b/sklearn/tree/_tree.pyx index 6f726e4fa25ee..be0e89d4dfa86 100644 --- a/sklearn/tree/_tree.pyx +++ b/sklearn/tree/_tree.pyx @@ -1591,7 +1591,7 @@ cdef class Tree: node_values = self._tree_accurate_node_values() - # ---- Pass B: wildcard routing counts (heavy, m-dependent) ----------- + # ---- Pass A: wildcard routing counts (heavy, m-dependent) ----------- cdef intp_t[::1] count = np.zeros(self.node_count, dtype=np.intp) cdef intp_t stack_capacity = 2 * (self.max_depth + 2) + 4 cdef intp_t[::1] stack_node = np.empty(stack_capacity, dtype=np.intp) @@ -1617,7 +1617,7 @@ cdef class Tree: stack_node[stack_size] = node.right_child stack_size += 1 - # ---- Pass C: per grid value, accumulate value * count --------------- + # ---- Pass B: per grid value, accumulate value * count --------------- for j in range(n_grid): g = grid[j] stack_size = 1 From 246de2a31bf3f8e1acd30dc182f92c61306cb339 Mon Sep 17 00:00:00 2001 From: Ron Wettenstein Date: Tue, 8 Sep 2026 11:17:20 +0300 Subject: [PATCH 07/11] Support categorical features --- .../tests/test_partial_dependence.py | 149 +++++++++++++++++- sklearn/tree/_tree.pyx | 27 +++- 2 files changed, 169 insertions(+), 7 deletions(-) diff --git a/sklearn/inspection/tests/test_partial_dependence.py b/sklearn/inspection/tests/test_partial_dependence.py index 57c4dfdb1f482..5b3b0a7781d91 100644 --- a/sklearn/inspection/tests/test_partial_dependence.py +++ b/sklearn/inspection/tests/test_partial_dependence.py @@ -38,7 +38,12 @@ StandardScaler, scale, ) -from sklearn.tree import DecisionTreeClassifier, DecisionTreeRegressor +from sklearn.tree import ( + DecisionTreeClassifier, + DecisionTreeRegressor, + ExtraTreeClassifier, + ExtraTreeRegressor, +) from sklearn.tree.tests.test_tree import assert_is_subtree from sklearn.utils._testing import assert_allclose, assert_array_equal from sklearn.utils.fixes import _IS_32BIT @@ -1445,3 +1450,145 @@ def test_tree_accurate_multi_output_matches_brute(): np.testing.assert_allclose( result_accurate["average"], result_brute["average"], atol=1e-6 ) + + +def _make_categorical_data(seed=0, n_samples=2000): + """Data whose first column is an unordered category with a non-monotone effect. + + The effect is deliberately not monotone in the category code, so a split test + that compares the code against a numeric threshold cannot reproduce it. + """ + rng = np.random.RandomState(seed) + category = rng.randint(0, 6, n_samples) + effect = np.array([0.0, 5.0, 1.0, 6.0, 2.0, 7.0]) + x1 = rng.normal(size=n_samples) + X = np.column_stack([category.astype(float), x1]) + y = effect[category] + 2.0 * x1 + 0.01 * rng.normal(size=n_samples) + return X, y + + +def _assert_split_kinds_present(est, expected_kind): + """Guard that the fitted tree really contains the split kind under test.""" + tree = est.tree_ if hasattr(est, "tree_") else est.estimators_[0].tree_ + internal = tree.children_left != -1 + assert expected_kind in set(np.asarray(tree.split_kind)[internal]) + + +@pytest.mark.parametrize("target_feature", [0, 1]) +def test_tree_accurate_categorical_bitset_matches_brute(target_feature): + """Bitset categorical splits must be routed with the prediction split test. + + Covers the category both as the PDP target feature and as a complementary + feature that is marginalised over. + """ + X, y = _make_categorical_data() + est = DecisionTreeRegressor( + categorical_features=[0], max_depth=6, random_state=0 + ).fit(X, y) + _assert_split_kinds_present(est, 1) # SPLIT_CATEGORICAL_BITSET + + result_ta = partial_dependence( + est, + X, + features=[target_feature], + method="tree_accurate", + categorical_features=[0], + ) + result_br = partial_dependence( + est, + X, + features=[target_feature], + method="brute", + categorical_features=[0], + ) + np.testing.assert_allclose( + result_ta["average"], result_br["average"], rtol=1e-7, atol=1e-9 + ) + + +def test_tree_accurate_categorical_hash_matches_brute(): + """Hash-routed categorical splits (ExtraTree) must match brute.""" + X, y = _make_categorical_data() + est = ExtraTreeRegressor(categorical_features=[0], max_depth=8, random_state=0).fit( + X, y + ) + _assert_split_kinds_present(est, 2) # SPLIT_CATEGORICAL_HASH + + result_ta = partial_dependence( + est, X, features=[0], method="tree_accurate", categorical_features=[0] + ) + result_br = partial_dependence( + est, X, features=[0], method="brute", categorical_features=[0] + ) + np.testing.assert_allclose( + result_ta["average"], result_br["average"], rtol=1e-7, atol=1e-9 + ) + + +def test_tree_accurate_categorical_binary_classifier_matches_brute(): + """Binary classification with categorical splits must match brute.""" + X, y = _make_categorical_data() + y_bin = (y > np.median(y)).astype(int) + est = DecisionTreeClassifier( + categorical_features=[0], max_depth=6, random_state=0 + ).fit(X, y_bin) + _assert_split_kinds_present(est, 1) + + result_ta = partial_dependence( + est, X, features=[0], method="tree_accurate", categorical_features=[0] + ) + result_br = partial_dependence( + est, X, features=[0], method="brute", categorical_features=[0] + ) + np.testing.assert_allclose( + result_ta["average"], result_br["average"], rtol=1e-7, atol=1e-9 + ) + + +def test_tree_accurate_categorical_multiclass_matches_brute(): + """Multiclass with categorical splits must match brute. + + ``splitter='best'`` rejects categorical features for multiclass, so this + uses ExtraTreeClassifier, which routes categories through the hash split. + """ + X, y = _make_categorical_data() + y_multi = np.digitize(y, np.quantile(y, [0.33, 0.66])) + est = ExtraTreeClassifier( + categorical_features=[0], max_depth=8, random_state=0 + ).fit(X, y_multi) + _assert_split_kinds_present(est, 2) + + result_ta = partial_dependence( + est, X, features=[0], method="tree_accurate", categorical_features=[0] + ) + result_br = partial_dependence( + est, X, features=[0], method="brute", categorical_features=[0] + ) + assert result_ta["average"].shape[0] == 3 + np.testing.assert_allclose( + result_ta["average"], result_br["average"], rtol=1e-7, atol=1e-9 + ) + + +@pytest.mark.parametrize("missing_feature", [0, 1]) +def test_tree_accurate_missing_values_match_brute(missing_feature): + """NaNs must follow ``missing_go_to_left`` rather than falling right. + + ``missing_feature`` is the column carrying NaNs; the PDP target is always + feature 0, so this covers NaNs in both the target and a marginalised feature. + """ + X, y = _make_categorical_data() + rng = np.random.RandomState(1) + X = X.copy() + X[rng.rand(X.shape[0]) < 0.2, missing_feature] = np.nan + + est = DecisionTreeRegressor(max_depth=6, random_state=0).fit(X, y) + tree = est.tree_ + internal = tree.children_left != -1 + assert np.asarray(tree.missing_go_to_left)[internal].sum() > 0 + + result_ta = partial_dependence(est, X, features=[0], method="tree_accurate") + result_br = partial_dependence(est, X, features=[0], method="brute") + np.testing.assert_allclose( + result_ta["average"], result_br["average"], rtol=1e-7, atol=1e-9 + ) diff --git a/sklearn/tree/_tree.pyx b/sklearn/tree/_tree.pyx index be0e89d4dfa86..278550aaa745c 100644 --- a/sklearn/tree/_tree.pyx +++ b/sklearn/tree/_tree.pyx @@ -1541,9 +1541,9 @@ cdef class Tree: ): """Partial dependence via background-data tree traversal (tree_accurate). - Two traversals are performed, both using the same split test as - prediction (so missing-value routing and any future split kind stay - consistent automatically): + Two traversals are performed, both routing through ``goes_left`` -- + the same split test as prediction -- so numeric, categorical and + missing-value routing stay consistent automatically: * Background pass -- the heavy, ``O(n_background * depth)`` part. Each sample is routed through the tree, branching both ways at splits on @@ -1582,7 +1582,8 @@ cdef class Tree: intp_t _TREE_LEAF_SENTINEL = TREE_LEAF intp_t K = out.shape[0] intp_t stack_size, sample_idx, node_idx, j, k, cnt - float64_t g + float32_t g + bint go_left Node* node float64_t[:, ::1] node_values @@ -1611,7 +1612,14 @@ cdef class Tree: stack_node[stack_size] = node.right_child stack_size += 1 else: - if X_bg[sample_idx, node.feature] <= node.threshold: + go_left = goes_left( + node.threshold, + node.left_cat_bitset, + node.missing_go_to_left, + node.split_kind, + X_bg[sample_idx, node.feature], + ) + if go_left: stack_node[stack_size] = node.left_child else: stack_node[stack_size] = node.right_child @@ -1632,7 +1640,14 @@ cdef class Tree: for k in range(K): out[k, j] += node_values[node_idx, k] * cnt elif node.feature == required_feature: - if g <= node.threshold: + go_left = goes_left( + node.threshold, + node.left_cat_bitset, + node.missing_go_to_left, + node.split_kind, + g, + ) + if go_left: stack_node[stack_size] = node.left_child else: stack_node[stack_size] = node.right_child From f6f20f384ebf1e6a2e548dace351d836f8f6670f Mon Sep 17 00:00:00 2001 From: Ron Wettenstein Date: Sat, 12 Sep 2026 12:50:15 +0300 Subject: [PATCH 08/11] make the tests more compact --- .../tests/test_partial_dependence.py | 291 +++++------------- 1 file changed, 82 insertions(+), 209 deletions(-) diff --git a/sklearn/inspection/tests/test_partial_dependence.py b/sklearn/inspection/tests/test_partial_dependence.py index 5b3b0a7781d91..1fb0c918fc0fb 100644 --- a/sklearn/inspection/tests/test_partial_dependence.py +++ b/sklearn/inspection/tests/test_partial_dependence.py @@ -44,6 +44,7 @@ ExtraTreeClassifier, ExtraTreeRegressor, ) +from sklearn.tree._utils import SPLIT_CATEGORICAL_BITSET, SPLIT_CATEGORICAL_HASH from sklearn.tree.tests.test_tree import assert_is_subtree from sklearn.utils._testing import assert_allclose, assert_array_equal from sklearn.utils.fixes import _IS_32BIT @@ -1265,6 +1266,28 @@ def test_partial_dependence_empty_categorical_features(): # ============================================================================= +def _assert_tree_accurate_matches_brute(est, X, feature, **pd_kwargs): + """Assert tree_accurate reproduces brute for one feature of ``X``. + + ``pd_kwargs`` is forwarded to both calls, so the two runs differ only in + ``method``. Returns the tree_accurate result so callers can assert more. + """ + result_ta = partial_dependence( + est, X, features=[feature], method="tree_accurate", **pd_kwargs + ) + result_br = partial_dependence( + est, X, features=[feature], method="brute", **pd_kwargs + ) + np.testing.assert_allclose( + result_ta["average"], + result_br["average"], + rtol=1e-7, + atol=1e-9, + err_msg=f"tree_accurate does not match brute for feature {feature}", + ) + return result_ta + + @pytest.mark.parametrize("seed", range(3)) @pytest.mark.parametrize( "Estimator", @@ -1287,49 +1310,7 @@ def test_tree_accurate_matches_brute(Estimator, seed): est = Estimator(**kwargs).fit(X, y) for feature in range(n_features): - pdp_brute = partial_dependence( - est, X, features=[feature], method="brute", grid_resolution=20 - ) - pdp_fast = partial_dependence( - est, X, features=[feature], method="tree_accurate", grid_resolution=20 - ) - np.testing.assert_allclose( - pdp_fast["average"], - pdp_brute["average"], - rtol=1e-4, - atol=1e-6, - err_msg=f"Mismatch on feature {feature} with seed {seed}", - ) - - -@pytest.mark.parametrize( - "Estimator", - [DecisionTreeRegressor, RandomForestRegressor], -) -def test_tree_accurate_output_shape(Estimator): - """Output shape should be (1, grid_resolution) for each feature.""" - rng = np.random.RandomState(0) - X = rng.randn(100, 4).astype(np.float64) - y = rng.randn(100) - - kwargs = dict(max_depth=3, random_state=0) - if Estimator is RandomForestRegressor: - kwargs["n_estimators"] = 2 - est = Estimator(**kwargs).fit(X, y) - - grid_resolution = 15 - for feature in range(4): - result = partial_dependence( - est, - X, - features=[feature], - method="tree_accurate", - grid_resolution=grid_resolution, - ) - assert result["average"].shape == (1, grid_resolution), ( - f"Expected (1, {grid_resolution}), got {result['average'].shape}" - ) - assert result["grid_values"][0].shape == (grid_resolution,) + _assert_tree_accurate_matches_brute(est, X, feature, grid_resolution=20) def test_tree_accurate_repeated_feature_in_path(): @@ -1345,22 +1326,11 @@ def test_tree_accurate_repeated_feature_in_path(): est = DecisionTreeRegressor(max_depth=6, random_state=0).fit(X, y) - pdp_brute = partial_dependence( - est, X, features=[0], method="brute", grid_resolution=30 - ) - pdp_fast = partial_dependence( - est, X, features=[0], method="tree_accurate", grid_resolution=30 - ) - np.testing.assert_allclose( - pdp_fast["average"], - pdp_brute["average"], - rtol=1e-4, - atol=1e-6, - ) + _assert_tree_accurate_matches_brute(est, X, 0, grid_resolution=30) -def test_tree_accurate_kind_not_average_raises(): - """kind != 'average' must raise ValueError.""" +def test_tree_accurate_kind_not_average_and_sample_weight_raise(): + """kind != 'average' and sample_weight != None must raise ValueError.""" rng = np.random.RandomState(0) X = rng.randn(50, 3).astype(np.float64) y = rng.randn(50) @@ -1370,6 +1340,11 @@ def test_tree_accurate_kind_not_average_raises(): with pytest.raises(ValueError, match="'tree_accurate' method only supports"): partial_dependence(est, X, features=[0], method="tree_accurate", kind=kind) + with pytest.raises(ValueError, match="'tree_accurate' method can only be applied"): + partial_dependence( + est, X, features=[0], method="tree_accurate", sample_weight=np.ones(50) + ) + def test_tree_accurate_unsupported_estimator_raises(): """Non-tree estimators must raise ValueError.""" @@ -1380,31 +1355,11 @@ def test_tree_accurate_unsupported_estimator_raises(): partial_dependence(est, X, features=[0], method="tree_accurate") -def test_tree_accurate_sample_weight_raises(): - """sample_weight != None must raise ValueError.""" - rng = np.random.RandomState(0) - X = rng.randn(50, 3).astype(np.float64) - y = rng.randn(50) - est = DecisionTreeRegressor(max_depth=3, random_state=0).fit(X, y) - - with pytest.raises(ValueError, match="'tree_accurate' method can only be applied"): - partial_dependence( - est, X, features=[0], method="tree_accurate", sample_weight=np.ones(50) - ) - - def test_tree_accurate_binary_classifier_matches_brute(): """Binary DecisionTreeClassifier: tree_accurate PDP must match brute.""" X, y = make_classification(n_samples=100, n_features=4, random_state=0) est = DecisionTreeClassifier(max_depth=3, random_state=0).fit(X, y) - - result_ta = partial_dependence( - est, X, features=[0], method="tree_accurate", grid_resolution=10 - ) - result_br = partial_dependence( - est, X, features=[0], method="brute", grid_resolution=10 - ) - np.testing.assert_allclose(result_ta["average"], result_br["average"], atol=1e-6) + _assert_tree_accurate_matches_brute(est, X, 0, grid_resolution=10) def test_tree_accurate_multiclass_classifier_matches_brute(): @@ -1418,14 +1373,7 @@ def test_tree_accurate_multiclass_classifier_matches_brute(): random_state=0, ) est = DecisionTreeClassifier(max_depth=3, random_state=0).fit(X, y) - - result_ta = partial_dependence( - est, X, features=[0], method="tree_accurate", grid_resolution=10 - ) - result_br = partial_dependence( - est, X, features=[0], method="brute", grid_resolution=10 - ) - np.testing.assert_allclose(result_ta["average"], result_br["average"], atol=1e-6) + _assert_tree_accurate_matches_brute(est, X, 0, grid_resolution=10) def test_tree_accurate_multi_output_matches_brute(): @@ -1438,18 +1386,10 @@ def test_tree_accurate_multi_output_matches_brute(): est = DecisionTreeRegressor(max_depth=4, random_state=0).fit(X, Y) assert est.n_outputs_ == 2 - result_accurate = partial_dependence( - est, X, features=[0], method="tree_accurate", grid_resolution=10 - ) - result_brute = partial_dependence( - est, X, features=[0], method="brute", grid_resolution=10 - ) + result_accurate = _assert_tree_accurate_matches_brute(est, X, 0, grid_resolution=10) # Both methods should return shape (n_outputs, n_grid) = (2, 10). assert result_accurate["average"].shape == (2, 10) - np.testing.assert_allclose( - result_accurate["average"], result_brute["average"], atol=1e-6 - ) def _make_categorical_data(seed=0, n_samples=2000): @@ -1474,121 +1414,54 @@ def _assert_split_kinds_present(est, expected_kind): assert expected_kind in set(np.asarray(tree.split_kind)[internal]) -@pytest.mark.parametrize("target_feature", [0, 1]) -def test_tree_accurate_categorical_bitset_matches_brute(target_feature): - """Bitset categorical splits must be routed with the prediction split test. - - Covers the category both as the PDP target feature and as a complementary - feature that is marginalised over. - """ - X, y = _make_categorical_data() - est = DecisionTreeRegressor( - categorical_features=[0], max_depth=6, random_state=0 - ).fit(X, y) - _assert_split_kinds_present(est, 1) # SPLIT_CATEGORICAL_BITSET - - result_ta = partial_dependence( - est, - X, - features=[target_feature], - method="tree_accurate", - categorical_features=[0], - ) - result_br = partial_dependence( - est, - X, - features=[target_feature], - method="brute", - categorical_features=[0], - ) - np.testing.assert_allclose( - result_ta["average"], result_br["average"], rtol=1e-7, atol=1e-9 - ) - - -def test_tree_accurate_categorical_hash_matches_brute(): - """Hash-routed categorical splits (ExtraTree) must match brute.""" - X, y = _make_categorical_data() - est = ExtraTreeRegressor(categorical_features=[0], max_depth=8, random_state=0).fit( - X, y - ) - _assert_split_kinds_present(est, 2) # SPLIT_CATEGORICAL_HASH - - result_ta = partial_dependence( - est, X, features=[0], method="tree_accurate", categorical_features=[0] - ) - result_br = partial_dependence( - est, X, features=[0], method="brute", categorical_features=[0] - ) - np.testing.assert_allclose( - result_ta["average"], result_br["average"], rtol=1e-7, atol=1e-9 - ) - - -def test_tree_accurate_categorical_binary_classifier_matches_brute(): - """Binary classification with categorical splits must match brute.""" - X, y = _make_categorical_data() - y_bin = (y > np.median(y)).astype(int) - est = DecisionTreeClassifier( - categorical_features=[0], max_depth=6, random_state=0 - ).fit(X, y_bin) - _assert_split_kinds_present(est, 1) - - result_ta = partial_dependence( - est, X, features=[0], method="tree_accurate", categorical_features=[0] - ) - result_br = partial_dependence( - est, X, features=[0], method="brute", categorical_features=[0] - ) - np.testing.assert_allclose( - result_ta["average"], result_br["average"], rtol=1e-7, atol=1e-9 - ) - - -def test_tree_accurate_categorical_multiclass_matches_brute(): - """Multiclass with categorical splits must match brute. - - ``splitter='best'`` rejects categorical features for multiclass, so this - uses ExtraTreeClassifier, which routes categories through the hash split. - """ - X, y = _make_categorical_data() - y_multi = np.digitize(y, np.quantile(y, [0.33, 0.66])) - est = ExtraTreeClassifier( - categorical_features=[0], max_depth=8, random_state=0 - ).fit(X, y_multi) - _assert_split_kinds_present(est, 2) - - result_ta = partial_dependence( - est, X, features=[0], method="tree_accurate", categorical_features=[0] - ) - result_br = partial_dependence( - est, X, features=[0], method="brute", categorical_features=[0] - ) - assert result_ta["average"].shape[0] == 3 - np.testing.assert_allclose( - result_ta["average"], result_br["average"], rtol=1e-7, atol=1e-9 - ) - - -@pytest.mark.parametrize("missing_feature", [0, 1]) -def test_tree_accurate_missing_values_match_brute(missing_feature): - """NaNs must follow ``missing_go_to_left`` rather than falling right. +@pytest.mark.parametrize( + "estimator_cls, make_target, split_kind, n_outputs", + [ + pytest.param( + DecisionTreeRegressor, + lambda y: y, + SPLIT_CATEGORICAL_BITSET, + 1, + id="bitset-category-as-target", + ), + pytest.param( + ExtraTreeRegressor, + lambda y: y, + SPLIT_CATEGORICAL_HASH, + 1, + id="hash-regressor", + ), + pytest.param( + DecisionTreeClassifier, + lambda y: (y > np.median(y)).astype(int), + SPLIT_CATEGORICAL_BITSET, + 1, + id="bitset-binary-classifier", + ), + pytest.param( + ExtraTreeClassifier, + lambda y: np.digitize(y, np.quantile(y, [0.33, 0.66])), + SPLIT_CATEGORICAL_HASH, + 3, + id="hash-multiclass-classifier", + ), + ], +) +def test_tree_accurate_categorical_matches_brute( + estimator_cls, make_target, split_kind, n_outputs +): + """Categorical splits must be routed with the same split test as prediction. - ``missing_feature`` is the column carrying NaNs; the PDP target is always - feature 0, so this covers NaNs in both the target and a marginalised feature. + Covers both routing kinds, regression and classification, and the category + as the PDP target as well as a feature that is marginalised over. + ``splitter='best'`` rejects categorical features for multiclass, so the + multiclass case uses ExtraTreeClassifier and its hash routing. """ X, y = _make_categorical_data() - rng = np.random.RandomState(1) - X = X.copy() - X[rng.rand(X.shape[0]) < 0.2, missing_feature] = np.nan - - est = DecisionTreeRegressor(max_depth=6, random_state=0).fit(X, y) - tree = est.tree_ - internal = tree.children_left != -1 - assert np.asarray(tree.missing_go_to_left)[internal].sum() > 0 + est = estimator_cls(categorical_features=[0], max_depth=8, random_state=0) + est.fit(X, make_target(y)) + _assert_split_kinds_present(est, split_kind) - result_ta = partial_dependence(est, X, features=[0], method="tree_accurate") - result_br = partial_dependence(est, X, features=[0], method="brute") - np.testing.assert_allclose( - result_ta["average"], result_br["average"], rtol=1e-7, atol=1e-9 - ) + result = _assert_tree_accurate_matches_brute(est, X, 0, categorical_features=[0]) + result = _assert_tree_accurate_matches_brute(est, X, 1, categorical_features=[0]) + assert result["average"].shape[0] == n_outputs From 1027faec7b72c86596eff06ea27f5e1c1c76f3d1 Mon Sep 17 00:00:00 2001 From: Ron Wettenstein Date: Sat, 12 Sep 2026 14:08:36 +0300 Subject: [PATCH 09/11] Support ExtraTrees --- sklearn/inspection/_partial_dependence.py | 40 ++++++++++++----------- 1 file changed, 21 insertions(+), 19 deletions(-) diff --git a/sklearn/inspection/_partial_dependence.py b/sklearn/inspection/_partial_dependence.py index e2c1a09bddfc4..fa32ea6e69553 100644 --- a/sklearn/inspection/_partial_dependence.py +++ b/sklearn/inspection/_partial_dependence.py @@ -10,13 +10,15 @@ from scipy.stats.mstats import mquantiles from sklearn.base import is_classifier, is_regressor -from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor +from sklearn.ensemble import RandomForestRegressor +from sklearn.ensemble._forest import ForestClassifier, ForestRegressor from sklearn.ensemble._gb import BaseGradientBoosting from sklearn.ensemble._hist_gradient_boosting.gradient_boosting import ( BaseHistGradientBoosting, ) from sklearn.inspection._pd_utils import _check_feature_names, _get_feature_index -from sklearn.tree import DecisionTreeClassifier, DecisionTreeRegressor +from sklearn.tree import DecisionTreeRegressor +from sklearn.tree._classes import BaseDecisionTree from sklearn.utils import Bunch, _safe_indexing, check_array from sklearn.utils._indexing import ( _determine_key_type, @@ -227,9 +229,9 @@ def _partial_dependence_tree_accurate(est, X, feature, grid): Parameters ---------- - est : DecisionTreeRegressor, RandomForestRegressor, DecisionTreeClassifier, - or RandomForestClassifier - A fitted tree-based estimator. + est : BaseDecisionTree, ForestRegressor or ForestClassifier + A fitted single tree (``DecisionTree*``, ``ExtraTree*``) or forest of + trees (``RandomForest*``, ``ExtraTrees*``). X : ndarray of shape (n_samples, n_features), dtype=np.float32 Background dataset used for marginalisation. feature : int @@ -251,10 +253,10 @@ def _partial_dependence_tree_accurate(est, X, feature, grid): out = np.zeros((n_effective_outputs, len(grid)), dtype=np.float64) - if isinstance(est, (DecisionTreeRegressor, DecisionTreeClassifier)): + if isinstance(est, BaseDecisionTree): est.tree_.compute_partial_dependence_tree_accurate(X, grid, feature, out) out /= m - elif isinstance(est, (RandomForestRegressor, RandomForestClassifier)): + elif isinstance(est, (ForestRegressor, ForestClassifier)): n_trees = len(est.estimators_) for tree_est in est.estimators_: tree_est.tree_.compute_partial_dependence_tree_accurate( @@ -557,11 +559,15 @@ def partial_dependence( - `'brute'` is supported for any estimator, but is more computationally intensive. - - `'tree_accurate'` is supported for + - `'tree_accurate'` is supported for the single-tree estimators :class:`~sklearn.tree.DecisionTreeRegressor`, + :class:`~sklearn.tree.DecisionTreeClassifier`, + :class:`~sklearn.tree.ExtraTreeRegressor`, + :class:`~sklearn.tree.ExtraTreeClassifier` and the forests :class:`~sklearn.ensemble.RandomForestRegressor`, - :class:`~sklearn.tree.DecisionTreeClassifier`, and - :class:`~sklearn.ensemble.RandomForestClassifier` + :class:`~sklearn.ensemble.RandomForestClassifier`, + :class:`~sklearn.ensemble.ExtraTreesRegressor`, + :class:`~sklearn.ensemble.ExtraTreesClassifier`, when `kind='average'` and `response_method='decision_function'`. Joint PDP (tuples in `features`) is not supported. This method is equivalent to the `'brute'` method, @@ -680,18 +686,14 @@ def partial_dependence( "sample_weight is None." ) if not isinstance( - estimator, - ( - DecisionTreeRegressor, - RandomForestRegressor, - DecisionTreeClassifier, - RandomForestClassifier, - ), + estimator, (BaseDecisionTree, ForestRegressor, ForestClassifier) ): raise ValueError( "The 'tree_accurate' method only supports DecisionTreeRegressor, " - "RandomForestRegressor, DecisionTreeClassifier, and " - "RandomForestClassifier. Use method='brute' for other estimators." + "DecisionTreeClassifier, ExtraTreeRegressor, ExtraTreeClassifier, " + "RandomForestRegressor, RandomForestClassifier, " + "ExtraTreesRegressor and ExtraTreesClassifier. " + "Use method='brute' for other estimators." ) if is_classifier(estimator): if response_method == "auto": From 4793d462fdd93d5bc64653ec8c7cd75bd9ae4030 Mon Sep 17 00:00:00 2001 From: Ron Wettenstein Date: Sat, 12 Sep 2026 14:09:55 +0300 Subject: [PATCH 10/11] Add doc to PartialDependenceDisplay --- sklearn/inspection/_plot/partial_dependence.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/sklearn/inspection/_plot/partial_dependence.py b/sklearn/inspection/_plot/partial_dependence.py index 958f988ff98ac..db7b5ebc8d30e 100644 --- a/sklearn/inspection/_plot/partial_dependence.py +++ b/sklearn/inspection/_plot/partial_dependence.py @@ -432,6 +432,20 @@ def from_estimator( the average of the ICEs by design, it is not compatible with ICE and thus `kind` must be `'average'`. + - `'tree_accurate'` is supported for the single-tree estimators + :class:`~sklearn.tree.DecisionTreeRegressor`, + :class:`~sklearn.tree.DecisionTreeClassifier`, + :class:`~sklearn.tree.ExtraTreeRegressor`, + :class:`~sklearn.tree.ExtraTreeClassifier` and the forests + :class:`~sklearn.ensemble.RandomForestRegressor`, + :class:`~sklearn.ensemble.RandomForestClassifier`, + :class:`~sklearn.ensemble.ExtraTreesRegressor`, + :class:`~sklearn.ensemble.ExtraTreesClassifier`, + when `kind='average'` and `response_method='decision_function'`. + Joint PDP (tuples in `features`) is not supported. + This method is equivalent to the `'brute'` method, + but significantly faster for tree-based estimators. + - `'brute'` is supported for any estimator, but is more computationally intensive. From e7f248193cbb826883519bc5d1d8b6138cc9e506 Mon Sep 17 00:00:00 2001 From: Ron Wettenstein Date: Sat, 12 Sep 2026 14:46:30 +0300 Subject: [PATCH 11/11] modesty --- sklearn/inspection/_plot/partial_dependence.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sklearn/inspection/_plot/partial_dependence.py b/sklearn/inspection/_plot/partial_dependence.py index db7b5ebc8d30e..1418aa28a0908 100644 --- a/sklearn/inspection/_plot/partial_dependence.py +++ b/sklearn/inspection/_plot/partial_dependence.py @@ -444,7 +444,7 @@ def from_estimator( when `kind='average'` and `response_method='decision_function'`. Joint PDP (tuples in `features`) is not supported. This method is equivalent to the `'brute'` method, - but significantly faster for tree-based estimators. + but faster for tree-based estimators. - `'brute'` is supported for any estimator, but is more computationally intensive.