Stop generation if all combinations are in

Author: jmafoster1

causal_testing/generation/abstract_causal_test_case.py

@@ -4,12 +4,15 @@
 import pandas as pd
 import z3
 from scipy import stats
+import itertools
 
 from causal_testing.specification.scenario import Scenario
 from causal_testing.specification.variable import Variable
 from causal_testing.testing.causal_test_case import CausalTestCase
 from causal_testing.testing.causal_test_outcome import CausalTestOutcome
 
+from enum import Enum
+
 logger = logging.getLogger(__name__)
 
 
@@ -24,21 +27,21 @@ def __init__(
         self,
         scenario: Scenario,
         intervention_constraints: set[z3.ExprRef],
-        treatment_variables: set[Variable],
+        treatment_variable: Variable,
         expected_causal_effect: dict[Variable:CausalTestOutcome],
         effect_modifiers: set[Variable] = None,
         estimate_type: str = "ate",
     ):
-        assert treatment_variables.issubset(scenario.variables.values()), (
+        assert treatment_variable in scenario.variables.values(), (
             "Treatment variables must be a subset of variables."
-            + f" Instead got:\ntreatment_variables={treatment_variables}\nvariables={scenario.variables}"
+            + f" Instead got:\ntreatment_variable={treatment_variable}\nvariables={scenario.variables}"
         )
 
         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.treatment_variable = treatment_variable
         self.expected_causal_effect = expected_causal_effect
         self.estimate_type = estimate_type
 
@@ -113,9 +116,9 @@ 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_variable]},
                 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_variable]
                 },
                 expected_causal_effect=list(self.expected_causal_effect.values())[0],
                 outcome_variables=list(self.expected_causal_effect.keys()),
@@ -208,7 +211,13 @@ def generate_concrete_tests(
                     for var in effect_modifier_configs.columns
                 }
             )
-            if target_ks_score and all((stat <= target_ks_score for stat in ks_stats.values())):
+            print("=== test ===")
+            control_values = [test.control_input_configuration[self.treatment_variable] for test in concrete_tests]
+            treatment_values = [test.treatment_input_configuration[self.treatment_variable] for test in concrete_tests]
+
+            if issubclass(self.treatment_variable.datatype, Enum) and set(zip(control_values, treatment_values)).issubset(itertools.product(self.treatment_variable.datatype, self.treatment_variable.datatype)):
+                break
+            elif 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())):

causal_testing/json_front/json_class.py

@@ -89,10 +89,12 @@ def setup(self):
         self._populate_metas()
 
     def _create_abstract_test_case(self, test, mutates, effects):
+        assert len(test["mutations"]) == 1
         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"]},
+            # TODO: Could change JSON to be treatment_var and mutation to it rather than a dict of mutations
+            treatment_variable=next(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()
@@ -121,7 +123,7 @@ def generate_tests(self, effects: dict, mutates: dict, estimators: dict, f_flag:
             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([(v.name, v.distribution) for v in [abstract_test.treatment_variable]])
             logger.info("Number of concrete tests for test case: %s", str(len(concrete_tests)))
             failures = self._execute_tests(concrete_tests, estimators, test, f_flag)
         logger.info("%s/%s failed", failures, len(concrete_tests))

causal_testing/testing/estimators.py

@@ -203,7 +203,7 @@ def estimate_ate(self) -> float:
         estimate = 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}"
+            f"Changing {self.treatment[0]} 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}"