All the tests pass and got rid of JSON front

Author: jmafoster1

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__)
 
 
@@ -30,21 +32,19 @@ class Estimator(ABC):
     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[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,21 @@ 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 = "",
     ):
         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 +52,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 +105,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 +117,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)

causal_testing/estimation/experimental_estimator.py

@@ -5,6 +5,7 @@
 import pandas as pd
 
 from causal_testing.estimation.abstract_estimator import Estimator
+from causal_testing.testing.base_test_case import BaseTestCase
 
 
 class ExperimentalEstimator(Estimator):
@@ -16,22 +17,20 @@ class ExperimentalEstimator(Estimator):
     def __init__(
         # pylint: disable=too-many-arguments
         self,
-        treatment: str,
+        base_test_case: BaseTestCase,
         treatment_value: float,
         control_value: float,
         adjustment_set: dict[str:Any],
-        outcome: str,
         effect_modifiers: dict[str:Any] = None,
         alpha: float = 0.05,
         repeats: int = 200,
     ):
         # 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,
             effect_modifiers=effect_modifiers,
             alpha=alpha,
         )
@@ -62,21 +61,40 @@ def estimate_ate(self) -> tuple[pd.Series, list[pd.Series, pd.Series]]:
 
         :return: The average treatment effect and the bootstrapped confidence intervals.
         """
-        control_configuration = self.adjustment_set | self.effect_modifiers | {self.treatment: self.control_value}
-        treatment_configuration = self.adjustment_set | self.effect_modifiers | {self.treatment: self.treatment_value}
+        control_configuration = (
+            self.adjustment_set
+            | self.effect_modifiers
+            | {self.base_test_case.treatment_variable.name: self.control_value}
+        )
+        treatment_configuration = (
+            self.adjustment_set
+            | self.effect_modifiers
+            | {self.base_test_case.treatment_variable.name: self.treatment_value}
+        )
 
         control_outcomes = pd.DataFrame([self.run_system(control_configuration) for _ in range(self.repeats)])
         treatment_outcomes = pd.DataFrame([self.run_system(treatment_configuration) for _ in range(self.repeats)])
 
-        difference = (treatment_outcomes[self.outcome] - control_outcomes[self.outcome]).sort_values().reset_index()
+        difference = (
+            (
+                treatment_outcomes[self.base_test_case.outcome_variable.name]
+                - control_outcomes[self.base_test_case.outcome_variable.name]
+            )
+            .sort_values()
+            .reset_index()
+        )
 
         ci_low_index = round(self.repeats * (self.alpha / 2))
         ci_low = difference.iloc[ci_low_index]
         ci_high = difference.iloc[self.repeats - ci_low_index]
 
-        return pd.Series({self.treatment: difference.mean()[self.outcome]}), [
-            pd.Series({self.treatment: ci_low[self.outcome]}),
-            pd.Series({self.treatment: ci_high[self.outcome]}),
+        return pd.Series(
+            {self.base_test_case.treatment_variable.name: difference.mean()[self.base_test_case.outcome_variable.name]}
+        ), [
+            pd.Series({self.base_test_case.treatment_variable.name: ci_low[self.base_test_case.outcome_variable.name]}),
+            pd.Series(
+                {self.base_test_case.treatment_variable.name: ci_high[self.base_test_case.outcome_variable.name]}
+            ),
         ]
 
     def estimate_risk_ratio(self) -> tuple[pd.Series, list[pd.Series, pd.Series]]:
@@ -85,19 +103,38 @@ def estimate_risk_ratio(self) -> tuple[pd.Series, list[pd.Series, pd.Series]]:
 
         :return: The average treatment effect and the bootstrapped confidence intervals.
         """
-        control_configuration = self.adjustment_set | self.effect_modifiers | {self.treatment: self.control_value}
-        treatment_configuration = self.adjustment_set | self.effect_modifiers | {self.treatment: self.treatment_value}
+        control_configuration = (
+            self.adjustment_set
+            | self.effect_modifiers
+            | {self.base_test_case.treatment_variable.name: self.control_value}
+        )
+        treatment_configuration = (
+            self.adjustment_set
+            | self.effect_modifiers
+            | {self.base_test_case.treatment_variable.name: self.treatment_value}
+        )
 
         control_outcomes = pd.DataFrame([self.run_system(control_configuration) for _ in range(self.repeats)])
         treatment_outcomes = pd.DataFrame([self.run_system(treatment_configuration) for _ in range(self.repeats)])
 
