Repository: shap/shap
Pull request #3273: Add experimental causalml support
Author: login: alexander-pv
Labels: name: enhancement
description: Indicates new feature requests
Created: 2023-09-23T18:54:28Z | Merged: 2025-03-23T20:16:04Z

## Overview

Hi!
This PR is reopened https://github.com/shap/shap/pull/2654

## Checklist

- [x] All [pre-commit checks](https://pre-commit.com/#install) pass.
- [x] Unit tests added (if fixing a bug or adding a new feature)



Diff:
diff --git a/pyproject.toml b/pyproject.toml
index 48c2e8176a..932bb2673a 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -89,6 +89,7 @@ test = [
   # Constraint to prevent the combination of tf<2.15 and numpy>=2.0.
   # See GH #3707, #3768, 3922
   "numpy<2.0",
+  "causalml",
 ]
 
 test_notebooks = [
@@ -171,6 +172,7 @@ module = [
   "torch.*",
   "tqdm.*",
   "xgboost.*",
+  "causalml.*",
 ]
 ignore_missing_imports = true
 
diff --git a/shap/explainers/_tree.py b/shap/explainers/_tree.py
index ebebf685b6..2c892e5277 100644
--- a/shap/explainers/_tree.py
+++ b/shap/explainers/_tree.py
@@ -24,6 +24,7 @@
     InvalidModelError,
 )
 from ..utils._legacy import DenseData
+from ..utils._warnings import ExperimentalWarning
 from ._explainer import Explainer
 from .other._ubjson import decode_ubjson_buffer
 
@@ -53,6 +54,41 @@
 }
 
 
+def _safe_check_tree_instance_experimental(tree_instance: Any) -> None:
+    """
+    This function checks if a tree instance has an experimental integration with shap TreeExplainer class.
+
+    To add experimental message support for your library add package name and its versions
+    verified to be used with shap to the 'experimental' dictionary below.
+
+    Parameters
+    ----------
+    tree_instance: object, tree instance from an external library
+    """
+    experimental = {
+        "causalml": "0.15.3",
+    }
+
+    safe_instance = None
+    if hasattr(tree_instance, "__class__"):
+        if hasattr(tree_instance.__class__, "__module__"):
+            safe_instance = tree_instance
+
+    if safe_instance:
+        library = safe_instance.__class__.__module__.split(".")[0]
+        if experimental.get(library):
+            warnings.warn(
+                f"You are using experimental integration with {library}. "
+                f"The {library} support is verified for the following versions: {experimental.get(library)}. "
+                f"As experimental functionality, this integration may be removed or significantly changed in future releases without following semantic versioning. Use in production systems at your own risk.",
+                ExperimentalWarning,
+            )
+    else:
+        warnings.warn(
+            f"Unable to check experimental integration status for {tree_instance} object", ExperimentalWarning
+        )
+
+
 def _check_xgboost_version(v: str):
     if version.parse(v) < version.parse("1.6"):  # pragma: no cover
         raise RuntimeError(f"SHAP requires XGBoost >= v1.6 , but found version {v}. Please upgrade XGBoost.")
@@ -233,6 +269,9 @@ def __init__(
                 "using shap.sample(data, 100) to create a smaller background data set."
             )
             warnings.warn(wmsg)
+
+        _safe_check_tree_instance_experimental(model)
+
         self.data_missing = None if self.data is None else pd.isna(self.data)
         self.feature_perturbation = feature_perturbation
         self.expected_value = None
@@ -912,6 +951,7 @@ def __init__(self, model, data=None, data_missing=None, model_output=None):
                 "sklearn.ensemble.RandomForestRegressor",
                 "sklearn.ensemble.forest.RandomForestRegressor",
                 "econml.grf._base_grf.BaseGRF",
+                "causalml.inference.tree.CausalRandomForestRegressor",
             ],
         ):
             assert hasattr(model, "estimators_"), "Model has no `estimators_`! Have you called `model.fit`?"
@@ -969,6 +1009,7 @@ def __init__(self, model, data=None, data_missing=None, model_output=None):
                 "sklearn.tree.DecisionTreeRegressor",
                 "sklearn.tree.tree.DecisionTreeRegressor",
                 "econml.grf._base_grftree.GRFTree",
+                "causalml.inference.tree.causal.causaltree.CausalTreeRegressor",
             ],
         ):
             self.internal_dtype = model.tree_.value.dtype.type
@@ -1656,7 +1697,14 @@ class SingleTree:
     def __init__(self, tree, normalize=False, scaling=1.0, data=None, data_missing=None):
         assert_import("cext")
 
-        if safe_isinstance(tree, ["sklearn.tree._tree.Tree", "econml.tree._tree.Tree"]):
+        if safe_isinstance(
+            tree,
+            [
+                "sklearn.tree._tree.Tree",
+                "econml.tree._tree.Tree",
+                "causalml.inference.tree._tree._tree.Tree",
+            ],
+        ):
             self.children_left = tree.children_left.astype(np.int32)
             self.children_right = tree.children_right.astype(np.int32)
             self.children_default = self.children_left
diff --git a/shap/utils/_warnings.py b/shap/utils/_warnings.py
new file mode 100644
index 0000000000..a384718427
--- /dev/null
+++ b/shap/utils/_warnings.py
@@ -0,0 +1,4 @@
+class ExperimentalWarning(Warning):
+    """Used to manage warning messages for any experimental integrations."""
+
+    pass
diff --git a/tests/explainers/conftest.py b/tests/explainers/conftest.py
index 4aef1c50a7..2037cf4b2f 100644
--- a/tests/explainers/conftest.py
+++ b/tests/explainers/conftest.py
@@ -36,3 +36,24 @@ def basic_translation_scenario():
     ]
 
     return model, tokenizer, data
+
+
+@pytest.fixture()
+def causalml_synth_data():
+    """
+    Generates synthetic data for causalml causal tree tests.
+
+    Unlike standard regression trees causal trees in causalml evaluate outcome conditioning on treatments
+    Thus, a causal tree estimates Y_hat|X,T=t, where t={0, 1,..., n}.
+    The simplest case is when T = {0, 1}, 0 - no treatment, 1 - some treatment.
+    """
+    dataset = pytest.importorskip("causalml.dataset")
+
+    data_mode = 1  # Basic synthetic data mode with a difficult nuisance components and an easy treatment effect
+    sigma = 0.1  # Synthetic standard deviation of the error term
+    n_observations = 100  # The number of samples to generate
+    n_features = 8  # X in (Y_hat|X, T=0, Y_hat|X, T=1)
+    n_outcomes = 2  # Treatment conditioned outcomes: (Y_hat|X,T=0, Y_hat|X,T=1)
+
+    data = dataset.synthetic_data(mode=data_mode, n=n_observations, p=n_features, sigma=sigma)
+    return data, n_outcomes
diff --git a/tests/explainers/test_tree.py b/tests/explainers/test_tree.py
index 84e0eab995..f9ac7fd159 100644
--- a/tests/explainers/test_tree.py
+++ b/tests/explainers/test_tree.py
@@ -13,6 +13,7 @@
 from sklearn.utils import check_array
 
 import shap
+from shap.explainers._explainer import Explanation
 from shap.explainers._tree import SingleTree
 from shap.utils._exceptions import InvalidModelError
 
@@ -1994,3 +1995,69 @@ def test_overflow_tree_path_dependent():
     clf.predict_proba(X)
     exp = shap.Explainer(clf, algorithm="tree", feature_perturbation="tree_path_dependent")
     exp(X)
+
+
+@pytest.mark.parametrize(
+    "n_estimators",
+    [
+        5,
+    ],
+)
+def test_check_consistent_outputs_for_causalml_causal_trees(causalml_synth_data, n_estimators, random_seed):
+    """
+    Causal trees predict individual treatment effect based on continuous outcomes Y|X,T
+    where T is the particular type of treatment. In the basic scenario we have T=0 and T=1.
+
+    Thus, causal tree terminal nodes separately contain multiple outcomes as conditioned sample means:
+        Y_hat|X,T=0 and Y_hat|X,T=1
+    in the same manner as sklearn DecisionTreeRegressor with multiple outputs: (n_samples, n_outputs).
+
+    However, unlike standard regression tree the final output of the predict() method in causal trees is
+    the individual treatment effect: Y_hat|X,T=1 - Y_hat|X,T=0 with an option of returning possible outcomes Y_hat|X,T
+
+    During research, it is important to analyze Y_hat|X,T=t, t={0,1,...t} aside from individual effects estimation.
+    That is why we should carefully track the shape of the following arrays along with other checks:
+        shap values:  (n_observations, n_features, n_outcomes)
+        base values:  (n_observations, n_outcomes) arrays
+    """
+    causalml = pytest.importorskip("causalml")
+
+    data, n_outcomes = causalml_synth_data
+    y, X, treatment, tau, b, e = data
+    n_observations, n_features = X.shape
+
+    ctree = causalml.inference.tree.CausalTreeRegressor(random_state=random_seed)
+    ctree.fit(X=X, treatment=treatment, y=y)
+    ctree_preds = ctree.predict(X)
+    ctree_explainer = shap.TreeExplainer(ctree)
+
+    cforest = causalml.inference.tree.CausalRandomForestRegressor(n_estimators=n_estimators, random_state=random_seed)
+    cforest.fit(X=X, treatment=treatment, y=y)
+    cforest_preds = cforest.predict(X)
+    cforest_explainer = shap.TreeExplainer(cforest)
+
+    for explainer, preds in zip([ctree_explainer, cforest_explainer], [ctree_preds, cforest_preds]):
+        explanation = explainer(X)
+        shap_values = explainer.shap_values(X)
+
+        assert isinstance(explanation, Explanation)
+        assert isinstance(explanation.data, np.ndarray)
+        assert isinstance(explanation.base_values, np.ndarray)
+        assert isinstance(explanation.values, np.ndarray)
+        assert isinstance(shap_values, np.ndarray)
+
+        # Explanation.values and the output of TreeExplainer.shap_values() are two ways to get shap values
+        np.testing.assert_allclose(explanation.values, shap_values)
+        np.testing.assert_allclose(explanation.data, X)
+
+        # Check Explanation class
+        assert explanation.data.shape == (n_observations, n_features)
+        assert explanation.base_values.shape == (n_observations, n_outcomes)
+        assert explanation.values.shape == (n_observations, n_features, n_outcomes)
+
+        # Check that shap values and base values can be collapsed into
+        # model prediction of individual treatment effects
+        y_outcomes = explanation.base_values + explanation.values.sum(axis=1)
+        individual_effects = y_outcomes[:, 1] - y_outcomes[:, 0]
+
+        np.testing.assert_allclose(preds, individual_effects, atol=1e-4)


Reviews:
user: login: CloseChoice
state: CHANGES_REQUESTED
author_association: COLLABORATOR
submitted_at: 2025-03-16T08:48:24Z
user: login: alexander-pv
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-22T23:13:55Z
user: login: CloseChoice
state: CHANGES_REQUESTED
author_association: COLLABORATOR
submitted_at: 2025-03-23T06:45:09Z
user: login: alexander-pv
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-23T11:58:13Z
user: login: alexander-pv
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-23T11:58:32Z
user: login: alexander-pv
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-23T12:02:58Z
user: login: CloseChoice
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-23T13:18:30Z
user: login: CloseChoice
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-23T13:21:55Z
user: login: alexander-pv
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-23T15:58:52Z
user: login: alexander-pv
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-23T16:00:16Z
user: login: alexander-pv
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-23T16:06:30Z
user: login: alexander-pv
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-23T16:13:26Z
user: login: CloseChoice
state: CHANGES_REQUESTED
author_association: COLLABORATOR
submitted_at: 2025-03-23T17:41:45Z
user: login: alexander-pv
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-23T17:49:27Z
user: login: alexander-pv
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-23T18:17:55Z
user: login: CloseChoice
body: Very nice work. Thanks for the effort and patience; this has been a long time coming.
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-03-23T19:01:19Z

