covasim examples update

Author: christopher-wild

examples/covasim_/doubling_beta/causal_test_beta.py

@@ -10,6 +10,7 @@
 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
 
 # Uncommenting the code below will make all graphs publication quality but requires a suitable latex installation
@@ -203,22 +204,26 @@ def engine_setup(observational_data_path):
     # 4. Construct a causal specification from the scenario and causal DAG
     causal_specification = CausalSpecification(scenario, causal_dag)
 
-    # 5. Create a causal test case
-    causal_test_case = CausalTestCase(control_input_configuration={beta: 0.016},
+    # 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,
-                                      treatment_input_configuration={beta: 0.032},
-                                      outcome_variables={cum_infections})
+                                      control_value=0.016,
+                                      treatment_value=0.032)
 
-    # 6. Create a data collector
+    # 7. Create a data collector
     data_collector = ObservationalDataCollector(scenario, observational_data_path)
 
-    # 7. Create an instance of the causal test engine
+    # 8. 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
-    causal_test_engine.identification(causal_test_case)
+    # 9. Obtain the minimal adjustment set for the base test case from the causal DAG
+    minimal_adjustment_set = causal_dag.identification(base_test_case)
 
-    return causal_test_engine.minimal_adjustment_set, causal_test_engine, causal_test_case
+    return minimal_adjustment_set, causal_test_engine, causal_test_case
 
 
 def plot_doubling_beta_CATEs(results_dict, title, figure=None, axes=None, row=None, col=None):

examples/covasim_/vaccinating_elderly/causal_test_vaccine.py

@@ -11,6 +11,7 @@
 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
 
 
 def experimental_causal_test_vaccinate_elderly(runs_per_test_per_config: int = 30, verbose: bool = False):
@@ -71,21 +72,21 @@ def experimental_causal_test_vaccinate_elderly(runs_per_test_per_config: int = 3
 
     # 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():
-        causal_test_case = CausalTestCase(control_input_configuration={vaccine: 0},
+        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,
-                                          treatment_input_configuration={vaccine: 1},
-                                          outcome_variables={outcome_variable})
-
-
-
+                                          control_value=0,
+                                          treatment_value=1)
 
         # 8. Obtain the minimal adjustment set for the causal test case from the causal DAG
-        causal_test_engine.identification(causal_test_case)
+        minimal_adjustment_set = causal_dag.identification(base_test_case)
 
         # 9. Build statistical model
         linear_regression_estimator = LinearRegressionEstimator((vaccine.name,), 1, 0,
-                                                                causal_test_engine.minimal_adjustment_set,
+                                                                minimal_adjustment_set,
                                                                 (outcome_variable.name,))
 
         # 10. Execute test and save results in dict

examples/lr91/causal_test_max_conductances_test_suite.py

@@ -72,20 +72,19 @@ def causal_testing_sensitivity_analysis():
         for treatment_value in treatment_values:
             test_list.append(CausalTestCase(base_test_case, oracle, control_value, treatment_value))
         test_suite.add_test_object(base_test_case=base_test_case,
-                                   test_list=test_list,
-                                   estimators=estimator_list,
+                                   causal_test_case_list=test_list,
+                                   estimators_classes=estimator_list,
                                    estimate_type='ate')
 
     causal_test_results = effects_on_APD90(OBSERVATIONAL_DATA_PATH, test_suite)
 
     # Extract data from causal_test_results needed for plotting
     for base_test_case in causal_test_results:
-
         # Place results of test_suite into format required for plotting
         results[base_test_case.treatment_variable.name] = \
             {"ate": [result.ate for result in causal_test_results[base_test_case]['LinearRegressionEstimator']],
              "cis": [result.confidence_intervals for result in
-                     causal_test_results[base_test_case][0]]}
+                     causal_test_results[base_test_case]['LinearRegressionEstimator']]}
 
     plot_ates_with_cis(results, treatment_values)