-        difference = (treatment_outcomes[self.outcome] / control_outcomes[self.outcome]).sort_values().reset_index()
+        difference = (
+            (
+                treatment_outcomes[self.base_test_case.outcome_variable.name]
+                / control_outcomes[self.base_test_case.outcome_variable.name]
+            )
+            .sort_values()
+            .reset_index()
+        )
 
         ci_low_index = round(self.repeats * (self.alpha / 2))
         ci_low = difference.iloc[ci_low_index]
         ci_high = difference.iloc[self.repeats - ci_low_index]
 
-        return pd.Series({self.treatment: difference.mean()[self.outcome]}), [
-            pd.Series({self.treatment: ci_low[self.outcome]}),
-            pd.Series({self.treatment: ci_high[self.outcome]}),
+        return pd.Series(
+            {self.base_test_case.treatment_variable.name: difference.mean()[self.base_test_case.outcome_variable.name]}
+        ), [
+            pd.Series({self.base_test_case.treatment_variable.name: ci_low[self.base_test_case.outcome_variable.name]}),
+            pd.Series(
+                {self.base_test_case.treatment_variable.name: ci_high[self.base_test_case.outcome_variable.name]}
+            ),
         ]

causal_testing/estimation/instrumental_variable_estimator.py

@@ -7,6 +7,7 @@
 import statsmodels.api as sm
 
 from causal_testing.estimation.abstract_estimator import Estimator
+from causal_testing.testing.base_test_case import BaseTestCase
 
 logger = logging.getLogger(__name__)
 
@@ -21,22 +22,20 @@ def __init__(
         # pylint: disable=too-many-arguments
         # pylint: disable=duplicate-code
         self,
-        treatment: str,
+        base_test_case: BaseTestCase,
         treatment_value: float,
         control_value: float,
         adjustment_set: set,
-        outcome: str,
         instrument: str,
         df: pd.DataFrame = None,
         alpha: float = 0.05,
         query: str = "",
     ):
         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=None,
             alpha=alpha,
@@ -68,10 +67,10 @@ def estimate_iv_coefficient(self, df) -> float:
         outcome.
         """
         # Estimate the total effect of instrument I on outcome Y = abI + c1
-        ab = sm.OLS(df[self.outcome], df[[self.instrument]]).fit().params[self.instrument]
+        ab = sm.OLS(df[self.base_test_case.outcome_variable.name], df[[self.instrument]]).fit().params[self.instrument]
 
         # Estimate the direct effect of instrument I on treatment X = aI + c1
-        a = sm.OLS(df[self.treatment], df[[self.instrument]]).fit().params[self.instrument]
+        a = sm.OLS(df[self.base_test_case.treatment_variable.name], df[[self.instrument]]).fit().params[self.instrument]
 
         # Estimate the coefficient of I on X by cancelling
         return ab / a

causal_testing/estimation/ipcw_estimator.py

@@ -11,6 +11,8 @@
 from lifelines import CoxPHFitter
 
 from causal_testing.estimation.abstract_estimator import Estimator
+from causal_testing.testing.base_test_case import BaseTestCase
+from causal_testing.specification.variable import Input, Output
 
 logger = logging.getLogger(__name__)
 
@@ -56,13 +58,12 @@ def __init__(
                             treatment) with the most elements multiplied by `timesteps_per_observation`.
         """
         super().__init__(
-            [var for _, var, _ in treatment_strategy],
-            [val for _, _, val in treatment_strategy],
-            [val for _, _, val in control_strategy],
-            None,
-            outcome,
-            df,
-            None,
+            base_test_case=BaseTestCase(Input("_", float), Output(outcome, float)),
+            treatment_value=[val for _, _, val in treatment_strategy],
+            control_value=[val for _, _, val in control_strategy],
+            adjustment_set=None,
+            df=df,
+            effect_modifiers=None,
             alpha=alpha,
             query="",
         )