Review comments:
diff_hunk: @@ -1641,6 +1642,64 @@ def test_lightgbm_call_explanation(self):
         assert (explanation.values == shap_values).all()
 
 
+class TestExplainerCausalML:
path: tests/explainers/test_tree.py
user: login: CloseChoice
body: we typically don't use classes in shap tests, would be great if you could refactor to use pytest.fixtures instead.

Also 
```python
        from causalml.dataset import synthetic_data
        from causalml.inference.tree import CausalTreeRegressor

        X, y, treatment = self.prepare_data(synthetic_data)
 ```       
could be handled in those. Pls also consider the usage of `pytest.importorskip` so that tests pass even if causalml is not installed.
created_at: 2025-03-16T08:48:01Z
author_association: COLLABORATOR
diff_hunk: @@ -1641,6 +1642,64 @@ def test_lightgbm_call_explanation(self):
         assert (explanation.values == shap_values).all()
 
 
+class TestExplainerCausalML:
path: tests/explainers/test_tree.py
user: login: alexander-pv
body: Hi, @CloseChoice :wave: , can you review my latest changes, please?
Thanks
created_at: 2025-03-22T23:13:55Z
author_association: CONTRIBUTOR
diff_hunk: @@ -912,6 +912,7 @@ def __init__(self, model, data=None, data_missing=None, model_output=None):
                 "sklearn.ensemble.RandomForestRegressor",
                 "sklearn.ensemble.forest.RandomForestRegressor",
                 "econml.grf._base_grf.BaseGRF",
+                "causalml.inference.tree.CausalRandomForestRegressor",
path: shap/explainers/_tree.py
user: login: CloseChoice
body: is it possible that we add an info that causalml support is experimental using the logging module?
created_at: 2025-03-23T06:41:13Z
author_association: COLLABORATOR
diff_hunk: @@ -1994,3 +1994,51 @@ def test_overflow_tree_path_dependent():
     clf.predict_proba(X)
     exp = shap.Explainer(clf, algorithm="tree", feature_perturbation="tree_path_dependent")
     exp(X)
+
+
+def test_causalml_causal_tree_explanation_output(causalml_synth_data, random_seed):
+    causalml = pytest.importorskip("causalml")
+    from causalml.inference.tree import CausalTreeRegressor
+
+    data, check_shape = causalml_synth_data
+    y, X, treatment, tau, b, e = data
+
+    ctree = CausalTreeRegressor(random_state=random_seed)
+    ctree.fit(X=X, treatment=treatment, y=y)
+    ctree_explainer = shap.TreeExplainer(ctree)
+    explanation = ctree_explainer(X)
+
+    shap_values_raw: list[np.ndarray] = ctree_explainer.shap_values(X)
path: tests/explainers/test_tree.py
user: login: CloseChoice
body: could you pls check the `__call__` interface as well?
created_at: 2025-03-23T06:42:05Z
author_association: COLLABORATOR
diff_hunk: @@ -1994,3 +1994,51 @@ def test_overflow_tree_path_dependent():
     clf.predict_proba(X)
     exp = shap.Explainer(clf, algorithm="tree", feature_perturbation="tree_path_dependent")
     exp(X)
+
+
+def test_causalml_causal_tree_explanation_output(causalml_synth_data, random_seed):
+    causalml = pytest.importorskip("causalml")
+    from causalml.inference.tree import CausalTreeRegressor
+
+    data, check_shape = causalml_synth_data
+    y, X, treatment, tau, b, e = data
+
+    ctree = CausalTreeRegressor(random_state=random_seed)
+    ctree.fit(X=X, treatment=treatment, y=y)
+    ctree_explainer = shap.TreeExplainer(ctree)
+    explanation = ctree_explainer(X)
+
+    shap_values_raw: list[np.ndarray] = ctree_explainer.shap_values(X)
+    shap_values: np.ndarray = np.stack(shap_values_raw, axis=-1)
path: tests/explainers/test_tree.py
user: login: CloseChoice
body: why is the output a list? For trees in general we moved to the shape
`[num_observations, ...features, n_classes]` for classifications and
`[num_observations, ...features, n_outputs]` for regression. Could you pls elaborate what each of these outputs mean? I couldn't find it at a quick glance at the causalml docs

created_at: 2025-03-23T06:45:02Z
author_association: COLLABORATOR
diff_hunk: @@ -912,6 +912,7 @@ def __init__(self, model, data=None, data_missing=None, model_output=None):
                 "sklearn.ensemble.RandomForestRegressor",
                 "sklearn.ensemble.forest.RandomForestRegressor",
                 "econml.grf._base_grf.BaseGRF",
+                "causalml.inference.tree.CausalRandomForestRegressor",
path: shap/explainers/_tree.py
user: login: alexander-pv
body: Definitely :+1: TreeExplainer class constructor seems to be the best place for it. Is it OK to use warnings module as in the other functions of this module?
created_at: 2025-03-23T11:58:13Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1994,3 +1994,51 @@ def test_overflow_tree_path_dependent():
     clf.predict_proba(X)
     exp = shap.Explainer(clf, algorithm="tree", feature_perturbation="tree_path_dependent")
     exp(X)
