Repository: CITCOM-project/CausalTestingFramework (language: python)

Issue:
Regression test examples
With large refactors of the framework happening, would be nice to regression test examples to see if they are stable


Discussion / hints:

Fix patch:
diff --git a/examples/.gitignore b/examples/.gitignore
index a1363379..041589cf 100644
--- a/examples/.gitignore
+++ b/examples/.gitignore
@@ -1,1 +1,2 @@
 *.pdf
+results/
diff --git a/examples/covasim_/doubling_beta/README.md b/examples/covasim_/doubling_beta/README.md
index 66cc2325..d8ac0ba2 100644
--- a/examples/covasim_/doubling_beta/README.md
+++ b/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.
\ No newline at end of file
+in Table 3 of the paper.
diff --git a/examples/covasim_/doubling_beta/causal_test_beta.py b/examples/covasim_/doubling_beta/example_beta.py
similarity index 61%
rename from examples/covasim_/doubling_beta/causal_test_beta.py
rename to examples/covasim_/doubling_beta/example_beta.py
index 32f6f8ad..6a312506 100644
--- a/examples/covasim_/doubling_beta/causal_test_beta.py
+++ b/examples/covasim_/doubling_beta/example_beta.py
@@ -13,6 +13,12 @@
 from causal_testing.testing.base_test_case import BaseTestCase
 from matplotlib.pyplot import rcParams
 
+import os
+import logging
+
+logger = logging.getLogger(__name__)
+logging.basicConfig(level=logging.DEBUG, format="%(message)s")
+
 # Uncommenting the code below will make all graphs publication quality but requires a suitable latex installation
 
 # plt.rcParams["figure.figsize"] = (8, 8)
@@ -25,12 +31,14 @@
 # }
 # rcParams.update(rc_fonts)
 
-OBSERVATIONAL_DATA_PATH = "data/10k_observational_data.csv"
+ROOT = os.path.realpath(os.path.dirname(__file__))
+OBSERVATIONAL_DATA_PATH = f"{ROOT}/data/10k_observational_data.csv"
 
 
-def doubling_beta_CATE_on_csv(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
+def doubling_beta_CATE_on_csv(
+    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
     after removing rows with beta==0.032.
 
@@ -95,13 +103,13 @@ def doubling_beta_CATE_on_csv(observational_data_path: str, simulate_counterfact
 
 
 def doubling_beta_CATEs(observational_data_path: str, simulate_counterfactual: bool = False, verbose: bool = False):
-    """ Compute the CATE for the effect of doubling beta across simulations with different age demographics.
+    """Compute the CATE for the effect of doubling beta across simulations with different age demographics.
     To compute the CATE, this method splits the observational data into high and low age data and computes the
     ATE using this data and a linear regression model.
 
     Since this method already adjusts for age, adding age as an adjustment to the LR model will have no impact.
     However, adding contacts as an adjustment should reduce bias and reveal the average causal effect of doubling beta
-    in simulations of a particular age demographic. """
+    in simulations of a particular age demographic."""
 
     # Create separate subplots for each more specific causal question
     all_fig, all_axes = plt.subplots(1, 1, figsize=(4, 3), squeeze=False)
@@ -110,71 +118,80 @@ def doubling_beta_CATEs(observational_data_path: str, simulate_counterfactual: b
 
     # Apply CT to get the ATE over all executions
     if verbose:
-        print("Running causal tests for all data...")
+        logger.info("Running causal tests for all data...")
     all_data_results_dict = doubling_beta_CATE_on_csv(observational_data_path, simulate_counterfactual, verbose=False)
     plot_doubling_beta_CATEs(all_data_results_dict, "All Data", all_fig, all_axes, row=0, col=0)
 
     # Split data into age-specific strata
     past_execution_df = pd.read_csv(observational_data_path)
-    min_age = np.floor(past_execution_df['avg_age'].min())
-    max_age = np.ceil(past_execution_df['avg_age'].max())
+    min_age = np.floor(past_execution_df["avg_age"].min())
+    max_age = np.ceil(past_execution_df["avg_age"].max())
     mid_age = (min_age + max_age) / 2
 
     # Split df into two age ranges
-    younger_population_df = past_execution_df.loc[past_execution_df['avg_age'] <= mid_age]
-    younger_population_df.to_csv("./data/younger_population.csv")
-    older_population_df = past_execution_df.loc[past_execution_df['avg_age'] > mid_age]
-    older_population_df.to_csv("./data/older_population.csv")
+    younger_population_df = past_execution_df.loc[past_execution_df["avg_age"] <= mid_age]
+    younger_population_df.to_csv(f"{ROOT}/data/younger_population.csv")
+    older_population_df = past_execution_df.loc[past_execution_df["avg_age"] > mid_age]
+    older_population_df.to_csv(f"{ROOT}/data/older_population.csv")
 
     # Repeat analysis on age-specific strata
-    separated_observational_data_paths = ["./data/younger_population.csv",
-                                          "./data/older_population.csv"]
+    separated_observational_data_paths = [f"{ROOT}/data/younger_population.csv", f"{ROOT}/data/older_population.csv"]
 
     for col, separated_observational_data_path in enumerate(separated_observational_data_paths):
-        age_data_results_dict = doubling_beta_CATE_on_csv(separated_observational_data_path, simulate_counterfactual,
-                                                          verbose=False)
+        age_data_results_dict = doubling_beta_CATE_on_csv(
+            separated_observational_data_path, simulate_counterfactual, verbose=False
+        )
         age_stratified_df = pd.read_csv(separated_observational_data_path)
         age_stratified_df_avg_age = round(age_stratified_df["avg_age"].mean(), 1)
         if verbose:
-            print(f"Running causal tests for data with average age of {age_stratified_df_avg_age}")
-        plot_doubling_beta_CATEs(age_data_results_dict, f"Age={age_stratified_df_avg_age}", age_fig, age_axes, row=0,
-                                 col=col)
+            logger.info("Running causal tests for data with average age of %s", age_stratified_df_avg_age)
+        plot_doubling_beta_CATEs(
+            age_data_results_dict, f"Age={age_stratified_df_avg_age}", age_fig, age_axes, row=0, col=col
+        )
 
         # Split df into contact-specific strata
-        min_contacts = np.floor(age_stratified_df['contacts'].min())
-        max_contacts = np.ceil(age_stratified_df['contacts'].max())
+        min_contacts = np.floor(age_stratified_df["contacts"].min())
+        max_contacts = np.ceil(age_stratified_df["contacts"].max())
         mid_contacts = (max_contacts + min_contacts) / 2
 
         # Save dfs to csv
-        low_contacts_df = age_stratified_df.loc[age_stratified_df['contacts'] <= mid_contacts]
-        low_contacts_df.to_csv(f"./data/low_contacts_avg_age_{age_stratified_df_avg_age}.csv")
-        high_contacts_df = age_stratified_df.loc[age_stratified_df['contacts'] > mid_contacts]
-        high_contacts_df.to_csv(f"./data/high_contacts_avg_age_{age_stratified_df_avg_age}.csv")
+        low_contacts_df = age_stratified_df.loc[age_stratified_df["contacts"] <= mid_contacts]
+        low_contacts_df.to_csv(f"{ROOT}/data/low_contacts_avg_age_{age_stratified_df_avg_age}.csv")
+        high_contacts_df = age_stratified_df.loc[age_stratified_df["contacts"] > mid_contacts]
+        high_contacts_df.to_csv(f"{ROOT}/data/high_contacts_avg_age_{age_stratified_df_avg_age}.csv")
 
-        contact_observational_data_paths = [f"./data/low_contacts_avg_age_"
-                                            f"{age_stratified_df_avg_age}.csv",
-                                            f"./data/high_contacts_avg_age_"
-                                            f"{age_stratified_df_avg_age}.csv"]
+        contact_observational_data_paths = [
+            f"{ROOT}/data/low_contacts_avg_age_" f"{age_stratified_df_avg_age}.csv",
+            f"{ROOT}/data/high_contacts_avg_age_" f"{age_stratified_df_avg_age}.csv",
+        ]
 
         # Compute the CATE for each age-contact group
         for row, age_contact_data_path in enumerate(contact_observational_data_paths):
-            age_contact_data_results_dict = doubling_beta_CATE_on_csv(age_contact_data_path, simulate_counterfactual,
-                                                                      verbose=False)
+            age_contact_data_results_dict = doubling_beta_CATE_on_csv(
+                age_contact_data_path, simulate_counterfactual, verbose=False
+            )
             age_contact_stratified_df = pd.read_csv(age_contact_data_path)
             age_contact_stratified_df_avg_contacts = round(age_contact_stratified_df["contacts"].mean(), 1)
             if verbose:
-                print(f"Running causal tests for data with average age of {age_stratified_df_avg_age} and "
-                      f"{age_contact_stratified_df_avg_contacts} average household contacts.")
-            plot_doubling_beta_CATEs(age_contact_data_results_dict,
-                                     f"Age={age_stratified_df_avg_age} "
-                                     f"Contacts={age_contact_stratified_df_avg_contacts}",
-                                     age_contact_fig, age_contact_axes, row=row, col=col)
+                logger.info(
+                    "Running causal tests for data with average age of %s and " "%s average household contacts.",
+                    age_stratified_df_avg_age,
+                    age_contact_stratified_df_avg_contacts,
+                )
+            plot_doubling_beta_CATEs(
+                age_contact_data_results_dict,
+                f"Age={age_stratified_df_avg_age} " f"Contacts={age_contact_stratified_df_avg_contacts}",
+                age_contact_fig,
+                age_contact_axes,
+                row=row,
+                col=col,
+            )
 
     # Save plots
     if simulate_counterfactual:
-        outpath_base_str = './counterfactuals_'
+        outpath_base_str = f"{ROOT}/counterfactuals_"
     else:
-        outpath_base_str = './'
+        outpath_base_str = f"{ROOT}/"
     all_fig.savefig(outpath_base_str + "all_executions.pdf", format="pdf")
     age_fig.savefig(outpath_base_str + "age_executions.pdf", format="pdf")
     age_contact_fig.savefig(outpath_base_str + "age_contact_executions.pdf", format="pdf")
@@ -182,40 +199,50 @@ def doubling_beta_CATEs(observational_data_path: str, simulate_counterfactual: b
 
 def engine_setup(observational_data_path):
     # 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)
-    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)
+    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})
+    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)
+    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)
+    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_path)
+    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)
@@ -228,57 +255,65 @@ def engine_setup(observational_data_path):
 
 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']
-    association_ate = results_dict['association']['ate']
+    ate = results_dict["causation"]["ate"]
+    association_ate = results_dict["association"]["ate"]
 
-    causation_df = results_dict['causation']['df']
-    association_df = results_dict['association']['df']
+    causation_df = results_dict["causation"]["df"]
+    association_df = results_dict["association"]["df"]
 
-    percentage_ate = round((ate / causation_df['cum_infections'].mean()) * 100, 3)
-    association_percentage_ate = round((association_ate / association_df['cum_infections'].mean()) * 100, 3)
+    percentage_ate = round((ate / causation_df["cum_infections"].mean()) * 100, 3)
+    association_percentage_ate = round((association_ate / association_df["cum_infections"].mean()) * 100, 3)
 
     # Get 95% confidence intervals for association and causation
-    ate_cis = results_dict['causation']['cis']
-    association_ate_cis = results_dict['association']['cis']
-    percentage_causal_ate_cis = [round(((ci / causation_df['cum_infections'].mean()) * 100), 3) for ci in ate_cis]
-    percentage_association_ate_cis = [round(((ci / association_df['cum_infections'].mean()) * 100), 3) for ci in
-                                      association_ate_cis]
+    ate_cis = results_dict["causation"]["cis"]
+    association_ate_cis = results_dict["association"]["cis"]
+    percentage_causal_ate_cis = [round(((ci / causation_df["cum_infections"].mean()) * 100), 3) for ci in ate_cis]
+    percentage_association_ate_cis = [
+        round(((ci / association_df["cum_infections"].mean()) * 100), 3) for ci in association_ate_cis
+    ]
 
     # Convert confidence intervals to errors for plotting
-    percentage_causal_errs = [percentage_ate - percentage_causal_ate_cis[0],
-                              percentage_causal_ate_cis[1] - percentage_ate]
-    percentage_association_errs = [association_percentage_ate - percentage_association_ate_cis[0],
-                                   percentage_association_ate_cis[1] - association_percentage_ate]
+    percentage_causal_errs = [
+        percentage_ate - percentage_causal_ate_cis[0],
+        percentage_causal_ate_cis[1] - percentage_ate,
+    ]
+    percentage_association_errs = [
+        association_percentage_ate - percentage_association_ate_cis[0],
+        percentage_association_ate_cis[1] - association_percentage_ate,
+    ]
 
     xs = [1, 2]
     ys = [association_percentage_ate, percentage_ate]
     yerrs = [percentage_association_errs, percentage_causal_errs]
-    xticks = ['Association', 'Causation']
-    print(f"Association ATE: {association_percentage_ate} {percentage_association_ate_cis}")
-    print(f"Association executions: {len(association_df)}")
-    print(f"Causal ATE: {percentage_ate} {percentage_causal_ate_cis}")
-    print(f"Causal executions: {len(causation_df)}")
-    if 'counterfactual' in results_dict.keys():
-        cf_ate = results_dict['counterfactual']['ate']
-        cf_df = results_dict['counterfactual']['df']
-        percentage_cf_ate = round((cf_ate / cf_df['cum_infections'].mean()) * 100, 3)
-        cf_ate_cis = results_dict['counterfactual']['cis']
-        percentage_cf_cis = [round(((ci / cf_df['cum_infections'].mean()) * 100), 3) for ci in cf_ate_cis]
-        percentage_cf_errs = [percentage_cf_ate - percentage_cf_cis[0],
-                              percentage_cf_cis[1] - percentage_cf_ate]
+    xticks = ["Association", "Causation"]
+    logger.info("Association ATE: %s %s", association_percentage_ate, percentage_association_ate_cis)
+    logger.info("Association executions: %s", len(association_df))
+    logger.info("Causal ATE: %s %s", percentage_ate, percentage_causal_ate_cis)
+    logger.info("Causal executions: %s", len(causation_df))
+    if "counterfactual" in results_dict.keys():
+        cf_ate = results_dict["counterfactual"]["ate"]
+        cf_df = results_dict["counterfactual"]["df"]
+        percentage_cf_ate = round((cf_ate / cf_df["cum_infections"].mean()) * 100, 3)
+        cf_ate_cis = results_dict["counterfactual"]["cis"]
+        percentage_cf_cis = [round(((ci / cf_df["cum_infections"].mean()) * 100), 3) for ci in cf_ate_cis]
+        percentage_cf_errs = [percentage_cf_ate - percentage_cf_cis[0], percentage_cf_cis[1] - percentage_cf_ate]
         xs = [0.5, 1.5, 2.5]
         ys = [association_percentage_ate, percentage_ate, percentage_cf_ate]
         yerrs = np.array([percentage_association_errs, percentage_causal_errs, percentage_cf_errs]).T
-        xticks = ['Association', 'Causation', 'Counterfactual']
-        print(f"Counterfactual ATE: {percentage_cf_ate} {percentage_cf_cis}")
-        print(f"Counterfactual executions: {len(cf_df)}")
+        xticks = ["Association", "Causation", "Counterfactual"]
+        logger.info("Counterfactual ATE: %s %s", percentage_cf_ate, percentage_cf_cis)
+        logger.info("Counterfactual executions: %s", len(cf_df))
     axes[row, col].set_ylim(0, 30)
     axes[row, col].set_xlim(0, 3)
     axes[row, col].set_xticks(xs, xticks)
     axes[row, col].set_title(title)
-    axes[row, col].errorbar(xs, ys, yerrs, fmt='o', markersize=3, capsize=3, markerfacecolor='red', color='black')
+    axes[row, col].errorbar(xs, ys, yerrs, fmt="o", markersize=3, capsize=3, markerfacecolor="red", color="black")
     figure.supylabel(r"\% Change in Cumulative Infections (ATE)", fontsize=10)
 
 
+def test_doubling_beta_CATEs():
+    doubling_beta_CATEs(OBSERVATIONAL_DATA_PATH, True, True)
+
+
 if __name__ == "__main__":
     doubling_beta_CATEs(OBSERVATIONAL_DATA_PATH, True, True)
