SomeEffect and NoneEffect applys now work with pd.Series

Author: christopher-wild

causal_testing/testing/causal_test_outcome.py

@@ -28,14 +28,13 @@ class SomeEffect(CausalTestOutcome):
     """An extension of TestOutcome representing that the expected causal effect should not be zero."""
 
     def apply(self, res: CausalTestResult) -> bool:
-        if res.test_value.type == "ate":
-            return (0 < res.ci_low() < res.ci_high()) or (res.ci_low() < res.ci_high() < 0)
-        if res.test_value.type == "coefficient":
-            ci_low = res.ci_low() if isinstance(res.ci_low(), Iterable) else [res.ci_low()]
-            ci_high = res.ci_high() if isinstance(res.ci_high(), Iterable) else [res.ci_high()]
-            return any(0 < ci_low < ci_high or ci_low < ci_high < 0 for ci_low, ci_high in zip(ci_low, ci_high))
         if res.test_value.type == "risk_ratio":
-            return (1 < res.ci_low() < res.ci_high()) or (res.ci_low() < res.ci_high() < 1)
+            return any(
+                1 < ci_low < ci_high or ci_low < ci_high < 1 for ci_low, ci_high in zip(res.ci_low(), res.ci_high()))
+        if res.test_value.type == "coefficient" or res.test_value.type == "ate":
+            return any(
+                0 < ci_low < ci_high or ci_low < ci_high < 0 for ci_low, ci_high in zip(res.ci_low(), res.ci_high()))
+
         raise ValueError(f"Test Value type {res.test_value.type} is not valid for this TestOutcome")
 
 
@@ -52,23 +51,21 @@ def __init__(self, atol: float = 1e-10, ctol: float = 0.05):
         self.ctol = ctol
 
     def apply(self, res: CausalTestResult) -> bool:
-        if res.test_value.type == "ate":
-            return (res.ci_low() < 0 < res.ci_high()) or (abs(res.test_value.value) < self.atol)[0]
-        if res.test_value.type == "coefficient":
-            ci_low = res.ci_low() if isinstance(res.ci_low(), Iterable) else [res.ci_low()]
-            ci_high = res.ci_high() if isinstance(res.ci_high(), Iterable) else [res.ci_high()]
+        if res.test_value.type == "risk_ratio":
+            return any(ci_low < 1 < ci_high or np.isclose(value, 1.0, atol=self.atol) for ci_low, ci_high, value in
+                       zip(res.ci_low(), res.ci_high(), res.test_value.value))
+        elif res.test_value.type == "coefficient" or res.test_value.type == "ate":
             value = res.test_value.value if isinstance(res.ci_high(), Iterable) else [res.test_value.value]
             value = value[0] if isinstance(value[0], pd.Series) else value
             return (
-                sum(
-                    not ((ci_low < 0 < ci_high) or abs(v) < self.atol)
-                    for ci_low, ci_high, v in zip(ci_low, ci_high, value)
-                )
-                / len(value)
-                < self.ctol
+                    sum(
+                        not ((ci_low < 0 < ci_high) or abs(v) < self.atol)
+                        for ci_low, ci_high, v in zip(res.ci_low(), res.ci_high(), value)
+                    )
+                    / len(value)
+                    < self.ctol
             )
-        if res.test_value.type == "risk_ratio":
-            return (res.ci_low() < 1 < res.ci_high()) or np.isclose(res.test_value.value, 1.0, atol=self.atol)
+
         raise ValueError(f"Test Value type {res.test_value.type} is not valid for this TestOutcome")