Deprecated cubic spline estimator - functionality subsumed by LR estimator

Author: jmafoster1

causal_testing/estimation/abstract_regression_estimator.py

@@ -6,7 +6,7 @@
 from typing import Any
 
 import pandas as pd
-from patsy import ModelDesc, dmatrices, dmatrix  # pylint: disable = no-name-in-module
+from patsy import ModelDesc, build_design_matrices, dmatrices  # pylint: disable = no-name-in-module
 from statsmodels.regression.linear_model import RegressionResultsWrapper
 
 from causal_testing.estimation.abstract_estimator import Estimator
@@ -114,7 +114,7 @@ def _setup_covariates(self, df: pd.DataFrame) -> pd.Series:
         _, covariate_data = dmatrices(self.formula, df, return_type="dataframe")
         df = pd.concat([df, covariate_data[[col for col in covariate_data.columns if col not in df]]], axis=1)
         covariates = covariate_data.columns.tolist()
-        return covariates, df.dropna(subset=covariates)
+        return df.dropna(subset=covariates), covariates, covariate_data.design_info
 
     @property
     @abstractmethod
@@ -143,8 +143,9 @@ def fit_model(self, df: pd.DataFrame) -> RegressionResultsWrapper:
         :param df: The data to use.
         :return: The model after fitting to data.
         """
-        covariates, df = self._setup_covariates(df)
+        df, covariates, design_info = self._setup_covariates(df)
         model = self.regressor(df[self.outcome_variable], df[covariates]).fit(disp=0)
+        model.design_info = design_info
         return model
 
     def treatment_columns(self, model: RegressionResultsWrapper) -> list[str]:
@@ -180,7 +181,7 @@ def _predict(self, df) -> pd.DataFrame:
 
         for k, v in self.adjustment_config.items():
             x[k] = v
-        x = dmatrix(self.formula.split("~")[1], x, return_type="dataframe")
+        x = build_design_matrices([model.design_info], x, return_type="dataframe")[0]
 
         return model.get_prediction(x).summary_frame()
 

causal_testing/estimation/cubic_spline_estimator.py

@@ -406,3 +406,28 @@ def test_categorical_confidence_intervals(self):
         self.assertTrue(
             effect_estimate.ci_high.round(2).equals(pd.Series({"color[T.grey]": 23.95, "color[T.orange]": 17.08}))
         )
+
+    def test_program_11_3_linear_regression(self):
+        """Test whether the cublic_spline regression implementation produces the same results as program 11.3 (p. 162).
+        https://www.hsph.harvard.edu/miguel-hernan/wp-content/uploads/sites/1268/2023/10/hernanrobins_WhatIf_30sep23.pdf
+        Slightly modified as Hernan et al. use linear regression for this example.
+        """
+
+        df = load_chapter_11_df()
+
+        cublic_spline_estimator = LinearRegressionEstimator(
+            treatment_variable="treatments",
+            outcome_variable="outcomes",
+            treatment_value=1,
+            control_value=0,
+            formula="outcomes ~ cr(treatments, df=3)",
+        )
+
+        ate_1 = cublic_spline_estimator.estimate_ate_calculated(df).value
+
+        cublic_spline_estimator.treatment_value = 2
+        ate_2 = cublic_spline_estimator.estimate_ate_calculated(df).value
+
+        # Doubling the treatemebnt value should roughly but not exactly double the ATE
+        self.assertNotEqual(ate_1[0] * 2, ate_2[0])
+        self.assertAlmostEqual(ate_1[0] * 2, ate_2[0])