Raise if intervention variable is used in the wrong place

Author: cthoyt

src/y0/graph.py

@@ -14,7 +14,7 @@
 from networkx.classes.reportviews import NodeView
 from networkx.utils import open_file
 
-from .dsl import CounterfactualVariable, Variable, vmap_adj, vmap_pairs
+from .dsl import CounterfactualVariable, Intervention, Variable, vmap_adj, vmap_pairs
 
 __all__ = [
     "NxMixedGraph",
@@ -553,7 +553,10 @@ def _get_latex(node) -> str:
 
 
 def _ensure_set(vertices: Union[Variable, Iterable[Variable]]) -> set[Variable]:
-    return {vertices} if isinstance(vertices, Variable) else set(vertices)
+    rv = {vertices} if isinstance(vertices, Variable) else set(vertices)
+    if any(isinstance(v, Intervention) for v in rv):
+        raise TypeError("can not use interventions here")
+    return rv
 
 
 class NoAnankeError(TypeError):

tests/test_graph.py

@@ -150,6 +150,10 @@ def test_intervention(self):
             undirected=[("Z", "Y")],
         )
         self.assertEqual(intervened_graph, graph.remove_in_edges({X}))
+        self.assertEqual(intervened_graph, graph.remove_in_edges(X))
+
+        with self.assertRaises(TypeError):
+            self.assertEqual(intervened_graph, graph.remove_in_edges({-X}))
 
     def test_remove_nodes_from(self):
         """Test generating a new graph without the given nodes."""