+
+
+def test_causalml_causal_tree_explanation_output(causalml_synth_data, random_seed):
+    causalml = pytest.importorskip("causalml")
+    from causalml.inference.tree import CausalTreeRegressor
+
+    data, check_shape = causalml_synth_data
+    y, X, treatment, tau, b, e = data
+
+    ctree = CausalTreeRegressor(random_state=random_seed)
+    ctree.fit(X=X, treatment=treatment, y=y)
+    ctree_explainer = shap.TreeExplainer(ctree)
+    explanation = ctree_explainer(X)
+
+    shap_values_raw: list[np.ndarray] = ctree_explainer.shap_values(X)
path: tests/explainers/test_tree.py
user: login: alexander-pv
body: Ok :+1: 
created_at: 2025-03-23T11:58:32Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1994,3 +1994,51 @@ def test_overflow_tree_path_dependent():
     clf.predict_proba(X)
     exp = shap.Explainer(clf, algorithm="tree", feature_perturbation="tree_path_dependent")
     exp(X)
+
+
+def test_causalml_causal_tree_explanation_output(causalml_synth_data, random_seed):
+    causalml = pytest.importorskip("causalml")
+    from causalml.inference.tree import CausalTreeRegressor
+
+    data, check_shape = causalml_synth_data
+    y, X, treatment, tau, b, e = data
+
+    ctree = CausalTreeRegressor(random_state=random_seed)
+    ctree.fit(X=X, treatment=treatment, y=y)
+    ctree_explainer = shap.TreeExplainer(ctree)
+    explanation = ctree_explainer(X)
+
+    shap_values_raw: list[np.ndarray] = ctree_explainer.shap_values(X)
+    shap_values: np.ndarray = np.stack(shap_values_raw, axis=-1)
path: tests/explainers/test_tree.py
user: login: alexander-pv
body: Thank you for the comment! Apologies, that's my mistake. It is in fact a numpy array. I'll fix it and add detailed explanation of the output.
created_at: 2025-03-23T12:02:58Z
author_association: CONTRIBUTOR
diff_hunk: @@ -912,6 +912,7 @@ def __init__(self, model, data=None, data_missing=None, model_output=None):
                 "sklearn.ensemble.RandomForestRegressor",
                 "sklearn.ensemble.forest.RandomForestRegressor",
                 "econml.grf._base_grf.BaseGRF",
+                "causalml.inference.tree.CausalRandomForestRegressor",
path: shap/explainers/_tree.py
user: login: CloseChoice
body: yes, that's fine
created_at: 2025-03-23T13:18:30Z
author_association: COLLABORATOR
diff_hunk: @@ -1994,3 +1994,51 @@ def test_overflow_tree_path_dependent():
     clf.predict_proba(X)
     exp = shap.Explainer(clf, algorithm="tree", feature_perturbation="tree_path_dependent")
     exp(X)
+
+
+def test_causalml_causal_tree_explanation_output(causalml_synth_data, random_seed):
+    causalml = pytest.importorskip("causalml")
+    from causalml.inference.tree import CausalTreeRegressor
+
+    data, check_shape = causalml_synth_data
+    y, X, treatment, tau, b, e = data
+
+    ctree = CausalTreeRegressor(random_state=random_seed)
+    ctree.fit(X=X, treatment=treatment, y=y)
+    ctree_explainer = shap.TreeExplainer(ctree)
+    explanation = ctree_explainer(X)
+
+    shap_values_raw: list[np.ndarray] = ctree_explainer.shap_values(X)
+    shap_values: np.ndarray = np.stack(shap_values_raw, axis=-1)
path: tests/explainers/test_tree.py
user: login: CloseChoice
body: First of all, I would never blame anyone here for any code contributed to shap (except for things that harm the project intentionally of course). Secondly we did the breaking change mid 2024, when this PR was already created and we just didn't check that due to capacity shortcomings and, I have to say, if anyone did a mistake here I would blame it on us, the maintainers. 

So long story short: thanks a lot for your contribution, we highly appreciate it.
created_at: 2025-03-23T13:21:55Z
author_association: COLLABORATOR
diff_hunk: @@ -912,6 +912,7 @@ def __init__(self, model, data=None, data_missing=None, model_output=None):
                 "sklearn.ensemble.RandomForestRegressor",
                 "sklearn.ensemble.forest.RandomForestRegressor",
                 "econml.grf._base_grf.BaseGRF",
+                "causalml.inference.tree.CausalRandomForestRegressor",
path: shap/explainers/_tree.py
user: login: alexander-pv
body: Added `_safe_check_tree_instance_experimental` for that purpose
created_at: 2025-03-23T15:58:52Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1994,3 +1994,51 @@ def test_overflow_tree_path_dependent():
     clf.predict_proba(X)
     exp = shap.Explainer(clf, algorithm="tree", feature_perturbation="tree_path_dependent")
     exp(X)
+
+
+def test_causalml_causal_tree_explanation_output(causalml_synth_data, random_seed):
+    causalml = pytest.importorskip("causalml")
+    from causalml.inference.tree import CausalTreeRegressor
+
+    data, check_shape = causalml_synth_data
+    y, X, treatment, tau, b, e = data
+
+    ctree = CausalTreeRegressor(random_state=random_seed)
+    ctree.fit(X=X, treatment=treatment, y=y)
+    ctree_explainer = shap.TreeExplainer(ctree)
+    explanation = ctree_explainer(X)
+
+    shap_values_raw: list[np.ndarray] = ctree_explainer.shap_values(X)
path: tests/explainers/test_tree.py
user: login: alexander-pv
body: Added `__call__` output check ( `Explanation` class) 
created_at: 2025-03-23T16:00:16Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1994,3 +1994,51 @@ def test_overflow_tree_path_dependent():
     clf.predict_proba(X)
     exp = shap.Explainer(clf, algorithm="tree", feature_perturbation="tree_path_dependent")
     exp(X)
