Tests passing again

Author: jmafoster1

causal_testing/json_front/json_class.py

@@ -11,10 +11,7 @@
 from statistics import StatisticsError
 
 import pandas as pd
-import scipy
-from fitter import Fitter, get_common_distributions
 
-from causal_testing.generation.abstract_causal_test_case import AbstractCausalTestCase
 from causal_testing.specification.causal_dag import CausalDAG
 from causal_testing.specification.causal_specification import CausalSpecification
 from causal_testing.specification.scenario import Scenario
@@ -90,33 +87,6 @@ def setup(self, scenario: Scenario, ignore_cycles=False):
             )
         self._populate_metas()
 
-    def _create_abstract_test_case(self, test, mutates, effects):
-        assert len(test["mutations"]) == 1
-        treatment_var = next(self.scenario.variables[v] for v in test["mutations"])
-
-        if not treatment_var.distribution:
-            fitter = Fitter(self.df[treatment_var.name], distributions=get_common_distributions())
-            fitter.fit()
-            (dist, params) = list(fitter.get_best(method="sumsquare_error").items())[0]
-            treatment_var.distribution = getattr(scipy.stats, dist)(**params)
-            self._append_to_file(treatment_var.name + f" {dist}({params})", logging.INFO)
-
-        abstract_test = AbstractCausalTestCase(
-            scenario=self.scenario,
-            intervention_constraints=[mutates[v](k) for k, v in test["mutations"].items()],
-            treatment_variable=treatment_var,
-            expected_causal_effect={
-                self.scenario.variables[variable]: effects[effect]
-                for variable, effect in test["expected_effect"].items()
-            },
-            effect_modifiers=(
-                {self.scenario.variables[v] for v in test["effect_modifiers"]} if "effect_modifiers" in test else {}
-            ),
-            estimate_type=test["estimate_type"],
-            effect=test.get("effect", "total"),
-        )
-        return abstract_test
-
     def run_json_tests(self, effects: dict, estimators: dict, f_flag: bool = False, mutates: dict = None):
         """Runs and evaluates each test case specified in the JSON input
 
@@ -139,9 +109,7 @@ def run_json_tests(self, effects: dict, estimators: dict, f_flag: bool = False,
                         test=test, f_flag=f_flag, effects=effects, estimate_type=test["estimate_type"]
                     )
                 else:
-                    failed, msg = self._run_metamorphic_tests(
-                        test=test, f_flag=f_flag, effects=effects, mutates=mutates
-                    )
+                    raise NotImplementedError("Tried to call deprecated method _run_metamorphic_tests")
             test["failed"] = failed
             test["result"] = msg
         return self.test_plan["tests"]
@@ -189,8 +157,6 @@ def _run_concrete_metamorphic_test(self, test: dict, f_flag: bool, effects: dict
         causal_test_case = CausalTestCase(
             base_test_case=base_test_case,
             expected_causal_effect=effects[test["expected_effect"][outcome_variable]],
-            control_value=test["control_value"],
-            treatment_value=test["treatment_value"],
             estimate_type=test["estimate_type"],
         )
         failed, msg = self._execute_test_case(causal_test_case=causal_test_case, test=test, f_flag=f_flag)
@@ -205,41 +171,6 @@ def _run_concrete_metamorphic_test(self, test: dict, f_flag: bool, effects: dict
         self._append_to_file(msg, logging.INFO)
         return failed, msg
 
-    def _run_metamorphic_tests(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']}"
-        )
-        self._append_to_file(msg, logging.INFO)
-        return failures, msg
-
     def _execute_tests(self, concrete_tests, test, f_flag):
         failures = 0
         details = []
@@ -274,7 +205,8 @@ def _execute_test_case(
         failed = False
 
         estimation_model = self._setup_test(causal_test_case=causal_test_case, test=test)
-        causal_test_result = causal_test_case.execute_test(estimator=estimation_model)
+        causal_test_case.estimator = estimation_model
+        causal_test_result = causal_test_case.execute_test()
         test_passes = causal_test_case.expected_causal_effect.apply(causal_test_result)
 
         if "coverage" in test and test["coverage"]:
@@ -307,6 +239,7 @@ def _setup_test(self, causal_test_case: CausalTestCase, test: Mapping) -> Estima
                 - estimation_model - Estimator instance for the test being run
         """
         estimator_kwargs = {}
