Merge pull request #164 from CITCOM-project/examples-as-tests

Examples now runnable with pytest. Closes #5. Closes #136.

Author: jmafoster1

.github/workflows/ci-tests.yaml

@@ -25,6 +25,7 @@ jobs:
           conda install -c conda-forge pygraphviz
           python --version
           pip install -e .
+          pip install -e .[test]
           pip install pytest pytest-cov
         shell: bash -l {0}
       - name: Test with pytest

examples/.gitignore

@@ -1 +1,2 @@
 *.pdf
+results/

examples/covasim_/doubling_beta/README.md

@@ -14,7 +14,7 @@ To run this case study:
 1. Ensure all project dependencies are installed by running `pip install .` from the top
 level of this directory (instructions are provided in the project README).
 2. Change directory to `causal_testing/examples/covasim_/doubling_beta`.
-3. Run the command `python causal_test_beta.py`.
+3. Run the command `python test_beta.py`.
 
 This will print out a series of test results covering a range of different causal questions that correspond to those
-in Table 3 of the paper.
+in Table 3 of the paper.

examples/covasim_/doubling_beta/causal_test_beta.py renamed to examples/covasim_/doubling_beta/example_beta.py

@@ -19,7 +19,7 @@ To run this case study:
 level of this directory (instructions are provided in the project README).
 2. Additionally, in order to run Covasim, install version 3.0.7 by running `pip install covasim==3.0.7`.
 3. Change directory to `causal_testing/examples/covasim_/vaccinating_elderly`.
-4. Run the command `python causal_test_vaccine.py`.
+4. Run the command `python test_vaccine.py`.
 
 This will run Covasim as described above and print out the causal test results for the effect of each input on each
 output.

examples/covasim_/vaccinating_elderly/README.md

@@ -1,3 +1,4 @@
+# -*- coding: utf-8 -*-
 import pandas as pd
 import numpy as np
 import covasim as cv  # Version used in our study is 3.07
@@ -13,9 +14,17 @@
 from causal_testing.testing.estimators import LinearRegressionEstimator
 from causal_testing.testing.base_test_case import BaseTestCase
 
+import os
+import logging
 
-def experimental_causal_test_vaccinate_elderly(runs_per_test_per_config: int = 30, verbose: bool = False):
-    """ Run the causal test case for the effect of changing vaccine to prioritise elderly. This uses the experimental
+logger = logging.getLogger(__name__)
+logging.basicConfig(level=logging.DEBUG, format="%(message)s")
+
+ROOT = os.path.realpath(os.path.dirname(__file__))
+
+
+def test_experimental_vaccinate_elderly(runs_per_test_per_config: int = 30, verbose: bool = False):
+    """Run the causal test case for the effect of changing vaccine to prioritise elderly. This uses the experimental
         data collector.
 
     :param runs_per_test_per_config: Number of times to run each input configuration (control and treatment) per test.