diff --git a/examples/covasim_/vaccinating_elderly/README.md b/examples/covasim_/vaccinating_elderly/README.md
index 1bbca341..5f2aaa62 100644
--- a/examples/covasim_/vaccinating_elderly/README.md
+++ b/examples/covasim_/vaccinating_elderly/README.md
@@ -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.
diff --git a/examples/covasim_/vaccinating_elderly/causal_test_vaccine.py b/examples/covasim_/vaccinating_elderly/example_vaccine.py
similarity index 56%
rename from examples/covasim_/vaccinating_elderly/causal_test_vaccine.py
rename to examples/covasim_/vaccinating_elderly/example_vaccine.py
index 03da8ac6..2789b09c 100644
--- a/examples/covasim_/vaccinating_elderly/causal_test_vaccine.py
+++ b/examples/covasim_/vaccinating_elderly/example_vaccine.py
@@ -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)
diff --git a/examples/lr91/README.md b/examples/lr91/README.md
index 729594e4..7a9329c9 100644
--- a/examples/lr91/README.md
+++ b/examples/lr91/README.md
@@ -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.
\ No newline at end of file
+to the inputs on APD90, and should also plot Figure 2 from the paper.
diff --git a/examples/lr91/causal_test_max_conductances.py b/examples/lr91/example_max_conductances.py
similarity index 63%
rename from examples/lr91/causal_test_max_conductances.py
rename to examples/lr91/example_max_conductances.py
index 3e039971..7818f4ad 100644
--- a/examples/lr91/causal_test_max_conductances.py
+++ b/examples/lr91/example_max_conductances.py
@@ -13,6 +13,12 @@
 from causal_testing.testing.base_test_case import BaseTestCase
 from matplotlib.pyplot import rcParams
 
+import os
+import logging
+
+logger = logging.getLogger(__name__)
+logging.basicConfig(level=logging.DEBUG, format="%(message)s")
+
 # Uncommenting the code below will make all graphs publication quality but requires a suitable latex installation
 
 # rc_fonts = {
@@ -23,38 +29,36 @@
 #     "text.latex.preamble": r"\usepackage{libertine}",
 # }
 # rcParams.update(rc_fonts)
-OBSERVATIONAL_DATA_PATH = "./data/normalised_results.csv"
+ROOT = os.path.realpath(os.path.dirname(__file__))
+OBSERVATIONAL_DATA_PATH = f"{ROOT}/data/normalised_results.csv"
 
 
-def causal_testing_sensitivity_analysis():
+def test_sensitivity_analysis():
     """Perform causal testing to evaluate the effect of six conductance inputs on one output, APD90, over the defined
-       (normalised) design distribution to quantify the extent to which each input affects the output, and plot as a
-       graph.
+    (normalised) design distribution to quantify the extent to which each input affects the output, and plot as a
+    graph.
     """
     # Read in the 200 model runs and define mean value and expected effect
-    model_runs = pd.read_csv("data/results.csv")
-    conductance_means = {'G_K': (.5, Positive),
-                         'G_b': (.5, Positive),
-                         'G_K1': (.5, Positive),
-                         'G_si': (.5, Negative),
-                         'G_Na': (.5, NoEffect),
-                         'G_Kp': (.5, NoEffect)}
+    model_runs = pd.read_csv(f"{ROOT}/data/results.csv")
+    conductance_means = {
+        "G_K": (0.5, Positive),
+        "G_b": (0.5, Positive),
+        "G_K1": (0.5, Positive),
+        "G_si": (0.5, Negative),
+        "G_Na": (0.5, NoEffect),
+        "G_Kp": (0.5, NoEffect),
+    }
 
     # Normalise the inputs as per the original study
     normalised_df = normalise_data(model_runs, columns=list(conductance_means.keys()))
-    normalised_df.to_csv("data/normalised_results.csv")
+    normalised_df.to_csv(f"{ROOT}/data/normalised_results.csv")
 
     # For each input, perform 10 causal tests that change the input from its mean value (0.5) to the equidistant values
     # [0, 0.1, 0.2, ..., 0.9, 1] over the input space of each input, as defined by the normalised design distribution.
     # For each input, this will yield 10 causal test results that measure the extent the input causes APD90 to change,
     # enabling us to compare the magnitude and direction of each inputs' effect.
     treatment_values = np.linspace(0, 1, 11)
-    results = {'G_K': {},
-               'G_b': {},
-               'G_K1': {},
-               'G_si': {},
-               'G_Na': {},
-               'G_Kp': {}}
+    results = {"G_K": {}, "G_b": {}, "G_K1": {}, "G_si": {}, "G_Na": {}, "G_Kp": {}}
     for conductance_param, mean_and_oracle in conductance_means.items():
         average_treatment_effects = []
         confidence_intervals = []
