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

Issue:
Logistic regression formula conversion issues
*Branch:* `functional_form`
*Steps to reproduce:* `pytest` from root directory
*Expected behaviour:* all tests should pass
*Observed behaviour:* Two tests in `tests/testing_tests/test_estimators.py` fail
*Further details:* I changed everything over in the `LogisticRegressionEstimator` class and the tests broke. I've isolated it to line 164 of `estimators.py`, i.e. `regression = smf.logit(formula=self.formula, data=self.df)`. Using the old `Logit` estimator from the line above works fine, so the models are somehow different...


Discussion / hints:

Fix patch:
diff --git a/examples/poisson-line-process/.gitignore b/examples/poisson-line-process/.gitignore
new file mode 100644
index 00000000..cfe602e6
--- /dev/null
+++ b/examples/poisson-line-process/.gitignore
@@ -0,0 +1,2 @@
+intensity_num_shapes_results_random_1000.csv
+width_num_shapes_results_random_1000.csv
diff --git a/examples/poisson-line-process/intensity_num_shapes_results_random_1000.csv b/examples/poisson-line-process/intensity_num_shapes_results_random_1000.csv
deleted file mode 100644
index 9f9d9208..00000000
--- a/examples/poisson-line-process/intensity_num_shapes_results_random_1000.csv
+++ /dev/null
@@ -1,41 +0,0 @@
-,width,height,control,treatment,smt_risk_ratio,obs_risk_ratio
-0,1,1,1,2,2.588785046728972,2.8279584525300003
-1,1,1,2,4,3.4461062403300664,3.1711041253654577
-2,1,1,4,8,3.617779108051482,3.477219389925258
-3,1,1,8,16,3.6186812277653675,3.699311115318508
-4,2,2,1,2,3.035917652212002,2.8279584525300003
-5,2,2,2,4,3.54104746789785,3.1711041253654577
-6,2,2,4,8,3.600252617854378,3.477219389925258
-7,2,2,8,16,3.7263838118626995,3.699311115318508
-8,3,3,1,2,3.502541942043722,2.8279584525300003
-9,3,3,2,4,3.5944553305755127,3.1711041253654577
-10,3,3,4,8,4.019120497496367,3.477219389925258
-11,3,3,8,16,3.654514490678642,3.699311115318508
-12,4,4,1,2,3.1137676292813965,2.8279584525300003
-13,4,4,2,4,3.5285135018548877,3.1711041253654577
-14,4,4,4,8,4.156177412650676,3.477219389925258
-15,4,4,8,16,3.729035465206917,3.699311115318508
-16,5,5,1,2,3.66858552631579,2.8279584525300003
-17,5,5,2,4,3.768577673167451,3.1711041253654577
-18,5,5,4,8,3.940844504918471,3.477219389925258
-19,5,5,8,16,3.8751226378071624,3.699311115318508
-20,6,6,1,2,3.693321680868109,2.8279584525300003
-21,6,6,2,4,3.698775461913118,3.1711041253654577
-22,6,6,4,8,3.9219249797415454,3.477219389925258
-23,6,6,8,16,3.9706956619846956,3.699311115318508
-24,7,7,1,2,3.712698412698412,2.8279584525300003
-25,7,7,2,4,3.6271483539974354,3.1711041253654577
-26,7,7,4,8,3.9861973856363226,3.477219389925258
-27,7,7,8,16,3.937031506423999,3.699311115318508
-28,8,8,1,2,3.4956584370373336,2.8279584525300003
-29,8,8,2,4,3.830002289377289,3.1711041253654577
-30,8,8,4,8,3.886060211184468,3.477219389925258
-31,8,8,8,16,4.011038790739813,3.699311115318508
-32,9,9,1,2,3.563344248874334,2.8279584525300003
-33,9,9,2,4,4.00093813120435,3.1711041253654577
-34,9,9,4,8,3.934243330541931,3.477219389925258
-35,9,9,8,16,3.9338491300643366,3.699311115318508
-36,10,10,1,2,3.8274726834113806,2.8279584525300003
-37,10,10,2,4,3.752515553050919,3.1711041253654577
-38,10,10,4,8,4.012820411486204,3.477219389925258
-39,10,10,8,16,4.018089779195966,3.699311115318508
diff --git a/examples/poisson-line-process/width_num_shapes_results_random_1000.csv b/examples/poisson-line-process/width_num_shapes_results_random_1000.csv
deleted file mode 100644
index 057e5087..00000000
--- a/examples/poisson-line-process/width_num_shapes_results_random_1000.csv
+++ /dev/null
@@ -1,154 +0,0 @@
-,control,treatment,intensity,ate,ci_low,ci_high
-0,1,2,0,-7.378642492215885,-16.05805863522675,1.3007736507949783
-1,2,3,0,-2.7096586545196,-12.873331738256546,7.454014429217346
-2,3,4,0,-1.5424126950955266,-14.72145204712045,11.6366266569294
-3,4,5,0,-1.0755143113258985,-17.578716168065007,15.42768754541321
-4,5,6,0,-0.8420651194410849,-20.770151444809144,19.086021205926976
-5,6,7,0,-0.7086655812211902,-24.105926511450654,22.688595349008274
-6,7,8,0,-0.6252908698337563,-27.51508980036924,26.264508060701722
-7,8,9,0,-0.5697077289088008,-30.965916499025454,29.826501041207848
-8,9,10,0,-0.5307995302613315,-34.44235723135286,33.38075817083019
-9,1,2,1,-7.378642492215885,-13.918238652043593,-0.8390463323881754
-10,2,3,1,-2.7096586545196004,-9.802882858188203,4.383565549149003
-11,3,4,1,-1.5424126950955266,-11.120887611821622,8.03606222163057
-12,4,5,1,-1.0755143113258985,-13.708422107725228,11.55739348507343
-13,5,6,1,-0.8420651194410844,-16.74129371640377,15.057163477521602
-14,6,7,1,-0.7086655812211902,-19.972927322591467,18.555596160149086
-15,7,8,1,-0.6252908698337558,-23.3084183840392,22.05783664437169
-16,8,9,1,-0.5697077289088011,-26.70434658441149,25.56493112659389
-17,9,10,1,-0.5307995302613318,-30.138282420998863,29.076683360476196
-18,1,2,2,-7.378642492215885,-16.38113570509901,1.6238507206672423
-19,2,3,2,-2.709658654519597,-11.125018484727406,5.705701175688212
-20,3,4,2,-1.5424126950955284,-10.885940125445654,7.801114735254597
-21,4,5,2,-1.0755143113259003,-12.372903285131162,10.221874662479362
-22,5,6,2,-0.8420651194410844,-14.717521484023713,13.033391245141544
-23,6,7,2,-0.7086655812211902,-17.509736580328116,16.092405417885736
-24,7,8,2,-0.6252908698337549,-20.549295407558162,19.29871366789065
-25,8,9,2,-0.5697077289088028,-23.735302054174575,22.59588659635697
-26,9,10,2,-0.53079953026133,-27.01370867477234,25.95210961424968
-27,1,2,3,-7.378642492215889,-20.043311732745266,5.286026748313489
-28,2,3,3,-2.7096586545196004,-14.257652928983724,8.838335619944523
-29,3,4,3,-1.5424126950955213,-12.86962491637475,9.784799526183704
-30,4,5,3,-1.0755143113259038,-13.130648317524077,10.979619694872273
-31,5,6,3,-0.8420651194410809,-14.473331702582268,12.789201463700103
-32,6,7,3,-0.7086655812211902,-16.53982963725459,15.122498474812215
-33,7,8,3,-0.625290869833762,-19.069141710952145,17.81855997128462
-34,8,9,3,-0.5697077289087957,-21.89167541250101,20.752259954683417
-35,9,10,3,-0.5307995302613335,-24.903950327600462,23.842351267077795
-36,1,2,4,-7.378642492215889,-23.37399716349333,8.616712179061551
-37,2,3,4,-2.7096586545196004,-17.34578972483895,11.92647241579975
-38,3,4,4,-1.5424126950955213,-15.337988269660414,12.253162879469372
-39,4,5,4,-1.0755143113259038,-14.765501080503121,12.614472457851306
-40,5,6,4,-0.8420651194410809,-15.241135322080375,13.557005083198213
-41,6,7,4,-0.7086655812211902,-16.55358588394165,15.136254721499277
-42,7,8,4,-0.625290869833762,-18.491752099772377,17.241170360104846
-43,8,9,4,-0.5697077289087957,-20.87015891353952,19.730743455721928
-44,9,10,4,-0.5307995302613335,-23.550598680286157,22.488999619763497
-45,1,2,5,-7.378642492215889,-26.03512013679405,11.277835152362272
-46,2,3,5,-2.7096586545196004,-19.863479936076885,14.444162627037684
-47,3,4,5,-1.5424126950955355,-17.49695503693016,14.41212964673909
-48,4,5,5,-1.0755143113258896,-16.403114838981367,14.252086216329587
-49,5,6,5,-0.842065119441088,-16.2618523148069,14.577722075924726
-50,6,7,5,-0.7086655812211973,-16.96002852563103,15.542697363188637
-51,7,8,5,-0.6252908698337478,-18.362111153622664,17.111529413955168
-52,8,9,5,-0.5697077289088099,-20.309364253540707,19.169948795723087
-53,9,10,5,-0.5307995302613193,-22.655986177258285,21.594387116735646
-54,1,2,6,-7.378642492215874,-27.948775665159417,13.191490680727668
-55,2,3,6,-2.7096586545196146,-21.671785689375753,16.252468380336524
-56,3,4,6,-1.5424126950955355,-19.070061978940387,15.985236588749302
-57,4,5,6,-1.0755143113258896,-17.635327009715894,15.484298387064129
-58,5,6,6,-0.842065119441088,-17.072003889461286,15.38787365057911
-59,6,7,6,-0.7086655812211973,-17.319524326107995,15.9021931636656
-60,7,8,6,-0.6252908698337336,-18.300276563558967,17.049694823891514
-61,8,9,6,-0.5697077289088099,-19.89101678248352,18.751601324665884
-62,9,10,6,-0.5307995302613335,-21.953616473882505,20.892017413359838
-63,1,2,7,-7.378642492215874,-29.11726513827736,14.359980153845612
-64,2,3,7,-2.709658654519586,-22.750759334589674,17.3314420255505
-65,3,4,7,-1.5424126950955213,-19.97785954523809,16.893034155047047
-66,4,5,7,-1.075514311325918,-18.30608567018305,16.155057047531216
-67,5,6,7,-0.842065119441088,-17.452181764775304,15.768051525893128
-68,6,7,7,-0.7086655812211688,-17.385478541472708,15.96814737903037
-69,7,8,7,-0.6252908698337762,-18.066675373156556,16.816093633489004
-70,8,9,7,-0.5697077289087815,-19.400709487950223,18.26129403013266
-71,9,10,7,-0.5307995302613335,-21.258871069217292,20.197272008694625
-72,1,2,8,-7.378642492215903,-29.589225422873483,14.831940438441677
-73,2,3,8,-2.709658654519586,-23.140754325536193,17.72143701649702
-74,3,4,8,-1.5424126950955213,-20.235252891690607,17.150427501499564
-75,4,5,8,-1.075514311325918,-18.39286072049805,16.241832097846213
-76,5,6,8,-0.842065119441088,-17.341020322637092,15.656890083754917
-77,6,7,8,-0.7086655812211688,-17.069946945430814,15.652615782988477
-78,7,8,8,-0.6252908698337762,-17.565333348018214,16.31475160835066
-79,8,9,8,-0.5697077289087815,-18.749087451671244,17.60967199385368
-80,9,10,8,-0.5307995302613335,-20.495979946318016,19.43438088579535
-81,1,2,9,-7.378642492215931,-29.4627601895898,14.705475205157938
-82,2,3,9,-2.709658654519558,-22.939924975757606,17.52060766671849
-83,3,4,9,-1.5424126950955497,-19.933724348774547,16.848898958583447
-84,4,5,9,-1.0755143113258896,-17.974722820937814,15.823694198286034
-85,5,6,9,-0.842065119441088,-16.802840553585952,15.118710314703776
-86,6,7,9,-0.7086655812211688,-16.425883438085236,15.008552275642899
-87,7,8,9,-0.6252908698338047,-16.844793222974374,15.594211483306765
-88,8,9,9,-0.5697077289087815,-17.986034056166602,16.84661859834904
-89,9,10,9,-0.5307995302613335,-19.71830733123204,18.656708270709373
-90,1,2,10,-7.378642492215874,-28.908492866963968,14.15120788253222
-91,2,3,10,-2.7096586545196146,-22.32918613514414,16.909868826104912
-92,3,4,10,-1.5424126950954928,-19.26530152311119,16.180476132920205
-93,4,5,10,-1.0755143113258896,-17.254314186285,15.103285563633221
-94,5,6,10,-0.842065119441088,-16.049664005016893,14.365533766134718
-95,6,7,10,-0.7086655812212257,-15.67238729850635,14.2550561360639
-96,7,8,10,-0.625290869833691,-16.12682189679697,14.876240157129587
-97,8,9,10,-0.5697077289088384,-17.329661181871984,16.190245724054307
-98,9,10,10,-0.5307995302613335,-19.13419906581737,18.072600005294703
-99,1,2,11,-7.378642492215818,-28.21103739755455,13.453752413122913
-100,2,3,11,-2.7096586545196715,-21.620809935711463,16.20149262667212
-101,3,4,11,-1.5424126950954928,-18.578314574479975,15.49348918428899
-102,4,5,11,-1.0755143113258896,-16.619351549555233,14.468322926903454
-103,5,6,11,-0.842065119441088,-15.504259980146685,13.820129741264509
-104,6,7,11,-0.7086655812212257,-15.252130048314655,13.834798885872203
-105,7,8,11,-0.625290869833691,-15.850915487393195,14.600333747725813
-106,8,9,11,-0.5697077289088384,-17.1947191980762,16.055303740258523
-107,9,10,11,-0.5307995302613335,-19.121187130685996,18.05958807016333
-108,1,2,12,-7.378642492215874,-27.816475373158653,13.059190388726904
-109,2,3,12,-2.7096586545196146,-21.31156409396158,15.89224678492235
-110,3,4,12,-1.5424126950955497,-18.431614755585258,15.346789365394159
-111,4,5,12,-1.0755143113258328,-16.68850845259334,14.537479829941674
-112,5,6,12,-0.8420651194411448,-15.823699617827174,14.139569378944884
-113,6,7,12,-0.7086655812211688,-15.821709691362003,14.404378528919665
-114,7,8,12,-0.6252908698337478,-16.633085203954636,15.38250346428714
-115,8,9,12,-0.5697077289087815,-18.13279022679427,16.993374768976707
-116,9,10,12,-0.5307995302613904,-20.161775051364486,19.100175990841706
-117,1,2,13,-7.378642492215931,-28.324569726565187,13.567284742133324
-118,2,3,13,-2.709658654519558,-22.052605303921496,16.63328799488238
-119,3,4,13,-1.542412695095436,-19.518911869390536,16.434086479199664
-120,4,5,13,-1.0755143113260601,-18.16831866989287,16.01729004724075
-121,5,6,13,-0.8420651194410311,-17.683477997062027,15.999347758179965
-122,6,7,13,-0.708665581221112,-17.98859788870118,16.571266726258955
-123,7,8,13,-0.6252908698337478,-19.000371980537466,17.74979024086997
-124,8,9,13,-0.569707728908952,-20.5992290332332,19.459813575415296
-125,9,10,13,-0.5307995302612198,-22.65487001642896,21.59327095590652
-126,1,2,14,-7.378642492216045,-30.32363663530839,15.5663516508763
-127,2,3,14,-2.709658654519558,-24.41936770763425,19.000050398595135
-128,3,4,14,-1.542412695095436,-22.35681312123677,19.271987731045897
-129,4,5,14,-1.0755143113259464,-21.474715641664602,19.32368701901271
-130,5,6,14,-0.8420651194411448,-21.383064154658655,19.698933915776365
-131,6,7,14,-0.708665581221112,-21.95823357837253,20.540902415930304
-132,7,8,14,-0.6252908698337478,-23.10765371588559,21.857071976218094
-133,8,9,14,-0.5697077289088384,-24.73540473325727,23.595989275439592
-134,9,10,14,-0.5307995302613335,-26.74746326216325,25.685864201640584
-135,1,2,15,-7.378642492216045,-34.13087918226552,19.373594197833427
-136,2,3,15,-2.709658654519558,-28.63698022963797,23.217662920598855
-137,3,4,15,-1.542412695095436,-27.042415426409207,23.957590036218335
-138,4,5,15,-1.0755143113259464,-26.58007596614152,24.42904734348963
-139,5,6,15,-0.8420651194411448,-26.807796364306114,25.123666125423824
-140,6,7,15,-0.708665581221112,-27.580783710990545,26.16345254854832
-141,7,8,15,-0.6252908698337478,-28.812496964083493,27.561915224415998
-142,8,9,15,-0.5697077289088384,-30.430425062116342,29.291009604298665
-143,9,10,15,-0.5307995302613335,-32.36915735084915,31.307558290326483
-144,1,2,16,-7.378642492216045,-39.72917583196795,24.97189084753586
-145,2,3,16,-2.709658654519558,-34.60204449801222,29.182727188973104
-146,3,4,16,-1.542412695095436,-33.39313515126673,30.30830976107586
-147,4,5,16,-1.0755143113259464,-33.255186525977024,31.10415790332513
-148,5,6,16,-0.8420651194410311,-33.717241878906066,32.033111640024
-149,6,7,16,-0.7086655812212257,-34.62845648255586,33.21112532011341
-150,7,8,16,-0.6252908698337478,-35.910796044480776,34.66021430481328
-151,8,9,16,-0.5697077289088384,-37.508216762277584,36.36880130445991
-152,9,10,16,-0.5307995302612198,-39.374044276785185,38.312445216262745
diff --git a/examples/poisson/README.md b/examples/poisson/README.md
index 9d75dfb0..5242ccd5 100644
--- a/examples/poisson/README.md
+++ b/examples/poisson/README.md
@@ -7,5 +7,3 @@ To run this case study:
    (instructions are provided in the project README).
 2. Change directory to `causal_testing/examples/poisson`.
 3. Run the command `python run_causal_tests.py --data_path data.csv --dag_path dag.dot --json_path causal_tests.json`
