From 790c12ba6fcff01f2f702bfc544268435d4bae9a Mon Sep 17 00:00:00 2001
From: Michael <jmafoster1@gmail.com>
Date: Tue, 6 Aug 2024 15:13:44 +0100
Subject: [PATCH 01/24] GP in

---
 causal_testing/gp/gp.py                | 312 +++++++++++++++++++++++++
 causal_testing/testing/estimators.py   |  20 ++
 pyproject.toml                         |   4 +-
 tests/testing_tests/test_estimators.py |  15 ++
 4 files changed, 350 insertions(+), 1 deletion(-)
 create mode 100644 causal_testing/gp/gp.py

diff --git a/causal_testing/gp/gp.py b/causal_testing/gp/gp.py
new file mode 100644
index 00000000..00737437
--- /dev/null
+++ b/causal_testing/gp/gp.py
@@ -0,0 +1,312 @@
+import random
+import warnings
+import patsy
+
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+import statsmodels.formula.api as smf
+import statsmodels
+import sympy
+import copy
+
+from functools import partial
+from deap import algorithms, base, creator, tools, gp
+
+from numpy import negative, exp, power, log, sin, cos, tan, sinh, cosh, tanh
+from inspect import isclass
+
+from operator import add, mul
+
+
+def root(x):
+    return power(x, 0.5)
+
+
+def square(x):
+    return power(x, 2)
+
+
+def cube(x):
+    return power(x, 3)
+
+
+def fourth_power(x):
+    return power(x, 4)
+
+
+def reciprocal(x):
+    return power(x, -1)
+
+
+def mutInsert(individual, pset):
+    """
+    Copied from gp.mutInsert, except that we import isclass from inspect, so we
+    won't have the "isclass not defined" bug.
+
+    Inserts a new branch at a random position in *individual*. The subtree
+    at the chosen position is used as child node of the created subtree, in
+    that way, it is really an insertion rather than a replacement. Note that
+    the original subtree will become one of the children of the new primitive
+    inserted, but not perforce the first (its position is randomly selected if
+    the new primitive has more than one child).
+
+    :param individual: The normal or typed tree to be mutated.
+    :returns: A tuple of one tree.
+    """
+    index = random.randrange(len(individual))
+    node = individual[index]
+    slice_ = individual.searchSubtree(index)
+    choice = random.choice
+
+    # As we want to keep the current node as children of the new one,
+    # it must accept the return value of the current node
+    primitives = [p for p in pset.primitives[node.ret] if node.ret in p.args]
+
+    if len(primitives) == 0:
+        return (individual,)
+
+    new_node = choice(primitives)
+    new_subtree = [None] * len(new_node.args)
+    position = choice([i for i, a in enumerate(new_node.args) if a == node.ret])
+
+    for i, arg_type in enumerate(new_node.args):
+        if i != position:
+            term = choice(pset.terminals[arg_type])
+            if isclass(term):
+                term = term()
+            new_subtree[i] = term
+
+    new_subtree[position : position + 1] = individual[slice_]
+    new_subtree.insert(0, new_node)
+    individual[slice_] = new_subtree
+    return (individual,)
+
+
+class GP:
+
+    def __init__(
+        self,
+        df: pd.DataFrame,
+        features: list,
+        outcome: str,
+        extra_operators: list = None,
+        sympy_conversions: dict = None,
+        seed=0,
+    ):
+        random.seed(seed)
+        self.df = df
+        self.features = features
+        self.outcome = outcome
+        self.seed = seed
+        self.pset = gp.PrimitiveSet("MAIN", len(self.features))
+        self.pset.renameArguments(**{f"ARG{i}": f for i, f in enumerate(self.features)})
+
+        standard_operators = [(add, 2), (mul, 2), (reciprocal, 1)]
+        if extra_operators is None:
+            extra_operators = [(log, 1), (reciprocal, 1)]
+        for operator, num_args in standard_operators + extra_operators:
+            self.pset.addPrimitive(operator, num_args)
+        if sympy_conversions is None:
+            sympy_conversions = {}
+        self.sympy_conversions = {
+            "mul": lambda *args_: "Mul({},{})".format(*args_),
+            "add": lambda *args_: "Add({},{})".format(*args_),
+            "reciprocal": lambda *args_: "Pow({},-1)".format(*args_),
+        } | sympy_conversions
+
+        creator.create("FitnessMin", base.Fitness, weights=(-1.0,))
+        creator.create("Individual", gp.PrimitiveTree, fitness=creator.FitnessMin)
+
+        self.toolbox = base.Toolbox()
+        self.toolbox.register("expr", gp.genHalfAndHalf, pset=self.pset, min_=1, max_=2)
+        self.toolbox.register("individual", tools.initIterate, creator.Individual, self.toolbox.expr)
+        self.toolbox.register("population", tools.initRepeat, list, self.toolbox.individual)
+        self.toolbox.register("compile", gp.compile, pset=self.pset)
+        self.toolbox.register("evaluate", self.evalSymbReg)
+        self.toolbox.register("repair", self.repair)
+        self.toolbox.register("select", tools.selBest)
+        self.toolbox.register("mate", gp.cxOnePoint)
+        self.toolbox.register("expr_mut", gp.genFull, min_=0, max_=2)
+        self.toolbox.register("mutate", self.mutate, expr=self.toolbox.expr_mut)
+        self.toolbox.decorate("mate", gp.staticLimit(key=lambda x: x.height + 1, max_value=17))
+        self.toolbox.decorate("mutate", gp.staticLimit(key=lambda x: x.height + 1, max_value=17))
+
+    def split(self, individual):
+        if len(individual) > 1:
+            terms = []
+            # Recurse over children if add/sub
+            if individual[0].name in ["add", "sub"]:
+                terms.extend(
+                    self.split(
+                        creator.Individual(
+                            gp.PrimitiveTree(
+                                individual[individual.searchSubtree(1).start : individual.searchSubtree(1).stop]
+                            )
+                        )
+                    )
+                )
+                terms.extend(
+                    self.split(creator.Individual(gp.PrimitiveTree(individual[individual.searchSubtree(1).stop :])))
+                )
+            else:
+                terms.append(individual)
+            return terms
+        return [individual]
+
+    def _convert_inverse_prim(self, prim, args):
+        """
+        Convert inverse prims according to:
+        [Dd]iv(a,b) -> Mul[a, 1/b]
+        [Ss]ub(a,b) -> Add[a, -b]
+        We achieve this by overwriting the corresponding format method of the sub and div prim.
+        """
+        prim = copy.copy(prim)
+        prim_formatter = self.sympy_conversions.get(prim.name, prim.format)
+
+        return prim_formatter(*args)
+
+    def _stringify_for_sympy(self, f):
+        """Return the expression in a human readable string."""
+        string = ""
+        stack = []
+        for node in f:
+            stack.append((node, []))
+            while len(stack[-1][1]) == stack[-1][0].arity:
+                prim, args = stack.pop()
+                string = self._convert_inverse_prim(prim, args)
+                if len(stack) == 0:
+                    break  # If stack is empty, all nodes should have been seen
+                stack[-1][1].append(string)
+        return string
+
+    def simplify(self, individual):
+        return sympy.simplify(self._stringify_for_sympy(individual))
+
+    def repair(self, individual):
+        eq = f"{self.outcome} ~ {' + '.join(str(x) for x in self.split(individual))}"
+        try:
+            # Create model, fit (run) it, give estimates from it]
+            model = smf.ols(eq, self.df)
+            res = model.fit()
+            y_estimates = res.predict(self.df)
+
+            eqn = f"{res.params['Intercept']}"
+            for term, coefficient in res.params.items():
+                if term != "Intercept":
+                    eqn = f"add({eqn}, mul({coefficient}, {term}))"
+            repaired = type(individual)(gp.PrimitiveTree.from_string(eqn, self.pset))
+            return repaired
+        except (
+            OverflowError,
+            ValueError,
+            ZeroDivisionError,
+            statsmodels.tools.sm_exceptions.MissingDataError,
+            patsy.PatsyError,
+        ) as e:
+            return individual
+
+    def evalSymbReg(self, individual):
+        old_settings = np.seterr(all="raise")
+        try:
+            # Create model, fit (run) it, give estimates from it]
+            func = gp.compile(individual, self.pset)
+            y_estimates = pd.Series([func(**x) for _, x in self.df[self.features].iterrows()])
+
+            # Calc errors using an improved normalised mean squared
+            sqerrors = (self.df[self.outcome] - y_estimates) ** 2
+            mean_squared = sqerrors.sum() / len(self.df)
+            nmse = mean_squared / (self.df[self.outcome].sum() / len(self.df))
+
+            return (nmse,)
+
+            # Fitness value of infinite if error - not return 1
+        except (
+            OverflowError,
+            ValueError,
+            ZeroDivisionError,
+            statsmodels.tools.sm_exceptions.MissingDataError,
+            patsy.PatsyError,
+            RuntimeWarning,
+            FloatingPointError,
+        ) as e:
+            return (float("inf"),)
+        finally:
+            np.seterr(**old_settings)  # Restore original settings
+
+    def make_offspring(self, population, lambda_):
+        offspring = []
+        for i in range(lambda_):
+            parent1, parent2 = tools.selTournament(population, 2, 2)
+            child, _ = self.toolbox.mate(self.toolbox.clone(parent1), self.toolbox.clone(parent2))
+            del child.fitness.values
+            (child,) = self.toolbox.mutate(child)
+            offspring.append(child)
+        return offspring
+
+    def eaMuPlusLambda(self, ngen, mu=20, lambda_=10, stats=None, verbose=False, seeds=None):
+        population = [self.toolbox.repair(ind) for ind in self.toolbox.population(n=mu)]
+        if seeds is not None:
+            for seed in seeds:
+                ind = creator.Individual(gp.PrimitiveTree.from_string(seed, self.pset))
+                ind.fitness.values = self.toolbox.evaluate(ind)
+                population.append(ind)
+
+        logbook = tools.Logbook()
+        logbook.header = ["gen", "nevals"] + (stats.fields if stats else [])
+
+        # Evaluate the individuals with an invalid fitness
+        for ind in population:
+            ind.fitness.values = self.toolbox.evaluate(ind)
+        population.sort(key=lambda x: (x.fitness.values, x.height))
+
+        record = stats.compile(population) if stats is not None else {}
+        logbook.record(gen=0, nevals=len(population), **record)
+        if verbose:
+            print(logbook.stream)
+
+        # Begin the generational process
+        for gen in range(1, ngen + 1):
+            # Vary the population
+            offspring = self.make_offspring(population, lambda_)
+            offspring = [self.toolbox.repair(ind) for ind in offspring]
+
+            # Evaluate the individuals with an invalid fitness
+            for ind in offspring:
+                ind.fitness.values = self.toolbox.evaluate(ind)
+
+            # Select the next generation population
+            population[:] = self.toolbox.select(population + offspring, mu)
+
+            # Update the statistics with the new population
+            record = stats.compile(population) if stats is not None else {}
+            logbook.record(gen=gen, nevals=len(offspring), **record)
+            if verbose:
+                print(logbook.stream)
+            population.sort(key=lambda x: (x.fitness.values, x.height))
+
+        return population[0]
+
+    def mutate(self, individual, expr):
+        choice = random.randint(1, 3)
+        if choice == 1:
+            mutated = gp.mutNodeReplacement(self.toolbox.clone(individual), self.pset)
+        elif choice == 2:
+            mutated = mutInsert(self.toolbox.clone(individual), self.pset)
+        elif choice == 3:
+            mutated = gp.mutShrink(self.toolbox.clone(individual))
+        else:
+            raise ValueError("Invalid mutation choice")
+        return mutated
+
+
+if __name__ == "__main__":
+    df = pd.DataFrame()
+    df["X"] = np.arange(10)
+    df["Y"] = 1 / (df.X + 1)
+
+    gp1 = GP(df.astype(float), ["X"], "Y", seed=1)
+    best = gp1.eaMuPlusLambda(ngen=100)
+    print(best, best.fitness.values[0])
+    simplified = gp1.simplify(best)
+    print(simplified)
diff --git a/causal_testing/testing/estimators.py b/causal_testing/testing/estimators.py
index 3920e1f2..d06cc14f 100644
--- a/causal_testing/testing/estimators.py
+++ b/causal_testing/testing/estimators.py
@@ -18,6 +18,7 @@
 
 from causal_testing.specification.variable import Variable
 from causal_testing.specification.capabilities import TreatmentSequence, Capability
+from causal_testing.gp.gp import GP
 
 logger = logging.getLogger(__name__)
 
@@ -333,6 +334,19 @@ def __init__(
         for term in self.effect_modifiers:
             self.adjustment_set.add(term)
 
+    def gp_formula(self, ngen=100, mu=20, lambda_=10, extra_operators=None, sympy_conversions=None, seeds=None, seed=0):
+        gp = GP(
+            df=self.df,
+            features=sorted(list(self.adjustment_set.union([self.treatment]))),
+            outcome=self.outcome,
+            extra_operators=extra_operators,
+            sympy_conversions=sympy_conversions,
+            seed=seed,
+        )
+        formula = gp.eaMuPlusLambda(ngen=ngen, mu=mu, lambda_=lambda_, seeds=seeds)
+        formula = gp.simplify(formula)
+        self.formula = f"{self.outcome} ~ I({formula}) - 1"
+
     def add_modelling_assumptions(self):
         """
         Add modelling assumptions to the estimator. This is a list of strings which list the modelling assumptions that
@@ -421,7 +435,13 @@ def estimate_control_treatment(self, adjustment_config: dict = None) -> tuple[pd
             if str(x.dtypes[col]) == "object":
                 x = pd.get_dummies(x, columns=[col], drop_first=True)
         x = x[model.params.index]
+
+        # This is a hack for "I(...)" equations
+        x[self.treatment] = [self.treatment_value, self.control_value]
+
         y = model.get_prediction(x).summary_frame()
+        print("=== Y ===")
+        print(y)
 
         return y.iloc[1], y.iloc[0]
 
diff --git a/pyproject.toml b/pyproject.toml
index e8e9fc68..9c3133b4 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -27,7 +27,9 @@ dependencies = [
     "statsmodels~=0.14",
     "tabulate~=0.9",
     "pydot~=2.0",
-    "pygad~=3.3"
+    "pygad~=3.3",
+    "deap~=1.4.1",
+    "sympy~=1.13.1",
 ]
 dynamic = ["version"]
 
diff --git a/tests/testing_tests/test_estimators.py b/tests/testing_tests/test_estimators.py
index b7604b86..3416f4e8 100644
--- a/tests/testing_tests/test_estimators.py
+++ b/tests/testing_tests/test_estimators.py
@@ -402,6 +402,21 @@ def test_program_11_2_with_robustness_validation(self):
         cv = CausalValidator()
         self.assertEqual(round(cv.estimate_robustness(model)["treatments"], 4), 0.7353)
 
+    def test_gp(self):
+        df = pd.DataFrame()
+        df["X"] = np.arange(10)
+        df["Y"] = 1 / (df["X"] + 1)
+        linear_regression_estimator = LinearRegressionEstimator("X", 0, 1, set(), "Y", df.astype(float))
+        linear_regression_estimator.gp_formula(seed=1)
+        self.assertEqual(
+            linear_regression_estimator.formula,
+            "Y ~ I((2.606801258739728e-17*X + 0.626132756132756)/(0.6261327561327561*X + 0.626132756132756)) - 1",
+        )
+        ate, (ci_low, ci_high) = linear_regression_estimator.estimate_ate_calculated()
+        self.assertEqual(round(ate[0], 2), 0.50)
+        self.assertEqual(round(ci_low[0], 2), 0.50)
+        self.assertEqual(round(ci_high[0], 2), 0.50)
+
 
 class TestCubicSplineRegressionEstimator(TestLinearRegressionEstimator):
     @classmethod

From c6f9d3148ec3bc13d93d71dadff5205a55df87e3 Mon Sep 17 00:00:00 2001
From: Michael <jmafoster1@gmail.com>
Date: Tue, 6 Aug 2024 15:37:43 +0100
Subject: [PATCH 02/24] Cleaned estimators a bit

---
 causal_testing/testing/estimators.py | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/causal_testing/testing/estimators.py b/causal_testing/testing/estimators.py
index d06cc14f..91d06bd6 100644
--- a/causal_testing/testing/estimators.py
+++ b/causal_testing/testing/estimators.py
@@ -424,7 +424,6 @@ def estimate_control_treatment(self, adjustment_config: dict = None) -> tuple[pd
         model = self._run_linear_regression()
 
         x = pd.DataFrame(columns=self.df.columns)
-        x[self.treatment] = [self.treatment_value, self.control_value]
         x["Intercept"] = 1  # self.intercept
         for k, v in adjustment_config.items():
             x[k] = v
@@ -436,12 +435,9 @@ def estimate_control_treatment(self, adjustment_config: dict = None) -> tuple[pd
                 x = pd.get_dummies(x, columns=[col], drop_first=True)
         x = x[model.params.index]
 
-        # This is a hack for "I(...)" equations
         x[self.treatment] = [self.treatment_value, self.control_value]
 
         y = model.get_prediction(x).summary_frame()
-        print("=== Y ===")
-        print(y)
 
         return y.iloc[1], y.iloc[0]
 

From a13d83b31b15c9fdee3a5c66860d3d251e8da215 Mon Sep 17 00:00:00 2001
From: Michael <jmafoster1@gmail.com>
Date: Tue, 6 Aug 2024 16:26:40 +0100
Subject: [PATCH 03/24] Moved estimation to a separate package

---
 .../estimation/cubic_spline_estimator.py      |  75 ++
 causal_testing/estimation/estimator.py        |  87 ++
 causal_testing/{gp => estimation}/gp.py       |   0
 causal_testing/estimation/ipcw_estimator.py   | 257 ++++++
 causal_testing/estimation/iv_estimator.py     | 104 +++
 .../estimation/linear_regression_estimator.py | 221 +++++
 .../logistic_regression_estimator.py          | 223 +++++
 causal_testing/json_front/json_class.py       |   5 +-
 .../surrogate/causal_surrogate_assisted.py    |   5 +-
 .../surrogate/surrogate_search_algorithms.py  |   2 +-
 .../testing/causal_test_adequacy.py           |   2 +-
 causal_testing/testing/causal_test_case.py    |   2 +-
 causal_testing/testing/causal_test_result.py  |   2 +-
 causal_testing/testing/causal_test_suite.py   |   2 +-
 causal_testing/testing/estimators.py          | 856 ------------------
 dafni/main_dafni.py                           |  72 +-
 docs/source/usage.rst                         |   2 +-
 .../covasim_/doubling_beta/example_beta.py    |   2 +-
 .../vaccinating_elderly/example_vaccine.py    |   2 +-
 examples/lr91/example_max_conductances.py     |   2 +-
 .../example_max_conductances_test_suite.py    |   2 +-
 .../example_poisson_process.py                |   2 +-
 examples/poisson/example_run_causal_tests.py  |   4 +-
 .../test_cubic_spline_estimator.py            |  46 +
 tests/estimation_tests/test_ipcw_estimator.py |  82 ++
 tests/estimation_tests/test_iv_estimator.py   |  53 ++
 .../test_linear_regression_estimator.py}      | 255 +-----
 .../test_logistic_regression_estimator.py     |  73 ++
 tests/json_front_tests/test_json_class.py     |   2 +-
 .../test_causal_surrogate_assisted.py         |  92 +-
 .../test_causal_test_adequacy.py              |   3 +-
 tests/testing_tests/test_causal_test_case.py  |   2 +-
 .../testing_tests/test_causal_test_outcome.py |   2 +-
 tests/testing_tests/test_causal_test_suite.py |   7 +-
 34 files changed, 1348 insertions(+), 1200 deletions(-)
 create mode 100644 causal_testing/estimation/cubic_spline_estimator.py
 create mode 100644 causal_testing/estimation/estimator.py
 rename causal_testing/{gp => estimation}/gp.py (100%)
 create mode 100644 causal_testing/estimation/ipcw_estimator.py
 create mode 100644 causal_testing/estimation/iv_estimator.py
 create mode 100644 causal_testing/estimation/linear_regression_estimator.py
 create mode 100644 causal_testing/estimation/logistic_regression_estimator.py
 delete mode 100644 causal_testing/testing/estimators.py
 create mode 100644 tests/estimation_tests/test_cubic_spline_estimator.py
 create mode 100644 tests/estimation_tests/test_ipcw_estimator.py
 create mode 100644 tests/estimation_tests/test_iv_estimator.py
 rename tests/{testing_tests/test_estimators.py => estimation_tests/test_linear_regression_estimator.py} (58%)
 create mode 100644 tests/estimation_tests/test_logistic_regression_estimator.py

diff --git a/causal_testing/estimation/cubic_spline_estimator.py b/causal_testing/estimation/cubic_spline_estimator.py
new file mode 100644
index 00000000..06ee2b44
--- /dev/null
+++ b/causal_testing/estimation/cubic_spline_estimator.py
@@ -0,0 +1,75 @@
+"""This module contains the CubicSplineRegressionEstimator class, for estimating continuous outcomes with changes in behaviour"""
+
+import logging
+from abc import ABC, abstractmethod
+from typing import Any
+from math import ceil
+
+import numpy as np
+import pandas as pd
+import statsmodels.api as sm
+import statsmodels.formula.api as smf
+from patsy import dmatrix  # pylint: disable = no-name-in-module
+from patsy import ModelDesc
+from statsmodels.regression.linear_model import RegressionResultsWrapper
+from statsmodels.tools.sm_exceptions import PerfectSeparationError
+from lifelines import CoxPHFitter
+
+from causal_testing.specification.variable import Variable
+from causal_testing.specification.capabilities import TreatmentSequence, Capability
+from causal_testing.estimation.estimator import Estimator
+from causal_testing.estimation.linear_regression_estimator import LinearRegressionEstimator
+
+logger = logging.getLogger(__name__)
+
+
+class CubicSplineRegressionEstimator(LinearRegressionEstimator):
+    """A Cubic Spline Regression Estimator is a parametric estimator which restricts the variables in the data to a
+    combination of parameters and basis functions of the variables.
+    """
+
+    def __init__(
+        # pylint: disable=too-many-arguments
+        self,
+        treatment: str,
+        treatment_value: float,
+        control_value: float,
+        adjustment_set: set,
+        outcome: str,
+        basis: int,
+        df: pd.DataFrame = None,
+        effect_modifiers: dict[Variable:Any] = None,
+        formula: str = None,
+        alpha: float = 0.05,
+        expected_relationship=None,
+    ):
+        super().__init__(
+            treatment, treatment_value, control_value, adjustment_set, outcome, df, effect_modifiers, formula, alpha
+        )
+
+        self.expected_relationship = expected_relationship
+
+        if effect_modifiers is None:
+            effect_modifiers = []
+
+        if formula is None:
+            terms = [treatment] + sorted(list(adjustment_set)) + sorted(list(effect_modifiers))
+            self.formula = f"{outcome} ~ cr({'+'.join(terms)}, df={basis})"
+
+    def estimate_ate_calculated(self, adjustment_config: dict = None) -> pd.Series:
+        model = self._run_linear_regression()
+
+        x = {"Intercept": 1, self.treatment: self.treatment_value}
+        if adjustment_config is not None:
+            for k, v in adjustment_config.items():
+                x[k] = v
+        if self.effect_modifiers is not None:
+            for k, v in self.effect_modifiers.items():
+                x[k] = v
+
+        treatment = model.predict(x).iloc[0]
+
+        x[self.treatment] = self.control_value
+        control = model.predict(x).iloc[0]
+
+        return pd.Series(treatment - control)
diff --git a/causal_testing/estimation/estimator.py b/causal_testing/estimation/estimator.py
new file mode 100644
index 00000000..81876518
--- /dev/null
+++ b/causal_testing/estimation/estimator.py
@@ -0,0 +1,87 @@
+"""This module contains the Estimator abstract class"""
+
+import logging
+from abc import ABC, abstractmethod
+from typing import Any
+from math import ceil
+
+import numpy as np
+import pandas as pd
+import statsmodels.api as sm
+import statsmodels.formula.api as smf
+from patsy import dmatrix  # pylint: disable = no-name-in-module
+from patsy import ModelDesc
+from statsmodels.regression.linear_model import RegressionResultsWrapper
+from statsmodels.tools.sm_exceptions import PerfectSeparationError
+from lifelines import CoxPHFitter
+
+from causal_testing.specification.variable import Variable
+from causal_testing.specification.capabilities import TreatmentSequence, Capability
+
+logger = logging.getLogger(__name__)
+
+
+class Estimator(ABC):
+    # pylint: disable=too-many-instance-attributes
+    """An estimator contains all of the information necessary to compute a causal estimate for the effect of changing
+    a set of treatment variables to a set of values.
+
+    All estimators must implement the following two methods:
+
+    1) add_modelling_assumptions: The validity of a model-assisted causal inference result depends on whether
+    the modelling assumptions imposed by a model actually hold. Therefore, for each model, is important to state
+    the modelling assumption upon which the validity of the results depend. To achieve this, the estimator object
+    maintains a list of modelling assumptions (as strings). If a user wishes to implement their own estimator, they
+    must implement this method and add all assumptions to the list of modelling assumptions.
+
+    2) estimate_ate: All estimators must be capable of returning the average treatment effect as a minimum. That is, the
+    average effect of the intervention (changing treatment from control to treated value) on the outcome of interest
+    adjusted for all confounders.
+    """
+
+    def __init__(
+        # pylint: disable=too-many-arguments
+        self,
+        treatment: str,
+        treatment_value: float,
+        control_value: float,
+        adjustment_set: set,
+        outcome: str,
+        df: pd.DataFrame = None,
+        effect_modifiers: dict[str:Any] = None,
+        alpha: float = 0.05,
+        query: str = "",
+    ):
+        self.treatment = treatment
+        self.treatment_value = treatment_value
+        self.control_value = control_value
+        self.adjustment_set = adjustment_set
+        self.outcome = outcome
+        self.alpha = alpha
+        self.df = df.query(query) if query else df
+
+        if effect_modifiers is None:
+            self.effect_modifiers = {}
+        elif isinstance(effect_modifiers, dict):
+            self.effect_modifiers = effect_modifiers
+        else:
+            raise ValueError(f"Unsupported type for effect_modifiers {effect_modifiers}. Expected iterable")
+        self.modelling_assumptions = []
+        if query:
+            self.modelling_assumptions.append(query)
+        self.add_modelling_assumptions()
+        logger.debug("Effect Modifiers: %s", self.effect_modifiers)
+
+    @abstractmethod
+    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.
+        """
+
+    def compute_confidence_intervals(self) -> list[float, float]:
+        """
+        Estimate the 95% Wald confidence intervals for the effect of changing the treatment from control values to
+        treatment values on the outcome.
+        :return: 95% Wald confidence intervals.
+        """
diff --git a/causal_testing/gp/gp.py b/causal_testing/estimation/gp.py
similarity index 100%
rename from causal_testing/gp/gp.py
rename to causal_testing/estimation/gp.py
diff --git a/causal_testing/estimation/ipcw_estimator.py b/causal_testing/estimation/ipcw_estimator.py
new file mode 100644
index 00000000..2330c2f9
--- /dev/null
+++ b/causal_testing/estimation/ipcw_estimator.py
@@ -0,0 +1,257 @@
+"""This module contains the IPCWEstimator class, for estimating the time to a particular event"""
+
+import logging
+from abc import ABC, abstractmethod
+from typing import Any
+from math import ceil
+
+import numpy as np
+import pandas as pd
+import statsmodels.api as sm
+import statsmodels.formula.api as smf
+from patsy import dmatrix  # pylint: disable = no-name-in-module
+from patsy import ModelDesc
+from statsmodels.regression.linear_model import RegressionResultsWrapper
+from statsmodels.tools.sm_exceptions import PerfectSeparationError
+from lifelines import CoxPHFitter
+
+from causal_testing.specification.variable import Variable
+from causal_testing.specification.capabilities import TreatmentSequence, Capability
+from causal_testing.estimation.estimator import Estimator
+
+logger = logging.getLogger(__name__)
+
+
+class IPCWEstimator(Estimator):
+    """
+    Class to perform inverse probability of censoring weighting (IPCW) estimation
+    for sequences of treatments over time-varying data.
+    """
+
+    # pylint: disable=too-many-arguments
+    # pylint: disable=too-many-instance-attributes
+    def __init__(
+        self,
+        df: pd.DataFrame,
+        timesteps_per_intervention: int,
+        control_strategy: TreatmentSequence,
+        treatment_strategy: TreatmentSequence,
+        outcome: str,
+        fault_column: str,
+        fit_bl_switch_formula: str,
+        fit_bltd_switch_formula: str,
+        eligibility=None,
+        alpha: float = 0.05,
+    ):
+        super().__init__(
+            [c.variable for c in treatment_strategy.capabilities],
+            [c.value for c in treatment_strategy.capabilities],
+            [c.value for c in control_strategy.capabilities],
+            None,
+            outcome,
+            df,
+            None,
+            alpha=alpha,
+            query="",
+        )
+        self.timesteps_per_intervention = timesteps_per_intervention
+        self.control_strategy = control_strategy
+        self.treatment_strategy = treatment_strategy
+        self.outcome = outcome
+        self.fault_column = fault_column
+        self.timesteps_per_intervention = timesteps_per_intervention
+        self.fit_bl_switch_formula = fit_bl_switch_formula
+        self.fit_bltd_switch_formula = fit_bltd_switch_formula
+        self.eligibility = eligibility
+        self.df = df
+        self.preprocess_data()
+
+    def add_modelling_assumptions(self):
+        self.modelling_assumptions.append("The variables in the data vary over time.")
+
+    def setup_xo_t_do(self, strategy_assigned: list, strategy_followed: list, eligible: pd.Series):
+        """
+        Return a binary sequence with each bit representing whether the current
+        index is the time point at which the individual diverted from the
+        assigned treatment strategy (and thus should be censored).
+
+        :param strategy_assigned - the assigned treatment strategy
+        :param strategy_followed - the strategy followed by the individual
+        :param eligible - binary sequence represnting the eligibility of the individual at each time step
+        """
+        strategy_assigned = [1] + strategy_assigned + [1]
+        strategy_followed = [1] + strategy_followed + [1]
+
+        mask = (
+            pd.Series(strategy_assigned, index=eligible.index) != pd.Series(strategy_followed, index=eligible.index)
+        ).astype("boolean")
+        mask = mask | ~eligible
+        mask.reset_index(inplace=True, drop=True)
+        false = mask.loc[mask]
+        if false.empty:
+            return np.zeros(len(mask))
+        mask = (mask * 1).tolist()
+        cutoff = false.index[0] + 1
+        return mask[:cutoff] + ([None] * (len(mask) - cutoff))
+
+    def setup_fault_t_do(self, individual: pd.DataFrame):
+        """
+        Return a binary sequence with each bit representing whether the current
+        index is the time point at which the event of interest (i.e. a fault)
+        occurred.
+        """
+        fault = individual[~individual[self.fault_column]]
+        fault_t_do = pd.Series(np.zeros(len(individual)), index=individual.index)
+
+        if not fault.empty:
+            fault_time = individual["time"].loc[fault.index[0]]
+            # Ceiling to nearest observation point
+            fault_time = ceil(fault_time / self.timesteps_per_intervention) * self.timesteps_per_intervention
+            # Set the correct observation point to be the fault time of doing (fault_t_do)
+            observations = individual.loc[
+                (individual["time"] % self.timesteps_per_intervention == 0) & (individual["time"] < fault_time)
+            ]
+            if not observations.empty:
+                fault_t_do.loc[observations.index[0]] = 1
+        assert sum(fault_t_do) <= 1, f"Multiple fault times for\n{individual}"
+
+        return pd.DataFrame({"fault_t_do": fault_t_do})
+
+    def setup_fault_time(self, individual: pd.DataFrame, perturbation: float = -0.001):
+        """
+        Return the time at which the event of interest (i.e. a fault) occurred.
+        """
+        fault = individual[~individual[self.fault_column]]
+        fault_time = (
+            individual["time"].loc[fault.index[0]]
+            if not fault.empty
+            else (individual["time"].max() + self.timesteps_per_intervention)
+        )
+        return pd.DataFrame({"fault_time": np.repeat(fault_time + perturbation, len(individual))})
+
+    def preprocess_data(self):
+        """
+        Set up the treatment-specific columns in the data that are needed to estimate the hazard ratio.
+        """
+        self.df["trtrand"] = None  # treatment/control arm
+        self.df["xo_t_do"] = None  # did the individual deviate from the treatment of interest here?
+        self.df["eligible"] = self.df.eval(self.eligibility) if self.eligibility is not None else True
+
+        # when did a fault occur?
+        self.df["fault_time"] = self.df.groupby("id")[[self.fault_column, "time"]].apply(self.setup_fault_time).values
+        self.df["fault_t_do"] = (
+            self.df.groupby("id")[["id", "time", self.fault_column]].apply(self.setup_fault_t_do).values
+        )
+        assert not pd.isnull(self.df["fault_time"]).any()
+
+        living_runs = self.df.query("fault_time > 0").loc[
+            (self.df["time"] % self.timesteps_per_intervention == 0)
+            & (self.df["time"] <= self.control_strategy.total_time())
+        ]
+
+        individuals = []
+        new_id = 0
+        logging.debug("  Preprocessing groups")
+        for _, individual in living_runs.groupby("id"):
+            assert sum(individual["fault_t_do"]) <= 1, (
+                f"Error initialising fault_t_do for individual\n"
+                f"{individual[['id', 'time', 'fault_time', 'fault_t_do']]}\n"
+                "with fault at {individual.fault_time.iloc[0]}"
+            )
+
+            strategy_followed = [
+                Capability(
+                    c.variable,
+                    individual.loc[individual["time"] == c.start_time, c.variable].values[0],
+                    c.start_time,
+                    c.end_time,
+                )
+                for c in self.treatment_strategy.capabilities
+            ]
+
+            # Control flow:
+            # Individuals that start off in both arms, need cloning (hence incrementing the ID within the if statement)
+            # Individuals that don't start off in either arm are left out
+            for inx, strategy_assigned in [(0, self.control_strategy), (1, self.treatment_strategy)]:
+                if strategy_assigned.capabilities[0] == strategy_followed[0] and individual.eligible.iloc[0]:
+                    individual["id"] = new_id
+                    new_id += 1
+                    individual["trtrand"] = inx
+                    individual["xo_t_do"] = self.setup_xo_t_do(
+                        strategy_assigned.capabilities, strategy_followed, individual["eligible"]
+                    )
+                    individuals.append(individual.loc[individual["time"] <= individual["fault_time"]].copy())
+        if len(individuals) == 0:
+            raise ValueError("No individuals followed either strategy.")
+
+        self.df = pd.concat(individuals)
+
+    def estimate_hazard_ratio(self):
+        """
+        Estimate the hazard ratio.
+        """
+
+        if self.df["fault_t_do"].sum() == 0:
+            raise ValueError("No recorded faults")
+
+        preprocessed_data = self.df.loc[self.df["xo_t_do"] == 0].copy()
+
+        # Use logistic regression to predict switching given baseline covariates
+        fit_bl_switch = smf.logit(self.fit_bl_switch_formula, data=self.df).fit()
+
+        preprocessed_data["pxo1"] = fit_bl_switch.predict(preprocessed_data)
+
+        # Use logistic regression to predict switching given baseline and time-updated covariates (model S12)
+        fit_bltd_switch = smf.logit(
+            self.fit_bltd_switch_formula,
+            data=self.df,
+        ).fit()
+
+        preprocessed_data["pxo2"] = fit_bltd_switch.predict(preprocessed_data)
+
+        # IPCW step 3: For each individual at each time, compute the inverse probability of remaining uncensored
+        # Estimate the probabilities of remaining ‘un-switched’ and hence the weights
+
+        preprocessed_data["num"] = 1 - preprocessed_data["pxo1"]
+        preprocessed_data["denom"] = 1 - preprocessed_data["pxo2"]
+        preprocessed_data[["num", "denom"]] = (
+            preprocessed_data.sort_values(["id", "time"]).groupby("id")[["num", "denom"]].cumprod()
+        )
+
+        assert (
+            not preprocessed_data["num"].isnull().any()
+        ), f"{len(preprocessed_data['num'].isnull())} null numerator values"
+        assert (
+            not preprocessed_data["denom"].isnull().any()
+        ), f"{len(preprocessed_data['denom'].isnull())} null denom values"
+
+        preprocessed_data["weight"] = 1 / preprocessed_data["denom"]
+        preprocessed_data["sweight"] = preprocessed_data["num"] / preprocessed_data["denom"]
+
+        preprocessed_data["tin"] = preprocessed_data["time"]
+        preprocessed_data["tout"] = pd.concat(
+            [(preprocessed_data["time"] + self.timesteps_per_intervention), preprocessed_data["fault_time"]],
+            axis=1,
+        ).min(axis=1)
+
+        assert (preprocessed_data["tin"] <= preprocessed_data["tout"]).all(), (
+            f"Left before joining\n" f"{preprocessed_data.loc[preprocessed_data['tin'] >= preprocessed_data['tout']]}"
+        )
+
+        #  IPCW step 4: Use these weights in a weighted analysis of the outcome model
+        # Estimate the KM graph and IPCW hazard ratio using Cox regression.
+        cox_ph = CoxPHFitter()
+        cox_ph.fit(
+            df=preprocessed_data,
+            duration_col="tout",
+            event_col="fault_t_do",
+            weights_col="weight",
+            cluster_col="id",
+            robust=True,
+            formula="trtrand",
+            entry_col="tin",
+        )
+
+        ci_low, ci_high = [np.exp(cox_ph.confidence_intervals_)[col] for col in cox_ph.confidence_intervals_.columns]
+
+        return (cox_ph.hazard_ratios_, (ci_low, ci_high))
diff --git a/causal_testing/estimation/iv_estimator.py b/causal_testing/estimation/iv_estimator.py
new file mode 100644
index 00000000..39bcf604
--- /dev/null
+++ b/causal_testing/estimation/iv_estimator.py
@@ -0,0 +1,104 @@
+"""This module contains the InstrumentalVariableEstimator class, for estimating continuous outcomes with unobservable confounding."""
+
+import logging
+from abc import ABC, abstractmethod
+from typing import Any
+from math import ceil
+
+import numpy as np
+import pandas as pd
+import statsmodels.api as sm
+import statsmodels.formula.api as smf
+from patsy import dmatrix  # pylint: disable = no-name-in-module
+from patsy import ModelDesc
+from statsmodels.regression.linear_model import RegressionResultsWrapper
+from statsmodels.tools.sm_exceptions import PerfectSeparationError
+from lifelines import CoxPHFitter
+
+from causal_testing.specification.variable import Variable
+from causal_testing.specification.capabilities import TreatmentSequence, Capability
+from causal_testing.estimation.estimator import Estimator
+
+logger = logging.getLogger(__name__)
+
+
+class InstrumentalVariableEstimator(Estimator):
+    """
+    Carry out estimation using instrumental variable adjustment rather than conventional adjustment. This means we do
+    not need to observe all confounders in order to adjust for them. A key assumption here is linearity.
+    """
+
+    def __init__(
+        # pylint: disable=too-many-arguments
+        self,
+        treatment: str,
+        treatment_value: float,
+        control_value: float,
+        adjustment_set: set,
+        outcome: str,
+        instrument: str,
+        df: pd.DataFrame = None,
+        intercept: int = 1,
+        effect_modifiers: dict = None,  # Not used (yet?). Needed for compatibility
+        alpha: float = 0.05,
+        query: str = "",
+    ):
+        super().__init__(
+            treatment=treatment,
+            treatment_value=treatment_value,
+            control_value=control_value,
+            adjustment_set=adjustment_set,
+            outcome=outcome,
+            df=df,
+            effect_modifiers=None,
+            alpha=alpha,
+            query=query,
+        )
+        self.intercept = intercept
+        self.model = None
+        self.instrument = instrument
+
+    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.append(
+            """The instrument and the treatment, and the treatment and the outcome must be
+        related linearly in the form Y = aX + b."""
+        )
+        self.modelling_assumptions.append(
+            """The three IV conditions must hold
+            (i) Instrument is associated with treatment
+            (ii) Instrument does not affect outcome except through its potential effect on treatment
+            (iii) Instrument and outcome do not share causes
+        """
+        )
+
+    def estimate_iv_coefficient(self, df) -> float:
+        """
+        Estimate the linear regression coefficient of the treatment on the
+        outcome.
+        """
+        # Estimate the total effect of instrument I on outcome Y = abI + c1
+        ab = sm.OLS(df[self.outcome], df[[self.instrument]]).fit().params[self.instrument]
+
+        # Estimate the direct effect of instrument I on treatment X = aI + c1
+        a = sm.OLS(df[self.treatment], df[[self.instrument]]).fit().params[self.instrument]
+
+        # Estimate the coefficient of I on X by cancelling
+        return ab / a
+
+    def estimate_coefficient(self, bootstrap_size=100) -> tuple[pd.Series, list[pd.Series, pd.Series]]:
+        """
+        Estimate the unit ate (i.e. coefficient) of the treatment on the
+        outcome.
+        """
+        bootstraps = sorted(
+            [self.estimate_iv_coefficient(self.df.sample(len(self.df), replace=True)) for _ in range(bootstrap_size)]
+        )
+        bound = ceil((bootstrap_size * self.alpha) / 2)
+        ci_low = pd.Series(bootstraps[bound])
+        ci_high = pd.Series(bootstraps[bootstrap_size - bound])
+
+        return pd.Series(self.estimate_iv_coefficient(self.df)), [ci_low, ci_high]
diff --git a/causal_testing/estimation/linear_regression_estimator.py b/causal_testing/estimation/linear_regression_estimator.py
new file mode 100644
index 00000000..e212dca7
--- /dev/null
+++ b/causal_testing/estimation/linear_regression_estimator.py
@@ -0,0 +1,221 @@
+"""This module contains the LinearRegressionEstimator for estimating continuous outcomes."""
+
+import logging
+from abc import ABC, abstractmethod
+from typing import Any
+from math import ceil
+
+import numpy as np
+import pandas as pd
+import statsmodels.api as sm
+import statsmodels.formula.api as smf
+from patsy import dmatrix  # pylint: disable = no-name-in-module
+from patsy import ModelDesc
+from statsmodels.regression.linear_model import RegressionResultsWrapper
+from statsmodels.tools.sm_exceptions import PerfectSeparationError
+from lifelines import CoxPHFitter
+
+from causal_testing.specification.variable import Variable
+from causal_testing.specification.capabilities import TreatmentSequence, Capability
+from causal_testing.estimation.gp import GP
+from causal_testing.estimation.estimator import Estimator
+
+logger = logging.getLogger(__name__)
+
+
+class LinearRegressionEstimator(Estimator):
+    """A Linear Regression Estimator is a parametric estimator which restricts the variables in the data to a linear
+    combination of parameters and functions of the variables (note these functions need not be linear).
+    """
+
+    def __init__(
+        # pylint: disable=too-many-arguments
+        self,
+        treatment: str,
+        treatment_value: float,
+        control_value: float,
+        adjustment_set: set,
+        outcome: str,
+        df: pd.DataFrame = None,
+        effect_modifiers: dict[Variable:Any] = None,
+        formula: str = None,
+        alpha: float = 0.05,
+        query: str = "",
+    ):
+        super().__init__(
+            treatment,
+            treatment_value,
+            control_value,
+            adjustment_set,
+            outcome,
+            df,
+            effect_modifiers,
+            alpha=alpha,
+            query=query,
+        )
+
+        self.model = None
+        if effect_modifiers is None:
+            effect_modifiers = []
+
+        if formula is not None:
+            self.formula = formula
+        else:
+            terms = [treatment] + sorted(list(adjustment_set)) + sorted(list(effect_modifiers))
+            self.formula = f"{outcome} ~ {'+'.join(terms)}"
+
+        for term in self.effect_modifiers:
+            self.adjustment_set.add(term)
+
+    def gp_formula(self, ngen=100, mu=20, lambda_=10, extra_operators=None, sympy_conversions=None, seeds=None, seed=0):
+        gp = GP(
+            df=self.df,
+            features=sorted(list(self.adjustment_set.union([self.treatment]))),
+            outcome=self.outcome,
+            extra_operators=extra_operators,
+            sympy_conversions=sympy_conversions,
+            seed=seed,
+        )
+        formula = gp.eaMuPlusLambda(ngen=ngen, mu=mu, lambda_=lambda_, seeds=seeds)
+        formula = gp.simplify(formula)
+        self.formula = f"{self.outcome} ~ I({formula}) - 1"
+
+    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.append(
+            "The variables in the data must fit a shape which can be expressed as a linear"
+            "combination of parameters and functions of variables. Note that these functions"
+            "do not need to be linear."
+        )
+
+    def estimate_coefficient(self) -> tuple[pd.Series, list[pd.Series, pd.Series]]:
+        """Estimate the unit average treatment effect of the treatment on the outcome. That is, the change in outcome
+        caused by a unit change in treatment.
+
+        :return: The unit average treatment effect and the 95% Wald confidence intervals.
+        """
+        model = self._run_linear_regression()
+        newline = "\n"
+        patsy_md = ModelDesc.from_formula(self.treatment)
+
+        if any(
+            (
+                self.df.dtypes[factor.name()] == "object"
+                for factor in patsy_md.rhs_termlist[1].factors
+                # We want to remove this long term as it prevents us from discovering categoricals within I(...) blocks
+                if factor.name() in self.df.dtypes
+            )
+        ):
+            design_info = dmatrix(self.formula.split("~")[1], self.df).design_info
+            treatment = design_info.column_names[design_info.term_name_slices[self.treatment]]
+        else:
+            treatment = [self.treatment]
+        assert set(treatment).issubset(
+            model.params.index.tolist()
+        ), f"{treatment} not in\n{'  ' + str(model.params.index).replace(newline, newline + '  ')}"
+        unit_effect = model.params[treatment]  # Unit effect is the coefficient of the treatment
+        [ci_low, ci_high] = self._get_confidence_intervals(model, treatment)
+        return unit_effect, [ci_low, ci_high]
+
+    def estimate_ate(self) -> tuple[pd.Series, list[pd.Series, pd.Series]]:
+        """Estimate the 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 average treatment effect and the 95% Wald confidence intervals.
+        """
+        model = self._run_linear_regression()
+
+        # Create an empty individual for the control and treated
+        individuals = pd.DataFrame(1, index=["control", "treated"], columns=model.params.index)
+
+        # For Pandas version > 2, we need to explicitly state that the dataframe takes floating-point values
+        individuals = individuals.astype(float)
+
+        # It is ABSOLUTELY CRITICAL that these go last, otherwise we can't index
+        # the effect with "ate = t_test_results.effect[0]"
+        individuals.loc["control", [self.treatment]] = self.control_value
+        individuals.loc["treated", [self.treatment]] = self.treatment_value
+
+        # Perform a t-test to compare the predicted outcome of the control and treated individual (ATE)
+        t_test_results = model.t_test(individuals.loc["treated"] - individuals.loc["control"])
+        ate = pd.Series(t_test_results.effect[0])
+        confidence_intervals = list(t_test_results.conf_int(alpha=self.alpha).flatten())
+        confidence_intervals = [pd.Series(interval) for interval in confidence_intervals]
+        return ate, confidence_intervals
+
+    def estimate_control_treatment(self, adjustment_config: dict = None) -> tuple[pd.Series, pd.Series]:
+        """Estimate the outcomes under control and treatment.
+
+        :return: The estimated outcome under control and treatment in the form
+        (control_outcome, treatment_outcome).
+        """
+        if adjustment_config is None:
+            adjustment_config = {}
+        model = self._run_linear_regression()
+
+        x = pd.DataFrame(columns=self.df.columns)
+        x[self.treatment] = [self.treatment_value, self.control_value]
+        x["Intercept"] = 1  # self.intercept
+        for k, v in adjustment_config.items():
+            x[k] = v
+        for k, v in self.effect_modifiers.items():
+            x[k] = v
+        x = dmatrix(self.formula.split("~")[1], x, return_type="dataframe")
+        for col in x:
+            if str(x.dtypes[col]) == "object":
+                x = pd.get_dummies(x, columns=[col], drop_first=True)
+        x = x[model.params.index]
+
+        x[self.treatment] = [self.treatment_value, self.control_value]
+
+        y = model.get_prediction(x).summary_frame()
+
+        return y.iloc[1], y.iloc[0]
+
+    def estimate_risk_ratio(self, adjustment_config: dict = None) -> tuple[pd.Series, list[pd.Series, pd.Series]]:
+        """Estimate the risk_ratio 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 average treatment effect and the 95% Wald confidence intervals.
+        """
+        if adjustment_config is None:
+            adjustment_config = {}
+        control_outcome, treatment_outcome = self.estimate_control_treatment(adjustment_config=adjustment_config)
+        ci_low = pd.Series(treatment_outcome["mean_ci_lower"] / control_outcome["mean_ci_upper"])
+        ci_high = pd.Series(treatment_outcome["mean_ci_upper"] / control_outcome["mean_ci_lower"])
+        return pd.Series(treatment_outcome["mean"] / control_outcome["mean"]), [ci_low, ci_high]
+
+    def estimate_ate_calculated(self, adjustment_config: dict = None) -> tuple[pd.Series, list[pd.Series, pd.Series]]:
+        """Estimate the ate 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. Here, we actually
+        calculate the expected outcomes under control and treatment and divide one by the other. This
+        allows for custom terms to be put in such as squares, inverses, products, etc.
+
+        :return: The average treatment effect and the 95% Wald confidence intervals.
+        """
+        if adjustment_config is None:
+            adjustment_config = {}
+        control_outcome, treatment_outcome = self.estimate_control_treatment(adjustment_config=adjustment_config)
+        ci_low = pd.Series(treatment_outcome["mean_ci_lower"] - control_outcome["mean_ci_upper"])
+        ci_high = pd.Series(treatment_outcome["mean_ci_upper"] - control_outcome["mean_ci_lower"])
+        return pd.Series(treatment_outcome["mean"] - control_outcome["mean"]), [ci_low, ci_high]
+
+    def _run_linear_regression(self) -> RegressionResultsWrapper:
+        """Run linear regression of the treatment and adjustment set against the outcome and return the model.
+
+        :return: The model after fitting to data.
+        """
+        model = smf.ols(formula=self.formula, data=self.df).fit()
+        self.model = model
+        return model
+
+    def _get_confidence_intervals(self, model, treatment):
+        confidence_intervals = model.conf_int(alpha=self.alpha, cols=None)
+        ci_low, ci_high = (
+            pd.Series(confidence_intervals[0].loc[treatment]),
+            pd.Series(confidence_intervals[1].loc[treatment]),
+        )
+        return [ci_low, ci_high]
diff --git a/causal_testing/estimation/logistic_regression_estimator.py b/causal_testing/estimation/logistic_regression_estimator.py
new file mode 100644
index 00000000..399d02e2
--- /dev/null
+++ b/causal_testing/estimation/logistic_regression_estimator.py
@@ -0,0 +1,223 @@
+"""This module contains the LogisticRegressionEstimator class for estimating categorical outcomes."""
+
+import logging
+from abc import ABC, abstractmethod
+from typing import Any
+from math import ceil
+
+import numpy as np
+import pandas as pd
+import statsmodels.api as sm
+import statsmodels.formula.api as smf
+from patsy import dmatrix  # pylint: disable = no-name-in-module
+from patsy import ModelDesc
+from statsmodels.regression.linear_model import RegressionResultsWrapper
+from statsmodels.tools.sm_exceptions import PerfectSeparationError
+from lifelines import CoxPHFitter
+
+from causal_testing.specification.variable import Variable
+from causal_testing.specification.capabilities import TreatmentSequence, Capability
+from causal_testing.estimation.estimator import Estimator
+
+logger = logging.getLogger(__name__)
+
+
+class LogisticRegressionEstimator(Estimator):
+    """A Logistic Regression Estimator is a parametric estimator which restricts the variables in the data to a linear
+    combination of parameters and functions of the variables (note these functions need not be linear). It is designed
+    for estimating categorical outcomes.
+    """
+
+    def __init__(
+        # pylint: disable=too-many-arguments
+        self,
+        treatment: str,
+        treatment_value: float,
+        control_value: float,
+        adjustment_set: set,
+        outcome: str,
+        df: pd.DataFrame = None,
+        effect_modifiers: dict[str:Any] = None,
+        formula: str = None,
+        query: str = "",
+    ):
+        super().__init__(
+            treatment=treatment,
+            treatment_value=treatment_value,
+            control_value=control_value,
+            adjustment_set=adjustment_set,
+            outcome=outcome,
+            df=df,
+            effect_modifiers=effect_modifiers,
+            query=query,
+        )
+
+        self.model = None
+
+        if formula is not None:
+            self.formula = formula
+        else:
+            terms = [treatment] + sorted(list(adjustment_set)) + sorted(list(self.effect_modifiers))
+            self.formula = f"{outcome} ~ {'+'.join(((terms)))}"
+
+    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.append(
+            "The variables in the data must fit a shape which can be expressed as a linear"
+            "combination of parameters and functions of variables. Note that these functions"
+            "do not need to be linear."
+        )
+        self.modelling_assumptions.append("The outcome must be binary.")
+        self.modelling_assumptions.append("Independently and identically distributed errors.")
+
+    def _run_logistic_regression(self, data) -> RegressionResultsWrapper:
+        """Run logistic regression of the treatment and adjustment set against the outcome and return the model.
+
+        :return: The model after fitting to data.
+        """
+        model = smf.logit(formula=self.formula, data=data).fit(disp=0)
+        self.model = model
+        return model
+
+    def estimate(self, data: pd.DataFrame, adjustment_config: dict = None) -> RegressionResultsWrapper:
+        """add terms to the dataframe and estimate the outcome from the data
+        :param data: A pandas dataframe containing execution data from the system-under-test.
+        :param adjustment_config: Dictionary containing the adjustment configuration of the adjustment set
+        """
+        if adjustment_config is None:
+            adjustment_config = {}
+        if set(self.adjustment_set) != set(adjustment_config):
+            raise ValueError(
+                f"Invalid adjustment configuration {adjustment_config}. Must specify values for {self.adjustment_set}"
+            )
+
+        model = self._run_logistic_regression(data)
+
+        x = pd.DataFrame(columns=self.df.columns)
+        x["Intercept"] = 1  # self.intercept
+        x[self.treatment] = [self.treatment_value, self.control_value]
+        for k, v in adjustment_config.items():
+            x[k] = v
+        for k, v in self.effect_modifiers.items():
+            x[k] = v
+        x = dmatrix(self.formula.split("~")[1], x, return_type="dataframe")
+        for col in x:
+            if str(x.dtypes[col]) == "object":
+                x = pd.get_dummies(x, columns=[col], drop_first=True)
+        # x = x[model.params.index]
+        return model.predict(x)
+
+    def estimate_control_treatment(
+        self, adjustment_config: dict = None, bootstrap_size: int = 100
+    ) -> tuple[pd.Series, pd.Series]:
+        """Estimate the outcomes under control and treatment.
+
+        :return: The estimated control and treatment values and their confidence
+        intervals in the form ((ci_low, control, ci_high), (ci_low, treatment, ci_high)).
+        """
+        if adjustment_config is None:
+            adjustment_config = {}
+        y = self.estimate(self.df, adjustment_config=adjustment_config)
+
+        try:
+            bootstrap_samples = [
+                self.estimate(self.df.sample(len(self.df), replace=True), adjustment_config=adjustment_config)
+                for _ in range(bootstrap_size)
+            ]
+            control, treatment = zip(*[(x.iloc[1], x.iloc[0]) for x in bootstrap_samples])
+        except PerfectSeparationError:
+            logger.warning(
+                "Perfect separation detected, results not available. Cannot calculate confidence intervals for such "
+                "a small dataset."
+            )
+            return (y.iloc[1], None), (y.iloc[0], None)
+        except np.linalg.LinAlgError:
+            logger.warning("Singular matrix detected. Confidence intervals not available. Try with a larger data set")
+            return (y.iloc[1], None), (y.iloc[0], None)
+
+        # Delta method confidence intervals from
+        # https://stackoverflow.com/questions/47414842/confidence-interval-of-probability-prediction-from-logistic-regression-statsmode
+        # cov = model.cov_params()
+        # gradient = (y * (1 - y) * x.T).T  # matrix of gradients for each observation
+        # std_errors = np.array([np.sqrt(np.dot(np.dot(g, cov), g)) for g in gradient.to_numpy()])
+        # c = 1.96  # multiplier for confidence interval
+        # upper = np.maximum(0, np.minimum(1, y + std_errors * c))
+        # lower = np.maximum(0, np.minimum(1, y - std_errors * c))
+
+        return (y.iloc[1], np.array(control)), (y.iloc[0], np.array(treatment))
+
+    def estimate_ate(self, adjustment_config: dict = None, bootstrap_size: int = 100) -> float:
+        """Estimate the ate 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. Here, we actually
+        calculate the expected outcomes under control and treatment and take one away from the other. This
+        allows for custom terms to be put in such as squares, inverses, products, etc.
+
+        :return: The estimated average treatment effect and 95% confidence intervals
+        """
+        if adjustment_config is None:
+            adjustment_config = {}
+        (control_outcome, control_bootstraps), (
+            treatment_outcome,
+            treatment_bootstraps,
+        ) = self.estimate_control_treatment(bootstrap_size=bootstrap_size, adjustment_config=adjustment_config)
+        estimate = treatment_outcome - control_outcome
+
+        if control_bootstraps is None or treatment_bootstraps is None:
+            return estimate, (None, None)
+
+        bootstraps = sorted(list(treatment_bootstraps - control_bootstraps))
+        bound = int((bootstrap_size * self.alpha) / 2)
+        ci_low = bootstraps[bound]
+        ci_high = bootstraps[bootstrap_size - bound]
+
+        logger.info(
+            f"Changing {self.treatment} from {self.control_value} to {self.treatment_value} gives an estimated "
+            f"ATE of {ci_low} < {estimate} < {ci_high}"
+        )
+        assert ci_low < estimate < ci_high, f"Expecting {ci_low} < {estimate} < {ci_high}"
+
+        return estimate, (ci_low, ci_high)
+
+    def estimate_risk_ratio(self, adjustment_config: dict = None, bootstrap_size: int = 100) -> float:
+        """Estimate the ate 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. Here, we actually
+        calculate the expected outcomes under control and treatment and divide one by the other. This
+        allows for custom terms to be put in such as squares, inverses, products, etc.
+
+        :return: The estimated risk ratio and 95% confidence intervals.
+        """
+        if adjustment_config is None:
+            adjustment_config = {}
+        (control_outcome, control_bootstraps), (
+            treatment_outcome,
+            treatment_bootstraps,
+        ) = self.estimate_control_treatment(bootstrap_size=bootstrap_size, adjustment_config=adjustment_config)
+        estimate = treatment_outcome / control_outcome
+
+        if control_bootstraps is None or treatment_bootstraps is None:
+            return estimate, (None, None)
+
+        bootstraps = sorted(list(treatment_bootstraps / control_bootstraps))
+        bound = ceil((bootstrap_size * self.alpha) / 2)
+        ci_low = bootstraps[bound]
+        ci_high = bootstraps[bootstrap_size - bound]
+
+        logger.info(
+            f"Changing {self.treatment} from {self.control_value} to {self.treatment_value} gives an estimated "
+            f"risk ratio of {ci_low} < {estimate} < {ci_high}"
+        )
+        assert ci_low < estimate < ci_high, f"Expecting {ci_low} < {estimate} < {ci_high}"
+
+        return estimate, (ci_low, ci_high)
+
+    def estimate_unit_odds_ratio(self) -> float:
+        """Estimate the odds ratio of increasing the treatment by one. In logistic regression, this corresponds to the
+        coefficient of the treatment of interest.
+
+        :return: The odds ratio. Confidence intervals are not yet supported.
+        """
+        model = self._run_logistic_regression(self.df)
+        return np.exp(model.params[self.treatment])
diff --git a/causal_testing/json_front/json_class.py b/causal_testing/json_front/json_class.py
index 8c6c265c..9a0de6e7 100644
--- a/causal_testing/json_front/json_class.py
+++ b/causal_testing/json_front/json_class.py
@@ -22,10 +22,13 @@
 from causal_testing.specification.variable import Input, Meta, Output
 from causal_testing.testing.causal_test_case import CausalTestCase
 from causal_testing.testing.causal_test_result import CausalTestResult