@@ -83,60 +87,73 @@ def effects_on_APD90(observational_data_path, treatment_var, control_val, treatm
     :return: ATE for the effect of G_K on APD90
     """
     # 1. Define Causal DAG
-    causal_dag = CausalDAG('./dag.dot')
+    causal_dag = CausalDAG(f"{ROOT}/dag.dot")
 
     # 2. Specify all inputs
-    g_na = Input('G_Na', float)
-    g_si = Input('G_si', float)
-    g_k = Input('G_K', float)
-    g_k1 = Input('G_K1', float)
-    g_kp = Input('G_Kp', float)
-    g_b = Input('G_b', float)
+    g_na = Input("G_Na", float)
+    g_si = Input("G_si", float)
+    g_k = Input("G_K", float)
+    g_k1 = Input("G_K1", float)
+    g_kp = Input("G_Kp", float)
+    g_b = Input("G_b", float)
 
     # 3. Specify all outputs
-    max_voltage = Output('max_voltage', float)
-    rest_voltage = Output('rest_voltage', float)
-    max_voltage_gradient = Output('max_voltage_gradient', float)
-    dome_voltage = Output('dome_voltage', float)
-    apd50 = Output('APD50', int)
-    apd90 = Output('APD90', int)
+    max_voltage = Output("max_voltage", float)
+    rest_voltage = Output("rest_voltage", float)
+    max_voltage_gradient = Output("max_voltage_gradient", float)
+    dome_voltage = Output("dome_voltage", float)
+    apd50 = Output("APD50", int)
+    apd90 = Output("APD90", int)
 
     # 4. Create scenario by applying constraints over a subset of the inputs
     scenario = Scenario(
-        variables={g_na, g_si, g_k, g_k1, g_kp, g_b,
-                   max_voltage, rest_voltage, max_voltage_gradient, dome_voltage, apd50, apd90},
-        constraints=set()
+        variables={
+            g_na,
+            g_si,
+            g_k,
+            g_k1,
+            g_kp,
+            g_b,
+            max_voltage,
+            rest_voltage,
+            max_voltage_gradient,
+            dome_voltage,
+            apd50,
+            apd90,
+        },
+        constraints=set(),
     )
 
     # 5. Create a causal specification from the scenario and causal DAG
     causal_specification = CausalSpecification(scenario, causal_dag)
     base_test_case = BaseTestCase(treatment_var, apd90)
     # 6. Create a causal test case
-    causal_test_case = CausalTestCase(base_test_case=base_test_case,
-                                      expected_causal_effect=expected_causal_effect,
-                                      control_value=control_val,
-                                      treatment_value=treatment_val)
+    causal_test_case = CausalTestCase(
+        base_test_case=base_test_case,
+        expected_causal_effect=expected_causal_effect,
+        control_value=control_val,
+        treatment_value=treatment_val,
+    )
 
     # 7. Create a data collector
-    data_collector = ObservationalDataCollector(scenario, observational_data_path)
+    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
     minimal_adjustment_set = causal_dag.identification(base_test_case)
-    linear_regression_estimator = LinearRegressionEstimator(treatment_var.name, treatment_val, control_val,
-                                                            minimal_adjustment_set,
-                                                            'APD90'
-                                                            )
+    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, 'ate')
-    print(causal_test_result)
+    causal_test_result = causal_test_engine.execute_test(linear_regression_estimator, causal_test_case, "ate")
+    logger.info("%s", causal_test_result)
     return causal_test_result.test_value.value, causal_test_result.confidence_intervals
 
 
-def plot_ates_with_cis(results_dict: dict, xs: list, save: bool = True):
+def plot_ates_with_cis(results_dict: dict, xs: list, save: bool = True, show: bool = False):
     """Plot the average treatment effects for a given treatment against a list of x-values with confidence intervals.
 
     :param results_dict: A dictionary containing results for sensitivity analysis of each input parameter.
@@ -144,40 +161,35 @@ def plot_ates_with_cis(results_dict: dict, xs: list, save: bool = True):
     :param save: Whether to save the plot.
     """
     fig, axes = plt.subplots()
-    input_colors = {'G_Na': 'red',
-                    'G_si': 'green',
-                    'G_K': 'blue',
-                    'G_K1': 'magenta',
-                    'G_Kp': 'cyan',
-                    'G_b': 'yellow'}
+    input_colors = {"G_Na": "red", "G_si": "green", "G_K": "blue", "G_K1": "magenta", "G_Kp": "cyan", "G_b": "yellow"}
     for treatment, test_results in results_dict.items():
-        ates = test_results['ate']
-        cis = test_results['cis']
-        before_underscore, after_underscore = treatment.split('_')
+        ates = test_results["ate"]
+        cis = test_results["cis"]
+        before_underscore, after_underscore = treatment.split("_")
         after_underscore_braces = f"{{{after_underscore}}}"
         latex_compatible_treatment_str = rf"${before_underscore}_{after_underscore_braces}$"
         cis_low = [c[0] for c in cis]
         cis_high = [c[1] for c in cis]
-        axes.fill_between(xs, cis_low, cis_high, alpha=.2, color=input_colors[treatment],
-                          label=latex_compatible_treatment_str)
+        axes.fill_between(
+            xs, cis_low, cis_high, alpha=0.2, color=input_colors[treatment], label=latex_compatible_treatment_str
+        )
         axes.plot(xs, ates, color=input_colors[treatment], linewidth=1)
-        axes.plot(xs, [0] * len(xs), color='black', alpha=.5, linestyle='--', linewidth=1)
+        axes.plot(xs, [0] * len(xs), color="black", alpha=0.5, linestyle="--", linewidth=1)
     axes.set_ylabel(r"ATE: Change in $APD_{90} (ms)$")
     axes.set_xlabel(r"Treatment value")
     axes.set_ylim(-80, 80)
     axes.set_xlim(min(xs), max(xs))
     box = axes.get_position()
-    axes.set_position([box.x0, box.y0 + box.height * 0.3,
-                       box.width * 0.85, box.height * 0.7])
-    plt.legend(loc='center left', bbox_to_anchor=(1.01, 0.5), fancybox=True, ncol=1,
-               title=r'Input (95\% CIs)')
+    axes.set_position([box.x0, box.y0 + box.height * 0.3, box.width * 0.85, box.height * 0.7])
+    plt.legend(loc="center left", bbox_to_anchor=(1.01, 0.5), fancybox=True, ncol=1, title=r"Input (95\% CIs)")
     if save:
         plt.savefig(f"APD90_sensitivity.pdf", format="pdf")
-    plt.show()
+    if show:
+        plt.show()
 
 
 def normalise_data(df, columns=None):
-    """ Normalise the data in the dataframe into the range [0, 1]. """
+    """Normalise the data in the dataframe into the range [0, 1]."""
     if columns:
         df[columns] = (df[columns] - df[columns].min()) / (df[columns].max() - df[columns].min())
         return df
@@ -185,5 +197,5 @@ def normalise_data(df, columns=None):
         return (df - df.min()) / (df.max() - df.min())
 
 
-if __name__ == '__main__':
-    causal_testing_sensitivity_analysis()
+if __name__ == "__main__":
+    test_sensitivity_analysis(show=True)
diff --git a/examples/poisson-line-process/README.md b/examples/poisson-line-process/README.md
index fbd70c04..2b8f1fdf 100644
--- a/examples/poisson-line-process/README.md
+++ b/examples/poisson-line-process/README.md
@@ -6,6 +6,6 @@ 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/poisson-line-process`.
-3. Run the command `python causal_test_poisson.py`
+3. Run the command `python test_poisson_process.py`
 
 This should print a series of causal test results and produce two CSV files. `intensity_num_shapes_results_random_1000.csv` corresponds to table 1, and `width_num_shapes_results_random_1000.csv` relates to our findings regarding the relationship of width and `P_u`.
diff --git a/examples/poisson-line-process/example_poisson_process.py b/examples/poisson-line-process/example_poisson_process.py
new file mode 100644
index 00000000..d8d941e0
--- /dev/null
+++ b/examples/poisson-line-process/example_poisson_process.py
@@ -0,0 +1,221 @@
+from causal_testing.specification.causal_dag import CausalDAG
+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 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
+
+import pandas as pd
+import os
+import logging
+
+logger = logging.getLogger(__name__)
+logging.basicConfig(level=logging.DEBUG, format="%(message)s")
+
+
+class EmpiricalMeanEstimator(Estimator):
+    def add_modelling_assumptions(self):
+        """
+        Add modelling assumptions to the estimator. This is a list of strings which list the modelling assumptions that
+        must hold if the resulting causal inference is to be considered valid.
+        """
+        self.modelling_assumptions += "The data must contain runs with the exact configuration of interest."
+
+    def estimate_ate(self) -> float:
+        """Estimate the outcomes under control and treatment.
+        :return: The empirical average treatment effect.
+        """
+        control_results = self.df.where(self.df[self.treatment[0]] == self.control_value)[self.outcome].dropna()
+        treatment_results = self.df.where(self.df[self.treatment[0]] == self.treatment_value)[self.outcome].dropna()
+        return treatment_results.mean()[0] - control_results.mean()[0], None
+
+    def estimate_risk_ratio(self) -> float:
+        """Estimate the outcomes under control and treatment.
+        :return: The empirical average treatment effect.
+        """
+        control_results = self.df.where(self.df[self.treatment[0]] == self.control_value)[self.outcome].dropna()
+        treatment_results = self.df.where(self.df[self.treatment[0]] == self.treatment_value)[self.outcome].dropna()
+        return treatment_results.mean()[0] / control_results.mean()[0], None
+
+
+# 1. Read in the Causal DAG
+ROOT = os.path.realpath(os.path.dirname(__file__))
+causal_dag = CausalDAG(f"{ROOT}/dag.dot")
+
+# 2. Create variables
+width = Input("width", float)
+height = Input("height", float)
+intensity = Input("intensity", float)
+
+num_lines_abs = Output("num_lines_abs", float)
+num_lines_unit = Output("num_lines_unit", float)
+num_shapes_abs = Output("num_shapes_abs", float)
+num_shapes_unit = Output("num_shapes_unit", float)
+
+# 3. Create scenario by applying constraints over a subset of the input variables
+scenario = Scenario(
+    variables={
+        width,
+        height,
+        intensity,
+        num_lines_abs,
+        num_lines_unit,
+        num_shapes_abs,
+        num_shapes_unit,
+    }
+)
+
+# 4. Construct a causal specification from the scenario and causal DAG
+causal_specification = CausalSpecification(scenario, causal_dag)
+
+observational_data_path = f"{ROOT}/data/random/data_random_1000.csv"
+
+
+def causal_test_intensity_num_shapes(
+    observational_data_path,
+    causal_test_case,
+    square_terms=[],
+    inverse_terms=[],
+    empirical=False,
+):
+    # 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
+    minimal_adjustment_set = causal_dag.identification(causal_test_case.base_test_case)
+
+    # 9. Set up an estimator
+    data = pd.read_csv(observational_data_path)
+
+    treatment = causal_test_case.get_treatment_variable()
+    outcome = causal_test_case.get_outcome_variable()
+
+    estimator = None
+    if empirical:
+        estimator = EmpiricalMeanEstimator(
+            treatment=[treatment],
+            control_value=causal_test_case.control_value,
+            treatment_value=causal_test_case.treatment_value,
+            adjustment_set=set(),
+            outcome=[outcome],
+            df=data,
+            effect_modifiers=causal_test_case.effect_modifier_configuration,
+        )
+    else:
+        square_terms = [f"np.power({t}, 2)" for t in square_terms]
+        inverse_terms = [f"np.float_power({t}, -1)" for t in inverse_terms]
+        estimator = LinearRegressionEstimator(
+            treatment=treatment,
+            control_value=causal_test_case.control_value,
+            treatment_value=causal_test_case.treatment_value,
+            adjustment_set=set(),
+            outcome=outcome,
+            df=data,
+            effect_modifiers=causal_test_case.effect_modifier_configuration,
+            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, causal_test_case.estimate_type
+    )
+
+    return causal_test_result
+
+
+def test_poisson_intensity_num_shapes(save=False):
+    intensity_num_shapes_results = []
+    for wh in range(1, 11):
+        smt_data_path = f"{ROOT}/data/smt_100/data_smt_wh{wh}_100.csv"
+        for control_value, treatment_value in [(1, 2), (2, 4), (4, 8), (8, 16)]:
+            logger.info("%s CAUSAL TEST %s", "=" * 33, "=" * 33)
+            logger.info("WIDTH = HEIGHT = %s", wh)
+            logger.info("Identifying")
+            base_test_case = BaseTestCase(treatment_variable=intensity, outcome_variable=num_shapes_unit)
+            causal_test_case = CausalTestCase(
+                base_test_case=base_test_case,
+                expected_causal_effect=ExactValue(4, tolerance=0.5),
+                treatment_value=treatment_value,
+                control_value=control_value,
+                estimate_type="risk_ratio",
+            )
+            obs_causal_test_result = causal_test_intensity_num_shapes(
+                observational_data_path,
+                causal_test_case,
+                square_terms=["intensity"],
+                empirical=False,
+            )
+            logger.info("Observational %s", obs_causal_test_result)
+            smt_causal_test_result = causal_test_intensity_num_shapes(
+                smt_data_path, causal_test_case, square_terms=["intensity"], empirical=True
+            )
+            logger.info("RCT %s", smt_causal_test_result)
+
+            results = {
+                "width": wh,
+                "height": wh,
+                "control": control_value,
+                "treatment": treatment_value,
+                "smt_risk_ratio": smt_causal_test_result.test_value.value,
+                "obs_risk_ratio": obs_causal_test_result.test_value.value,
+            }
+            intensity_num_shapes_results.append(results)
+
+    intensity_num_shapes_results = pd.DataFrame(intensity_num_shapes_results)
+    if save:
+        intensity_num_shapes_results.to_csv("intensity_num_shapes_results_random_1000.csv")
+    logger.info("%s", intensity_num_shapes_results)
+
+
+def test_poisson_width_num_shapes(save=False):
+    width_num_shapes_results = []
+    for i in range(17):
+        for w in range(1, 10):
+            logger.info("%s CAUSAL TEST %s", "=" * 33, "=" * 33)
+            logger.info("Identifying")
+            # 5. Create a causal test case
+            control_value = w
+            treatment_value = w + 1
+            base_test_case = BaseTestCase(treatment_variable=width, outcome_variable=num_shapes_unit)
+            causal_test_case = CausalTestCase(
+                base_test_case=base_test_case,
+                expected_causal_effect=Positive(),
+                control_value=control_value,
+                treatment_value=treatment_value,
+                estimate_type="ate_calculated",
+                effect_modifier_configuration={"intensity": i},
+            )
+            causal_test_result = causal_test_intensity_num_shapes(
+                observational_data_path,
+                causal_test_case,
+                square_terms=["intensity"],
+                inverse_terms=["width"],
+            )
+            logger.info("%s", causal_test_result)
+            results = {
+                "control": control_value,
+                "treatment": treatment_value,
+                "intensity": i,
+                "ate": causal_test_result.test_value.value,
+                "ci_low": min(causal_test_result.confidence_intervals),
+                "ci_high": max(causal_test_result.confidence_intervals),
+            }
+            width_num_shapes_results.append(results)
+    width_num_shapes_results = pd.DataFrame(width_num_shapes_results)
+    if save:
+        width_num_shapes_results.to_csv("width_num_shapes_results_random_1000.csv")
+    logger.info("%s", width_num_shapes_results)
+
+
+if __name__ == "__main__":
+    # test_poisson_intensity_num_shapes(save=True)
+    test_poisson_width_num_shapes(save=True)
diff --git a/examples/poisson/README.md b/examples/poisson/README.md
index 5242ccd5..f4e27a62 100644
--- a/examples/poisson/README.md
+++ b/examples/poisson/README.md
@@ -6,4 +6,6 @@ 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/poisson`.
-3. Run the command `python run_causal_tests.py --data_path data.csv --dag_path dag.dot --json_path causal_tests.json`
+3. Run the command `python test_run_causal_tests.py --data_path data.csv --dag_path dag.dot --json_path causal_tests.json`
+
+This should print a series of causal test results and produce two CSV files. `intensity_num_shapes_results_random_1000.csv` corresponds to table 1, and `width_num_shapes_results_random_1000.csv` relates to our findings regarding the relationship of width and `P_u`.
diff --git a/pyproject.toml b/pyproject.toml
index 0bdc7a57..6e1406e8 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -41,6 +41,9 @@ dev = [
     "sphinx-autoapi",
     "sphinx_rtd_theme",
 ]
+test = [
+  "covasim~=3.0.7",
+]
 
 [project.urls]
 Bug_Tracker = "https://github.com/CITCOM-project/CausalTestingFramework/issues"
@@ -61,4 +64,10 @@ target-version = ["py39"]
 max_line_length = 120
 in-place = true
 recursive = true
-aggressive = 3
\ No newline at end of file
+aggressive = 3
+
+[tool.pytest.ini_options]
+minversion = "6.0"
+python_files=[
+    "test_*.py",
+    "example_*.py"]


Test patch:
diff --git a/.github/workflows/ci-tests.yaml b/.github/workflows/ci-tests.yaml
index 35101591..db190220 100644
--- a/.github/workflows/ci-tests.yaml
+++ b/.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
diff --git a/examples/lr91/causal_test_max_conductances_test_suite.py b/examples/lr91/example_max_conductances_test_suite.py
similarity index 62%
rename from examples/lr91/causal_test_max_conductances_test_suite.py
rename to examples/lr91/example_max_conductances_test_suite.py
index 8c8320a2..e7367163 100644
--- a/examples/lr91/causal_test_max_conductances_test_suite.py
+++ b/examples/lr91/example_max_conductances_test_suite.py
@@ -14,6 +14,8 @@
 from causal_testing.testing.causal_test_suite import CausalTestSuite
 from matplotlib.pyplot import rcParams
 
+import os
+
 # Uncommenting the code below will make all graphs publication quality but requires a suitable latex installation
 
 # rc_fonts = {
@@ -24,40 +26,38 @@
 #     "text.latex.preamble": r"\usepackage{libertine}",
 # }
 # rcParams.update(rc_fonts)
-OBSERVATIONAL_DATA_PATH = "./data/normalised_results.csv"
+ROOT = os.path.realpath(os.path.dirname(__file__))
+OBSERVATIONAL_DATA_PATH = f"{ROOT}/data/normalised_results.csv"
 
 
-def causal_testing_sensitivity_analysis():
+def test_sensitivity_analysis():
     """Perform causal testing to evaluate the effect of six conductance inputs on one output, APD90, over the defined
-       (normalised) design distribution to quantify the extent to which each input affects the output, and plot as a
-       graph.
+    (normalised) design distribution to quantify the extent to which each input affects the output, and plot as a
+    graph.
     """
     # Read in the 200 model runs and define mean value and expected effect
-    model_runs = pd.read_csv("data/results.csv")
-    conductance_means = {'G_K': (.5, Positive),
-                         'G_b': (.5, Positive),
-                         'G_K1': (.5, Positive),
-                         'G_si': (.5, Negative),
-                         'G_Na': (.5, NoEffect),
-                         'G_Kp': (.5, NoEffect)}
+    model_runs = pd.read_csv(f"{ROOT}/data/results.csv")
+    conductance_means = {
+        "G_K": (0.5, Positive),
+        "G_b": (0.5, Positive),
+        "G_K1": (0.5, Positive),
+        "G_si": (0.5, Negative),
+        "G_Na": (0.5, NoEffect),
+        "G_Kp": (0.5, NoEffect),
+    }
 
     # Normalise the inputs as per the original study
     normalised_df = normalise_data(model_runs, columns=list(conductance_means.keys()))
-    normalised_df.to_csv("data/normalised_results.csv")
+    normalised_df.to_csv(f"{ROOT}/data/normalised_results.csv")
 
     # For each input, perform 10 causal tests that change the input from its mean value (0.5) to the equidistant values
     # [0, 0.1, 0.2, ..., 0.9, 1] over the input space of each input, as defined by the normalised design distribution.
     # For each input, this will yield 10 causal test results that measure the extent the input causes APD90 to change,
     # enabling us to compare the magnitude and direction of each inputs' effect.
     treatment_values = np.linspace(0, 1, 11)
-    results = {'G_K': {},
-               'G_b': {},
-               'G_K1': {},
-               'G_si': {},
-               'G_Na': {},
-               'G_Kp': {}}
-
-    apd90 = Output('APD90', int)
+    results = {"G_K": {}, "G_b": {}, "G_K1": {}, "G_si": {}, "G_Na": {}, "G_Kp": {}}
+
+    apd90 = Output("APD90", int)
     outcome_variable = apd90
     test_suite = CausalTestSuite()
     estimator_list = [LinearRegressionEstimator]
@@ -71,20 +71,27 @@ def causal_testing_sensitivity_analysis():
         mean, oracle = mean_and_oracle
         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,
-                                   causal_test_case_list=test_list,
-                                   estimators_classes=estimator_list,
-                                   estimate_type='ate')
+        test_suite.add_test_object(
+            base_test_case=base_test_case,
+            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.test_value.value for result in causal_test_results[base_test_case]['LinearRegressionEstimator']],
-             "cis": [result.confidence_intervals for result in
-                     causal_test_results[base_test_case]['LinearRegressionEstimator']]}
+        results[base_test_case.treatment_variable.name] = {
+            "ate": [
+                result.test_value.value for result in causal_test_results[base_test_case]["LinearRegressionEstimator"]
+            ],
+            "cis": [
+                result.confidence_intervals
+                for result in causal_test_results[base_test_case]["LinearRegressionEstimator"]
+            ],
+        }
 
     plot_ates_with_cis(results, treatment_values)
 
@@ -97,36 +104,48 @@ def effects_on_APD90(observational_data_path, test_suite):
     :return: causal_test_results containing a list of causal_test_result objects
     """
     # 1. Define Causal DAG
-    causal_dag = CausalDAG('./dag.dot')
+    causal_dag = CausalDAG(f"{ROOT}/dag.dot")
 
     # 2. Specify all inputs
-    g_na = Input('G_Na', float)
-    g_si = Input('G_si', float)
-    g_k = Input('G_K', float)
-    g_k1 = Input('G_K1', float)
-    g_kp = Input('G_Kp', float)
-    g_b = Input('G_b', float)
+    g_na = Input("G_Na", float)
+    g_si = Input("G_si", float)
+    g_k = Input("G_K", float)
+    g_k1 = Input("G_K1", float)
+    g_kp = Input("G_Kp", float)
+    g_b = Input("G_b", float)
 
     # 3. Specify all outputs
-    max_voltage = Output('max_voltage', float)
-    rest_voltage = Output('rest_voltage', float)
-    max_voltage_gradient = Output('max_voltage_gradient', float)
-    dome_voltage = Output('dome_voltage', float)
-    apd50 = Output('APD50', int)
-    apd90 = Output('APD90', int)
+    max_voltage = Output("max_voltage", float)
+    rest_voltage = Output("rest_voltage", float)
+    max_voltage_gradient = Output("max_voltage_gradient", float)
+    dome_voltage = Output("dome_voltage", float)
+    apd50 = Output("APD50", int)
+    apd90 = Output("APD90", int)
 
     # 4. Create scenario by applying constraints over a subset of the inputs
     scenario = Scenario(
-        variables={g_na, g_si, g_k, g_k1, g_kp, g_b,
-                   max_voltage, rest_voltage, max_voltage_gradient, dome_voltage, apd50, apd90},
-        constraints=set()
+        variables={
+            g_na,
+            g_si,
+            g_k,
+            g_k1,
+            g_kp,
+            g_b,
+            max_voltage,
+            rest_voltage,
+            max_voltage_gradient,
+            dome_voltage,
+            apd50,
+            apd90,
+        },
+        constraints=set(),
     )
 
     # 5. Create a causal specification from the scenario and causal DAG
     causal_specification = CausalSpecification(scenario, causal_dag)
 
     # 7. Create a data collector
-    data_collector = ObservationalDataCollector(scenario, observational_data_path)
+    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)
@@ -136,7 +155,7 @@ def effects_on_APD90(observational_data_path, test_suite):
     return causal_test_results
 
 
