Merge pull request #219 from CITCOM-project/test-engine-refactor

Test engine refactor

Author: christopher-wild

causal_testing/data_collection/data_collector.py

@@ -149,7 +149,6 @@ def collect_data(self, **kwargs) -> pd.DataFrame:
 
         :return: A pandas dataframe containing execution data that is valid for the scenario-under-test.
         """
-
         execution_data_df = self.data
         for meta in self.scenario.metas():
             if meta.name not in self.data:

causal_testing/json_front/json_class.py

@@ -5,7 +5,7 @@
 import json
 import logging
 
-from collections.abc import Iterable, Mapping
+from collections.abc import Mapping
 from dataclasses import dataclass
 from pathlib import Path
 from statistics import StatisticsError
@@ -22,7 +22,6 @@
 from causal_testing.specification.variable import Input, Meta, Output
 from causal_testing.testing.causal_test_case import CausalTestCase
 from causal_testing.testing.causal_test_result import CausalTestResult
-from causal_testing.testing.causal_test_engine import CausalTestEngine
 from causal_testing.testing.estimators import Estimator
 from causal_testing.testing.base_test_case import BaseTestCase
 
@@ -49,12 +48,12 @@ class JsonUtility:
     def __init__(self, output_path: str, output_overwrite: bool = False):
         self.input_paths = None
         self.variables = {"inputs": {}, "outputs": {}, "metas": {}}
-        self.data = []
         self.test_plan = None
         self.scenario = None
         self.causal_specification = None
         self.output_path = Path(output_path)
         self.check_file_exists(self.output_path, output_overwrite)
+        self.data_collector = None
 
     def set_paths(self, json_path: str, dag_path: str, data_paths: list[str] = None):
         """
@@ -69,6 +68,7 @@ def set_paths(self, json_path: str, dag_path: str, data_paths: list[str] = None)
 
     def setup(self, scenario: Scenario):
         """Function to populate all the necessary parts of the json_class needed to execute tests"""
+        data = []
         self.scenario = scenario
         self._get_scenario_variables()
         self.scenario.setup_treatment_variables()
@@ -80,20 +80,21 @@ def setup(self, scenario: Scenario):
             self.test_plan = json.load(f)
         # Populate the data
         if self.input_paths.data_paths:
-            self.data = pd.concat([pd.read_csv(data_file, header=0) for data_file in self.input_paths.data_paths])
-        if len(self.data) == 0:
+            data = pd.concat([pd.read_csv(data_file, header=0) for data_file in self.input_paths.data_paths])
+        if len(data) == 0:
             raise ValueError(
                 "No data found, either provide a path to a file containing data or manually populate the .data "
                 "attribute with a dataframe before calling .setup()"
             )
+        self.data_collector = ObservationalDataCollector(self.scenario, data)
         self._populate_metas()
 
     def _create_abstract_test_case(self, test, mutates, effects):
         assert len(test["mutations"]) == 1
         treatment_var = next(self.scenario.variables[v] for v in test["mutations"])
 
         if not treatment_var.distribution:
-            fitter = Fitter(self.data[treatment_var.name], distributions=get_common_distributions())
+            fitter = Fitter(self.data_collector.data[treatment_var.name], distributions=get_common_distributions())
             fitter.fit()
             (dist, params) = list(fitter.get_best(method="sumsquare_error").items())[0]
             treatment_var.distribution = getattr(scipy.stats, dist)(**params)
@@ -149,6 +150,7 @@ def run_json_tests(self, effects: dict, estimators: dict, f_flag: bool = False,
                     treatment_value=test["treatment_value"],
                     estimate_type=test["estimate_type"],
                 )