-
-This should print a series of causal test results and produce two CSV files. `intensity_num_shapes_results_random_1000.csv` corresponds to table 1, and `width_num_shapes_results_random_1000.csv` relates to our findings regarding the relationship of width and `P_u`.


Test patch:
diff --git a/causal_testing/testing/estimators.py b/causal_testing/testing/estimators.py
index 5bfba43e..c401b8fd 100644
--- a/causal_testing/testing/estimators.py
+++ b/causal_testing/testing/estimators.py
@@ -3,11 +3,15 @@
 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 econml.dml import CausalForestDML
+from patsy import dmatrix
+
 from sklearn.ensemble import GradientBoostingRegressor
 from statsmodels.regression.linear_model import RegressionResultsWrapper
 from statsmodels.tools.sm_exceptions import PerfectSeparationError
@@ -44,7 +48,7 @@ def __init__(
         adjustment_set: set,
         outcome: str,
         df: pd.DataFrame = None,
-        effect_modifiers: dict[Variable:Any] = None,
+        effect_modifiers: dict[str:Any] = None,
     ):
         self.treatment = treatment
         self.treatment_value = treatment_value
@@ -54,10 +58,8 @@ def __init__(
         self.df = df
         if effect_modifiers is None:
             self.effect_modifiers = {}
-        elif isinstance(effect_modifiers, (list, set)):
-            self.effect_modifiers = {k.name for k in effect_modifiers}
         elif isinstance(effect_modifiers, dict):
-            self.effect_modifiers = {k.name: v for k, v in effect_modifiers.items()}
+            self.effect_modifiers = effect_modifiers
         else:
             raise ValueError(f"Unsupported type for effect_modifiers {effect_modifiers}. Expected iterable")
         self.modelling_assumptions = []
@@ -102,20 +104,19 @@ def __init__(
         adjustment_set: set,
         outcome: str,
         df: pd.DataFrame = None,
-        effect_modifiers: dict[Variable:Any] = None,
-        intercept: int = 1,
+        effect_modifiers: dict[str:Any] = None,
+        formula: str = None,
     ):
         super().__init__(treatment, treatment_value, control_value, adjustment_set, outcome, df, effect_modifiers)
 
-        for term in self.effect_modifiers:
-            self.adjustment_set.add(term)
-
-        self.product_terms = []
-        self.square_terms = []
-        self.inverse_terms = []
-        self.intercept = intercept
         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
@@ -143,60 +144,62 @@ def _run_logistic_regression(self, data) -> RegressionResultsWrapper:
         logger.debug(reduced_df[necessary_cols])
 
         # 2. Add intercept
-        reduced_df["Intercept"] = self.intercept
+        reduced_df["Intercept"] = 1  # self.intercept
 
         # 3. Estimate the unit difference in outcome caused by unit difference in treatment
         cols = [self.treatment]
         cols += [x for x in self.adjustment_set if x not in cols]
         treatment_and_adjustments_cols = reduced_df[cols + ["Intercept"]]
-        outcome_col = reduced_df[[self.outcome]]
         for col in treatment_and_adjustments_cols:
             if str(treatment_and_adjustments_cols.dtypes[col]) == "object":
                 treatment_and_adjustments_cols = pd.get_dummies(
                     treatment_and_adjustments_cols, columns=[col], drop_first=True
                 )
-        model = sm.Logit(outcome_col, treatment_and_adjustments_cols).fit(disp=0)
+        model = smf.logit(formula=self.formula, data=data).fit(disp=0)
         return model
 
-    def estimate(self, data: pd.DataFrame) -> RegressionResultsWrapper:
+    def estimate(self, data: pd.DataFrame, adjustment_config=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.
 
         """
+        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)
         self.model = model
 
-        x = pd.DataFrame()
+        x = pd.DataFrame(columns=self.df.columns)
+        x["Intercept"] = 1  # self.intercept
         x[self.treatment] = [self.treatment_value, self.control_value]
-        x["Intercept"] = self.intercept
+        for k, v in adjustment_config.items():
+            x[k] = v
         for k, v in self.effect_modifiers.items():
             x[k] = v
-        for t in self.square_terms:
-            x[t + "^2"] = x[t] ** 2
-        for t in self.inverse_terms:
-            x["1/" + t] = 1 / x[t]
-        for a, b in self.product_terms:
-            x[f"{a}*{b}"] = x[a] * x[b]
-
+        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 = x[model.params.index]
         return model.predict(x)
 
-    def estimate_control_treatment(self, bootstrap_size=100) -> tuple[pd.Series, pd.Series]:
+    def estimate_control_treatment(self, bootstrap_size=100, adjustment_config=None) -> 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)).
         """
 
-        y = self.estimate(self.df)
+        y = self.estimate(self.df, adjustment_config=adjustment_config)
 
         try:
             bootstrap_samples = [
-                self.estimate(self.df.sample(len(self.df), replace=True)) for _ in range(bootstrap_size)
+                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:
@@ -205,6 +208,9 @@ def estimate_control_treatment(self, bootstrap_size=100) -> tuple[pd.Series, pd.
                 "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
@@ -217,7 +223,7 @@ def estimate_control_treatment(self, bootstrap_size=100) -> tuple[pd.Series, pd.
 
         return (y.iloc[1], np.array(control)), (y.iloc[0], np.array(treatment))
 
-    def estimate_ate(self, bootstrap_size=100) -> float:
+    def estimate_ate(self, bootstrap_size=100, adjustment_config=None) -> 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
@@ -228,7 +234,7 @@ def estimate_ate(self, bootstrap_size=100) -> float:
         (control_outcome, control_bootstraps), (
             treatment_outcome,
             treatment_bootstraps,
-        ) = self.estimate_control_treatment()
+        ) = 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:
@@ -247,7 +253,7 @@ def estimate_ate(self, bootstrap_size=100) -> float:
 
         return estimate, (ci_low, ci_high)
 
-    def estimate_risk_ratio(self, bootstrap_size=100) -> float:
+    def estimate_risk_ratio(self, bootstrap_size=100, adjustment_config=None) -> 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
@@ -258,14 +264,14 @@ def estimate_risk_ratio(self, bootstrap_size=100) -> float:
         (control_outcome, control_bootstraps), (
             treatment_outcome,
             treatment_bootstraps,
-        ) = self.estimate_control_treatment()
+        ) = 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 * 0.05) / 2)
+        bound = ceil((bootstrap_size * 0.05) / 2)
         ci_low = bootstraps[bound]
         ci_high = bootstraps[bootstrap_size - bound]
 
@@ -302,20 +308,20 @@ def __init__(
         outcome: str,
         df: pd.DataFrame = None,
         effect_modifiers: dict[Variable:Any] = None,
-        product_terms: list[tuple[Variable, Variable]] = None,
-        intercept: int = 1,
+        formula: str = None,
     ):
         super().__init__(treatment, treatment_value, control_value, adjustment_set, outcome, df, effect_modifiers)
 
-        self.product_terms = []
-        self.square_terms = []
-        self.inverse_terms = []
-        self.intercept = intercept
         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)))}"
 
