Merge branch 'main' into windows_wheel_install

Author: christopher-wild

causal_testing/json_front/json_class.py

@@ -123,54 +123,15 @@ def run_json_tests(self, effects: dict, estimators: dict, f_flag: bool = False,
         :param estimators: Dictionary mapping estimator classes to string representations.
         :param f_flag: Failure flag that if True the script will stop executing when a test fails.
         """
-        failures = 0
-        msg = ""
         for test in self.test_plan["tests"]:
             if "skip" in test and test["skip"]:
                 continue
             test["estimator"] = estimators[test["estimator"]]
             if "mutations" in test:
                 if test["estimate_type"] == "coefficient":
-                    base_test_case = BaseTestCase(
-                        treatment_variable=next(self.scenario.variables[v] for v in test["mutations"]),
-                        outcome_variable=next(self.scenario.variables[v] for v in test["expected_effect"]),
-                        effect=test.get("effect", "direct"),
-                    )
-                    assert len(test["expected_effect"]) == 1, "Can only have one expected effect."
-                    causal_test_case = CausalTestCase(
-                        base_test_case=base_test_case,
-                        expected_causal_effect=next(
-                            effects[effect] for variable, effect in test["expected_effect"].items()
-                        ),
-                        estimate_type="coefficient",
-                        effect_modifier_configuration={
-                            self.scenario.variables[v] for v in test.get("effect_modifiers", [])
-                        },
-                    )
-                    result = self._execute_test_case(causal_test_case=causal_test_case, test=test, f_flag=f_flag)
-                    msg = (
-                        f"Executing test: {test['name']} \n"
-                        + f"  {causal_test_case} \n"
-                        + "  "
-                        + ("\n  ").join(str(result[1]).split("\n"))
-                        + "==============\n"
-                        + f"  Result: {'FAILED' if result[0] else 'Passed'}"
-                    )
-                    print(msg)
+                    msg = self._run_coefficient_test(test=test, f_flag=f_flag, effects=effects)
                 else:
-                    abstract_test = self._create_abstract_test_case(test, mutates, effects)
-                    concrete_tests, _ = abstract_test.generate_concrete_tests(5, 0.05)
-                    failures, _ = self._execute_tests(concrete_tests, test, f_flag)
-
-                    msg = (
-                        f"Executing test: {test['name']} \n"
-                        + "  abstract_test \n"
-                        + f"  {abstract_test} \n"
-                        + f"  {abstract_test.treatment_variable.name},"
-                        + f"  {abstract_test.treatment_variable.distribution} \n"
-                        + f"  Number of concrete tests for test case: {str(len(concrete_tests))} \n"
-                        + f"  {failures}/{len(concrete_tests)} failed for {test['name']}"
-                    )
+                    msg = self._run_ate_test(test=test, f_flag=f_flag, effects=effects, mutates=mutates)
                 self._append_to_file(msg, logging.INFO)
             else:
                 outcome_variable = next(
@@ -197,8 +158,74 @@ def run_json_tests(self, effects: dict, estimators: dict, f_flag: bool = False,
                     + f"control value = {test['control_value']}, treatment value = {test['treatment_value']} \n"
                     + f"Result: {'FAILED' if failed else 'Passed'}"
                 )
+                print(msg)
                 self._append_to_file(msg, logging.INFO)
 
+    def _run_coefficient_test(self, test: dict, f_flag: bool, effects: dict):
+        """Builds structures and runs test case for tests with an estimate_type of 'coefficient'.
+
+        :param test: Single JSON test definition stored in a mapping (dict)
+        :param f_flag: Failure flag that if True the script will stop executing when a test fails.
+        :param effects: Dictionary mapping effect class instances to string representations.
+        :return: String containing the message to be outputted
+        """
+        base_test_case = BaseTestCase(
+            treatment_variable=next(self.scenario.variables[v] for v in test["mutations"]),
+            outcome_variable=next(self.scenario.variables[v] for v in test["expected_effect"]),
+            effect=test.get("effect", "direct"),
+        )
+        assert len(test["expected_effect"]) == 1, "Can only have one expected effect."
+        causal_test_case = CausalTestCase(
+            base_test_case=base_test_case,
+            expected_causal_effect=next(effects[effect] for variable, effect in test["expected_effect"].items()),
+            estimate_type="coefficient",
+            effect_modifier_configuration={self.scenario.variables[v] for v in test.get("effect_modifiers", [])},
+        )
+        result = self._execute_test_case(causal_test_case=causal_test_case, test=test, f_flag=f_flag)
+        msg = (
+            f"Executing test: {test['name']} \n"
+            + f"  {causal_test_case} \n"
+            + "  "
+            + ("\n  ").join(str(result[1]).split("\n"))
+            + "==============\n"
+            + f"  Result: {'FAILED' if result[0] else 'Passed'}"
+        )
+        return msg
+
+    def _run_ate_test(self, test: dict, f_flag: bool, effects: dict, mutates: dict):
+        """Builds structures and runs test case for tests with an estimate_type of 'ate'.
+
+        :param test: Single JSON test definition stored in a mapping (dict)
+        :param f_flag: Failure flag that if True the script will stop executing when a test fails.
+        :param effects: Dictionary mapping effect class instances to string representations.
+        :param mutates: Dictionary mapping mutation functions to string representations.
+        :return: String containing the message to be outputted
+        """
+        if "sample_size" in test:
+            sample_size = test["sample_size"]
+        else:
+            sample_size = 5
+        if "target_ks_score" in test:
+            target_ks_score = test["target_ks_score"]
+        else:
+            target_ks_score = 0.05
+        abstract_test = self._create_abstract_test_case(test, mutates, effects)
+        concrete_tests, _ = abstract_test.generate_concrete_tests(
+            sample_size=sample_size, target_ks_score=target_ks_score
+        )
+        failures, _ = self._execute_tests(concrete_tests, test, f_flag)
+
+        msg = (
+            f"Executing test: {test['name']} \n"
+            + "  abstract_test \n"
+            + f"  {abstract_test} \n"
+            + f"  {abstract_test.treatment_variable.name},"
+            + f"  {abstract_test.treatment_variable.distribution} \n"
+            + f"  Number of concrete tests for test case: {str(len(concrete_tests))} \n"
+            + f"  {failures}/{len(concrete_tests)} failed for {test['name']}"
+        )
+        return msg
+
     def _execute_tests(self, concrete_tests, test, f_flag):
         failures = 0
         details = []
@@ -286,6 +313,7 @@ def _setup_test(
             "outcome": causal_test_case.outcome_variable.name,
             "df": causal_test_engine.scenario_execution_data_df,
             "effect_modifiers": causal_test_case.effect_modifier_configuration,
+            "alpha": test["alpha"] if "alpha" in test else 0.05,
         }
         if "formula" in test:
             estimator_kwargs["formula"] = test["formula"]

causal_testing/specification/metamorphic_relation.py

@@ -181,6 +181,7 @@ def to_json_stub(self, skip=True) -> dict:
             "mutations": [self.treatment_var],
             "expected_effect": {self.output_var: "NoEffect"},
             "formula": f"{self.output_var} ~ {' + '.join([self.treatment_var] + self.adjustment_vars)}",
+            "alpha": 0.05,
             "skip": skip,
         }
 

causal_testing/testing/causal_test_case.py

@@ -30,6 +30,7 @@ def __init__(
         control_value: Any = None,
         treatment_value: Any = None,
         estimate_type: str = "ate",
+        estimate_params: dict = None,
         effect_modifier_configuration: dict[Variable:Any] = None,
     ):
         """
@@ -47,6 +48,8 @@ def __init__(
         self.treatment_variable = base_test_case.treatment_variable
         self.treatment_value = treatment_value
         self.estimate_type = estimate_type
+        if estimate_params is None:
+            self.estimate_params = {}
         self.effect = base_test_case.effect
 
         if effect_modifier_configuration:

causal_testing/testing/causal_test_engine.py

@@ -162,7 +162,8 @@ def _return_causal_test_results(self, estimator, causal_test_case):
             )
         elif causal_test_case.estimate_type == "risk_ratio":
             logger.debug("calculating risk_ratio")
-            risk_ratio, confidence_intervals = estimator.estimate_risk_ratio()
+            risk_ratio, confidence_intervals = estimator.estimate_risk_ratio(**causal_test_case.estimate_params)
+
             causal_test_result = CausalTestResult(
                 estimator=estimator,
                 test_value=TestValue("risk_ratio", risk_ratio),

causal_testing/testing/causal_test_outcome.py

@@ -51,6 +51,13 @@ def apply(self, res: CausalTestResult) -> bool:
             ci_low = res.ci_low() if isinstance(res.ci_low(), Iterable) else [res.ci_low()]
             ci_high = res.ci_high() if isinstance(res.ci_high(), Iterable) else [res.ci_high()]
             value = res.test_value.value if isinstance(res.ci_high(), Iterable) else [res.test_value.value]
+
+            if not all(ci_low < 0 < ci_high for ci_low, ci_high in zip(ci_low, ci_high)):
+                print(
+                    "FAILING ON",
+                    [(ci_low, ci_high) for ci_low, ci_high in zip(ci_low, ci_high) if not ci_low < 0 < ci_high],
+                )
+
             return all(ci_low < 0 < ci_high for ci_low, ci_high in zip(ci_low, ci_high)) or all(
                 abs(v) < self.atol for v in value
             )

causal_testing/testing/causal_test_result.py

@@ -56,6 +56,7 @@ def push(s, inc="  "):
             f"Treatment value: {self.estimator.treatment_value}\n"
             f"Outcome: {self.estimator.outcome}\n"
             f"Adjustment set: {self.adjustment_set}\n"
+            f"Formula: {self.estimator.formula}\n"
             f"{self.test_value.type}: {result_str}\n"
         )
         confidence_str = ""
@@ -64,6 +65,7 @@ def push(s, inc="  "):
             if "\n" in ci_str:
                 ci_str = " " + push(pd.DataFrame(self.confidence_intervals).transpose().to_string(header=False))
             confidence_str += f"Confidence intervals:{ci_str}\n"
+            confidence_str += f"Alpha:{self.estimator.alpha}\n"
         return base_str + confidence_str
 
     def to_dict(self):

causal_testing/testing/estimators.py

@@ -49,13 +49,15 @@ def __init__(
         outcome: str,
         df: pd.DataFrame = None,
         effect_modifiers: dict[str:Any] = None,
+        alpha: float = 0.05,
     ):
         self.treatment = treatment
         self.treatment_value = treatment_value
         self.control_value = control_value
         self.adjustment_set = adjustment_set
         self.outcome = outcome
         self.df = df
+        self.alpha = alpha
         if effect_modifiers is None:
             self.effect_modifiers = {}
         elif isinstance(effect_modifiers, dict):
@@ -179,7 +181,7 @@ def estimate(self, data: pd.DataFrame, adjustment_config: dict = None) -> Regres
         # x = x[model.params.index]
         return model.predict(x)
 
-    def estimate_control_treatment(self, bootstrap_size=100, adjustment_config=None) -> tuple[pd.Series, pd.Series]:
+    def estimate_control_treatment(self, bootstrap_size, adjustment_config) -> tuple[pd.Series, pd.Series]:
         """Estimate the outcomes under control and treatment.
 
         :return: The estimated control and treatment values and their confidence
@@ -215,14 +217,18 @@ def estimate_control_treatment(self, bootstrap_size=100, adjustment_config=None)
 
         return (y.iloc[1], np.array(control)), (y.iloc[0], np.array(treatment))
 
-    def estimate_ate(self, bootstrap_size=100, adjustment_config=None) -> float:
+    def estimate_ate(self, estimator_params: dict = None) -> 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 estimated average treatment effect and 95% confidence intervals
         """
+        if estimator_params is None:
+            estimator_params = {}
+        bootstrap_size = estimator_params.get("bootstrap_size", 100)
+        adjustment_config = estimator_params.get("adjustment_config", None)
         (control_outcome, control_bootstraps), (
             treatment_outcome,
             treatment_bootstraps,
@@ -233,7 +239,7 @@ def estimate_ate(self, bootstrap_size=100, adjustment_config=None) -> float:
             return estimate, (None, None)
 
         bootstraps = sorted(list(treatment_bootstraps - control_bootstraps))
-        bound = int((bootstrap_size * 0.05) / 2)
+        bound = int((bootstrap_size * self.alpha) / 2)
         ci_low = bootstraps[bound]
         ci_high = bootstraps[bootstrap_size - bound]
 
@@ -245,14 +251,18 @@ def estimate_ate(self, bootstrap_size=100, adjustment_config=None) -> float:
 
         return estimate, (ci_low, ci_high)
 
-    def estimate_risk_ratio(self, bootstrap_size=100, adjustment_config=None) -> float:
+    def estimate_risk_ratio(self, estimator_params: dict = None) -> 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 estimated risk ratio and 95% confidence intervals.
         """
+        if estimator_params is None:
+            estimator_params = {}
+        bootstrap_size = estimator_params.get("bootstrap_size", 100)
+        adjustment_config = estimator_params.get("adjustment_config", None)
         (control_outcome, control_bootstraps), (
             treatment_outcome,
             treatment_bootstraps,
@@ -263,7 +273,7 @@ def estimate_risk_ratio(self, bootstrap_size=100, adjustment_config=None) -> flo
             return estimate, (None, None)
 
         bootstraps = sorted(list(treatment_bootstraps / control_bootstraps))
-        bound = ceil((bootstrap_size * 0.05) / 2)
+        bound = ceil((bootstrap_size * self.alpha) / 2)
         ci_low = bootstraps[bound]
         ci_high = bootstraps[bootstrap_size - bound]
 
@@ -301,8 +311,11 @@ def __init__(
         df: pd.DataFrame = None,
         effect_modifiers: dict[Variable:Any] = None,
         formula: str = None,
+        alpha: float = 0.05,
     ):
-        super().__init__(treatment, treatment_value, control_value, adjustment_set, outcome, df, effect_modifiers)
+        super().__init__(
+            treatment, treatment_value, control_value, adjustment_set, outcome, df, effect_modifiers, alpha=alpha
+        )
 
         self.model = None
         if effect_modifiers is None:
@@ -336,7 +349,6 @@ def estimate_unit_ate(self) -> float:
         """
         model = self._run_linear_regression()
         newline = "\n"
-        print(model.conf_int())
         treatment = [self.treatment]
         if str(self.df.dtypes[self.treatment]) == "object":
             design_info = dmatrix(self.formula.split("~")[1], self.df).design_info
@@ -372,7 +384,7 @@ def estimate_ate(self) -> tuple[float, list[float, float], float]:
         # Perform a t-test to compare the predicted outcome of the control and treated individual (ATE)
         t_test_results = model.t_test(individuals.loc["treated"] - individuals.loc["control"])
         ate = t_test_results.effect[0]
-        confidence_intervals = list(t_test_results.conf_int().flatten())
+        confidence_intervals = list(t_test_results.conf_int(alpha=self.alpha).flatten())
         return ate, confidence_intervals
 
     def estimate_control_treatment(self, adjustment_config: dict = None) -> tuple[pd.Series, pd.Series]:
@@ -434,25 +446,11 @@ def _run_linear_regression(self) -> RegressionResultsWrapper:
 
         :return: The model after fitting to data.
         """
-        # 1. Reduce dataframe to contain only the necessary columns
-        reduced_df = self.df.copy()
-        necessary_cols = [self.treatment] + list(self.adjustment_set) + [self.outcome]
-        missing_rows = reduced_df[necessary_cols].isnull().any(axis=1)
-        reduced_df = reduced_df[~missing_rows]
-        reduced_df = reduced_df.sort_values([self.treatment])
-        logger.debug(reduced_df[necessary_cols])
-
-        # 2. Add intercept
-        reduced_df["Intercept"] = 1  # self.intercept
-
-        # 3. Estimate the unit difference in outcome caused by unit difference in treatment
-        cols = [self.treatment]
-        cols += [x for x in self.adjustment_set if x not in cols]
         model = smf.ols(formula=self.formula, data=self.df).fit()
         return model
 
     def _get_confidence_intervals(self, model, treatment):
-        confidence_intervals = model.conf_int(alpha=0.05, cols=None)
+        confidence_intervals = model.conf_int(alpha=self.alpha, cols=None)
         ci_low, ci_high = (
             confidence_intervals[0].loc[treatment],
             confidence_intervals[1].loc[treatment],
@@ -519,7 +517,7 @@ def estimate_unit_ate(self, bootstrap_size=100):
         bootstraps = sorted(
             [self.estimate_coefficient(self.df.sample(len(self.df), replace=True)) for _ in range(bootstrap_size)]
         )
-        bound = ceil((bootstrap_size * 0.05) / 2)
+        bound = ceil((bootstrap_size * self.alpha) / 2)
         ci_low = bootstraps[bound]
         ci_high = bootstraps[bootstrap_size - bound]
 
@@ -610,7 +608,7 @@ def estimate_cates(self) -> pd.DataFrame:
         # Obtain CATES and confidence intervals
         conditional_ates = model.effect(effect_modifier_df, T0=self.control_value, T1=self.treatment_value).flatten()
         [ci_low, ci_high] = model.effect_interval(
-            effect_modifier_df, T0=self.control_value, T1=self.treatment_value, alpha=0.05
+            effect_modifier_df, T0=self.control_value, T1=self.treatment_value, alpha=self.alpha
         )
 
         # Merge results into a dataframe (CATE, confidence intervals, and effect modifier values)