Add propensity clipping diagnostics to DRTester

econml/tests/test_validate_utils.py

@@ -0,0 +1,159 @@
+import unittest
+import warnings
+from unittest.mock import Mock
+
+import numpy as np
+
+from econml.validate.drtester import DRTester
+from econml.validate.utils import _calculate_dr_outcomes_with_diagnostics, calculate_dr_outcomes
+
+
+class TestValidateUtils(unittest.TestCase):
+    def test_dr_outcome_clipping_configuration_and_warning(self):
+        D = np.array([0, 1, 0, 1])
+        y = np.array([1.0, 2.0, 3.0, 4.0])
+        reg_preds = np.zeros((4, 2))
+        prop_preds = np.array(
+            [
+                [0.005, 0.995],
+                [0.995, 0.005],
+                [0.5, 0.5],
+                [0.5, 0.5],
+            ]
+        )
+
+        with self.assertWarnsRegex(UserWarning, r"2 values.*0: 1, 1: 1"):
+            actual = calculate_dr_outcomes(D, y, reg_preds, prop_preds)
+
+        y_dr_0 = (D == 0) / np.clip(prop_preds[:, 0], 0.01, np.inf) * y
+        y_dr_1 = (D == 1) / np.clip(prop_preds[:, 1], 0.01, np.inf) * y
+        np.testing.assert_array_equal(actual[:, 0], y_dr_1 - y_dr_0)
+
+        with warnings.catch_warnings():
+            warnings.simplefilter("error")
+            custom = calculate_dr_outcomes(D, y, reg_preds, prop_preds, min_propensity=0.1, warn_on_clip=False)
+        custom_y_dr_0 = (D == 0) / np.clip(prop_preds[:, 0], 0.1, np.inf) * y
+        custom_y_dr_1 = (D == 1) / np.clip(prop_preds[:, 1], 0.1, np.inf) * y
+        np.testing.assert_array_equal(custom[:, 0], custom_y_dr_1 - custom_y_dr_0)
+
+    def test_dr_outcome_default_is_unchanged_without_clipping(self):
+        D = np.array([0, 1, 2])
+        y = np.array([1.0, 2.0, 3.0])
+        reg_preds = np.arange(9, dtype=float).reshape(3, 3) / 10
+        prop_preds = np.array(
+            [
+                [0.5, 0.3, 0.2],
+                [0.2, 0.6, 0.2],
+                [0.3, 0.2, 0.5],
+            ]
+        )
+
+        d0_mask = np.where(D == 0, 1, 0)
+        y_dr_0 = reg_preds[:, 0] + (d0_mask / np.clip(prop_preds[:, 0], 0.01, np.inf)) * (y - reg_preds[:, 0])
+        expected = []
+        for k in (1, 2):
+            dk_mask = np.where(D == k, 1, 0)
+            y_dr_k = reg_preds[:, k] + (dk_mask / np.clip(prop_preds[:, k], 0.01, np.inf)) * (y - reg_preds[:, k])
+            expected.append(y_dr_k - y_dr_0)
+
+        with warnings.catch_warnings():
+            warnings.simplefilter("error")
+            actual = calculate_dr_outcomes(D, y, reg_preds, prop_preds)
+        np.testing.assert_array_equal(actual, np.column_stack(expected))
+
+    def test_unassigned_low_propensities_do_not_warn(self):
+        D = np.array([0, 0, 1, 1])
+        y = np.ones(4)
+        reg_preds = np.zeros((4, 2))
+        prop_preds = np.array(
+            [
+                [0.9, 0.001],
+                [0.9, 0.001],
+                [0.5, 0.5],
+                [0.5, 0.5],
+            ]
+        )
+
+        with warnings.catch_warnings():
+            warnings.simplefilter("error")
+            calculate_dr_outcomes(D, y, reg_preds, prop_preds)
+
+        _, diagnostics = _calculate_dr_outcomes_with_diagnostics(D, y, reg_preds, prop_preds)
+        self.assertEqual(diagnostics["below_threshold_by_treatment"], {0: 0, 1: 2})
+        self.assertEqual(diagnostics["n_below_threshold"], 2)
+        self.assertEqual(diagnostics["clipped_by_treatment"], {0: 0, 1: 0})
+        self.assertEqual(diagnostics["n_clipped"], 0)
+
+    def test_dr_outcome_input_validation(self):
+        D = np.array([0, 1])
+        y = np.array([1.0, 2.0])
+        reg_preds = np.zeros((2, 2))
+        prop_preds = np.full((2, 2), 0.5)
+
+        for threshold in (0, 0.5, np.nan, np.inf, True, "0.01"):
+            with self.subTest(threshold=threshold), self.assertRaisesRegex(ValueError, "min_propensity"):
+                calculate_dr_outcomes(D, y, reg_preds, prop_preds, min_propensity=threshold)
+        with self.assertRaisesRegex(ValueError, "same number of rows"):
+            calculate_dr_outcomes(D, y[:-1], reg_preds, prop_preds)
+        with self.assertRaisesRegex(ValueError, "same number of treatment columns"):
+            calculate_dr_outcomes(D, y, reg_preds[:, :1], prop_preds)
+        with self.assertRaisesRegex(ValueError, "finite values"):
+            calculate_dr_outcomes(D, y, reg_preds, np.array([[np.nan, 0.5], [0.5, 0.5]]))
+        with self.assertRaisesRegex(ValueError, "column for every treatment label"):
+            calculate_dr_outcomes(np.array([0, 2]), y, reg_preds, prop_preds)
+
+    def test_clip_diagnostics_distinguish_train_and_validation(self):
+        Dtrain = np.array([0, 1, 2])
+        Dval = np.array([0, 1])
+        ytrain = np.ones(3)
+        yval = np.ones(2)
+        train_preds = (
+            np.zeros((3, 3)),
+            np.array([[0.005, 0.5, 0.5], [0.5, 0.004, 0.5], [0.5, 0.5, 0.003]]),
+        )
+        validation_preds = (
+            np.zeros((2, 3)),
+            np.array([[0.002, 0.5, 0.5], [0.5, 0.5, 0.5]]),
+        )
+        tester = DRTester(model_regression=None, model_propensity=None, cate=None)
+        tester.fit_nuisance_cv = Mock(return_value=train_preds)
+        tester.fit_nuisance_train = Mock(return_value=validation_preds)
+
+        with self.assertWarnsRegex(UserWarning, r"train \(0: 1, 1: 1, 2: 1\); validation \(0: 1\)"):
+            tester.fit_nuisance(
+                np.zeros((2, 1)),
+                Dval,
+                yval,
+                np.zeros((3, 1)),
+                Dtrain,
+                ytrain,
+            )
+
+        self.assertEqual(set(tester.clip_diagnostics_), {"train", "validation"})
+        self.assertEqual(tester.clip_diagnostics_["train"]["treatment_labels"], (0, 1, 2))
+        self.assertEqual(tester.clip_diagnostics_["validation"]["treatment_labels"], (0, 1))
+        self.assertNotIn(2, tester.clip_diagnostics_["validation"]["clipped_by_treatment"])
+
+    def test_clip_diagnostics_are_available_when_warning_is_disabled(self):
+        Dval = np.array([0, 1])
+        tester = DRTester(
+            model_regression=None,
+            model_propensity=None,
+            cate=None,
+            min_propensity=0.1,
+            warn_on_clip=False,
+        )
+        tester.fit_nuisance_cv = Mock(
+            return_value=(
+                np.zeros((2, 2)),
+                np.array([[0.05, 0.95], [0.95, 0.05]]),
+            )
+        )
+
+        with warnings.catch_warnings():
+            warnings.simplefilter("error")
+            tester.fit_nuisance(np.zeros((2, 1)), Dval, np.ones(2))
+
+        self.assertEqual(set(tester.clip_diagnostics_), {"validation"})
+        self.assertEqual(tester.clip_diagnostics_["validation"]["n_clipped"], 2)
+        self.assertEqual(tester.clip_diagnostics_["validation"]["min_propensity"], 0.1)

