Share tree kernel: hoist group criterion base + unify builders (#945) - #974
Open
jeongyoonlee wants to merge 2 commits into
Open
Share tree kernel: hoist group criterion base + unify builders (#945)#974jeongyoonlee wants to merge 2 commits into
jeongyoonlee wants to merge 2 commits into
Conversation
jeongyoonlee
marked this pull request as ready for review
July 26, 2026 06:21
jeongyoonlee
requested review from
alexander-pv,
huigangchen,
paullo0106,
ppstacy,
ras44,
t-tte,
vincewu51 and
zhenyuz0500
July 26, 2026 06:21
10 tasks
Move NodeState / NodeSplitState / CausalRegressionCriterion (the per-group control-vs-treatment sufficient-statistics base shared by the causal- and uplift-tree criteria) out of causal/_criterion into a new shared kernel module _tree/_group_criterion. The uplift criterion now depends on the shared _tree kernel instead of the sibling causal package, making _tree the single home of the treatment-group base that both tree families extend. Pure relocation: the concrete causal criteria (StandardMSE / CausalMSE / TTest) and the uplift criteria are unchanged aside from cimport paths. Fitted trees are byte-identical (verified across causal MSE/standard/t-test x depth-/best-first and uplift KL/ED/Chi/CTS/DDP/IT/CIT/IDDP x plain/regularized/normalized/honest). Part of #945.
The causal and uplift trees each carried a ~570-line fork of scikit-learn's DepthFirstTreeBuilder / BestFirstTreeBuilder that were ~76% identical. Move the common skeleton into a new shared kernel module _tree/_group_builder (GroupDepthFirstTreeBuilder / GroupBestFirstTreeBuilder) and express the family-specific differences as four overridable cdef hooks (all no-ops by default, so the base reproduces the stock traversal): - _before_build(tree) once, after splitter.init - _before_node_split(splitter, parent) after node_reset, before the split search - _after_node_value(splitter, node_id) after node_value - _node_is_leaf_extra(splitter) extra leaf predicate OR-ed into stopping The family builders collapse to thin subclasses: causal overrides _node_is_leaf_extra (per-group minimum-size stopping); uplift overrides the three threading hooks (parent-summary shrinkage). Their now-empty _builder.pxd files are removed (the subclasses are self-contained). Only the depth-first uplift builder is reachable from the public API (_KernelUpliftTreeClassifier fixes max_leaf_nodes=None); the best-first uplift subclass is kept for symmetry. The shared best-first retains sklearn's impurity <= EPSILON early-leaf, a no-op for uplift (node_impurity is 1.0). Behavior-preserving: fitted trees are byte-identical before/after across causal (causal_mse/standard_mse/t_test x depth-/best-first x group-size stopping) and uplift (KL/ED/Chi/CTS/DDP/IT/CIT/IDDP x plain/regularized/normalized/honest); test_causal_trees, test_uplift_trees and test_uplift_trees_kernel (incl. the legacy-vs-kernel parity tests) all pass. Part of #945.
jeongyoonlee
force-pushed
the
feature/tree-kernel-sharing
branch
from
July 26, 2026 07:23
c31eb60 to
f1f32ca
Compare
ras44
approved these changes
Jul 27, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Further modularize the tree kernel so the causal and uplift trees share more code, per the direction in #945. Two focused, behavior-preserving refactors:
1. Hoist the shared treatment-group criterion base into the
_treekernel.NodeState/NodeSplitState/CausalRegressionCriterion— the per-group (control-vs-treatment) sufficient-statistics base that both families' criteria extend — moves out ofcausal/_criterioninto a new shared module_tree/_group_criterion. The uplift criterion now depends on the shared_treekernel instead of reaching sideways into the siblingcausalpackage.2. Unify the two forked builder skeletons.
causal/_builderand_uplift/_builderwere each a ~570-line fork of scikit-learn'sDepthFirstTreeBuilder/BestFirstTreeBuilder(~76% identical). The common skeleton now lives once in_tree/_group_builderasGroupDepthFirstTreeBuilder/GroupBestFirstTreeBuilder, with the family differences expressed as four overridablecdefhooks (all no-ops by default):_node_is_leaf_extra_before_node_split/_after_node_value/_before_buildThe family builders collapse to thin subclasses (~90 lines each) overriding only their hooks. The two now-empty
_builder.pxdfiles are removed (the subclasses are self-contained).Why
The criterion base was already shared by inheritance (uplift's criterion subclasses the causal one), but it physically lived under
causal/, so uplift depended on a sibling rather than on the kernel. And the builder skeleton was copied verbatim in both families. This centralizes both in_tree/, the single home of the shared tree machinery — the layering the #945 epic was driving toward.Behavior preservation
Both changes are behavior-preserving. A parity harness fits 37 trees — causal (
causal_mse/standard_mse/t_test× depth-/best-first × group-size stopping) and uplift (KL/ED/Chi/CTS/DDP/IT/CIT/IDDP× plain / regularized / normalized / honest) — and asserts the fittedtree_arrays (feature, threshold, value, children, impurity, counts) are byte-identical before and after each tier. Result: 37/37 identical after Tier 1 and after Tier 2.Note: only the depth-first uplift builder is reachable from the public API (
_KernelUpliftTreeClassifierfixesmax_leaf_nodes=None); the best-first uplift subclass is kept for symmetry. The shared best-first retains sklearn'simpurity <= EPSILONearly-leaf — a no-op for uplift, whosenode_impurityis the constant1.0.Verification
tests/test_causal_trees.py,tests/test_uplift_trees.py,tests/test_uplift_trees_kernel.py(incl. the legacy-vs-kernel parity tests) — all pass.blackclean.Net effect
~1140 lines of forked builder collapse to one ~600-line shared skeleton plus two ~90-line subclasses; the criterion base is a pure relocation.
_tree/becomes the single source of the shared criterion base and builder traversal for both tree families.Part of #945.