Revert changes to data checks in observational data collector

Author: christopher-wild

causal_testing/data_collection/data_collector.py

@@ -18,8 +18,6 @@ class DataCollector(ABC):
 
     def __init__(self, scenario: Scenario):
         self.scenario = scenario
-        self.data_checked = False
-
     @abstractmethod
     def collect_data(self, **kwargs) -> pd.DataFrame:
         """
@@ -158,5 +156,4 @@ def collect_data(self, **kwargs) -> pd.DataFrame:
         for var_name, var in self.scenario.variables.items():
             if issubclass(var.datatype, Enum):
                 scenario_execution_data_df[var_name] = [var.datatype(x) for x in scenario_execution_data_df[var_name]]
-        self.data_checked = True
-        self.data = scenario_execution_data_df
+        return scenario_execution_data_df

causal_testing/testing/causal_test_case.py

@@ -81,17 +81,15 @@ 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: ObservationalDataCollector, causal_specification: CausalSpecification) -> CausalTestResult:
+    def execute_test(self, estimator: type(Estimator), data_collector: ObservationalDataCollector) -> CausalTestResult:
         """Execute a causal test case and return the causal test result.
 
         :param estimator: A reference to an Estimator class.
         :param causal_test_case: The CausalTestCase object to be tested
         :return causal_test_result: A CausalTestResult for the executed causal test case.
         """
-        if not data_collector.data_checked:
-            data_collector.collect_data()
         if estimator.df is None:
-            estimator.df = data_collector.data
+            estimator.df = data_collector.collect_data()
         treatment_variable = self.treatment_variable
         treatments = treatment_variable.name
         outcome_variable = self.outcome_variable

examples/poisson-line-process/example_poisson_process.py

@@ -119,7 +119,7 @@ def causal_test_intensity_num_shapes(
         )
 
     # 9. Execute the test
-    causal_test_result = causal_test_case.execute_test(estimator, data_collector, causal_specification)
+    causal_test_result = causal_test_case.execute_test(estimator, data_collector)
 
     return causal_test_result
 

tests/data_collection_tests/test_observational_data_collector.py

@@ -44,15 +44,13 @@ def test_not_all_variables_in_data(self):
     def test_all_variables_in_data(self):
         scenario = Scenario({self.X1, self.X2, self.X3, self.Y1, self.Y2})
         observational_data_collector = ObservationalDataCollector(scenario, self.observational_df)
-        observational_data_collector.collect_data(index_col=0)
-        df = observational_data_collector.data
+        df = observational_data_collector.collect_data(index_col=0)
         assert df.equals(self.observational_df), f"\n{df}\nwas not equal to\n{self.observational_df}"
 
     def test_data_constraints(self):
         scenario = Scenario({self.X1, self.X2, self.X3, self.Y1, self.Y2}, {self.X1.z3 > 2})
         observational_data_collector = ObservationalDataCollector(scenario, self.observational_df)
-        observational_data_collector.collect_data(index_col=0)
-        df = observational_data_collector.data
+        df = observational_data_collector.collect_data(index_col=0)
         expected = self.observational_df.loc[[2, 3]]
         assert df.equals(expected), f"\n{df}\nwas not equal to\n{expected}"
 
@@ -64,7 +62,7 @@ def populate_m(data):
         scenario = Scenario({self.X1, meta})
         observational_data_collector = ObservationalDataCollector(scenario, self.observational_df)
         observational_data_collector.collect_data()
-        data = observational_data_collector.data
+        data = observational_data_collector.collect_data()
         assert all((m == 2 * x1 for x1, m in zip(data["X1"], data["M"])))
 
     def tearDown(self) -> None: