diff --git a/causalml/inference/tree/_uplift/upliftforest.py b/causalml/inference/tree/_uplift/upliftforest.py index 2782611f..e02ecb4d 100644 --- a/causalml/inference/tree/_uplift/upliftforest.py +++ b/causalml/inference/tree/_uplift/upliftforest.py @@ -23,6 +23,8 @@ class is switched over. from sklearn.ensemble._forest import ForestRegressor, MAX_INT from sklearn.utils import check_array, check_random_state +from causalml.inference.serialization import SerializableLearner + from .uplifttree import _KernelUpliftTreeClassifier @@ -71,7 +73,7 @@ def _parallel_build_tree(tree, X, treatment, y, control_name, sample_weight): return tree -class _KernelUpliftRandomForestClassifier(ForestRegressor): +class _KernelUpliftRandomForestClassifier(SerializableLearner, ForestRegressor): """A random forest of kernel-backed uplift trees. Bags ``_KernelUpliftTreeClassifier`` on sklearn's ``ForestRegressor`` @@ -79,6 +81,11 @@ class _KernelUpliftRandomForestClassifier(ForestRegressor): ``P(Y=1|T=t) - P(Y=1|control)``; ``predict(..., full_output=True)`` returns the full frame (per-group probabilities, recommended treatment, deltas, ``max_delta``). + + Inherits :class:`~causalml.inference.serialization.SerializableLearner` for + ``save`` / ``load`` and is a scikit-learn ``BaseEstimator`` via + ``ForestRegressor`` (``get_params`` / ``clone`` round-trip). Full + ``check_estimator`` does not apply -- ``fit`` takes ``(X, treatment, y)``. """ def __init__( diff --git a/causalml/inference/tree/_uplift/uplifttree.py b/causalml/inference/tree/_uplift/uplifttree.py index ee19ea70..afcf7408 100644 --- a/causalml/inference/tree/_uplift/uplifttree.py +++ b/causalml/inference/tree/_uplift/uplifttree.py @@ -21,6 +21,7 @@ from sklearn.utils.validation import check_is_fitted from causalml.inference.meta.utils import check_treatment_vector +from causalml.inference.serialization import SerializableLearner from ._tree import BaseUpliftDecisionTree from .._tree._tree import Tree, build_pruned_tree_from_mask @@ -56,8 +57,16 @@ def __init__( self.summary = summary -class _KernelUpliftTreeClassifier(BaseUpliftDecisionTree): - """A single uplift tree grown on the shared ``_tree`` Cython kernel.""" +class _KernelUpliftTreeClassifier(SerializableLearner, BaseUpliftDecisionTree): + """A single uplift tree grown on the shared ``_tree`` Cython kernel. + + Inherits :class:`~causalml.inference.serialization.SerializableLearner` for + ``save`` / ``load`` and is a scikit-learn ``BaseEstimator`` via + ``BaseDecisionTree`` (``get_params`` / ``clone`` round-trip). ``__init__`` + stores its arguments verbatim per the sklearn convention. Full + ``check_estimator`` is not applicable: the supervised-uplift ``fit`` takes + ``(X, treatment, y)``, so sklearn's ``fit(X, y)``-shaped checks cannot run. + """ def __init__( self, @@ -80,9 +89,6 @@ def __init__( self.min_samples_treatment = min_samples_treatment self.n_reg = n_reg self.normalization = normalization - # IDDP requires the honest approach (legacy uplift.pyx ~468-469). - if criterion == "IDDP" and not honesty: - honesty = True self.honesty = honesty self.estimation_sample_size = estimation_sample_size super().__init__( @@ -120,7 +126,11 @@ def fit( """ X_enc, y_2dim = self._prepare_data(X=X, treatment=treatment, y=y) - if not self.honesty: + # IDDP requires the honest approach (legacy uplift.pyx ~468-469). Resolved + # here rather than in __init__ so the constructor stores its arguments + # verbatim (sklearn get_params / clone round-trip). + honest = self.honesty or self.criterion == "IDDP" + if not honest: super().fit( X=X_enc, y=y_2dim, sample_weight=sample_weight, check_input=check_input ) diff --git a/tests/test_uplift_trees_kernel.py b/tests/test_uplift_trees_kernel.py index 2ba2b256..12fe5ee8 100644 --- a/tests/test_uplift_trees_kernel.py +++ b/tests/test_uplift_trees_kernel.py @@ -28,6 +28,7 @@ import pytest from joblib import parallel_backend from numpy.testing import assert_array_almost_equal +from sklearn.base import clone from sklearn.model_selection import train_test_split from causalml.dataset import make_uplift_classification @@ -678,9 +679,25 @@ def test_kernel_uplift_parity_iddp_normalization_effective(): def test_kernel_uplift_iddp_forces_honesty(): - """IDDP forces the honest approach on, even when honesty=False is requested.""" - kern = _KernelUpliftTreeClassifier(criterion="IDDP", honesty=False) - assert kern.honesty is True + """IDDP forces the honest approach at fit, without mutating stored honesty.""" + X, treatment, y = _make_binary_feature_data() + common = dict( + criterion="IDDP", + control_name=CONTROL_NAME, + max_depth=3, + min_samples_leaf=100, + random_state=RANDOM_SEED, + ) + forced = _KernelUpliftTreeClassifier(honesty=False, **common).fit(X, treatment, y) + explicit = _KernelUpliftTreeClassifier(honesty=True, **common).fit(X, treatment, y) + # honesty=False is forced onto the honest path in fit, matching honesty=True, + # while the constructor keeps the argument verbatim (sklearn clone/get_params). + assert forced.get_params()["honesty"] is False + assert_array_almost_equal( + forced.predict_proba_by_group(X), + explicit.predict_proba_by_group(X), + decimal=12, + ) def test_kernel_uplift_iddp_guard_rejects_multi_treatment(): @@ -1175,3 +1192,100 @@ def test_kernel_uplift_forest_feature_importances(): # Uplift split gains are not a monotone impurity decrease, so per-feature # importances can be signed; they just have to be finite and averaged. assert np.isfinite(fi).all() + + +# --------------------------------------------------------------------------- +# Serialization + BaseEstimator conformance (issue #954) +# --------------------------------------------------------------------------- +# The kernel classes inherit SerializableLearner (save/load) and are sklearn +# BaseEstimators (get_params / clone) via BaseDecisionTree / ForestRegressor. +# Full check_estimator does not apply: fit takes (X, treatment, y), so sklearn's +# fit(X, y)-shaped checks cannot run. These cover the achievable conformance. + + +def _make_serialization_estimator(kind, **overrides): + """Construct an unfitted kernel tree or forest with test-friendly defaults.""" + if kind == "tree": + params = dict( + criterion="KL", + control_name=CONTROL_NAME, + max_depth=3, + min_samples_leaf=100, + random_state=RANDOM_SEED, + ) + params.update(overrides) + return _KernelUpliftTreeClassifier(**params) + params = dict( + control_name=CONTROL_NAME, + n_estimators=5, + criterion="KL", + max_depth=3, + min_samples_leaf=100, + random_state=RANDOM_SEED, + ) + params.update(overrides) + return _KernelUpliftRandomForestClassifier(**params) + + +@pytest.mark.parametrize("kind", ["tree", "forest"]) +def test_kernel_uplift_save_load_round_trip(kind, tmp_path): + """save() then load() restores an estimator with identical predictions.""" + X, treatment, y = _make_binary_feature_data() + est = _make_serialization_estimator(kind).fit(X, treatment, y) + path = str(tmp_path / f"{kind}.causalml") + est.save(path) + + loaded = type(est).load(path) + assert type(loaded) is type(est) + assert_array_almost_equal( + np.asarray(loaded.predict(X)), np.asarray(est.predict(X)), decimal=12 + ) + + +@pytest.mark.parametrize("kind", ["tree", "forest"]) +def test_kernel_uplift_unfitted_save_raises(kind, tmp_path): + """Saving before fit is rejected.""" + est = _make_serialization_estimator(kind) + with pytest.raises(ValueError): + est.save(str(tmp_path / "unfitted.causalml")) + + +@pytest.mark.parametrize("kind", ["tree", "forest"]) +def test_kernel_uplift_get_params_clone_round_trip(kind): + """clone() reproduces the constructor params exactly (BaseEstimator).""" + est = _make_serialization_estimator(kind) + assert clone(est).get_params() == est.get_params() + + +def test_kernel_uplift_load_class_mismatch_raises(tmp_path): + """Loading a saved tree as the forest class is rejected.""" + X, treatment, y = _make_binary_feature_data() + tree = _make_serialization_estimator("tree").fit(X, treatment, y) + path = str(tmp_path / "tree.causalml") + tree.save(path) + with pytest.raises(ValueError): + _KernelUpliftRandomForestClassifier.load(path) + + +def test_kernel_uplift_init_stores_args_verbatim(): + """__init__ stores its arguments unchanged (sklearn get_params / clone contract). + + In particular ``criterion="IDDP"`` must not mutate the stored ``honesty`` -- + IDDP forces the honest approach in fit(), not in the constructor (see + ``test_kernel_uplift_iddp_forces_honesty``). + """ + args = dict( + criterion="IDDP", + control_name="ctrl", + max_depth=7, + min_samples_leaf=42, + min_samples_treatment=3, + n_reg=5, + normalization=False, + honesty=False, + estimation_sample_size=0.3, + random_state=123, + ) + params = _KernelUpliftTreeClassifier(**args).get_params() + for name, value in args.items(): + assert params[name] == value