Native missing-value (NaN) support for kernel uplift/causal trees (prereq for #955) - #972
Merged
Merged
Conversation
jeongyoonlee
marked this pull request as ready for review
July 25, 2026 21:12
jeongyoonlee
requested review from
alexander-pv,
huigangchen,
paullo0106,
ppstacy,
ras44,
t-tte,
vincewu51 and
zhenyuz0500
July 25, 2026 21:12
…945) Prerequisite for the #955 switchover: mirror scikit-learn's tree missing-value design so the kernel uplift and causal trees/forests accept NaNs in X natively, rather than rejecting them. The shared CausalRegressionCriterion (inherited by every uplift and causal criterion) now implements scikit-learn's per-feature missing-value machinery: * NodeSplitState gains a per-group `missing` NodeState; init_missing accumulates the tail (sample_indices[end - n_missing:end]) into both the flat sum_missing (used by StandardMSE) and the per-group state, respecting the group-encoding NaN mask. * reset/reverse_reset fold the missing block into the child chosen by missing_go_to_left; update scans only the non-missing range; StandardMSE's children_impurity adds the missing squared-sum when routed left. These are no-ops when n_missing == 0, so clean-data behavior is unchanged. The Python tree layer follows sklearn 1.7: `_support_missing_values` and an `allow_nan` tag (dense best splitter only), fit validates with ensure_all_finite=False and threads `_compute_missing_values_in_feature_mask` into the builder. `_prepare_data` and the forests allow NaN through validation; non-NaN non-finite values (inf) are still rejected. Adds missing-value tests (acceptance, tag/support, all-NaN column/row, learned NaN routing, inf-rejection, forest) for both the kernel uplift and causal trees/forests. Full uplift-kernel and causal-tree suites pass (154).
jeongyoonlee
force-pushed
the
feature/955-native-missing-values
branch
from
July 25, 2026 21:24
f88624b to
3310316
Compare
This was referenced Jul 26, 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
Prerequisite for the #955 switchover (tree-unification epic #945). Adds native missing-value (NaN) support to the kernel-backed uplift and causal trees/forests, mirroring scikit-learn's tree implementation closely. Today these estimators reject
NaNinX; the switched-over publicUpliftTreeClassifier/UpliftRandomForestClassifierneed to accept it, so this lands the criterion-level work first, isolated from the compat/switchover mechanics.How (follows scikit-learn's design)
The shared
CausalRegressionCriterion(inherited by every uplift and causal criterion) now implements sklearn's per-feature missing-value machinery:NodeSplitStategains a per-groupmissingNodeState.init_missing(n_missing)accumulates the missing tail (sample_indices[end - n_missing:end], the sklearn contract) into both the flatsum_missing(used byStandardMSE) and the per-group state — still respecting the group-encoding NaN mask (isnan(y[i, k])= "sample not in group k").reset/reverse_resetfold the missing block into the child chosen bymissing_go_to_left;updatescans only the non-missing range (end - n_missing);StandardMSE.children_impurityadds the missing squared-sum when routed left. All are no-ops whenn_missing == 0, so clean-data behaviour is unchanged.sum_missingallocation (init_sum_missing) is inherited fromRegressionCriterionunchanged.The Python tree layer follows sklearn 1.7:
_support_missing_values+ anallow_nantag via__sklearn_tags__(dense best splitter only — random/sparse don't carry the machinery). The vendored base's_get_tags()["allow_nan"]is dead under sklearn 1.7; the causal/uplift layers now use__sklearn_tags__().input_tags.allow_nan.fitvalidates withensure_all_finite=Falseand threads_compute_missing_values_in_feature_mask(X)intobuilder.build(...)._prepare_dataand both forests allow NaN through validation. inf is still rejected. The forestallow_nantag is delegated automatically by sklearn'sBaseForest.__sklearn_tags__.Scope
_KernelUpliftTreeClassifier,_KernelUpliftRandomForestClassifier,CausalTreeRegressor,CausalRandomForestRegressor(the criterion is shared, so causal benefits for free — as discussed for the epic).Tests
New missing-value tests for both the uplift and causal trees/forests:
allow_nantag +_support_missing_values(dense True / sparse False),Full
tests/test_uplift_trees_kernel.py+tests/test_causal_trees.pypass (154);blackclean.Part of #945. Prerequisite for #955.