Improved categorical treatment handling

Author: jmafoster1

causal_testing/testing/causal_test_outcome.py

@@ -41,18 +41,21 @@ def apply(self, res: CausalTestResult) -> bool:
 class NoEffect(CausalTestOutcome):
     """An extension of TestOutcome representing that the expected causal effect should be zero."""
 
-    def apply(self, res: CausalTestResult, atol: float = 1e-10) -> bool:
+    def __init__(self, atol: float = 1e-10):
+        self.atol = atol
+
+    def apply(self, res: CausalTestResult) -> bool:
         if res.test_value.type == "ate":
-            return (res.ci_low() < 0 < res.ci_high()) or (abs(res.test_value.value) < atol)
+            return (res.ci_low() < 0 < res.ci_high()) or (abs(res.test_value.value) < self.atol)
         if res.test_value.type == "coefficient":
             ci_low = res.ci_low() if isinstance(res.ci_low(), Iterable) else [res.ci_low()]
             ci_high = res.ci_high() if isinstance(res.ci_high(), Iterable) else [res.ci_high()]
             value = res.test_value.value if isinstance(res.ci_high(), Iterable) else [res.test_value.value]
             return all(ci_low < 0 < ci_high for ci_low, ci_high in zip(ci_low, ci_high)) or all(
-                abs(v) < 1e-10 for v in value
+                abs(v) < self.atol for v in value
             )
         if res.test_value.type == "risk_ratio":
-            return (res.ci_low() < 1 < res.ci_high()) or np.isclose(res.test_value.value, 1.0, atol=atol)
+            return (res.ci_low() < 1 < res.ci_high()) or np.isclose(res.test_value.value, 1.0, atol=self.atol)
         raise ValueError(f"Test Value type {res.test_value.type} is not valid for this TestOutcome")
 
 

causal_testing/testing/estimators.py

@@ -339,13 +339,8 @@ def estimate_unit_ate(self) -> float:
         print(model.conf_int())
         treatment = [self.treatment]
         if str(self.df.dtypes[self.treatment]) == "object":
-            reference = min(self.df[self.treatment])
-            treatment = [
-                x.replace("[", "[T.")
-                for x in dmatrix(
-                    f"{self.treatment}-1", self.df.query(f"{self.treatment} != '{reference}'"), return_type="dataframe"
-                ).columns
-            ]
+            design_info = dmatrix(self.formula.split("~")[1], self.df).design_info
+            treatment = design_info.column_names[design_info.term_name_slices[self.treatment]]
         assert set(treatment).issubset(
             model.params.index.tolist()
         ), f"{treatment} not in\n{'  '+str(model.params.index).replace(newline, newline+'  ')}"

examples/poisson/example_run_causal_tests.py

@@ -163,7 +163,9 @@ def test_run_causal_tests():
     )  # Set the path to the data.csv, dag.dot and causal_tests.json file
 
     # Load the Causal Variables into the JsonUtility class ready to be used in the tests
-    json_utility.setup(scenario=modelling_scenario)  # Sets up all the necessary parts of the json_class needed to execute tests
+    json_utility.setup(
+        scenario=modelling_scenario
+    )  # Sets up all the necessary parts of the json_class needed to execute tests
 
     json_utility.run_json_tests(effects=effects, mutates=mutates, estimators=estimators, f_flag=False)
 
@@ -176,6 +178,8 @@ def test_run_causal_tests():
     )  # Set the path to the data.csv, dag.dot and causal_tests.json file
 
     # Load the Causal Variables into the JsonUtility class ready to be used in the tests
-    json_utility.setup(scenario=modelling_scenario)  # Sets up all the necessary parts of the json_class needed to execute tests
+    json_utility.setup(
+        scenario=modelling_scenario
+    )  # Sets up all the necessary parts of the json_class needed to execute tests
 
     json_utility.run_json_tests(effects=effects, mutates=mutates, estimators=estimators, f_flag=args.f)

tests/json_front_tests/test_json_class.py

@@ -48,8 +48,6 @@ def test_setting_no_path(self):
         json_class.set_paths(self.json_path, self.dag_path, None)
         self.assertEqual(json_class.input_paths.data_paths, [])  # Needs to be list of Paths
 
-
-
     def test_setting_paths(self):
         self.assertEqual(self.json_class.input_paths.json_path, Path(self.json_path))
         self.assertEqual(self.json_class.input_paths.dag_path, Path(self.dag_path))