Reduced size of _return_causal_test_results

Author: jmafoster1

causal_testing/testing/causal_test_engine.py

@@ -148,63 +148,17 @@ def _return_causal_test_results(self, estimator, causal_test_case):
         :param causal_test_case: The concrete test case to be executed
         :return: a CausalTestResult object containing the confidence intervals
         """
-        if causal_test_case.estimate_type == "cate":
-            logger.debug("calculating cate")
-            if not hasattr(estimator, "estimate_cates"):
-                raise NotImplementedError(f"{estimator.__class__} has no CATE method.")
-
-            cates_df, confidence_intervals = estimator.estimate_cates()
-            causal_test_result = CausalTestResult(
-                estimator=estimator,
-                test_value=TestValue("ate", cates_df),
-                effect_modifier_configuration=causal_test_case.effect_modifier_configuration,
-                confidence_intervals=confidence_intervals,
-            )
-        elif causal_test_case.estimate_type == "risk_ratio":
-            logger.debug("calculating risk_ratio")
-            risk_ratio, confidence_intervals = estimator.estimate_risk_ratio(**causal_test_case.estimate_params)
-
-            causal_test_result = CausalTestResult(
-                estimator=estimator,
-                test_value=TestValue("risk_ratio", risk_ratio),
-                effect_modifier_configuration=causal_test_case.effect_modifier_configuration,
-                confidence_intervals=confidence_intervals,
-            )
-        elif causal_test_case.estimate_type == "coefficient":
-            logger.debug("calculating coefficient")
-            coefficient, confidence_intervals = estimator.estimate_unit_ate()
-            causal_test_result = CausalTestResult(
-                estimator=estimator,
-                test_value=TestValue("coefficient", coefficient),
-                effect_modifier_configuration=causal_test_case.effect_modifier_configuration,
-                confidence_intervals=confidence_intervals,
-            )
-        elif causal_test_case.estimate_type == "ate":
-            logger.debug("calculating ate")
-            ate, confidence_intervals = estimator.estimate_ate()
-            causal_test_result = CausalTestResult(
-                estimator=estimator,
-                test_value=TestValue("ate", ate),
-                effect_modifier_configuration=causal_test_case.effect_modifier_configuration,
-                confidence_intervals=confidence_intervals,
-            )
-            # causal_test_result = CausalTestResult(minimal_adjustment_set, ate, confidence_intervals)
-            # causal_test_result.apply_test_oracle_procedure(self.causal_test_case.expected_causal_effect)
-        elif causal_test_case.estimate_type == "ate_calculated":
-            logger.debug("calculating ate")
-            ate, confidence_intervals = estimator.estimate_ate_calculated()
-            causal_test_result = CausalTestResult(
-                estimator=estimator,
-                test_value=TestValue("ate", ate),
-                effect_modifier_configuration=causal_test_case.effect_modifier_configuration,
-                confidence_intervals=confidence_intervals,
-            )
-            # causal_test_result = CausalTestResult(minimal_adjustment_set, ate, confidence_intervals)
-            # causal_test_result.apply_test_oracle_procedure(self.causal_test_case.expected_causal_effect)
-        else:
-            raise ValueError(
-                f"Invalid estimate type {causal_test_case.estimate_type}, expected 'ate', 'cate', or 'risk_ratio'"
-            )
+        if not hasattr(estimator, f"estimate_{causal_test_case.estimate_type}"):
+            raise NotImplementedError(f"{estimator.__class__} has no {causal_test_case.estimate_type} method.")
+        estimate_effect = getattr(estimator, f"estimate_{causal_test_case.estimate_type}")
+        effect, confidence_intervals = estimate_effect(**causal_test_case.estimate_params)
+        causal_test_result = CausalTestResult(
+            estimator=estimator,
+            test_value=TestValue(causal_test_case.estimate_type, effect),
+            effect_modifier_configuration=causal_test_case.effect_modifier_configuration,
+            confidence_intervals=confidence_intervals,
+        )
+
         return causal_test_result
 
     def _check_positivity_violation(self, variables_list):

causal_testing/testing/estimators.py

@@ -341,7 +341,7 @@ def add_modelling_assumptions(self):
             "do not need to be linear."
         )
 
-    def estimate_unit_ate(self) -> float:
+    def estimate_coefficient(self) -> float:
         """Estimate the unit average treatment effect of the treatment on the outcome. That is, the change in outcome
         caused by a unit change in treatment.
 
