Merge pull request #116 from CITCOM-project/logistic-regression-confidence-intervals

Added confidence intervals to the logistic regression estimator

Author: jmafoster1

causal_testing/json_front/json_class.py

@@ -176,9 +176,7 @@ def _execute_test_case(self, causal_test_case: CausalTestCase, estimator: Estima
 
         result_string = str()
         if causal_test_result.ci_low() and causal_test_result.ci_high():
-            result_string = (
-                f"{causal_test_result.ci_low()} < {causal_test_result.test_value.value} <  {causal_test_result.ci_high()}"
-            )
+            result_string = f"{causal_test_result.ci_low()} < {causal_test_result.test_value.value} <  {causal_test_result.ci_high()}"
         else:
             result_string = causal_test_result.test_value.value
         if f_flag:

causal_testing/testing/estimators.py

@@ -149,7 +149,8 @@ def _run_logistic_regression(self) -> RegressionResultsWrapper:
     def estimate_control_treatment(self) -> tuple[pd.Series, pd.Series]:
         """Estimate the outcomes under control and treatment.
 
-        :return: The average treatment effect and the 95% Wald confidence intervals.
+        :return: The estimated control and treatment values and their confidence
+        intervals in the form ((ci_low, control, ci_high), (ci_low, treatment, ci_high)).
         """
         model = self._run_logistic_regression()
         self.model = model
@@ -168,31 +169,53 @@ def estimate_control_treatment(self) -> tuple[pd.Series, pd.Series]:
         x = x[model.params.index]
 
         y = model.predict(x)
-        return y.iloc[1], y.iloc[0]
+
+        # Delta method confidence intervals from
+        # https://stackoverflow.com/questions/47414842/confidence-interval-of-probability-prediction-from-logistic-regression-statsmode
+        cov = model.cov_params()
+        gradient = (y * (1 - y) * x.T).T  # matrix of gradients for each observation
+        std_errors = np.array([np.sqrt(np.dot(np.dot(g, cov), g)) for g in gradient.to_numpy()])
+        c = 1.96  # multiplier for confidence interval
+        upper = np.maximum(0, np.minimum(1, y + std_errors * c))
+        lower = np.maximum(0, np.minimum(1, y - std_errors * c))
+
+        return (lower.iloc[1], y.iloc[1], upper.iloc[1]), (lower.iloc[0], y.iloc[0], upper.iloc[0])
 
     def estimate_ate(self) -> 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 average treatment effect. Confidence intervals are not yet supported.
+        :return: The estimated average treatment effect and 95% confidence intervals
         """
-        control_outcome, treatment_outcome = self.estimate_control_treatment()
+        (cci_low, control_outcome, cci_high), (tci_low, treatment_outcome, tci_high) = self.estimate_control_treatment()
+
+        ci_low = tci_low - cci_high
+        ci_high = tci_high - cci_low
+        estimate = treatment_outcome - control_outcome
 
-        return treatment_outcome - control_outcome
+        logger.info(
+            f"Changing {self.treatment} from {self.control_values} to {self.treatment_values} gives an estimated ATE of {ci_low} < {estimate} < {ci_high}"
+        )
+        assert ci_low < estimate < ci_high, f"Expecting {ci_low} < {estimate} < {ci_high}"
+
+        return estimate, (ci_low, ci_high)
 
     def estimate_risk_ratio(self) -> 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 average treatment effect. Confidence intervals are not yet supported.
+        :return: The estimated risk ratio and 95% confidence intervals.
         """
-        control_outcome, treatment_outcome = self.estimate_control_treatment()
+        (cci_low, control_outcome, cci_high), (tci_low, treatment_outcome, tci_high) = self.estimate_control_treatment()
 
-        return treatment_outcome / control_outcome
+        ci_low = tci_low / cci_high
+        ci_high = tci_high / cci_low
+
+        return treatment_outcome / control_outcome, (ci_low, ci_high)
 
     def estimate_unit_odds_ratio(self) -> float:
         """Estimate the odds ratio of increasing the treatment by one. In logistic regression, this corresponds to the
@@ -214,7 +237,7 @@ def __init__(
         treatment: tuple,
         treatment_values: float,
         control_values: float,
-        adjustment_set: set,
+        adjustment_set: list[float],
         outcome: tuple,
         df: pd.DataFrame = None,
         effect_modifiers: dict[Variable:Any] = None,
@@ -332,7 +355,8 @@ def estimate_ate(self) -> tuple[float, list[float, float], float]:
     def estimate_control_treatment(self) -> tuple[pd.Series, pd.Series]:
         """Estimate the outcomes under control and treatment.
 
-        :return: The average treatment effect and the 95% Wald confidence intervals.
+        :return: The estimated outcome under control and treatment in the form
+        (control_outcome, treatment_outcome).
         """
         model = self._run_linear_regression()
         self.model = model

tests/data_collection_tests/test_observational_data_collector.py

@@ -9,15 +9,15 @@
 
 
 class TestObservationalDataCollector(unittest.TestCase):
-
     def setUp(self) -> None:
         temp_dir_path = create_temp_dir_if_non_existent()
         self.dag_dot_path = os.path.join(temp_dir_path, "dag.dot")
         self.observational_df_path = os.path.join(temp_dir_path, "observational_data.csv")
         # Y = 3*X1 + X2*X3 + 10
         self.observational_df = pd.DataFrame({"X1": [1, 2, 3, 4], "X2": [5, 6, 7, 8], "X3": [10, 20, 30, 40]})
         self.observational_df["Y"] = self.observational_df.apply(
-                lambda row: (3 * row.X1) + (row.X2 * row.X3) + 10, axis=1)
+            lambda row: (3 * row.X1) + (row.X2 * row.X3) + 10, axis=1
+        )
         self.observational_df.to_csv(self.observational_df_path)
         self.X1 = Input("X1", int, uniform(1, 4))
         self.X2 = Input("X2", int, rv_discrete(values=([7], [1])))
@@ -45,12 +45,13 @@ def test_data_constraints(self):
 
     def test_meta_population(self):
         def populate_m(data):
-            data['M'] = data['X1'] * 2
+            data["M"] = data["X1"] * 2
+
         meta = Meta("M", int, populate_m)
         scenario = Scenario({self.X1, meta})
         observational_data_collector = ObservationalDataCollector(scenario, self.observational_df_path)
         data = observational_data_collector.collect_data()
-        assert all((m == 2*x1 for x1, m in zip(data['X1'], data['M'])))
+        assert all((m == 2 * x1 for x1, m in zip(data["X1"], data["M"])))
 
     def tearDown(self) -> None:
         remove_temp_dir_if_existent()

tests/generation_tests/test_abstract_test_case.py

@@ -8,18 +8,21 @@
 from tests.test_helpers import create_temp_dir_if_non_existent, remove_temp_dir_if_existent
 from causal_testing.testing.causal_test_outcome import Positive
 
+
 class TestAbstractTestCase(unittest.TestCase):
     """
-        Class to test abstract test cases.
+    Class to test abstract test cases.
     """
+
     def setUp(self) -> None:
         temp_dir_path = create_temp_dir_if_non_existent()
         self.dag_dot_path = os.path.join(temp_dir_path, "dag.dot")
         self.observational_df_path = os.path.join(temp_dir_path, "observational_data.csv")
         # Y = 3*X1 + X2*X3 + 10
         self.observational_df = pd.DataFrame({"X1": [1, 2, 3, 4], "X2": [5, 6, 7, 8], "X3": [10, 20, 30, 40]})
-        self.observational_df["Y"] = self.observational_df.apply(lambda row: (3 * row.X1) + (row.X2 * row.X3) + 10,
-                                                                 axis=1)
+        self.observational_df["Y"] = self.observational_df.apply(
+            lambda row: (3 * row.X1) + (row.X2 * row.X3) + 10, axis=1
+        )
         self.observational_df.to_csv(self.observational_df_path)
         self.X1 = Input("X1", float, uniform(1, 4))
         self.X2 = Input("X2", int, rv_discrete(values=([7], [1])))
@@ -51,9 +54,9 @@ def test_str(self):
             expected_causal_effect={self.Y: Positive()},
             effect_modifiers=None,
         )
-        assert str(abstract) == \
-        "When we apply intervention {X1' > X1}, the effect on Output: Y::int should be Positive", \
-        f"Unexpected string {str(abstract)}"
+        assert (
+            str(abstract) == "When we apply intervention {X1' > X1}, the effect on Output: Y::int should be Positive"
+        ), f"Unexpected string {str(abstract)}"
 
     def test_datapath(self):
         scenario = Scenario({self.X1, self.X2, self.X3, self.X4})
@@ -109,7 +112,6 @@ def test_generate_concrete_test_cases_rct(self):
         assert len(concrete_tests) == 2, "Expected 2 concrete tests"
         assert len(runs) == 4, "Expected 4 runs"
 
-
     def test_infeasible_constraints(self):
         scenario = Scenario({self.X1, self.X2, self.X3, self.X4}, [self.X1.z3 > 2])
         scenario.setup_treatment_variables()
@@ -125,10 +127,9 @@ def test_infeasible_constraints(self):
 
         with self.assertWarns(Warning):
             concrete_tests, runs = abstract.generate_concrete_tests(4, rct=True, target_ks_score=0.1, hard_max=HARD_MAX)
-        self.assertTrue(all((x > 2 for x in runs['X1'])))
+        self.assertTrue(all((x > 2 for x in runs["X1"])))
         self.assertEqual(len(concrete_tests), HARD_MAX * NUM_STRATA)
 
-
     def test_feasible_constraints(self):
         scenario = Scenario({self.X1, self.X2, self.X3, self.X4})
         scenario.setup_treatment_variables()
@@ -142,7 +143,6 @@ def test_feasible_constraints(self):
         concrete_tests, _ = abstract.generate_concrete_tests(4, rct=True, target_ks_score=0.1, hard_max=1000)
         assert len(concrete_tests) < 1000
 
-
     def tearDown(self) -> None:
         remove_temp_dir_if_existent()
 

tests/json_front_tests/test_json_class.py

@@ -93,12 +93,10 @@ def test_generate_tests_from_json(self):
         self.json_class.test_plan = example_test
         effects = {"NoEffect": NoEffect()}
         mutates = {
-            "Increase": lambda x: self.json_class.modelling_scenario.treatment_variables[x].z3 >
-                                  self.json_class.modelling_scenario.variables[x].z3
-        }
-        estimators = {
-            "LinearRegressionEstimator": LinearRegressionEstimator
+            "Increase": lambda x: self.json_class.modelling_scenario.treatment_variables[x].z3
+            > self.json_class.modelling_scenario.variables[x].z3
         }
+        estimators = {"LinearRegressionEstimator": LinearRegressionEstimator}
 
         with self.assertLogs() as captured:
             self.json_class.generate_tests(effects, mutates, estimators, False)
@@ -108,9 +106,8 @@ def test_generate_tests_from_json(self):
 
     def tearDown(self) -> None:
         pass
-        #remove_temp_dir_if_existent()
+        # remove_temp_dir_if_existent()
 
 
 def populate_example(*args, **kwargs):
     pass
-