-        if product_terms:
-            for term_a, term_b in product_terms:
-                self.add_product_term_to_df(term_a, term_b)
         for term in self.effect_modifiers:
             self.adjustment_set.add(term)
 
@@ -330,54 +336,6 @@ def add_modelling_assumptions(self):
             "do not need to be linear."
         )
 
-    def add_squared_term_to_df(self, term_to_square: str):
-        """Add a squared term to the linear regression model and df.
-
-        This enables the user to capture curvilinear relationships with a linear regression model, not just straight
-        lines, while automatically adding the modelling assumption imposed by the addition of this term.
-
-        :param term_to_square: The term (column in data and variable in DAG) which is to be squared.
-        """
-        new_term = str(term_to_square) + "^2"
-        self.df[new_term] = self.df[term_to_square] ** 2
-        self.adjustment_set.add(new_term)
-        self.modelling_assumptions += (
-            f"Relationship between {self.treatment} and {self.outcome} varies quadratically" f"with {term_to_square}."
-        )
-        self.square_terms.append(term_to_square)
-
-    def add_inverse_term_to_df(self, term_to_invert: str):
-        """Add an inverse term to the linear regression model and df.
-
-        This enables the user to capture curvilinear relationships with a linear regression model, not just straight
-        lines, while automatically adding the modelling assumption imposed by the addition of this term.
-
-        :param term_to_square: The term (column in data and variable in DAG) which is to be squared.
-        """
-        new_term = "1/" + str(term_to_invert)
-        self.df[new_term] = 1 / self.df[term_to_invert]
-        self.adjustment_set.add(new_term)
-        self.modelling_assumptions += (
-            f"Relationship between {self.treatment} and {self.outcome} varies inversely" f"with {term_to_invert}."
-        )
-        self.inverse_terms.append(term_to_invert)
-
-    def add_product_term_to_df(self, term_a: str, term_b: str):
-        """Add a product term to the linear regression model and df.
-
-        This enables the user to capture interaction between a pair of variables in the model. In other words, while
-        each covariate's contribution to the mean is assumed to be independent of the other covariates, the pair of
-        product terms term_a*term_b a are restricted to vary linearly with each other.
-
-        :param term_a: The first term of the product term.
-        :param term_b: The second term of the product term.
-        """
-        new_term = str(term_a) + "*" + str(term_b)
-        self.df[new_term] = self.df[term_a] * self.df[term_b]
-        self.adjustment_set.add(new_term)
-        self.modelling_assumptions += f"{term_a} and {term_b} vary linearly with each other."
-        self.product_terms.append((term_a, term_b))
-
     def estimate_unit_ate(self) -> float:
         """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.
@@ -402,10 +360,10 @@ def estimate_ate(self) -> tuple[float, list[float, float], float]:
         individuals = pd.DataFrame(1, index=["control", "treated"], columns=model.params.index)
 
         # This is a temporary hack
-        for t in self.square_terms:
-            individuals[t + "^2"] = individuals[t] ** 2
-        for a, b in self.product_terms:
-            individuals[f"{a}*{b}"] = individuals[a] * individuals[b]
+        # for t in self.square_terms:
+        #     individuals[t + "^2"] = individuals[t] ** 2
+        # for a, b in self.product_terms:
+        #     individuals[f"{a}*{b}"] = individuals[a] * individuals[b]
 
         # It is ABSOLUTELY CRITICAL that these go last, otherwise we can't index
         # the effect with "ate = t_test_results.effect[0]"
@@ -430,19 +388,14 @@ def estimate_control_treatment(self, adjustment_config: dict = None) -> tuple[pd
         model = self._run_linear_regression()
         self.model = model
 
-        x = pd.DataFrame()
+        x = pd.DataFrame(columns=self.df.columns)
         x[self.treatment] = [self.treatment_value, self.control_value]
-        x["Intercept"] = self.intercept
+        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
-        for t in self.square_terms:
-            x[t + "^2"] = x[t] ** 2
-        for t in self.inverse_terms:
-            x["1/" + t] = 1 / x[t]
-        for a, b in self.product_terms:
-            x[f"{a}*{b}"] = x[a] * x[b]
+        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)
@@ -488,7 +441,7 @@ def estimate_cates(self) -> tuple[float, list[float, float]]:
         ), f"Must have at least one effect modifier to compute CATE - {self.effect_modifiers}."
         x = pd.DataFrame()
         x[self.treatment] = [self.treatment_value, self.control_value]
-        x["Intercept"] = self.intercept
+        x["Intercept"] = 1  # self.intercept
         for k, v in self.effect_modifiers.items():
             self.adjustment_set.add(k)
             x[k] = v
@@ -520,20 +473,18 @@ def _run_linear_regression(self) -> RegressionResultsWrapper:
         logger.debug(reduced_df[necessary_cols])
 
         # 2. Add intercept
-        reduced_df["Intercept"] = self.intercept
+        reduced_df["Intercept"] = 1  # self.intercept
 
         # 3. Estimate the unit difference in outcome caused by unit difference in treatment
         cols = [self.treatment]
         cols += [x for x in self.adjustment_set if x not in cols]
         treatment_and_adjustments_cols = reduced_df[cols + ["Intercept"]]
-        outcome_col = reduced_df[[self.outcome]]
         for col in treatment_and_adjustments_cols:
             if str(treatment_and_adjustments_cols.dtypes[col]) == "object":
                 treatment_and_adjustments_cols = pd.get_dummies(
                     treatment_and_adjustments_cols, columns=[col], drop_first=True
                 )
-        regression = sm.OLS(outcome_col, treatment_and_adjustments_cols)
-        model = regression.fit()
+        model = smf.ols(formula=self.formula, data=self.df).fit()
         return model
 
     def _get_confidence_intervals(self, model):
diff --git a/examples/covasim_/doubling_beta/causal_test_beta.py b/examples/covasim_/doubling_beta/causal_test_beta.py
index cf6a6b48..32f6f8ad 100644
--- a/examples/covasim_/doubling_beta/causal_test_beta.py
+++ b/examples/covasim_/doubling_beta/causal_test_beta.py
@@ -50,18 +50,18 @@ def doubling_beta_CATE_on_csv(observational_data_path: str, simulate_counterfact
     linear_regression_estimator = LinearRegressionEstimator('beta', 0.032, 0.016,
                                                             {'avg_age', 'contacts'},  # We use custom adjustment set
                                                             'cum_infections',
-                                                            df=past_execution_df)
+                                                            df=past_execution_df,
+                                                            formula="cum_infections ~ beta + np.power(beta, 2) + avg_age + contacts")
 
     # Add squared terms for beta, since it has a quadratic relationship with cumulative infections
-    linear_regression_estimator.add_squared_term_to_df('beta')
     causal_test_result = causal_test_engine.execute_test(linear_regression_estimator, causal_test_case, 'ate')
 
     # Repeat for association estimate (no adjustment)
     no_adjustment_linear_regression_estimator = LinearRegressionEstimator('beta', 0.032, 0.016,
                                                                           set(),
                                                                           'cum_infections',
-                                                                          df=past_execution_df)
-    no_adjustment_linear_regression_estimator.add_squared_term_to_df('beta')
+                                                                          df=past_execution_df,
+                                                                          formula="cum_infections ~ beta + np.power(beta, 2)")
     association_test_result = causal_test_engine.execute_test(no_adjustment_linear_regression_estimator, causal_test_case, 'ate')
 
     # Store results for plotting
@@ -82,8 +82,8 @@ def doubling_beta_CATE_on_csv(observational_data_path: str, simulate_counterfact
         counterfactual_linear_regression_estimator = LinearRegressionEstimator('beta', 0.032, 0.016,
                                                                                {'avg_age', 'contacts'},
                                                                                'cum_infections',
-                                                                               df=counterfactual_past_execution_df)
-        counterfactual_linear_regression_estimator.add_squared_term_to_df('beta')
+                                                                               df=counterfactual_past_execution_df,
+                                                                               formula="cum_infections ~ beta + np.power(beta, 2) + avg_age + contacts")
         counterfactual_causal_test_result = causal_test_engine.execute_test(linear_regression_estimator, causal_test_case, 'ate')
         results_dict['counterfactual'] = {'ate': counterfactual_causal_test_result.test_value.value,
                                           'cis': counterfactual_causal_test_result.confidence_intervals,
diff --git a/examples/poisson-line-process/causal_test_poisson.py b/examples/poisson-line-process/causal_test_poisson.py
index acc992e3..c6c1fa87 100644
--- a/examples/poisson-line-process/causal_test_poisson.py
+++ b/examples/poisson-line-process/causal_test_poisson.py
@@ -117,6 +117,8 @@ def 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]
         estimator = LinearRegressionEstimator(
             treatment=treatment,
             control_value=causal_test_case.control_value,
@@ -124,13 +126,9 @@ def test_intensity_num_shapes(
             adjustment_set=set(),
             outcome=outcome,
             df=data,
-            intercept=0,
             effect_modifiers=causal_test_case.effect_modifier_configuration,
+            formula=f"{outcome} ~ {treatment} + {'+'.join(square_terms + inverse_terms + list([e.name for e in causal_test_case.effect_modifier_configuration]))} -1"
         )
-        for t in square_terms:
-            estimator.add_squared_term_to_df(t)
-        for t in inverse_terms:
-            estimator.add_inverse_term_to_df(t)
 
     # 10. Execute the test
     causal_test_result = causal_test_engine.execute_test(
diff --git a/examples/poisson/run_causal_tests.py b/examples/poisson/run_causal_tests.py
index 93f754d3..9dcdf8e8 100644
--- a/examples/poisson/run_causal_tests.py
+++ b/examples/poisson/run_causal_tests.py
@@ -150,11 +150,7 @@ class MyJsonUtility(JsonUtility):
     def add_modelling_assumptions(self, estimation_model: Estimator):
         # Add squared intensity term as a modelling assumption if intensity is the treatment of the test
         if "intensity" in estimation_model.treatment[0]:
-            estimation_model.add_squared_term_to_df("intensity")
             estimation_model.intercept = 0
-        if isinstance(estimation_model, WidthHeightEstimator):
-            estimation_model.add_product_term_to_df("width", "intensity")
-            estimation_model.add_product_term_to_df("height", "intensity")
 
 
 if __name__ == "__main__":
diff --git a/tests/json_front_tests/test_json_class.py b/tests/json_front_tests/test_json_class.py
index 25b56613..5f3d1b5c 100644
--- a/tests/json_front_tests/test_json_class.py
+++ b/tests/json_front_tests/test_json_class.py
@@ -67,8 +67,6 @@ def test_argparse(self):
 
     def test_setup_modelling_scenario(self):
         self.json_class.setup()
-        print(type(self.json_class.modelling_scenario))
-        print(self.json_class.modelling_scenario)
         self.assertIsInstance(self.json_class.modelling_scenario, Scenario)
 
     def test_setup_causal_specification(self):
diff --git a/tests/testing_tests/test_causal_test_engine.py b/tests/testing_tests/test_causal_test_engine.py
index 6e1c11c4..3e6ba39d 100644
--- a/tests/testing_tests/test_causal_test_engine.py
+++ b/tests/testing_tests/test_causal_test_engine.py
@@ -218,8 +218,8 @@ def test_execute_test_observational_linear_regression_estimator_squared_term(sel
             self.minimal_adjustment_set,
             "C",
             self.causal_test_engine.scenario_execution_data_df,
+            formula=f"C ~ A + {'+'.join(self.minimal_adjustment_set)} + (D ** 2)",
         )
-        estimation_model.add_squared_term_to_df("D")
         causal_test_result = self.causal_test_engine.execute_test(estimation_model, self.causal_test_case)
         self.assertAlmostEqual(round(causal_test_result.test_value.value, 1), 4, delta=1)
 
@@ -240,7 +240,7 @@ def test_execute_observational_causal_forest_estimator_cates(self):
             self.minimal_adjustment_set,
             "C",
             self.causal_test_engine.scenario_execution_data_df,
-            effect_modifiers={Input("M", int): None},
+            effect_modifiers={"M": None},
         )
         causal_test_result = self.causal_test_engine.execute_test(
             estimation_model, self.causal_test_case, estimate_type="cate"
diff --git a/tests/testing_tests/test_causal_test_outcome.py b/tests/testing_tests/test_causal_test_outcome.py
index afa1896e..361edd80 100644
--- a/tests/testing_tests/test_causal_test_outcome.py
+++ b/tests/testing_tests/test_causal_test_outcome.py
@@ -1,5 +1,5 @@
 import unittest
-from causal_testing.testing.causal_test_outcome import ExactValue, SomeEffect
+from causal_testing.testing.causal_test_outcome import ExactValue, SomeEffect, Positive, Negative
 from causal_testing.testing.causal_test_result import CausalTestResult, TestValue
 from causal_testing.testing.estimators import LinearRegressionEstimator
 
@@ -63,6 +63,50 @@ def test_empty_adjustment_set(self):
             ),
         )
 
+    def test_Positive_pass(self):
+        test_value = TestValue(type="ate", value=5.05)
+        ctr = CausalTestResult(
+            estimator=self.estimator,
+            test_value=test_value,
+            confidence_intervals=None,
+            effect_modifier_configuration=None,
+        )
+        ev = Positive()
+        self.assertTrue(ev.apply(ctr))
+
+    def test_Positive_fail(self):
+        test_value = TestValue(type="ate", value=0)
+        ctr = CausalTestResult(
+            estimator=self.estimator,
+            test_value=test_value,
+            confidence_intervals=None,
+            effect_modifier_configuration=None,
+        )
+        ev = Positive()
+        self.assertFalse(ev.apply(ctr))
+
+    def test_Negative_pass(self):
+        test_value = TestValue(type="ate", value=-5.05)
+        ctr = CausalTestResult(
+            estimator=self.estimator,
+            test_value=test_value,
+            confidence_intervals=None,
+            effect_modifier_configuration=None,
+        )
+        ev = Negative()
+        self.assertTrue(ev.apply(ctr))
+
+    def test_Negative_fail(self):
+        test_value = TestValue(type="ate", value=0)
+        ctr = CausalTestResult(
+            estimator=self.estimator,
+            test_value=test_value,
+            confidence_intervals=None,
+            effect_modifier_configuration=None,
+        )
+        ev = Negative()
+        self.assertFalse(ev.apply(ctr))
+
     def test_exactValue_pass(self):
         test_value = TestValue(type="ate", value=5.05)
         ctr = CausalTestResult(
diff --git a/tests/testing_tests/test_estimators.py b/tests/testing_tests/test_estimators.py
index 4d1bed9b..719f2200 100644
--- a/tests/testing_tests/test_estimators.py
+++ b/tests/testing_tests/test_estimators.py
@@ -78,38 +78,75 @@ class TestLogisticRegressionEstimator(unittest.TestCase):
     def setUpClass(cls) -> None:
         cls.scarf_df = pd.DataFrame(
             [
-                {"length_in": 55, "completed": 1},
-                {"length_in": 55, "completed": 1},
-                {"length_in": 55, "completed": 1},
-                {"length_in": 60, "completed": 1},
-                {"length_in": 60, "completed": 0},
-                {"length_in": 70, "completed": 1},
-                {"length_in": 70, "completed": 0},
-                {"length_in": 82, "completed": 1},
-                {"length_in": 82, "completed": 0},
-                {"length_in": 82, "completed": 0},
-                {"length_in": 82, "completed": 0},
+                {"length_in": 55, "large_gauge": 1, "color": "orange", "completed": 1},
+                {"length_in": 55, "large_gauge": 0, "color": "orange", "completed": 1},
+                {"length_in": 55, "large_gauge": 0, "color": "brown", "completed": 1},
+                {"length_in": 60, "large_gauge": 0, "color": "brown", "completed": 1},
+                {"length_in": 60, "large_gauge": 0, "color": "grey", "completed": 0},
+                {"length_in": 70, "large_gauge": 0, "color": "grey", "completed": 1},
+                {"length_in": 70, "large_gauge": 0, "color": "orange", "completed": 0},
+                {"length_in": 82, "large_gauge": 1, "color": "grey", "completed": 1},
+                {"length_in": 82, "large_gauge": 0, "color": "brown", "completed": 0},
+                {"length_in": 82, "large_gauge": 0, "color": "orange", "completed": 0},
+                {"length_in": 82, "large_gauge": 1, "color": "brown", "completed": 0},
             ]
         )
 
     def test_ate(self):
-        df = self.scarf_df
+        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
+        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
+        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):
     """