@@ -25,78 +34,79 @@ def experimental_causal_test_vaccinate_elderly(runs_per_test_per_config: int = 3
     """
 
     # 1. Read in the Causal DAG
-    causal_dag = CausalDAG('dag.dot')
+    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)
-    vaccine = Input('vaccine', int)
-    cum_infections = Output('cum_infections', int)
-    cum_vaccinations = Output('cum_vaccinations', int)
-    cum_vaccinated = Output('cum_vaccinated', int)
-    max_doses = Output('max_doses', int)
+    pop_size = Input("pop_size", int)
+    pop_infected = Input("pop_infected", int)
+    n_days = Input("n_days", int)
+    vaccine = Input("vaccine", int)
+    cum_infections = Output("cum_infections", int)
+    cum_vaccinations = Output("cum_vaccinations", int)
+    cum_vaccinated = Output("cum_vaccinated", int)
+    max_doses = Output("max_doses", int)
 
     # 3. Create scenario by applying constraints over a subset of the input variables
-    scenario = Scenario(variables={pop_size, pop_infected, n_days, cum_infections, vaccine,
-                                   cum_vaccinated, cum_vaccinations, max_doses},
-                        constraints={pop_size.z3 == 50000, pop_infected.z3 == 1000, n_days.z3 == 50})
+    scenario = Scenario(
+        variables={
+            pop_size,
+            pop_infected,
+            n_days,
+            cum_infections,
+            vaccine,
+            cum_vaccinated,
+            cum_vaccinations,
+            max_doses,
+        },
+        constraints={pop_size.z3 == 50000, pop_infected.z3 == 1000, n_days.z3 == 50},
+    )
 
     # 4. Construct a causal specification from the scenario and causal DAG
     causal_specification = CausalSpecification(scenario, causal_dag)
 
     # 5. Instantiate the experimental data collector for Covasim
-    covasim_parameters_dict = {'pop_size': 50000,
-                               'pop_type': 'hybrid',
-                               'pop_infected': 1000,
-                               'n_days': 50}
-    control_input_configuration = {'covasim_parameters_dict': covasim_parameters_dict,
-                                   'target_elderly': False}
-    treatment_input_configuration = {'covasim_parameters_dict': covasim_parameters_dict,
-                                     'target_elderly': True}
-    data_collector = CovasimVaccineDataCollector(scenario, control_input_configuration,
-                                                 treatment_input_configuration,
-                                                 runs_per_test_per_config)
+    covasim_parameters_dict = {"pop_size": 50000, "pop_type": "hybrid", "pop_infected": 1000, "n_days": 50}
+    control_input_configuration = {"covasim_parameters_dict": covasim_parameters_dict, "target_elderly": False}
+    treatment_input_configuration = {"covasim_parameters_dict": covasim_parameters_dict, "target_elderly": True}
+    data_collector = CovasimVaccineDataCollector(
+        scenario, control_input_configuration, treatment_input_configuration, runs_per_test_per_config
+    )
 
     # 6. Express expected outcomes
-    expected_outcome_effects = {cum_infections: Positive(),
-                                cum_vaccinations: Negative(),
-                                cum_vaccinated: Negative(),
-                                max_doses: NoEffect()
-                                }
-    results_dict = {'cum_infections': {},
-                    'cum_vaccinations': {},
-                    'cum_vaccinated': {},
-                    'max_doses': {}
-                    }
+    expected_outcome_effects = {
+        cum_infections: Positive(),
+        cum_vaccinations: Negative(),
+        cum_vaccinated: Negative(),
+        max_doses: NoEffect(),
+    }
+    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)
+        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
         minimal_adjustment_set = causal_dag.identification(base_test_case)
 
         # 9. Build statistical model
-        linear_regression_estimator = LinearRegressionEstimator(vaccine.name, 1, 0,
-                                                                minimal_adjustment_set,
-                                                                outcome_variable.name)
+        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, 'ate')
+        causal_test_result = causal_test_engine.execute_test(linear_regression_estimator, causal_test_case, "ate")
         if verbose:
-            print(f"Causation:\n{causal_test_result}")
-        results_dict[outcome_variable.name]['ate'] = causal_test_result.test_value.value
-        results_dict[outcome_variable.name]['cis'] = causal_test_result.confidence_intervals
-        results_dict[outcome_variable.name]['test_passes'] = causal_test_case.expected_causal_effect.apply(
-            causal_test_result)
+            logging.info("Causation:\n%s", causal_test_result)
+        results_dict[outcome_variable.name]["ate"] = causal_test_result.test_value.value
+        results_dict[outcome_variable.name]["cis"] = causal_test_result.confidence_intervals
+        results_dict[outcome_variable.name]["test_passes"] = causal_test_case.expected_causal_effect.apply(
+            causal_test_result
+        )
     return results_dict
 
 
@@ -108,19 +118,19 @@ class CovasimVaccineDataCollector(ExperimentalDataCollector):
     """
 
     def run_system_with_input_configuration(self, input_configuration: dict) -> pd.DataFrame:
-        """ Run the system with a given input configuration.
+        """Run the system with a given input configuration.
 
         :param input_configuration: A nested dictionary containing Covasim parameters, desired number of repeats, and
         a bool to determine whether elderly should be prioritised for vaccination.
         :return: A dataframe containing results for this input configuration.
         """
-        results_df = self.simulate_vaccine(input_configuration['covasim_parameters_dict'],
-                                           self.n_repeats,
-                                           input_configuration['target_elderly'])
+        results_df = self.simulate_vaccine(
+            input_configuration["covasim_parameters_dict"], self.n_repeats, input_configuration["target_elderly"]
+        )
         return results_df
 
     def simulate_vaccine(self, pars_dict: dict, n_simulations: int = 100, target_elderly: bool = False):
-        """ Simulate observational data that contains a vaccine that is optionally given preferentially to the elderly.
+        """Simulate observational data that contains a vaccine that is optionally given preferentially to the elderly.
 
         :param pars_dict: A dictionary containing simulation parameters.
         :param n_simulations: Number of simulations to run.
@@ -129,38 +139,49 @@ def simulate_vaccine(self, pars_dict: dict, n_simulations: int = 100, target_eld
         """
         simulations_results_dfs = []
         for sim_n in range(n_simulations):
-            print(f'Simulation {sim_n + 1}/{n_simulations}.')
+            logging.info("Simulation %s/%s.", sim_n + 1, n_simulations)
 
             # Update simulation parameters with vaccine and optionally sub-target
             if target_elderly:
-                print("Prioritising the elderly for vaccination")
-                vaccine = cv.vaccinate_prob(vaccine="Pfizer", label="prioritise_elderly",
-                                            subtarget=self.vaccinate_by_age, days=list(range(7, pars_dict['n_days'])))
+                logger.info("Prioritising the elderly for vaccination")
+                vaccine = cv.vaccinate_prob(
+                    vaccine="Pfizer",
+                    label="prioritise_elderly",
+                    subtarget=self.vaccinate_by_age,
+                    days=list(range(7, pars_dict["n_days"])),
+                )
             else:
-                print("Using standard vaccination protocol")
-                vaccine = cv.vaccinate_prob(vaccine="Pfizer", label="regular", days=list(range(7, pars_dict['n_days'])))
-
-            pars_dict['interventions'] = vaccine
-            pars_dict['use_waning'] = True  # Must be set to true for vaccination
-            sim_results_df = self.run_sim_with_pars(pars_dict=pars_dict,
-                                                    desired_outputs=['cum_infections', 'cum_deaths', 'cum_recoveries',
-                                                                     'cum_vaccinations', 'cum_vaccinated'],
-                                                    n_runs=1)
-
-            sim_results_df['interventions'] = vaccine.label  # Store label in results instead of vaccine object
-            sim_results_df['target_elderly'] = target_elderly
-            sim_results_df['vaccine'] = int(target_elderly)  # 0 if standard vaccine, 1 if target elderly vaccine
-            sim_results_df['max_doses'] = vaccine.p['doses']  # Get max doses for the vaccine
+                logger.info("Using standard vaccination protocol")
+                vaccine = cv.vaccinate_prob(vaccine="Pfizer", label="regular", days=list(range(7, pars_dict["n_days"])))
+
+            pars_dict["interventions"] = vaccine
+            pars_dict["use_waning"] = True  # Must be set to true for vaccination
+            sim_results_df = self.run_sim_with_pars(
+                pars_dict=pars_dict,
+                desired_outputs=[
+                    "cum_infections",
+                    "cum_deaths",
+                    "cum_recoveries",
+                    "cum_vaccinations",
+                    "cum_vaccinated",
+                ],
+                n_runs=1,
+            )
+
+            sim_results_df["interventions"] = vaccine.label  # Store label in results instead of vaccine object
+            sim_results_df["target_elderly"] = target_elderly
+            sim_results_df["vaccine"] = int(target_elderly)  # 0 if standard vaccine, 1 if target elderly vaccine
+            sim_results_df["max_doses"] = vaccine.p["doses"]  # Get max doses for the vaccine
             simulations_results_dfs.append(sim_results_df)
 
         # Create a single dataframe containing a row for every execution
         obs_df = pd.concat(simulations_results_dfs, ignore_index=True)
-        obs_df.rename(columns={'interventions': 'vaccine_type'}, inplace=True)
+        obs_df.rename(columns={"interventions": "vaccine_type"}, inplace=True)
         return obs_df
 
     @staticmethod
     def run_sim_with_pars(pars_dict: dict, desired_outputs: [str], n_runs: int = 1, verbose: int = -1):
-        """ Runs a Covasim COVID-19 simulation with a given dict of parameters and collects the desired outputs,
+        """Runs a Covasim COVID-19 simulation with a given dict of parameters and collects the desired outputs,
             which are given as a list of output names.
 
         :param pars_dict: A dictionary containing the parameters and their values for the run.
@@ -170,17 +191,17 @@ def run_sim_with_pars(pars_dict: dict, desired_outputs: [str], n_runs: int = 1,
 
         :return results_df: A pandas df containing the results for each run
         """
-        results_dict = {k: [] for k in list(pars_dict.keys()) + desired_outputs + ['rand_seed']}
+        results_dict = {k: [] for k in list(pars_dict.keys()) + desired_outputs + ["rand_seed"]}
         for _ in range(n_runs):
             # For every run, generate and use a new a random seed.
             # This is to avoid using Covasim's sequential random seeds.
             random.seed()
             rand_seed = random.randint(0, 10000)
-            pars_dict['rand_seed'] = rand_seed
-            print(f"Rand Seed: {rand_seed}")
+            pars_dict["rand_seed"] = rand_seed
+            logger.info("Rand Seed: %s", rand_seed)
             sim = cv.Sim(pars=pars_dict)
             m_sim = cv.MultiSim(sim)
-            m_sim.run(n_runs=1, verbose=verbose, n_cpus=1)
+            m_sim.run(n_runs=1, verbose=False, n_cpus=1)
 
             for run in m_sim.sims:
                 results = run.results
@@ -191,14 +212,15 @@ def run_sim_with_pars(pars_dict: dict, desired_outputs: [str], n_runs: int = 1,
                 # Append outputs to results
                 for output in desired_outputs:
                     if output not in results:
-                        raise IndexError(f'{output} is not in the Covasim outputs.')
+                        raise IndexError(f"{output} is not in the Covasim outputs.")
                     results_dict[output].append(
-                        results[output][-1])  # Append the final recorded value for each variable
+                        results[output][-1]
+                    )  # Append the final recorded value for each variable
 
         # Any parameters without results are assigned np.nan for each execution
         for param, results in results_dict.items():
             if not results:
-                results_dict[param] = [np.nan] * len(results_dict['rand_seed'])
+                results_dict[param] = [np.nan] * len(results_dict["rand_seed"])
         return pd.DataFrame(results_dict)
 
     @staticmethod
@@ -223,5 +245,5 @@ def vaccinate_by_age(simulation):
 
 
 if __name__ == "__main__":
-    test_results = experimental_causal_test_vaccinate_elderly(runs_per_test_per_config=30, verbose=True)
-    print(test_results)
+    test_results = test_experimental_vaccinate_elderly(runs_per_test_per_config=30, verbose=True)
+    logging.info("%s", test_results)

examples/covasim_/vaccinating_elderly/causal_test_vaccine.py renamed to examples/covasim_/vaccinating_elderly/example_vaccine.py

@@ -18,7 +18,7 @@ To run this case study:
 1. Ensure all project dependencies are installed by running `pip install .` in the top level directory
    (instructions are provided in the project README).
 2. Change directory to `causal_testing/examples/lr91`.
-3. Run the command `python causal_test_max_conductances.py` or `python causal_test_max_conductances_test_suite.py`
+3. Run the command `python test_max_conductances.py` or `python test_max_conductances_test_suite.py`
 
 This should print a series of causal test results covering the effects of a range of different sized interventions made
-to the inputs on APD90, and should also plot Figure 2 from the paper.
+to the inputs on APD90, and should also plot Figure 2 from the paper.