Removed scenario from CausalDAG

Author: jmafoster1

causal_testing/main.py

@@ -313,7 +313,10 @@ def create_causal_test(self, test: dict, base_test: BaseTestCase) -> CausalTestC
             base_test_case=base_test,
             treatment_value=test.get("treatment_value"),
             control_value=test.get("control_value"),
-            adjustment_set=test.get("adjustment_set", self.causal_specification.causal_dag.identification(base_test)),
+            adjustment_set=test.get(
+                "adjustment_set",
+                self.causal_specification.causal_dag.identification(base_test, self.scenario.hidden_variables()),
+            ),
             df=filtered_df,
             effect_modifiers=None,
             formula=test.get("formula"),

causal_testing/specification/causal_dag.py

@@ -10,8 +10,7 @@
 
 from causal_testing.testing.base_test_case import BaseTestCase
 
-from .scenario import Scenario
-from .variable import Output
+from .variable import Variable
 
 Node = Union[str, int]  # Node type hint: A node is a string or an int
 
@@ -489,16 +488,7 @@ def get_backdoor_graph(self, treatments: list[str]) -> CausalDAG:
         backdoor_graph.add_edges_from(filter(lambda x: x not in outgoing_edges, self.edges))
         return backdoor_graph
 
-    @staticmethod
-    def remove_hidden_adjustment_sets(minimal_adjustment_sets: list[str], scenario: Scenario):
-        """Remove variables labelled as hidden from adjustment set(s)
-
-        :param minimal_adjustment_sets: list of minimal adjustment set(s) to have hidden variables removed from
-        :param scenario: The modelling scenario which informs the variables that are hidden
-        """
-        return [adj for adj in minimal_adjustment_sets if all(not scenario.variables.get(x).hidden for x in adj)]
-
-    def identification(self, base_test_case: BaseTestCase, scenario: Scenario = None):
+    def identification(self, base_test_case: BaseTestCase, avoid_variables: set[Variable] = None):
         """Identify and return the minimum adjustment set
 
         :param base_test_case: A base test case instance containing the outcome_variable and the
@@ -523,8 +513,10 @@ def identification(self, base_test_case: BaseTestCase, scenario: Scenario = None
         else:
             raise ValueError("Causal effect should be 'total' or 'direct'")
 
-        if scenario is not None:
-            minimal_adjustment_sets = self.remove_hidden_adjustment_sets(minimal_adjustment_sets, scenario)
+        if avoid_variables is not None:
+            minimal_adjustment_sets = [
+                adj for adj in minimal_adjustment_sets if not {x.name for x in avoid_variables}.intersection(adj)
+            ]
 
         minimal_adjustment_set = min(minimal_adjustment_sets, key=len, default=set())
         return set(minimal_adjustment_set)

causal_testing/specification/scenario.py

@@ -31,3 +31,11 @@ def __init__(self, variables: Iterable[Variable], constraints: set[str] = None):
             self.constraints = set(constraints)
         else:
             self.constraints = set()
+
+    def hidden_variables(self) -> set[Variable]:
+        """Get the set of hidden variables
+
+        :return The variables marked as hidden.
+        :rtype: {Variable}
+        """
+        return {v for v in self.variables.values() if v.hidden}

causal_testing/surrogate/causal_surrogate_assisted.py

@@ -133,7 +133,9 @@ def generate_surrogates(
                 to_var = specification.scenario.variables.get(v)
                 base_test_case = BaseTestCase(from_var, to_var)
 
-                minimal_adjustment_set = specification.causal_dag.identification(base_test_case, specification.scenario)
+                minimal_adjustment_set = specification.causal_dag.identification(
+                    base_test_case, specification.scenario.hidden_variables()
+                )
 
                 surrogate = CubicSplineRegressionEstimator(
                     base_test_case,

tests/specification_tests/test_causal_dag.py

@@ -1,6 +1,7 @@
 import unittest
 import os
-import shutil, tempfile
+import shutil
+import tempfile
 import networkx as nx
 from causal_testing.specification.causal_dag import CausalDAG, close_separator, list_all_min_sep
 from causal_testing.specification.scenario import Scenario
@@ -420,10 +421,10 @@ def test_hidden_varaible_adjustment_sets(self):
         m = Input("M", int)
 
         scenario = Scenario(variables={z, x, m})
-        adjustment_sets = causal_dag.identification(BaseTestCase(x, m), scenario)
+        adjustment_sets = causal_dag.identification(BaseTestCase(x, m), scenario.hidden_variables())
 
         z.hidden = True
-        adjustment_sets_with_hidden = causal_dag.identification(BaseTestCase(x, m), scenario)
+        adjustment_sets_with_hidden = causal_dag.identification(BaseTestCase(x, m), scenario.hidden_variables())
 
         self.assertNotEqual(adjustment_sets, adjustment_sets_with_hidden)