Pylint

Author: jmafoster1

causal_testing/estimation/iv_estimator.py

@@ -19,6 +19,7 @@ class InstrumentalVariableEstimator(Estimator):
 
     def __init__(
         # pylint: disable=too-many-arguments
+        # pylint: disable=duplicate-code
         self,
         treatment: str,
         treatment_value: float,
@@ -27,8 +28,6 @@ def __init__(
         outcome: str,
         instrument: str,
         df: pd.DataFrame = None,
-        intercept: int = 1,
-        effect_modifiers: dict = None,  # Not used (yet?). Needed for compatibility
         alpha: float = 0.05,
         query: str = "",
     ):
@@ -43,8 +42,7 @@ def __init__(
             alpha=alpha,
             query=query,
         )
-        self.intercept = intercept
-        self.model = None
+
         self.instrument = instrument
 
     def add_modelling_assumptions(self):

causal_testing/estimation/logistic_regression_estimator.py

@@ -34,6 +34,24 @@ def add_modelling_assumptions(self):
         self.modelling_assumptions.append("The outcome must be binary.")
         self.modelling_assumptions.append("Independently and identically distributed errors.")
 
+    def _predict(self, data=None, adjustment_config: dict = None) -> pd.DataFrame:
+        """Estimate the outcomes under control and treatment.
+
+        :param data: The data to use, defaults to `self.df`. Controllable for boostrap sampling.
+        :param: adjustment_config: The values of the adjustment variables to use.
+
+        :return: The estimated outcome under control and treatment, with confidence intervals in the form of a
+                 dataframe with columns "predicted", "se", "ci_lower", and "ci_upper".
+        """
+        if adjustment_config is None:
+            adjustment_config = {}
+        if set(self.adjustment_set) != set(adjustment_config):
+            raise ValueError(
+                f"Invalid adjustment configuration {adjustment_config}. Must specify values for {self.adjustment_set}"
+            )
+
+        return super()._predict(data, adjustment_config)
+
     def estimate_control_treatment(
         self, adjustment_config: dict = None, bootstrap_size: int = 100
     ) -> tuple[pd.Series, pd.Series]:

causal_testing/estimation/regression_estimator.py

@@ -88,7 +88,7 @@ def _run_regression(self, data=None) -> RegressionResultsWrapper:
         self.model = model
         return model
 
-    def _predict(self, data=None, adjustment_config: dict = None) -> tuple[pd.Series, pd.Series]:
+    def _predict(self, data=None, adjustment_config: dict = None) -> pd.DataFrame:
         """Estimate the outcomes under control and treatment.
 
         :param data: The data to use, defaults to `self.df`. Controllable for boostrap sampling.