Examples now use causal_test_case.execute_test

Author: christopher-wild

examples/covasim_/doubling_beta/example_beta.py

@@ -8,7 +8,6 @@
 from causal_testing.data_collection.data_collector import ObservationalDataCollector
 from causal_testing.testing.causal_test_case import CausalTestCase
 from causal_testing.testing.causal_test_outcome import Positive
-from causal_testing.testing.causal_test_engine import CausalTestEngine
 from causal_testing.testing.estimators import LinearRegressionEstimator
 from causal_testing.testing.base_test_case import BaseTestCase
 from matplotlib.pyplot import rcParams
@@ -36,7 +35,7 @@
 
 
 def doubling_beta_CATE_on_csv(
-    observational_data_path: str, simulate_counterfactuals: bool = False, verbose: bool = False
+        observational_data_path: str, simulate_counterfactuals: bool = False, verbose: bool = False
 ):
     """Compute the CATE of increasing beta from 0.016 to 0.032 on cum_infections using the dataframe
     loaded from the specified path. Additionally simulate the counterfactuals by repeating the analysis
@@ -50,9 +49,9 @@ def doubling_beta_CATE_on_csv(
     """
     results_dict = {"association": {}, "causation": {}}
 
-    # Read in the observational data, perform identification, and setup the causal_test_engine
+    # Read in the observational data, perform identification
     past_execution_df = pd.read_csv(observational_data_path)
-    _, causal_test_engine, causal_test_case = engine_setup(observational_data_path)
+    data_collector, _, causal_test_case, causal_specification = setup(observational_data_path)
 
     linear_regression_estimator = LinearRegressionEstimator(
         "beta",
@@ -65,7 +64,9 @@ def doubling_beta_CATE_on_csv(
     )
 
     # Add squared terms for beta, since it has a quadratic relationship with cumulative infections
-    causal_test_result = causal_test_engine.execute_test(linear_regression_estimator, causal_test_case)
+    causal_test_result = causal_test_case.execute_test(estimator=linear_regression_estimator,
+                                                       data_collector=data_collector,
+                                                       causal_specification=causal_specification)
 
     # Repeat for association estimate (no adjustment)
     no_adjustment_linear_regression_estimator = LinearRegressionEstimator(
@@ -77,9 +78,9 @@ def doubling_beta_CATE_on_csv(
         df=past_execution_df,
         formula="cum_infections ~ beta + np.power(beta, 2)",
     )
-    association_test_result = causal_test_engine.execute_test(
-        no_adjustment_linear_regression_estimator, causal_test_case
-    )
+    association_test_result = causal_test_case.execute_test(estimator=no_adjustment_linear_regression_estimator,
+                                                            data_collector=data_collector,
+                                                            causal_specification=causal_specification)
 
     # Store results for plotting
     results_dict["association"] = {
@@ -109,8 +110,9 @@ def doubling_beta_CATE_on_csv(
             df=counterfactual_past_execution_df,
             formula="cum_infections ~ beta + np.power(beta, 2) + avg_age + contacts",
         )
-        counterfactual_causal_test_result = causal_test_engine.execute_test(
-            linear_regression_estimator, causal_test_case
+        counterfactual_causal_test_result = causal_test_case.execute_test(
+            estimator=linear_regression_estimator, data_collector=data_collector,
+            causal_specification=causal_specification
         )
         results_dict["counterfactual"] = {
             "ate": counterfactual_causal_test_result.test_value.value,
@@ -218,7 +220,7 @@ def doubling_beta_CATEs(observational_data_path: str, simulate_counterfactual: b
     age_contact_fig.savefig(outpath_base_str + "age_contact_executions.pdf", format="pdf")
 
 
-def engine_setup(observational_data_path):
+def setup(observational_data_path):
     # 1. Read in the Causal DAG
     causal_dag = CausalDAG(f"{ROOT}/dag.dot")
 
@@ -265,13 +267,10 @@ def engine_setup(observational_data_path):
     # 7. Create a data collector
     data_collector = ObservationalDataCollector(scenario, pd.read_csv(observational_data_path))
 
-    # 8. Create an instance of the causal test engine
-    causal_test_engine = CausalTestEngine(causal_specification, data_collector)
-
-    # 9. Obtain the minimal adjustment set for the base test case from the causal DAG
+    # 8. Obtain the minimal adjustment set for the base test case from the causal DAG
     minimal_adjustment_set = causal_dag.identification(base_test_case)
 
-    return minimal_adjustment_set, causal_test_engine, causal_test_case
+    return data_collector, minimal_adjustment_set, causal_test_case, causal_specification
 
 
 def plot_doubling_beta_CATEs(results_dict, title, figure=None, axes=None, row=None, col=None):

examples/covasim_/vaccinating_elderly/example_vaccine.py

@@ -10,7 +10,6 @@
 from causal_testing.data_collection.data_collector import ExperimentalDataCollector
 from causal_testing.testing.causal_test_case import CausalTestCase
 from causal_testing.testing.causal_test_outcome import Positive, Negative, NoEffect
-from causal_testing.testing.causal_test_engine import CausalTestEngine
 from causal_testing.testing.estimators import LinearRegressionEstimator
 from causal_testing.testing.base_test_case import BaseTestCase
 
@@ -81,25 +80,23 @@ def test_experimental_vaccinate_elderly(runs_per_test_per_config: int = 30, verb
     }
     results_dict = {"cum_infections": {}, "cum_vaccinations": {}, "cum_vaccinated": {}, "max_doses": {}}
 
-    # 7. Create an instance of the causal test engine
-    causal_test_engine = CausalTestEngine(causal_specification, data_collector, index_col=0)
-
     for outcome_variable, expected_effect in expected_outcome_effects.items():
         base_test_case = BaseTestCase(treatment_variable=vaccine, outcome_variable=outcome_variable)
         causal_test_case = CausalTestCase(
             base_test_case=base_test_case, expected_causal_effect=expected_effect, control_value=0, treatment_value=1
         )
 
-        # 8. Obtain the minimal adjustment set for the causal test case from the causal DAG
+        # 7. Obtain the minimal adjustment set for the causal test case from the causal DAG
         minimal_adjustment_set = causal_dag.identification(base_test_case)
 
-        # 9. Build statistical model
+        # 8. Build statistical model
         linear_regression_estimator = LinearRegressionEstimator(
             vaccine.name, 1, 0, minimal_adjustment_set, outcome_variable.name
         )
 
-        # 10. Execute test and save results in dict
-        causal_test_result = causal_test_engine.execute_test(linear_regression_estimator, causal_test_case)
+        # 9. Execute test and save results in dict
+        causal_test_result = causal_test_case.execute_test(linear_regression_estimator, data_collector,
+                                                           causal_specification)
         if verbose:
             logging.info("Causation:\n%s", causal_test_result)
         results_dict[outcome_variable.name]["ate"] = causal_test_result.test_value.value

examples/lr91/example_max_conductances.py

@@ -8,7 +8,6 @@
 from causal_testing.data_collection.data_collector import ObservationalDataCollector
 from causal_testing.testing.causal_test_case import CausalTestCase
 from causal_testing.testing.causal_test_outcome import Positive, Negative, NoEffect
-from causal_testing.testing.causal_test_engine import CausalTestEngine
 from causal_testing.testing.estimators import LinearRegressionEstimator
 from causal_testing.testing.base_test_case import BaseTestCase
 from matplotlib.pyplot import rcParams
@@ -138,17 +137,14 @@ def effects_on_APD90(observational_data_path, treatment_var, control_val, treatm
     # 7. Create a data collector
     data_collector = ObservationalDataCollector(scenario, pd.read_csv(observational_data_path))
 
-    # 8. Create an instance of the causal test engine
-    causal_test_engine = CausalTestEngine(causal_specification, data_collector)
-
-    # 9. Obtain the minimal adjustment set from the causal DAG
+    # 8. Obtain the minimal adjustment set from the causal DAG
     minimal_adjustment_set = causal_dag.identification(base_test_case)
     linear_regression_estimator = LinearRegressionEstimator(
         treatment_var.name, treatment_val, control_val, minimal_adjustment_set, "APD90"
     )
 
-    # 10. Run the causal test and print results
-    causal_test_result = causal_test_engine.execute_test(linear_regression_estimator, causal_test_case)
+    # 9. Run the causal test and print results
+    causal_test_result = causal_test_case.execute_test(linear_regression_estimator, data_collector, causal_specification)
     logger.info("%s", causal_test_result)
     return causal_test_result.test_value.value, causal_test_result.confidence_intervals
 
@@ -198,4 +194,4 @@ def normalise_data(df, columns=None):
 
 
 if __name__ == "__main__":
-    test_sensitivity_analysis(show=True)
+    test_sensitivity_analysis()

examples/lr91/example_max_conductances_test_suite.py

@@ -8,7 +8,6 @@
 from causal_testing.data_collection.data_collector import ObservationalDataCollector
 from causal_testing.testing.causal_test_case import CausalTestCase
 from causal_testing.testing.causal_test_outcome import Positive, Negative, NoEffect
-from causal_testing.testing.causal_test_engine import CausalTestEngine
 from causal_testing.testing.estimators import LinearRegressionEstimator
 from causal_testing.testing.base_test_case import BaseTestCase
 from causal_testing.testing.causal_test_suite import CausalTestSuite
@@ -147,11 +146,9 @@ def effects_on_APD90(observational_data_path, test_suite):
     # 7. Create a data collector
     data_collector = ObservationalDataCollector(scenario, pd.read_csv(observational_data_path))
 
-    # 8. Create an instance of the causal test engine
-    causal_test_engine = CausalTestEngine(causal_specification, data_collector)
 
-    # 9. Run the causal test suite
-    causal_test_results = causal_test_engine.execute_test_suite(test_suite)
+    # 8. Run the causal test suite
+    causal_test_results = test_suite.execute_test_suite(data_collector, causal_specification)
     return causal_test_results
 
 
@@ -200,4 +197,4 @@ def normalise_data(df, columns=None):
 
 
 if __name__ == "__main__":
-    test_sensitivity_analysis(show=True, save=True)
+    test_sensitivity_analysis()

examples/poisson-line-process/example_poisson_process.py

@@ -5,7 +5,6 @@
 from causal_testing.data_collection.data_collector import ObservationalDataCollector
 from causal_testing.testing.causal_test_case import CausalTestCase
 from causal_testing.testing.causal_test_outcome import ExactValue, Positive
-from causal_testing.testing.causal_test_engine import CausalTestEngine
 from causal_testing.testing.estimators import LinearRegressionEstimator, Estimator
 from causal_testing.testing.base_test_case import BaseTestCase
 
@@ -85,13 +84,10 @@ def causal_test_intensity_num_shapes(
     # 6. Create a data collector
     data_collector = ObservationalDataCollector(scenario, pd.read_csv(observational_data_path))
 
-    # 7. Create an instance of the causal test engine
-    causal_test_engine = CausalTestEngine(causal_specification, data_collector)
-
-    # 8. Obtain the minimal adjustment set for the causal test case from the causal DAG
+    # 7. Obtain the minimal adjustment set for the causal test case from the causal DAG
     minimal_adjustment_set = causal_dag.identification(causal_test_case.base_test_case)
 
-    # 9. Set up an estimator
+    # 8. Set up an estimator
     data = pd.read_csv(observational_data_path)
 
     treatment = causal_test_case.get_treatment_variable()
@@ -122,8 +118,8 @@ def causal_test_intensity_num_shapes(
             formula=f"{outcome} ~ {treatment} + {'+'.join(square_terms + inverse_terms + list([e for e in causal_test_case.effect_modifier_configuration]))} -1",
         )
 
-    # 10. Execute the test
-    causal_test_result = causal_test_engine.execute_test(estimator, causal_test_case)
+    # 9. Execute the test
+    causal_test_result = causal_test_case.execute_test(estimator, data_collector, causal_specification)
 
     return causal_test_result