757: add return_ci=True in sensitivity (#758)

* 757: add return_ci=True in sensitivity * remove if-else logic - resolved in #295 * black
uber · Apr 12, 2024 · 12747ee · 12747ee
1 parent d4a19a2
commit 12747ee
Showing 1 changed file with 7 additions and 12 deletions.
diff --git a/causalml/metrics/sensitivity.py b/causalml/metrics/sensitivity.py
@@ -147,19 +147,14 @@ def get_ate_ci(self, X, p, treatment, y):
         learner = self.learner
         from ..inference.meta.tlearner import BaseTLearner
 
-        if isinstance(learner, BaseTLearner):
-            ate, ate_lower, ate_upper = learner.estimate_ate(
-                X=X, treatment=treatment, y=y
+        try:
+            ate, ate_lower, ate_upper = self.learner.estimate_ate(
+                X=X, p=p, treatment=treatment, y=y, return_ci=True
+            )
+        except TypeError:
+            ate, ate_lower, ate_upper = self.learner.estimate_ate(
+                X=X, p=p, treatment=treatment, y=y
             )
-        else:
-            try:
-                ate, ate_lower, ate_upper = learner.estimate_ate(
-                    X=X, p=p, treatment=treatment, y=y
-                )
-            except TypeError:
-                ate, ate_lower, ate_upper = learner.estimate_ate(
-                    X=X, treatment=treatment, y=y, return_ci=True
-                )
         return ate[0], ate_lower[0], ate_upper[0]
 
     @staticmethod