Merge pull request #233 from CITCOM-project/adjustment_set_formula_check

Remove Redundant Adjustment Set Calc when Formula

Author: christopher-wild

causal_testing/json_front/json_class.py

@@ -22,7 +22,7 @@
 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
+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
 
@@ -307,20 +307,29 @@ def _setup_test(self, causal_test_case: CausalTestCase, test: Mapping) -> Estima
         :returns:
                 - estimation_model - Estimator instance for the test being run
         """
-        minimal_adjustment_set = self.causal_specification.causal_dag.identification(causal_test_case.base_test_case)
-        treatment_var = causal_test_case.treatment_variable
-        minimal_adjustment_set = minimal_adjustment_set - {treatment_var}
-        estimator_kwargs = {
-            "treatment": treatment_var.name,
-            "treatment_value": causal_test_case.treatment_value,
-            "control_value": causal_test_case.control_value,
-            "adjustment_set": minimal_adjustment_set,
-            "outcome": causal_test_case.outcome_variable.name,
-            "effect_modifiers": causal_test_case.effect_modifier_configuration,
-            "alpha": test["alpha"] if "alpha" in test else 0.05,
-        }
+        estimator_kwargs = {}
         if "formula" in test:
+            if test["estimator"] != (LinearRegressionEstimator or LogisticRegressionEstimator):
+                raise TypeError(
+                    "Currently only LinearRegressionEstimator and LogisticRegressionEstimator supports the use of "
+                    "formulas"
+                )
             estimator_kwargs["formula"] = test["formula"]
+            estimator_kwargs["adjustment_set"] = None
+        else:
+            minimal_adjustment_set = self.causal_specification.causal_dag.identification(
+                causal_test_case.base_test_case
+            )
+            minimal_adjustment_set = minimal_adjustment_set - {causal_test_case.treatment_variable}
+            estimator_kwargs["adjustment_set"] = minimal_adjustment_set
+
+        estimator_kwargs["treatment"] = causal_test_case.treatment_variable.name
+        estimator_kwargs["treatment_value"] = causal_test_case.treatment_value
+        estimator_kwargs["control_value"] = causal_test_case.control_value
+        estimator_kwargs["outcome"] = causal_test_case.outcome_variable.name
+        estimator_kwargs["effect_modifiers"] = causal_test_case.effect_modifier_configuration
+        estimator_kwargs["alpha"] = test["alpha"] if "alpha" in test else 0.05
+
         estimation_model = test["estimator"](**estimator_kwargs)
         return estimation_model
 

causal_testing/testing/estimators.py

@@ -334,7 +334,7 @@ def estimate_coefficient(self) -> float:
             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+'  ')}"
+        ), 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)
         if str(self.df.dtypes[self.treatment]) != "object":

tests/json_front_tests/test_json_class.py

@@ -4,7 +4,7 @@
 import scipy
 import os
 
-from causal_testing.testing.estimators import LinearRegressionEstimator
+from causal_testing.testing.estimators import LinearRegressionEstimator, CausalForestEstimator
 from causal_testing.testing.causal_test_outcome import NoEffect, Positive
 from tests.test_helpers import remove_temp_dir_if_existent
 from causal_testing.json_front.json_class import JsonUtility, CausalVariables
@@ -291,6 +291,31 @@ def test_no_data_provided(self):
         with self.assertRaises(ValueError):
             json_class.setup(self.scenario)
 
+    def test_estimator_formula_type_check(self):
+        example_test = {
+            "tests": [
+                {
+                    "name": "test1",
+                    "mutations": {"test_input": "Increase"},
+                    "estimator": "CausalForestEstimator",
+                    "estimate_type": "ate",
+                    "effect_modifiers": [],
+                    "expected_effect": {"test_output": "Positive"},
+                    "skip": False,
+                    "formula": "test_output ~ test_input",
+                }
+            ]
+        }
+        self.json_class.test_plan = example_test
+        effects = {"Positive": Positive()}
+        mutates = {
+            "Increase": lambda x: self.json_class.scenario.treatment_variables[x].z3
+                                  > self.json_class.scenario.variables[x].z3
+        }
+        estimators = {"CausalForestEstimator": CausalForestEstimator}
+        with self.assertRaises(TypeError):
+            self.json_class.run_json_tests(effects=effects, mutates=mutates, estimators=estimators, f_flag=False)
+
     def tearDown(self) -> None:
         remove_temp_dir_if_existent()
         if os.path.exists("temp_out.txt"):