@@ -158,6 +195,7 @@ def test_program_11_2(self):
         model = linear_regression_estimator._run_linear_regression()
         ate, _ = linear_regression_estimator.estimate_unit_ate()
 
+        print(model.summary())
         self.assertEqual(round(model.params["Intercept"] + 90 * model.params["treatments"], 1), 216.9)
 
         # Increasing treatments from 90 to 100 should be the same as 10 times the unit ATE
@@ -166,13 +204,17 @@ def test_program_11_2(self):
     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", 100, 90, set(), "outcomes", df)
-        linear_regression_estimator.add_squared_term_to_df("treatments")
+        linear_regression_estimator = LinearRegressionEstimator(
+            "treatments", 100, 90, set(), "outcomes", df, formula="outcomes ~ treatments + np.power(treatments, 2)"
+        )
         model = linear_regression_estimator._run_linear_regression()
         ate, _ = linear_regression_estimator.estimate_unit_ate()
         self.assertEqual(
             round(
-                model.params["Intercept"] + 90 * model.params["treatments"] + 90 * 90 * model.params["treatments^2"], 1
+                model.params["Intercept"]
+                + 90 * model.params["treatments"]
+                + 90 * 90 * model.params["np.power(treatments, 2)"],
+                1,
             ),
             197.1,
         )
