black

Author: christopher-wild

causal_testing/__init__.py

@@ -1,4 +1,5 @@
 import logging
+
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.INFO)
 logger.addHandler(logging.StreamHandler())

causal_testing/data_collection/data_collector.py

@@ -21,7 +21,9 @@ 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.
@@ -37,7 +39,9 @@ def filter_valid_data(self, data: pd.DataFrame, check_pos: bool = True) -> pd.Da
 
         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()
@@ -72,8 +76,10 @@ def filter_valid_data(self, data: pd.DataFrame, check_pos: bool = True) -> pd.Da
         # How many rows did we drop?
         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}")
+            logger.warning(
+                f"Discarded {size_diff}/{len(data)} values due to constraint violations.\n"
+                f"For example{unsat_core}"
+            )
         return satisfying_data
 
 
@@ -83,8 +89,13 @@ class ExperimentalDataCollector(DataCollector):
     Users should implement these methods to collect data from their system.
     """
 
-    def __init__(self, scenario: Scenario, control_input_configuration: dict, treatment_input_configuration: dict,
-                 n_repeats: int = 1):
+    def __init__(
+        self,
+        scenario: Scenario,
+        control_input_configuration: dict,
+        treatment_input_configuration: dict,
+        n_repeats: int = 1,
+    ):
         super().__init__(scenario)
         self.control_input_configuration = control_input_configuration
         self.treatment_input_configuration = treatment_input_configuration
@@ -97,13 +108,21 @@ 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

@@ -25,45 +25,52 @@ def __init__(
         scenario: Scenario,
         intervention_constraints: set[z3.ExprRef],
         treatment_variables: set[Variable],
-        expected_causal_effect: dict[Variable: CausalTestOutcome],
+        expected_causal_effect: dict[Variable:CausalTestOutcome],
         effect_modifiers: set[Variable] = None,
-        estimate_type: str = "ate"
+        estimate_type: str = "ate",
     ):
         assert treatment_variables.issubset(scenario.variables.values()), (
             "Treatment variables must be a subset of variables."
             + 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
         self.treatment_variables = treatment_variables
         self.expected_causal_effect = expected_causal_effect
-        self.estimate_type=estimate_type
+        self.estimate_type = estimate_type
 
         if effect_modifiers is not None:
             self.effect_modifiers = effect_modifiers
         else:
             self.effect_modifiers = {}
 
     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()])
-        return (
-            f"When we apply intervention {self.intervention_constraints}, {outcome_string}"
+        outcome_string = " and ".join(
+            [
+                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}"
 
     def datapath(self):
         def sanitise(string):
             return "".join([x for x in string if x.isalnum()])
 
         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"
         )
 
-
     def _generate_concrete_tests(
         self, sample_size: int, rct: bool = False, seed: int = 0
     ) -> tuple[list[CausalTestCase], pd.DataFrame]:
@@ -80,8 +87,9 @@ 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)
@@ -103,7 +111,9 @@ 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,9 +132,9 @@ def _generate_concrete_tests(
                 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 = {
+                effect_modifier_configuration={
                     v: v.cast(model[v.z3]) for v in self.effect_modifiers
-                }
+                },
             )
 
             for v in self.scenario.inputs():
@@ -160,9 +170,13 @@ def _generate_concrete_tests(
 
         return concrete_tests, pd.DataFrame(runs, columns=run_columns + ["bin"])
 
-
     def generate_concrete_tests(
-        self, sample_size: int, target_ks_score: float = None, rct: bool = False, seed: int = 0, hard_max: int = 1000
+        self,
+        sample_size: int,
+        target_ks_score: float = None,
+        rct: bool = False,
+        seed: int = 0,
+        hard_max: int = 1000,
     ) -> tuple[list[CausalTestCase], pd.DataFrame]:
         """Generates a list of `num` concrete test cases.
 
@@ -189,14 +203,22 @@ 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])
-            ks_stats = {var: stats.kstest(control_configs[var], var.distribution.cdf).statistic for var in control_configs.columns}
+            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
+            }
             # Putting treatment and control values in messes it up because the two are not independent...
             # This is potentially problematic as constraints might mean we don't get good coverage if we use control values alone
             # We might then need to carefully craft our _control value_ generating distributions so that we can get good coverage
@@ -205,11 +227,29 @@ 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])
-            ks_stats.update({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())):
+            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
+                    for var in effect_modifier_configs.columns
+                }
+            )
+            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())):
-            logger.error("Hard max of %s reached but could not achieve target ks_score of %s. Got %s.", hard_max, target_ks_score, ks_stats)
+        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,
+                target_ks_score,
+                ks_stats,
+            )
         return concrete_tests, runs