+        treatment_variable = next(self.scenario.variables[v] for v in test["mutations"])
         if "formula" in test:
             if test["estimator"] != (LinearRegressionEstimator or LogisticRegressionEstimator):
                 raise TypeError(
@@ -319,14 +252,14 @@ def _setup_test(self, causal_test_case: CausalTestCase, test: Mapping) -> Estima
             minimal_adjustment_set = self.causal_specification.causal_dag.identification(
                 causal_test_case.base_test_case
             )
-            minimal_adjustment_set = minimal_adjustment_set - {causal_test_case.treatment_variable}
+            minimal_adjustment_set = minimal_adjustment_set - {treatment_variable}
             estimator_kwargs["adjustment_set"] = minimal_adjustment_set
 
         estimator_kwargs["query"] = test["query"] if "query" in test else ""
-        estimator_kwargs["treatment"] = causal_test_case.treatment_variable.name
-        estimator_kwargs["treatment_value"] = causal_test_case.treatment_value
-        estimator_kwargs["control_value"] = causal_test_case.control_value
-        estimator_kwargs["outcome"] = causal_test_case.outcome_variable.name
+        estimator_kwargs["treatment"] = treatment_variable.name
+        estimator_kwargs["treatment_value"] = test.get("treatment_value")
+        estimator_kwargs["control_value"] = test.get("control_value")
+        estimator_kwargs["outcome"] = next(v for v in test["expected_effect"])
         estimator_kwargs["effect_modifiers"] = causal_test_case.effect_modifier_configuration
         estimator_kwargs["df"] = self.df
         estimator_kwargs["alpha"] = test["alpha"] if "alpha" in test else 0.05

causal_testing/testing/causal_test_case.py

@@ -27,8 +27,6 @@ def __init__(
         self,
         base_test_case: BaseTestCase,
         expected_causal_effect: CausalTestOutcome,
-        control_value: Any = None,
-        treatment_value: Any = None,
         estimate_type: str = "ate",
         estimate_params: dict = None,
         effect_modifier_configuration: dict[Variable:Any] = None,
@@ -37,18 +35,14 @@ def __init__(
         """
         :param base_test_case: A BaseTestCase object consisting of a treatment variable, outcome variable and effect
         :param expected_causal_effect: The expected causal effect (Positive, Negative, No Effect).
-        :param control_value: The control value for the treatment variable (before intervention).
-        :param treatment_value: The treatment value for the treatment variable (after intervention).
         :param estimate_type: A string which denotes the type of estimate to return.
         :param effect_modifier_configuration: The assignment of the effect modifiers to use for estimates.
         :param estimator: An Estimator class object
         """
         self.base_test_case = base_test_case
-        self.control_value = control_value
         self.expected_causal_effect = expected_causal_effect
         self.outcome_variable = base_test_case.outcome_variable
         self.treatment_variable = base_test_case.treatment_variable
-        self.treatment_value = treatment_value
         self.estimate_type = estimate_type
         self.estimator = estimator
         if estimate_params is None:
@@ -60,27 +54,34 @@ def __init__(
         else:
             self.effect_modifier_configuration = {}
 
-    def execute_test(self) -> CausalTestResult:
-        """Execute a causal test case and return the causal test result.
+    def execute_test(self, estimator: type(Estimator) = None) -> CausalTestResult:
+        """
+        Execute a causal test case and return the causal test result.
+        :param estimator: An alternative estimator. Defaults to `self.estimator`. This parameter is useful when you want
+        to execute a test with different data or a different equational form, but don't want to redefine the whole test
+        case.
 
         :return causal_test_result: A CausalTestResult for the executed causal test case.
         """
 
-        if not hasattr(self.estimator, f"estimate_{self.estimate_type}"):
-            raise AttributeError(f"{self.estimator.__class__} has no {self.estimate_type} method.")
-        estimate_effect = getattr(self.estimator, f"estimate_{self.estimate_type}")
+        if estimator is None:
+            estimator = self.estimator
+
+        if not hasattr(estimator, f"estimate_{self.estimate_type}"):
+            raise AttributeError(f"{estimator.__class__} has no {self.estimate_type} method.")
+        estimate_effect = getattr(estimator, f"estimate_{self.estimate_type}")
         effect, confidence_intervals = estimate_effect(**self.estimate_params)
         return CausalTestResult(
-            estimator=self.estimator,
+            estimator=estimator,
             test_value=TestValue(self.estimate_type, effect),
             effect_modifier_configuration=self.effect_modifier_configuration,
             confidence_intervals=confidence_intervals,
         )
 
     def __str__(self):
-        treatment_config = {self.treatment_variable.name: self.treatment_value}
-        control_config = {self.treatment_variable.name: self.control_value}
-        outcome_variable = {self.outcome_variable}
+        treatment_config = {self.treatment_variable.name: self.estimator.treatment_value}
+        control_config = {self.treatment_variable.name: self.estimator.control_value}
+        outcome_variable = {self.outcome_variable.name}
         return (
             f"Running {treatment_config} instead of {control_config} should cause the following "
             f"changes to {outcome_variable}: {self.expected_causal_effect}."

examples/covasim_/doubling_beta/example_beta.py

@@ -69,33 +69,37 @@ def doubling_beta_CATE_on_csv(
 
     # 6. Create a causal test case
     causal_test_case = CausalTestCase(
-        base_test_case=base_test_case, expected_causal_effect=Positive, control_value=0.016, treatment_value=0.032
-    )
-
-    linear_regression_estimator = LinearRegressionEstimator(
-        "beta",
-        0.032,
-        0.016,
-        {"avg_age", "contacts"},  # We use custom adjustment set
-        "cum_infections",
-        df=past_execution_df,
-        formula="cum_infections ~ beta + I(beta ** 2) + avg_age + contacts",
+        base_test_case=base_test_case,
+        expected_causal_effect=Positive,
+        estimator=LinearRegressionEstimator(
+            "beta",
+            0.032,
+            0.016,
+            {"avg_age", "contacts"},  # We use custom adjustment set
+            "cum_infections",
+            df=past_execution_df,
+            formula="cum_infections ~ beta + I(beta ** 2) + avg_age + contacts",
+        ),
     )
 
     # Add squared terms for beta, since it has a quadratic relationship with cumulative infections
-    causal_test_result = causal_test_case.execute_test(estimator=linear_regression_estimator)
+    causal_test_result = causal_test_case.execute_test()
 
     # Repeat for association estimate (no adjustment)
-    no_adjustment_linear_regression_estimator = LinearRegressionEstimator(
-        "beta",
-        0.032,
-        0.016,
-        set(),
-        "cum_infections",
-        df=past_execution_df,
-        formula="cum_infections ~ beta + I(beta ** 2)",
+    causal_test_case = CausalTestCase(
+        base_test_case=base_test_case,
+        expected_causal_effect=Positive,
+        estimator=LinearRegressionEstimator(
+            "beta",
+            0.032,
+            0.016,
+            set(),
+            "cum_infections",
+            df=past_execution_df,
+            formula="cum_infections ~ beta + I(beta ** 2)",
+        ),
     )
-    association_test_result = causal_test_case.execute_test(estimator=no_adjustment_linear_regression_estimator)
+    association_test_result = causal_test_case.execute_test()
 
     # Store results for plotting
     results_dict["association"] = {
@@ -116,7 +120,7 @@ def doubling_beta_CATE_on_csv(
     # Repeat causal inference after deleting all rows with treatment value to obtain counterfactual inferences
     if simulate_counterfactuals:
         counterfactual_past_execution_df = past_execution_df[past_execution_df["beta"] != 0.032]
-        counterfactual_causal_test_result = causal_test_case.execute_test(estimator=linear_regression_estimator)
+        counterfactual_causal_test_result = causal_test_case.execute_test()
 
         results_dict["counterfactual"] = {
             "ate": counterfactual_causal_test_result.test_value.value,

examples/lr91/example_max_conductances.py

@@ -129,23 +129,18 @@ def effects_on_APD90(observational_data_path, treatment_var, control_val, treatm
     causal_test_case = CausalTestCase(
         base_test_case=base_test_case,
         expected_causal_effect=expected_causal_effect,
-        control_value=control_val,
-        treatment_value=treatment_val,
-    )
-
-    # 8. Obtain the minimal adjustment set from the causal DAG
-    minimal_adjustment_set = causal_dag.identification(base_test_case)
-    linear_regression_estimator = LinearRegressionEstimator(
-        treatment_var.name,
-        treatment_val,
-        control_val,
-        minimal_adjustment_set,
-        "APD90",
-        df=pd.read_csv(observational_data_path),
+        estimator=LinearRegressionEstimator(
+            treatment=treatment_var.name,
+            treatment_value=treatment_val,
+            control_value=control_val,
+            adjustment_set=causal_dag.identification(base_test_case),
+            outcome="APD90",
+            df=pd.read_csv(observational_data_path),
+        ),
     )
 
     # 9. Run the causal test and print results
-    causal_test_result = causal_test_case.execute_test(linear_regression_estimator)
+    causal_test_result = causal_test_case.execute_test()
     logger.info("%s", causal_test_result)
     return causal_test_result.test_value.value, causal_test_result.confidence_intervals