Merge pull request #118 from CITCOM-project/causal_test_case_refactor

Causal test case refactor

Author: christopher-wild

causal_testing/generation/abstract_causal_test_case.py

@@ -19,7 +19,7 @@
 
 class AbstractCausalTestCase:
     """
-    An abstract test case serves as a generator for concrete test cases. Instead of having concrete conctrol
+    An abstract test case serves as a generator for concrete test cases. Instead of having concrete control
     and treatment values, we instead just specify the intervention and the treatment variables. This then
     enables potentially infinite concrete test cases to be generated between different values of the treatment.
     """
@@ -157,7 +157,7 @@ def _generate_concrete_tests(
                 # 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({concrete_test.treatment_variable.name: concrete_test.treatment_value})
                     treatment_run["bin"] = index
                     runs.append(treatment_run)
 
@@ -204,7 +204,7 @@ def generate_concrete_tests(
             runs = pd.concat([runs, runs_])
             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.treatment_variable: test.control_value} for test in concrete_tests])
             ks_stats = {
                 var: stats.kstest(control_configs[var], var.distribution.cdf).statistic
                 for var in control_configs.columns
@@ -227,8 +227,8 @@ def generate_concrete_tests(
                     for var in effect_modifier_configs.columns
                 }
             )
-            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]
+            control_values = [test.control_value for test in concrete_tests]
+            treatment_values = [test.treatment_value for test in concrete_tests]
 
             if self.treatment_variable.datatype is bool and set([(True, False), (False, True)]).issubset(
                 set(zip(control_values, treatment_values))

causal_testing/json_front/json_class.py

@@ -202,12 +202,12 @@ def _setup_test(self, causal_test_case: CausalTestCase, estimator: Estimator) ->
         data_collector = ObservationalDataCollector(self.modelling_scenario, self.data_path)
         causal_test_engine = CausalTestEngine(self.causal_specification, data_collector, index_col=0)
         minimal_adjustment_set = self.causal_specification.causal_dag.identification(causal_test_case.base_test_case)
-        treatment_vars = list(causal_test_case.treatment_input_configuration)
-        minimal_adjustment_set = minimal_adjustment_set - {v.name for v in treatment_vars}
+        treatment_var = causal_test_case.treatment_variable
+        minimal_adjustment_set = minimal_adjustment_set - {treatment_var}
         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],
+            (treatment_var.name,),
+            causal_test_case.treatment_value,
+            causal_test_case.control_value,
             minimal_adjustment_set,
             (causal_test_case.outcome_variable.name,),
             causal_test_engine.scenario_execution_data_df,

causal_testing/specification/causal_dag.py

