Moved effect type to causal test case rather than causal test engine

Author: jmafoster1

causal_testing/testing/causal_test_case.py

@@ -18,7 +18,7 @@ class CausalTestCase:
 
     def __init__(self, control_input_configuration: dict[Variable: Any], expected_causal_effect: CausalTestOutcome,
                  outcome_variables: dict[Variable], treatment_input_configuration: dict[Variable: Any] = None,
-                 estimate_type: str = "ate", effect_modifier_configuration: dict[Variable: Any] = None):
+                 estimate_type: str = "ate", effect_modifier_configuration: dict[Variable: Any] = None, effect: str = "total"):
         """
         When a CausalTestCase is initialised, it takes the intervention and applies it to the input configuration to
         create two distinct input configurations: a control input configuration and a treatment input configuration.
@@ -35,6 +35,7 @@ def __init__(self, control_input_configuration: dict[Variable: Any], expected_ca
         self.outcome_variables = outcome_variables
         self.treatment_input_configuration = treatment_input_configuration
         self.estimate_type = estimate_type
+        self.effect = effect
         if effect_modifier_configuration:
             self.effect_modifier_configuration = effect_modifier_configuration
         else:

causal_testing/testing/causal_test_engine.py

@@ -30,13 +30,12 @@ class CausalTestEngine:
     """
 
     def __init__(self, causal_test_case: CausalTestCase, causal_specification: CausalSpecification,
-                 data_collector: DataCollector, effect: str = "total"):
+                 data_collector: DataCollector):
         self.causal_test_case = causal_test_case
         self.treatment_variables = list(self.causal_test_case.control_input_configuration)
         self.casual_dag, self.scenario = causal_specification.causal_dag, causal_specification.scenario
         self.data_collector = data_collector
         self.scenario_execution_data_df = pd.DataFrame()
-        self.effect = effect
 
     def load_data(self, **kwargs):
         """ Load execution data corresponding to the causal test case into a pandas dataframe and return the minimal
@@ -61,12 +60,12 @@ def load_data(self, **kwargs):
         self.scenario_execution_data_df = self.data_collector.collect_data(**kwargs)
 
         minimal_adjustment_sets = []
-        if self.effect == "total":
+        if self.causal_test_case.effect == "total":
             minimal_adjustment_sets = self.casual_dag.enumerate_minimal_adjustment_sets(
                     [v.name for v in self.treatment_variables],
                     [v.name for v in self.causal_test_case.outcome_variables]
                 )
-        elif self.effect == "direct":
+        elif self.causal_test_case.effect == "direct":
             minimal_adjustment_sets = self.casual_dag.direct_effect_adjustment_sets(
                     [v.name for v in self.treatment_variables],
                     [v.name for v in self.causal_test_case.outcome_variables]

tests/testing_tests/test_causal_test_engine.py

@@ -145,12 +145,17 @@ def test_execute_test_observational_causal_forest_estimator(self):
     def test_invalid_causal_effect(self):
         """ Check that executing the causal test case returns the correct results for dummy data using a linear
         regression estimator. """
+        causal_test_case = CausalTestCase(
+            control_input_configuration={self.A: 0},
+            expected_causal_effect=self.expected_causal_effect,
+            treatment_input_configuration={self.A: 1},
+            outcome_variables={self.C},
+            effect="error")
         # 5. Create causal test engine
         causal_test_engine = CausalTestEngine(
-            self.causal_test_case,
+            causal_test_case,
             self.causal_specification,
-            self.data_collector,
-            effect="error"
+            self.data_collector
         )
         with self.assertRaises(Exception):
             causal_test_engine.load_data()
@@ -172,12 +177,18 @@ def test_execute_test_observational_linear_regression_estimator(self):
     def test_execute_test_observational_linear_regression_estimator_direct_effect(self):
         """ Check that executing the causal test case returns the correct results for dummy data using a linear
         regression estimator. """
+        causal_test_case = CausalTestCase(
+            control_input_configuration={self.A: 0},
+            expected_causal_effect=self.expected_causal_effect,
+            treatment_input_configuration={self.A: 1},
+            outcome_variables={self.C},
+            effect="direct")
+
         # 5. Create causal test engine
         causal_test_engine = CausalTestEngine(
-            self.causal_test_case,
+            causal_test_case,
             self.causal_specification,
-            self.data_collector,
-            effect="direct"
+            self.data_collector
         )
         self.minimal_adjustment_set = causal_test_engine.load_data()