-def plot_ates_with_cis(results_dict: dict, xs: list, save: bool = True):
+def plot_ates_with_cis(results_dict: dict, xs: list, save: bool = False, show=False):
     """Plot the average treatment effects for a given treatment against a list of x-values with confidence intervals.
 
     :param results_dict: A dictionary containing results for sensitivity analysis of each input parameter.
@@ -144,40 +163,35 @@ def plot_ates_with_cis(results_dict: dict, xs: list, save: bool = True):
     :param save: Whether to save the plot.
     """
     fig, axes = plt.subplots()
-    input_colors = {'G_Na': 'red',
-                    'G_si': 'green',
-                    'G_K': 'blue',
-                    'G_K1': 'magenta',
-                    'G_Kp': 'cyan',
-                    'G_b': 'yellow'}
+    input_colors = {"G_Na": "red", "G_si": "green", "G_K": "blue", "G_K1": "magenta", "G_Kp": "cyan", "G_b": "yellow"}
     for treatment, test_results in results_dict.items():
-        ates = test_results['ate']
-        cis = test_results['cis']
-        before_underscore, after_underscore = treatment.split('_')
+        ates = test_results["ate"]
+        cis = test_results["cis"]
+        before_underscore, after_underscore = treatment.split("_")
         after_underscore_braces = f"{{{after_underscore}}}"
         latex_compatible_treatment_str = rf"${before_underscore}_{after_underscore_braces}$"
         cis_low = [c[0] for c in cis]
         cis_high = [c[1] for c in cis]