@@ -471,7 +471,7 @@ def depends_on_outputs(self, node: Node, scenario: Scenario) -> bool:
     def identification(self, base_test_case):
         """Identify and return the minimum adjustment set
 
-        :param base_test_case: A base test case class instance containing the outcome_variable and the
+        :param base_test_case: A base test case instance containing the outcome_variable and the
         treatment_variable required for identification.
         :return minimal_adjustment_set: The smallest set of variables which can be adjusted for to obtain a causal
         estimate as opposed to a purely associational estimate.

causal_testing/testing/causal_test_case.py

@@ -9,11 +9,15 @@
 
 
 class CausalTestCase:
-    """)
-    A causal test case is a triple (X, Delta, Y), where X is an input configuration, Delta is an intervention, and
-    Y is the expected causal effect on a particular output. The goal of a causal test case is to test whether the
-    intervention Delta made to the input configuration X causes the model-under-test to produce the expected change
-    in Y.
+    """
+    A CausalTestCase extends the information held in a BaseTestCase. As well as storing the treatment and outcome
+    variables, a CausalTestCase stores the values of these variables. Also the outcome variable and value are
+    specified.
+
+    The goal of a CausalTestCase is to test whether the intervention made to the control via the treatment causes the
+    model-under-test to produce the expected change. The CausalTestCase structure is designed for execution using the
+    CausalTestEngine, using either execute_test() function to execute a single test case or packing CausalTestCases into
+    a CausalTestSuite and executing them as a batch using the execute_test_suite() function.
     """
 
     def __init__(
@@ -26,50 +30,46 @@ def __init__(
         effect_modifier_configuration: dict[Variable:Any] = None,
     ):
         """
-        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.
-        The former is the input configuration before applying the intervention and the latter is the input configuration
-        after applying the intervention.
-
-        :param control_input_configuration: The input configuration representing the control values of the treatment
-        variables.
-        :param treatment_input_configuration: The input configuration representing the treatment values of the treatment
-        variables. That is, the input configuration *after* applying the intervention.
+        :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:
         """
         self.base_test_case = base_test_case
-        self.control_input_configuration = {base_test_case.treatment_variable: control_value}
+        self.control_value = control_value
         self.expected_causal_effect = expected_causal_effect
         self.outcome_variable = base_test_case.outcome_variable
-        self.treatment_input_configuration = {base_test_case.treatment_variable: treatment_value}
+        self.treatment_variable = base_test_case.treatment_variable
+        self.treatment_value = treatment_value
         self.estimate_type = estimate_type
         self.effect = base_test_case.effect
+
         if effect_modifier_configuration:
             self.effect_modifier_configuration = effect_modifier_configuration
         else:
             self.effect_modifier_configuration = dict()
-        assert (
-            self.control_input_configuration.keys() == self.treatment_input_configuration.keys()
-        ), "Control and treatment input configurations must have the same keys."
 
-    def get_treatment_variables(self):
-        """Return a list of the treatment variables (as strings) for this causal test case."""
-        return [v.name for v in self.control_input_configuration]
+    def get_treatment_variable(self):
+        """Return the treatment variable name (as string) for this causal test case"""
+        return self.treatment_variable.name
 
-    def get_outcome_variables(self):
-        """Return a list of the outcome variables (as strings) for this causal test case."""
-        return [self.outcome_variable.name]
+    def get_outcome_variable(self):
+        """Return the outcome variable name (as string) for this causal test case."""
+        return self.outcome_variable.name
 
-    def get_control_values(self):
-        """Return a list of the control values for each treatment variable in this causal test case."""
-        return list(self.control_input_configuration.values())
+    def get_control_value(self):
+        """Return a the control value of the treatment variable in this causal test case."""
+        return self.control_value
 
-    def get_treatment_values(self):
-        """Return a list of the treatment values for each treatment variable in this causal test case."""
-        return list(self.treatment_input_configuration.values())
+    def get_treatment_value(self):
+        """Return the treatment value of the treatment variable in this causal test case."""
+        return self.treatment_value
 
     def __str__(self):
-        treatment_config = {k.name: v for k, v in self.treatment_input_configuration.items()}
-        control_config = {k.name: v for k, v in self.control_input_configuration.items()}
+        treatment_config = {self.treatment_variable.name: self.treatment_value}
+        control_config = {self.treatment_variable.name: self.control_value}
         outcome_variable = {self.outcome_variable}
         return (
             f"Running {treatment_config} instead of {control_config} should cause the following "

causal_testing/testing/causal_test_engine.py

@@ -88,9 +88,9 @@ def execute_test_suite(self, test_suite: CausalTestSuite) -> list[CausalTestResu
                 causal_test_results = []
 
                 for test in tests:
-                    treatment_variable = list(test.treatment_input_configuration.keys())[0]
-                    treatment_value = list(test.treatment_input_configuration.values())[0]
-                    control_value = list(test.control_input_configuration.values())[0]
+                    treatment_variable = test.treatment_variable
+                    treatment_value = test.treatment_value
+                    control_value = test.control_value
                     estimator = EstimatorClass(
                         (treatment_variable.name,),
                         treatment_value,
@@ -131,19 +131,16 @@ def execute_test(
             raise Exception("No data has been loaded. Please call load_data prior to executing a causal test case.")
         if estimator.df is None:
             estimator.df = self.scenario_execution_data_df
-        treatment_variable = list(causal_test_case.control_input_configuration.keys())[0]
+        treatment_variable = causal_test_case.treatment_variable
         treatments = treatment_variable.name
         outcome_variable = causal_test_case.outcome_variable
 
         logger.info("treatments: %s", treatments)
         logger.info("outcomes: %s", outcome_variable)
         minimal_adjustment_set = self.causal_dag.identification(BaseTestCase(treatment_variable, outcome_variable))
-        minimal_adjustment_set = minimal_adjustment_set - {v.name for v in causal_test_case.control_input_configuration}
+        minimal_adjustment_set = minimal_adjustment_set - set(treatment_variable.name)
         minimal_adjustment_set = minimal_adjustment_set - set(outcome_variable.name)
-        assert all(
-            (v.name not in minimal_adjustment_set for v in causal_test_case.control_input_configuration)
-        ), "Treatment vars in adjustment set"
-        assert outcome_variable not in minimal_adjustment_set, "Outcome vars in adjustment set"
+
         variables_for_positivity = list(minimal_adjustment_set) + [treatment_variable.name] + [outcome_variable.name]
 
         if self._check_positivity_violation(variables_for_positivity):
@@ -172,8 +169,8 @@ def _return_causal_test_results(self, estimate_type, estimator, causal_test_case
                 causal_test_result = CausalTestResult(
                     treatment=estimator.treatment,
                     outcome=estimator.outcome,
-                    treatment_value=estimator.treatment_values,
-                    control_value=estimator.control_values,
+                    treatment_value=estimator.treatment_value,
+                    control_value=estimator.control_value,
                     adjustment_set=estimator.adjustment_set,
                     test_value=TestValue("ate", cates_df),
                     effect_modifier_configuration=causal_test_case.effect_modifier_configuration,
@@ -185,8 +182,8 @@ def _return_causal_test_results(self, estimate_type, estimator, causal_test_case
             causal_test_result = CausalTestResult(
                 treatment=estimator.treatment,
                 outcome=estimator.outcome,
-                treatment_value=estimator.treatment_values,
-                control_value=estimator.control_values,
+                treatment_value=estimator.treatment_value,
+                control_value=estimator.control_value,
                 adjustment_set=estimator.adjustment_set,
                 test_value=TestValue("risk_ratio", risk_ratio),
                 effect_modifier_configuration=causal_test_case.effect_modifier_configuration,
@@ -198,8 +195,8 @@ def _return_causal_test_results(self, estimate_type, estimator, causal_test_case
             causal_test_result = CausalTestResult(
                 treatment=estimator.treatment,
                 outcome=estimator.outcome,
-                treatment_value=estimator.treatment_values,
-                control_value=estimator.control_values,
+                treatment_value=estimator.treatment_value,
+                control_value=estimator.control_value,
                 adjustment_set=estimator.adjustment_set,
                 test_value=TestValue("ate", ate),
                 effect_modifier_configuration=causal_test_case.effect_modifier_configuration,
@@ -213,8 +210,8 @@ def _return_causal_test_results(self, estimate_type, estimator, causal_test_case
             causal_test_result = CausalTestResult(
                 treatment=estimator.treatment,
                 outcome=estimator.outcome,
-                treatment_value=estimator.treatment_values,
-                control_value=estimator.control_values,
+                treatment_value=estimator.treatment_value,
+                control_value=estimator.control_value,
                 adjustment_set=estimator.adjustment_set,
                 test_value=TestValue("ate", ate),
                 effect_modifier_configuration=causal_test_case.effect_modifier_configuration,

causal_testing/testing/causal_test_suite.py

@@ -1,4 +1,8 @@
 from collections import UserDict
+from typing import Type, Iterable
+from causal_testing.testing.base_test_case import BaseTestCase
+from causal_testing.testing.causal_test_case import CausalTestCase
+from causal_testing.testing.estimators import Estimator
 
 
 class CausalTestSuite(UserDict):
@@ -13,15 +17,21 @@ class CausalTestSuite(UserDict):
     base_test_case's and execute each causal_test_case with each iterator.
     """
 
-    def add_test_object(self, base_test_case, causal_test_case_list, estimators, estimate_type: str = "ate"):
+    def add_test_object(
+        self,
+        base_test_case: BaseTestCase,
+        causal_test_case_list: Iterable[CausalTestCase],
+        estimators_classes: Iterable[Type[Estimator]],
+        estimate_type: str = "ate",
+    ):
         """
         A setter object to allow for the easy construction of the dictionary test suite structure
 
         :param base_test_case: A BaseTestCase object consisting of a treatment variable, outcome variable and effect
         :param causal_test_case_list: A list of causal test cases to be executed
-        :param estimators: A list of estimators, the execute_test_suite function in the TestEngine will produce a list
+        :param estimators_classes: A list of estimator class references, the execute_test_suite function in the TestEngine will produce a list
                             of test results for each estimator
         :param estimate_type: A string which denotes the type of estimate to return
         """
-        test_object = {"tests": causal_test_case_list, "estimators": estimators, "estimate_type": estimate_type}
+        test_object = {"tests": causal_test_case_list, "estimators": estimators_classes, "estimate_type": estimate_type}
         self.data[base_test_case] = test_object