+
+
+def test_causalml_causal_tree_explanation_output(causalml_synth_data, random_seed):
+    causalml = pytest.importorskip("causalml")
+    from causalml.inference.tree import CausalTreeRegressor
+
+    data, check_shape = causalml_synth_data
+    y, X, treatment, tau, b, e = data
+
+    ctree = CausalTreeRegressor(random_state=random_seed)
+    ctree.fit(X=X, treatment=treatment, y=y)
+    ctree_explainer = shap.TreeExplainer(ctree)
+    explanation = ctree_explainer(X)
+
+    shap_values_raw: list[np.ndarray] = ctree_explainer.shap_values(X)
+    shap_values: np.ndarray = np.stack(shap_values_raw, axis=-1)
path: tests/explainers/test_tree.py
user: login: alexander-pv
body: I think everything is correct now. Please, add comments if you notice something else during the review.
Thanks again for a quick feedback! :handshake: 
created_at: 2025-03-23T16:06:30Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,6 @@
+class ExperimentalWarning(Warning):
path: shap/utils/_warnings.py
user: login: alexander-pv
body: Decided to suggest custom warnings similar to existing `_exceptions.py`. Of course, it is not very consistent to add a new module in this commit, but this class can clearly be reused during the initial integration of new models and I feel blue placing it inside `_tree.py` :sweat_smile: .
Let me know if it is better to avoid it at this time.
created_at: 2025-03-23T16:13:26Z
author_association: CONTRIBUTOR
diff_hunk: @@ -53,6 +54,41 @@
 }
 
 
