Jmafoster1/remove data collector

causal_testing/estimation/abstract_estimator.py

@@ -6,6 +6,8 @@
 
 import pandas as pd
 
+from causal_testing.testing.base_test_case import BaseTestCase
+
 logger = logging.getLogger(__name__)
 
 
@@ -29,22 +31,21 @@ class Estimator(ABC):
 
     def __init__(
         # pylint: disable=too-many-arguments
+        # pylint: disable=R0801
         self,
-        treatment: str,
+        base_test_case: BaseTestCase,
         treatment_value: float,
         control_value: float,
         adjustment_set: set,
-        outcome: str,
         df: pd.DataFrame = None,
         effect_modifiers: dict[str:Any] = None,
         alpha: float = 0.05,
         query: str = "",
     ):
-        self.treatment = treatment
+        self.base_test_case = base_test_case
         self.treatment_value = treatment_value
         self.control_value = control_value
         self.adjustment_set = adjustment_set
-        self.outcome = outcome
         self.alpha = alpha
         self.df = df.query(query) if query else df
 

causal_testing/estimation/abstract_regression_estimator.py

@@ -10,6 +10,7 @@
 
 from causal_testing.specification.variable import Variable
 from causal_testing.estimation.abstract_estimator import Estimator
+from causal_testing.testing.base_test_case import BaseTestCase
 
 logger = logging.getLogger(__name__)
 
@@ -22,23 +23,22 @@ class RegressionEstimator(Estimator):
     def __init__(
         # pylint: disable=too-many-arguments
         self,
-        treatment: str,
+        base_test_case: BaseTestCase,
         treatment_value: float,
         control_value: float,
         adjustment_set: set,
-        outcome: str,
         df: pd.DataFrame = None,
         effect_modifiers: dict[Variable:Any] = None,
         formula: str = None,
         alpha: float = 0.05,
         query: str = "",
     ):
+        # pylint: disable=R0801
         super().__init__(
-            treatment=treatment,
+            base_test_case=base_test_case,
             treatment_value=treatment_value,
             control_value=control_value,
             adjustment_set=adjustment_set,
-            outcome=outcome,
             df=df,
             effect_modifiers=effect_modifiers,
             alpha=alpha,
@@ -53,8 +53,10 @@ def __init__(
         if formula is not None:
             self.formula = formula
         else:
-            terms = [treatment] + sorted(list(adjustment_set)) + sorted(list(effect_modifiers))
-            self.formula = f"{outcome} ~ {'+'.join(terms)}"
+            terms = (
+                [base_test_case.treatment_variable.name] + sorted(list(adjustment_set)) + sorted(list(effect_modifiers))
+            )
+            self.formula = f"{base_test_case.outcome_variable.name} ~ {'+'.join(terms)}"
 
     @property
     @abstractmethod
@@ -104,7 +106,7 @@ def _predict(self, data=None, adjustment_config: dict = None) -> pd.DataFrame:
 
         x = pd.DataFrame(columns=self.df.columns)
         x["Intercept"] = 1  # self.intercept
-        x[self.treatment] = [self.treatment_value, self.control_value]
+        x[self.base_test_case.treatment_variable.name] = [self.treatment_value, self.control_value]
 
         for k, v in adjustment_config.items():
             x[k] = v
@@ -116,5 +118,5 @@ def _predict(self, data=None, adjustment_config: dict = None) -> pd.DataFrame:
                 x = pd.get_dummies(x, columns=[col], drop_first=True)
 
         # This has to be here in case the treatment variable is in an I(...) block in the self.formula
-        x[self.treatment] = [self.treatment_value, self.control_value]
+        x[self.base_test_case.treatment_variable.name] = [self.treatment_value, self.control_value]
         return model.get_prediction(x).summary_frame()

causal_testing/estimation/cubic_spline_estimator.py

@@ -8,6 +8,7 @@
 
 from causal_testing.specification.variable import Variable
 from causal_testing.estimation.linear_regression_estimator import LinearRegressionEstimator
+from causal_testing.testing.base_test_case import BaseTestCase
 
 logger = logging.getLogger(__name__)
 
@@ -20,11 +21,10 @@ class CubicSplineRegressionEstimator(LinearRegressionEstimator):
     def __init__(
         # pylint: disable=too-many-arguments
         self,
-        treatment: str,
+        base_test_case: BaseTestCase,
         treatment_value: float,
         control_value: float,
         adjustment_set: set,
-        outcome: str,
         basis: int,
         df: pd.DataFrame = None,
         effect_modifiers: dict[Variable:Any] = None,
@@ -33,7 +33,7 @@ def __init__(
         expected_relationship=None,
     ):
         super().__init__(
-            treatment, treatment_value, control_value, adjustment_set, outcome, df, effect_modifiers, formula, alpha
+            base_test_case, treatment_value, control_value, adjustment_set, df, effect_modifiers, formula, alpha
         )
 
         self.expected_relationship = expected_relationship
@@ -42,8 +42,10 @@ def __init__(
             effect_modifiers = []
 
         if formula is None:
-            terms = [treatment] + sorted(list(adjustment_set)) + sorted(list(effect_modifiers))
-            self.formula = f"{outcome} ~ cr({'+'.join(terms)}, df={basis})"
+            terms = (
+                [base_test_case.treatment_variable.name] + sorted(list(adjustment_set)) + sorted(list(effect_modifiers))
+            )
+            self.formula = f"{base_test_case.outcome_variable.name} ~ cr({'+'.join(terms)}, df={basis})"
 
     def estimate_ate_calculated(self, adjustment_config: dict = None) -> pd.Series:
         """Estimate the ate effect of the treatment on the outcome. That is, the change in outcome caused
@@ -59,7 +61,7 @@ def estimate_ate_calculated(self, adjustment_config: dict = None) -> pd.Series:
         """
         model = self._run_regression()
 
-        x = {"Intercept": 1, self.treatment: self.treatment_value}
+        x = {"Intercept": 1, self.base_test_case.treatment_variable.name: self.treatment_value}
         if adjustment_config is not None:
             for k, v in adjustment_config.items():
                 x[k] = v
@@ -69,7 +71,7 @@ def estimate_ate_calculated(self, adjustment_config: dict = None) -> pd.Series:
 
         treatment = model.predict(x).iloc[0]
 
-        x[self.treatment] = self.control_value
+        x[self.base_test_case.treatment_variable.name] = self.control_value
         control = model.predict(x).iloc[0]
 
         return pd.Series(treatment - control)