diff --git a/causalml/metrics/sensitivity.py b/causalml/metrics/sensitivity.py index 42f1e74e..6317dada 100644 --- a/causalml/metrics/sensitivity.py +++ b/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 diff --git a/tests/test_sensitivity.py b/tests/test_sensitivity.py index 9a611642..50ba7001 100644 --- a/tests/test_sensitivity.py +++ b/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 )