+def _safe_check_tree_instance_experimental(tree_instance: Any) -> None:
+    """
+    This function checks if a tree instance has an experimental integration with shap TreeExplainer class.
+
+    To add experimental message support for your library add package name and its versions
+    verified to be used with shap to the 'experimental' dictionary below.
+
+    Parameters
+    ----------
+    tree_instance: object, tree instance from an external library
+    """
+    experimental = {
+        "causalml": "0.15.3",
+    }
+
+    safe_instance = None
+    if hasattr(tree_instance, "__class__"):
+        if hasattr(tree_instance.__class__, "__module__"):
+            safe_instance = tree_instance
+
+    if safe_instance:
+        library = safe_instance.__class__.__module__.split(".")[0]
+        if experimental.get(library):
+            warnings.warn(
+                f"You are using experimental integration with {library}. "
+                f"The {library} support is verified for the following versions: {experimental.get(library)}.",
path: shap/explainers/_tree.py
user: login: CloseChoice
body: Can we make clear that we do not commit ourselves to supporting this long term
```suggestion
                f"The {library} support is verified for the following versions: {experimental.get(library)}. "
                 f"As experimental functionality, this integration may be removed or significantly changed in future releases without following semantic versioning. Use in production systems at your own risk.",
```
created_at: 2025-03-23T17:30:37Z
author_association: COLLABORATOR
diff_hunk: @@ -89,16 +125,16 @@ class TreeExplainer(Explainer):
     """
 
     def __init__(
-        self,
-        model,
-        data=None,
-        model_output="raw",
-        feature_perturbation="auto",
-        feature_names=None,
-        approximate=DEPRECATED_APPROX,
-        # FIXME: The `link` and `linearize_link` arguments are ignored. GH #3513
-        link=None,
-        linearize_link=None,
+            self,
path: shap/explainers/_tree.py
user: login: CloseChoice
body: is there a particular reason for the formatting changes? I would to limit changes to either formatting or adding functionality, and not mix both into one PR.
created_at: 2025-03-23T17:31:37Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,6 @@
+class ExperimentalWarning(Warning):
path: shap/utils/_warnings.py
user: login: CloseChoice
body: I think that's a good idea, let's go with it.
created_at: 2025-03-23T17:32:24Z
author_association: COLLABORATOR
diff_hunk: @@ -2000,45 +2001,96 @@ def test_causalml_causal_tree_explanation_output(causalml_synth_data, random_see
     causalml = pytest.importorskip("causalml")
     from causalml.inference.tree import CausalTreeRegressor
 
-    data, check_shape = causalml_synth_data
+    data, n_outcomes = causalml_synth_data
     y, X, treatment, tau, b, e = data
+    n_observations, n_features = X.shape
 
     ctree = CausalTreeRegressor(random_state=random_seed)
     ctree.fit(X=X, treatment=treatment, y=y)
     ctree_explainer = shap.TreeExplainer(ctree)
-    explanation = ctree_explainer(X)
 
-    shap_values_raw: list[np.ndarray] = ctree_explainer.shap_values(X)
-    shap_values: np.ndarray = np.stack(shap_values_raw, axis=-1)
+    explanation = ctree_explainer(X)
+    shap_values = ctree_explainer.shap_values(X)
 
-    assert shap_values.shape == check_shape
-    assert len(explanation.base_values) == len(y)
+    assert isinstance(explanation.data, np.ndarray)
+    assert isinstance(explanation.base_values, np.ndarray)
     assert isinstance(explanation.values, np.ndarray)
     assert isinstance(shap_values, np.ndarray)
 
+    # Explanation.values attribute and the output of TreeExplainer.shap_values() method are two ways to get shap values
+    np.testing.assert_allclose(explanation.values, shap_values)
+
+    assert explanation.data.shape == (n_observations, n_features)
+    assert explanation.base_values.shape == (n_observations, n_outcomes)
+    assert explanation.values.shape == (n_observations, n_features, n_outcomes)
+
 
 @pytest.mark.parametrize(
     "n_estimators",
     [
         5,
     ],
 )
-def test_causalml_causal_random_forest_explanation_output(causalml_synth_data, n_estimators, random_seed):
+def test_check_consistent_outputs_for_causalml_causal_trees(causalml_synth_data, n_estimators, random_seed):
+    """
+    Causal trees predict individual treatment effect based on continuous outcomes Y|X,T
+    where T is the particular type of treatment. In the basic scenario we have T=0 and T=1.
+
+    Thus, causal tree terminal nodes separately contain multiple outcomes as conditioned sample means:
+        Y_hat|X,T=0 and Y_hat|X,T=1
+    in the same manner as sklearn DecisionTreeRegressor with multiple outputs: (n_samples, n_outputs).
+
+    However, unlike standard regression tree the final output of the predict() method in causal trees is
+    the individual treatment effect: Y_hat|X,T=1 - Y_hat|X,T=0 with an option of returning possible outcomes Y_hat|X,T
+
+    During research, it is important to analyze Y_hat|X,T=t, t={0,1,...t} aside from individual effects estimation.
+    That is why we should carefully track the shape of the following arrays along with other checks:
+        shap values:  (n_observations, n_features, n_outcomes)
+        base values:  (n_observations, n_outcomes) arrays
+    """
     causalml = pytest.importorskip("causalml")
+    from causalml.inference.tree import CausalTreeRegressor
     from causalml.inference.tree import CausalRandomForestRegressor
 
-    data, check_shape = causalml_synth_data
+    data, n_outcomes = causalml_synth_data
     y, X, treatment, tau, b, e = data
+    n_observations, n_features = X.shape
+
+    ctree = CausalTreeRegressor(random_state=random_seed)
+    ctree.fit(X=X, treatment=treatment, y=y)
+    ctree_preds = ctree.predict(X)
+    ctree_explainer = shap.TreeExplainer(ctree)
 
     cforest = CausalRandomForestRegressor(n_estimators=n_estimators, random_state=random_seed)
     cforest.fit(X=X, treatment=treatment, y=y)
+    cforest_preds = cforest.predict(X)
     cforest_explainer = shap.TreeExplainer(cforest)
-    explanation = cforest_explainer(X)
 
-    shap_values_raw: list[np.ndarray] = cforest_explainer.shap_values(X)
-    shap_values: np.ndarray = np.stack(shap_values_raw, axis=-1)
+    for explainer, preds in zip(
+            [ctree_explainer, cforest_explainer],
+            [ctree_preds, cforest_preds]
+    ):
+        explanation = explainer(X)
+        shap_values = explainer.shap_values(X)
 
-    assert shap_values.shape == check_shape
-    assert len(explanation.base_values) == len(y)
-    assert isinstance(explanation.values, np.ndarray)
-    assert isinstance(shap_values, np.ndarray)
+        assert isinstance(explanation, Explanation)
+        assert isinstance(explanation.data, np.ndarray)
+        assert isinstance(explanation.base_values, np.ndarray)
+        assert isinstance(explanation.values, np.ndarray)
+        assert isinstance(shap_values, np.ndarray)
+
+        # Explanation.values and the output of TreeExplainer.shap_values() are two ways to get shap values
+        np.testing.assert_allclose(explanation.values, shap_values)
+        np.testing.assert_allclose(explanation.data, X)
+
+        # Check Explanation class
+        assert explanation.data.shape == (n_observations, n_features)
+        assert explanation.base_values.shape == (n_observations, n_outcomes)
+        assert explanation.values.shape == (n_observations, n_features, n_outcomes)
+
+        # Check that shap values and base values can be collapsed into
+        # model prediction of individual treatment effects
+        y_outcomes = explanation.base_values + explanation.values.sum(axis=1)
+        individual_effects = y_outcomes[:, 1] - y_outcomes[:, 0]
path: tests/explainers/test_tree.py
user: login: CloseChoice
body: If I understand it correctly the difference between the treatment effect and the non-treatment effect is the uplift, correct? This is exactly what we're calculating here, therefore it must match `preds`
created_at: 2025-03-23T17:41:36Z
author_association: COLLABORATOR
diff_hunk: @@ -2000,45 +2001,96 @@ def test_causalml_causal_tree_explanation_output(causalml_synth_data, random_see
     causalml = pytest.importorskip("causalml")
     from causalml.inference.tree import CausalTreeRegressor
 
-    data, check_shape = causalml_synth_data
+    data, n_outcomes = causalml_synth_data
     y, X, treatment, tau, b, e = data
+    n_observations, n_features = X.shape
 
     ctree = CausalTreeRegressor(random_state=random_seed)
     ctree.fit(X=X, treatment=treatment, y=y)
     ctree_explainer = shap.TreeExplainer(ctree)
-    explanation = ctree_explainer(X)
 
-    shap_values_raw: list[np.ndarray] = ctree_explainer.shap_values(X)
-    shap_values: np.ndarray = np.stack(shap_values_raw, axis=-1)
+    explanation = ctree_explainer(X)
+    shap_values = ctree_explainer.shap_values(X)
 
-    assert shap_values.shape == check_shape
-    assert len(explanation.base_values) == len(y)
+    assert isinstance(explanation.data, np.ndarray)
+    assert isinstance(explanation.base_values, np.ndarray)
     assert isinstance(explanation.values, np.ndarray)
     assert isinstance(shap_values, np.ndarray)
 
+    # Explanation.values attribute and the output of TreeExplainer.shap_values() method are two ways to get shap values
+    np.testing.assert_allclose(explanation.values, shap_values)
+
+    assert explanation.data.shape == (n_observations, n_features)
+    assert explanation.base_values.shape == (n_observations, n_outcomes)
+    assert explanation.values.shape == (n_observations, n_features, n_outcomes)
+
 
 @pytest.mark.parametrize(
     "n_estimators",
     [
         5,
     ],
 )
-def test_causalml_causal_random_forest_explanation_output(causalml_synth_data, n_estimators, random_seed):
+def test_check_consistent_outputs_for_causalml_causal_trees(causalml_synth_data, n_estimators, random_seed):
+    """
+    Causal trees predict individual treatment effect based on continuous outcomes Y|X,T
+    where T is the particular type of treatment. In the basic scenario we have T=0 and T=1.
+
+    Thus, causal tree terminal nodes separately contain multiple outcomes as conditioned sample means:
+        Y_hat|X,T=0 and Y_hat|X,T=1
+    in the same manner as sklearn DecisionTreeRegressor with multiple outputs: (n_samples, n_outputs).
+
+    However, unlike standard regression tree the final output of the predict() method in causal trees is
+    the individual treatment effect: Y_hat|X,T=1 - Y_hat|X,T=0 with an option of returning possible outcomes Y_hat|X,T
+
+    During research, it is important to analyze Y_hat|X,T=t, t={0,1,...t} aside from individual effects estimation.
+    That is why we should carefully track the shape of the following arrays along with other checks:
+        shap values:  (n_observations, n_features, n_outcomes)
+        base values:  (n_observations, n_outcomes) arrays
+    """
     causalml = pytest.importorskip("causalml")
+    from causalml.inference.tree import CausalTreeRegressor
     from causalml.inference.tree import CausalRandomForestRegressor
 
-    data, check_shape = causalml_synth_data
+    data, n_outcomes = causalml_synth_data
     y, X, treatment, tau, b, e = data
+    n_observations, n_features = X.shape
+
+    ctree = CausalTreeRegressor(random_state=random_seed)
+    ctree.fit(X=X, treatment=treatment, y=y)
+    ctree_preds = ctree.predict(X)
+    ctree_explainer = shap.TreeExplainer(ctree)
 
     cforest = CausalRandomForestRegressor(n_estimators=n_estimators, random_state=random_seed)
     cforest.fit(X=X, treatment=treatment, y=y)
+    cforest_preds = cforest.predict(X)
     cforest_explainer = shap.TreeExplainer(cforest)
-    explanation = cforest_explainer(X)
 
-    shap_values_raw: list[np.ndarray] = cforest_explainer.shap_values(X)
-    shap_values: np.ndarray = np.stack(shap_values_raw, axis=-1)
+    for explainer, preds in zip(
+            [ctree_explainer, cforest_explainer],
+            [ctree_preds, cforest_preds]
+    ):
+        explanation = explainer(X)
+        shap_values = explainer.shap_values(X)
 
-    assert shap_values.shape == check_shape
-    assert len(explanation.base_values) == len(y)
-    assert isinstance(explanation.values, np.ndarray)
-    assert isinstance(shap_values, np.ndarray)
+        assert isinstance(explanation, Explanation)
+        assert isinstance(explanation.data, np.ndarray)
+        assert isinstance(explanation.base_values, np.ndarray)
+        assert isinstance(explanation.values, np.ndarray)
+        assert isinstance(shap_values, np.ndarray)
+
+        # Explanation.values and the output of TreeExplainer.shap_values() are two ways to get shap values
+        np.testing.assert_allclose(explanation.values, shap_values)
+        np.testing.assert_allclose(explanation.data, X)
+
+        # Check Explanation class
+        assert explanation.data.shape == (n_observations, n_features)
+        assert explanation.base_values.shape == (n_observations, n_outcomes)
+        assert explanation.values.shape == (n_observations, n_features, n_outcomes)
+
+        # Check that shap values and base values can be collapsed into
+        # model prediction of individual treatment effects
+        y_outcomes = explanation.base_values + explanation.values.sum(axis=1)
+        individual_effects = y_outcomes[:, 1] - y_outcomes[:, 0]
path: tests/explainers/test_tree.py
user: login: alexander-pv
body: Yes, that's correct
created_at: 2025-03-23T17:49:27Z
author_association: CONTRIBUTOR
diff_hunk: @@ -89,16 +125,16 @@ class TreeExplainer(Explainer):
     """
 
     def __init__(
-        self,
-        model,
-        data=None,
-        model_output="raw",
-        feature_perturbation="auto",
-        feature_names=None,
-        approximate=DEPRECATED_APPROX,
-        # FIXME: The `link` and `linearize_link` arguments are ignored. GH #3513
-        link=None,
-        linearize_link=None,
+            self,
path: shap/explainers/_tree.py
user: login: alexander-pv
body: Hmm, I didn't change formatting although it is showed in the merge commit https://github.com/shap/shap/pull/3273/commits/cc326cc3ecb9d655443af66ca69c0b1ab5701066. Anyway the latest diff shows that nothing changed.
created_at: 2025-03-23T18:17:55Z
author_association: CONTRIBUTOR

PR comments:
user: login: connortann
created_at: 2023-09-26T13:54:39Z
body: The tests are failing due to an unrelated HuggingFace error: apologies. We'll get that sorted asap (on #3286), and then we can get this PR passing.

Edit: done.
author_association: COLLABORATOR
user: login: codecov[bot]
created_at: 2023-09-27T00:05:39Z
body: ## [Codecov](https://app.codecov.io/gh/shap/shap/pull/3273?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=shap) Report
Attention: Patch coverage is `93.75000%` with `1 line` in your changes missing coverage. Please review.
> Project coverage is 64.81%. Comparing base [(`d5ce4c7`)](https://app.codecov.io/gh/shap/shap/commit/d5ce4c7ccfa02b59879c2fb29004c7c1cfb0d793?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=shap) to head [(`d06ff25`)](https://app.codecov.io/gh/shap/shap/commit/d06ff253c7b79e22cd500701815994955e97712f?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=shap).
> Report is 29 commits behind head on master.

| [Files with missing lines](https://app.codecov.io/gh/shap/shap/pull/3273?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=shap) | Patch % | Lines |
|---|---|---|
| [shap/explainers/\_tree.py](https://app.codecov.io/gh/shap/shap/pull/3273?src=pr&el=tree&filepath=shap%2Fexplainers%2F_tree.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=shap#diff-c2hhcC9leHBsYWluZXJzL190cmVlLnB5) | 92.85% | [1 Missing :warning: ](https://app.codecov.io/gh/shap/shap/pull/3273?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=shap) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master    #3273      +/-   ##
==========================================
+ Coverage   64.78%   64.81%   +0.03%     
==========================================
  Files          92       93       +1     
  Lines       12895    12910      +15     
==========================================
+ Hits         8354     8368      +14     
- Misses       4541     4542       +1     
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/shap/shap/pull/3273?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=shap).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=shap).

<details><summary> :rocket: New features to boost your workflow: </summary>

- :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
- :package: [JS Bundle Analysis](https://docs.codecov.com/docs/javascript-bundle-analysis): Save yourself from yourself by tracking and limiting bundle sizes in JS merges.
</details>
author_association: NONE
user: login: connortann
created_at: 2023-09-27T01:49:37Z
body: This PR has quite a large diff, as there are many unrelated style changes alongside the functional changes. Please would you remove the style changes, leaving just the functional changes?

We plan to introduce python black soon to remove this kind of hassle for contributors.

Nb. I inadvertently briefly closed this PR (currently on mobile), apologies.
author_association: COLLABORATOR
user: login: alexander-pv
created_at: 2023-09-28T21:35:48Z
body: Apologies for style change. I will remove it.

Python black usage would be great!
author_association: CONTRIBUTOR
user: login: thatlittleboy
created_at: 2023-10-01T08:24:08Z
body: Hi @alexander-pv , do you think you could also add in some tests to `test_tree.py` (just a simple function or two to demonstrate that the proposed changes here work as expected?)

Thanks!
author_association: COLLABORATOR
user: login: connortann
created_at: 2023-12-04T11:47:37Z
body: I'll add the "stale" label for now as there hasn't been recent activity. Please feel free to comment if you're still interested in the PR and are happy to work on adding a minimal unit test.
author_association: COLLABORATOR
user: login: alexander-pv
created_at: 2023-12-08T23:34:54Z
body: Hi, @connortann!  I added several test to check if shap works as intended with causalml trees.
author_association: CONTRIBUTOR
user: login: alexander-pv
created_at: 2023-12-08T23:47:58Z
body: About current scikit-learn building issue in recent tests: https://github.com/scikit-learn/scikit-learn/issues/26858
author_association: CONTRIBUTOR
user: login: connortann
created_at: 2023-12-21T13:29:41Z
body: Apologies that we haven't had time to look further at this PR, as many maintainers are away over the Christmas period. Thanks again for the PR - I hope to get to this early in the new year. 
author_association: COLLABORATOR
user: login: alexander-pv
created_at: 2024-04-21T11:25:42Z
body: Hi, @connortann! I updated the PR resolving a couple of conflicts. Tests work as expected on my computer.
So, awaiting feedback from label can be removed now. I will be glad to contribute if you think anything else is needed for this PR!

Thanks!
author_association: CONTRIBUTOR
user: login: jeongyoonlee
created_at: 2024-05-10T17:16:57Z
body: Hi, @connortann. Can you take a look at this PR? This has been a long-standing issue (e.g., uber/causalml#265, uber/causalml#527, uber/causalml#615, uber/causalml#727, uber/causalml#735) for the `causalml` package, and we hope to resolve it soon.

Your support will be greatly appreciated.
author_association: NONE
user: login: jeongyoonlee
created_at: 2024-09-13T16:05:55Z
body: @alexander-pv, can you look into t[he test failures](https://github.com/shap/shap/actions/runs/10618029654/job/29432295181?pr=3273#step:6:797) in < Python3.12? Weirdly, it passes in Python 3.12.
author_association: NONE
user: login: CloseChoice
created_at: 2025-03-16T08:49:53Z
body: Thanks for the PR and sorry that this took so, so long. Will try to respond in a timely manner, pls feel free to ping me whenever you have questions or want another review.
author_association: COLLABORATOR
user: login: alexander-pv
created_at: 2025-03-16T16:13:51Z
body: > Thanks for the PR and sorry that this took so, so long. Will try to respond in a timely manner, pls feel free to ping me whenever you have questions or want another review.

Ok, thanks! 
author_association: CONTRIBUTOR
user: login: CloseChoice
created_at: 2025-03-23T17:42:44Z
body: @alexander-pv I think we're almost through. Could you pls just remove the formatting changes (or explain why we need them) and add the comment regarding future support in the warning and we're good to merge.
author_association: COLLABORATOR
user: login: alexander-pv
created_at: 2025-03-23T18:19:25Z
body: > @alexander-pv I think we're almost through. Could you pls just remove the formatting changes (or explain why we need them) and add the comment regarding future support in the warning and we're good to merge.

Done. Answered in the comments above.
author_association: CONTRIBUTOR
user: login: CloseChoice
created_at: 2025-03-23T19:09:37Z
body: @alexander-pv : I created a branch based on yours and fixed the pre-commit issues: https://github.com/shap/shap/tree/feature/causalml_trees. Pls merge that one in and then we'll merge once the tests pass
author_association: COLLABORATOR
user: login: alexander-pv
created_at: 2025-03-23T19:49:46Z
body: > @alexander-pv : I created a branch based on yours and fixed the pre-commit issues: https://github.com/shap/shap/tree/feature/causalml_trees. Pls merge that one in and then we'll merge once the tests pass

Merged :ok_hand: 
author_association: CONTRIBUTOR