Improve SensitivityMSM robustness and test coverage

causalml/metrics/sensitivity.py

@@ -255,19 +255,9 @@ def sensitivity_analysis(
 
         return summary_df
 
-    # Learners whose fit_predict(return_components=True) does NOT return
-    # potential-outcome regressions (mu0, mu1). X-learner returns two CATE
-    # estimates from its tau models instead (xlearner.py); R-learner has no
-    # outcome-regression decomposition at all. Both are rejected explicitly
-    # rather than silently misinterpreted.
-    _UNSUPPORTED_POTENTIAL_OUTCOME_LEARNERS = (
-        "BaseXLearner",
-        "BaseXRegressor",
-        "BaseXClassifier",
-        "BaseRLearner",
-        "BaseRRegressor",
-        "BaseRClassifier",
-    )
+    # Learner families whose fit_predict(return_components=True)
+    # exposes potential-outcome regressions (mu0_hat, mu1_hat).
+    # Unknown learner types are rejected rather than assumed compatible.
 
     def get_potential_outcome_predictions(self, X, p, treatment, y):
         """Return separate potential-outcome predictions mu1_hat, mu0_hat.
@@ -290,8 +280,13 @@ def get_potential_outcome_predictions(self, X, p, treatment, y):
         """
         learner = self.learner
         learner_name = type(learner).__name__
+        mro_names = {cls.__name__ for cls in type(learner).__mro__}
 
-        if learner_name in self._UNSUPPORTED_POTENTIAL_OUTCOME_LEARNERS:
+        if not (
+            "BaseSLearner" in mro_names
+            or "BaseTLearner" in mro_names
+            or "BaseDRLearner" in mro_names
+        ):
             raise NotImplementedError(
                 "SensitivityMSM does not support {} yet: it needs potential-"
                 "outcome regressions (mu0_hat, mu1_hat), which this learner's "
@@ -744,6 +739,11 @@ def _bounds_for_gamma(self, mu1_hat, mu0_hat, p, t, y, gamma):
             (tuple of float): (ate_lower, ate_upper)
         """
         p_lower, p_upper = msm_propensity_bounds(p, gamma)
+
+        eps = np.finfo(float).eps
+        p_lower = np.clip(p_lower, eps, 1.0 - eps)
+        p_upper = np.clip(p_upper, eps, 1.0 - eps)
+
         resid_t = y - mu1_hat
         resid_c = y - mu0_hat
 

tests/test_sensitivity.py

@@ -7,6 +7,7 @@
 from causalml.inference.meta import (
     BaseSLearner,
     BaseTLearner,
+    BaseDRLearner,
     XGBTRegressor,
     BaseXLearner,
     BaseRLearner,
@@ -28,7 +29,13 @@
     alignment_att,
 )
 
-from .const import TREATMENT_COL, SCORE_COL, OUTCOME_COL, NUM_FEATURES
+from .const import (
+    TREATMENT_COL,
+    SCORE_COL,
+    OUTCOME_COL,
+    NUM_FEATURES,
+    RANDOM_SEED,
+)
 
 
 @pytest.mark.parametrize(
@@ -261,7 +268,17 @@ def test_alignment_att():
     assert y.shape == adj.shape
 
 
-def test_SensitivityMSM():
+@pytest.mark.parametrize(
+    "learner",
+    [
+        BaseSLearner(LinearRegression()),
+        BaseTLearner(LinearRegression()),
+        BaseDRLearner(LinearRegression()),
+        XGBTRegressor(),
+    ],
+)
+def test_SensitivityMSM(learner):
+    np.random.seed(RANDOM_SEED)
     y, X, treatment, tau, b, e = synthetic_data(
         mode=1, n=100000, p=NUM_FEATURES, sigma=1.0
     )
@@ -271,7 +288,6 @@ def test_SensitivityMSM():
     df[OUTCOME_COL] = y
     df[SCORE_COL] = e
 
-    learner = BaseTLearner(LinearRegression())
     sens = SensitivityMSM(
         df=df,
         inference_features=INFERENCE_FEATURES,
@@ -298,6 +314,8 @@ def test_SensitivityMSM():
 
 
 def test_SensitivityMSM_unsupported_learner():
+    np.random.seed(RANDOM_SEED)
+
     y, X, treatment, tau, b, e = synthetic_data(
         mode=1, n=100000, p=NUM_FEATURES, sigma=1.0
     )