Skip to content

Share tree kernel: hoist group criterion base + unify builders (#945) - #974

Open
jeongyoonlee wants to merge 2 commits into
masterfrom
feature/tree-kernel-sharing
Open

Share tree kernel: hoist group criterion base + unify builders (#945)#974
jeongyoonlee wants to merge 2 commits into
masterfrom
feature/tree-kernel-sharing

Conversation

@jeongyoonlee

Copy link
Copy Markdown
Collaborator

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 _tree kernel.
NodeState / NodeSplitState / CausalRegressionCriterion — the per-group (control-vs-treatment) sufficient-statistics base that both families' criteria extend — moves out of causal/_criterion into a new shared module _tree/_group_criterion. The uplift criterion now depends on the shared _tree kernel instead of reaching sideways into the sibling causal package.

2. Unify the two forked builder skeletons.
causal/_builder and _uplift/_builder were each a ~570-line fork of scikit-learn's DepthFirstTreeBuilder / BestFirstTreeBuilder (~76% identical). The common skeleton now lives once in _tree/_group_builder as GroupDepthFirstTreeBuilder / GroupBestFirstTreeBuilder, with the family differences expressed as four overridable cdef hooks (all no-ops by default):

hook causal uplift
_node_is_leaf_extra per-group min-size stopping
_before_node_split / _after_node_value / _before_build parent-summary shrinkage threading

The family builders collapse to thin subclasses (~90 lines each) overriding only their hooks. The two now-empty _builder.pxd files 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 fitted tree_ 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 (_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, whose node_impurity is the constant 1.0.

Verification

  • Parity harness: 37/37 byte-identical trees (Tier 1 and Tier 2).
  • 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.
  • black clean.

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.

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
jeongyoonlee force-pushed the feature/tree-kernel-sharing branch from c31eb60 to f1f32ca Compare July 26, 2026 07:23
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

refactoring Code refactoring

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants