Removed datacollector from testing

Author: jmafoster1

examples/covasim_/doubling_beta/example_beta.py

@@ -6,7 +6,6 @@
 from causal_testing.specification.scenario import Scenario
 from causal_testing.specification.variable import Input, Output
 from causal_testing.specification.causal_specification import CausalSpecification
-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.estimation.linear_regression_estimator import LinearRegressionEstimator
@@ -52,7 +51,26 @@ def doubling_beta_CATE_on_csv(
 
     # Read in the observational data, perform identification
     past_execution_df = pd.read_csv(observational_data_path)
-    data_collector, _, causal_test_case, causal_specification = setup(past_execution_df)
+
+    # 2. Create variables
+    pop_size = Input("pop_size", int)
+    pop_infected = Input("pop_infected", int)
+    n_days = Input("n_days", int)
+    cum_infections = Output("cum_infections", int)
+    cum_deaths = Output("cum_deaths", int)
+    location = Input("location", str)
+    variants = Input("variants", str)
+    avg_age = Input("avg_age", float)
+    beta = Input("beta", float)
+    contacts = Input("contacts", float)
+
+    # 5. Create a base test case
+    base_test_case = BaseTestCase(treatment_variable=beta, outcome_variable=cum_infections)
+
+    # 6. Create a causal test case
+    causal_test_case = CausalTestCase(
+        base_test_case=base_test_case, expected_causal_effect=Positive, control_value=0.016, treatment_value=0.032
+    )
 
     linear_regression_estimator = LinearRegressionEstimator(
         "beta",
@@ -98,15 +116,6 @@ def doubling_beta_CATE_on_csv(
     # Repeat causal inference after deleting all rows with treatment value to obtain counterfactual inferences
     if simulate_counterfactuals:
         counterfactual_past_execution_df = past_execution_df[past_execution_df["beta"] != 0.032]
-        counterfactual_linear_regression_estimator = LinearRegressionEstimator(
-            "beta",
-            0.032,
-            0.016,
-            {"avg_age", "contacts"},
-            "cum_infections",
-            df=counterfactual_past_execution_df,
-            formula="cum_infections ~ beta + I(beta ** 2) + avg_age + contacts",
-        )
         counterfactual_causal_test_result = causal_test_case.execute_test(estimator=linear_regression_estimator)
 
         results_dict["counterfactual"] = {
@@ -215,59 +224,6 @@ 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 setup(observational_data):
-    # 1. Read in the Causal DAG
-    causal_dag = CausalDAG(f"{ROOT}/dag.dot")
-
-    # 2. Create variables
-    pop_size = Input("pop_size", int)
-    pop_infected = Input("pop_infected", int)
-    n_days = Input("n_days", int)
-    cum_infections = Output("cum_infections", int)
-    cum_deaths = Output("cum_deaths", int)
-    location = Input("location", str)
-    variants = Input("variants", str)
-    avg_age = Input("avg_age", float)
-    beta = Input("beta", float)
-    contacts = Input("contacts", float)
-
-    # 3. Create scenario by applying constraints over a subset of the input variables
-    scenario = Scenario(
-        variables={
-            pop_size,
-            pop_infected,
-            n_days,
-            cum_infections,
-            cum_deaths,
-            location,
-            variants,
-            avg_age,
-            beta,
-            contacts,
-        },
-        constraints={pop_size.z3 == 51633, pop_infected.z3 == 1000, n_days.z3 == 216},
-    )
-
-    # 4. Construct a causal specification from the scenario and causal DAG
-    causal_specification = CausalSpecification(scenario, causal_dag)
-
-    # 5. Create a base test case
-    base_test_case = BaseTestCase(treatment_variable=beta, outcome_variable=cum_infections)
-
-    # 6. Create a causal test case
-    causal_test_case = CausalTestCase(
-        base_test_case=base_test_case, expected_causal_effect=Positive, control_value=0.016, treatment_value=0.032
-    )
-
-    # 7. Create a data collector
-    data_collector = ObservationalDataCollector(scenario, observational_data)
-
-    # 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 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):
     # Get the CATE as a percentage for association and causation
     ate = results_dict["causation"]["ate"][0]

examples/covasim_/vaccinating_elderly/example_vaccine.py

@@ -6,7 +6,6 @@
 from causal_testing.specification.scenario import Scenario
 from causal_testing.specification.variable import Input, Output
 from causal_testing.specification.causal_specification import CausalSpecification
-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.estimation.linear_regression_estimator import LinearRegressionEstimator
@@ -19,8 +18,8 @@
 
 
 def setup_test_case(verbose: bool = False):
-    """Run the causal test case for the effect of changing vaccine to prioritise elderly from an observational
-    data collector that was previously simulated.
+    """Run the causal test case for the effect of changing vaccine to prioritise elderly from observational
+    data that was previously simulated.
 
     :param verbose: Whether to print verbose details (causal test results).
     :return results_dict: A dictionary containing ATE, 95% CIs, and Test Pass/Fail
@@ -57,11 +56,9 @@ def setup_test_case(verbose: bool = False):
     # 4. Construct a causal specification from the scenario and causal DAG
     causal_specification = CausalSpecification(scenario, causal_dag)
 
-    # 5. Instantiate the observational data collector using the previously simulated data
+    # 5. Read the previously simulated data
     obs_df = pd.read_csv("simulated_data.csv")
 
-    data_collector = ObservationalDataCollector(scenario, obs_df)
-
     # 6. Express expected outcomes
     expected_outcome_effects = {
         cum_infections: Positive(),
@@ -90,7 +87,7 @@ def setup_test_case(verbose: bool = False):
         )
 
         # 9. Execute test and save results in dict
-        causal_test_result = causal_test_case.execute_test(linear_regression_estimator, data_collector)
+        causal_test_result = causal_test_case.execute_test(linear_regression_estimator, obs_df)
 
         if verbose:
             logging.info("Causation:\n%s", causal_test_result)
@@ -110,4 +107,4 @@ def setup_test_case(verbose: bool = False):
 
     test_results = setup_test_case(verbose=True)
 
-    logging.info("%s", test_results)
+    logging.info("%s", test_results)