@@ -495,7 +495,7 @@ def add_modelling_assumptions(self):
             (iii) Instrument and outcome do not share causes
         """
 
-    def estimate_coefficient(self, df):
+    def estimate_coefficient_aux(self, df):
         """
         Estimate the linear regression coefficient of the treatment on the
         outcome.
@@ -509,19 +509,19 @@ def estimate_coefficient(self, df):
         # Estimate the coefficient of I on X by cancelling
         return ab / a
 
-    def estimate_unit_ate(self, bootstrap_size=100):
+    def estimate_coefficient(self, bootstrap_size=100):
         """
         Estimate the unit ate (i.e. coefficient) of the treatment on the
         outcome.
         """
         bootstraps = sorted(
-            [self.estimate_coefficient(self.df.sample(len(self.df), replace=True)) for _ in range(bootstrap_size)]
+            [self.estimate_coefficient_aux(self.df.sample(len(self.df), replace=True)) for _ in range(bootstrap_size)]
         )
         bound = ceil((bootstrap_size * self.alpha) / 2)
         ci_low = bootstraps[bound]
         ci_high = bootstraps[bootstrap_size - bound]
 
-        return self.estimate_coefficient(self.df), (ci_low, ci_high)
+        return self.estimate_coefficient_aux(self.df), (ci_low, ci_high)
 
 
 class CausalForestEstimator(Estimator):

tests/json_front_tests/test_json_class.py

@@ -98,7 +98,7 @@ def test_f_flag(self):
         effects = {"NoEffect": NoEffect()}
         mutates = {
             "Increase": lambda x: self.json_class.scenario.treatment_variables[x].z3
-                                  > self.json_class.scenario.variables[x].z3
+            > self.json_class.scenario.variables[x].z3
         }
         estimators = {"LinearRegressionEstimator": LinearRegressionEstimator}
         with self.assertRaises(StatisticsError):
@@ -147,7 +147,7 @@ def test_run_json_tests_from_json(self):
         effects = {"NoEffect": NoEffect()}
         mutates = {
             "Increase": lambda x: self.json_class.scenario.treatment_variables[x].z3
-                                  > self.json_class.scenario.variables[x].z3
+            > self.json_class.scenario.variables[x].z3
         }
         estimators = {"LinearRegressionEstimator": LinearRegressionEstimator}
 
@@ -177,7 +177,7 @@ def test_generate_tests_from_json_no_dist(self):
         effects = {"NoEffect": NoEffect()}
         mutates = {
             "Increase": lambda x: self.json_class.scenario.treatment_variables[x].z3
-                                  > self.json_class.scenario.variables[x].z3
+            > self.json_class.scenario.variables[x].z3
         }
         estimators = {"LinearRegressionEstimator": LinearRegressionEstimator}
 
@@ -207,7 +207,7 @@ def test_formula_in_json_test(self):
         effects = {"Positive": Positive()}
         mutates = {
             "Increase": lambda x: self.json_class.scenario.treatment_variables[x].z3
-                                  > self.json_class.scenario.variables[x].z3
+            > self.json_class.scenario.variables[x].z3
         }
         estimators = {"LinearRegressionEstimator": LinearRegressionEstimator}
 
@@ -239,6 +239,7 @@ def test_run_concrete_json_testcase(self):
         with open("temp_out.txt", "r") as reader:
             temp_out = reader.readlines()
         self.assertIn("FAILED", temp_out[-1])