econml/validate/drtester.py

@@ -1,3 +1,4 @@
+import warnings
 from typing import Tuple, Union, List
 
 import numpy as np
@@ -8,7 +9,7 @@
 from econml._lazy import _LazyModule
 from econml.utilities import check_input_arrays, deprecated, add_constant
 from .results import CalibrationEvaluationResults, BLPEvaluationResults, UpliftEvaluationResults, EvaluationResults
-from .utils import calculate_dr_outcomes, calc_uplift
+from .utils import _calculate_dr_outcomes_with_diagnostics, _format_clip_counts, calc_uplift
 
 _statsmodels_api = _LazyModule("statsmodels.api")  # lazy: only needed for evaluate_blp()
 
@@ -101,6 +102,24 @@ class DRTester:
         Splitter used for cross-validation. Can be either an integer (corresponding to the number of desired folds)
         or a list of indices corresponding to membership in each fold.
 
+    min_propensity: float, default 0.01
+        Lower bound applied independently to each treatment propensity used in a doubly robust denominator.
+
+    warn_on_clip: bool, default True
+        Whether ``fit_nuisance`` warns when propensities are clipped. Structured counts remain available in
+        ``clip_diagnostics_`` when warnings are disabled. Clipping is a numerical stabilization signal and does not
+        by itself prove that identification failed.
+
+    Attributes
+    ----------
+    clip_diagnostics_ : dict
+        Populated by ``fit_nuisance`` with a ``validation`` entry and, when training data are supplied, a ``train``
+        entry. Each split contains ``min_propensity`` (float), ``n_samples`` (int), ``treatment_labels`` (tuple of
+        labels observed in that split), ``below_threshold_by_treatment`` (counts across all rows),
+        ``n_below_threshold`` (their total), ``clipped_by_treatment`` (counts on rows assigned to the corresponding
+        treatment, where clipping changes a doubly robust denominator), and ``n_clipped`` (their total). Treatment
+        keys can differ between splits when an arm is absent from one split.
+
     References
     ----------
     [Chernozhukov2022] V. Chernozhukov et al.
@@ -124,12 +143,16 @@ def __init__(
         model_regression,
         model_propensity,
         cate,
-        cv: Union[int, List] = 5
+        cv: Union[int, List] = 5,
+        min_propensity: float = 0.01,
+        warn_on_clip: bool = True,
     ):
         self.model_regression = model_regression
         self.model_propensity = model_propensity
         self.cate = cate
         self.cv = cv
+        self.min_propensity = min_propensity
+        self.warn_on_clip = warn_on_clip
 
     def get_cv_splitter(self, random_state: int = 123):
         """
@@ -237,15 +260,37 @@ def fit_nuisance(
         if self.fit_on_train:
             # Get DR outcomes in training sample
             reg_preds_train, prop_preds_train = self.fit_nuisance_cv(Xtrain, Dtrain, ytrain)
-            self.dr_train_ = calculate_dr_outcomes(Dtrain, ytrain, reg_preds_train, prop_preds_train)
+            self.dr_train_, train_diagnostics = _calculate_dr_outcomes_with_diagnostics(
+                Dtrain, ytrain, reg_preds_train, prop_preds_train, min_propensity=self.min_propensity
+            )
 
             # Get DR outcomes in validation sample
             reg_preds_val, prop_preds_val = self.fit_nuisance_train(Xtrain, Dtrain, ytrain, Xval)
-            self.dr_val_ = calculate_dr_outcomes(Dval, yval, reg_preds_val, prop_preds_val)
+            self.dr_val_, validation_diagnostics = _calculate_dr_outcomes_with_diagnostics(
+                Dval, yval, reg_preds_val, prop_preds_val, min_propensity=self.min_propensity
+            )
+            self.clip_diagnostics_ = {"train": train_diagnostics, "validation": validation_diagnostics}
         else:
             # Get DR outcomes in validation sample
             reg_preds_val, prop_preds_val = self.fit_nuisance_cv(Xval, Dval, yval)
-            self.dr_val_ = calculate_dr_outcomes(Dval, yval, reg_preds_val, prop_preds_val)
+            self.dr_val_, validation_diagnostics = _calculate_dr_outcomes_with_diagnostics(
+                Dval, yval, reg_preds_val, prop_preds_val, min_propensity=self.min_propensity
+            )
+            self.clip_diagnostics_ = {"validation": validation_diagnostics}
+
+        clipped_splits = [
+            f"{split} ({_format_clip_counts(diagnostics)})"
+            for split, diagnostics in self.clip_diagnostics_.items()
+            if diagnostics["n_clipped"]
+        ]
+        if self.warn_on_clip and clipped_splits:
+            warnings.warn(
+                f"Propensity scores were clipped below min_propensity={float(self.min_propensity):g} in "
+                f"{'; '.join(clipped_splits)}. Clipping is a numerical stabilization signal and does not by itself "
+                "prove that identification failed.",
+                UserWarning,
+                stacklevel=2,
+            )
 
         # Calculate ATE in the validation sample
         self.ate_val = self.dr_val_.mean(axis=0)