Estimator return types

causal_testing/testing/causal_test_outcome.py

@@ -5,6 +5,7 @@
 from abc import ABC, abstractmethod
 from collections.abc import Iterable
 import numpy as np
+import pandas as pd
 
 from causal_testing.testing.causal_test_result import CausalTestResult
 
@@ -57,7 +58,7 @@ def apply(self, res: CausalTestResult) -> bool:
             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()]
             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)

causal_testing/testing/causal_test_result.py

@@ -99,12 +99,16 @@ def to_dict(self, json=False):
     def ci_low(self):
         """Return the lower bracket of the confidence intervals."""
         if self.confidence_intervals:
+            if isinstance(self.confidence_intervals[0], pd.Series):
+                return self.confidence_intervals[0][0]
             return self.confidence_intervals[0]
         return None
 
     def ci_high(self):
         """Return the higher bracket of the confidence intervals."""
         if self.confidence_intervals:
+            if isinstance(self.confidence_intervals[1], pd.Series):
+                return self.confidence_intervals[1][0]
             return self.confidence_intervals[1]
         return None
 

causal_testing/testing/estimators.py

@@ -11,7 +11,7 @@
 import statsmodels.formula.api as smf
 from econml.dml import CausalForestDML
 from patsy import dmatrix  # pylint: disable = no-name-in-module
-
+from patsy import ModelDesc
 from sklearn.ensemble import GradientBoostingRegressor
 from statsmodels.regression.linear_model import RegressionResultsWrapper
 from statsmodels.tools.sm_exceptions import PerfectSeparationError
@@ -351,19 +351,17 @@ def estimate_coefficient(self) -> float:
         """
         model = self._run_linear_regression()
         newline = "\n"
-        treatment = [self.treatment]
-        if str(self.df.dtypes[self.treatment]) == "object":
+        patsy_md = ModelDesc.from_formula(self.treatment)
+        if any((self.df.dtypes[factor.name()] == 'object' for factor in patsy_md.rhs_termlist[1].factors)):
             design_info = dmatrix(self.formula.split("~")[1], self.df).design_info
             treatment = design_info.column_names[design_info.term_name_slices[self.treatment]]
+        else:
+            treatment = [self.treatment]
         assert set(treatment).issubset(
             model.params.index.tolist()
         ), f"{treatment} not in\n{'  ' + str(model.params.index).replace(newline, newline + '  ')}"
         unit_effect = model.params[treatment]  # Unit effect is the coefficient of the treatment
         [ci_low, ci_high] = self._get_confidence_intervals(model, treatment)
-        if str(self.df.dtypes[self.treatment]) != "object":
-            unit_effect = unit_effect[0]
-            ci_low = ci_low[0]
-            ci_high = ci_high[0]
         return unit_effect, [ci_low, ci_high]
 
     def estimate_ate(self) -> tuple[float, list[float, float], float]:

tests/testing_tests/test_estimators.py

@@ -231,7 +231,7 @@ def test_program_11_2(self):
         self.assertEqual(round(model.params["Intercept"] + 90 * model.params["treatments"], 1), 216.9)
 
         # Increasing treatments from 90 to 100 should be the same as 10 times the unit ATE
-        self.assertEqual(round(model.params["treatments"], 1), round(ate, 1))
+        self.assertEqual(round(model.params["treatments"], 1), round(ate[0], 1))
 
     def test_program_11_3(self):
         """Test whether our linear regression implementation produces the same results as program 11.3 (p. 144)."""
@@ -251,7 +251,7 @@ def test_program_11_3(self):
             197.1,
         )
         # Increasing treatments from 90 to 100 should be the same as 10 times the unit ATE
-        self.assertEqual(round(model.params["treatments"], 3), round(ate, 3))
+        self.assertEqual(round(model.params["treatments"], 3), round(ate[0], 3))
 
     def test_program_15_1A(self):
         """Test whether our linear regression implementation produces the same results as program 15.1 (p. 163, 184)."""
@@ -329,8 +329,8 @@ def test_program_15_no_interaction(self):
         # terms_to_square = ["age", "wt71", "smokeintensity", "smokeyrs"]
         # for term_to_square in terms_to_square:
         ate, [ci_low, ci_high] = linear_regression_estimator.estimate_coefficient()
-        self.assertEqual(round(ate, 1), 3.5)
-        self.assertEqual([round(ci_low, 1), round(ci_high, 1)], [2.6, 4.3])
+        self.assertEqual(round(ate[0], 1), 3.5)
+        self.assertEqual([round(ci_low[0], 1), round(ci_high[0], 1)], [2.6, 4.3])
 
     def test_program_15_no_interaction_ate(self):
         """Test whether our linear regression implementation produces the same results as program 15.1 (p. 163, 184)