-from causal_testing.testing.estimators import Estimator, LinearRegressionEstimator, LogisticRegressionEstimator
 from causal_testing.testing.base_test_case import BaseTestCase
 from causal_testing.testing.causal_test_adequacy import DataAdequacy
 
+from causal_testing.estimation.estimator import Estimator
+from causal_testing.estimation.linear_regression_estimator import LinearRegressionEstimator
+from causal_testing.estimation.logistic_regression_estimator import LogisticRegressionEstimator
+
 logger = logging.getLogger(__name__)
 
 
diff --git a/causal_testing/surrogate/causal_surrogate_assisted.py b/causal_testing/surrogate/causal_surrogate_assisted.py
index c30b2086..4fba5371 100644
--- a/causal_testing/surrogate/causal_surrogate_assisted.py
+++ b/causal_testing/surrogate/causal_surrogate_assisted.py
@@ -7,7 +7,8 @@
 from causal_testing.data_collection.data_collector import ObservationalDataCollector
 from causal_testing.specification.causal_specification import CausalSpecification
 from causal_testing.testing.base_test_case import BaseTestCase
-from causal_testing.testing.estimators import CubicSplineRegressionEstimator
+from causal_testing.estimation.cubic_spline_estimator import CubicSplineRegressionEstimator
+
 
 @dataclass
 class SimulationResult:
@@ -19,7 +20,7 @@ class SimulationResult:
 
     def to_dataframe(self) -> pd.DataFrame:
         """Convert the simulation result data to a pandas DataFrame"""
-        data_as_lists = {k: v if isinstance(v, list) else [v] for k,v in self.data.items()}
+        data_as_lists = {k: v if isinstance(v, list) else [v] for k, v in self.data.items()}
         return pd.DataFrame(data_as_lists)
 
 
diff --git a/causal_testing/surrogate/surrogate_search_algorithms.py b/causal_testing/surrogate/surrogate_search_algorithms.py
index 495c2f86..54e7bb48 100644
--- a/causal_testing/surrogate/surrogate_search_algorithms.py
+++ b/causal_testing/surrogate/surrogate_search_algorithms.py
@@ -6,7 +6,7 @@
 from pygad import GA
 
 from causal_testing.specification.causal_specification import CausalSpecification
-from causal_testing.testing.estimators import CubicSplineRegressionEstimator
+from causal_testing.estimation.cubic_spline_estimator import CubicSplineRegressionEstimator
 from causal_testing.surrogate.causal_surrogate_assisted import SearchAlgorithm
 
 
diff --git a/causal_testing/testing/causal_test_adequacy.py b/causal_testing/testing/causal_test_adequacy.py
index 740bba5e..003ac714 100644
--- a/causal_testing/testing/causal_test_adequacy.py
+++ b/causal_testing/testing/causal_test_adequacy.py
@@ -11,7 +11,7 @@
 
 from causal_testing.testing.causal_test_suite import CausalTestSuite
 from causal_testing.specification.causal_dag import CausalDAG
-from causal_testing.testing.estimators import Estimator
+from causal_testing.estimation.estimator import Estimator
 from causal_testing.testing.causal_test_case import CausalTestCase
 
 logger = logging.getLogger(__name__)
diff --git a/causal_testing/testing/causal_test_case.py b/causal_testing/testing/causal_test_case.py
index 52daf040..a648d7a9 100644
--- a/causal_testing/testing/causal_test_case.py
+++ b/causal_testing/testing/causal_test_case.py
@@ -7,7 +7,7 @@
 from causal_testing.specification.variable import Variable
 from causal_testing.testing.causal_test_outcome import CausalTestOutcome
 from causal_testing.testing.base_test_case import BaseTestCase
-from causal_testing.testing.estimators import Estimator
+from causal_testing.estimation.estimator import Estimator
 from causal_testing.testing.causal_test_result import CausalTestResult, TestValue
 from causal_testing.data_collection.data_collector import DataCollector
 
diff --git a/causal_testing/testing/causal_test_result.py b/causal_testing/testing/causal_test_result.py
index 65a2085e..11c6790d 100644
--- a/causal_testing/testing/causal_test_result.py
+++ b/causal_testing/testing/causal_test_result.py
@@ -6,7 +6,7 @@
 from dataclasses import dataclass
 import pandas as pd
 
-from causal_testing.testing.estimators import Estimator
+from causal_testing.estimation.estimator import Estimator
 from causal_testing.specification.variable import Variable
 
 
diff --git a/causal_testing/testing/causal_test_suite.py b/causal_testing/testing/causal_test_suite.py
index 47c5ef98..062cd7e5 100644
--- a/causal_testing/testing/causal_test_suite.py
+++ b/causal_testing/testing/causal_test_suite.py
@@ -7,7 +7,7 @@
 from typing import Type, Iterable
 from causal_testing.testing.base_test_case import BaseTestCase
 from causal_testing.testing.causal_test_case import CausalTestCase
-from causal_testing.testing.estimators import Estimator
+from causal_testing.estimation.estimator import Estimator
 from causal_testing.testing.causal_test_result import CausalTestResult
 from causal_testing.data_collection.data_collector import DataCollector
 from causal_testing.specification.causal_specification import CausalSpecification
diff --git a/causal_testing/testing/estimators.py b/causal_testing/testing/estimators.py
deleted file mode 100644
index 91d06bd6..00000000
--- a/causal_testing/testing/estimators.py
+++ /dev/null
@@ -1,856 +0,0 @@
-"""This module contains the Estimator abstract class, as well as its concrete extensions: LogisticRegressionEstimator,
-LinearRegressionEstimator"""
-
-import logging
-from abc import ABC, abstractmethod
-from typing import Any
-from math import ceil
-
-import numpy as np
-import pandas as pd
-import statsmodels.api as sm
-import statsmodels.formula.api as smf
-from patsy import dmatrix  # pylint: disable = no-name-in-module
-from patsy import ModelDesc
-from statsmodels.regression.linear_model import RegressionResultsWrapper
-from statsmodels.tools.sm_exceptions import PerfectSeparationError
-from lifelines import CoxPHFitter
-
-from causal_testing.specification.variable import Variable
-from causal_testing.specification.capabilities import TreatmentSequence, Capability
-from causal_testing.gp.gp import GP
-
-logger = logging.getLogger(__name__)
-
-
-class Estimator(ABC):
-    # pylint: disable=too-many-instance-attributes
-    """An estimator contains all of the information necessary to compute a causal estimate for the effect of changing
-    a set of treatment variables to a set of values.
-
-    All estimators must implement the following two methods:
-
-    1) add_modelling_assumptions: The validity of a model-assisted causal inference result depends on whether
-    the modelling assumptions imposed by a model actually hold. Therefore, for each model, is important to state
-    the modelling assumption upon which the validity of the results depend. To achieve this, the estimator object
-    maintains a list of modelling assumptions (as strings). If a user wishes to implement their own estimator, they
-    must implement this method and add all assumptions to the list of modelling assumptions.
-
-    2) estimate_ate: All estimators must be capable of returning the average treatment effect as a minimum. That is, the
-    average effect of the intervention (changing treatment from control to treated value) on the outcome of interest
-    adjusted for all confounders.
-    """
-
-    def __init__(
-        # pylint: disable=too-many-arguments
-        self,
-        treatment: str,
-        treatment_value: float,
-        control_value: float,
-        adjustment_set: set,
-        outcome: str,
-        df: pd.DataFrame = None,
-        effect_modifiers: dict[str:Any] = None,
-        alpha: float = 0.05,
-        query: str = "",
-    ):
-        self.treatment = treatment
-        self.treatment_value = treatment_value
-        self.control_value = control_value
-        self.adjustment_set = adjustment_set
-        self.outcome = outcome
-        self.alpha = alpha
-        self.df = df.query(query) if query else df
-
-        if effect_modifiers is None:
-            self.effect_modifiers = {}
-        elif isinstance(effect_modifiers, dict):
-            self.effect_modifiers = effect_modifiers
-        else:
-            raise ValueError(f"Unsupported type for effect_modifiers {effect_modifiers}. Expected iterable")
-        self.modelling_assumptions = []
-        if query:
-            self.modelling_assumptions.append(query)
-        self.add_modelling_assumptions()
-        logger.debug("Effect Modifiers: %s", self.effect_modifiers)
-
-    @abstractmethod
-    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.
-        """
-
-    def compute_confidence_intervals(self) -> list[float, float]:
-        """
-        Estimate the 95% Wald confidence intervals for the effect of changing the treatment from control values to
-        treatment values on the outcome.
-        :return: 95% Wald confidence intervals.
-        """
-
-
-class LogisticRegressionEstimator(Estimator):
-    """A Logistic Regression Estimator is a parametric estimator which restricts the variables in the data to a linear
-    combination of parameters and functions of the variables (note these functions need not be linear). It is designed
-    for estimating categorical outcomes.
-    """
-
-    def __init__(
-        # pylint: disable=too-many-arguments
-        self,
-        treatment: str,
-        treatment_value: float,
-        control_value: float,
-        adjustment_set: set,
-        outcome: str,
-        df: pd.DataFrame = None,
-        effect_modifiers: dict[str:Any] = None,
-        formula: str = None,
-        query: str = "",
-    ):
-        super().__init__(
-            treatment=treatment,
-            treatment_value=treatment_value,
-            control_value=control_value,
-            adjustment_set=adjustment_set,
-            outcome=outcome,
-            df=df,
-            effect_modifiers=effect_modifiers,
-            query=query,
-        )
-
-        self.model = None
-
-        if formula is not None:
-            self.formula = formula
-        else:
-            terms = [treatment] + sorted(list(adjustment_set)) + sorted(list(self.effect_modifiers))
-            self.formula = f"{outcome} ~ {'+'.join(((terms)))}"
-
-    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.append(
-            "The variables in the data must fit a shape which can be expressed as a linear"
-            "combination of parameters and functions of variables. Note that these functions"
-            "do not need to be linear."
-        )
-        self.modelling_assumptions.append("The outcome must be binary.")
-        self.modelling_assumptions.append("Independently and identically distributed errors.")
-
-    def _run_logistic_regression(self, data) -> RegressionResultsWrapper:
-        """Run logistic regression of the treatment and adjustment set against the outcome and return the model.
-
-        :return: The model after fitting to data.
-        """
-        model = smf.logit(formula=self.formula, data=data).fit(disp=0)
-        self.model = model
-        return model
-
-    def estimate(self, data: pd.DataFrame, adjustment_config: dict = None) -> RegressionResultsWrapper:
-        """add terms to the dataframe and estimate the outcome from the data
-        :param data: A pandas dataframe containing execution data from the system-under-test.
-        :param adjustment_config: Dictionary containing the adjustment configuration of the adjustment set
-        """
-        if adjustment_config is None:
-            adjustment_config = {}
-        if set(self.adjustment_set) != set(adjustment_config):
-            raise ValueError(
-                f"Invalid adjustment configuration {adjustment_config}. Must specify values for {self.adjustment_set}"
-            )
-
-        model = self._run_logistic_regression(data)
-
-        x = pd.DataFrame(columns=self.df.columns)
-        x["Intercept"] = 1  # self.intercept
-        x[self.treatment] = [self.treatment_value, self.control_value]
-        for k, v in adjustment_config.items():
-            x[k] = v
-        for k, v in self.effect_modifiers.items():
-            x[k] = v
-        x = dmatrix(self.formula.split("~")[1], x, return_type="dataframe")
-        for col in x:
-            if str(x.dtypes[col]) == "object":
-                x = pd.get_dummies(x, columns=[col], drop_first=True)
-        # x = x[model.params.index]
-        return model.predict(x)
-
-    def estimate_control_treatment(
-        self, adjustment_config: dict = None, bootstrap_size: int = 100
-    ) -> tuple[pd.Series, pd.Series]:
-        """Estimate the outcomes under control and treatment.
-
-        :return: The estimated control and treatment values and their confidence
-        intervals in the form ((ci_low, control, ci_high), (ci_low, treatment, ci_high)).
-        """
-        if adjustment_config is None:
-            adjustment_config = {}
-        y = self.estimate(self.df, adjustment_config=adjustment_config)
-
-        try:
-            bootstrap_samples = [
-                self.estimate(self.df.sample(len(self.df), replace=True), adjustment_config=adjustment_config)
-                for _ in range(bootstrap_size)
-            ]
-            control, treatment = zip(*[(x.iloc[1], x.iloc[0]) for x in bootstrap_samples])
-        except PerfectSeparationError:
-            logger.warning(
-                "Perfect separation detected, results not available. Cannot calculate confidence intervals for such "
-                "a small dataset."
-            )
-            return (y.iloc[1], None), (y.iloc[0], None)
-        except np.linalg.LinAlgError:
-            logger.warning("Singular matrix detected. Confidence intervals not available. Try with a larger data set")
-            return (y.iloc[1], None), (y.iloc[0], None)
-
-        # Delta method confidence intervals from
-        # https://stackoverflow.com/questions/47414842/confidence-interval-of-probability-prediction-from-logistic-regression-statsmode
-        # cov = model.cov_params()
-        # gradient = (y * (1 - y) * x.T).T  # matrix of gradients for each observation
-        # std_errors = np.array([np.sqrt(np.dot(np.dot(g, cov), g)) for g in gradient.to_numpy()])
-        # c = 1.96  # multiplier for confidence interval
-        # upper = np.maximum(0, np.minimum(1, y + std_errors * c))
-        # lower = np.maximum(0, np.minimum(1, y - std_errors * c))
-
-        return (y.iloc[1], np.array(control)), (y.iloc[0], np.array(treatment))
-
-    def estimate_ate(self, adjustment_config: dict = None, bootstrap_size: int = 100) -> float:
-        """Estimate the ate 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. Here, we actually
-        calculate the expected outcomes under control and treatment and take one away from the other. This
-        allows for custom terms to be put in such as squares, inverses, products, etc.
-
-        :return: The estimated average treatment effect and 95% confidence intervals
-        """
-        if adjustment_config is None:
-            adjustment_config = {}
-        (control_outcome, control_bootstraps), (
-            treatment_outcome,
-            treatment_bootstraps,
-        ) = self.estimate_control_treatment(bootstrap_size=bootstrap_size, adjustment_config=adjustment_config)
-        estimate = treatment_outcome - control_outcome
-
-        if control_bootstraps is None or treatment_bootstraps is None:
-            return estimate, (None, None)
-
-        bootstraps = sorted(list(treatment_bootstraps - control_bootstraps))
-        bound = int((bootstrap_size * self.alpha) / 2)
-        ci_low = bootstraps[bound]
-        ci_high = bootstraps[bootstrap_size - bound]
-
-        logger.info(
-            f"Changing {self.treatment} from {self.control_value} to {self.treatment_value} gives an estimated "
-            f"ATE of {ci_low} < {estimate} < {ci_high}"
-        )
-        assert ci_low < estimate < ci_high, f"Expecting {ci_low} < {estimate} < {ci_high}"
-
-        return estimate, (ci_low, ci_high)
-
-    def estimate_risk_ratio(self, adjustment_config: dict = None, bootstrap_size: int = 100) -> float:
-        """Estimate the ate 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. Here, we actually
-        calculate the expected outcomes under control and treatment and divide one by the other. This
-        allows for custom terms to be put in such as squares, inverses, products, etc.
-
-        :return: The estimated risk ratio and 95% confidence intervals.
-        """
-        if adjustment_config is None:
-            adjustment_config = {}
-        (control_outcome, control_bootstraps), (
-            treatment_outcome,
-            treatment_bootstraps,
-        ) = self.estimate_control_treatment(bootstrap_size=bootstrap_size, adjustment_config=adjustment_config)
-        estimate = treatment_outcome / control_outcome
-
-        if control_bootstraps is None or treatment_bootstraps is None:
-            return estimate, (None, None)
-
-        bootstraps = sorted(list(treatment_bootstraps / control_bootstraps))
-        bound = ceil((bootstrap_size * self.alpha) / 2)
-        ci_low = bootstraps[bound]
-        ci_high = bootstraps[bootstrap_size - bound]
-
-        logger.info(
-            f"Changing {self.treatment} from {self.control_value} to {self.treatment_value} gives an estimated "
-            f"risk ratio of {ci_low} < {estimate} < {ci_high}"
-        )
-        assert ci_low < estimate < ci_high, f"Expecting {ci_low} < {estimate} < {ci_high}"
-
-        return estimate, (ci_low, ci_high)
-
-    def estimate_unit_odds_ratio(self) -> float:
-        """Estimate the odds ratio of increasing the treatment by one. In logistic regression, this corresponds to the
-        coefficient of the treatment of interest.
-
-        :return: The odds ratio. Confidence intervals are not yet supported.
-        """
-        model = self._run_logistic_regression(self.df)
-        return np.exp(model.params[self.treatment])
-
-
-class LinearRegressionEstimator(Estimator):
-    """A Linear Regression Estimator is a parametric estimator which restricts the variables in the data to a linear
-    combination of parameters and functions of the variables (note these functions need not be linear).
-    """
-
-    def __init__(
-        # pylint: disable=too-many-arguments
-        self,
-        treatment: str,
-        treatment_value: float,
-        control_value: float,
-        adjustment_set: set,
-        outcome: str,
-        df: pd.DataFrame = None,
-        effect_modifiers: dict[Variable:Any] = None,
-        formula: str = None,
-        alpha: float = 0.05,
-        query: str = "",
-    ):
-        super().__init__(
-            treatment,
-            treatment_value,
-            control_value,
-            adjustment_set,
-            outcome,
-            df,
-            effect_modifiers,
-            alpha=alpha,
-            query=query,
-        )
-
-        self.model = None
-        if effect_modifiers is None:
-            effect_modifiers = []
-
-        if formula is not None:
-            self.formula = formula
-        else:
-            terms = [treatment] + sorted(list(adjustment_set)) + sorted(list(effect_modifiers))
-            self.formula = f"{outcome} ~ {'+'.join(terms)}"
-
-        for term in self.effect_modifiers:
-            self.adjustment_set.add(term)
-
-    def gp_formula(self, ngen=100, mu=20, lambda_=10, extra_operators=None, sympy_conversions=None, seeds=None, seed=0):
-        gp = GP(
-            df=self.df,
-            features=sorted(list(self.adjustment_set.union([self.treatment]))),
-            outcome=self.outcome,
-            extra_operators=extra_operators,
-            sympy_conversions=sympy_conversions,
-            seed=seed,
-        )
-        formula = gp.eaMuPlusLambda(ngen=ngen, mu=mu, lambda_=lambda_, seeds=seeds)
-        formula = gp.simplify(formula)
-        self.formula = f"{self.outcome} ~ I({formula}) - 1"
-
-    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.append(
-            "The variables in the data must fit a shape which can be expressed as a linear"
-            "combination of parameters and functions of variables. Note that these functions"
-            "do not need to be linear."
-        )
-
-    def estimate_coefficient(self) -> tuple[pd.Series, list[pd.Series, pd.Series]]:
-        """Estimate the unit average treatment effect of the treatment on the outcome. That is, the change in outcome
-        caused by a unit change in treatment.
-
-        :return: The unit average treatment effect and the 95% Wald confidence intervals.
-        """
-        model = self._run_linear_regression()
-        newline = "\n"
-        patsy_md = ModelDesc.from_formula(self.treatment)
-
-        if any(
-            (
-                self.df.dtypes[factor.name()] == "object"
-                for factor in patsy_md.rhs_termlist[1].factors
-                # We want to remove this long term as it prevents us from discovering categoricals within I(...) blocks
-                if factor.name() in self.df.dtypes
-            )
-        ):
-            design_info = dmatrix(self.formula.split("~")[1], self.df).design_info
-            treatment = design_info.column_names[design_info.term_name_slices[self.treatment]]
-        else:
-            treatment = [self.treatment]
-        assert set(treatment).issubset(
-            model.params.index.tolist()
-        ), f"{treatment} not in\n{'  ' + str(model.params.index).replace(newline, newline + '  ')}"
-        unit_effect = model.params[treatment]  # Unit effect is the coefficient of the treatment
-        [ci_low, ci_high] = self._get_confidence_intervals(model, treatment)
-        return unit_effect, [ci_low, ci_high]
-
-    def estimate_ate(self) -> tuple[pd.Series, list[pd.Series, pd.Series]]:
-        """Estimate the 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 average treatment effect and the 95% Wald confidence intervals.
-        """
-        model = self._run_linear_regression()
-
-        # Create an empty individual for the control and treated
-        individuals = pd.DataFrame(1, index=["control", "treated"], columns=model.params.index)
-
-        # For Pandas version > 2, we need to explicitly state that the dataframe takes floating-point values
-        individuals = individuals.astype(float)
-
-        # It is ABSOLUTELY CRITICAL that these go last, otherwise we can't index
-        # the effect with "ate = t_test_results.effect[0]"
-        individuals.loc["control", [self.treatment]] = self.control_value
-        individuals.loc["treated", [self.treatment]] = self.treatment_value
-
-        # Perform a t-test to compare the predicted outcome of the control and treated individual (ATE)
-        t_test_results = model.t_test(individuals.loc["treated"] - individuals.loc["control"])
-        ate = pd.Series(t_test_results.effect[0])
-        confidence_intervals = list(t_test_results.conf_int(alpha=self.alpha).flatten())
-        confidence_intervals = [pd.Series(interval) for interval in confidence_intervals]
-        return ate, confidence_intervals
-
-    def estimate_control_treatment(self, adjustment_config: dict = None) -> tuple[pd.Series, pd.Series]:
-        """Estimate the outcomes under control and treatment.
-
-        :return: The estimated outcome under control and treatment in the form
-        (control_outcome, treatment_outcome).
-        """
-        if adjustment_config is None:
-            adjustment_config = {}
-        model = self._run_linear_regression()
-
-        x = pd.DataFrame(columns=self.df.columns)
-        x["Intercept"] = 1  # self.intercept
-        for k, v in adjustment_config.items():
-            x[k] = v
-        for k, v in self.effect_modifiers.items():
-            x[k] = v
-        x = dmatrix(self.formula.split("~")[1], x, return_type="dataframe")
-        for col in x:
-            if str(x.dtypes[col]) == "object":
-                x = pd.get_dummies(x, columns=[col], drop_first=True)
-        x = x[model.params.index]
-
-        x[self.treatment] = [self.treatment_value, self.control_value]
-
-        y = model.get_prediction(x).summary_frame()
-
-        return y.iloc[1], y.iloc[0]
-
-    def estimate_risk_ratio(self, adjustment_config: dict = None) -> tuple[pd.Series, list[pd.Series, pd.Series]]:
-        """Estimate the risk_ratio 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 average treatment effect and the 95% Wald confidence intervals.
-        """
-        if adjustment_config is None:
-            adjustment_config = {}
-        control_outcome, treatment_outcome = self.estimate_control_treatment(adjustment_config=adjustment_config)
-        ci_low = pd.Series(treatment_outcome["mean_ci_lower"] / control_outcome["mean_ci_upper"])
-        ci_high = pd.Series(treatment_outcome["mean_ci_upper"] / control_outcome["mean_ci_lower"])
-        return pd.Series(treatment_outcome["mean"] / control_outcome["mean"]), [ci_low, ci_high]
-
-    def estimate_ate_calculated(self, adjustment_config: dict = None) -> tuple[pd.Series, list[pd.Series, pd.Series]]:
-        """Estimate the ate 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. Here, we actually
-        calculate the expected outcomes under control and treatment and divide one by the other. This
-        allows for custom terms to be put in such as squares, inverses, products, etc.
-
-        :return: The average treatment effect and the 95% Wald confidence intervals.
-        """
-        if adjustment_config is None:
-            adjustment_config = {}
-        control_outcome, treatment_outcome = self.estimate_control_treatment(adjustment_config=adjustment_config)
-        ci_low = pd.Series(treatment_outcome["mean_ci_lower"] - control_outcome["mean_ci_upper"])
-        ci_high = pd.Series(treatment_outcome["mean_ci_upper"] - control_outcome["mean_ci_lower"])
-        return pd.Series(treatment_outcome["mean"] - control_outcome["mean"]), [ci_low, ci_high]
-
-    def _run_linear_regression(self) -> RegressionResultsWrapper:
-        """Run linear regression of the treatment and adjustment set against the outcome and return the model.
-
-        :return: The model after fitting to data.
-        """
-        model = smf.ols(formula=self.formula, data=self.df).fit()
-        self.model = model
-        return model
-
-    def _get_confidence_intervals(self, model, treatment):
-        confidence_intervals = model.conf_int(alpha=self.alpha, cols=None)
-        ci_low, ci_high = (
-            pd.Series(confidence_intervals[0].loc[treatment]),
-            pd.Series(confidence_intervals[1].loc[treatment]),
-        )
-        return [ci_low, ci_high]
-
-
-class CubicSplineRegressionEstimator(LinearRegressionEstimator):
-    """A Cubic Spline Regression Estimator is a parametric estimator which restricts the variables in the data to a
-    combination of parameters and basis functions of the variables.
-    """
-
-    def __init__(
-        # pylint: disable=too-many-arguments
-        self,
-        treatment: str,
-        treatment_value: float,
-        control_value: float,
-        adjustment_set: set,
-        outcome: str,
-        basis: int,
-        df: pd.DataFrame = None,
-        effect_modifiers: dict[Variable:Any] = None,
-        formula: str = None,
-        alpha: float = 0.05,
-        expected_relationship=None,
-    ):
-        super().__init__(
-            treatment, treatment_value, control_value, adjustment_set, outcome, df, effect_modifiers, formula, alpha
-        )
-
-        self.expected_relationship = expected_relationship
-
-        if effect_modifiers is None:
-            effect_modifiers = []
-
-        if formula is None:
-            terms = [treatment] + sorted(list(adjustment_set)) + sorted(list(effect_modifiers))
-            self.formula = f"{outcome} ~ cr({'+'.join(terms)}, df={basis})"
-
-    def estimate_ate_calculated(self, adjustment_config: dict = None) -> pd.Series:
-        model = self._run_linear_regression()
-
-        x = {"Intercept": 1, self.treatment: self.treatment_value}
-        if adjustment_config is not None:
-            for k, v in adjustment_config.items():
-                x[k] = v
-        if self.effect_modifiers is not None:
-            for k, v in self.effect_modifiers.items():
-                x[k] = v
-
-        treatment = model.predict(x).iloc[0]
-
-        x[self.treatment] = self.control_value
-        control = model.predict(x).iloc[0]
-
-        return pd.Series(treatment - control)
-
-
-class InstrumentalVariableEstimator(Estimator):
-    """
-    Carry out estimation using instrumental variable adjustment rather than conventional adjustment. This means we do
-    not need to observe all confounders in order to adjust for them. A key assumption here is linearity.
-    """
-
-    def __init__(
-        # pylint: disable=too-many-arguments
-        self,
-        treatment: str,
-        treatment_value: float,
-        control_value: float,
-        adjustment_set: set,
-        outcome: str,
-        instrument: str,
-        df: pd.DataFrame = None,
-        intercept: int = 1,
-        effect_modifiers: dict = None,  # Not used (yet?). Needed for compatibility
-        alpha: float = 0.05,
-        query: str = "",
-    ):
-        super().__init__(
-            treatment=treatment,
-            treatment_value=treatment_value,
-            control_value=control_value,
-            adjustment_set=adjustment_set,
-            outcome=outcome,
-            df=df,
-            effect_modifiers=None,
-            alpha=alpha,
-            query=query,
-        )
-        self.intercept = intercept
-        self.model = None
-        self.instrument = instrument
-
-    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.append(
-            """The instrument and the treatment, and the treatment and the outcome must be
-        related linearly in the form Y = aX + b."""
-        )
-        self.modelling_assumptions.append(
-            """The three IV conditions must hold
-            (i) Instrument is associated with treatment
-            (ii) Instrument does not affect outcome except through its potential effect on treatment
-            (iii) Instrument and outcome do not share causes
-        """
-        )
-
-    def estimate_iv_coefficient(self, df) -> float:
-        """
-        Estimate the linear regression coefficient of the treatment on the
-        outcome.
-        """
-        # Estimate the total effect of instrument I on outcome Y = abI + c1
-        ab = sm.OLS(df[self.outcome], df[[self.instrument]]).fit().params[self.instrument]
-
-        # Estimate the direct effect of instrument I on treatment X = aI + c1
-        a = sm.OLS(df[self.treatment], df[[self.instrument]]).fit().params[self.instrument]
-
-        # Estimate the coefficient of I on X by cancelling
-        return ab / a
-
-    def estimate_coefficient(self, bootstrap_size=100) -> tuple[pd.Series, list[pd.Series, pd.Series]]:
-        """
-        Estimate the unit ate (i.e. coefficient) of the treatment on the
-        outcome.
-        """
-        bootstraps = sorted(
-            [self.estimate_iv_coefficient(self.df.sample(len(self.df), replace=True)) for _ in range(bootstrap_size)]
-        )
-        bound = ceil((bootstrap_size * self.alpha) / 2)
-        ci_low = pd.Series(bootstraps[bound])
-        ci_high = pd.Series(bootstraps[bootstrap_size - bound])
-
-        return pd.Series(self.estimate_iv_coefficient(self.df)), [ci_low, ci_high]
-
-
-class IPCWEstimator(Estimator):
-    """
-    Class to perform inverse probability of censoring weighting (IPCW) estimation
-    for sequences of treatments over time-varying data.
-    """
-
-    # pylint: disable=too-many-arguments
-    # pylint: disable=too-many-instance-attributes
-    def __init__(
-        self,
-        df: pd.DataFrame,
-        timesteps_per_intervention: int,
-        control_strategy: TreatmentSequence,
-        treatment_strategy: TreatmentSequence,
-        outcome: str,
-        fault_column: str,
-        fit_bl_switch_formula: str,
-        fit_bltd_switch_formula: str,
-        eligibility=None,
-        alpha: float = 0.05,
-    ):
-        super().__init__(
-            [c.variable for c in treatment_strategy.capabilities],
-            [c.value for c in treatment_strategy.capabilities],
-            [c.value for c in control_strategy.capabilities],
-            None,
-            outcome,
-            df,
-            None,
-            alpha=alpha,
-            query="",
-        )
-        self.timesteps_per_intervention = timesteps_per_intervention
-        self.control_strategy = control_strategy
-        self.treatment_strategy = treatment_strategy
-        self.outcome = outcome
-        self.fault_column = fault_column
-        self.timesteps_per_intervention = timesteps_per_intervention
-        self.fit_bl_switch_formula = fit_bl_switch_formula
-        self.fit_bltd_switch_formula = fit_bltd_switch_formula
-        self.eligibility = eligibility
-        self.df = df
-        self.preprocess_data()
-
-    def add_modelling_assumptions(self):
-        self.modelling_assumptions.append("The variables in the data vary over time.")
-
-    def setup_xo_t_do(self, strategy_assigned: list, strategy_followed: list, eligible: pd.Series):
-        """
-        Return a binary sequence with each bit representing whether the current
-        index is the time point at which the individual diverted from the
-        assigned treatment strategy (and thus should be censored).
-
-        :param strategy_assigned - the assigned treatment strategy
-        :param strategy_followed - the strategy followed by the individual
-        :param eligible - binary sequence represnting the eligibility of the individual at each time step
-        """
-        strategy_assigned = [1] + strategy_assigned + [1]
-        strategy_followed = [1] + strategy_followed + [1]
-
-        mask = (
-            pd.Series(strategy_assigned, index=eligible.index) != pd.Series(strategy_followed, index=eligible.index)
-        ).astype("boolean")
-        mask = mask | ~eligible
-        mask.reset_index(inplace=True, drop=True)
-        false = mask.loc[mask]
-        if false.empty:
-            return np.zeros(len(mask))
-        mask = (mask * 1).tolist()
-        cutoff = false.index[0] + 1
-        return mask[:cutoff] + ([None] * (len(mask) - cutoff))
-
-    def setup_fault_t_do(self, individual: pd.DataFrame):
-        """
-        Return a binary sequence with each bit representing whether the current
-        index is the time point at which the event of interest (i.e. a fault)
-        occurred.
-        """
-        fault = individual[~individual[self.fault_column]]
-        fault_t_do = pd.Series(np.zeros(len(individual)), index=individual.index)
-
-        if not fault.empty:
-            fault_time = individual["time"].loc[fault.index[0]]
-            # Ceiling to nearest observation point
-            fault_time = ceil(fault_time / self.timesteps_per_intervention) * self.timesteps_per_intervention
-            # Set the correct observation point to be the fault time of doing (fault_t_do)
-            observations = individual.loc[
-                (individual["time"] % self.timesteps_per_intervention == 0) & (individual["time"] < fault_time)
-            ]
-            if not observations.empty:
-                fault_t_do.loc[observations.index[0]] = 1
-        assert sum(fault_t_do) <= 1, f"Multiple fault times for\n{individual}"
-
-        return pd.DataFrame({"fault_t_do": fault_t_do})
-
-    def setup_fault_time(self, individual: pd.DataFrame, perturbation: float = -0.001):
-        """
-        Return the time at which the event of interest (i.e. a fault) occurred.
-        """
-        fault = individual[~individual[self.fault_column]]
-        fault_time = (
-            individual["time"].loc[fault.index[0]]
-            if not fault.empty
-            else (individual["time"].max() + self.timesteps_per_intervention)
-        )
-        return pd.DataFrame({"fault_time": np.repeat(fault_time + perturbation, len(individual))})
-
-    def preprocess_data(self):
-        """
-        Set up the treatment-specific columns in the data that are needed to estimate the hazard ratio.
-        """
-        self.df["trtrand"] = None  # treatment/control arm
-        self.df["xo_t_do"] = None  # did the individual deviate from the treatment of interest here?
-        self.df["eligible"] = self.df.eval(self.eligibility) if self.eligibility is not None else True
-
-        # when did a fault occur?
-        self.df["fault_time"] = self.df.groupby("id")[[self.fault_column, "time"]].apply(self.setup_fault_time).values
-        self.df["fault_t_do"] = (
-            self.df.groupby("id")[["id", "time", self.fault_column]].apply(self.setup_fault_t_do).values
-        )
-        assert not pd.isnull(self.df["fault_time"]).any()
-
-        living_runs = self.df.query("fault_time > 0").loc[
-            (self.df["time"] % self.timesteps_per_intervention == 0)
-            & (self.df["time"] <= self.control_strategy.total_time())
-        ]
-
-        individuals = []
-        new_id = 0
-        logging.debug("  Preprocessing groups")
-        for _, individual in living_runs.groupby("id"):
-            assert sum(individual["fault_t_do"]) <= 1, (
-                f"Error initialising fault_t_do for individual\n"
-                f"{individual[['id', 'time', 'fault_time', 'fault_t_do']]}\n"
-                "with fault at {individual.fault_time.iloc[0]}"
-            )
-
-            strategy_followed = [
-                Capability(
-                    c.variable,
-                    individual.loc[individual["time"] == c.start_time, c.variable].values[0],
-                    c.start_time,
-                    c.end_time,
-                )
-                for c in self.treatment_strategy.capabilities
-            ]
-
-            # Control flow:
-            # Individuals that start off in both arms, need cloning (hence incrementing the ID within the if statement)
-            # Individuals that don't start off in either arm are left out
-            for inx, strategy_assigned in [(0, self.control_strategy), (1, self.treatment_strategy)]:
-                if strategy_assigned.capabilities[0] == strategy_followed[0] and individual.eligible.iloc[0]:
-                    individual["id"] = new_id
-                    new_id += 1
-                    individual["trtrand"] = inx
-                    individual["xo_t_do"] = self.setup_xo_t_do(
-                        strategy_assigned.capabilities, strategy_followed, individual["eligible"]
-                    )
-                    individuals.append(individual.loc[individual["time"] <= individual["fault_time"]].copy())
-        if len(individuals) == 0:
-            raise ValueError("No individuals followed either strategy.")
-
-        self.df = pd.concat(individuals)
-
-    def estimate_hazard_ratio(self):
-        """
-        Estimate the hazard ratio.
-        """
-
-        if self.df["fault_t_do"].sum() == 0:
-            raise ValueError("No recorded faults")
-
-        preprocessed_data = self.df.loc[self.df["xo_t_do"] == 0].copy()
-
-        # Use logistic regression to predict switching given baseline covariates
-        fit_bl_switch = smf.logit(self.fit_bl_switch_formula, data=self.df).fit()
-
-        preprocessed_data["pxo1"] = fit_bl_switch.predict(preprocessed_data)
-
-        # Use logistic regression to predict switching given baseline and time-updated covariates (model S12)
-        fit_bltd_switch = smf.logit(
-            self.fit_bltd_switch_formula,
-            data=self.df,
-        ).fit()
-
-        preprocessed_data["pxo2"] = fit_bltd_switch.predict(preprocessed_data)
-
-        # IPCW step 3: For each individual at each time, compute the inverse probability of remaining uncensored
-        # Estimate the probabilities of remaining ‘un-switched’ and hence the weights
-
-        preprocessed_data["num"] = 1 - preprocessed_data["pxo1"]
-        preprocessed_data["denom"] = 1 - preprocessed_data["pxo2"]
-        preprocessed_data[["num", "denom"]] = (
-            preprocessed_data.sort_values(["id", "time"]).groupby("id")[["num", "denom"]].cumprod()
-        )
-
-        assert (
-            not preprocessed_data["num"].isnull().any()
-        ), f"{len(preprocessed_data['num'].isnull())} null numerator values"
-        assert (
-            not preprocessed_data["denom"].isnull().any()
-        ), f"{len(preprocessed_data['denom'].isnull())} null denom values"
-
-        preprocessed_data["weight"] = 1 / preprocessed_data["denom"]
-        preprocessed_data["sweight"] = preprocessed_data["num"] / preprocessed_data["denom"]
-
-        preprocessed_data["tin"] = preprocessed_data["time"]
-        preprocessed_data["tout"] = pd.concat(
-            [(preprocessed_data["time"] + self.timesteps_per_intervention), preprocessed_data["fault_time"]],
-            axis=1,
-        ).min(axis=1)
-
-        assert (preprocessed_data["tin"] <= preprocessed_data["tout"]).all(), (
-            f"Left before joining\n" f"{preprocessed_data.loc[preprocessed_data['tin'] >= preprocessed_data['tout']]}"
-        )
-
-        #  IPCW step 4: Use these weights in a weighted analysis of the outcome model
-        # Estimate the KM graph and IPCW hazard ratio using Cox regression.
-        cox_ph = CoxPHFitter()
-        cox_ph.fit(
-            df=preprocessed_data,
-            duration_col="tout",
-            event_col="fault_t_do",
-            weights_col="weight",
-            cluster_col="id",
-            robust=True,
-            formula="trtrand",
-            entry_col="tin",
-        )
-
-        ci_low, ci_high = [np.exp(cox_ph.confidence_intervals_)[col] for col in cox_ph.confidence_intervals_.columns]
-
-        return (cox_ph.hazard_ratios_, (ci_low, ci_high))
diff --git a/dafni/main_dafni.py b/dafni/main_dafni.py
index e6b142f3..5fa66b0a 100644
--- a/dafni/main_dafni.py
+++ b/dafni/main_dafni.py
@@ -11,7 +11,7 @@
 from causal_testing.specification.scenario import Scenario
 from causal_testing.specification.variable import Input, Output
 from causal_testing.testing.causal_test_outcome import Positive, Negative, NoEffect, SomeEffect
-from causal_testing.testing.estimators import LinearRegressionEstimator
+from causal_testing.estimation.linear_regression_estimator import LinearRegressionEstimator
 from causal_testing.json_front.json_class import JsonUtility
 
 
@@ -30,35 +30,38 @@ def get_args(test_args=None) -> argparse.Namespace:
     """
     parser = argparse.ArgumentParser(description="A script for running the CTF on DAFNI.")
 
-    parser.add_argument(
-        "--data_path", required=True,
-        help="Path to the input runtime data (.csv)", nargs="+")
+    parser.add_argument("--data_path", required=True, help="Path to the input runtime data (.csv)", nargs="+")
 
-    parser.add_argument('--tests_path', required=True,
-                        help='Input configuration file path '
-                             'containing the causal tests (.json)')
+    parser.add_argument(
+        "--tests_path", required=True, help="Input configuration file path " "containing the causal tests (.json)"
+    )
 
-    parser.add_argument('--variables_path', required=True,
-                        help='Input configuration file path '
-                             'containing the predefined variables (.json)')
+    parser.add_argument(
+        "--variables_path",
+        required=True,
+        help="Input configuration file path " "containing the predefined variables (.json)",
+    )
 
-    parser.add_argument("--dag_path", required=True,
-                        help="Input configuration file path containing a valid DAG (.dot). "
-                             "Note: this must be supplied if the --tests argument isn't provided.")
+    parser.add_argument(
+        "--dag_path",
+        required=True,
+        help="Input configuration file path containing a valid DAG (.dot). "
+        "Note: this must be supplied if the --tests argument isn't provided.",
+    )
 
-    parser.add_argument('--output_path', required=False, help='Path to the output directory.')
+    parser.add_argument("--output_path", required=False, help="Path to the output directory.")
 
     parser.add_argument(
-        "-f",
-        default=False,
-        help="(Optional) Failure flag to step the framework from running if a test has failed.")
+        "-f", default=False, help="(Optional) Failure flag to step the framework from running if a test has failed."
+    )
 
     parser.add_argument(
         "-w",
         default=False,
         help="(Optional) Specify to overwrite any existing output files. "
-             "This can lead to the loss of existing outputs if not "
-             "careful")
+        "This can lead to the loss of existing outputs if not "
+        "careful",
+    )
 
     args = parser.parse_args(test_args)
 
@@ -100,7 +103,7 @@ def read_variables(variables_path: Path) -> FileNotFoundError | dict:
 
         raise FileNotFoundError
 
-    with variables_path.open('r') as file:
+    with variables_path.open("r") as file:
 
         inputs = json.load(file)
 
@@ -118,13 +121,17 @@ def validate_variables(data_dict: dict) -> tuple:
 
         variables = data_dict["variables"]
 
-        inputs = [Input(variable["name"], eval(variable["datatype"]))
-                  for variable in variables if
-                  variable["typestring"] == "Input"]
+        inputs = [
+            Input(variable["name"], eval(variable["datatype"]))
+            for variable in variables
+            if variable["typestring"] == "Input"
+        ]
 
-        outputs = [Output(variable["name"], eval(variable["datatype"]))
-                   for variable in variables if
-                   variable["typestring"] == "Output"]
+        outputs = [
+            Output(variable["name"], eval(variable["datatype"]))
+            for variable in variables
+            if variable["typestring"] == "Output"
+        ]
 
         constraints = set()
 
@@ -172,7 +179,8 @@ def main():
             "Positive": Positive(),
             "Negative": Negative(),
             "NoEffect": NoEffect(),
-            "SomeEffect": SomeEffect()}
+            "SomeEffect": SomeEffect(),
+        }
 
         # Step 4: Call the JSONUtility class to perform the causal tests
 
@@ -185,9 +193,9 @@ def main():
         json_utility.setup(scenario=modelling_scenario, data=data_frame)
 
         # Step 7: Run the causal tests
-        test_outcomes = json_utility.run_json_tests(effects=expected_outcome_effects,
-                                                    mutates={}, estimators=estimators,
-                                                    f_flag=args.f)
+        test_outcomes = json_utility.run_json_tests(
+            effects=expected_outcome_effects, mutates={}, estimators=estimators, f_flag=args.f
+        )
 
         # Step 8: Update, print and save the final outputs
 
@@ -201,7 +209,6 @@ def main():
 
             test["result"].pop("control_value")
 
-
         with open(args.output_path, "w", encoding="utf-8") as f:
 
             print(json.dumps(test_outcomes, indent=2), file=f)
@@ -214,8 +221,7 @@ def main():
 
     else:
 
-        print(f"Execution successful. "
-              f"Output file saved at {Path(args.output_path).parent.resolve()}")
+        print(f"Execution successful. " f"Output file saved at {Path(args.output_path).parent.resolve()}")
 
 
 if __name__ == "__main__":
diff --git a/docs/source/usage.rst b/docs/source/usage.rst
index 7f1e701d..f2f95bf8 100644
--- a/docs/source/usage.rst
+++ b/docs/source/usage.rst
@@ -87,7 +87,7 @@ the `documentation <https://causal-testing-framework.readthedocs.io/en/latest/au
 
 .. code-block:: python
 
-   from causal_testing.testing.estimators import LinearRegressionEstimator
+   from causal_testing.estimation.linear_regression_estimator import LinearRegressionEstimator
 
    minimal_adjustment_set = causal_dag.identification(base_test_case)
    estimation_model =  LinearRegressionEstimator(treatment=treatment, control=control, treatment_value=1, control_value=0, adjustment_set = minimal_adjustment_set, df = obs_df)
diff --git a/examples/covasim_/doubling_beta/example_beta.py b/examples/covasim_/doubling_beta/example_beta.py
index d4ae15c0..15152b7c 100644
--- a/examples/covasim_/doubling_beta/example_beta.py
+++ b/examples/covasim_/doubling_beta/example_beta.py
@@ -9,7 +9,7 @@
 from causal_testing.data_collection.data_collector import ObservationalDataCollector
 from causal_testing.testing.causal_test_case import CausalTestCase
 from causal_testing.testing.causal_test_outcome import Positive
-from causal_testing.testing.estimators import LinearRegressionEstimator
+from causal_testing.estimation.linear_regression_estimator import LinearRegressionEstimator
 from causal_testing.testing.base_test_case import BaseTestCase
 from matplotlib.pyplot import rcParams
 
diff --git a/examples/covasim_/vaccinating_elderly/example_vaccine.py b/examples/covasim_/vaccinating_elderly/example_vaccine.py
index 6742409a..6c4ba0d2 100644
--- a/examples/covasim_/vaccinating_elderly/example_vaccine.py
+++ b/examples/covasim_/vaccinating_elderly/example_vaccine.py
@@ -9,7 +9,7 @@
 from causal_testing.data_collection.data_collector import ObservationalDataCollector
 from causal_testing.testing.causal_test_case import CausalTestCase
 from causal_testing.testing.causal_test_outcome import Positive, Negative, NoEffect
-from causal_testing.testing.estimators import LinearRegressionEstimator
+from causal_testing.estimation.linear_regression_estimator import LinearRegressionEstimator
 from causal_testing.testing.base_test_case import BaseTestCase
 
 
diff --git a/examples/lr91/example_max_conductances.py b/examples/lr91/example_max_conductances.py
index bef94f04..d235abee 100644
--- a/examples/lr91/example_max_conductances.py
+++ b/examples/lr91/example_max_conductances.py
@@ -8,7 +8,7 @@
 from causal_testing.data_collection.data_collector import ObservationalDataCollector
 from causal_testing.testing.causal_test_case import CausalTestCase
 from causal_testing.testing.causal_test_outcome import Positive, Negative, NoEffect
-from causal_testing.testing.estimators import LinearRegressionEstimator
+from causal_testing.estimation.linear_regression_estimator import LinearRegressionEstimator
 from causal_testing.testing.base_test_case import BaseTestCase
 from matplotlib.pyplot import rcParams
 
diff --git a/examples/lr91/example_max_conductances_test_suite.py b/examples/lr91/example_max_conductances_test_suite.py
index fe9c1e2c..79b49ad7 100644
--- a/examples/lr91/example_max_conductances_test_suite.py
+++ b/examples/lr91/example_max_conductances_test_suite.py
@@ -8,7 +8,7 @@
 from causal_testing.data_collection.data_collector import ObservationalDataCollector
 from causal_testing.testing.causal_test_case import CausalTestCase
 from causal_testing.testing.causal_test_outcome import Positive, Negative, NoEffect
-from causal_testing.testing.estimators import LinearRegressionEstimator
+from causal_testing.estimation.linear_regression_estimator import LinearRegressionEstimator
 from causal_testing.testing.base_test_case import BaseTestCase
 from causal_testing.testing.causal_test_suite import CausalTestSuite
 from matplotlib.pyplot import rcParams
diff --git a/examples/poisson-line-process/example_poisson_process.py b/examples/poisson-line-process/example_poisson_process.py
index 820c3537..8a2f12a8 100644
--- a/examples/poisson-line-process/example_poisson_process.py
+++ b/examples/poisson-line-process/example_poisson_process.py
@@ -5,7 +5,7 @@
 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.estimators import LinearRegressionEstimator, Estimator
+from causal_testing.estimation.linear_regression_estimator import LinearRegressionEstimator, Estimator
 from causal_testing.testing.base_test_case import BaseTestCase
 
 import pandas as pd
diff --git a/examples/poisson/example_run_causal_tests.py b/examples/poisson/example_run_causal_tests.py
index 50ee8458..787ab6ac 100644
--- a/examples/poisson/example_run_causal_tests.py
+++ b/examples/poisson/example_run_causal_tests.py
@@ -3,11 +3,11 @@
 import scipy
 import os
 
-from causal_testing.testing.estimators import LinearRegressionEstimator
+from causal_testing.estimation.linear_regression_estimator import LinearRegressionEstimator
 from causal_testing.testing.causal_test_outcome import ExactValue, Positive, Negative, NoEffect, CausalTestOutcome
 from causal_testing.testing.causal_test_result import CausalTestResult
 from causal_testing.json_front.json_class import JsonUtility
-from causal_testing.testing.estimators import Estimator
+from causal_testing.estimation.estimator import Estimator
 from causal_testing.specification.scenario import Scenario
 from causal_testing.specification.variable import Input, Output, Meta
 
diff --git a/tests/estimation_tests/test_cubic_spline_estimator.py b/tests/estimation_tests/test_cubic_spline_estimator.py
new file mode 100644
index 00000000..f8ec8117
--- /dev/null
+++ b/tests/estimation_tests/test_cubic_spline_estimator.py
@@ -0,0 +1,46 @@
+import unittest
+import pandas as pd
+import numpy as np
+import matplotlib.pyplot as plt
+from causal_testing.specification.variable import Input
+from causal_testing.utils.validation import CausalValidator
+from causal_testing.specification.capabilities import TreatmentSequence
+
+from causal_testing.estimation.cubic_spline_estimator import CubicSplineRegressionEstimator
+from causal_testing.estimation.linear_regression_estimator import LinearRegressionEstimator
+
+from tests.estimation_tests.test_linear_regression_estimator import TestLinearRegressionEstimator
+
+
+class TestCubicSplineRegressionEstimator(TestLinearRegressionEstimator):
+    @classmethod
+    def setUpClass(cls):
+        super().setUpClass()
+
+    def test_program_11_3_cublic_spline(self):
+        """Test whether the cublic_spline regression implementation produces the same results as program 11.3 (p. 162).
+        https://www.hsph.harvard.edu/miguel-hernan/wp-content/uploads/sites/1268/2023/10/hernanrobins_WhatIf_30sep23.pdf
+        Slightly modified as Hernan et al. use linear regression for this example.
+        """
+
+        df = self.chapter_11_df.copy()
+
+        cublic_spline_estimator = CubicSplineRegressionEstimator("treatments", 1, 0, set(), "outcomes", 3, df)
+
+        model = cublic_spline_estimator._run_linear_regression()
+
+        self.assertEqual(
+            round(
+                cublic_spline_estimator.model.predict({"Intercept": 1, "treatments": 90}).iloc[0],
+                1,
+            ),
+            195.6,
+        )
+
+        ate_1 = cublic_spline_estimator.estimate_ate_calculated()
+        cublic_spline_estimator.treatment_value = 2
+        ate_2 = cublic_spline_estimator.estimate_ate_calculated()
+
+        # Doubling the treatemebnt value should roughly but not exactly double the ATE
+        self.assertNotEqual(ate_1[0] * 2, ate_2[0])
+        self.assertAlmostEqual(ate_1[0] * 2, ate_2[0])
diff --git a/tests/estimation_tests/test_ipcw_estimator.py b/tests/estimation_tests/test_ipcw_estimator.py
new file mode 100644
index 00000000..d4383d6e
--- /dev/null
+++ b/tests/estimation_tests/test_ipcw_estimator.py
@@ -0,0 +1,82 @@
+import unittest
+import pandas as pd
+import numpy as np
+import matplotlib.pyplot as plt
+from causal_testing.specification.variable import Input
+from causal_testing.utils.validation import CausalValidator
+from causal_testing.specification.capabilities import TreatmentSequence
+
+from causal_testing.estimation.ipcw_estimator import IPCWEstimator
+
+
+class TestIPCWEstimator(unittest.TestCase):
+    """
+    Test the IPCW estimator class
+    """
+
+    def test_estimate_hazard_ratio(self):
+        timesteps_per_intervention = 1
+        control_strategy = TreatmentSequence(timesteps_per_intervention, [("t", 0), ("t", 0), ("t", 0)])
+        treatment_strategy = TreatmentSequence(timesteps_per_intervention, [("t", 1), ("t", 1), ("t", 1)])
+        outcome = "outcome"
+        fit_bl_switch_formula = "xo_t_do ~ time"
+        df = pd.read_csv("tests/resources/data/temporal_data.csv")
+        df["ok"] = df["outcome"] == 1
+        estimation_model = IPCWEstimator(
+            df,
+            timesteps_per_intervention,
+            control_strategy,
+            treatment_strategy,
+            outcome,
+            "ok",
+            fit_bl_switch_formula=fit_bl_switch_formula,
+            fit_bltd_switch_formula=fit_bl_switch_formula,
+            eligibility=None,
+        )
+        estimate, intervals = estimation_model.estimate_hazard_ratio()
+        self.assertEqual(estimate["trtrand"], 1.0)
+
+    def test_invalid_treatment_strategies(self):
+        timesteps_per_intervention = 1
+        control_strategy = TreatmentSequence(timesteps_per_intervention, [("t", 0), ("t", 0), ("t", 0)])
+        treatment_strategy = TreatmentSequence(timesteps_per_intervention, [("t", 1), ("t", 1), ("t", 1)])
+        outcome = "outcome"
+        fit_bl_switch_formula = "xo_t_do ~ time"
+        df = pd.read_csv("tests/resources/data/temporal_data.csv")
+        df["t"] = (["1", "0"] * len(df))[: len(df)]
+        df["ok"] = df["outcome"] == 1
+        with self.assertRaises(ValueError):
+            estimation_model = IPCWEstimator(
+                df,
+                timesteps_per_intervention,
+                control_strategy,
+                treatment_strategy,
+                outcome,
+                "ok",
+                fit_bl_switch_formula=fit_bl_switch_formula,
+                fit_bltd_switch_formula=fit_bl_switch_formula,
+                eligibility=None,
+            )
+
+    def test_invalid_fault_t_do(self):
+        timesteps_per_intervention = 1
+        control_strategy = TreatmentSequence(timesteps_per_intervention, [("t", 0), ("t", 0), ("t", 0)])
+        treatment_strategy = TreatmentSequence(timesteps_per_intervention, [("t", 1), ("t", 1), ("t", 1)])
+        outcome = "outcome"
+        fit_bl_switch_formula = "xo_t_do ~ time"
+        df = pd.read_csv("tests/resources/data/temporal_data.csv")
+        df["ok"] = df["outcome"] == 1
+        estimation_model = IPCWEstimator(
+            df,
+            timesteps_per_intervention,
+            control_strategy,
+            treatment_strategy,
+            outcome,
+            "ok",
+            fit_bl_switch_formula=fit_bl_switch_formula,
+            fit_bltd_switch_formula=fit_bl_switch_formula,
+            eligibility=None,
+        )
+        estimation_model.df["fault_t_do"] = 0
+        with self.assertRaises(ValueError):
+            estimate, intervals = estimation_model.estimate_hazard_ratio()
diff --git a/tests/estimation_tests/test_iv_estimator.py b/tests/estimation_tests/test_iv_estimator.py
new file mode 100644
index 00000000..9711f555
--- /dev/null
+++ b/tests/estimation_tests/test_iv_estimator.py
@@ -0,0 +1,53 @@
+import unittest
+import pandas as pd
+import numpy as np
+import matplotlib.pyplot as plt
+from causal_testing.specification.variable import Input
+from causal_testing.utils.validation import CausalValidator
+from causal_testing.specification.capabilities import TreatmentSequence
+
+from causal_testing.estimation.iv_estimator import InstrumentalVariableEstimator
+
+
+class TestInstrumentalVariableEstimator(unittest.TestCase):
+    """
+    Test the instrumental variable estimator.
+    """
+
+    @classmethod
+    def setUpClass(cls) -> None:
+        Z = np.linspace(0, 10)
+        X = 2 * Z
+        Y = 2 * X
+        cls.df = pd.DataFrame({"Z": Z, "X": X, "Y": Y})
+
+    def test_estimate_coefficient(self):
+        """
+        Test we get the correct coefficient.
+        """
+        iv_estimator = InstrumentalVariableEstimator(
+            df=self.df,
+            treatment="X",
+            treatment_value=None,
+            control_value=None,
+            adjustment_set=set(),
+            outcome="Y",
+            instrument="Z",
+        )
+        self.assertEqual(iv_estimator.estimate_coefficient(self.df), 2)
+
+    def test_estimate_coefficient(self):
+        """
+        Test we get the correct coefficient.
+        """
+        iv_estimator = InstrumentalVariableEstimator(
+            df=self.df,
+            treatment="X",
+            treatment_value=None,
+            control_value=None,
+            adjustment_set=set(),
+            outcome="Y",
+            instrument="Z",
+        )
+        coefficient, [low, high] = iv_estimator.estimate_coefficient()
+        self.assertEqual(coefficient[0], 2)
diff --git a/tests/testing_tests/test_estimators.py b/tests/estimation_tests/test_linear_regression_estimator.py
similarity index 58%
rename from tests/testing_tests/test_estimators.py
rename to tests/estimation_tests/test_linear_regression_estimator.py
index 3416f4e8..e6ea4206 100644
--- a/tests/testing_tests/test_estimators.py
+++ b/tests/estimation_tests/test_linear_regression_estimator.py
@@ -2,33 +2,11 @@
 import pandas as pd
 import numpy as np
 import matplotlib.pyplot as plt
-from causal_testing.testing.estimators import (
-    LinearRegressionEstimator,
-    LogisticRegressionEstimator,
-    InstrumentalVariableEstimator,
-    CubicSplineRegressionEstimator,
-    IPCWEstimator,
-)
 from causal_testing.specification.variable import Input
 from causal_testing.utils.validation import CausalValidator
 from causal_testing.specification.capabilities import TreatmentSequence
 
-
-def plot_results_df(df):
-    """A helper method to plot results dataframe for estimators, where the df parameter must have columns for the cate,
-    ci_low, and ci_high.
-
-    :param df: A dataframe containing the columns cate, ci_low, and ci_high, where each row is an observation.
-    :return: Plot the treatment effect with confidence intervals for each observation.
-    """
-
-    df.sort_values("smokeintensity", inplace=True, ascending=True)
-    df.reset_index(inplace=True, drop=True)
-    plt.scatter(df["smokeintensity"], df["cate"], label="CATE", color="black")
-    plt.fill_between(df["smokeintensity"], df["ci_low"], df["ci_high"], alpha=0.2)
-    plt.ylabel("Weight Change (kg) caused by stopping smoking")
-    plt.xlabel("Smoke intensity (cigarettes smoked per day)")
-    plt.show()
+from causal_testing.estimation.linear_regression_estimator import LinearRegressionEstimator
 
 
 def load_nhefs_df():
@@ -72,123 +50,6 @@ def load_chapter_11_df():
     return chapter_11_df
 
 
-class TestLogisticRegressionEstimator(unittest.TestCase):
-    """Test the logistic regression estimator against the scarf example from
-    https://investigate.ai/regression/logistic-regression/.
-    """
-
-    @classmethod
-    def setUpClass(cls) -> None:
-        cls.scarf_df = pd.read_csv("tests/resources/data/scarf_data.csv")
-
-    # Yes, this probably shouldn't be in here, but it uses the scarf data so it makes more sense to put it
-    # here than duplicating the scarf data for a single test
-    def test_linear_regression_categorical_ate(self):
-        df = self.scarf_df.copy()
-        logistic_regression_estimator = LinearRegressionEstimator("color", None, None, set(), "completed", df)
-        ate, confidence = logistic_regression_estimator.estimate_coefficient()
-        self.assertTrue(all([ci_low < 0 < ci_high for ci_low, ci_high in zip(confidence[0], confidence[1])]))
-
-    def test_ate(self):
-        df = self.scarf_df.copy()
-        logistic_regression_estimator = LogisticRegressionEstimator("length_in", 65, 55, set(), "completed", df)
-        ate, _ = logistic_regression_estimator.estimate_ate()
-        self.assertEqual(round(ate, 4), -0.1987)
-
-    def test_risk_ratio(self):
-        df = self.scarf_df.copy()
-        logistic_regression_estimator = LogisticRegressionEstimator("length_in", 65, 55, set(), "completed", df)
-        rr, _ = logistic_regression_estimator.estimate_risk_ratio()
-        self.assertEqual(round(rr, 4), 0.7664)
-
-    def test_odds_ratio(self):
-        df = self.scarf_df.copy()
-        logistic_regression_estimator = LogisticRegressionEstimator("length_in", 65, 55, set(), "completed", df)
-        odds = logistic_regression_estimator.estimate_unit_odds_ratio()
-        self.assertEqual(round(odds, 4), 0.8948)
-
-    def test_ate_adjustment(self):
-        df = self.scarf_df.copy()
-        logistic_regression_estimator = LogisticRegressionEstimator(
-            "length_in", 65, 55, {"large_gauge"}, "completed", df
-        )
-        ate, _ = logistic_regression_estimator.estimate_ate(adjustment_config={"large_gauge": 0})
-        self.assertEqual(round(ate, 4), -0.3388)
-
-    def test_ate_invalid_adjustment(self):
-        df = self.scarf_df.copy()
-        logistic_regression_estimator = LogisticRegressionEstimator("length_in", 65, 55, {}, "completed", df)
-        with self.assertRaises(ValueError):
-            ate, _ = logistic_regression_estimator.estimate_ate(adjustment_config={"large_gauge": 0})
-
-    def test_ate_effect_modifiers(self):
-        df = self.scarf_df.copy()
-        logistic_regression_estimator = LogisticRegressionEstimator(
-            "length_in", 65, 55, set(), "completed", df, effect_modifiers={"large_gauge": 0}
-        )
-        ate, _ = logistic_regression_estimator.estimate_ate()
-        self.assertEqual(round(ate, 4), -0.3388)
-
-    def test_ate_effect_modifiers_formula(self):
-        df = self.scarf_df.copy()
-        logistic_regression_estimator = LogisticRegressionEstimator(
-            "length_in",
-            65,
-            55,
-            set(),
-            "completed",
-            df,
-            effect_modifiers={"large_gauge": 0},
-            formula="completed ~ length_in + large_gauge",
-        )
-        ate, _ = logistic_regression_estimator.estimate_ate()
-        self.assertEqual(round(ate, 4), -0.3388)
-
-
-class TestInstrumentalVariableEstimator(unittest.TestCase):
-    """
-    Test the instrumental variable estimator.
-    """
-
-    @classmethod
-    def setUpClass(cls) -> None:
-        Z = np.linspace(0, 10)
-        X = 2 * Z
-        Y = 2 * X
-        cls.df = pd.DataFrame({"Z": Z, "X": X, "Y": Y})
-
-    def test_estimate_coefficient(self):
-        """
-        Test we get the correct coefficient.
-        """
-        iv_estimator = InstrumentalVariableEstimator(
-            df=self.df,
-            treatment="X",
-            treatment_value=None,
-            control_value=None,
-            adjustment_set=set(),
-            outcome="Y",
-            instrument="Z",
-        )
-        self.assertEqual(iv_estimator.estimate_coefficient(self.df), 2)
-
-    def test_estimate_coefficient(self):
-        """
-        Test we get the correct coefficient.
-        """
-        iv_estimator = InstrumentalVariableEstimator(
-            df=self.df,
-            treatment="X",
-            treatment_value=None,
-            control_value=None,
-            adjustment_set=set(),
-            outcome="Y",
-            instrument="Z",
-        )
-        coefficient, [low, high] = iv_estimator.estimate_coefficient()
-        self.assertEqual(coefficient[0], 2)
-
-
 class TestLinearRegressionEstimator(unittest.TestCase):
     """Test the linear regression estimator against the programming exercises in Section 2 of Hernán and Robins [1].
 
@@ -200,6 +61,7 @@ class TestLinearRegressionEstimator(unittest.TestCase):
     def setUpClass(cls) -> None:
         cls.nhefs_df = load_nhefs_df()
         cls.chapter_11_df = load_chapter_11_df()
+        cls.scarf_df = pd.read_csv("tests/resources/data/scarf_data.csv")
 
     def test_query(self):
         df = self.nhefs_df
@@ -208,6 +70,12 @@ def test_query(self):
         )
         self.assertTrue(linear_regression_estimator.df.sex.all())
 
+    def test_linear_regression_categorical_ate(self):
+        df = self.scarf_df.copy()
+        logistic_regression_estimator = LinearRegressionEstimator("color", None, None, set(), "completed", df)
+        ate, confidence = logistic_regression_estimator.estimate_coefficient()
+        self.assertTrue(all([ci_low < 0 < ci_high for ci_low, ci_high in zip(confidence[0], confidence[1])]))
+
     def test_program_11_2(self):
         """Test whether our linear regression implementation produces the same results as program 11.2 (p. 141)."""
         df = self.chapter_11_df
@@ -418,40 +286,6 @@ def test_gp(self):
         self.assertEqual(round(ci_high[0], 2), 0.50)
 
 
-class TestCubicSplineRegressionEstimator(TestLinearRegressionEstimator):
-    @classmethod
-    def setUpClass(cls):
-        super().setUpClass()
-
-    def test_program_11_3_cublic_spline(self):
-        """Test whether the cublic_spline regression implementation produces the same results as program 11.3 (p. 162).
-        https://www.hsph.harvard.edu/miguel-hernan/wp-content/uploads/sites/1268/2023/10/hernanrobins_WhatIf_30sep23.pdf
-        Slightly modified as Hernan et al. use linear regression for this example.
-        """
-
-        df = self.chapter_11_df.copy()
-
-        cublic_spline_estimator = CubicSplineRegressionEstimator("treatments", 1, 0, set(), "outcomes", 3, df)
-
-        model = cublic_spline_estimator._run_linear_regression()
-
-        self.assertEqual(
-            round(
-                cublic_spline_estimator.model.predict({"Intercept": 1, "treatments": 90}).iloc[0],
-                1,
-            ),
-            195.6,
-        )
-
-        ate_1 = cublic_spline_estimator.estimate_ate_calculated()
-        cublic_spline_estimator.treatment_value = 2
-        ate_2 = cublic_spline_estimator.estimate_ate_calculated()
-
-        # Doubling the treatemebnt value should roughly but not exactly double the ATE
-        self.assertNotEqual(ate_1[0] * 2, ate_2[0])
-        self.assertAlmostEqual(ate_1[0] * 2, ate_2[0])
-
-
 class TestLinearRegressionInteraction(unittest.TestCase):
     """Test linear regression for estimating effects involving interaction."""
 
@@ -487,76 +321,3 @@ def test_categorical_confidence_intervals(self):
         self.assertTrue(coefficients.round(2).equals(pd.Series({"color[T.grey]": 0.92, "color[T.orange]": -4.25})))
         self.assertTrue(ci_low.round(2).equals(pd.Series({"color[T.grey]": -22.12, "color[T.orange]": -25.58})))
         self.assertTrue(ci_high.round(2).equals(pd.Series({"color[T.grey]": 23.95, "color[T.orange]": 17.08})))
-
-
-class TestIPCWEstimator(unittest.TestCase):
-    """
-    Test the IPCW estimator class
-    """
-
-    def test_estimate_hazard_ratio(self):
-        timesteps_per_intervention = 1
-        control_strategy = TreatmentSequence(timesteps_per_intervention, [("t", 0), ("t", 0), ("t", 0)])
-        treatment_strategy = TreatmentSequence(timesteps_per_intervention, [("t", 1), ("t", 1), ("t", 1)])
-        outcome = "outcome"
-        fit_bl_switch_formula = "xo_t_do ~ time"
-        df = pd.read_csv("tests/resources/data/temporal_data.csv")
-        df["ok"] = df["outcome"] == 1
-        estimation_model = IPCWEstimator(
-            df,
-            timesteps_per_intervention,
-            control_strategy,
-            treatment_strategy,
-            outcome,
-            "ok",
-            fit_bl_switch_formula=fit_bl_switch_formula,
-            fit_bltd_switch_formula=fit_bl_switch_formula,
-            eligibility=None,
-        )
-        estimate, intervals = estimation_model.estimate_hazard_ratio()
-        self.assertEqual(estimate["trtrand"], 1.0)
-
-    def test_invalid_treatment_strategies(self):
-        timesteps_per_intervention = 1
-        control_strategy = TreatmentSequence(timesteps_per_intervention, [("t", 0), ("t", 0), ("t", 0)])
-        treatment_strategy = TreatmentSequence(timesteps_per_intervention, [("t", 1), ("t", 1), ("t", 1)])
-        outcome = "outcome"
-        fit_bl_switch_formula = "xo_t_do ~ time"
-        df = pd.read_csv("tests/resources/data/temporal_data.csv")
-        df["t"] = (["1", "0"] * len(df))[: len(df)]
-        df["ok"] = df["outcome"] == 1
-        with self.assertRaises(ValueError):
-            estimation_model = IPCWEstimator(
-                df,
-                timesteps_per_intervention,
-                control_strategy,
-                treatment_strategy,
-                outcome,
-                "ok",
-                fit_bl_switch_formula=fit_bl_switch_formula,
-                fit_bltd_switch_formula=fit_bl_switch_formula,
-                eligibility=None,
-            )
-
-    def test_invalid_fault_t_do(self):
-        timesteps_per_intervention = 1
-        control_strategy = TreatmentSequence(timesteps_per_intervention, [("t", 0), ("t", 0), ("t", 0)])
-        treatment_strategy = TreatmentSequence(timesteps_per_intervention, [("t", 1), ("t", 1), ("t", 1)])
-        outcome = "outcome"
-        fit_bl_switch_formula = "xo_t_do ~ time"
-        df = pd.read_csv("tests/resources/data/temporal_data.csv")
-        df["ok"] = df["outcome"] == 1
-        estimation_model = IPCWEstimator(
-            df,
-            timesteps_per_intervention,
-            control_strategy,
-            treatment_strategy,
-            outcome,
-            "ok",
-            fit_bl_switch_formula=fit_bl_switch_formula,
-            fit_bltd_switch_formula=fit_bl_switch_formula,
-            eligibility=None,
-        )
-        estimation_model.df["fault_t_do"] = 0
-        with self.assertRaises(ValueError):
-            estimate, intervals = estimation_model.estimate_hazard_ratio()
diff --git a/tests/estimation_tests/test_logistic_regression_estimator.py b/tests/estimation_tests/test_logistic_regression_estimator.py
new file mode 100644
index 00000000..5dc4df24
--- /dev/null
+++ b/tests/estimation_tests/test_logistic_regression_estimator.py
@@ -0,0 +1,73 @@
+import unittest
+import pandas as pd
+import numpy as np
+import matplotlib.pyplot as plt
+from causal_testing.specification.variable import Input
+from causal_testing.utils.validation import CausalValidator
+from causal_testing.specification.capabilities import TreatmentSequence
+from causal_testing.estimation.logistic_regression_estimator import LogisticRegressionEstimator
+
+
+class TestLogisticRegressionEstimator(unittest.TestCase):
+    """Test the logistic regression estimator against the scarf example from
+    https://investigate.ai/regression/logistic-regression/.
+    """
+
+    @classmethod
+    def setUpClass(cls) -> None:
+        cls.scarf_df = pd.read_csv("tests/resources/data/scarf_data.csv")
+
+    def test_ate(self):
+        df = self.scarf_df.copy()
+        logistic_regression_estimator = LogisticRegressionEstimator("length_in", 65, 55, set(), "completed", df)
+        ate, _ = logistic_regression_estimator.estimate_ate()
+        self.assertEqual(round(ate, 4), -0.1987)
+
+    def test_risk_ratio(self):
+        df = self.scarf_df.copy()
+        logistic_regression_estimator = LogisticRegressionEstimator("length_in", 65, 55, set(), "completed", df)
+        rr, _ = logistic_regression_estimator.estimate_risk_ratio()
+        self.assertEqual(round(rr, 4), 0.7664)
+
+    def test_odds_ratio(self):
+        df = self.scarf_df.copy()
+        logistic_regression_estimator = LogisticRegressionEstimator("length_in", 65, 55, set(), "completed", df)
+        odds = logistic_regression_estimator.estimate_unit_odds_ratio()
+        self.assertEqual(round(odds, 4), 0.8948)
+
+    def test_ate_adjustment(self):
+        df = self.scarf_df.copy()
+        logistic_regression_estimator = LogisticRegressionEstimator(
+            "length_in", 65, 55, {"large_gauge"}, "completed", df
+        )
+        ate, _ = logistic_regression_estimator.estimate_ate(adjustment_config={"large_gauge": 0})
+        self.assertEqual(round(ate, 4), -0.3388)
+
+    def test_ate_invalid_adjustment(self):
+        df = self.scarf_df.copy()
+        logistic_regression_estimator = LogisticRegressionEstimator("length_in", 65, 55, {}, "completed", df)
+        with self.assertRaises(ValueError):
+            ate, _ = logistic_regression_estimator.estimate_ate(adjustment_config={"large_gauge": 0})
+
+    def test_ate_effect_modifiers(self):
+        df = self.scarf_df.copy()
+        logistic_regression_estimator = LogisticRegressionEstimator(
+            "length_in", 65, 55, set(), "completed", df, effect_modifiers={"large_gauge": 0}
+        )
+        ate, _ = logistic_regression_estimator.estimate_ate()
+        self.assertEqual(round(ate, 4), -0.3388)
+
+    def test_ate_effect_modifiers_formula(self):
+        df = self.scarf_df.copy()
+        logistic_regression_estimator = LogisticRegressionEstimator(
+            "length_in",
+            65,
+            55,
+            set(),
+            "completed",
+            df,
+            effect_modifiers={"large_gauge": 0},
+            formula="completed ~ length_in + large_gauge",
+        )
+        ate, _ = logistic_regression_estimator.estimate_ate()
+        self.assertEqual(round(ate, 4), -0.3388)
diff --git a/tests/json_front_tests/test_json_class.py b/tests/json_front_tests/test_json_class.py
index 8fa49194..b6585485 100644
--- a/tests/json_front_tests/test_json_class.py
+++ b/tests/json_front_tests/test_json_class.py
@@ -4,7 +4,7 @@
 import scipy
 import os
 
-from causal_testing.testing.estimators import LinearRegressionEstimator, Estimator
+from causal_testing.estimation.linear_regression_estimator import LinearRegressionEstimator, Estimator
 from causal_testing.testing.causal_test_outcome import NoEffect, Positive
 from causal_testing.json_front.json_class import JsonUtility, CausalVariables
 from causal_testing.specification.variable import Input, Output, Meta
diff --git a/tests/surrogate_tests/test_causal_surrogate_assisted.py b/tests/surrogate_tests/test_causal_surrogate_assisted.py
index 54c93af1..5d408a85 100644
--- a/tests/surrogate_tests/test_causal_surrogate_assisted.py
+++ b/tests/surrogate_tests/test_causal_surrogate_assisted.py
@@ -4,20 +4,25 @@
 from causal_testing.specification.causal_specification import CausalSpecification
 from causal_testing.specification.scenario import Scenario
 from causal_testing.specification.variable import Input, Output
-from causal_testing.surrogate.causal_surrogate_assisted import SimulationResult, CausalSurrogateAssistedTestCase, Simulator
+from causal_testing.surrogate.causal_surrogate_assisted import (
+    SimulationResult,
+    CausalSurrogateAssistedTestCase,
+    Simulator,
+)
 from causal_testing.surrogate.surrogate_search_algorithms import GeneticSearchAlgorithm
-from causal_testing.testing.estimators import CubicSplineRegressionEstimator
+from causal_testing.estimation.cubic_spline_estimator import CubicSplineRegressionEstimator
 
 import os
 import shutil, tempfile
 import pandas as pd
 import numpy as np
 
+
 class TestSimulationResult(unittest.TestCase):
 
     def setUp(self):
 
-        self.data = {'key': 'value'}
+        self.data = {"key": "value"}
 
     def test_inputs(self):
 
@@ -37,6 +42,7 @@ def test_inputs(self):
 
                     self.assertEqual(result.relationship, relationship)
 
+
 class TestCausalSurrogate(unittest.TestCase):
 
     @classmethod
@@ -79,20 +85,18 @@ def test_causal_surrogate_assisted_execution(self):
         x = Input("X", float)
         m = Input("M", int)
         y = Output("Y", float)
-        scenario = Scenario(variables={z, x, m, y}, constraints={
-            z <= 0, z >= 3,
-            x <= 0, x >= 3,
-            m <= 0, m >= 3
-        })
+        scenario = Scenario(variables={z, x, m, y}, constraints={z <= 0, z >= 3, x <= 0, x >= 3, m <= 0, m >= 3})
         specification = CausalSpecification(scenario, causal_dag)
 
-        search_algorithm = GeneticSearchAlgorithm(config= {
+        search_algorithm = GeneticSearchAlgorithm(
+            config={
                 "parent_selection_type": "tournament",
                 "K_tournament": 4,
                 "mutation_type": "random",
                 "mutation_percent_genes": 50,
                 "mutation_by_replacement": True,
-            })
+            }
+        )
         simulator = TestSimulator()
 
         c_s_a_test_case = CausalSurrogateAssistedTestCase(specification, search_algorithm, simulator)
@@ -111,20 +115,18 @@ def test_causal_surrogate_assisted_execution_failure(self):
         x = Input("X", float)
         m = Input("M", int)
         y = Output("Y", float)
-        scenario = Scenario(variables={z, x, m, y}, constraints={
-            z <= 0, z >= 3,
-            x <= 0, x >= 3,
-            m <= 0, m >= 3
-        })
+        scenario = Scenario(variables={z, x, m, y}, constraints={z <= 0, z >= 3, x <= 0, x >= 3, m <= 0, m >= 3})
         specification = CausalSpecification(scenario, causal_dag)
 
-        search_algorithm = GeneticSearchAlgorithm(config= {
+        search_algorithm = GeneticSearchAlgorithm(
+            config={
                 "parent_selection_type": "tournament",
                 "K_tournament": 4,
                 "mutation_type": "random",
                 "mutation_percent_genes": 50,
                 "mutation_by_replacement": True,
-            })
+            }
+        )
         simulator = TestSimulatorFailing()
 
         c_s_a_test_case = CausalSurrogateAssistedTestCase(specification, search_algorithm, simulator)
@@ -143,26 +145,25 @@ def test_causal_surrogate_assisted_execution_custom_aggregator(self):
         x = Input("X", float)
         m = Input("M", int)
         y = Output("Y", float)
-        scenario = Scenario(variables={z, x, m, y}, constraints={
-            z <= 0, z >= 3,
-            x <= 0, x >= 3,
-            m <= 0, m >= 3
-        })
+        scenario = Scenario(variables={z, x, m, y}, constraints={z <= 0, z >= 3, x <= 0, x >= 3, m <= 0, m >= 3})
         specification = CausalSpecification(scenario, causal_dag)
 
-        search_algorithm = GeneticSearchAlgorithm(config= {
+        search_algorithm = GeneticSearchAlgorithm(
+            config={
                 "parent_selection_type": "tournament",
                 "K_tournament": 4,
                 "mutation_type": "random",
                 "mutation_percent_genes": 50,
                 "mutation_by_replacement": True,
-            })
+            }
+        )
         simulator = TestSimulator()
 
         c_s_a_test_case = CausalSurrogateAssistedTestCase(specification, search_algorithm, simulator)
 
-        result, iterations, result_data = c_s_a_test_case.execute(ObservationalDataCollector(scenario, df), 
-                                                                    custom_data_aggregator=data_double_aggregator)
+        result, iterations, result_data = c_s_a_test_case.execute(
+            ObservationalDataCollector(scenario, df), custom_data_aggregator=data_double_aggregator
+        )
 
         self.assertIsInstance(result, SimulationResult)
         self.assertEqual(iterations, 1)
@@ -176,62 +177,69 @@ def test_causal_surrogate_assisted_execution_incorrect_search_config(self):
         x = Input("X", float)
         m = Input("M", int)
         y = Output("Y", float)
-        scenario = Scenario(variables={z, x, m, y}, constraints={
-            z <= 0, z >= 3,
-            x <= 0, x >= 3,
-            m <= 0, m >= 3
-        })
+        scenario = Scenario(variables={z, x, m, y}, constraints={z <= 0, z >= 3, x <= 0, x >= 3, m <= 0, m >= 3})
         specification = CausalSpecification(scenario, causal_dag)
 
-        search_algorithm = GeneticSearchAlgorithm(config= {
+        search_algorithm = GeneticSearchAlgorithm(
+            config={
                 "parent_selection_type": "tournament",
                 "K_tournament": 4,
                 "mutation_type": "random",
                 "mutation_percent_genes": 50,
                 "mutation_by_replacement": True,
-                "gene_space": "Something"
-            })
+                "gene_space": "Something",
+            }
+        )
         simulator = TestSimulator()
 
         c_s_a_test_case = CausalSurrogateAssistedTestCase(specification, search_algorithm, simulator)
 
-        self.assertRaises(ValueError, c_s_a_test_case.execute, 
-                          data_collector=ObservationalDataCollector(scenario, df),
-                          custom_data_aggregator=data_double_aggregator)
+        self.assertRaises(
+            ValueError,
+            c_s_a_test_case.execute,
+            data_collector=ObservationalDataCollector(scenario, df),
+            custom_data_aggregator=data_double_aggregator,
+        )
 
     def tearDown(self) -> None:
         shutil.rmtree(self.temp_dir_path)
 
+
 def load_class_df():
     """Get the testing data and put into a dataframe."""
 
-    class_df = pd.DataFrame({"Z": np.arange(16), "X": np.arange(16), "M": np.arange(16, 32), "Y": np.arange(32,16,-1)})
+    class_df = pd.DataFrame(
+        {"Z": np.arange(16), "X": np.arange(16), "M": np.arange(16, 32), "Y": np.arange(32, 16, -1)}
+    )
     return class_df
 
+
 class TestSimulator(Simulator):
 
     def run_with_config(self, configuration: dict) -> SimulationResult:
         return SimulationResult({"Z": 1, "X": 1, "M": 1, "Y": 1}, True, None)
-    
+
     def startup(self):
         pass
 
     def shutdown(self):
         pass
 
+
 class TestSimulatorFailing(Simulator):
 
     def run_with_config(self, configuration: dict) -> SimulationResult:
         return SimulationResult({"Z": 1, "X": 1, "M": 1, "Y": 1}, False, None)
-    
+
     def startup(self):
         pass
 
     def shutdown(self):
         pass
 
+
 def data_double_aggregator(data, new_data):
-    """Previously used data.append(new_data), however, pandas version >2 requires pd.concat() since append is now a private method. 
-       Converting new_data to a pd.DataFrame is required to use pd.concat(). """
+    """Previously used data.append(new_data), however, pandas version >2 requires pd.concat() since append is now a private method.
+    Converting new_data to a pd.DataFrame is required to use pd.concat()."""
     new_data = pd.DataFrame([new_data])
     return pd.concat([data, new_data, new_data], ignore_index=True)
diff --git a/tests/testing_tests/test_causal_test_adequacy.py b/tests/testing_tests/test_causal_test_adequacy.py
index e29a8207..d5aeb5b1 100644
--- a/tests/testing_tests/test_causal_test_adequacy.py
+++ b/tests/testing_tests/test_causal_test_adequacy.py
@@ -5,7 +5,8 @@
 import os
 import pandas as pd
 
-from causal_testing.testing.estimators import LinearRegressionEstimator, IPCWEstimator
+from causal_testing.estimation.linear_regression_estimator import LinearRegressionEstimator
+from causal_testing.estimation.ipcw_estimator import IPCWEstimator
 from causal_testing.testing.base_test_case import BaseTestCase
 from causal_testing.testing.causal_test_case import CausalTestCase
 from causal_testing.testing.causal_test_suite import CausalTestSuite
diff --git a/tests/testing_tests/test_causal_test_case.py b/tests/testing_tests/test_causal_test_case.py
index 4d081a62..01b55cc1 100644
--- a/tests/testing_tests/test_causal_test_case.py
+++ b/tests/testing_tests/test_causal_test_case.py
@@ -11,7 +11,7 @@
 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
-from causal_testing.testing.estimators import LinearRegressionEstimator
+from causal_testing.estimation.linear_regression_estimator import LinearRegressionEstimator
 from causal_testing.testing.base_test_case import BaseTestCase
 
 
diff --git a/tests/testing_tests/test_causal_test_outcome.py b/tests/testing_tests/test_causal_test_outcome.py
index 235cc724..0bdbe4fa 100644
--- a/tests/testing_tests/test_causal_test_outcome.py
+++ b/tests/testing_tests/test_causal_test_outcome.py
@@ -2,7 +2,7 @@
 import pandas as pd
 from causal_testing.testing.causal_test_outcome import ExactValue, SomeEffect, Positive, Negative, NoEffect
 from causal_testing.testing.causal_test_result import CausalTestResult, TestValue
-from causal_testing.testing.estimators import LinearRegressionEstimator
+from causal_testing.estimation.linear_regression_estimator import LinearRegressionEstimator
 from causal_testing.utils.validation import CausalValidator
 
 
diff --git a/tests/testing_tests/test_causal_test_suite.py b/tests/testing_tests/test_causal_test_suite.py
index d6af77cb..151f7af2 100644
--- a/tests/testing_tests/test_causal_test_suite.py
+++ b/tests/testing_tests/test_causal_test_suite.py
@@ -9,7 +9,8 @@
 from causal_testing.testing.base_test_case import BaseTestCase
 from causal_testing.specification.variable import Input, Output
 from causal_testing.testing.causal_test_outcome import ExactValue
-from causal_testing.testing.estimators import LinearRegressionEstimator, LogisticRegressionEstimator
+from causal_testing.estimation.linear_regression_estimator import LinearRegressionEstimator
+from causal_testing.estimation.logistic_regression_estimator import LogisticRegressionEstimator
 from causal_testing.specification.causal_specification import CausalSpecification, Scenario
 from causal_testing.data_collection.data_collector import ObservationalDataCollector
 from causal_testing.specification.causal_dag import CausalDAG
@@ -102,7 +103,9 @@ def test_execute_test_suite_single_base_test_case(self):
 
         causal_test_results = self.test_suite.execute_test_suite(self.data_collector, self.causal_specification)
         causal_test_case_result = causal_test_results[self.base_test_case]
-        self.assertAlmostEqual(causal_test_case_result["LinearRegressionEstimator"][0].test_value.value[0], 4, delta=1e-10)
+        self.assertAlmostEqual(
+            causal_test_case_result["LinearRegressionEstimator"][0].test_value.value[0], 4, delta=1e-10
+        )
 
     # Without CausalForestEstimator we now only have 2 estimators. Unfortunately LogicisticRegressionEstimator does not
     # currently work with TestSuite. So for now removed test

From 3e252567b277aee94482e9d790017b2fe1a6c157 Mon Sep 17 00:00:00 2001
From: Michael <jmafoster1@gmail.com>
Date: Wed, 7 Aug 2024 15:02:38 +0100
Subject: [PATCH 04/24] pylint

---
 .../estimation/cubic_spline_estimator.py      |  15 +-
 causal_testing/estimation/estimator.py        |  12 -
 causal_testing/estimation/gp.py               | 237 ++++++++++--------
 causal_testing/estimation/ipcw_estimator.py   |   8 -
 causal_testing/estimation/iv_estimator.py     |  15 +-
 .../estimation/linear_regression_estimator.py |  37 ++-
 .../logistic_regression_estimator.py          |   6 -
 causal_testing/specification/capabilities.py  |   1 +
 .../covasim_/doubling_beta/example_beta.py    |   6 +-
 .../example_poisson_process.py                |  16 +-
 .../test_linear_regression_estimator.py       |  18 +-
 11 files changed, 186 insertions(+), 185 deletions(-)

diff --git a/causal_testing/estimation/cubic_spline_estimator.py b/causal_testing/estimation/cubic_spline_estimator.py
index 06ee2b44..d1ca72b5 100644
--- a/causal_testing/estimation/cubic_spline_estimator.py
+++ b/causal_testing/estimation/cubic_spline_estimator.py
@@ -1,23 +1,12 @@
-"""This module contains the CubicSplineRegressionEstimator class, for estimating continuous outcomes with changes in behaviour"""
+"""This module contains the CubicSplineRegressionEstimator class, for estimating
+continuous outcomes with changes in behaviour"""
 
 import logging
-from abc import ABC, abstractmethod
 from typing import Any
-from math import ceil
 
-import numpy as np
 import pandas as pd
-import statsmodels.api as sm
-import statsmodels.formula.api as smf
-from patsy import dmatrix  # pylint: disable = no-name-in-module
-from patsy import ModelDesc
-from statsmodels.regression.linear_model import RegressionResultsWrapper
-from statsmodels.tools.sm_exceptions import PerfectSeparationError
-from lifelines import CoxPHFitter
 
 from causal_testing.specification.variable import Variable
-from causal_testing.specification.capabilities import TreatmentSequence, Capability
-from causal_testing.estimation.estimator import Estimator
 from causal_testing.estimation.linear_regression_estimator import LinearRegressionEstimator
 
 logger = logging.getLogger(__name__)
diff --git a/causal_testing/estimation/estimator.py b/causal_testing/estimation/estimator.py
index 81876518..ab73723f 100644
--- a/causal_testing/estimation/estimator.py
+++ b/causal_testing/estimation/estimator.py
@@ -3,20 +3,8 @@
 import logging
 from abc import ABC, abstractmethod
 from typing import Any
-from math import ceil
 
-import numpy as np
 import pandas as pd
-import statsmodels.api as sm
-import statsmodels.formula.api as smf
-from patsy import dmatrix  # pylint: disable = no-name-in-module
-from patsy import ModelDesc
-from statsmodels.regression.linear_model import RegressionResultsWrapper
-from statsmodels.tools.sm_exceptions import PerfectSeparationError
-from lifelines import CoxPHFitter
-
-from causal_testing.specification.variable import Variable
-from causal_testing.specification.capabilities import TreatmentSequence, Capability
 
 logger = logging.getLogger(__name__)
 
diff --git a/causal_testing/estimation/gp.py b/causal_testing/estimation/gp.py
index 00737437..f1030c8d 100644
--- a/causal_testing/estimation/gp.py
+++ b/causal_testing/estimation/gp.py
@@ -1,62 +1,54 @@
+"""
+This module contains a genetic programming implementation to infer the functional
+form between the adjustment set and the outcome.
+"""
+
+import copy
+from inspect import isclass
+from operator import add, mul
 import random
-import warnings
-import patsy
 
-import matplotlib.pyplot as plt
+import patsy
 import numpy as np
 import pandas as pd
 import statsmodels.formula.api as smf
 import statsmodels
 import sympy
-import copy
-
-from functools import partial
-from deap import algorithms, base, creator, tools, gp
-
-from numpy import negative, exp, power, log, sin, cos, tan, sinh, cosh, tanh
-from inspect import isclass
-
-from operator import add, mul
-
-
-def root(x):
-    return power(x, 0.5)
-
-
-def square(x):
-    return power(x, 2)
 
+from deap import base, creator, tools, gp
 
-def cube(x):
-    return power(x, 3)
+from numpy import power, log
 
 
-def fourth_power(x):
-    return power(x, 4)
-
-
-def reciprocal(x):
+def reciprocal(x: float) -> float:
+    """
+    Return the reciprocal of the input.
+    :param x: Float to reciprocate.
+    :return: 1/x
+    """
     return power(x, -1)
 
 
-def mutInsert(individual, pset):
+def mut_insert(expression: gp.PrimitiveTree, pset: gp.PrimitiveSet):
     """
     Copied from gp.mutInsert, except that we import isclass from inspect, so we
     won't have the "isclass not defined" bug.
 
-    Inserts a new branch at a random position in *individual*. The subtree
+    Inserts a new branch at a random position in *expression*. The subtree
     at the chosen position is used as child node of the created subtree, in
     that way, it is really an insertion rather than a replacement. Note that
     the original subtree will become one of the children of the new primitive
     inserted, but not perforce the first (its position is randomly selected if
     the new primitive has more than one child).
 
-    :param individual: The normal or typed tree to be mutated.
-    :returns: A tuple of one tree.
+    :param expression: The normal or typed tree to be mutated.
+    :param pset: The pset object defining the variables and constants.
+
+    :return: A tuple of one tree.
     """
-    index = random.randrange(len(individual))
-    node = individual[index]
-    slice_ = individual.searchSubtree(index)
+    index = random.randrange(len(expression))
+    node = expression[index]
+    expr_slice = expression.searchSubtree(index)
     choice = random.choice
 
     # As we want to keep the current node as children of the new one,
@@ -64,7 +56,7 @@ def mutInsert(individual, pset):
     primitives = [p for p in pset.primitives[node.ret] if node.ret in p.args]
 
     if len(primitives) == 0:
-        return (individual,)
+        return (expression,)
 
     new_node = choice(primitives)
     new_subtree = [None] * len(new_node.args)
@@ -77,13 +69,16 @@ def mutInsert(individual, pset):
                 term = term()
             new_subtree[i] = term
 
-    new_subtree[position : position + 1] = individual[slice_]
+    new_subtree[position : position + 1] = expression[expr_slice]
     new_subtree.insert(0, new_node)
-    individual[slice_] = new_subtree
-    return (individual,)
+    expression[expr_slice] = new_subtree
+    return (expression,)
 
 
 class GP:
+    """
+    Object to perform genetic programming.
+    """
 
     def __init__(
         self,
@@ -94,7 +89,9 @@ def __init__(
         sympy_conversions: dict = None,
         seed=0,
     ):
+        # pylint: disable=too-many-arguments
         random.seed(seed)
+        np.random.seed(seed)
         self.df = df
         self.features = features
         self.outcome = outcome
@@ -110,9 +107,9 @@ def __init__(
         if sympy_conversions is None:
             sympy_conversions = {}
         self.sympy_conversions = {
-            "mul": lambda *args_: "Mul({},{})".format(*args_),
-            "add": lambda *args_: "Add({},{})".format(*args_),
-            "reciprocal": lambda *args_: "Pow({},-1)".format(*args_),
+            "mul": lambda x1, x2: f"Mul({x1},{x2})",
+            "add": lambda x1, x2: f"Add({x1},{x2})",
+            "reciprocal": lambda x1: f"Pow({x1},-1)",
         } | sympy_conversions
 
         creator.create("FitnessMin", base.Fitness, weights=(-1.0,))
@@ -123,16 +120,21 @@ def __init__(
         self.toolbox.register("individual", tools.initIterate, creator.Individual, self.toolbox.expr)
         self.toolbox.register("population", tools.initRepeat, list, self.toolbox.individual)
         self.toolbox.register("compile", gp.compile, pset=self.pset)
-        self.toolbox.register("evaluate", self.evalSymbReg)
+        self.toolbox.register("evaluate", self.fitness)
         self.toolbox.register("repair", self.repair)
         self.toolbox.register("select", tools.selBest)
         self.toolbox.register("mate", gp.cxOnePoint)
-        self.toolbox.register("expr_mut", gp.genFull, min_=0, max_=2)
-        self.toolbox.register("mutate", self.mutate, expr=self.toolbox.expr_mut)
+        self.toolbox.register("mutate", self.mutate)
         self.toolbox.decorate("mate", gp.staticLimit(key=lambda x: x.height + 1, max_value=17))
         self.toolbox.decorate("mutate", gp.staticLimit(key=lambda x: x.height + 1, max_value=17))
 
-    def split(self, individual):
+    def split(self, individual: gp.PrimitiveTree) -> list:
+        """
+        Split an expression into its components, e.g. 2x + 4y - xy -> [2x, 4y, xy].
+
+        :param individual: The expression to be split.
+        :return: A list of the equations components that are linearly combined into the full equation.
+        """
         if len(individual) > 1:
             terms = []
             # Recurse over children if add/sub
@@ -154,48 +156,69 @@ def split(self, individual):
             return terms
         return [individual]
 
-    def _convert_inverse_prim(self, prim, args):
+    def _convert_prim(self, prim: gp.Primitive, args: list) -> str:
         """
-        Convert inverse prims according to:
-        [Dd]iv(a,b) -> Mul[a, 1/b]
-        [Ss]ub(a,b) -> Add[a, -b]
-        We achieve this by overwriting the corresponding format method of the sub and div prim.
+        Convert primitives to sympy format.
+
+        :param prim: A GP primitive, e.g. add
+        :param args: The list of arguments
+
+        :return: A sympy compatible string representing the function, e.g. add(x, y) -> Add(x, y).
         """
         prim = copy.copy(prim)
         prim_formatter = self.sympy_conversions.get(prim.name, prim.format)
-
         return prim_formatter(*args)
 
-    def _stringify_for_sympy(self, f):
-        """Return the expression in a human readable string."""
+    def _stringify_for_sympy(self, expression: gp.PrimitiveTree) -> str:
+        """
+        Return the expression in a sympy compatible string.
+
+        :param expression: The expression to be simplified.
+
+        :return: A sympy compatible string representing the equation.
+        """
         string = ""
         stack = []
-        for node in f:
+        for node in expression:
             stack.append((node, []))
             while len(stack[-1][1]) == stack[-1][0].arity:
                 prim, args = stack.pop()
-                string = self._convert_inverse_prim(prim, args)
+                string = self._convert_prim(prim, args)
                 if len(stack) == 0:
                     break  # If stack is empty, all nodes should have been seen
                 stack[-1][1].append(string)
         return string
 
-    def simplify(self, individual):
-        return sympy.simplify(self._stringify_for_sympy(individual))
+    def simplify(self, expression: gp.PrimitiveTree) -> sympy.core.Expr:
+        """
+        Simplify an expression by appling mathematical equivalences.
+
+        :param expression: The expression to simplify.
+
+        :return: The simplified expression as a sympy Expr object.
+        """
+        return sympy.simplify(self._stringify_for_sympy(expression))
+
+    def repair(self, expression: gp.PrimitiveTree) -> gp.PrimitiveTree:
+        """
+        Use linear regression to infer the coefficients of the linear components of the expression.
+        Named "repair" since a "repair operator" is quite common in GP.
+
+        :param expression: The expression to process.
 
-    def repair(self, individual):
-        eq = f"{self.outcome} ~ {' + '.join(str(x) for x in self.split(individual))}"
+        :return: The expression with constant coefficients, or the original expression if that fails.
+        """
+        eq = f"{self.outcome} ~ {' + '.join(str(x) for x in self.split(expression))}"
         try:
             # Create model, fit (run) it, give estimates from it]
             model = smf.ols(eq, self.df)
             res = model.fit()
-            y_estimates = res.predict(self.df)
 
             eqn = f"{res.params['Intercept']}"
             for term, coefficient in res.params.items():
                 if term != "Intercept":
                     eqn = f"add({eqn}, mul({coefficient}, {term}))"
-            repaired = type(individual)(gp.PrimitiveTree.from_string(eqn, self.pset))
+            repaired = type(expression)(gp.PrimitiveTree.from_string(eqn, self.pset))
             return repaired
         except (
             OverflowError,
@@ -203,14 +226,22 @@ def repair(self, individual):
             ZeroDivisionError,
             statsmodels.tools.sm_exceptions.MissingDataError,
             patsy.PatsyError,
-        ) as e:
-            return individual
+        ):
+            return expression
+
+    def fitness(self, expression: gp.PrimitiveTree) -> float:
+        """
+        Evaluate the fitness of an candidate expression according to the error between the estimated and observed
+        values. Low values are better.
 
-    def evalSymbReg(self, individual):
+        :param expression: The candidate expression to evaluate.
+
+        :return: The fitness of the individual.
+        """
         old_settings = np.seterr(all="raise")
         try:
             # Create model, fit (run) it, give estimates from it]
-            func = gp.compile(individual, self.pset)
+            func = gp.compile(expression, self.pset)
             y_estimates = pd.Series([func(**x) for _, x in self.df[self.features].iterrows()])
 
             # Calc errors using an improved normalised mean squared
@@ -229,14 +260,22 @@ def evalSymbReg(self, individual):
             patsy.PatsyError,
             RuntimeWarning,
             FloatingPointError,
-        ) as e:
+        ):
             return (float("inf"),)
         finally:
             np.seterr(**old_settings)  # Restore original settings
 
-    def make_offspring(self, population, lambda_):
+    def make_offspring(self, population: list, num_offspring: int) -> list:
+        """
+        Create the next generation of individuals.
+
+        :param population: The current population.
+        :param num_offspring: The number of new individuals to generate.
+
+        :return: A list of num_offspring new individuals generated through crossover and mutation.
+        """
         offspring = []
-        for i in range(lambda_):
+        for _ in range(num_offspring):
             parent1, parent2 = tools.selTournament(population, 2, 2)
             child, _ = self.toolbox.mate(self.toolbox.clone(parent1), self.toolbox.clone(parent2))
             del child.fitness.values
@@ -244,69 +283,63 @@ def make_offspring(self, population, lambda_):
             offspring.append(child)
         return offspring
 
-    def eaMuPlusLambda(self, ngen, mu=20, lambda_=10, stats=None, verbose=False, seeds=None):
-        population = [self.toolbox.repair(ind) for ind in self.toolbox.population(n=mu)]
+    def run_gp(self, ngen: int, pop_size: int = 20, num_offspring: int = 10, seeds: list = None) -> gp.PrimitiveTree:
+        """
+        Execute Genetic Programming to find the best expression using a mu+lambda algorithm.
+
+        :param ngen: The maximum number of generations.
+        :param pop_size: The population size.
+        :param num_offspring: The number of new individuals per generation.
+        :param seeds: Seed individuals for the initial population.
+
+        :return: The best candididate expression.
+        """
+        population = [self.toolbox.repair(ind) for ind in self.toolbox.population(n=pop_size)]
         if seeds is not None:
             for seed in seeds:
                 ind = creator.Individual(gp.PrimitiveTree.from_string(seed, self.pset))
                 ind.fitness.values = self.toolbox.evaluate(ind)
                 population.append(ind)
 
-        logbook = tools.Logbook()
-        logbook.header = ["gen", "nevals"] + (stats.fields if stats else [])
-
         # Evaluate the individuals with an invalid fitness
         for ind in population:
             ind.fitness.values = self.toolbox.evaluate(ind)
         population.sort(key=lambda x: (x.fitness.values, x.height))
 
-        record = stats.compile(population) if stats is not None else {}
-        logbook.record(gen=0, nevals=len(population), **record)
-        if verbose:
-            print(logbook.stream)
-
         # Begin the generational process
-        for gen in range(1, ngen + 1):
+        for _ in range(1, ngen + 1):
             # Vary the population
-            offspring = self.make_offspring(population, lambda_)
+            offspring = self.make_offspring(population, num_offspring)
             offspring = [self.toolbox.repair(ind) for ind in offspring]
 
             # Evaluate the individuals with an invalid fitness
             for ind in offspring:
                 ind.fitness.values = self.toolbox.evaluate(ind)
 
-            # Select the next generation population
-            population[:] = self.toolbox.select(population + offspring, mu)
+            # Select the best pop_size individuals to continue to the next generation
+            population[:] = self.toolbox.select(population + offspring, pop_size)
 
             # Update the statistics with the new population
-            record = stats.compile(population) if stats is not None else {}
-            logbook.record(gen=gen, nevals=len(offspring), **record)
-            if verbose:
-                print(logbook.stream)
             population.sort(key=lambda x: (x.fitness.values, x.height))
 
         return population[0]
 
-    def mutate(self, individual, expr):
+    def mutate(self, expression: gp.PrimitiveTree) -> gp.PrimitiveTree:
+        """
+        mutate individuals to replicate the small changes in DNA that occur in natural reproduction.
+        A node will randomly be inserted, removed, or replaced.
+
+        :param expression: The expression to mutate.
+
+        :return: The mutated expression.
+        """
         choice = random.randint(1, 3)
         if choice == 1:
-            mutated = gp.mutNodeReplacement(self.toolbox.clone(individual), self.pset)
+            mutated = gp.mutNodeReplacement(self.toolbox.clone(expression), self.pset)
         elif choice == 2:
-            mutated = mutInsert(self.toolbox.clone(individual), self.pset)
+            mutated = mut_insert(self.toolbox.clone(expression), self.pset)
         elif choice == 3:
-            mutated = gp.mutShrink(self.toolbox.clone(individual))
+            mutated = gp.mutShrink(self.toolbox.clone(expression))
         else:
             raise ValueError("Invalid mutation choice")
         return mutated
-
-
-if __name__ == "__main__":
-    df = pd.DataFrame()
-    df["X"] = np.arange(10)
-    df["Y"] = 1 / (df.X + 1)
-
-    gp1 = GP(df.astype(float), ["X"], "Y", seed=1)
-    best = gp1.eaMuPlusLambda(ngen=100)
-    print(best, best.fitness.values[0])
-    simplified = gp1.simplify(best)
-    print(simplified)
diff --git a/causal_testing/estimation/ipcw_estimator.py b/causal_testing/estimation/ipcw_estimator.py
index 2330c2f9..a0c5a819 100644
--- a/causal_testing/estimation/ipcw_estimator.py
+++ b/causal_testing/estimation/ipcw_estimator.py
@@ -1,21 +1,13 @@
 """This module contains the IPCWEstimator class, for estimating the time to a particular event"""
 
 import logging
-from abc import ABC, abstractmethod
-from typing import Any
 from math import ceil
 
 import numpy as np
 import pandas as pd
-import statsmodels.api as sm
 import statsmodels.formula.api as smf
-from patsy import dmatrix  # pylint: disable = no-name-in-module
-from patsy import ModelDesc
-from statsmodels.regression.linear_model import RegressionResultsWrapper
-from statsmodels.tools.sm_exceptions import PerfectSeparationError
 from lifelines import CoxPHFitter
 
-from causal_testing.specification.variable import Variable
 from causal_testing.specification.capabilities import TreatmentSequence, Capability
 from causal_testing.estimation.estimator import Estimator
 
diff --git a/causal_testing/estimation/iv_estimator.py b/causal_testing/estimation/iv_estimator.py
index 39bcf604..5a30dfd6 100644
--- a/causal_testing/estimation/iv_estimator.py
+++ b/causal_testing/estimation/iv_estimator.py
@@ -1,22 +1,11 @@
-"""This module contains the InstrumentalVariableEstimator class, for estimating continuous outcomes with unobservable confounding."""
+"""This module contains the InstrumentalVariableEstimator class, for estimating
+continuous outcomes with unobservable confounding."""
 
 import logging
-from abc import ABC, abstractmethod
-from typing import Any
 from math import ceil
-
-import numpy as np
 import pandas as pd
 import statsmodels.api as sm
-import statsmodels.formula.api as smf
-from patsy import dmatrix  # pylint: disable = no-name-in-module
-from patsy import ModelDesc
-from statsmodels.regression.linear_model import RegressionResultsWrapper
-from statsmodels.tools.sm_exceptions import PerfectSeparationError
-from lifelines import CoxPHFitter
 
-from causal_testing.specification.variable import Variable
-from causal_testing.specification.capabilities import TreatmentSequence, Capability
 from causal_testing.estimation.estimator import Estimator
 
 logger = logging.getLogger(__name__)
diff --git a/causal_testing/estimation/linear_regression_estimator.py b/causal_testing/estimation/linear_regression_estimator.py
index e212dca7..34ea31ce 100644
--- a/causal_testing/estimation/linear_regression_estimator.py
+++ b/causal_testing/estimation/linear_regression_estimator.py
@@ -1,22 +1,15 @@
 """This module contains the LinearRegressionEstimator for estimating continuous outcomes."""
 
 import logging
-from abc import ABC, abstractmethod
 from typing import Any
-from math import ceil
 
-import numpy as np
 import pandas as pd
-import statsmodels.api as sm
 import statsmodels.formula.api as smf
 from patsy import dmatrix  # pylint: disable = no-name-in-module
 from patsy import ModelDesc
 from statsmodels.regression.linear_model import RegressionResultsWrapper
-from statsmodels.tools.sm_exceptions import PerfectSeparationError
-from lifelines import CoxPHFitter
 
 from causal_testing.specification.variable import Variable
-from causal_testing.specification.capabilities import TreatmentSequence, Capability
 from causal_testing.estimation.gp import GP
 from causal_testing.estimation.estimator import Estimator
 
@@ -67,7 +60,31 @@ def __init__(
         for term in self.effect_modifiers:
             self.adjustment_set.add(term)
 
-    def gp_formula(self, ngen=100, mu=20, lambda_=10, extra_operators=None, sympy_conversions=None, seeds=None, seed=0):
+    def gp_formula(
+        self,
+        ngen: int = 100,
+        mu: int = 20,
+        lambda_: int = 10,
+        extra_operators: list = None,
+        sympy_conversions: dict = None,
+        seeds: list = None,
+        seed: int = 0,
+    ):
+        # pylint: disable=too-many-arguments,invalid-name
+        """
+        Use Genetic Programming (GP) to infer the regression equation from the data.
+
+        :param ngen: The maximum number of GP generations to run for.
+        :param mu: The GP population size.
+        :param lambda_: The number of offspring per generation.
+        :param extra_operators: Additional operators for the GP (defaults are +, *, and 1/x). Operations should be of
+                                the form (fun, numArgs), e.g. (add, 2).
+        :param sympy_conversions: Dictionary of conversions of extra_operators for sympy,
+                                  e.g. `"mul": lambda *args_: "Mul({},{})".format(*args_)`.
+        :param seeds: Seed individuals for the population (e.g. if you think that the relationship between X and Y is
+                      probably logarithmic, you can put that in).
+        :param seed: Random seed for the GP.
+        """
         gp = GP(
             df=self.df,
             features=sorted(list(self.adjustment_set.union([self.treatment]))),
@@ -76,7 +93,7 @@ def gp_formula(self, ngen=100, mu=20, lambda_=10, extra_operators=None, sympy_co
             sympy_conversions=sympy_conversions,
             seed=seed,
         )
-        formula = gp.eaMuPlusLambda(ngen=ngen, mu=mu, lambda_=lambda_, seeds=seeds)
+        formula = gp.run_gp(ngen=ngen, pop_size=mu, num_offspring=lambda_, seeds=seeds)
         formula = gp.simplify(formula)
         self.formula = f"{self.outcome} ~ I({formula}) - 1"
 
@@ -159,6 +176,8 @@ def estimate_control_treatment(self, adjustment_config: dict = None) -> tuple[pd
         x = pd.DataFrame(columns=self.df.columns)
         x[self.treatment] = [self.treatment_value, self.control_value]
         x["Intercept"] = 1  # self.intercept
+
+        print(x[self.treatment])
         for k, v in adjustment_config.items():
             x[k] = v
         for k, v in self.effect_modifiers.items():
diff --git a/causal_testing/estimation/logistic_regression_estimator.py b/causal_testing/estimation/logistic_regression_estimator.py
index 399d02e2..1e180c49 100644
--- a/causal_testing/estimation/logistic_regression_estimator.py
+++ b/causal_testing/estimation/logistic_regression_estimator.py
@@ -1,22 +1,16 @@
 """This module contains the LogisticRegressionEstimator class for estimating categorical outcomes."""
 
 import logging
-from abc import ABC, abstractmethod
 from typing import Any
 from math import ceil
 
 import numpy as np
 import pandas as pd
-import statsmodels.api as sm
 import statsmodels.formula.api as smf
 from patsy import dmatrix  # pylint: disable = no-name-in-module
-from patsy import ModelDesc
 from statsmodels.regression.linear_model import RegressionResultsWrapper
 from statsmodels.tools.sm_exceptions import PerfectSeparationError
-from lifelines import CoxPHFitter
 
-from causal_testing.specification.variable import Variable
-from causal_testing.specification.capabilities import TreatmentSequence, Capability
 from causal_testing.estimation.estimator import Estimator
 
 logger = logging.getLogger(__name__)
diff --git a/causal_testing/specification/capabilities.py b/causal_testing/specification/capabilities.py
index ed692e6d..1cb6d30b 100644
--- a/causal_testing/specification/capabilities.py
+++ b/causal_testing/specification/capabilities.py
@@ -2,6 +2,7 @@
 This module contains the Capability and TreatmentSequence classes to implement
 treatment sequences that operate over time.
 """
+
 from typing import Any
 from causal_testing.specification.variable import Variable
 
diff --git a/examples/covasim_/doubling_beta/example_beta.py b/examples/covasim_/doubling_beta/example_beta.py
index 15152b7c..fb3ebb59 100644
--- a/examples/covasim_/doubling_beta/example_beta.py
+++ b/examples/covasim_/doubling_beta/example_beta.py
@@ -61,7 +61,7 @@ def doubling_beta_CATE_on_csv(
         {"avg_age", "contacts"},  # We use custom adjustment set
         "cum_infections",
         df=past_execution_df,
-        formula="cum_infections ~ beta + np.power(beta, 2) + avg_age + contacts",
+        formula="cum_infections ~ beta + I(beta ** 2) + avg_age + contacts",
     )
 
     # Add squared terms for beta, since it has a quadratic relationship with cumulative infections
@@ -77,7 +77,7 @@ def doubling_beta_CATE_on_csv(
         set(),
         "cum_infections",
         df=past_execution_df,
-        formula="cum_infections ~ beta + np.power(beta, 2)",
+        formula="cum_infections ~ beta + I(beta ** 2)",
     )
     association_test_result = causal_test_case.execute_test(
         estimator=no_adjustment_linear_regression_estimator, data_collector=data_collector
@@ -109,7 +109,7 @@ def doubling_beta_CATE_on_csv(
             {"avg_age", "contacts"},
             "cum_infections",
             df=counterfactual_past_execution_df,
-            formula="cum_infections ~ beta + np.power(beta, 2) + avg_age + contacts",
+            formula="cum_infections ~ beta + I(beta ** 2) + avg_age + contacts",
         )
         counterfactual_causal_test_result = causal_test_case.execute_test(
             estimator=linear_regression_estimator, data_collector=data_collector
diff --git a/examples/poisson-line-process/example_poisson_process.py b/examples/poisson-line-process/example_poisson_process.py
index 8a2f12a8..da6ce4d9 100644
--- a/examples/poisson-line-process/example_poisson_process.py
+++ b/examples/poisson-line-process/example_poisson_process.py
@@ -2,7 +2,6 @@
 from causal_testing.specification.scenario import Scenario
 from causal_testing.specification.variable import Input, Output
 from causal_testing.specification.causal_specification import CausalSpecification
-from causal_testing.data_collection.data_collector import ObservationalDataCollector
 from causal_testing.testing.causal_test_case import CausalTestCase
 from causal_testing.testing.causal_test_outcome import ExactValue, Positive
 from causal_testing.estimation.linear_regression_estimator import LinearRegressionEstimator, Estimator
@@ -81,14 +80,11 @@ def causal_test_intensity_num_shapes(
     inverse_terms=[],
     empirical=False,
 ):
-    # 6. Create a data collector
-    data_collector = ObservationalDataCollector(scenario, pd.read_csv(observational_data_path))
-
     # 7. Obtain the minimal adjustment set for the causal test case from the causal DAG
     minimal_adjustment_set = causal_dag.identification(causal_test_case.base_test_case)
 
     # 8. Set up an estimator
-    data = pd.read_csv(observational_data_path)
+    data = pd.read_csv(observational_data_path, index_col=0).astype(float)
 
     treatment = causal_test_case.treatment_variable.name
     outcome = causal_test_case.outcome_variable.name
@@ -105,8 +101,8 @@ def causal_test_intensity_num_shapes(
             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]
+        square_terms = [f"I({t} ** 2)" for t in square_terms]
+        inverse_terms = [f"I({t} ** -1)" for t in inverse_terms]
         estimator = LinearRegressionEstimator(
             treatment=treatment,
             control_value=causal_test_case.control_value,
@@ -119,7 +115,7 @@ def causal_test_intensity_num_shapes(
         )
 
     # 9. Execute the test
-    causal_test_result = causal_test_case.execute_test(estimator, data_collector)
+    causal_test_result = causal_test_case.execute_test(estimator, None)
 
     return causal_test_result
 
@@ -175,8 +171,8 @@ def test_poisson_width_num_shapes(save=False):
             logger.info("%s CAUSAL TEST %s", "=" * 33, "=" * 33)
             logger.info("Identifying")
             # 5. Create a causal test case
-            control_value = w
-            treatment_value = w + 1
+            control_value = float(w)
+            treatment_value = w + 1.0
             base_test_case = BaseTestCase(treatment_variable=width, outcome_variable=num_shapes_unit)
             causal_test_case = CausalTestCase(
                 base_test_case=base_test_case,
diff --git a/tests/estimation_tests/test_linear_regression_estimator.py b/tests/estimation_tests/test_linear_regression_estimator.py
index e6ea4206..264a7278 100644
--- a/tests/estimation_tests/test_linear_regression_estimator.py
+++ b/tests/estimation_tests/test_linear_regression_estimator.py
@@ -92,7 +92,7 @@ def test_program_11_3(self):
         """Test whether our linear regression implementation produces the same results as program 11.3 (p. 144)."""
         df = self.chapter_11_df.copy()
         linear_regression_estimator = LinearRegressionEstimator(
-            "treatments", None, None, set(), "outcomes", df, formula="outcomes ~ treatments + np.power(treatments, 2)"
+            "treatments", None, None, set(), "outcomes", df, formula="outcomes ~ treatments + I(treatments ** 2)"
         )
         model = linear_regression_estimator._run_linear_regression()
         ate, _ = linear_regression_estimator.estimate_coefficient()
@@ -100,7 +100,7 @@ def test_program_11_3(self):
             round(
                 model.params["Intercept"]
                 + 90 * model.params["treatments"]
-                + 90 * 90 * model.params["np.power(treatments, 2)"],
+                + 90 * 90 * model.params["I(treatments ** 2)"],
                 1,
             ),
             197.1,
@@ -136,10 +136,10 @@ def test_program_15_1A(self):
             df,
             formula=f"""wt82_71 ~ qsmk +
                              {'+'.join(sorted(list(covariates)))} +
-                             np.power(age, 2) +
-                             np.power(wt71, 2) +
-                             np.power(smokeintensity, 2) +
-                             np.power(smokeyrs, 2) +
+                             I(age ** 2) +
+                             I(wt71 ** 2) +
+                             I(smokeintensity ** 2) +
+                             I(smokeyrs ** 2) +
                              (qsmk * smokeintensity)""",
         )
         # terms_to_square = ["age", "wt71", "smokeintensity", "smokeyrs"]
@@ -179,7 +179,7 @@ def test_program_15_no_interaction(self):
             covariates,
             "wt82_71",
             df,
-            formula="wt82_71 ~ qsmk + age + np.power(age, 2) + wt71 + np.power(wt71, 2) + smokeintensity + np.power(smokeintensity, 2) + smokeyrs + np.power(smokeyrs, 2)",
+            formula="wt82_71 ~ qsmk + age + I(age ** 2) + wt71 + I(wt71 ** 2) + smokeintensity + I(smokeintensity ** 2) + smokeyrs + I(smokeyrs ** 2)",
         )
         # terms_to_square = ["age", "wt71", "smokeintensity", "smokeyrs"]
         # for term_to_square in terms_to_square:
@@ -215,7 +215,7 @@ def test_program_15_no_interaction_ate(self):
             covariates,
             "wt82_71",
             df,
-            formula="wt82_71 ~ qsmk + age + np.power(age, 2) + wt71 + np.power(wt71, 2) + smokeintensity + np.power(smokeintensity, 2) + smokeyrs + np.power(smokeyrs, 2)",
+            formula="wt82_71 ~ qsmk + age + I(age ** 2) + wt71 + I(wt71 ** 2) + smokeintensity + I(smokeintensity ** 2) + smokeyrs + I(smokeyrs ** 2)",
         )
         # terms_to_square = ["age", "wt71", "smokeintensity", "smokeyrs"]
         # for term_to_square in terms_to_square:
@@ -250,7 +250,7 @@ def test_program_15_no_interaction_ate_calculated(self):
             covariates,
             "wt82_71",
             df,
-            formula="wt82_71 ~ qsmk + age + np.power(age, 2) + wt71 + np.power(wt71, 2) + smokeintensity + np.power(smokeintensity, 2) + smokeyrs + np.power(smokeyrs, 2)",
+            formula="wt82_71 ~ qsmk + age + I(age ** 2) + wt71 + I(wt71 ** 2) + smokeintensity + I(smokeintensity ** 2) + smokeyrs + I(smokeyrs ** 2)",
         )
         # terms_to_square = ["age", "wt71", "smokeintensity", "smokeyrs"]
         # for term_to_square in terms_to_square:

From 62e6b3d2e5d4fce6d1f27795e6be687ace1f411e Mon Sep 17 00:00:00 2001
From: Michael <jmafoster1@gmail.com>
Date: Wed, 7 Aug 2024 15:45:43 +0100
Subject: [PATCH 05/24] pytest

---
 causal_testing/estimation/gp.py               | 21 +++++++++++++---
 .../estimation/linear_regression_estimator.py | 16 ++++++------
 .../test_linear_regression_estimator.py       | 25 +++++++++++++++----
 3 files changed, 46 insertions(+), 16 deletions(-)

diff --git a/causal_testing/estimation/gp.py b/causal_testing/estimation/gp.py
index f1030c8d..4f3326cc 100644
--- a/causal_testing/estimation/gp.py
+++ b/causal_testing/estimation/gp.py
@@ -85,33 +85,46 @@ def __init__(
         df: pd.DataFrame,
         features: list,
         outcome: str,
+        max_order: int = 0,
         extra_operators: list = None,
         sympy_conversions: dict = None,
         seed=0,
     ):
         # pylint: disable=too-many-arguments
         random.seed(seed)
-        np.random.seed(seed)
         self.df = df
         self.features = features
         self.outcome = outcome
+        self.max_order = max_order
         self.seed = seed
         self.pset = gp.PrimitiveSet("MAIN", len(self.features))
         self.pset.renameArguments(**{f"ARG{i}": f for i, f in enumerate(self.features)})
 
-        standard_operators = [(add, 2), (mul, 2), (reciprocal, 1)]
+        standard_operators = [(add, 2), (mul, 2)]
         if extra_operators is None:
             extra_operators = [(log, 1), (reciprocal, 1)]
-        for operator, num_args in standard_operators + extra_operators:
-            self.pset.addPrimitive(operator, num_args)
         if sympy_conversions is None:
             sympy_conversions = {}
+        for operator, num_args in standard_operators + extra_operators:
+            self.pset.addPrimitive(operator, num_args)
+
         self.sympy_conversions = {
             "mul": lambda x1, x2: f"Mul({x1},{x2})",
             "add": lambda x1, x2: f"Add({x1},{x2})",
             "reciprocal": lambda x1: f"Pow({x1},-1)",
         } | sympy_conversions
 
+        for i in range(self.max_order):
+            print("Adding in order", i)
+            name = f"power_{i}"
+            self.pset.addPrimitive(lambda x: power(x, i), 1, name=name)
+            if name in self.sympy_conversions:
+                raise ValueError(
+                    f"You have provided a function called {name}, which is reserved for raising to power"
+                    f"{i}. Please choose a different name for your function."
+                )
+            self.sympy_conversions[name] = lambda x1: f"Pow({x1},{i})"
+
         creator.create("FitnessMin", base.Fitness, weights=(-1.0,))
         creator.create("Individual", gp.PrimitiveTree, fitness=creator.FitnessMin)
 
diff --git a/causal_testing/estimation/linear_regression_estimator.py b/causal_testing/estimation/linear_regression_estimator.py
index 34ea31ce..cb42fc22 100644
--- a/causal_testing/estimation/linear_regression_estimator.py
+++ b/causal_testing/estimation/linear_regression_estimator.py
@@ -63,20 +63,21 @@ def __init__(
     def gp_formula(
         self,
         ngen: int = 100,
-        mu: int = 20,
-        lambda_: int = 10,
+        pop_size: int = 20,
+        num_offspring: int = 10,
         extra_operators: list = None,
         sympy_conversions: dict = None,
+        max_order: int = 0,
         seeds: list = None,
         seed: int = 0,
     ):
-        # pylint: disable=too-many-arguments,invalid-name
+        # pylint: disable=too-many-arguments
         """
         Use Genetic Programming (GP) to infer the regression equation from the data.
 
         :param ngen: The maximum number of GP generations to run for.
         :param mu: The GP population size.
-        :param lambda_: The number of offspring per generation.
+        :param num_offspring: The number of offspring per generation.
         :param extra_operators: Additional operators for the GP (defaults are +, *, and 1/x). Operations should be of
                                 the form (fun, numArgs), e.g. (add, 2).
         :param sympy_conversions: Dictionary of conversions of extra_operators for sympy,
@@ -85,16 +86,17 @@ def gp_formula(
                       probably logarithmic, you can put that in).
         :param seed: Random seed for the GP.
         """
-        gp = GP(
+        self.gp = GP(
             df=self.df,
             features=sorted(list(self.adjustment_set.union([self.treatment]))),
             outcome=self.outcome,
             extra_operators=extra_operators,
             sympy_conversions=sympy_conversions,
             seed=seed,
+            max_order=max_order,
         )
-        formula = gp.run_gp(ngen=ngen, pop_size=mu, num_offspring=lambda_, seeds=seeds)
-        formula = gp.simplify(formula)
+        formula = self.gp.run_gp(ngen=ngen, pop_size=num_offspring, num_offspring=num_offspring, seeds=seeds)
+        formula = self.gp.simplify(formula)
         self.formula = f"{self.outcome} ~ I({formula}) - 1"
 
     def add_modelling_assumptions(self):
diff --git a/tests/estimation_tests/test_linear_regression_estimator.py b/tests/estimation_tests/test_linear_regression_estimator.py
index 264a7278..0eda0d21 100644
--- a/tests/estimation_tests/test_linear_regression_estimator.py
+++ b/tests/estimation_tests/test_linear_regression_estimator.py
@@ -7,6 +7,7 @@
 from causal_testing.specification.capabilities import TreatmentSequence
 
 from causal_testing.estimation.linear_regression_estimator import LinearRegressionEstimator
+from causal_testing.estimation.gp import reciprocal
 
 
 def load_nhefs_df():
@@ -275,16 +276,30 @@ def test_gp(self):
         df["X"] = np.arange(10)
         df["Y"] = 1 / (df["X"] + 1)
         linear_regression_estimator = LinearRegressionEstimator("X", 0, 1, set(), "Y", df.astype(float))
-        linear_regression_estimator.gp_formula(seed=1)
-        self.assertEqual(
-            linear_regression_estimator.formula,
-            "Y ~ I((2.606801258739728e-17*X + 0.626132756132756)/(0.6261327561327561*X + 0.626132756132756)) - 1",
-        )
+        linear_regression_estimator.gp_formula(seeds=["reciprocal(add(X, 1))"])
+        print("MAPPING")
+        print(linear_regression_estimator.gp.pset.mapping)
+        self.assertEqual(linear_regression_estimator.formula, "Y ~ I(1/(X + 1)) - 1")
         ate, (ci_low, ci_high) = linear_regression_estimator.estimate_ate_calculated()
         self.assertEqual(round(ate[0], 2), 0.50)
         self.assertEqual(round(ci_low[0], 2), 0.50)
         self.assertEqual(round(ci_high[0], 2), 0.50)
 
+    def test_gp_power(self):
+        df = pd.DataFrame()
+        df["X"] = np.arange(10)
+        df["Y"] = 2 * (df["X"] ** 2)
+        linear_regression_estimator = LinearRegressionEstimator("X", 0, 1, set(), "Y", df.astype(float))
+        linear_regression_estimator.gp_formula(seed=1, max_order=0)
+        self.assertEqual(
+            linear_regression_estimator.formula,
+            "Y ~ I(2.0*X**2 + 3.8205100524608823e-31) - 1",
+        )
+        ate, (ci_low, ci_high) = linear_regression_estimator.estimate_ate_calculated()
+        self.assertEqual(round(ate[0], 2), -2.00)
+        self.assertEqual(round(ci_low[0], 2), -2.00)
+        self.assertEqual(round(ci_high[0], 2), -2.00)
+
 
 class TestLinearRegressionInteraction(unittest.TestCase):
     """Test linear regression for estimating effects involving interaction."""

From 5d915ed53a0756c1da4e95cad7b3b9c9b098c633 Mon Sep 17 00:00:00 2001
From: Michael <jmafoster1@gmail.com>
Date: Thu, 8 Aug 2024 09:24:32 +0100
Subject: [PATCH 06/24] pylint

---
 causal_testing/estimation/gp.py                | 18 ++++++++++++++----
 .../estimation/linear_regression_estimator.py  | 13 +++++++------
 .../test_linear_regression_estimator.py        |  2 --
 3 files changed, 21 insertions(+), 12 deletions(-)

diff --git a/causal_testing/estimation/gp.py b/causal_testing/estimation/gp.py
index 4f3326cc..fbac0353 100644
--- a/causal_testing/estimation/gp.py
+++ b/causal_testing/estimation/gp.py
@@ -75,6 +75,16 @@ def mut_insert(expression: gp.PrimitiveTree, pset: gp.PrimitiveSet):
     return (expression,)
 
 
+def create_power_function(power):
+    def power_func(x):
+        return power(x, power)
+
+    def sympy_conversion(x1):
+        return f"Pow({x1},{i})"
+
+    return power_func, sympy_conversion
+
+
 class GP:
     """
     Object to perform genetic programming.
@@ -90,7 +100,7 @@ def __init__(
         sympy_conversions: dict = None,
         seed=0,
     ):
-        # pylint: disable=too-many-arguments
+        # pylint: disable=too-many-arguments,too-many-instance-attributes
         random.seed(seed)
         self.df = df
         self.features = features
@@ -115,15 +125,15 @@ def __init__(
         } | sympy_conversions
 
         for i in range(self.max_order):
-            print("Adding in order", i)
             name = f"power_{i}"
-            self.pset.addPrimitive(lambda x: power(x, i), 1, name=name)
+            func, conversion = create_power_function(i)
+            self.pset.addPrimitive(func, 1, name=name)
             if name in self.sympy_conversions:
                 raise ValueError(
                     f"You have provided a function called {name}, which is reserved for raising to power"
                     f"{i}. Please choose a different name for your function."
                 )
-            self.sympy_conversions[name] = lambda x1: f"Pow({x1},{i})"
+            self.sympy_conversions[name] = conversion
 
         creator.create("FitnessMin", base.Fitness, weights=(-1.0,))
         creator.create("Individual", gp.PrimitiveTree, fitness=creator.FitnessMin)
diff --git a/causal_testing/estimation/linear_regression_estimator.py b/causal_testing/estimation/linear_regression_estimator.py
index cb42fc22..3c94e274 100644
--- a/causal_testing/estimation/linear_regression_estimator.py
+++ b/causal_testing/estimation/linear_regression_estimator.py
@@ -65,9 +65,9 @@ def gp_formula(
         ngen: int = 100,
         pop_size: int = 20,
         num_offspring: int = 10,
+        max_order: int = 0,
         extra_operators: list = None,
         sympy_conversions: dict = None,
-        max_order: int = 0,
         seeds: list = None,
         seed: int = 0,
     ):
@@ -76,9 +76,10 @@ def gp_formula(
         Use Genetic Programming (GP) to infer the regression equation from the data.
 
         :param ngen: The maximum number of GP generations to run for.
-        :param mu: The GP population size.
+        :param pop_size: The GP population size.
         :param num_offspring: The number of offspring per generation.
-        :param extra_operators: Additional operators for the GP (defaults are +, *, and 1/x). Operations should be of
+        :param max_order: The maximum polynomial order to use, e.g. `max_order=2` will give polynomials of the form `ax^2 + bx + c`.
+        :param extra_operators: Additional operators for the GP (defaults are +, *, log(x), and 1/x). Operations should be of
                                 the form (fun, numArgs), e.g. (add, 2).
         :param sympy_conversions: Dictionary of conversions of extra_operators for sympy,
                                   e.g. `"mul": lambda *args_: "Mul({},{})".format(*args_)`.
@@ -86,7 +87,7 @@ def gp_formula(
                       probably logarithmic, you can put that in).
         :param seed: Random seed for the GP.
         """
-        self.gp = GP(
+        gp = GP(
             df=self.df,
             features=sorted(list(self.adjustment_set.union([self.treatment]))),
             outcome=self.outcome,
@@ -95,8 +96,8 @@ def gp_formula(
             seed=seed,
             max_order=max_order,
         )
-        formula = self.gp.run_gp(ngen=ngen, pop_size=num_offspring, num_offspring=num_offspring, seeds=seeds)
-        formula = self.gp.simplify(formula)
+        formula = gp.run_gp(ngen=ngen, pop_size=pop_size, num_offspring=num_offspring, seeds=seeds)
+        formula = gp.simplify(formula)
         self.formula = f"{self.outcome} ~ I({formula}) - 1"
 
     def add_modelling_assumptions(self):
diff --git a/tests/estimation_tests/test_linear_regression_estimator.py b/tests/estimation_tests/test_linear_regression_estimator.py
index 0eda0d21..42d652f3 100644
--- a/tests/estimation_tests/test_linear_regression_estimator.py
+++ b/tests/estimation_tests/test_linear_regression_estimator.py
@@ -277,8 +277,6 @@ def test_gp(self):
         df["Y"] = 1 / (df["X"] + 1)
         linear_regression_estimator = LinearRegressionEstimator("X", 0, 1, set(), "Y", df.astype(float))
         linear_regression_estimator.gp_formula(seeds=["reciprocal(add(X, 1))"])
-        print("MAPPING")
-        print(linear_regression_estimator.gp.pset.mapping)
         self.assertEqual(linear_regression_estimator.formula, "Y ~ I(1/(X + 1)) - 1")
         ate, (ci_low, ci_high) = linear_regression_estimator.estimate_ate_calculated()
         self.assertEqual(round(ate[0], 2), 0.50)

From 5b0661d7150fe20442d03b4a3e7016d6c021f4fe Mon Sep 17 00:00:00 2001
From: Michael <jmafoster1@gmail.com>
Date: Thu, 8 Aug 2024 09:30:49 +0100
Subject: [PATCH 07/24] pylint

---
 causal_testing/estimation/gp.py               | 19 ++++++++++++++-----
 .../estimation/linear_regression_estimator.py |  7 ++++---
 2 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/causal_testing/estimation/gp.py b/causal_testing/estimation/gp.py
index fbac0353..567143dc 100644
--- a/causal_testing/estimation/gp.py
+++ b/causal_testing/estimation/gp.py
@@ -75,12 +75,20 @@ def mut_insert(expression: gp.PrimitiveTree, pset: gp.PrimitiveSet):
     return (expression,)
 
 
-def create_power_function(power):
+def create_power_function(order: int):
+    """
+    Creates a power operator and its corresponding sympy conversion.
+
+    :param order: The order of the power, e.g. `order=2` will give x^2.
+
+    :return: A pair consisting of the power function and the sympy conversion
+    """
+
     def power_func(x):
-        return power(x, power)
+        return power(x, order)
 
-    def sympy_conversion(x1):
-        return f"Pow({x1},{i})"
+    def sympy_conversion(x):
+        return f"Pow({x},{order})"
 
     return power_func, sympy_conversion
 
@@ -100,7 +108,8 @@ def __init__(
         sympy_conversions: dict = None,
         seed=0,
     ):
-        # pylint: disable=too-many-arguments,too-many-instance-attributes
+        # pylint: disable=too-many-arguments
+        # pylint: disable=too-many-instance-attributes
         random.seed(seed)
         self.df = df
         self.features = features
diff --git a/causal_testing/estimation/linear_regression_estimator.py b/causal_testing/estimation/linear_regression_estimator.py
index 3c94e274..f902afcc 100644
--- a/causal_testing/estimation/linear_regression_estimator.py
+++ b/causal_testing/estimation/linear_regression_estimator.py
@@ -78,9 +78,10 @@ def gp_formula(
         :param ngen: The maximum number of GP generations to run for.
         :param pop_size: The GP population size.
         :param num_offspring: The number of offspring per generation.
-        :param max_order: The maximum polynomial order to use, e.g. `max_order=2` will give polynomials of the form `ax^2 + bx + c`.
-        :param extra_operators: Additional operators for the GP (defaults are +, *, log(x), and 1/x). Operations should be of
-                                the form (fun, numArgs), e.g. (add, 2).
+        :param max_order: The maximum polynomial order to use, e.g. `max_order=2` will give
+                          polynomials of the form `ax^2 + bx + c`.
+        :param extra_operators: Additional operators for the GP (defaults are +, *, log(x), and 1/x).
+                                Operations should be of the form (fun, numArgs), e.g. (add, 2).
         :param sympy_conversions: Dictionary of conversions of extra_operators for sympy,
                                   e.g. `"mul": lambda *args_: "Mul({},{})".format(*args_)`.
         :param seeds: Seed individuals for the population (e.g. if you think that the relationship between X and Y is

From e7deb788c3de756ef67ebcca9375fee01fdc82b9 Mon Sep 17 00:00:00 2001
From: Michael <jmafoster1@gmail.com>
Date: Thu, 8 Aug 2024 10:54:19 +0100
Subject: [PATCH 08/24] all tests pass

---
 causal_testing/estimation/gp.py               |  3 +-
 .../logistic_regression_estimator.py          |  3 +-
 .../estimation/regression_estimator.py        | 86 +++++++++++++++++++
 .../example_poisson_process.py                |  3 +-
 .../test_linear_regression_estimator.py       | 43 +++++++---
 tests/json_front_tests/test_json_class.py     |  5 +-
 6 files changed, 124 insertions(+), 19 deletions(-)
 create mode 100644 causal_testing/estimation/regression_estimator.py

diff --git a/causal_testing/estimation/gp.py b/causal_testing/estimation/gp.py
index 567143dc..ab53d026 100644
--- a/causal_testing/estimation/gp.py
+++ b/causal_testing/estimation/gp.py
@@ -98,6 +98,8 @@ class GP:
     Object to perform genetic programming.
     """
 
+    # pylint: disable=too-many-instance-attributes
+
     def __init__(
         self,
         df: pd.DataFrame,
@@ -109,7 +111,6 @@ def __init__(
         seed=0,
     ):
         # pylint: disable=too-many-arguments
-        # pylint: disable=too-many-instance-attributes
         random.seed(seed)
         self.df = df
         self.features = features
diff --git a/causal_testing/estimation/logistic_regression_estimator.py b/causal_testing/estimation/logistic_regression_estimator.py
index 1e180c49..6d08fa94 100644
--- a/causal_testing/estimation/logistic_regression_estimator.py
+++ b/causal_testing/estimation/logistic_regression_estimator.py
@@ -47,7 +47,8 @@ def __init__(
         )
 
         self.model = None
-
+        if effect_modifiers is None:
+            effect_modifiers = []
         if formula is not None:
             self.formula = formula
         else:
diff --git a/causal_testing/estimation/regression_estimator.py b/causal_testing/estimation/regression_estimator.py
new file mode 100644
index 00000000..d713840a
--- /dev/null
+++ b/causal_testing/estimation/regression_estimator.py
@@ -0,0 +1,86 @@
+"""This module contains the RegressionEstimator, which is an abstract class for concrete regression estimators."""
+
+import logging
+from typing import Any
+from abc import abstractmethod, abstractmethod
+
+import pandas as pd
+import statsmodels.formula.api as smf
+from patsy import dmatrix  # pylint: disable = no-name-in-module
+from patsy import ModelDesc
+from statsmodels.regression.linear_model import RegressionResultsWrapper
+
+from causal_testing.specification.variable import Variable
+from causal_testing.estimation.gp import GP
+from causal_testing.estimation.estimator import Estimator
+
+logger = logging.getLogger(__name__)
+
+
+class RegressionEstimator(Estimator):
+    """A Linear Regression Estimator is a parametric estimator which restricts the variables in the data to a linear
+    combination of parameters and functions of the variables (note these functions need not be linear).
+    """
+
+    def __init__(
+        # pylint: disable=too-many-arguments
+        self,
+        treatment: str,
+        treatment_value: float,
+        control_value: float,
+        adjustment_set: set,
+        outcome: str,
+        df: pd.DataFrame = None,
+        effect_modifiers: dict[Variable:Any] = None,
+        formula: str = None,
+        alpha: float = 0.05,
+        query: str = "",
+    ):
+        super().__init__(
+            treatment=treatment,
+            treatment_value=treatment_value,
+            control_value=control_value,
+            adjustment_set=adjustment_set,
+            outcome=outcome,
+            df=df,
+            effect_modifiers=effect_modifiers,
+            query=query,
+        )
+
+        self.model = None
+        if effect_modifiers is None:
+            effect_modifiers = []
+        if formula is not None:
+            self.formula = formula
+        else:
+            terms = [treatment] + sorted(list(adjustment_set)) + sorted(list(effect_modifiers))
+            self.formula = f"{outcome} ~ {'+'.join(terms)}"
+        for term in self.effect_modifiers:
+            self.adjustment_set.add(term)
+
+    @property
+    @abstractmethod
+    def regression(self):
+        raise NotImplementedError("Subclasses must implement the 'model' property.")
+
+    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.append(
+            "The variables in the data must fit a shape which can be expressed as a linear"
+            "combination of parameters and functions of variables. Note that these functions"
+            "do not need to be linear."
+        )
+
+    def _run_regression(self, data=None) -> RegressionResultsWrapper:
+        """Run logistic regression of the treatment and adjustment set against the outcome and return the model.
+
+        :return: The model after fitting to data.
+        """
+        if data is None:
+            data = self.df
+        model = self.regression(formula=self.formula, data=data).fit(disp=0)
+        self.model = model
+        return model
diff --git a/examples/poisson-line-process/example_poisson_process.py b/examples/poisson-line-process/example_poisson_process.py
index da6ce4d9..a8d4addd 100644
--- a/examples/poisson-line-process/example_poisson_process.py
+++ b/examples/poisson-line-process/example_poisson_process.py
@@ -4,7 +4,8 @@
 from causal_testing.specification.causal_specification import CausalSpecification
 from causal_testing.testing.causal_test_case import CausalTestCase
 from causal_testing.testing.causal_test_outcome import ExactValue, Positive
-from causal_testing.estimation.linear_regression_estimator import LinearRegressionEstimator, Estimator
+from causal_testing.estimation.linear_regression_estimator import LinearRegressionEstimator
+from causal_testing.estimation.estimator import Estimator
 from causal_testing.testing.base_test_case import BaseTestCase
 
 import pandas as pd
diff --git a/tests/estimation_tests/test_linear_regression_estimator.py b/tests/estimation_tests/test_linear_regression_estimator.py
index 42d652f3..6a6dec10 100644
--- a/tests/estimation_tests/test_linear_regression_estimator.py
+++ b/tests/estimation_tests/test_linear_regression_estimator.py
@@ -81,13 +81,24 @@ def test_program_11_2(self):
         """Test whether our linear regression implementation produces the same results as program 11.2 (p. 141)."""
         df = self.chapter_11_df
         linear_regression_estimator = LinearRegressionEstimator("treatments", None, None, set(), "outcomes", df)
-        model = linear_regression_estimator._run_linear_regression()
         ate, _ = linear_regression_estimator.estimate_coefficient()
 
-        self.assertEqual(round(model.params["Intercept"] + 90 * model.params["treatments"], 1), 216.9)
+        self.assertEqual(
+            round(
+                linear_regression_estimator.model.params["Intercept"]
+                + 90 * linear_regression_estimator.model.params["treatments"],
+                1,
+            ),
+            216.9,
+        )
 
         # Increasing treatments from 90 to 100 should be the same as 10 times the unit ATE
-        self.assertTrue(all(round(model.params["treatments"], 1) == round(ate_single, 1) for ate_single in ate))
+        self.assertTrue(
+            all(
+                round(linear_regression_estimator.model.params["treatments"], 1) == round(ate_single, 1)
+                for ate_single in ate
+            )
+        )
 
     def test_program_11_3(self):
         """Test whether our linear regression implementation produces the same results as program 11.3 (p. 144)."""
@@ -95,19 +106,23 @@ def test_program_11_3(self):
         linear_regression_estimator = LinearRegressionEstimator(
             "treatments", None, None, set(), "outcomes", df, formula="outcomes ~ treatments + I(treatments ** 2)"
         )
-        model = linear_regression_estimator._run_linear_regression()
         ate, _ = linear_regression_estimator.estimate_coefficient()
         self.assertEqual(
             round(
-                model.params["Intercept"]
-                + 90 * model.params["treatments"]
-                + 90 * 90 * model.params["I(treatments ** 2)"],
+                linear_regression_estimator.model.params["Intercept"]
+                + 90 * linear_regression_estimator.model.params["treatments"]
+                + 90 * 90 * linear_regression_estimator.model.params["I(treatments ** 2)"],
                 1,
             ),
             197.1,
         )
         # Increasing treatments from 90 to 100 should be the same as 10 times the unit ATE
-        self.assertTrue(all(round(model.params["treatments"], 3) == round(ate_single, 3) for ate_single in ate))
+        self.assertTrue(
+            all(
+                round(linear_regression_estimator.model.params["treatments"], 3) == round(ate_single, 3)
+                for ate_single in ate
+            )
+        )
 
     def test_program_15_1A(self):
         """Test whether our linear regression implementation produces the same results as program 15.1 (p. 163, 184)."""
@@ -149,9 +164,9 @@ def test_program_15_1A(self):
         # for term_a, term_b in terms_to_product:
         #     linear_regression_estimator.add_product_term_to_df(term_a, term_b)
 
-        model = linear_regression_estimator._run_linear_regression()
-        self.assertEqual(round(model.params["qsmk"], 1), 2.6)
-        self.assertEqual(round(model.params["qsmk:smokeintensity"], 2), 0.05)
+        linear_regression_estimator.estimate_coefficient()
+        self.assertEqual(round(linear_regression_estimator.model.params["qsmk"], 1), 2.6)
+        self.assertEqual(round(linear_regression_estimator.model.params["qsmk:smokeintensity"], 2), 0.05)
 
     def test_program_15_no_interaction(self):
         """Test whether our linear regression implementation produces the same results as program 15.1 (p. 163, 184)
@@ -266,10 +281,10 @@ def test_program_11_2_with_robustness_validation(self):
         """Test whether our linear regression estimator, as used in test_program_11_2 can correctly estimate robustness."""
         df = self.chapter_11_df.copy()
         linear_regression_estimator = LinearRegressionEstimator("treatments", 100, 90, set(), "outcomes", df)
-        model = linear_regression_estimator._run_linear_regression()
+        linear_regression_estimator.estimate_coefficient()
 
         cv = CausalValidator()
-        self.assertEqual(round(cv.estimate_robustness(model)["treatments"], 4), 0.7353)
+        self.assertEqual(round(cv.estimate_robustness(linear_regression_estimator.model)["treatments"], 4), 0.7353)
 
     def test_gp(self):
         df = pd.DataFrame()
@@ -291,7 +306,7 @@ def test_gp_power(self):
         linear_regression_estimator.gp_formula(seed=1, max_order=0)
         self.assertEqual(
             linear_regression_estimator.formula,
-            "Y ~ I(2.0*X**2 + 3.8205100524608823e-31) - 1",
+            "Y ~ I(1.9999999999999999*X**2 - 1.0043240235058056e-116*X + 2.6645352591003757e-15) - 1",
         )
         ate, (ci_low, ci_high) = linear_regression_estimator.estimate_ate_calculated()
         self.assertEqual(round(ate[0], 2), -2.00)
diff --git a/tests/json_front_tests/test_json_class.py b/tests/json_front_tests/test_json_class.py
index b6585485..3aff9782 100644
--- a/tests/json_front_tests/test_json_class.py
+++ b/tests/json_front_tests/test_json_class.py
@@ -4,7 +4,8 @@
 import scipy
 import os
 
-from causal_testing.estimation.linear_regression_estimator import LinearRegressionEstimator, Estimator
+from causal_testing.estimation.linear_regression_estimator import LinearRegressionEstimator
+from causal_testing.estimation.estimator import Estimator
 from causal_testing.testing.causal_test_outcome import NoEffect, Positive
 from causal_testing.json_front.json_class import JsonUtility, CausalVariables
 from causal_testing.specification.variable import Input, Output, Meta
@@ -313,7 +314,7 @@ def add_modelling_assumptions(self):
         effects = {"Positive": Positive()}
         mutates = {
             "Increase": lambda x: self.json_class.scenario.treatment_variables[x].z3
-                                  > self.json_class.scenario.variables[x].z3
+            > self.json_class.scenario.variables[x].z3
         }
         estimators = {"ExampleEstimator": ExampleEstimator}
         with self.assertRaises(TypeError):

From 05c5499629e6c181f185c9d8c6ad4c7db9cf351d Mon Sep 17 00:00:00 2001
From: Michael <jmafoster1@gmail.com>
Date: Thu, 8 Aug 2024 12:19:19 +0100
Subject: [PATCH 09/24] RegressionEstimator class to combine common elements of
 Linear and Logistic Regression Estimator classes.

---
 .../estimation/cubic_spline_estimator.py      |  2 +-
 .../estimation/linear_regression_estimator.py | 77 ++++-------------
 .../logistic_regression_estimator.py          | 85 ++-----------------
 .../estimation/regression_estimator.py        | 52 ++++++++++--
 .../test_cubic_spline_estimator.py            |  3 +-
 5 files changed, 68 insertions(+), 151 deletions(-)

diff --git a/causal_testing/estimation/cubic_spline_estimator.py b/causal_testing/estimation/cubic_spline_estimator.py
index d1ca72b5..3467156e 100644
--- a/causal_testing/estimation/cubic_spline_estimator.py
+++ b/causal_testing/estimation/cubic_spline_estimator.py
@@ -46,7 +46,7 @@ def __init__(
             self.formula = f"{outcome} ~ cr({'+'.join(terms)}, df={basis})"
 
     def estimate_ate_calculated(self, adjustment_config: dict = None) -> pd.Series:
-        model = self._run_linear_regression()
+        model = self._run_regression()
 
         x = {"Intercept": 1, self.treatment: self.treatment_value}
         if adjustment_config is not None:
diff --git a/causal_testing/estimation/linear_regression_estimator.py b/causal_testing/estimation/linear_regression_estimator.py
index f902afcc..1ab00f45 100644
--- a/causal_testing/estimation/linear_regression_estimator.py
+++ b/causal_testing/estimation/linear_regression_estimator.py
@@ -5,22 +5,22 @@
 
 import pandas as pd
 import statsmodels.formula.api as smf
-from patsy import dmatrix  # pylint: disable = no-name-in-module
-from patsy import ModelDesc
-from statsmodels.regression.linear_model import RegressionResultsWrapper
+from patsy import dmatrix, ModelDesc  # pylint: disable = no-name-in-module
 
 from causal_testing.specification.variable import Variable
 from causal_testing.estimation.gp import GP
-from causal_testing.estimation.estimator import Estimator
+from causal_testing.estimation.regression_estimator import RegressionEstimator
 
 logger = logging.getLogger(__name__)
 
 
-class LinearRegressionEstimator(Estimator):
+class LinearRegressionEstimator(RegressionEstimator):
     """A Linear Regression Estimator is a parametric estimator which restricts the variables in the data to a linear
     combination of parameters and functions of the variables (note these functions need not be linear).
     """
 
+    regressor = smf.ols
+
     def __init__(
         # pylint: disable=too-many-arguments
         self,
@@ -35,6 +35,7 @@ def __init__(
         alpha: float = 0.05,
         query: str = "",
     ):
+        # pylint: disable=too-many-arguments
         super().__init__(
             treatment,
             treatment_value,
@@ -43,20 +44,10 @@ def __init__(
             outcome,
             df,
             effect_modifiers,
-            alpha=alpha,
-            query=query,
+            formula,
+            alpha,
+            query,
         )
-
-        self.model = None
-        if effect_modifiers is None:
-            effect_modifiers = []
-
-        if formula is not None:
-            self.formula = formula
-        else:
-            terms = [treatment] + sorted(list(adjustment_set)) + sorted(list(effect_modifiers))
-            self.formula = f"{outcome} ~ {'+'.join(terms)}"
-
         for term in self.effect_modifiers:
             self.adjustment_set.add(term)
 
@@ -118,7 +109,7 @@ def estimate_coefficient(self) -> tuple[pd.Series, list[pd.Series, pd.Series]]:
 
         :return: The unit average treatment effect and the 95% Wald confidence intervals.
         """
-        model = self._run_linear_regression()
+        model = self._run_regression()
         newline = "\n"
         patsy_md = ModelDesc.from_formula(self.treatment)
 
@@ -147,7 +138,7 @@ def estimate_ate(self) -> tuple[pd.Series, list[pd.Series, pd.Series]]:
 
         :return: The average treatment effect and the 95% Wald confidence intervals.
         """
-        model = self._run_linear_regression()
+        model = self._run_regression()
 
         # Create an empty individual for the control and treated
         individuals = pd.DataFrame(1, index=["control", "treated"], columns=model.params.index)
@@ -167,37 +158,6 @@ def estimate_ate(self) -> tuple[pd.Series, list[pd.Series, pd.Series]]:
         confidence_intervals = [pd.Series(interval) for interval in confidence_intervals]
         return ate, confidence_intervals
 
-    def estimate_control_treatment(self, adjustment_config: dict = None) -> tuple[pd.Series, pd.Series]:
-        """Estimate the outcomes under control and treatment.
-
-        :return: The estimated outcome under control and treatment in the form
-        (control_outcome, treatment_outcome).
-        """
-        if adjustment_config is None:
-            adjustment_config = {}
-        model = self._run_linear_regression()
-
-        x = pd.DataFrame(columns=self.df.columns)
-        x[self.treatment] = [self.treatment_value, self.control_value]
-        x["Intercept"] = 1  # self.intercept
-
-        print(x[self.treatment])
-        for k, v in adjustment_config.items():
-            x[k] = v
-        for k, v in self.effect_modifiers.items():
-            x[k] = v
-        x = dmatrix(self.formula.split("~")[1], x, return_type="dataframe")
-        for col in x:
-            if str(x.dtypes[col]) == "object":
-                x = pd.get_dummies(x, columns=[col], drop_first=True)
-        x = x[model.params.index]
-
-        x[self.treatment] = [self.treatment_value, self.control_value]
-
-        y = model.get_prediction(x).summary_frame()
-
-        return y.iloc[1], y.iloc[0]
-
     def estimate_risk_ratio(self, adjustment_config: dict = None) -> tuple[pd.Series, list[pd.Series, pd.Series]]:
         """Estimate the risk_ratio 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.
@@ -206,7 +166,8 @@ def estimate_risk_ratio(self, adjustment_config: dict = None) -> tuple[pd.Series
         """
         if adjustment_config is None:
             adjustment_config = {}
-        control_outcome, treatment_outcome = self.estimate_control_treatment(adjustment_config=adjustment_config)
+        prediction = self._predict(adjustment_config=adjustment_config)
+        control_outcome, treatment_outcome = prediction.iloc[1], prediction.iloc[0]
         ci_low = pd.Series(treatment_outcome["mean_ci_lower"] / control_outcome["mean_ci_upper"])
         ci_high = pd.Series(treatment_outcome["mean_ci_upper"] / control_outcome["mean_ci_lower"])
         return pd.Series(treatment_outcome["mean"] / control_outcome["mean"]), [ci_low, ci_high]
@@ -221,20 +182,12 @@ def estimate_ate_calculated(self, adjustment_config: dict = None) -> tuple[pd.Se
         """
         if adjustment_config is None:
             adjustment_config = {}
-        control_outcome, treatment_outcome = self.estimate_control_treatment(adjustment_config=adjustment_config)
+        prediction = self._predict(adjustment_config=adjustment_config)
+        control_outcome, treatment_outcome = prediction.iloc[1], prediction.iloc[0]
         ci_low = pd.Series(treatment_outcome["mean_ci_lower"] - control_outcome["mean_ci_upper"])
         ci_high = pd.Series(treatment_outcome["mean_ci_upper"] - control_outcome["mean_ci_lower"])
         return pd.Series(treatment_outcome["mean"] - control_outcome["mean"]), [ci_low, ci_high]
 
-    def _run_linear_regression(self) -> RegressionResultsWrapper:
-        """Run linear regression of the treatment and adjustment set against the outcome and return the model.
-
-        :return: The model after fitting to data.
-        """
-        model = smf.ols(formula=self.formula, data=self.df).fit()
-        self.model = model
-        return model
-
     def _get_confidence_intervals(self, model, treatment):
         confidence_intervals = model.conf_int(alpha=self.alpha, cols=None)
         ci_low, ci_high = (
diff --git a/causal_testing/estimation/logistic_regression_estimator.py b/causal_testing/estimation/logistic_regression_estimator.py
index 6d08fa94..3616b27d 100644
--- a/causal_testing/estimation/logistic_regression_estimator.py
+++ b/causal_testing/estimation/logistic_regression_estimator.py
@@ -1,59 +1,25 @@
 """This module contains the LogisticRegressionEstimator class for estimating categorical outcomes."""
 
 import logging
-from typing import Any
 from math import ceil
 
 import numpy as np
 import pandas as pd
 import statsmodels.formula.api as smf
-from patsy import dmatrix  # pylint: disable = no-name-in-module
-from statsmodels.regression.linear_model import RegressionResultsWrapper
 from statsmodels.tools.sm_exceptions import PerfectSeparationError
 
-from causal_testing.estimation.estimator import Estimator
+from causal_testing.estimation.regression_estimator import RegressionEstimator
 
 logger = logging.getLogger(__name__)
 
 
-class LogisticRegressionEstimator(Estimator):
+class LogisticRegressionEstimator(RegressionEstimator):
     """A Logistic Regression Estimator is a parametric estimator which restricts the variables in the data to a linear
     combination of parameters and functions of the variables (note these functions need not be linear). It is designed
     for estimating categorical outcomes.
     """
 
-    def __init__(
-        # pylint: disable=too-many-arguments
-        self,
-        treatment: str,
-        treatment_value: float,
-        control_value: float,
-        adjustment_set: set,
-        outcome: str,
-        df: pd.DataFrame = None,
-        effect_modifiers: dict[str:Any] = None,
-        formula: str = None,
-        query: str = "",
-    ):
-        super().__init__(
-            treatment=treatment,
-            treatment_value=treatment_value,
-            control_value=control_value,
-            adjustment_set=adjustment_set,
-            outcome=outcome,
-            df=df,
-            effect_modifiers=effect_modifiers,
-            query=query,
-        )
-
-        self.model = None
-        if effect_modifiers is None:
-            effect_modifiers = []
-        if formula is not None:
-            self.formula = formula
-        else:
-            terms = [treatment] + sorted(list(adjustment_set)) + sorted(list(self.effect_modifiers))
-            self.formula = f"{outcome} ~ {'+'.join(((terms)))}"
+    regressor = smf.logit
 
     def add_modelling_assumptions(self):
         """
@@ -68,43 +34,6 @@ def add_modelling_assumptions(self):
         self.modelling_assumptions.append("The outcome must be binary.")
         self.modelling_assumptions.append("Independently and identically distributed errors.")
 
-    def _run_logistic_regression(self, data) -> RegressionResultsWrapper:
-        """Run logistic regression of the treatment and adjustment set against the outcome and return the model.
-
-        :return: The model after fitting to data.
-        """
-        model = smf.logit(formula=self.formula, data=data).fit(disp=0)
-        self.model = model
-        return model
-
-    def estimate(self, data: pd.DataFrame, adjustment_config: dict = None) -> RegressionResultsWrapper:
-        """add terms to the dataframe and estimate the outcome from the data
-        :param data: A pandas dataframe containing execution data from the system-under-test.
-        :param adjustment_config: Dictionary containing the adjustment configuration of the adjustment set
-        """
-        if adjustment_config is None:
-            adjustment_config = {}
-        if set(self.adjustment_set) != set(adjustment_config):
-            raise ValueError(
-                f"Invalid adjustment configuration {adjustment_config}. Must specify values for {self.adjustment_set}"
-            )
-
-        model = self._run_logistic_regression(data)
-
-        x = pd.DataFrame(columns=self.df.columns)
-        x["Intercept"] = 1  # self.intercept
-        x[self.treatment] = [self.treatment_value, self.control_value]
-        for k, v in adjustment_config.items():
-            x[k] = v
-        for k, v in self.effect_modifiers.items():
-            x[k] = v
-        x = dmatrix(self.formula.split("~")[1], x, return_type="dataframe")
-        for col in x:
-            if str(x.dtypes[col]) == "object":
-                x = pd.get_dummies(x, columns=[col], drop_first=True)
-        # x = x[model.params.index]
-        return model.predict(x)
-
     def estimate_control_treatment(
         self, adjustment_config: dict = None, bootstrap_size: int = 100
     ) -> tuple[pd.Series, pd.Series]:
@@ -115,11 +44,13 @@ def estimate_control_treatment(
         """
         if adjustment_config is None:
             adjustment_config = {}
-        y = self.estimate(self.df, adjustment_config=adjustment_config)
+        y = self._predict(self.df, adjustment_config=adjustment_config)["predicted"]
 
         try:
             bootstrap_samples = [
-                self.estimate(self.df.sample(len(self.df), replace=True), adjustment_config=adjustment_config)
+                self._predict(self.df.sample(len(self.df), replace=True), adjustment_config=adjustment_config)[
+                    "predicted"
+                ]
                 for _ in range(bootstrap_size)
             ]
             control, treatment = zip(*[(x.iloc[1], x.iloc[0]) for x in bootstrap_samples])
@@ -214,5 +145,5 @@ def estimate_unit_odds_ratio(self) -> float:
 
         :return: The odds ratio. Confidence intervals are not yet supported.
         """
-        model = self._run_logistic_regression(self.df)
+        model = self._run_regression(self.df)
         return np.exp(model.params[self.treatment])
diff --git a/causal_testing/estimation/regression_estimator.py b/causal_testing/estimation/regression_estimator.py
index d713840a..9f0eed68 100644
--- a/causal_testing/estimation/regression_estimator.py
+++ b/causal_testing/estimation/regression_estimator.py
@@ -2,16 +2,13 @@
 
 import logging
 from typing import Any
-from abc import abstractmethod, abstractmethod
+from abc import abstractmethod
 
 import pandas as pd
-import statsmodels.formula.api as smf
-from patsy import dmatrix  # pylint: disable = no-name-in-module
-from patsy import ModelDesc
 from statsmodels.regression.linear_model import RegressionResultsWrapper
+from patsy import dmatrix  # pylint: disable = no-name-in-module
 
 from causal_testing.specification.variable import Variable
-from causal_testing.estimation.gp import GP
 from causal_testing.estimation.estimator import Estimator
 
 logger = logging.getLogger(__name__)
@@ -50,17 +47,23 @@ def __init__(
         self.model = None
         if effect_modifiers is None:
             effect_modifiers = []
+        if adjustment_set is None:
+            adjustment_set = []
         if formula is not None:
             self.formula = formula
         else:
             terms = [treatment] + sorted(list(adjustment_set)) + sorted(list(effect_modifiers))
             self.formula = f"{outcome} ~ {'+'.join(terms)}"
-        for term in self.effect_modifiers:
-            self.adjustment_set.add(term)
 
     @property
     @abstractmethod
-    def regression(self):
+    def regressor(self):
+        """
+        The regressor to use, e.g. ols or logit.
+        This should be a property accessible with self.regressor.
+        Define as `regressor = ...`` outside of __init__, not as `self.regressor = ...`, otherwise
+        you'll get an "cannot instantiate with abstract method" error.
+        """
         raise NotImplementedError("Subclasses must implement the 'model' property.")
 
     def add_modelling_assumptions(self):
@@ -81,6 +84,37 @@ def _run_regression(self, data=None) -> RegressionResultsWrapper:
         """
         if data is None:
             data = self.df
-        model = self.regression(formula=self.formula, data=data).fit(disp=0)
+        model = self.regressor(formula=self.formula, data=data).fit(disp=0)
         self.model = model
         return model
+
+    def _predict(self, data=None, adjustment_config: dict = None) -> tuple[pd.Series, pd.Series]:
+        """Estimate the outcomes under control and treatment.
+
+        :param data: The data to use, defaults to `self.df`. Controllable for boostrap sampling.
+        :param: adjustment_config: The values of the adjustment variables to use.
+
+        :return: The estimated outcome under control and treatment, with confidence intervals in the form of a
+                 dataframe with columns "predicted", "se", "ci_lower", and "ci_upper".
+        """
+        if adjustment_config is None:
+            adjustment_config = {}
+
+        model = self._run_regression(data)
+
+        x = pd.DataFrame(columns=self.df.columns)
+        x["Intercept"] = 1  # self.intercept
+        x[self.treatment] = [self.treatment_value, self.control_value]
+
+        for k, v in adjustment_config.items():
+            x[k] = v
+        for k, v in self.effect_modifiers.items():
+            x[k] = v
+        x = dmatrix(self.formula.split("~")[1], x, return_type="dataframe")
+        for col in x:
+            if str(x.dtypes[col]) == "object":
+                x = pd.get_dummies(x, columns=[col], drop_first=True)
+
+        # This has to be here in case the treatment variable is in an I(...) block in the self.formula
+        x[self.treatment] = [self.treatment_value, self.control_value]
+        return model.get_prediction(x).summary_frame()
diff --git a/tests/estimation_tests/test_cubic_spline_estimator.py b/tests/estimation_tests/test_cubic_spline_estimator.py
index f8ec8117..e8560b88 100644
--- a/tests/estimation_tests/test_cubic_spline_estimator.py
+++ b/tests/estimation_tests/test_cubic_spline_estimator.py
@@ -27,7 +27,7 @@ def test_program_11_3_cublic_spline(self):
 
         cublic_spline_estimator = CubicSplineRegressionEstimator("treatments", 1, 0, set(), "outcomes", 3, df)
 