@@ -198,17 +240,30 @@ def test_program_15_1A(self):
             "smokeintensity",
             "smokeyrs",
         }
-        linear_regression_estimator = LinearRegressionEstimator("qsmk", 1, 0, covariates, "wt82_71", df)
-        terms_to_square = ["age", "wt71", "smokeintensity", "smokeyrs"]
-        terms_to_product = [("qsmk", "smokeintensity")]
-        for term_to_square in terms_to_square:
-            linear_regression_estimator.add_squared_term_to_df(term_to_square)
-        for term_a, term_b in terms_to_product:
-            linear_regression_estimator.add_product_term_to_df(term_a, term_b)
+        linear_regression_estimator = LinearRegressionEstimator(
+            "qsmk",
+            1,
+            0,
+            covariates,
+            "wt82_71",
+            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) +
+                             (qsmk * smokeintensity)""",
+        )
+        # terms_to_square = ["age", "wt71", "smokeintensity", "smokeyrs"]
+        # terms_to_product = [("qsmk", "smokeintensity")]
+        # for term_to_square in terms_to_square:
+        # 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)
+        self.assertEqual(round(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)
@@ -230,10 +285,17 @@ def test_program_15_no_interaction(self):
             "smokeintensity",
             "smokeyrs",
         }
-        linear_regression_estimator = LinearRegressionEstimator("qsmk", 1, 0, covariates, "wt82_71", df)
-        terms_to_square = ["age", "wt71", "smokeintensity", "smokeyrs"]
-        for term_to_square in terms_to_square:
-            linear_regression_estimator.add_squared_term_to_df(term_to_square)
+        linear_regression_estimator = LinearRegressionEstimator(
+            "qsmk",
+            1,
+            0,
+            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)",
+        )
+        # terms_to_square = ["age", "wt71", "smokeintensity", "smokeyrs"]
+        # for term_to_square in terms_to_square:
         ate, [ci_low, ci_high] = linear_regression_estimator.estimate_unit_ate()
         self.assertEqual(round(ate, 1), 3.5)
         self.assertEqual([round(ci_low, 1), round(ci_high, 1)], [2.6, 4.3])
@@ -258,10 +320,17 @@ def test_program_15_no_interaction_ate(self):
             "smokeintensity",
             "smokeyrs",
         }
-        linear_regression_estimator = LinearRegressionEstimator("qsmk", 1, 0, covariates, "wt82_71", df)
-        terms_to_square = ["age", "wt71", "smokeintensity", "smokeyrs"]
-        for term_to_square in terms_to_square:
-            linear_regression_estimator.add_squared_term_to_df(term_to_square)
+        linear_regression_estimator = LinearRegressionEstimator(
+            "qsmk",
+            1,
+            0,
+            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)",
+        )
+        # terms_to_square = ["age", "wt71", "smokeintensity", "smokeyrs"]
+        # for term_to_square in terms_to_square:
         ate, [ci_low, ci_high] = linear_regression_estimator.estimate_ate()
         self.assertEqual(round(ate, 1), 3.5)
         self.assertEqual([round(ci_low, 1), round(ci_high, 1)], [2.6, 4.3])
@@ -286,10 +355,17 @@ def test_program_15_no_interaction_ate_calculated(self):
             "smokeintensity",
             "smokeyrs",
         }
-        linear_regression_estimator = LinearRegressionEstimator("qsmk", 1, 0, covariates, "wt82_71", df)
-        terms_to_square = ["age", "wt71", "smokeintensity", "smokeyrs"]
-        for term_to_square in terms_to_square:
-            linear_regression_estimator.add_squared_term_to_df(term_to_square)
+        linear_regression_estimator = LinearRegressionEstimator(
+            "qsmk",
+            1,
+            0,
+            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)",
+        )
+        # terms_to_square = ["age", "wt71", "smokeintensity", "smokeyrs"]
+        # for term_to_square in terms_to_square:
         ate, [ci_low, ci_high] = linear_regression_estimator.estimate_ate_calculated(
             {k: self.nhefs_df.mean()[k] for k in covariates}
         )
@@ -328,9 +404,7 @@ def test_program_15_ate(self):
             "smokeintensity",
             "smokeyrs",
         }
-        causal_forest = CausalForestEstimator(
-            "qsmk", 1, 0, covariates, "wt82_71", df, {Input("smokeintensity", int): 40}
-        )
+        causal_forest = CausalForestEstimator("qsmk", 1, 0, covariates, "wt82_71", df, {"smokeintensity": 40})
         ate, _ = causal_forest.estimate_ate()
         self.assertGreater(round(ate, 1), 2.5)
         self.assertLess(round(ate, 1), 4.5)
@@ -356,7 +430,7 @@ def test_program_15_cate(self):
             "smokeyrs",
         }
         causal_forest = CausalForestEstimator(
-            "qsmk", 1, 0, covariates, "wt82_71", smoking_intensity_5_and_40_df, {Input("smokeintensity", int): 40}
+            "qsmk", 1, 0, covariates, "wt82_71", smoking_intensity_5_and_40_df, {"smokeintensity": 40}
         )
         cates_df, _ = causal_forest.estimate_cates()
         self.assertGreater(cates_df["cate"].mean(), 0)
@@ -376,7 +450,7 @@ def test_X1_effect(self):
         """When we fix the value of X2 to 0, the effect of X1 on Y should become ~2 (because X2 terms are cancelled)."""
         x2 = Input("X2", float)
         lr_model = LinearRegressionEstimator(
-            "X1", 1, 0, {"X2"}, "Y", effect_modifiers={x2: 0}, product_terms=[("X1", "X2")], df=self.df
+            "X1", 1, 0, {"X2"}, "Y", effect_modifiers={x2.name: 0}, formula="Y ~ X1 + X2 + (X1 * X2)", df=self.df
         )
         test_results = lr_model.estimate_ate()
         ate = test_results[0]