black

Author: christopher-wild

causal_testing/data_collection/data_collector.py

@@ -21,9 +21,7 @@ def collect_data(self, **kwargs) -> pd.DataFrame:
         """
         pass
 
-    def filter_valid_data(
-        self, data: pd.DataFrame, check_pos: bool = True
-    ) -> pd.DataFrame:
+    def filter_valid_data(self, data: pd.DataFrame, check_pos: bool = True) -> pd.DataFrame:
         """Check is execution data is valid for the scenario-under-test.
 
         Data is invalid if it does not meet the constraints specified in the scenario-under-test.
@@ -39,9 +37,7 @@ def filter_valid_data(
 
         if check_pos and not scenario_variables.issubset(data.columns):
             missing_variables = scenario_variables - set(data.columns)
-            raise IndexError(
-                f"Positivity violation: missing data for variables {missing_variables}."
-            )
+            raise IndexError(f"Positivity violation: missing data for variables {missing_variables}.")
 
         # For each row, does it satisfy the constraints?
         solver = z3.Solver()
@@ -54,9 +50,7 @@ def filter_valid_data(
             # Need to explicitly cast variables to their specified type. Z3 will not take e.g. np.int64 to be an int.
             model = [
                 self.scenario.variables[var].z3
-                == self.scenario.variables[var].z3_val(
-                    self.scenario.variables[var].z3, row[var]
-                )
+                == self.scenario.variables[var].z3_val(self.scenario.variables[var].z3, row[var])
                 for var in self.scenario.variables
             ]
             for c in model:
@@ -77,8 +71,7 @@ def filter_valid_data(
         size_diff = len(data) - len(satisfying_data)
         if size_diff > 0:
             logger.warning(
-                f"Discarded {size_diff}/{len(data)} values due to constraint violations.\n"
-                f"For example{unsat_core}"
+                f"Discarded {size_diff}/{len(data)} values due to constraint violations.\n" f"For example{unsat_core}"
             )
         return satisfying_data
 
@@ -108,21 +101,13 @@ def collect_data(self, **kwargs) -> pd.DataFrame:
         :return: A pandas dataframe containing execution data for the system-under-test in both control and treatment
         executions.
         """
-        control_results_df = self.run_system_with_input_configuration(
-            self.control_input_configuration
-        )
-        treatment_results_df = self.run_system_with_input_configuration(
-            self.treatment_input_configuration
-        )
-        results_df = pd.concat(
-            [control_results_df, treatment_results_df], ignore_index=True
-        )
+        control_results_df = self.run_system_with_input_configuration(self.control_input_configuration)
+        treatment_results_df = self.run_system_with_input_configuration(self.treatment_input_configuration)
+        results_df = pd.concat([control_results_df, treatment_results_df], ignore_index=True)
         return results_df
 
     @abstractmethod
-    def run_system_with_input_configuration(
-        self, input_configuration: dict
-    ) -> pd.DataFrame:
+    def run_system_with_input_configuration(self, input_configuration: dict) -> pd.DataFrame:
         """Run the system with a given input configuration and return the resulting execution data.
 
         :param input_configuration: A dictionary which maps a subset of inputs to values.

causal_testing/generation/abstract_causal_test_case.py

@@ -34,9 +34,7 @@ def __init__(
             + f" Instead got:\ntreatment_variables={treatment_variables}\nvariables={scenario.variables}"
         )
 
-        assert (
-            len(expected_causal_effect) == 1
-        ), "We currently only support tests with one causal outcome"
+        assert len(expected_causal_effect) == 1, "We currently only support tests with one causal outcome"
 
         self.scenario = scenario
         self.intervention_constraints = intervention_constraints
@@ -51,10 +49,7 @@ def __init__(
 
     def __str__(self):
         outcome_string = " and ".join(
-            [
-                f"the effect on {var} should be {str(effect)}"
-                for var, effect in self.expected_causal_effect.items()
-            ]
+            [f"the effect on {var} should be {str(effect)}" for var, effect in self.expected_causal_effect.items()]
         )
         return f"When we apply intervention {self.intervention_constraints}, {outcome_string}"
 
@@ -65,9 +60,7 @@ def sanitise(string):
         return (
             sanitise("-".join([str(c) for c in self.intervention_constraints]))
             + "_"
-            + "-".join(
-                [f"{v.name}_{e}" for v, e in self.expected_causal_effect.items()]
-            )
+            + "-".join([f"{v.name}_{e}" for v, e in self.expected_causal_effect.items()])
             + ".csv"
         )
 
@@ -87,9 +80,7 @@ def _generate_concrete_tests(
 
         concrete_tests = []
         runs = []
-        run_columns = sorted(
-            [v.name for v in self.scenario.variables.values() if v.distribution]
-        )
+        run_columns = sorted([v.name for v in self.scenario.variables.values() if v.distribution])
 
         # Generate the Latin Hypercube samples and put into a dataframe
         # lhsmdu.setRandomSeed(seed+i)
@@ -100,9 +91,7 @@ def _generate_concrete_tests(
         # Project the samples to the variables' distributions
         for name in run_columns:
             var = self.scenario.variables[name]
-            samples[var.name] = lhsmdu.inverseTransformSample(
-                var.distribution, samples[var.name]
-            )
+            samples[var.name] = lhsmdu.inverseTransformSample(var.distribution, samples[var.name])
 
         for index, row in samples.iterrows():
             optimizer = z3.Optimize()
@@ -111,9 +100,7 @@ def _generate_concrete_tests(
             for c in self.intervention_constraints:
                 optimizer.assert_and_track(c, str(c))
 
-            optimizer.add_soft(
-                [self.scenario.variables[v].z3 == row[v] for v in run_columns]
-            )
+            optimizer.add_soft([self.scenario.variables[v].z3 == row[v] for v in run_columns])
             if optimizer.check() == z3.unsat:
                 logger.warning(
                     "Satisfiability of test case was unsat.\n"
@@ -122,26 +109,19 @@ def _generate_concrete_tests(
             model = optimizer.model()
 
             concrete_test = CausalTestCase(
-                control_input_configuration={
-                    v: v.cast(model[v.z3]) for v in self.treatment_variables
-                },
+                control_input_configuration={v: v.cast(model[v.z3]) for v in self.treatment_variables},
                 treatment_input_configuration={
-                    v: v.cast(model[self.scenario.treatment_variables[v.name].z3])
-                    for v in self.treatment_variables
+                    v: v.cast(model[self.scenario.treatment_variables[v.name].z3]) for v in self.treatment_variables
                 },
                 expected_causal_effect=list(self.expected_causal_effect.values())[0],
                 outcome_variables=list(self.expected_causal_effect.keys()),
                 estimate_type=self.estimate_type,
-                effect_modifier_configuration={
-                    v: v.cast(model[v.z3]) for v in self.effect_modifiers
-                },
+                effect_modifier_configuration={v: v.cast(model[v.z3]) for v in self.effect_modifiers},
             )
 
             for v in self.scenario.inputs():
                 if row[v.name] != v.cast(model[v.z3]):
-                    constraints = "\n  ".join(
-                        [str(c) for c in self.scenario.constraints if v.name in str(c)]
-                    )
+                    constraints = "\n  ".join([str(c) for c in self.scenario.constraints if v.name in str(c)])
                     logger.warning(
                         f"Unable to set variable {v.name} to {row[v.name]} because of constraints\n"
                         + f"{constraints}\nUsing value {v.cast(model[v.z3])} instead in test\n{concrete_test}"
@@ -150,21 +130,14 @@ def _generate_concrete_tests(
             concrete_tests.append(concrete_test)
             # Control run
             control_run = {
-                v.name: v.cast(model[v.z3])
-                for v in self.scenario.variables.values()
-                if v.name in run_columns
+                v.name: v.cast(model[v.z3]) for v in self.scenario.variables.values() if v.name in run_columns
             }
             control_run["bin"] = index
             runs.append(control_run)
             # Treatment run
             if rct:
                 treatment_run = control_run.copy()
-                treatment_run.update(
-                    {
-                        k.name: v
-                        for k, v in concrete_test.treatment_input_configuration.items()
-                    }
-                )
+                treatment_run.update({k.name: v for k, v in concrete_test.treatment_input_configuration.items()})
                 treatment_run["bin"] = index
                 runs.append(treatment_run)
 
@@ -203,18 +176,12 @@ def generate_concrete_tests(
         ks_stats = []
 
         for i in range(hard_max):
-            concrete_tests_, runs_ = self._generate_concrete_tests(
-                sample_size, rct, seed + i
-            )
+            concrete_tests_, runs_ = self._generate_concrete_tests(sample_size, rct, seed + i)
             concrete_tests += concrete_tests_
             runs = pd.concat([runs, runs_])
-            assert (
-                concrete_tests_ not in concrete_tests
-            ), "Duplicate entries unlikely unless something went wrong"
+            assert concrete_tests_ not in concrete_tests, "Duplicate entries unlikely unless something went wrong"
 
-            control_configs = pd.DataFrame(
-                [test.control_input_configuration for test in concrete_tests]
-            )
+            control_configs = pd.DataFrame([test.control_input_configuration for test in concrete_tests])
             ks_stats = {
                 var: stats.kstest(control_configs[var], var.distribution.cdf).statistic
                 for var in control_configs.columns
@@ -227,25 +194,17 @@ def generate_concrete_tests(
             # treatment_configs = pd.DataFrame([test.treatment_input_configuration for test in concrete_tests])
             # both_configs = pd.concat([control_configs, treatment_configs])
             # ks_stats = {var: stats.kstest(both_configs[var], var.distribution.cdf).statistic for var in both_configs.columns}
-            effect_modifier_configs = pd.DataFrame(
-                [test.effect_modifier_configuration for test in concrete_tests]
-            )
+            effect_modifier_configs = pd.DataFrame([test.effect_modifier_configuration for test in concrete_tests])
             ks_stats.update(
                 {
-                    var: stats.kstest(
-                        effect_modifier_configs[var], var.distribution.cdf
-                    ).statistic
+                    var: stats.kstest(effect_modifier_configs[var], var.distribution.cdf).statistic
                     for var in effect_modifier_configs.columns
                 }
             )
-            if target_ks_score and all(
-                (stat <= target_ks_score for stat in ks_stats.values())
-            ):
+            if target_ks_score and all((stat <= target_ks_score for stat in ks_stats.values())):
                 break
 
-        if target_ks_score is not None and not all(
-            (stat <= target_ks_score for stat in ks_stats.values())
-        ):
+        if target_ks_score is not None and not all((stat <= target_ks_score for stat in ks_stats.values())):
             logger.error(
                 "Hard max of %s reached but could not achieve target ks_score of %s. Got %s.",
                 hard_max,

causal_testing/json_front/json_class.py

@@ -74,27 +74,19 @@ def set_variables(self, inputs: dict, outputs: dict, metas: dict):
         """
         self.inputs = [Input(i["name"], i["type"], i["distribution"]) for i in inputs]
         self.outputs = [Output(i["name"], i["type"]) for i in outputs]
-        self.metas = (
-            [Meta(i["name"], i["type"], i["populate"]) for i in metas]
-            if metas
-            else []
-        )
+        self.metas = [Meta(i["name"], i["type"], i["populate"]) for i in metas] if metas else []
 
     def setup(self):
         """Function to populate all the necessary parts of the json_class needed to execute tests"""
-        self.modelling_scenario = Scenario(
-            self.inputs + self.outputs + self.metas, None
-        )
+        self.modelling_scenario = Scenario(self.inputs + self.outputs + self.metas, None)
         self.modelling_scenario.setup_treatment_variables()
         self.causal_specification = CausalSpecification(
             scenario=self.modelling_scenario, causal_dag=CausalDAG(self.dag_path)
         )
         self._json_parse()
         self._populate_metas()
 
-    def execute_tests(
-        self, effects: dict, mutates: dict, estimators: dict, f_flag: bool
-    ):
+    def execute_tests(self, effects: dict, mutates: dict, estimators: dict, f_flag: bool):
         """Runs and evaluates each test case specified in the JSON input
 
         :param effects: Dictionary mapping effect class instances to string representations.
@@ -110,20 +102,13 @@ def execute_tests(
 
             abstract_test = AbstractCausalTestCase(
                 scenario=self.modelling_scenario,
-                intervention_constraints=[
-                    mutates[v](k) for k, v in test["mutations"].items()
-                ],
-                treatment_variables={
-                    self.modelling_scenario.variables[v] for v in test["mutations"]
-                },
+                intervention_constraints=[mutates[v](k) for k, v in test["mutations"].items()],
+                treatment_variables={self.modelling_scenario.variables[v] for v in test["mutations"]},
                 expected_causal_effect={
                     self.modelling_scenario.variables[variable]: effects[effect]
                     for variable, effect in test["expectedEffect"].items()
                 },
-                effect_modifiers={
-                    self.modelling_scenario.variables[v]
-                    for v in test["effect_modifiers"]
-                }
+                effect_modifiers={self.modelling_scenario.variables[v] for v in test["effect_modifiers"]}
                 if "effect_modifiers" in test
                 else {},
                 estimate_type=test["estimate_type"],
@@ -132,17 +117,11 @@ def execute_tests(
             concrete_tests, dummy = abstract_test.generate_concrete_tests(5, 0.05)
             logger.info("Executing test: %s", test["name"])
             logger.info(abstract_test)
-            logger.info(
-                [(v.name, v.distribution) for v in abstract_test.treatment_variables]
-            )
-            logger.info(
-                "Number of concrete tests for test case: %s", str(len(concrete_tests))
-            )
+            logger.info([(v.name, v.distribution) for v in abstract_test.treatment_variables])
+            logger.info("Number of concrete tests for test case: %s", str(len(concrete_tests)))
             for concrete_test in concrete_tests:
                 executed_tests += 1
-                failed = self._execute_test_case(
-                    concrete_test, estimators[test["estimator"]], f_flag
-                )
+                failed = self._execute_test_case(concrete_test, estimators[test["estimator"]], f_flag)
                 if failed:
                     failures += 1
 
@@ -170,9 +149,7 @@ def _populate_metas(self):
             var.distribution = getattr(scipy.stats, dist)(**params)
             logger.info(var.name + f"{dist}({params})")
 
-    def _execute_test_case(
-        self, causal_test_case: CausalTestCase, estimator: Estimator, f_flag: bool
-    ) -> bool:
+    def _execute_test_case(self, causal_test_case: CausalTestCase, estimator: Estimator, f_flag: bool) -> bool:
         """Executes a singular test case, prints the results and returns the test case result
         :param causal_test_case: The concrete test case to be executed
         :param f_flag: Failure flag that if True the script will stop executing when a test fails.
@@ -181,9 +158,7 @@ def _execute_test_case(
         """
         failed = False
 
-        causal_test_engine, estimation_model = self._setup_test(
-            causal_test_case, estimator
-        )
+        causal_test_engine, estimation_model = self._setup_test(causal_test_case, estimator)
         causal_test_result = causal_test_engine.execute_test(
             estimation_model, estimate_type=causal_test_case.estimate_type
         )
@@ -192,7 +167,9 @@ def _execute_test_case(
 
         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.ate} <  {causal_test_result.ci_high()}"
+            result_string = (
+                f"{causal_test_result.ci_low()} < {causal_test_result.ate} <  {causal_test_result.ci_high()}"
+            )
         else:
             result_string = causal_test_result.ate
         if f_flag:
@@ -209,34 +186,22 @@ def _execute_test_case(
             )
         return failed
 
-    def _setup_test(
-        self, causal_test_case: CausalTestCase, estimator: Estimator
-    ) -> tuple[CausalTestEngine, Estimator]:
+    def _setup_test(self, causal_test_case: CausalTestCase, estimator: Estimator) -> tuple[CausalTestEngine, Estimator]:
         """Create the necessary inputs for a single test case
         :param causal_test_case: The concrete test case to be executed
         :returns:
                 - causal_test_engine - Test Engine instance for the test being run
                 - estimation_model - Estimator instance for the test being run
         """
-        data_collector = ObservationalDataCollector(
-            self.modelling_scenario, self.data_path
-        )
-        causal_test_engine = CausalTestEngine(
-            causal_test_case, self.causal_specification, data_collector
-        )
+        data_collector = ObservationalDataCollector(self.modelling_scenario, self.data_path)
+        causal_test_engine = CausalTestEngine(causal_test_case, self.causal_specification, data_collector)
         minimal_adjustment_set = causal_test_engine.load_data(index_col=0)
         treatment_vars = list(causal_test_case.treatment_input_configuration)
-        minimal_adjustment_set = minimal_adjustment_set - {
-            v.name for v in treatment_vars
-        }
+        minimal_adjustment_set = minimal_adjustment_set - {v.name for v in treatment_vars}
         estimation_model = estimator(
             (list(treatment_vars)[0].name,),
-            [causal_test_case.treatment_input_configuration[v] for v in treatment_vars][
-                0
-            ],
-            [causal_test_case.control_input_configuration[v] for v in treatment_vars][
-                0
-            ],
+            [causal_test_case.treatment_input_configuration[v] for v in treatment_vars][0],
+            [causal_test_case.control_input_configuration[v] for v in treatment_vars][0],
             minimal_adjustment_set,
             (list(causal_test_case.outcome_variables)[0].name,),
             causal_test_engine.scenario_execution_data_df,