Merge branch 'main' into json_concrete_param

Author: christopher-wild

causal_testing/testing/causal_test_case.py

@@ -30,6 +30,7 @@ def __init__(
         control_value: Any = None,
         treatment_value: Any = None,
         estimate_type: str = "ate",
+        estimate_params: dict = None,
         effect_modifier_configuration: dict[Variable:Any] = None,
     ):
         """
@@ -47,6 +48,8 @@ def __init__(
         self.treatment_variable = base_test_case.treatment_variable
         self.treatment_value = treatment_value
         self.estimate_type = estimate_type
+        if estimate_params is None:
+            self.estimate_params = {}
         self.effect = base_test_case.effect
 
         if effect_modifier_configuration:

causal_testing/testing/causal_test_engine.py

@@ -162,7 +162,8 @@ def _return_causal_test_results(self, estimator, causal_test_case):
             )
         elif causal_test_case.estimate_type == "risk_ratio":
             logger.debug("calculating risk_ratio")
-            risk_ratio, confidence_intervals = estimator.estimate_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),

causal_testing/testing/estimators.py

@@ -179,7 +179,7 @@ def estimate(self, data: pd.DataFrame, adjustment_config: dict = None) -> Regres
         # x = x[model.params.index]
         return model.predict(x)
 
-    def estimate_control_treatment(self, bootstrap_size=100, adjustment_config=None) -> tuple[pd.Series, pd.Series]:
+    def estimate_control_treatment(self, bootstrap_size, adjustment_config) -> tuple[pd.Series, pd.Series]:
         """Estimate the outcomes under control and treatment.
 
         :return: The estimated control and treatment values and their confidence
@@ -215,14 +215,18 @@ def estimate_control_treatment(self, bootstrap_size=100, adjustment_config=None)
 
         return (y.iloc[1], np.array(control)), (y.iloc[0], np.array(treatment))
 
-    def estimate_ate(self, bootstrap_size=100, adjustment_config=None) -> float:
+    def estimate_ate(self, estimator_params: dict = None) -> float:
         """Estimate the ate effect of the treatment on the outcome. That is, the change in outcome caused
         by changing the treatment variable from the control value to the treatment value. Here, we actually
         calculate the expected outcomes under control and treatment and take one away from the other. This
         allows for custom terms to be put in such as squares, inverses, products, etc.
 
         :return: The estimated average treatment effect and 95% confidence intervals
         """
+        if estimator_params is None:
+            estimator_params = {}
+        bootstrap_size = estimator_params.get("bootstrap_size", 100)
+        adjustment_config = estimator_params.get("adjustment_config", None)
         (control_outcome, control_bootstraps), (
             treatment_outcome,
             treatment_bootstraps,
@@ -245,14 +249,18 @@ def estimate_ate(self, bootstrap_size=100, adjustment_config=None) -> float:
 
         return estimate, (ci_low, ci_high)
 
-    def estimate_risk_ratio(self, bootstrap_size=100, adjustment_config=None) -> float:
+    def estimate_risk_ratio(self, estimator_params: dict = None) -> float:
         """Estimate the ate effect of the treatment on the outcome. That is, the change in outcome caused
         by changing the treatment variable from the control value to the treatment value. Here, we actually
         calculate the expected outcomes under control and treatment and divide one by the other. This
         allows for custom terms to be put in such as squares, inverses, products, etc.
 
         :return: The estimated risk ratio and 95% confidence intervals.
         """
+        if estimator_params is None:
+            estimator_params = {}
+        bootstrap_size = estimator_params.get("bootstrap_size", 100)
+        adjustment_config = estimator_params.get("adjustment_config", None)
         (control_outcome, control_bootstraps), (
             treatment_outcome,
             treatment_bootstraps,

tests/testing_tests/test_estimators.py

@@ -124,14 +124,14 @@ def test_ate_adjustment(self):
         logistic_regression_estimator = LogisticRegressionEstimator(
             "length_in", 65, 55, {"large_gauge"}, "completed", df
         )
-        ate, _ = logistic_regression_estimator.estimate_ate(adjustment_config={"large_gauge": 0})
+        ate, _ = logistic_regression_estimator.estimate_ate(estimator_params={"adjustment_config": {"large_gauge": 0}})
         self.assertEqual(round(ate, 4), -0.3388)
 
     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(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()