-        model = cublic_spline_estimator._run_linear_regression()
+        ate_1 = cublic_spline_estimator.estimate_ate_calculated()
 
         self.assertEqual(
             round(
@@ -37,7 +37,6 @@ def test_program_11_3_cublic_spline(self):
             195.6,
         )
 
-        ate_1 = cublic_spline_estimator.estimate_ate_calculated()
         cublic_spline_estimator.treatment_value = 2
         ate_2 = cublic_spline_estimator.estimate_ate_calculated()
 

From 03f1bd8fad12c253f6d5fb10103cc5b67911b828 Mon Sep 17 00:00:00 2001
From: Michael <jmafoster1@gmail.com>
Date: Thu, 8 Aug 2024 13:01:13 +0100
Subject: [PATCH 10/24] Pylint

---
 causal_testing/estimation/iv_estimator.py      |  6 ++----
 .../logistic_regression_estimator.py           | 18 ++++++++++++++++++
 .../estimation/regression_estimator.py         |  2 +-
 3 files changed, 21 insertions(+), 5 deletions(-)

diff --git a/causal_testing/estimation/iv_estimator.py b/causal_testing/estimation/iv_estimator.py
index 5a30dfd6..b8b56170 100644
--- a/causal_testing/estimation/iv_estimator.py
+++ b/causal_testing/estimation/iv_estimator.py
@@ -19,6 +19,7 @@ class InstrumentalVariableEstimator(Estimator):
 
     def __init__(
         # pylint: disable=too-many-arguments
+        # pylint: disable=duplicate-code
         self,
         treatment: str,
         treatment_value: float,
@@ -27,8 +28,6 @@ def __init__(
         outcome: str,
         instrument: str,
         df: pd.DataFrame = None,
-        intercept: int = 1,
-        effect_modifiers: dict = None,  # Not used (yet?). Needed for compatibility
         alpha: float = 0.05,
         query: str = "",
     ):
@@ -43,8 +42,7 @@ def __init__(
             alpha=alpha,
             query=query,
         )
-        self.intercept = intercept
-        self.model = None
+
         self.instrument = instrument
 
     def add_modelling_assumptions(self):
diff --git a/causal_testing/estimation/logistic_regression_estimator.py b/causal_testing/estimation/logistic_regression_estimator.py
index 3616b27d..45ab8233 100644
--- a/causal_testing/estimation/logistic_regression_estimator.py
+++ b/causal_testing/estimation/logistic_regression_estimator.py
@@ -34,6 +34,24 @@ def add_modelling_assumptions(self):
         self.modelling_assumptions.append("The outcome must be binary.")
         self.modelling_assumptions.append("Independently and identically distributed errors.")
 
+    def _predict(self, data=None, adjustment_config: dict = None) -> pd.DataFrame:
+        """Estimate the outcomes under control and treatment.
+
+        :param data: The data to use, defaults to `self.df`. Controllable for boostrap sampling.
+        :param: adjustment_config: The values of the adjustment variables to use.
+
+        :return: The estimated outcome under control and treatment, with confidence intervals in the form of a
+                 dataframe with columns "predicted", "se", "ci_lower", and "ci_upper".
+        """
+        if adjustment_config is None:
+            adjustment_config = {}
+        if set(self.adjustment_set) != set(adjustment_config):
+            raise ValueError(
+                f"Invalid adjustment configuration {adjustment_config}. Must specify values for {self.adjustment_set}"
+            )
+
+        return super()._predict(data, adjustment_config)
+
     def estimate_control_treatment(
         self, adjustment_config: dict = None, bootstrap_size: int = 100
     ) -> tuple[pd.Series, pd.Series]:
diff --git a/causal_testing/estimation/regression_estimator.py b/causal_testing/estimation/regression_estimator.py
index 9f0eed68..613b427b 100644
--- a/causal_testing/estimation/regression_estimator.py
+++ b/causal_testing/estimation/regression_estimator.py
@@ -88,7 +88,7 @@ def _run_regression(self, data=None) -> RegressionResultsWrapper:
         self.model = model
         return model
 
-    def _predict(self, data=None, adjustment_config: dict = None) -> tuple[pd.Series, pd.Series]:
+    def _predict(self, data=None, adjustment_config: dict = None) -> pd.DataFrame:
         """Estimate the outcomes under control and treatment.
 
         :param data: The data to use, defaults to `self.df`. Controllable for boostrap sampling.

From 968fc89dfa3454bf5b0ff6baf6a1c30b27dd5c76 Mon Sep 17 00:00:00 2001
From: Michael <jmafoster1@gmail.com>
Date: Thu, 8 Aug 2024 14:27:54 +0100
Subject: [PATCH 11/24] Seeding gp power

---
 causal_testing/estimation/gp.py                            | 3 ++-
 tests/estimation_tests/test_linear_regression_estimator.py | 4 ++--
 2 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/causal_testing/estimation/gp.py b/causal_testing/estimation/gp.py
index ab53d026..f5777cd7 100644
--- a/causal_testing/estimation/gp.py
+++ b/causal_testing/estimation/gp.py
@@ -134,7 +134,7 @@ def __init__(
             "reciprocal": lambda x1: f"Pow({x1},-1)",
         } | sympy_conversions
 
-        for i in range(self.max_order):
+        for i in range(self.max_order + 1):
             name = f"power_{i}"
             func, conversion = create_power_function(i)
             self.pset.addPrimitive(func, 1, name=name)
@@ -145,6 +145,7 @@ def __init__(
                 )
             self.sympy_conversions[name] = conversion
 
+        print(self.pset.mapping)
         creator.create("FitnessMin", base.Fitness, weights=(-1.0,))
         creator.create("Individual", gp.PrimitiveTree, fitness=creator.FitnessMin)
 
diff --git a/tests/estimation_tests/test_linear_regression_estimator.py b/tests/estimation_tests/test_linear_regression_estimator.py
index 6a6dec10..8778c086 100644
--- a/tests/estimation_tests/test_linear_regression_estimator.py
+++ b/tests/estimation_tests/test_linear_regression_estimator.py
@@ -303,10 +303,10 @@ def test_gp_power(self):
         df["X"] = np.arange(10)
         df["Y"] = 2 * (df["X"] ** 2)
         linear_regression_estimator = LinearRegressionEstimator("X", 0, 1, set(), "Y", df.astype(float))
-        linear_regression_estimator.gp_formula(seed=1, max_order=0)
+        linear_regression_estimator.gp_formula(seed=1, max_order=2, seeds=["mul(2, power_2(X))"])
         self.assertEqual(
             linear_regression_estimator.formula,
-            "Y ~ I(1.9999999999999999*X**2 - 1.0043240235058056e-116*X + 2.6645352591003757e-15) - 1",
+            "Y ~ I(2*X**2) - 1",
         )
         ate, (ci_low, ci_high) = linear_regression_estimator.estimate_ate_calculated()
         self.assertEqual(round(ate[0], 2), -2.00)

From c0fbbddf146015dfc8ed9517c19d0c540650d417 Mon Sep 17 00:00:00 2001
From: Michael <jmafoster1@gmail.com>
Date: Thu, 8 Aug 2024 15:05:32 +0100
Subject: [PATCH 12/24] Removed ate and risk ratio from logistic regression
 estimator, since they're not statistically meaningful

---
 .../logistic_regression_estimator.py          | 123 ------------------
 .../test_logistic_regression_estimator.py     |  49 -------
 2 files changed, 172 deletions(-)

diff --git a/causal_testing/estimation/logistic_regression_estimator.py b/causal_testing/estimation/logistic_regression_estimator.py
index 45ab8233..bd8e3cc0 100644
--- a/causal_testing/estimation/logistic_regression_estimator.py
+++ b/causal_testing/estimation/logistic_regression_estimator.py
@@ -34,129 +34,6 @@ def add_modelling_assumptions(self):
         self.modelling_assumptions.append("The outcome must be binary.")
         self.modelling_assumptions.append("Independently and identically distributed errors.")
 
-    def _predict(self, data=None, adjustment_config: dict = None) -> pd.DataFrame:
-        """Estimate the outcomes under control and treatment.
-
-        :param data: The data to use, defaults to `self.df`. Controllable for boostrap sampling.
-        :param: adjustment_config: The values of the adjustment variables to use.
-
-        :return: The estimated outcome under control and treatment, with confidence intervals in the form of a
-                 dataframe with columns "predicted", "se", "ci_lower", and "ci_upper".
-        """
-        if adjustment_config is None:
-            adjustment_config = {}
-        if set(self.adjustment_set) != set(adjustment_config):
-            raise ValueError(
-                f"Invalid adjustment configuration {adjustment_config}. Must specify values for {self.adjustment_set}"
-            )
-
-        return super()._predict(data, adjustment_config)
-
-    def estimate_control_treatment(
-        self, adjustment_config: dict = None, bootstrap_size: int = 100
-    ) -> tuple[pd.Series, pd.Series]:
-        """Estimate the outcomes under control and treatment.
-
-        :return: The estimated control and treatment values and their confidence
-        intervals in the form ((ci_low, control, ci_high), (ci_low, treatment, ci_high)).
-        """
-        if adjustment_config is None:
-            adjustment_config = {}
-        y = self._predict(self.df, adjustment_config=adjustment_config)["predicted"]
-
-        try:
-            bootstrap_samples = [
-                self._predict(self.df.sample(len(self.df), replace=True), adjustment_config=adjustment_config)[
-                    "predicted"
-                ]
-                for _ in range(bootstrap_size)
-            ]
-            control, treatment = zip(*[(x.iloc[1], x.iloc[0]) for x in bootstrap_samples])
-        except PerfectSeparationError:
-            logger.warning(
-                "Perfect separation detected, results not available. Cannot calculate confidence intervals for such "
-                "a small dataset."
-            )
-            return (y.iloc[1], None), (y.iloc[0], None)
-        except np.linalg.LinAlgError:
-            logger.warning("Singular matrix detected. Confidence intervals not available. Try with a larger data set")
-            return (y.iloc[1], None), (y.iloc[0], None)
-
-        # Delta method confidence intervals from
-        # https://stackoverflow.com/questions/47414842/confidence-interval-of-probability-prediction-from-logistic-regression-statsmode
-        # cov = model.cov_params()
-        # gradient = (y * (1 - y) * x.T).T  # matrix of gradients for each observation
-        # std_errors = np.array([np.sqrt(np.dot(np.dot(g, cov), g)) for g in gradient.to_numpy()])
-        # c = 1.96  # multiplier for confidence interval
-        # upper = np.maximum(0, np.minimum(1, y + std_errors * c))
-        # lower = np.maximum(0, np.minimum(1, y - std_errors * c))
-
-        return (y.iloc[1], np.array(control)), (y.iloc[0], np.array(treatment))
-
-    def estimate_ate(self, adjustment_config: dict = None, bootstrap_size: int = 100) -> float:
-        """Estimate the ate 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. Here, we actually
-        calculate the expected outcomes under control and treatment and take one away from the other. This
-        allows for custom terms to be put in such as squares, inverses, products, etc.
-
-        :return: The estimated average treatment effect and 95% confidence intervals
-        """
-        if adjustment_config is None:
-            adjustment_config = {}
-        (control_outcome, control_bootstraps), (
-            treatment_outcome,
-            treatment_bootstraps,
-        ) = self.estimate_control_treatment(bootstrap_size=bootstrap_size, adjustment_config=adjustment_config)
-        estimate = treatment_outcome - control_outcome
-
-        if control_bootstraps is None or treatment_bootstraps is None:
-            return estimate, (None, None)
-
-        bootstraps = sorted(list(treatment_bootstraps - control_bootstraps))
-        bound = int((bootstrap_size * self.alpha) / 2)
-        ci_low = bootstraps[bound]
-        ci_high = bootstraps[bootstrap_size - bound]
-
-        logger.info(
-            f"Changing {self.treatment} from {self.control_value} to {self.treatment_value} gives an estimated "
-            f"ATE of {ci_low} < {estimate} < {ci_high}"
-        )
-        assert ci_low < estimate < ci_high, f"Expecting {ci_low} < {estimate} < {ci_high}"
-
-        return estimate, (ci_low, ci_high)
-
-    def estimate_risk_ratio(self, adjustment_config: dict = None, bootstrap_size: int = 100) -> float:
-        """Estimate the ate 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. Here, we actually
-        calculate the expected outcomes under control and treatment and divide one by the other. This
-        allows for custom terms to be put in such as squares, inverses, products, etc.
-
-        :return: The estimated risk ratio and 95% confidence intervals.
-        """
-        if adjustment_config is None:
-            adjustment_config = {}
-        (control_outcome, control_bootstraps), (
-            treatment_outcome,
-            treatment_bootstraps,
-        ) = self.estimate_control_treatment(bootstrap_size=bootstrap_size, adjustment_config=adjustment_config)
-        estimate = treatment_outcome / control_outcome
-
-        if control_bootstraps is None or treatment_bootstraps is None:
-            return estimate, (None, None)
-
-        bootstraps = sorted(list(treatment_bootstraps / control_bootstraps))
-        bound = ceil((bootstrap_size * self.alpha) / 2)
-        ci_low = bootstraps[bound]
-        ci_high = bootstraps[bootstrap_size - bound]
-
-        logger.info(
-            f"Changing {self.treatment} from {self.control_value} to {self.treatment_value} gives an estimated "
-            f"risk ratio of {ci_low} < {estimate} < {ci_high}"
-        )
-        assert ci_low < estimate < ci_high, f"Expecting {ci_low} < {estimate} < {ci_high}"
-
-        return estimate, (ci_low, ci_high)
-
     def estimate_unit_odds_ratio(self) -> float:
         """Estimate the odds ratio of increasing the treatment by one. In logistic regression, this corresponds to the
         coefficient of the treatment of interest.
diff --git a/tests/estimation_tests/test_logistic_regression_estimator.py b/tests/estimation_tests/test_logistic_regression_estimator.py
index 5dc4df24..a5d104e3 100644
--- a/tests/estimation_tests/test_logistic_regression_estimator.py
+++ b/tests/estimation_tests/test_logistic_regression_estimator.py
@@ -17,57 +17,8 @@ class TestLogisticRegressionEstimator(unittest.TestCase):
     def setUpClass(cls) -> None:
         cls.scarf_df = pd.read_csv("tests/resources/data/scarf_data.csv")
 
-    def test_ate(self):
-        df = self.scarf_df.copy()
-        logistic_regression_estimator = LogisticRegressionEstimator("length_in", 65, 55, set(), "completed", df)
-        ate, _ = logistic_regression_estimator.estimate_ate()
-        self.assertEqual(round(ate, 4), -0.1987)
-
-    def test_risk_ratio(self):
-        df = self.scarf_df.copy()
-        logistic_regression_estimator = LogisticRegressionEstimator("length_in", 65, 55, set(), "completed", df)
-        rr, _ = logistic_regression_estimator.estimate_risk_ratio()
-        self.assertEqual(round(rr, 4), 0.7664)
-
     def test_odds_ratio(self):
         df = self.scarf_df.copy()
         logistic_regression_estimator = LogisticRegressionEstimator("length_in", 65, 55, set(), "completed", df)
         odds = logistic_regression_estimator.estimate_unit_odds_ratio()
         self.assertEqual(round(odds, 4), 0.8948)
-
-    def test_ate_adjustment(self):
-        df = self.scarf_df.copy()
-        logistic_regression_estimator = LogisticRegressionEstimator(
-            "length_in", 65, 55, {"large_gauge"}, "completed", df
-        )
-        ate, _ = logistic_regression_estimator.estimate_ate(adjustment_config={"large_gauge": 0})
-        self.assertEqual(round(ate, 4), -0.3388)
-
-    def test_ate_invalid_adjustment(self):
-        df = self.scarf_df.copy()
-        logistic_regression_estimator = LogisticRegressionEstimator("length_in", 65, 55, {}, "completed", df)
-        with self.assertRaises(ValueError):
-            ate, _ = logistic_regression_estimator.estimate_ate(adjustment_config={"large_gauge": 0})
-
-    def test_ate_effect_modifiers(self):
-        df = self.scarf_df.copy()
-        logistic_regression_estimator = LogisticRegressionEstimator(
-            "length_in", 65, 55, set(), "completed", df, effect_modifiers={"large_gauge": 0}
-        )
-        ate, _ = logistic_regression_estimator.estimate_ate()
-        self.assertEqual(round(ate, 4), -0.3388)
-
-    def test_ate_effect_modifiers_formula(self):
-        df = self.scarf_df.copy()
-        logistic_regression_estimator = LogisticRegressionEstimator(
-            "length_in",
-            65,
-            55,
-            set(),
-            "completed",
-            df,
-            effect_modifiers={"large_gauge": 0},
-            formula="completed ~ length_in + large_gauge",
-        )
-        ate, _ = logistic_regression_estimator.estimate_ate()
-        self.assertEqual(round(ate, 4), -0.3388)

From b1575835f8ad6ac784b7ebc7700c030aa8932349 Mon Sep 17 00:00:00 2001
From: Michael <jmafoster1@gmail.com>
Date: Thu, 8 Aug 2024 15:10:01 +0100
Subject: [PATCH 13/24] removed unnecessary raise NotImplementedError

---
 causal_testing/estimation/regression_estimator.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/causal_testing/estimation/regression_estimator.py b/causal_testing/estimation/regression_estimator.py
index 613b427b..44ef67da 100644
--- a/causal_testing/estimation/regression_estimator.py
+++ b/causal_testing/estimation/regression_estimator.py
@@ -64,7 +64,6 @@ def regressor(self):
         Define as `regressor = ...`` outside of __init__, not as `self.regressor = ...`, otherwise
         you'll get an "cannot instantiate with abstract method" error.
         """
-        raise NotImplementedError("Subclasses must implement the 'model' property.")
 
     def add_modelling_assumptions(self):
         """

From 5886480e5161c12fe96b3c00458c24e7f048023c Mon Sep 17 00:00:00 2001
From: Michael <jmafoster1@gmail.com>
Date: Thu, 8 Aug 2024 15:12:19 +0100
Subject: [PATCH 14/24] pylint

---
 causal_testing/estimation/logistic_regression_estimator.py | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/causal_testing/estimation/logistic_regression_estimator.py b/causal_testing/estimation/logistic_regression_estimator.py
index bd8e3cc0..0546ba6a 100644
--- a/causal_testing/estimation/logistic_regression_estimator.py
+++ b/causal_testing/estimation/logistic_regression_estimator.py
@@ -1,12 +1,9 @@
 """This module contains the LogisticRegressionEstimator class for estimating categorical outcomes."""
 
 import logging
-from math import ceil
 
 import numpy as np
-import pandas as pd
 import statsmodels.formula.api as smf
-from statsmodels.tools.sm_exceptions import PerfectSeparationError
 
 from causal_testing.estimation.regression_estimator import RegressionEstimator
 

From c66bb1531745e0f2b2697797c71525f70c711eba Mon Sep 17 00:00:00 2001
From: Michael Foster <m.foster@sheffield.ac.uk>
Date: Wed, 4 Sep 2024 14:42:10 +0100
Subject: [PATCH 15/24] pyproject

---
 pyproject.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pyproject.toml b/pyproject.toml
index 8e45dd21..083a4e3b 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -30,7 +30,7 @@ dependencies = [
     "pygad~=3.3",
     "deap~=1.4.1",
     "sympy~=1.13.1",
-    "deap"~=1.4.1"
+    "deap"~=1.4.1",
 ]
 dynamic = ["version"]
 

From 232cde2083ca3041aecc72c1886e338996439d93 Mon Sep 17 00:00:00 2001
From: Michael Foster <m.foster@sheffield.ac.uk>
Date: Wed, 4 Sep 2024 14:43:35 +0100
Subject: [PATCH 16/24] pyproject

---
 pyproject.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pyproject.toml b/pyproject.toml
index 083a4e3b..738c2d06 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -30,7 +30,7 @@ dependencies = [
     "pygad~=3.3",
     "deap~=1.4.1",
     "sympy~=1.13.1",
-    "deap"~=1.4.1",
+    "deap~=1.4.1",
 ]
 dynamic = ["version"]
 

From 1d6b712ce72896dec968e0918f8b238f523b83b4 Mon Sep 17 00:00:00 2001
From: Michael Foster <m.foster@sheffield.ac.uk>
Date: Wed, 4 Sep 2024 14:48:50 +0100
Subject: [PATCH 17/24] fixed pylintrc for deprecated exceptions

---
 .pylintrc | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/.pylintrc b/.pylintrc
index f1e2a61c..4ab81c8c 100644
--- a/.pylintrc
+++ b/.pylintrc
@@ -371,8 +371,8 @@ min-public-methods=2
 [EXCEPTIONS]
 
 # Exceptions that will emit a warning when caught.
-overgeneral-exceptions=BaseException,
-                       Exception
+overgeneral-exceptions=builtins.BaseException,
+                       builtins.Exception
 
 
 [FORMAT]

From 2a7ce30ebbd12ef180c62860b170aad35a025912 Mon Sep 17 00:00:00 2001
From: Michael Foster <m.foster@sheffield.ac.uk>
Date: Wed, 4 Sep 2024 14:48:55 +0100
Subject: [PATCH 18/24] black

---
 causal_testing/surrogate/causal_surrogate_assisted.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/causal_testing/surrogate/causal_surrogate_assisted.py b/causal_testing/surrogate/causal_surrogate_assisted.py
index 11380d4d..4fba5371 100644
--- a/causal_testing/surrogate/causal_surrogate_assisted.py
+++ b/causal_testing/surrogate/causal_surrogate_assisted.py
@@ -10,7 +10,6 @@
 from causal_testing.estimation.cubic_spline_estimator import CubicSplineRegressionEstimator
 
 
-
 @dataclass
 class SimulationResult:
     """Data class holding the data and result metadata of a simulation"""

From 673165cfbb3296e04dd40d65942cbdd4675c9cf5 Mon Sep 17 00:00:00 2001
From: Michael Foster <m.foster@sheffield.ac.uk>
Date: Wed, 4 Sep 2024 15:00:03 +0100
Subject: [PATCH 19/24] pylint

---
 causal_testing/estimation/cubic_spline_estimator.py   | 11 +++++++++++
 .../estimation/linear_regression_estimator.py         |  4 ++++
 2 files changed, 15 insertions(+)

diff --git a/causal_testing/estimation/cubic_spline_estimator.py b/causal_testing/estimation/cubic_spline_estimator.py
index 3467156e..b8ceb2fd 100644
--- a/causal_testing/estimation/cubic_spline_estimator.py
+++ b/causal_testing/estimation/cubic_spline_estimator.py
@@ -46,6 +46,17 @@ def __init__(
             self.formula = f"{outcome} ~ cr({'+'.join(terms)}, df={basis})"
 
     def estimate_ate_calculated(self, adjustment_config: dict = None) -> pd.Series:
+        """Estimate the ate 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. Here, we actually
+        calculate the expected outcomes under control and treatment and divide one by the other. This
+        allows for custom terms to be put in such as squares, inverses, products, etc.
+
+        :param: adjustment_config: The configuration of the adjustment set as a dict mapping variable names to
+                                   their values. N.B. Every variable in the adjustment set MUST have a value in
+                                   order to estimate the outcome under control and treatment.
+
+        :return: The average treatment effect.
+        """
         model = self._run_regression()
 
         x = {"Intercept": 1, self.treatment: self.treatment_value}
diff --git a/causal_testing/estimation/linear_regression_estimator.py b/causal_testing/estimation/linear_regression_estimator.py
index 1ab00f45..bf6ce6d4 100644
--- a/causal_testing/estimation/linear_regression_estimator.py
+++ b/causal_testing/estimation/linear_regression_estimator.py
@@ -178,6 +178,10 @@ def estimate_ate_calculated(self, adjustment_config: dict = None) -> tuple[pd.Se
         calculate the expected outcomes under control and treatment and divide one by the other. This
         allows for custom terms to be put in such as squares, inverses, products, etc.
 
+        :param: adjustment_config: The configuration of the adjustment set as a dict mapping variable names to
+                                   their values. N.B. Every variable in the adjustment set MUST have a value in
+                                   order to estimate the outcome under control and treatment.
+
         :return: The average treatment effect and the 95% Wald confidence intervals.
         """
         if adjustment_config is None:

From 6efe5946e455c89177364edf9c017783137a92c7 Mon Sep 17 00:00:00 2001
From: Michael <jmafoster1@gmail.com>
Date: Wed, 11 Sep 2024 09:27:57 +0100
Subject: [PATCH 20/24] increased coverage a little

---
 causal_testing/estimation/estimator.py                |  4 +---
 causal_testing/estimation/gp.py                       | 10 ++++++----
 .../estimation/linear_regression_estimator.py         |  4 ----
 tests/estimation_tests/test_gp.py                     | 11 +++++++++++
 tests/specification_tests/test_capabilities.py        |  2 --
 tests/specification_tests/test_causal_dag.py          |  9 +++++----
 tests/specification_tests/test_variable.py            |  2 --
 tests/testing_tests/test_causal_test_case.py          |  3 +--
 8 files changed, 24 insertions(+), 21 deletions(-)
 create mode 100644 tests/estimation_tests/test_gp.py

diff --git a/causal_testing/estimation/estimator.py b/causal_testing/estimation/estimator.py
index ab73723f..47ab1efe 100644
--- a/causal_testing/estimation/estimator.py
+++ b/causal_testing/estimation/estimator.py
@@ -50,10 +50,8 @@ def __init__(
 
         if effect_modifiers is None:
             self.effect_modifiers = {}
-        elif isinstance(effect_modifiers, dict):
-            self.effect_modifiers = effect_modifiers
         else:
-            raise ValueError(f"Unsupported type for effect_modifiers {effect_modifiers}. Expected iterable")
+            self.effect_modifiers = effect_modifiers
         self.modelling_assumptions = []
         if query:
             self.modelling_assumptions.append(query)
diff --git a/causal_testing/estimation/gp.py b/causal_testing/estimation/gp.py
index f5777cd7..49aa4dfc 100644
--- a/causal_testing/estimation/gp.py
+++ b/causal_testing/estimation/gp.py
@@ -31,8 +31,12 @@ def reciprocal(x: float) -> float:
 
 def mut_insert(expression: gp.PrimitiveTree, pset: gp.PrimitiveSet):
     """
-    Copied from gp.mutInsert, except that we import isclass from inspect, so we
-    won't have the "isclass not defined" bug.
+    NOTE: This is a temporary workaround. This method is copied verbatim from
+    gp.mutInsert. It seems they forgot to import isclass from inspect, so their
+    method throws an error, saying that "isclass is not defined". A couple of
+    lines are not covered by tests, but since this is 1. a temporary workaround
+    until they release a new version of DEAP, and 2. not our code, I don't think
+    that matters.
 
     Inserts a new branch at a random position in *expression*. The subtree
     at the chosen position is used as child node of the created subtree, in
@@ -374,6 +378,4 @@ def mutate(self, expression: gp.PrimitiveTree) -> gp.PrimitiveTree:
             mutated = mut_insert(self.toolbox.clone(expression), self.pset)
         elif choice == 3:
             mutated = gp.mutShrink(self.toolbox.clone(expression))
-        else:
-            raise ValueError("Invalid mutation choice")
         return mutated
diff --git a/causal_testing/estimation/linear_regression_estimator.py b/causal_testing/estimation/linear_regression_estimator.py
index bf6ce6d4..bfefd445 100644
--- a/causal_testing/estimation/linear_regression_estimator.py
+++ b/causal_testing/estimation/linear_regression_estimator.py
@@ -164,8 +164,6 @@ def estimate_risk_ratio(self, adjustment_config: dict = None) -> tuple[pd.Series
 
         :return: The average treatment effect and the 95% Wald confidence intervals.
         """
-        if adjustment_config is None:
-            adjustment_config = {}
         prediction = self._predict(adjustment_config=adjustment_config)
         control_outcome, treatment_outcome = prediction.iloc[1], prediction.iloc[0]
         ci_low = pd.Series(treatment_outcome["mean_ci_lower"] / control_outcome["mean_ci_upper"])
@@ -184,8 +182,6 @@ def estimate_ate_calculated(self, adjustment_config: dict = None) -> tuple[pd.Se
 
         :return: The average treatment effect and the 95% Wald confidence intervals.
         """
-        if adjustment_config is None:
-            adjustment_config = {}
         prediction = self._predict(adjustment_config=adjustment_config)
         control_outcome, treatment_outcome = prediction.iloc[1], prediction.iloc[0]
         ci_low = pd.Series(treatment_outcome["mean_ci_lower"] - control_outcome["mean_ci_upper"])
diff --git a/tests/estimation_tests/test_gp.py b/tests/estimation_tests/test_gp.py
new file mode 100644
index 00000000..4946a871
--- /dev/null
+++ b/tests/estimation_tests/test_gp.py
@@ -0,0 +1,11 @@
+import unittest
+import pandas as pd
+
+from causal_testing.estimation.gp import GP
+
+
+class TestGP(unittest.TestCase):
+
+    def test_init_invalid_fun_name(self):
+        with self.assertRaises(ValueError):
+            GP(df=pd.DataFrame(), features=[], outcome="", max_order=2, sympy_conversions={"power_1": ""})
diff --git a/tests/specification_tests/test_capabilities.py b/tests/specification_tests/test_capabilities.py
index 505c7eff..98fd466c 100644
--- a/tests/specification_tests/test_capabilities.py
+++ b/tests/specification_tests/test_capabilities.py
@@ -3,7 +3,6 @@
 
 
 class TestCapability(unittest.TestCase):
-
     """
     Test the Capability class for basic methods.
     """
@@ -17,7 +16,6 @@ def test_repr(self):
 
 
 class TestTreatmentSequence(unittest.TestCase):
-
     """
     Test the TreatmentSequence class for basic methods.
     """
diff --git a/tests/specification_tests/test_causal_dag.py b/tests/specification_tests/test_causal_dag.py
index c020ae67..bd01d11c 100644
--- a/tests/specification_tests/test_causal_dag.py
+++ b/tests/specification_tests/test_causal_dag.py
@@ -8,8 +8,6 @@
 from causal_testing.testing.base_test_case import BaseTestCase
 
 
-
-
 class TestCausalDAGIssue90(unittest.TestCase):
     """
     Test the CausalDAG class for the resolution of Issue 90.
@@ -63,10 +61,11 @@ def test_common_cause(self):
         causal_dag.graph.add_edge("U", "I")
         with self.assertRaises(ValueError):
             causal_dag.check_iv_assumptions("X", "Y", "I")
-    
+
     def tearDown(self) -> None:
         shutil.rmtree(self.temp_dir_path)
 
+
 class TestCausalDAG(unittest.TestCase):
     """
     Test the CausalDAG class for creation of Causal Directed Acyclic Graphs (DAGs).
@@ -154,10 +153,11 @@ def test_direct_effect_adjustment_sets_no_adjustment(self):
         causal_dag = CausalDAG(self.dag_dot_path)
         adjustment_sets = causal_dag.direct_effect_adjustment_sets(["X2"], ["D1"])
         self.assertEqual(list(adjustment_sets), [set()])
-    
+
     def tearDown(self) -> None:
         shutil.rmtree(self.temp_dir_path)
 
+
 class TestDAGIdentification(unittest.TestCase):
     """
     Test the Causal DAG identification algorithms and supporting algorithms.
@@ -345,6 +345,7 @@ def test_dag_with_non_character_nodes(self):
     def tearDown(self) -> None:
         shutil.rmtree(self.temp_dir_path)
 
+
 class TestDependsOnOutputs(unittest.TestCase):
     """
     Test the depends_on_outputs method.
diff --git a/tests/specification_tests/test_variable.py b/tests/specification_tests/test_variable.py
index cc76df6a..b35724a2 100644
--- a/tests/specification_tests/test_variable.py
+++ b/tests/specification_tests/test_variable.py
@@ -7,7 +7,6 @@
 
 
 class TestVariable(unittest.TestCase):
-
     """
     Test the Variable class for basic methods.
     """
@@ -143,7 +142,6 @@ def test_copy(self):
 
 
 class TestZ3Methods(unittest.TestCase):
-
     """
     Test the Variable class for Z3 methods.
 
diff --git a/tests/testing_tests/test_causal_test_case.py b/tests/testing_tests/test_causal_test_case.py
index 01b55cc1..600191d3 100644
--- a/tests/testing_tests/test_causal_test_case.py
+++ b/tests/testing_tests/test_causal_test_case.py
@@ -107,7 +107,6 @@ def test_check_minimum_adjustment_set(self):
         minimal_adjustment_set = self.causal_dag.identification(self.base_test_case)
         self.assertEqual(minimal_adjustment_set, {"D"})
 
-
     def test_invalid_causal_effect(self):
         """Check that executing the causal test case returns the correct results for dummy data using a linear
         regression estimator."""
@@ -170,7 +169,7 @@ def test_execute_test_observational_linear_regression_estimator_coefficient(self
         )
         self.causal_test_case.estimate_type = "coefficient"
         causal_test_result = self.causal_test_case.execute_test(estimation_model, self.data_collector)
-        pd.testing.assert_series_equal(causal_test_result.test_value.value, pd.Series({'D': 0.0}), atol=1e-1)
+        pd.testing.assert_series_equal(causal_test_result.test_value.value, pd.Series({"D": 0.0}), atol=1e-1)
 
     def test_execute_test_observational_linear_regression_estimator_risk_ratio(self):
         """Check that executing the causal test case returns the correct results for dummy data using a linear

From 47fc503153fbad91ce0ba6144eae4bcbbc3e2bdb Mon Sep 17 00:00:00 2001
From: Michael <jmafoster1@gmail.com>
Date: Wed, 11 Sep 2024 09:34:55 +0100
Subject: [PATCH 21/24] actions/upload-artifact updated to v3

---
 .github/workflows/lint-format.yaml | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/lint-format.yaml b/.github/workflows/lint-format.yaml
index d6b49c37..3c2ea9b7 100644
--- a/.github/workflows/lint-format.yaml
+++ b/.github/workflows/lint-format.yaml
@@ -25,9 +25,9 @@ jobs:
 
       - name: Archive production artifacts
         if: ${{ success() }} || ${{ failure() }}
-        uses: actions/upload-artifact@v2
+        uses: actions/upload-artifact@v3
         with:
           name: MegaLinter reports
           path: |
             megalinter-reports
-            mega-linter.log
\ No newline at end of file
+            mega-linter.log

From 9e2b4baa39828f8f5470614f681ab938228393e9 Mon Sep 17 00:00:00 2001
From: Michael <jmafoster1@gmail.com>
Date: Wed, 11 Sep 2024 09:42:07 +0100
Subject: [PATCH 22/24] deleted estimators.py

---
 causal_testing/testing/estimators.py | 840 ---------------------------
 1 file changed, 840 deletions(-)
 delete mode 100644 causal_testing/testing/estimators.py

diff --git a/causal_testing/testing/estimators.py b/causal_testing/testing/estimators.py
deleted file mode 100644
index ffe76387..00000000
--- a/causal_testing/testing/estimators.py
+++ /dev/null
@@ -1,840 +0,0 @@
-"""This module contains the Estimator abstract class, as well as its concrete extensions: LogisticRegressionEstimator,
-LinearRegressionEstimator"""
-
-import logging
-from abc import ABC, abstractmethod
-from typing import Any
-from math import ceil
-
-import numpy as np
-import pandas as pd
-import statsmodels.api as sm
-import statsmodels.formula.api as smf
-from patsy import dmatrix  # pylint: disable = no-name-in-module
-from patsy import ModelDesc
-from statsmodels.regression.linear_model import RegressionResultsWrapper
-from statsmodels.tools.sm_exceptions import PerfectSeparationError
-from lifelines import CoxPHFitter
-
-from causal_testing.specification.variable import Variable
-from causal_testing.specification.capabilities import TreatmentSequence, Capability
-
-logger = logging.getLogger(__name__)
-
-
-class Estimator(ABC):
-    # pylint: disable=too-many-instance-attributes
-    """An estimator contains all of the information necessary to compute a causal estimate for the effect of changing
-    a set of treatment variables to a set of values.
-
-    All estimators must implement the following two methods:
-
-    1) add_modelling_assumptions: The validity of a model-assisted causal inference result depends on whether
-    the modelling assumptions imposed by a model actually hold. Therefore, for each model, is important to state
-    the modelling assumption upon which the validity of the results depend. To achieve this, the estimator object
-    maintains a list of modelling assumptions (as strings). If a user wishes to implement their own estimator, they
-    must implement this method and add all assumptions to the list of modelling assumptions.
-
-    2) estimate_ate: All estimators must be capable of returning the average treatment effect as a minimum. That is, the
-    average effect of the intervention (changing treatment from control to treated value) on the outcome of interest
-    adjusted for all confounders.
-    """
-
-    def __init__(
-        # pylint: disable=too-many-arguments
-        self,
-        treatment: str,
-        treatment_value: float,
-        control_value: float,
-        adjustment_set: set,
-        outcome: str,
-        df: pd.DataFrame = None,
-        effect_modifiers: dict[str:Any] = None,
-        alpha: float = 0.05,
-        query: str = "",
-    ):
-        self.treatment = treatment
-        self.treatment_value = treatment_value
-        self.control_value = control_value
-        self.adjustment_set = adjustment_set
-        self.outcome = outcome
-        self.alpha = alpha
-        self.df = df.query(query) if query else df
-
-        if effect_modifiers is None:
-            self.effect_modifiers = {}
-        elif isinstance(effect_modifiers, dict):
-            self.effect_modifiers = effect_modifiers
-        else:
-            raise ValueError(f"Unsupported type for effect_modifiers {effect_modifiers}. Expected iterable")
-        self.modelling_assumptions = []
-        if query:
-            self.modelling_assumptions.append(query)
-        self.add_modelling_assumptions()
-        logger.debug("Effect Modifiers: %s", self.effect_modifiers)
-
-    @abstractmethod
-    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.
-        """
-
-    def compute_confidence_intervals(self) -> list[float, float]:
-        """
-        Estimate the 95% Wald confidence intervals for the effect of changing the treatment from control values to
-        treatment values on the outcome.
-        :return: 95% Wald confidence intervals.
-        """
-
-
-class LogisticRegressionEstimator(Estimator):
-    """A Logistic Regression Estimator is a parametric estimator which restricts the variables in the data to a linear
-    combination of parameters and functions of the variables (note these functions need not be linear). It is designed
-    for estimating categorical outcomes.
-    """
-
-    def __init__(
-        # pylint: disable=too-many-arguments
-        self,
-        treatment: str,
-        treatment_value: float,
-        control_value: float,
-        adjustment_set: set,
-        outcome: str,
-        df: pd.DataFrame = None,
-        effect_modifiers: dict[str:Any] = None,
-        formula: str = None,
-        query: str = "",
-    ):
-        super().__init__(
-            treatment=treatment,
-            treatment_value=treatment_value,
-            control_value=control_value,
-            adjustment_set=adjustment_set,
-            outcome=outcome,
-            df=df,
-            effect_modifiers=effect_modifiers,
-            query=query,
-        )
-
-        self.model = None
-
-        if formula is not None:
-            self.formula = formula
-        else:
-            terms = [treatment] + sorted(list(adjustment_set)) + sorted(list(self.effect_modifiers))
-            self.formula = f"{outcome} ~ {'+'.join(((terms)))}"
-
-    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.append(
-            "The variables in the data must fit a shape which can be expressed as a linear"
-            "combination of parameters and functions of variables. Note that these functions"
-            "do not need to be linear."
-        )
-        self.modelling_assumptions.append("The outcome must be binary.")
-        self.modelling_assumptions.append("Independently and identically distributed errors.")
-
-    def _run_logistic_regression(self, data) -> RegressionResultsWrapper:
-        """Run logistic regression of the treatment and adjustment set against the outcome and return the model.
-
-        :return: The model after fitting to data.
-        """
-        model = smf.logit(formula=self.formula, data=data).fit(disp=0)
-        self.model = model
-        return model
-
-    def estimate(self, data: pd.DataFrame, adjustment_config: dict = None) -> RegressionResultsWrapper:
-        """add terms to the dataframe and estimate the outcome from the data
-        :param data: A pandas dataframe containing execution data from the system-under-test.
-        :param adjustment_config: Dictionary containing the adjustment configuration of the adjustment set
-        """
-        if adjustment_config is None:
-            adjustment_config = {}
-        if set(self.adjustment_set) != set(adjustment_config):
-            raise ValueError(
-                f"Invalid adjustment configuration {adjustment_config}. Must specify values for {self.adjustment_set}"
-            )
-
-        model = self._run_logistic_regression(data)
-
-        x = pd.DataFrame(columns=self.df.columns)
-        x["Intercept"] = 1  # self.intercept
-        x[self.treatment] = [self.treatment_value, self.control_value]
-        for k, v in adjustment_config.items():
-            x[k] = v
-        for k, v in self.effect_modifiers.items():
-            x[k] = v
-        x = dmatrix(self.formula.split("~")[1], x, return_type="dataframe")
-        for col in x:
-            if str(x.dtypes[col]) == "object":
-                x = pd.get_dummies(x, columns=[col], drop_first=True)
-        # x = x[model.params.index]
-        return model.predict(x)
-
-    def estimate_control_treatment(
-        self, adjustment_config: dict = None, bootstrap_size: int = 100
-    ) -> tuple[pd.Series, pd.Series]:
-        """Estimate the outcomes under control and treatment.
-
-        :return: The estimated control and treatment values and their confidence
-        intervals in the form ((ci_low, control, ci_high), (ci_low, treatment, ci_high)).
-        """
-        if adjustment_config is None:
-            adjustment_config = {}
-        y = self.estimate(self.df, adjustment_config=adjustment_config)
-
-        try:
-            bootstrap_samples = [
-                self.estimate(self.df.sample(len(self.df), replace=True), adjustment_config=adjustment_config)
-                for _ in range(bootstrap_size)
-            ]
-            control, treatment = zip(*[(x.iloc[1], x.iloc[0]) for x in bootstrap_samples])
-        except PerfectSeparationError:
-            logger.warning(
-                "Perfect separation detected, results not available. Cannot calculate confidence intervals for such "
-                "a small dataset."
-            )
-            return (y.iloc[1], None), (y.iloc[0], None)
-        except np.linalg.LinAlgError:
-            logger.warning("Singular matrix detected. Confidence intervals not available. Try with a larger data set")
-            return (y.iloc[1], None), (y.iloc[0], None)
-
-        # Delta method confidence intervals from
-        # https://stackoverflow.com/questions/47414842/confidence-interval-of-probability-prediction-from-logistic-regression-statsmode
-        # cov = model.cov_params()
-        # gradient = (y * (1 - y) * x.T).T  # matrix of gradients for each observation
-        # std_errors = np.array([np.sqrt(np.dot(np.dot(g, cov), g)) for g in gradient.to_numpy()])
-        # c = 1.96  # multiplier for confidence interval
-        # upper = np.maximum(0, np.minimum(1, y + std_errors * c))
-        # lower = np.maximum(0, np.minimum(1, y - std_errors * c))
-
-        return (y.iloc[1], np.array(control)), (y.iloc[0], np.array(treatment))
-
-    def estimate_ate(self, adjustment_config: dict = None, bootstrap_size: int = 100) -> float:
-        """Estimate the ate 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. Here, we actually
-        calculate the expected outcomes under control and treatment and take one away from the other. This
-        allows for custom terms to be put in such as squares, inverses, products, etc.
-
-        :return: The estimated average treatment effect and 95% confidence intervals
-        """
-        if adjustment_config is None:
-            adjustment_config = {}
-        (control_outcome, control_bootstraps), (
-            treatment_outcome,
-            treatment_bootstraps,
-        ) = self.estimate_control_treatment(bootstrap_size=bootstrap_size, adjustment_config=adjustment_config)
-        estimate = treatment_outcome - control_outcome
-
-        if control_bootstraps is None or treatment_bootstraps is None:
-            return estimate, (None, None)
-
-        bootstraps = sorted(list(treatment_bootstraps - control_bootstraps))
-        bound = int((bootstrap_size * self.alpha) / 2)
-        ci_low = bootstraps[bound]
-        ci_high = bootstraps[bootstrap_size - bound]
-
-        logger.info(
-            f"Changing {self.treatment} from {self.control_value} to {self.treatment_value} gives an estimated "
-            f"ATE of {ci_low} < {estimate} < {ci_high}"
-        )
-        assert ci_low < estimate < ci_high, f"Expecting {ci_low} < {estimate} < {ci_high}"
-
-        return estimate, (ci_low, ci_high)
-
-    def estimate_risk_ratio(self, adjustment_config: dict = None, bootstrap_size: int = 100) -> float:
-        """Estimate the ate 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. Here, we actually
-        calculate the expected outcomes under control and treatment and divide one by the other. This
-        allows for custom terms to be put in such as squares, inverses, products, etc.
-
-        :return: The estimated risk ratio and 95% confidence intervals.
-        """
-        if adjustment_config is None:
-            adjustment_config = {}
-        (control_outcome, control_bootstraps), (
-            treatment_outcome,
-            treatment_bootstraps,
-        ) = self.estimate_control_treatment(bootstrap_size=bootstrap_size, adjustment_config=adjustment_config)
-        estimate = treatment_outcome / control_outcome
-
-        if control_bootstraps is None or treatment_bootstraps is None:
-            return estimate, (None, None)
-
-        bootstraps = sorted(list(treatment_bootstraps / control_bootstraps))
-        bound = ceil((bootstrap_size * self.alpha) / 2)
-        ci_low = bootstraps[bound]
-        ci_high = bootstraps[bootstrap_size - bound]
-
-        logger.info(
-            f"Changing {self.treatment} from {self.control_value} to {self.treatment_value} gives an estimated "
-            f"risk ratio of {ci_low} < {estimate} < {ci_high}"
-        )
-        assert ci_low < estimate < ci_high, f"Expecting {ci_low} < {estimate} < {ci_high}"
-
-        return estimate, (ci_low, ci_high)
-
-    def estimate_unit_odds_ratio(self) -> float:
-        """Estimate the odds ratio of increasing the treatment by one. In logistic regression, this corresponds to the
-        coefficient of the treatment of interest.
-
-        :return: The odds ratio. Confidence intervals are not yet supported.
-        """
-        model = self._run_logistic_regression(self.df)
-        return np.exp(model.params[self.treatment])
-
-
-class LinearRegressionEstimator(Estimator):
-    """A Linear Regression Estimator is a parametric estimator which restricts the variables in the data to a linear
-    combination of parameters and functions of the variables (note these functions need not be linear).
-    """
-
-    def __init__(
-        # pylint: disable=too-many-arguments
-        self,
-        treatment: str,
-        treatment_value: float,
-        control_value: float,
-        adjustment_set: set,
-        outcome: str,
-        df: pd.DataFrame = None,
-        effect_modifiers: dict[Variable:Any] = None,
-        formula: str = None,
-        alpha: float = 0.05,
-        query: str = "",
-    ):
-        super().__init__(
-            treatment,
-            treatment_value,
-            control_value,
-            adjustment_set,
-            outcome,
-            df,
-            effect_modifiers,
-            alpha=alpha,
-            query=query,
-        )
-
-        self.model = None
-        if effect_modifiers is None:
-            effect_modifiers = []
-
-        if formula is not None:
-            self.formula = formula
-        else:
-            terms = [treatment] + sorted(list(adjustment_set)) + sorted(list(effect_modifiers))
-            self.formula = f"{outcome} ~ {'+'.join(terms)}"
-
-        for term in self.effect_modifiers:
-            self.adjustment_set.add(term)
-
-    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.append(
-            "The variables in the data must fit a shape which can be expressed as a linear"
-            "combination of parameters and functions of variables. Note that these functions"
-            "do not need to be linear."
-        )
-
-    def estimate_coefficient(self) -> tuple[pd.Series, list[pd.Series, pd.Series]]:
-        """Estimate the unit average treatment effect of the treatment on the outcome. That is, the change in outcome
-        caused by a unit change in treatment.
-
-        :return: The unit average treatment effect and the 95% Wald confidence intervals.
-        """
-        model = self._run_linear_regression()
-        newline = "\n"
-        patsy_md = ModelDesc.from_formula(self.treatment)
-
-        if any(
-            (
-                self.df.dtypes[factor.name()] == "object"
-                for factor in patsy_md.rhs_termlist[1].factors
-                # We want to remove this long term as it prevents us from discovering categoricals within I(...) blocks
-                if factor.name() in self.df.dtypes
-            )
-        ):
-            design_info = dmatrix(self.formula.split("~")[1], self.df).design_info
-            treatment = design_info.column_names[design_info.term_name_slices[self.treatment]]
-        else:
-            treatment = [self.treatment]
-        assert set(treatment).issubset(
-            model.params.index.tolist()
-        ), f"{treatment} not in\n{'  ' + str(model.params.index).replace(newline, newline + '  ')}"
-        unit_effect = model.params[treatment]  # Unit effect is the coefficient of the treatment
-        [ci_low, ci_high] = self._get_confidence_intervals(model, treatment)
-        return unit_effect, [ci_low, ci_high]
-
-    def estimate_ate(self) -> tuple[pd.Series, list[pd.Series, pd.Series]]:
-        """Estimate the 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 average treatment effect and the 95% Wald confidence intervals.
-        """
-        model = self._run_linear_regression()
-
-        # Create an empty individual for the control and treated
-        individuals = pd.DataFrame(1, index=["control", "treated"], columns=model.params.index)
-
-        # For Pandas version > 2, we need to explicitly state that the dataframe takes floating-point values
-        individuals = individuals.astype(float)
-
-        # It is ABSOLUTELY CRITICAL that these go last, otherwise we can't index
-        # the effect with "ate = t_test_results.effect[0]"
-        individuals.loc["control", [self.treatment]] = self.control_value
-        individuals.loc["treated", [self.treatment]] = self.treatment_value
-
-        # Perform a t-test to compare the predicted outcome of the control and treated individual (ATE)
-        t_test_results = model.t_test(individuals.loc["treated"] - individuals.loc["control"])
-        ate = pd.Series(t_test_results.effect[0])
-        confidence_intervals = list(t_test_results.conf_int(alpha=self.alpha).flatten())
-        confidence_intervals = [pd.Series(interval) for interval in confidence_intervals]
-        return ate, confidence_intervals
-
-    def estimate_control_treatment(self, adjustment_config: dict = None) -> tuple[pd.Series, pd.Series]:
-        """Estimate the outcomes under control and treatment.
-
-        :return: The estimated outcome under control and treatment in the form
-        (control_outcome, treatment_outcome).
-        """
-        if adjustment_config is None:
-            adjustment_config = {}
-        model = self._run_linear_regression()
-
-        x = pd.DataFrame(columns=self.df.columns)
-        x[self.treatment] = [self.treatment_value, self.control_value]
-        x["Intercept"] = 1  # self.intercept
-        for k, v in adjustment_config.items():
-            x[k] = v
-        for k, v in self.effect_modifiers.items():
-            x[k] = v
-        x = dmatrix(self.formula.split("~")[1], x, return_type="dataframe")
-        for col in x:
-            if str(x.dtypes[col]) == "object":
-                x = pd.get_dummies(x, columns=[col], drop_first=True)
-        x = x[model.params.index]
-        y = model.get_prediction(x).summary_frame()
-
-        return y.iloc[1], y.iloc[0]
-
-    def estimate_risk_ratio(self, adjustment_config: dict = None) -> tuple[pd.Series, list[pd.Series, pd.Series]]:
-        """Estimate the risk_ratio 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 average treatment effect and the 95% Wald confidence intervals.
-        """
-        if adjustment_config is None:
-            adjustment_config = {}
-        control_outcome, treatment_outcome = self.estimate_control_treatment(adjustment_config=adjustment_config)
-        ci_low = pd.Series(treatment_outcome["mean_ci_lower"] / control_outcome["mean_ci_upper"])
-        ci_high = pd.Series(treatment_outcome["mean_ci_upper"] / control_outcome["mean_ci_lower"])
-        return pd.Series(treatment_outcome["mean"] / control_outcome["mean"]), [ci_low, ci_high]
-
-    def estimate_ate_calculated(self, adjustment_config: dict = None) -> tuple[pd.Series, list[pd.Series, pd.Series]]:
-        """Estimate the ate 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. Here, we actually
-        calculate the expected outcomes under control and treatment and divide one by the other. This
-        allows for custom terms to be put in such as squares, inverses, products, etc.
-
-        :return: The average treatment effect and the 95% Wald confidence intervals.
-        """
-        if adjustment_config is None:
-            adjustment_config = {}
-        control_outcome, treatment_outcome = self.estimate_control_treatment(adjustment_config=adjustment_config)
-        ci_low = pd.Series(treatment_outcome["mean_ci_lower"] - control_outcome["mean_ci_upper"])
-        ci_high = pd.Series(treatment_outcome["mean_ci_upper"] - control_outcome["mean_ci_lower"])
-        return pd.Series(treatment_outcome["mean"] - control_outcome["mean"]), [ci_low, ci_high]
-
-    def _run_linear_regression(self) -> RegressionResultsWrapper:
-        """Run linear regression of the treatment and adjustment set against the outcome and return the model.
-
-        :return: The model after fitting to data.
-        """
-        model = smf.ols(formula=self.formula, data=self.df).fit()
-        self.model = model
-        return model
-
-    def _get_confidence_intervals(self, model, treatment):
-        confidence_intervals = model.conf_int(alpha=self.alpha, cols=None)
-        ci_low, ci_high = (
-            pd.Series(confidence_intervals[0].loc[treatment]),
-            pd.Series(confidence_intervals[1].loc[treatment]),
-        )
-        return [ci_low, ci_high]
-
-
-class CubicSplineRegressionEstimator(LinearRegressionEstimator):
-    """A Cubic Spline Regression Estimator is a parametric estimator which restricts the variables in the data to a
-    combination of parameters and basis functions of the variables.
-    """
-
-    def __init__(
-        # pylint: disable=too-many-arguments
-        self,
-        treatment: str,
-        treatment_value: float,
-        control_value: float,
-        adjustment_set: set,
-        outcome: str,
-        basis: int,
-        df: pd.DataFrame = None,
-        effect_modifiers: dict[Variable:Any] = None,
-        formula: str = None,
-        alpha: float = 0.05,
-        expected_relationship=None,
-    ):
-        super().__init__(
-            treatment, treatment_value, control_value, adjustment_set, outcome, df, effect_modifiers, formula, alpha
-        )
-
-        self.expected_relationship = expected_relationship
-
-        if effect_modifiers is None:
-            effect_modifiers = []
-
-        if formula is None:
-            terms = [treatment] + sorted(list(adjustment_set)) + sorted(list(effect_modifiers))
-            self.formula = f"{outcome} ~ cr({'+'.join(terms)}, df={basis})"
-
-    def estimate_ate_calculated(self, adjustment_config: dict = None) -> pd.Series:
-        model = self._run_linear_regression()
-
-        x = {"Intercept": 1, self.treatment: self.treatment_value}
-        if adjustment_config is not None:
-            for k, v in adjustment_config.items():
-                x[k] = v
-        if self.effect_modifiers is not None:
-            for k, v in self.effect_modifiers.items():
-                x[k] = v
-
-        treatment = model.predict(x).iloc[0]
-
-        x[self.treatment] = self.control_value
-        control = model.predict(x).iloc[0]
-
-        return pd.Series(treatment - control)
-
-
-class InstrumentalVariableEstimator(Estimator):
-    """
-    Carry out estimation using instrumental variable adjustment rather than conventional adjustment. This means we do
-    not need to observe all confounders in order to adjust for them. A key assumption here is linearity.
-    """
-
-    def __init__(
-        # pylint: disable=too-many-arguments
-        self,
-        treatment: str,
-        treatment_value: float,
-        control_value: float,
-        adjustment_set: set,
-        outcome: str,
-        instrument: str,
-        df: pd.DataFrame = None,
-        intercept: int = 1,
-        effect_modifiers: dict = None,  # Not used (yet?). Needed for compatibility
-        alpha: float = 0.05,
-        query: str = "",
-    ):
-        super().__init__(
-            treatment=treatment,
-            treatment_value=treatment_value,
-            control_value=control_value,
-            adjustment_set=adjustment_set,
-            outcome=outcome,
-            df=df,
-            effect_modifiers=None,
-            alpha=alpha,
-            query=query,
-        )
-        self.intercept = intercept
-        self.model = None
-        self.instrument = instrument
-
-    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.append(
-            """The instrument and the treatment, and the treatment and the outcome must be
-        related linearly in the form Y = aX + b."""
-        )
-        self.modelling_assumptions.append(
-            """The three IV conditions must hold
-            (i) Instrument is associated with treatment
-            (ii) Instrument does not affect outcome except through its potential effect on treatment
-            (iii) Instrument and outcome do not share causes
-        """
-        )
-
-    def estimate_iv_coefficient(self, df) -> float:
-        """
-        Estimate the linear regression coefficient of the treatment on the
-        outcome.
-        """
-        # Estimate the total effect of instrument I on outcome Y = abI + c1
-        ab = sm.OLS(df[self.outcome], df[[self.instrument]]).fit().params[self.instrument]
-
-        # Estimate the direct effect of instrument I on treatment X = aI + c1
-        a = sm.OLS(df[self.treatment], df[[self.instrument]]).fit().params[self.instrument]
-
-        # Estimate the coefficient of I on X by cancelling
-        return ab / a
-
-    def estimate_coefficient(self, bootstrap_size=100) -> tuple[pd.Series, list[pd.Series, pd.Series]]:
-        """
-        Estimate the unit ate (i.e. coefficient) of the treatment on the
-        outcome.
-        """
-        bootstraps = sorted(
-            [self.estimate_iv_coefficient(self.df.sample(len(self.df), replace=True)) for _ in range(bootstrap_size)]
-        )
-        bound = ceil((bootstrap_size * self.alpha) / 2)
-        ci_low = pd.Series(bootstraps[bound])
-        ci_high = pd.Series(bootstraps[bootstrap_size - bound])
-
-        return pd.Series(self.estimate_iv_coefficient(self.df)), [ci_low, ci_high]
-
-
-class IPCWEstimator(Estimator):
-    """
-    Class to perform inverse probability of censoring weighting (IPCW) estimation
-    for sequences of treatments over time-varying data.
-    """
-
-    # pylint: disable=too-many-arguments
-    # pylint: disable=too-many-instance-attributes
-    def __init__(
-        self,
-        df: pd.DataFrame,
-        timesteps_per_intervention: int,
-        control_strategy: TreatmentSequence,
-        treatment_strategy: TreatmentSequence,
-        outcome: str,
-        fault_column: str,
-        fit_bl_switch_formula: str,
-        fit_bltd_switch_formula: str,
-        eligibility=None,
-        alpha: float = 0.05,
-    ):
-        super().__init__(
-            [c.variable for c in treatment_strategy.capabilities],
-            [c.value for c in treatment_strategy.capabilities],
-            [c.value for c in control_strategy.capabilities],
-            None,
-            outcome,
-            df,
-            None,
-            alpha=alpha,
-            query="",
-        )
-        self.timesteps_per_intervention = timesteps_per_intervention
-        self.control_strategy = control_strategy
-        self.treatment_strategy = treatment_strategy
-        self.outcome = outcome
-        self.fault_column = fault_column
-        self.timesteps_per_intervention = timesteps_per_intervention
-        self.fit_bl_switch_formula = fit_bl_switch_formula
-        self.fit_bltd_switch_formula = fit_bltd_switch_formula
-        self.eligibility = eligibility
-        self.df = df
-        self.preprocess_data()
-
-    def add_modelling_assumptions(self):
-        self.modelling_assumptions.append("The variables in the data vary over time.")
-
-    def setup_xo_t_do(self, strategy_assigned: list, strategy_followed: list, eligible: pd.Series):
-        """
-        Return a binary sequence with each bit representing whether the current
-        index is the time point at which the individual diverted from the
-        assigned treatment strategy (and thus should be censored).
-
-        :param strategy_assigned - the assigned treatment strategy
-        :param strategy_followed - the strategy followed by the individual
-        :param eligible - binary sequence represnting the eligibility of the individual at each time step
-        """
-        strategy_assigned = [1] + strategy_assigned + [1]
-        strategy_followed = [1] + strategy_followed + [1]
-
-        mask = (
-            pd.Series(strategy_assigned, index=eligible.index) != pd.Series(strategy_followed, index=eligible.index)
-        ).astype("boolean")
-        mask = mask | ~eligible
-        mask.reset_index(inplace=True, drop=True)
-        false = mask.loc[mask]
-        if false.empty:
-            return np.zeros(len(mask))
-        mask = (mask * 1).tolist()
-        cutoff = false.index[0] + 1
-        return mask[:cutoff] + ([None] * (len(mask) - cutoff))
-
-    def setup_fault_t_do(self, individual: pd.DataFrame):
-        """
-        Return a binary sequence with each bit representing whether the current
-        index is the time point at which the event of interest (i.e. a fault)
-        occurred.
-        """
-        fault = individual[~individual[self.fault_column]]
-        fault_t_do = pd.Series(np.zeros(len(individual)), index=individual.index)
-
-        if not fault.empty:
-            fault_time = individual["time"].loc[fault.index[0]]
-            # Ceiling to nearest observation point
-            fault_time = ceil(fault_time / self.timesteps_per_intervention) * self.timesteps_per_intervention
-            # Set the correct observation point to be the fault time of doing (fault_t_do)
-            observations = individual.loc[
-                (individual["time"] % self.timesteps_per_intervention == 0) & (individual["time"] < fault_time)
-            ]
-            if not observations.empty:
-                fault_t_do.loc[observations.index[0]] = 1
-        assert sum(fault_t_do) <= 1, f"Multiple fault times for\n{individual}"
-
-        return pd.DataFrame({"fault_t_do": fault_t_do})
-
-    def setup_fault_time(self, individual: pd.DataFrame, perturbation: float = -0.001):
-        """
-        Return the time at which the event of interest (i.e. a fault) occurred.
-        """
-        fault = individual[~individual[self.fault_column]]
-        fault_time = (
-            individual["time"].loc[fault.index[0]]
-            if not fault.empty
-            else (individual["time"].max() + self.timesteps_per_intervention)
-        )
-        return pd.DataFrame({"fault_time": np.repeat(fault_time + perturbation, len(individual))})
-
-    def preprocess_data(self):
-        """
-        Set up the treatment-specific columns in the data that are needed to estimate the hazard ratio.
-        """
-        self.df["trtrand"] = None  # treatment/control arm
-        self.df["xo_t_do"] = None  # did the individual deviate from the treatment of interest here?
-        self.df["eligible"] = self.df.eval(self.eligibility) if self.eligibility is not None else True
-
-        # when did a fault occur?
-        self.df["fault_time"] = self.df.groupby("id")[[self.fault_column, "time"]].apply(self.setup_fault_time).values
-        self.df["fault_t_do"] = (
-            self.df.groupby("id")[["id", "time", self.fault_column]].apply(self.setup_fault_t_do).values
-        )
-        assert not pd.isnull(self.df["fault_time"]).any()
-
-        living_runs = self.df.query("fault_time > 0").loc[
-            (self.df["time"] % self.timesteps_per_intervention == 0)
-            & (self.df["time"] <= self.control_strategy.total_time())
-        ]
-
-        individuals = []
-        new_id = 0
-        logging.debug("  Preprocessing groups")
-        for _, individual in living_runs.groupby("id"):
-            assert sum(individual["fault_t_do"]) <= 1, (
-                f"Error initialising fault_t_do for individual\n"
-                f"{individual[['id', 'time', 'fault_time', 'fault_t_do']]}\n"
-                "with fault at {individual.fault_time.iloc[0]}"
-            )
-
-            strategy_followed = [
-                Capability(
-                    c.variable,
-                    individual.loc[individual["time"] == c.start_time, c.variable].values[0],
-                    c.start_time,
-                    c.end_time,
-                )
-                for c in self.treatment_strategy.capabilities
-            ]
-
-            # Control flow:
-            # Individuals that start off in both arms, need cloning (hence incrementing the ID within the if statement)
-            # Individuals that don't start off in either arm are left out
-            for inx, strategy_assigned in [(0, self.control_strategy), (1, self.treatment_strategy)]:
-                if strategy_assigned.capabilities[0] == strategy_followed[0] and individual.eligible.iloc[0]:
-                    individual["id"] = new_id
-                    new_id += 1
-                    individual["trtrand"] = inx
-                    individual["xo_t_do"] = self.setup_xo_t_do(
-                        strategy_assigned.capabilities, strategy_followed, individual["eligible"]
-                    )
-                    individuals.append(individual.loc[individual["time"] <= individual["fault_time"]].copy())
-        if len(individuals) == 0:
-            raise ValueError("No individuals followed either strategy.")
-
-        self.df = pd.concat(individuals)
-
-    def estimate_hazard_ratio(self):
-        """
-        Estimate the hazard ratio.
-        """
-
-        if self.df["fault_t_do"].sum() == 0:
-            raise ValueError("No recorded faults")
-
-        preprocessed_data = self.df.loc[self.df["xo_t_do"] == 0].copy()
-
-        # Use logistic regression to predict switching given baseline covariates
-        fit_bl_switch = smf.logit(self.fit_bl_switch_formula, data=self.df).fit()
-
-        preprocessed_data["pxo1"] = fit_bl_switch.predict(preprocessed_data)
-
-        # Use logistic regression to predict switching given baseline and time-updated covariates (model S12)
-        fit_bltd_switch = smf.logit(
-            self.fit_bltd_switch_formula,
-            data=self.df,
-        ).fit()
-
-        preprocessed_data["pxo2"] = fit_bltd_switch.predict(preprocessed_data)
-
-        # IPCW step 3: For each individual at each time, compute the inverse probability of remaining uncensored
-        # Estimate the probabilities of remaining ‘un-switched’ and hence the weights
-
-        preprocessed_data["num"] = 1 - preprocessed_data["pxo1"]
-        preprocessed_data["denom"] = 1 - preprocessed_data["pxo2"]
-        preprocessed_data[["num", "denom"]] = (
-            preprocessed_data.sort_values(["id", "time"]).groupby("id")[["num", "denom"]].cumprod()
-        )
-
-        assert (
-            not preprocessed_data["num"].isnull().any()
-        ), f"{len(preprocessed_data['num'].isnull())} null numerator values"
-        assert (
-            not preprocessed_data["denom"].isnull().any()
-        ), f"{len(preprocessed_data['denom'].isnull())} null denom values"
-
-        preprocessed_data["weight"] = 1 / preprocessed_data["denom"]
-        preprocessed_data["sweight"] = preprocessed_data["num"] / preprocessed_data["denom"]
-
-        preprocessed_data["tin"] = preprocessed_data["time"]
-        preprocessed_data["tout"] = pd.concat(
-            [(preprocessed_data["time"] + self.timesteps_per_intervention), preprocessed_data["fault_time"]],
-            axis=1,
-        ).min(axis=1)
-
-        assert (preprocessed_data["tin"] <= preprocessed_data["tout"]).all(), (
-            f"Left before joining\n" f"{preprocessed_data.loc[preprocessed_data['tin'] >= preprocessed_data['tout']]}"
-        )
-
-        #  IPCW step 4: Use these weights in a weighted analysis of the outcome model
-        # Estimate the KM graph and IPCW hazard ratio using Cox regression.
-        cox_ph = CoxPHFitter(alpha=self.alpha)
-        cox_ph.fit(
-            df=preprocessed_data,
-            duration_col="tout",
-            event_col="fault_t_do",
-            weights_col="weight",
-            cluster_col="id",
-            robust=True,
-            formula="trtrand",
-            entry_col="tin",
-        )
-
-        ci_low, ci_high = [np.exp(cox_ph.confidence_intervals_)[col] for col in cox_ph.confidence_intervals_.columns]
-
-        return (cox_ph.hazard_ratios_, (ci_low, ci_high))

From 97fb6e79398618845d9aa1a7d45dda0bc69ad03f Mon Sep 17 00:00:00 2001
From: Farhad Allian <allianfarhad@gmail.com>
Date: Fri, 13 Sep 2024 17:55:19 +0100
Subject: [PATCH 23/24] fix: code coverage

---
 .../estimation/linear_regression_estimator.py          | 10 ----------
 1 file changed, 10 deletions(-)

diff --git a/causal_testing/estimation/linear_regression_estimator.py b/causal_testing/estimation/linear_regression_estimator.py
index bfefd445..096d2960 100644
--- a/causal_testing/estimation/linear_regression_estimator.py
+++ b/causal_testing/estimation/linear_regression_estimator.py
@@ -92,16 +92,6 @@ def gp_formula(
         formula = gp.simplify(formula)
         self.formula = f"{self.outcome} ~ I({formula}) - 1"
 
-    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.append(
-            "The variables in the data must fit a shape which can be expressed as a linear"
-            "combination of parameters and functions of variables. Note that these functions"
-            "do not need to be linear."
-        )
 
     def estimate_coefficient(self) -> tuple[pd.Series, list[pd.Series, pd.Series]]:
         """Estimate the unit average treatment effect of the treatment on the outcome. That is, the change in outcome

From ed832dadcc3778f26f90f785e086f508948d0ea9 Mon Sep 17 00:00:00 2001
From: Michael Foster <m.foster@sheffield.ac.uk>
Date: Tue, 17 Sep 2024 15:23:31 +0100
Subject: [PATCH 24/24] file renameing as suggested by @f-allian

---
 .../estimation/{estimator.py => abstract_estimator.py}       | 0
 ...ression_estimator.py => abstract_regression_estimator.py} | 2 +-
 .../{gp.py => genetic_programming_regression_fitter.py}      | 0
 .../{iv_estimator.py => instrumental_variable_estimator.py}  | 2 +-
 causal_testing/estimation/ipcw_estimator.py                  | 2 +-
 causal_testing/estimation/linear_regression_estimator.py     | 5 ++---
 causal_testing/estimation/logistic_regression_estimator.py   | 2 +-
 causal_testing/json_front/json_class.py                      | 2 +-
 causal_testing/testing/causal_test_adequacy.py               | 2 +-
 causal_testing/testing/causal_test_case.py                   | 2 +-
 causal_testing/testing/causal_test_result.py                 | 2 +-
 causal_testing/testing/causal_test_suite.py                  | 2 +-
 examples/poisson-line-process/example_poisson_process.py     | 2 +-
 examples/poisson/example_run_causal_tests.py                 | 2 +-
 ...t_gp.py => test_genetic_programming_regression_fitter.py} | 3 +--
 ..._estimator.py => test_instrumental_variable_estimator.py} | 2 +-
 tests/estimation_tests/test_linear_regression_estimator.py   | 2 +-
 tests/json_front_tests/test_json_class.py                    | 2 +-
 18 files changed, 17 insertions(+), 19 deletions(-)
 rename causal_testing/estimation/{estimator.py => abstract_estimator.py} (100%)
 rename causal_testing/estimation/{regression_estimator.py => abstract_regression_estimator.py} (98%)
 rename causal_testing/estimation/{gp.py => genetic_programming_regression_fitter.py} (100%)
 rename causal_testing/estimation/{iv_estimator.py => instrumental_variable_estimator.py} (98%)
 rename tests/estimation_tests/{test_gp.py => test_genetic_programming_regression_fitter.py} (77%)
 rename tests/estimation_tests/{test_iv_estimator.py => test_instrumental_variable_estimator.py} (93%)

diff --git a/causal_testing/estimation/estimator.py b/causal_testing/estimation/abstract_estimator.py
similarity index 100%
rename from causal_testing/estimation/estimator.py
rename to causal_testing/estimation/abstract_estimator.py
diff --git a/causal_testing/estimation/regression_estimator.py b/causal_testing/estimation/abstract_regression_estimator.py
similarity index 98%
rename from causal_testing/estimation/regression_estimator.py
rename to causal_testing/estimation/abstract_regression_estimator.py
index 44ef67da..c6786d20 100644
--- a/causal_testing/estimation/regression_estimator.py
+++ b/causal_testing/estimation/abstract_regression_estimator.py
@@ -9,7 +9,7 @@
 from patsy import dmatrix  # pylint: disable = no-name-in-module
 
 from causal_testing.specification.variable import Variable
-from causal_testing.estimation.estimator import Estimator
+from causal_testing.estimation.abstract_estimator import Estimator
 
 logger = logging.getLogger(__name__)
 
diff --git a/causal_testing/estimation/gp.py b/causal_testing/estimation/genetic_programming_regression_fitter.py
similarity index 100%
rename from causal_testing/estimation/gp.py
rename to causal_testing/estimation/genetic_programming_regression_fitter.py
diff --git a/causal_testing/estimation/iv_estimator.py b/causal_testing/estimation/instrumental_variable_estimator.py
similarity index 98%
rename from causal_testing/estimation/iv_estimator.py
rename to causal_testing/estimation/instrumental_variable_estimator.py
index b8b56170..38d0fc1b 100644
--- a/causal_testing/estimation/iv_estimator.py
+++ b/causal_testing/estimation/instrumental_variable_estimator.py
@@ -6,7 +6,7 @@
 import pandas as pd
 import statsmodels.api as sm
 
-from causal_testing.estimation.estimator import Estimator
+from causal_testing.estimation.abstract_estimator import Estimator
 
 logger = logging.getLogger(__name__)
 
diff --git a/causal_testing/estimation/ipcw_estimator.py b/causal_testing/estimation/ipcw_estimator.py
index a0c5a819..584182ff 100644
--- a/causal_testing/estimation/ipcw_estimator.py
+++ b/causal_testing/estimation/ipcw_estimator.py
@@ -9,7 +9,7 @@
 from lifelines import CoxPHFitter
 
 from causal_testing.specification.capabilities import TreatmentSequence, Capability
-from causal_testing.estimation.estimator import Estimator
+from causal_testing.estimation.abstract_estimator import Estimator
 
 logger = logging.getLogger(__name__)
 
diff --git a/causal_testing/estimation/linear_regression_estimator.py b/causal_testing/estimation/linear_regression_estimator.py
index 096d2960..85a4b178 100644
--- a/causal_testing/estimation/linear_regression_estimator.py
+++ b/causal_testing/estimation/linear_regression_estimator.py
@@ -8,8 +8,8 @@
 from patsy import dmatrix, ModelDesc  # pylint: disable = no-name-in-module
 
 from causal_testing.specification.variable import Variable
-from causal_testing.estimation.gp import GP
-from causal_testing.estimation.regression_estimator import RegressionEstimator
+from causal_testing.estimation.genetic_programming_regression_fitter import GP
+from causal_testing.estimation.abstract_regression_estimator import RegressionEstimator
 
 logger = logging.getLogger(__name__)
 
@@ -92,7 +92,6 @@ def gp_formula(
         formula = gp.simplify(formula)
         self.formula = f"{self.outcome} ~ I({formula}) - 1"
 
-
     def estimate_coefficient(self) -> tuple[pd.Series, list[pd.Series, pd.Series]]:
         """Estimate the unit average treatment effect of the treatment on the outcome. That is, the change in outcome
         caused by a unit change in treatment.
diff --git a/causal_testing/estimation/logistic_regression_estimator.py b/causal_testing/estimation/logistic_regression_estimator.py
index 0546ba6a..ca5537d4 100644
--- a/causal_testing/estimation/logistic_regression_estimator.py
+++ b/causal_testing/estimation/logistic_regression_estimator.py
@@ -5,7 +5,7 @@
 import numpy as np
 import statsmodels.formula.api as smf
 
-from causal_testing.estimation.regression_estimator import RegressionEstimator
+from causal_testing.estimation.abstract_regression_estimator import RegressionEstimator
 
 logger = logging.getLogger(__name__)
 
diff --git a/causal_testing/json_front/json_class.py b/causal_testing/json_front/json_class.py
index 9a0de6e7..6be7fa68 100644
--- a/causal_testing/json_front/json_class.py
+++ b/causal_testing/json_front/json_class.py
@@ -25,7 +25,7 @@
 from causal_testing.testing.base_test_case import BaseTestCase
 from causal_testing.testing.causal_test_adequacy import DataAdequacy
 
-from causal_testing.estimation.estimator import Estimator
+from causal_testing.estimation.abstract_estimator import Estimator
 from causal_testing.estimation.linear_regression_estimator import LinearRegressionEstimator
 from causal_testing.estimation.logistic_regression_estimator import LogisticRegressionEstimator
 
diff --git a/causal_testing/testing/causal_test_adequacy.py b/causal_testing/testing/causal_test_adequacy.py
index 003ac714..5fb043eb 100644
--- a/causal_testing/testing/causal_test_adequacy.py
+++ b/causal_testing/testing/causal_test_adequacy.py
@@ -11,7 +11,7 @@
 
 from causal_testing.testing.causal_test_suite import CausalTestSuite
 from causal_testing.specification.causal_dag import CausalDAG
-from causal_testing.estimation.estimator import Estimator
+from causal_testing.estimation.abstract_estimator import Estimator
 from causal_testing.testing.causal_test_case import CausalTestCase
 
 logger = logging.getLogger(__name__)
diff --git a/causal_testing/testing/causal_test_case.py b/causal_testing/testing/causal_test_case.py
index a648d7a9..08ad54f1 100644
--- a/causal_testing/testing/causal_test_case.py
+++ b/causal_testing/testing/causal_test_case.py
@@ -7,7 +7,7 @@
 from causal_testing.specification.variable import Variable
 from causal_testing.testing.causal_test_outcome import CausalTestOutcome
 from causal_testing.testing.base_test_case import BaseTestCase
-from causal_testing.estimation.estimator import Estimator
+from causal_testing.estimation.abstract_estimator import Estimator
 from causal_testing.testing.causal_test_result import CausalTestResult, TestValue
 from causal_testing.data_collection.data_collector import DataCollector
 
diff --git a/causal_testing/testing/causal_test_result.py b/causal_testing/testing/causal_test_result.py
index 11c6790d..bfcfe826 100644
--- a/causal_testing/testing/causal_test_result.py
+++ b/causal_testing/testing/causal_test_result.py
@@ -6,7 +6,7 @@
 from dataclasses import dataclass
 import pandas as pd
 
-from causal_testing.estimation.estimator import Estimator
+from causal_testing.estimation.abstract_estimator import Estimator
 from causal_testing.specification.variable import Variable
 
 
diff --git a/causal_testing/testing/causal_test_suite.py b/causal_testing/testing/causal_test_suite.py
index 062cd7e5..14099143 100644
--- a/causal_testing/testing/causal_test_suite.py
+++ b/causal_testing/testing/causal_test_suite.py
@@ -7,7 +7,7 @@
 from typing import Type, Iterable
 from causal_testing.testing.base_test_case import BaseTestCase
 from causal_testing.testing.causal_test_case import CausalTestCase
-from causal_testing.estimation.estimator import Estimator
+from causal_testing.estimation.abstract_estimator import Estimator
 from causal_testing.testing.causal_test_result import CausalTestResult
 from causal_testing.data_collection.data_collector import DataCollector
 from causal_testing.specification.causal_specification import CausalSpecification
diff --git a/examples/poisson-line-process/example_poisson_process.py b/examples/poisson-line-process/example_poisson_process.py
index a8d4addd..be8bc906 100644
--- a/examples/poisson-line-process/example_poisson_process.py
+++ b/examples/poisson-line-process/example_poisson_process.py
@@ -5,7 +5,7 @@
 from causal_testing.testing.causal_test_case import CausalTestCase
 from causal_testing.testing.causal_test_outcome import ExactValue, Positive
 from causal_testing.estimation.linear_regression_estimator import LinearRegressionEstimator
-from causal_testing.estimation.estimator import Estimator
+from causal_testing.estimation.abstract_estimator import Estimator
 from causal_testing.testing.base_test_case import BaseTestCase
 
 import pandas as pd
diff --git a/examples/poisson/example_run_causal_tests.py b/examples/poisson/example_run_causal_tests.py
index 787ab6ac..2ae72e20 100644
--- a/examples/poisson/example_run_causal_tests.py
+++ b/examples/poisson/example_run_causal_tests.py
@@ -7,7 +7,7 @@
 from causal_testing.testing.causal_test_outcome import ExactValue, Positive, Negative, NoEffect, CausalTestOutcome
 from causal_testing.testing.causal_test_result import CausalTestResult
 from causal_testing.json_front.json_class import JsonUtility
-from causal_testing.estimation.estimator import Estimator
+from causal_testing.estimation.abstract_estimator import Estimator
 from causal_testing.specification.scenario import Scenario
 from causal_testing.specification.variable import Input, Output, Meta
 
diff --git a/tests/estimation_tests/test_gp.py b/tests/estimation_tests/test_genetic_programming_regression_fitter.py
similarity index 77%
rename from tests/estimation_tests/test_gp.py
rename to tests/estimation_tests/test_genetic_programming_regression_fitter.py
index 4946a871..4382b819 100644
--- a/tests/estimation_tests/test_gp.py
+++ b/tests/estimation_tests/test_genetic_programming_regression_fitter.py
@@ -1,11 +1,10 @@
 import unittest
 import pandas as pd
 
-from causal_testing.estimation.gp import GP
+from causal_testing.estimation.genetic_programming_regression_fitter import GP
 
 
 class TestGP(unittest.TestCase):
-
     def test_init_invalid_fun_name(self):
         with self.assertRaises(ValueError):
             GP(df=pd.DataFrame(), features=[], outcome="", max_order=2, sympy_conversions={"power_1": ""})
diff --git a/tests/estimation_tests/test_iv_estimator.py b/tests/estimation_tests/test_instrumental_variable_estimator.py
similarity index 93%
rename from tests/estimation_tests/test_iv_estimator.py
rename to tests/estimation_tests/test_instrumental_variable_estimator.py
index 9711f555..4cfa7dfe 100644
--- a/tests/estimation_tests/test_iv_estimator.py
+++ b/tests/estimation_tests/test_instrumental_variable_estimator.py
@@ -6,7 +6,7 @@
 from causal_testing.utils.validation import CausalValidator
 from causal_testing.specification.capabilities import TreatmentSequence
 
-from causal_testing.estimation.iv_estimator import InstrumentalVariableEstimator
+from causal_testing.estimation.instrumental_variable_estimator import InstrumentalVariableEstimator
 
 
 class TestInstrumentalVariableEstimator(unittest.TestCase):
diff --git a/tests/estimation_tests/test_linear_regression_estimator.py b/tests/estimation_tests/test_linear_regression_estimator.py
index 8778c086..c3749fdb 100644
--- a/tests/estimation_tests/test_linear_regression_estimator.py
+++ b/tests/estimation_tests/test_linear_regression_estimator.py
@@ -7,7 +7,7 @@
 from causal_testing.specification.capabilities import TreatmentSequence
 
 from causal_testing.estimation.linear_regression_estimator import LinearRegressionEstimator
-from causal_testing.estimation.gp import reciprocal
+from causal_testing.estimation.genetic_programming_regression_fitter import reciprocal
 
 
 def load_nhefs_df():
diff --git a/tests/json_front_tests/test_json_class.py b/tests/json_front_tests/test_json_class.py
index 3aff9782..ab565da7 100644
--- a/tests/json_front_tests/test_json_class.py
+++ b/tests/json_front_tests/test_json_class.py
@@ -5,7 +5,7 @@
 import os
 
 from causal_testing.estimation.linear_regression_estimator import LinearRegressionEstimator
-from causal_testing.estimation.estimator import Estimator
+from causal_testing.estimation.abstract_estimator import Estimator
 from causal_testing.testing.causal_test_outcome import NoEffect, Positive
 from causal_testing.json_front.json_class import JsonUtility, CausalVariables
 from causal_testing.specification.variable import Input, Output, Meta