+
     def test_concrete_generate_params(self):
         example_test = {
             "tests": [
@@ -259,7 +260,7 @@ def test_concrete_generate_params(self):
         effects = {"NoEffect": NoEffect()}
         mutates = {
             "Increase": lambda x: self.json_class.scenario.treatment_variables[x].z3
-                                  > self.json_class.scenario.variables[x].z3
+            > self.json_class.scenario.variables[x].z3
         }
         estimators = {"LinearRegressionEstimator": LinearRegressionEstimator}
 

tests/testing_tests/test_causal_test_engine.py

@@ -220,7 +220,7 @@ def test_invalid_estimate_type(self):
             self.causal_test_engine.scenario_execution_data_df,
         )
         self.causal_test_case.estimate_type = "invalid"
-        with self.assertRaises(ValueError):
+        with self.assertRaises(NotImplementedError):
             self.causal_test_engine.execute_test(estimation_model, self.causal_test_case)
 
     def test_execute_test_observational_linear_regression_estimator_squared_term(self):
@@ -257,7 +257,7 @@ def test_execute_observational_causal_forest_estimator_cates(self):
             self.causal_test_engine.scenario_execution_data_df,
             effect_modifiers={"M": None},
         )
-        self.causal_test_case.estimate_type = "cate"
+        self.causal_test_case.estimate_type = "cates"
         causal_test_result = self.causal_test_engine.execute_test(estimation_model, self.causal_test_case)
         causal_test_result = causal_test_result.test_value.value
         # Check that each effect modifier's strata has a greater ATE than the last (ascending order)

tests/testing_tests/test_estimators.py

@@ -98,7 +98,7 @@ def setUpClass(cls) -> None:
     def test_linear_regression_categorical_ate(self):
         df = self.scarf_df.copy()
         logistic_regression_estimator = LinearRegressionEstimator("color", None, None, set(), "completed", df)
-        ate, confidence = logistic_regression_estimator.estimate_unit_ate()
+        ate, confidence = logistic_regression_estimator.estimate_coefficient()
         self.assertTrue(all([ci_low < 0 < ci_high for ci_low, ci_high in zip(confidence[0], confidence[1])]))
 
     def test_ate(self):
@@ -131,7 +131,9 @@ def test_ate_invalid_adjustment(self):
         df = self.scarf_df.copy()
         logistic_regression_estimator = LogisticRegressionEstimator("length_in", 65, 55, {}, "completed", df)
         with self.assertRaises(ValueError):
-            ate, _ = logistic_regression_estimator.estimate_ate(estimator_params={"adjustment_config": {"large_gauge": 0}})
+            ate, _ = logistic_regression_estimator.estimate_ate(
+                estimator_params={"adjustment_config": {"large_gauge": 0}}
+            )
 
     def test_ate_effect_modifiers(self):
         df = self.scarf_df.copy()
@@ -184,7 +186,7 @@ def test_estimate_coefficient(self):
         )
         self.assertEqual(iv_estimator.estimate_coefficient(self.df), 2)
 
-    def test_estimate_unit_ate(self):
+    def test_estimate_coefficient(self):
         """
         Test we get the correct coefficient.
         """
@@ -197,8 +199,8 @@ def test_estimate_unit_ate(self):
             outcome="Y",
             instrument="Z",
         )
-        unit_ate, [low, high] = iv_estimator.estimate_unit_ate()
-        self.assertEqual(unit_ate, 2)
+        coefficient, [low, high] = iv_estimator.estimate_coefficient()
+        self.assertEqual(coefficient, 2)
 
 
 class TestLinearRegressionEstimator(unittest.TestCase):
@@ -218,7 +220,7 @@ def test_program_11_2(self):
         df = self.chapter_11_df
         linear_regression_estimator = LinearRegressionEstimator("treatments", None, None, set(), "outcomes", df)
         model = linear_regression_estimator._run_linear_regression()
-        ate, _ = linear_regression_estimator.estimate_unit_ate()
+        ate, _ = linear_regression_estimator.estimate_coefficient()
 
         self.assertEqual(round(model.params["Intercept"] + 90 * model.params["treatments"], 1), 216.9)
 
@@ -232,7 +234,7 @@ def test_program_11_3(self):
             "treatments", None, None, set(), "outcomes", df, formula="outcomes ~ treatments + np.power(treatments, 2)"
         )
         model = linear_regression_estimator._run_linear_regression()
-        ate, _ = linear_regression_estimator.estimate_unit_ate()
+        ate, _ = linear_regression_estimator.estimate_coefficient()
         self.assertEqual(
             round(
                 model.params["Intercept"]
@@ -320,7 +322,7 @@ 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_unit_ate()
+        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])