+
                 failed, _ = self._execute_test_case(causal_test_case=causal_test_case, test=test, f_flag=f_flag)
 
                 msg = (
@@ -231,6 +233,7 @@ def _execute_tests(self, concrete_tests, test, f_flag):
         details = []
         if "formula" in test:
             self._append_to_file(f"Estimator formula used for test: {test['formula']}")
+
         for concrete_test in concrete_tests:
             failed, result = self._execute_test_case(concrete_test, test, f_flag)
             details.append(result)
@@ -243,10 +246,10 @@ def _populate_metas(self):
         Populate data with meta-variable values and add distributions to Causal Testing Framework Variables
         """
         for meta in self.scenario.variables_of_type(Meta):
-            meta.populate(self.data)
+            meta.populate(self.data_collector.data)
 
     def _execute_test_case(
-        self, causal_test_case: CausalTestCase, test: Iterable[Mapping], f_flag: bool
+        self, causal_test_case: CausalTestCase, test: Mapping, f_flag: bool
     ) -> (bool, CausalTestResult):
         """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
@@ -258,10 +261,10 @@ def _execute_test_case(
         """
         failed = False
 
-        causal_test_engine, estimation_model = self._setup_test(
-            causal_test_case, test, test["conditions"] if "conditions" in test else None
+        estimation_model = self._setup_test(causal_test_case=causal_test_case, test=test)
+        causal_test_result = causal_test_case.execute_test(
+            estimator=estimation_model, data_collector=self.data_collector
         )
-        causal_test_result = causal_test_engine.execute_test(estimation_model, causal_test_case)
 
         test_passes = causal_test_case.expected_causal_effect.apply(causal_test_result)
 
@@ -283,9 +286,7 @@ def _execute_test_case(
             logger.warning("   FAILED- expected %s, got %s", causal_test_case.expected_causal_effect, result_string)
         return failed, causal_test_result
 
-    def _setup_test(
-        self, causal_test_case: CausalTestCase, test: Mapping, conditions: list[str] = None
-    ) -> tuple[CausalTestEngine, Estimator]:
+    def _setup_test(self, causal_test_case: CausalTestCase, test: Mapping) -> Estimator:
         """Create the necessary inputs for a single test case
         :param causal_test_case: The concrete test case to be executed
         :param test: Single JSON test definition stored in a mapping (dict)
@@ -296,12 +297,6 @@ def _setup_test(
                 - causal_test_engine - Test Engine instance for the test being run
                 - estimation_model - Estimator instance for the test being run
         """
-
-        data_collector = ObservationalDataCollector(
-            self.scenario, self.data.query(" & ".join(conditions)) if conditions else self.data
-        )
-        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_var = causal_test_case.treatment_variable
         minimal_adjustment_set = minimal_adjustment_set - {treatment_var}
@@ -311,14 +306,13 @@ def _setup_test(
             "control_value": causal_test_case.control_value,
             "adjustment_set": minimal_adjustment_set,
             "outcome": causal_test_case.outcome_variable.name,
-            "df": causal_test_engine.scenario_execution_data_df,
             "effect_modifiers": causal_test_case.effect_modifier_configuration,
             "alpha": test["alpha"] if "alpha" in test else 0.05,
         }
         if "formula" in test:
             estimator_kwargs["formula"] = test["formula"]
         estimation_model = test["estimator"](**estimator_kwargs)
-        return causal_test_engine, estimation_model
+        return estimation_model
 
     def _append_to_file(self, line: str, log_level: int = None):
         """Appends given line(s) to the current output file. If log_level is specified it also logs that message to the

causal_testing/specification/causal_specification.py

@@ -1,6 +1,5 @@
 """This module holds the abstract CausalSpecification data class, which holds a Scenario and CausalDag"""
 
-from abc import ABC
 from dataclasses import dataclass
 from typing import Union
 
@@ -11,9 +10,9 @@
 
 
 @dataclass
-class CausalSpecification(ABC):
+class CausalSpecification:
     """
-    Abstract Class for the Causal Specification (combination of Scenario and Causal Dag)
+    Data class storing the Causal Specification (combination of Scenario and Causal Dag)
     """
 
     scenario: Scenario

causal_testing/testing/causal_test_case.py

@@ -5,6 +5,10 @@
 from causal_testing.specification.variable import Variable
 from causal_testing.testing.causal_test_outcome import CausalTestOutcome
 from causal_testing.testing.base_test_case import BaseTestCase
+from causal_testing.testing.estimators import Estimator
+from causal_testing.testing.causal_test_result import CausalTestResult, TestValue
+from causal_testing.data_collection.data_collector import DataCollector
+
 
 logger = logging.getLogger(__name__)
 
@@ -73,6 +77,41 @@ def get_treatment_value(self):
         """Return the treatment value of the treatment variable in this causal test case."""
         return self.treatment_value
 
+    def execute_test(self, estimator: type(Estimator), data_collector: DataCollector) -> CausalTestResult:
+        """Execute a causal test case and return the causal test result.
+
+        :param estimator: A reference to an Estimator class.
+        :param data_collector: The data collector to be used which provides a dataframe for the Estimator
+        :return causal_test_result: A CausalTestResult for the executed causal test case.
+        """
+        if estimator.df is None:
+            estimator.df = data_collector.collect_data()
+
+        logger.info("treatments: %s", self.treatment_variable.name)
+        logger.info("outcomes: %s", self.outcome_variable)
+
+        causal_test_result = self._return_causal_test_results(estimator)
+        return causal_test_result
+
+    def _return_causal_test_results(self, estimator) -> CausalTestResult:
+        """Depending on the estimator used, calculate the 95% confidence intervals and return in a causal_test_result
+
+        :param estimator: An Estimator class object
+        :return: a CausalTestResult object containing the confidence intervals
+        """
+        if not hasattr(estimator, f"estimate_{self.estimate_type}"):
+            raise AttributeError(f"{estimator.__class__} has no {self.estimate_type} method.")
+        estimate_effect = getattr(estimator, f"estimate_{self.estimate_type}")
+        effect, confidence_intervals = estimate_effect(**self.estimate_params)
+        causal_test_result = CausalTestResult(
+            estimator=estimator,
+            test_value=TestValue(self.estimate_type, effect),
+            effect_modifier_configuration=self.effect_modifier_configuration,
+            confidence_intervals=confidence_intervals,
+        )
+
+        return causal_test_result
+
     def __str__(self):
         treatment_config = {self.treatment_variable.name: self.treatment_value}
         control_config = {self.treatment_variable.name: self.control_value}