-        axes.fill_between(xs, cis_low, cis_high, alpha=.2, color=input_colors[treatment],
-                          label=latex_compatible_treatment_str)
+        axes.fill_between(
+            xs, cis_low, cis_high, alpha=0.2, color=input_colors[treatment], label=latex_compatible_treatment_str
+        )
         axes.plot(xs, ates, color=input_colors[treatment], linewidth=1)
-        axes.plot(xs, [0] * len(xs), color='black', alpha=.5, linestyle='--', linewidth=1)
+        axes.plot(xs, [0] * len(xs), color="black", alpha=0.5, linestyle="--", linewidth=1)
     axes.set_ylabel(r"ATE: Change in $APD_{90} (ms)$")
     axes.set_xlabel(r"Treatment value")
     axes.set_ylim(-80, 80)
     axes.set_xlim(min(xs), max(xs))
     box = axes.get_position()
-    axes.set_position([box.x0, box.y0 + box.height * 0.3,
-                       box.width * 0.85, box.height * 0.7])
-    plt.legend(loc='center left', bbox_to_anchor=(1.01, 0.5), fancybox=True, ncol=1,
-               title=r'Input (95\% CIs)')
+    axes.set_position([box.x0, box.y0 + box.height * 0.3, box.width * 0.85, box.height * 0.7])
+    plt.legend(loc="center left", bbox_to_anchor=(1.01, 0.5), fancybox=True, ncol=1, title=r"Input (95\% CIs)")
     if save:
         plt.savefig(f"APD90_sensitivity.pdf", format="pdf")
-    plt.show()
+    if show:
+        plt.show()
 
 
 def normalise_data(df, columns=None):
-    """ Normalise the data in the dataframe into the range [0, 1]. """
+    """Normalise the data in the dataframe into the range [0, 1]."""
     if columns:
         df[columns] = (df[columns] - df[columns].min()) / (df[columns].max() - df[columns].min())
         return df
@@ -185,5 +199,5 @@ def normalise_data(df, columns=None):
         return (df - df.min()) / (df.max() - df.min())
 
 
-if __name__ == '__main__':
-    causal_testing_sensitivity_analysis()
+if __name__ == "__main__":
+    test_sensitivity_analysis(show=True, save=True)
diff --git a/examples/poisson-line-process/causal_test_poisson.py b/examples/poisson-line-process/causal_test_poisson.py
deleted file mode 100644
index c6c1fa87..00000000
--- a/examples/poisson-line-process/causal_test_poisson.py
+++ /dev/null
@@ -1,224 +0,0 @@
-from causal_testing.specification.causal_dag import CausalDAG
-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 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
-
-import pandas as pd
-
-
-class EmpiricalMeanEstimator(Estimator):
-    def add_modelling_assumptions(self):
-        """
-        Add modelling assumptions to the estimator. This is a list of strings which list the modelling assumptions that
-        must hold if the resulting causal inference is to be considered valid.
-        """
-        self.modelling_assumptions += (
-            "The data must contain runs with the exact configuration of interest."
-        )
-
-    def estimate_ate(self) -> float:
-        """ Estimate the outcomes under control and treatment.
-        :return: The empirical average treatment effect.
-        """
-        control_results = self.df.where(
-            self.df[self.treatment[0]] == self.control_value
-        )[self.outcome].dropna()
-        treatment_results = self.df.where(
-            self.df[self.treatment[0]] == self.treatment_value
-        )[self.outcome].dropna()
-        return treatment_results.mean()[0] - control_results.mean()[0], None
-
-    def estimate_risk_ratio(self) -> float:
-        """ Estimate the outcomes under control and treatment.
-        :return: The empirical average treatment effect.
-        """
-        control_results = self.df.where(
-            self.df[self.treatment[0]] == self.control_value
-        )[self.outcome].dropna()
-        treatment_results = self.df.where(
-            self.df[self.treatment[0]] == self.treatment_value
-        )[self.outcome].dropna()
-        return treatment_results.mean()[0] / control_results.mean()[0], None
-
-
-# 1. Read in the Causal DAG
-causal_dag = CausalDAG("./dag.dot")
-
-# 2. Create variables
-width = Input("width", float)
-height = Input("height", float)
-intensity = Input("intensity", float)
-
-num_lines_abs = Output("num_lines_abs", float)
-num_lines_unit = Output("num_lines_unit", float)
-num_shapes_abs = Output("num_shapes_abs", float)
-num_shapes_unit = Output("num_shapes_unit", float)
-
-# 3. Create scenario by applying constraints over a subset of the input variables
-scenario = Scenario(
-    variables={
-        width,
-        height,
-        intensity,
-        num_lines_abs,
-        num_lines_unit,
-        num_shapes_abs,
-        num_shapes_unit,
-    }
-)
-
-# 4. Construct a causal specification from the scenario and causal DAG
-causal_specification = CausalSpecification(scenario, causal_dag)
-
-observational_data_path = "data/random/data_random_1000.csv"
-
-intensity_num_shapes_results = []
-
-
-def test_intensity_num_shapes(
-    observational_data_path,
-    causal_test_case,
-    square_terms=[],
-    inverse_terms=[],
-    empirical=False,
-):
-    # 6. Create a data collector
-    data_collector = ObservationalDataCollector(scenario, 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
-    minimal_adjustment_set = causal_dag.identification(causal_test_case.base_test_case)
-
-    # 9. Set up an estimator
-    data = pd.read_csv(observational_data_path)
-
-    treatment = causal_test_case.get_treatment_variable()
-    outcome = causal_test_case.get_outcome_variable()
-
-    estimator = None
-    if empirical:
-        estimator = EmpiricalMeanEstimator(
-            treatment=[treatment],
-            control_value=causal_test_case.control_value,
-            treatment_value=causal_test_case.treatment_value,
-            adjustment_set=set(),
-            outcome=[outcome],
-            df=data,
-            effect_modifiers=causal_test_case.effect_modifier_configuration,
-        )
-    else:
-        square_terms = [f"np.power({t}, 2)" for t in square_terms]
-        inverse_terms = [f"np.float_power({t}, -1)" for t in inverse_terms]
-        estimator = LinearRegressionEstimator(
-            treatment=treatment,
-            control_value=causal_test_case.control_value,
-            treatment_value=causal_test_case.treatment_value,
-            adjustment_set=set(),
-            outcome=outcome,
-            df=data,
-            effect_modifiers=causal_test_case.effect_modifier_configuration,
-            formula=f"{outcome} ~ {treatment} + {'+'.join(square_terms + inverse_terms + list([e.name 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, causal_test_case.estimate_type
-    )
-
-    return causal_test_result
-
-
-
-for wh in range(1, 11):
-    smt_data_path = f"data/smt_100/data_smt_wh{wh}_100.csv"
-    for control_value, treatment_value in [(1, 2), (2, 4), (4, 8), (8, 16)]:
-        print("=" * 33, "CAUSAL TEST", "=" * 33)
-        print(f"WIDTH = HEIGHT = {wh}")
-
-        print("Identifying")
-        base_test_case = BaseTestCase(treatment_variable=intensity,
-                                      outcome_variable=num_shapes_unit)
-        causal_test_case = CausalTestCase(
-            base_test_case=base_test_case,
-            expected_causal_effect=ExactValue(4, tolerance=0.5),
-            treatment_value=treatment_value,
-            control_value=control_value,
-            estimate_type="risk_ratio"
-        )
-        obs_causal_test_result = test_intensity_num_shapes(
-            observational_data_path,
-            causal_test_case,
-            square_terms=["intensity"],
-            empirical=False,
-        )
-        print("Observational", end=" ")
-        print(obs_causal_test_result)
-        smt_causal_test_result = test_intensity_num_shapes(
-            smt_data_path, causal_test_case, square_terms=["intensity"], empirical=True
-        )
-        print("RCT", end=" ")
-        print(smt_causal_test_result)
-
-        results = {
-            "width": wh,
-            "height": wh,
-            "control": control_value,
-            "treatment": treatment_value,
-            "smt_risk_ratio": smt_causal_test_result.test_value.value,
-            "obs_risk_ratio": obs_causal_test_result.test_value.value,
-        }
-        intensity_num_shapes_results.append(results)
-
-intensity_num_shapes_results = pd.DataFrame(intensity_num_shapes_results)
-intensity_num_shapes_results.to_csv("intensity_num_shapes_results_random_1000.csv")
-print(intensity_num_shapes_results)
-
-width_num_shapes_results = []
-for i in range(17):
-    for w in range(1, 10):
-        print("=" * 37, "CAUSAL TEST", "=" * 37)
-        print("Identifying")
-        # 5. Create a causal test case
-        control_value = w
-        treatment_value = w + 1
-        base_test_case = BaseTestCase(
-            treatment_variable=width,
-            outcome_variable=num_shapes_unit
-        )
-        causal_test_case = CausalTestCase(
-            base_test_case=base_test_case,
-            expected_causal_effect=Positive(),
-            control_value=control_value,
-            treatment_value=treatment_value,
-            estimate_type="ate_calculated",
-            effect_modifier_configuration={intensity: i}
-        )
-        causal_test_result = test_intensity_num_shapes(
-            observational_data_path,
-            causal_test_case,
-            square_terms=["intensity"],
-            inverse_terms=["width"],
-        )
-        print(causal_test_result)
-        results = {
-            "control": control_value,
-            "treatment": treatment_value,
-            "intensity": i,
-            "ate": causal_test_result.test_value.value,
-            "ci_low": min(causal_test_result.confidence_intervals),
-            "ci_high": max(causal_test_result.confidence_intervals),
-        }
-        width_num_shapes_results.append(results)
-width_num_shapes_results = pd.DataFrame(width_num_shapes_results)
-width_num_shapes_results.to_csv("width_num_shapes_results_random_1000.csv")
-print(width_num_shapes_results)
diff --git a/examples/poisson/run_causal_tests.py b/examples/poisson/example_run_causal_tests.py
similarity index 58%
rename from examples/poisson/run_causal_tests.py
rename to examples/poisson/example_run_causal_tests.py
index 9dcdf8e8..ff5d203a 100644
--- a/examples/poisson/run_causal_tests.py
+++ b/examples/poisson/example_run_causal_tests.py
@@ -1,6 +1,7 @@
 import numpy as np
 import pandas as pd
 import scipy
+import os
 
 from causal_testing.testing.estimators import LinearRegressionEstimator, CausalForestEstimator
 from causal_testing.testing.causal_test_outcome import ExactValue, Positive, Negative, NoEffect, CausalTestOutcome
@@ -10,6 +11,11 @@
 from causal_testing.specification.scenario import Scenario
 from causal_testing.specification.variable import Input, Output, Meta
 
+import logging
+
+logger = logging.getLogger(__name__)
+logging.basicConfig(level=logging.DEBUG, format="%(message)s")
+
 
 class WidthHeightEstimator(LinearRegressionEstimator):
     """
@@ -21,17 +27,19 @@ def add_modelling_assumptions(self):
         Add modelling assumptions to the estimator. This is a list of strings which list the modelling assumptions that
         must hold if the resulting causal inference is to be considered valid.
         """
-        self.modelling_assumptions += 'The output varies according to 2i(w+h)'
+        self.modelling_assumptions += "The output varies according to 2i(w+h)"
 
     def estimate_ate(self) -> (float, [float, float], float):
-        """ Estimate the conditional average treatment effect of the treatment on the outcome. That is, the change
+        """Estimate the conditional average treatment effect of the treatment on the outcome. That is, the change
         in outcome caused by changing the treatment variable from the control value to the treatment value.
         :return: The conditional average treatment effect and the 95% Wald confidence intervals.
         """
-        assert self.effect_modifiers, f"Must have at least one effect modifier to compute CATE - {self.effect_modifiers}."
+        assert (
+            self.effect_modifiers
+        ), f"Must have at least one effect modifier to compute CATE - {self.effect_modifiers}."
         x = pd.DataFrame()
         x[self.treatment[0]] = [self.treatment_values, self.control_values]
-        x['Intercept'] = 1
+        x["Intercept"] = 1
         for k, v in self.effect_modifiers.items():
             self.adjustment_set.add(k)
             x[k] = v
@@ -39,14 +47,14 @@ def estimate_ate(self) -> (float, [float, float], float):
             for a, b in self.product_terms:
                 x[f"{a}*{b}"] = x[a] * x[b]
 
-        x.drop(['width', 'height'], axis=1, inplace=True)
+        x.drop(["width", "height"], axis=1, inplace=True)
         self.adjustment_set = {"width*intensity", "height*intensity"}
 
-        print(x)
-        print(self.adjustment_set)
+        logger.info("%s", x)
+        logger.info("%s", self.adjustment_set)
 
         model = self._run_linear_regression()
-        print(model.summary())
+        logger.info("%s", model.summary())
         y = model.predict(x)
         treatment_outcome = y.iloc[0]
         control_outcome = y.iloc[1]
@@ -63,14 +71,14 @@ def __init__(self, tolerance=0.5):
 
     def apply(self, res: CausalTestResult) -> bool:
         # TODO: confidence intervals?
-        print("=== APPLYING ===")
-        print("effect_modifier_configuration", res.effect_modifier_configuration)
+        logger.info("=== APPLYING ===")
+        logger.info("effect_modifier_configuration", res.effect_modifier_configuration)
         effect_modifier_configuration = {k.name: v for k, v in res.effect_modifier_configuration.items()}
         c = res.treatment_value - res.control_value
-        i = effect_modifier_configuration['intensity']
+        i = effect_modifier_configuration["intensity"]
         self.i2c = i * 2 * c
-        print("2ic:", f"2*{i}*{c}={self.i2c}")
-        print("ate:", res.test_value.value)
+        logger.info("2ic: 2 * %s * %s = %s", i, c, self.i2c)
+        logger.info("ate: %s", res.test_value.value)
         return np.isclose(res.test_value.value, self.i2c, atol=self.tolerance)
 
     def __str__(self):
@@ -80,34 +88,31 @@ def __str__(self):
 
 
 def populate_width_height(data):
-    data['width_plus_height'] = data['width'] + data['height']
+    data["width_plus_height"] = data["width"] + data["height"]
 
 
 def populate_num_lines_unit(data):
-    area = data['width'] * data['height']
-    data['num_lines_unit'] = data['num_lines_abs'] / area
+    area = data["width"] * data["height"]
+    data["num_lines_unit"] = data["num_lines_abs"] / area
 
 
 def populate_num_shapes_unit(data):
-    area = data['width'] * data['height']
-    data['num_shapes_unit'] = data['num_shapes_abs'] / area
+    area = data["width"] * data["height"]
+    data["num_shapes_unit"] = data["num_shapes_abs"] / area
 
 
 inputs = [
-    {"name": "width", "type": float, "distribution": scipy.stats.uniform(0, 10)},
-    {"name": "height", "type": float, "distribution": scipy.stats.uniform(0, 10)},
-    {"name": "intensity", "type": float, "distribution": scipy.stats.uniform(0, 10)}
+    {"name": "width", "datatype": float, "distribution": scipy.stats.uniform(0, 10)},
+    {"name": "height", "datatype": float, "distribution": scipy.stats.uniform(0, 10)},
+    {"name": "intensity", "datatype": float, "distribution": scipy.stats.uniform(0, 10)},
 ]
 
-outputs = [
-    {"name": "num_lines_abs", "type": float},
-    {"name": "num_shapes_abs", "type": float}
-]
+outputs = [{"name": "num_lines_abs", "datatype": float}, {"name": "num_shapes_abs", "datatype": float}]
 
 metas = [
-    {"name": "num_lines_unit", "type": float, "populate": populate_num_lines_unit},
-    {"name": "num_shapes_unit", "type": float, "populate": populate_num_shapes_unit},
-    {"name": "width_plus_height", "type": float, "populate": populate_width_height}
+    {"name": "num_lines_unit", "datatype": float, "populate": populate_num_lines_unit},
+    {"name": "num_shapes_unit", "datatype": float, "populate": populate_num_shapes_unit},
+    {"name": "width_plus_height", "datatype": float, "populate": populate_width_height},
 ]
 
 constraints = ["width > 0", "height > 0", "intensity > 0"]
@@ -117,7 +122,7 @@ def populate_num_shapes_unit(data):
     "Positive": Positive(),
     "Negative": Negative(),
     "ExactValue4_05": ExactValue(4, tolerance=0.5),
-    "NoEffect": NoEffect()
+    "NoEffect": NoEffect(),
 }
 
 estimators = {
@@ -127,20 +132,22 @@ def populate_num_shapes_unit(data):
 }
 
 # Create input structure required to create a modelling scenario
-modelling_inputs = [Input(i['name'], i['type'], i['distribution']) for i in inputs] + \
-                   [Output(i['name'], i['type']) for i in outputs] + \
-                   [Meta(i['name'], i['type'], [i['populate']]) for i in metas] if metas else list()
+modelling_inputs = (
+    [Input(i["name"], i["datatype"], i["distribution"]) for i in inputs]
+    + [Output(i["name"], i["datatype"]) for i in outputs]
+    + [Meta(i["name"], i["datatype"], [i["populate"]]) for i in metas]
+    if metas
+    else list()
+)
 
 # Create modelling scenario to access z3 variable mirrors
 modelling_scenario = Scenario(modelling_inputs, None)
 modelling_scenario.setup_treatment_variables()
 
 mutates = {
-    "Increase": lambda x: modelling_scenario.treatment_variables[x].z3 >
-                          modelling_scenario.variables[x].z3,
-    "ChangeByFactor(2)": lambda x: modelling_scenario.treatment_variables[x].z3 ==
-                                   modelling_scenario.variables[
-                                       x].z3 * 2
+    "Increase": lambda x: modelling_scenario.treatment_variables[x].z3 > modelling_scenario.variables[x].z3,
+    "ChangeByFactor(2)": lambda x: modelling_scenario.treatment_variables[x].z3
+    == modelling_scenario.variables[x].z3 * 2,
 }
 
 
@@ -153,11 +160,32 @@ def add_modelling_assumptions(self, estimation_model: Estimator):
             estimation_model.intercept = 0
 
 
+def test_run_causal_tests():
+    ROOT = os.path.realpath(os.path.dirname(__file__))
+
+    log_path = f"{ROOT}/json_frontend.log"
+    json_path = f"{ROOT}/causal_tests.json"
+    dag_path = f"{ROOT}/dag.dot"
+    data_path = f"{ROOT}/data.csv"
+
+    json_utility = MyJsonUtility(log_path)  # Create an instance of the extended JsonUtility class
+    json_utility.set_paths(
+        json_path, dag_path, [data_path]
+    )  # Set the path to the data.csv, dag.dot and causal_tests.json file
+
+    # Load the Causal Variables into the JsonUtility class ready to be used in the tests
+    json_utility.set_variables(inputs, outputs, metas)
+    json_utility.setup()  # Sets up all the necessary parts of the json_class needed to execute tests
+
+    json_utility.generate_tests(effects, mutates, estimators, False)
+
+
 if __name__ == "__main__":
     args = MyJsonUtility.get_args()
     json_utility = MyJsonUtility(args.log_path)  # Create an instance of the extended JsonUtility class
-    json_utility.set_paths(args.json_path, args.dag_path,
-                           args.data_path)  # Set the path to the data.csv, dag.dot and causal_tests.json file
+    json_utility.set_paths(
+        args.json_path, args.dag_path, args.data_path
+    )  # Set the path to the data.csv, dag.dot and causal_tests.json file
 
     # Load the Causal Variables into the JsonUtility class ready to be used in the tests
     json_utility.set_variables(inputs, outputs, metas)