Give functions better names

Co-Authored-By: Jeremy Zucker <[email protected]>

Author: cthoyt

src/y0/algorithm/identify/id_c.py

@@ -49,8 +49,7 @@ def rule_2_of_do_calculus_applies(identification: Identification, condition: Var
     treatments = identification.treatments
     conditions = treatments | (identification.conditions - {condition})
 
-    # TODO give a better name
-    graph_mod = graph.intervene(treatments).remove_outgoing_edges_from([condition])
+    graph_mod = graph.remove_in_edges(treatments).remove_out_edges(condition)
 
     judgements = [
         are_d_separated(graph_mod, outcome, condition, conditions=conditions)

src/y0/algorithm/identify/id_std.py

@@ -31,7 +31,7 @@ def identify(identification: Identification) -> Expression:
         return identify(line_2(identification))
 
     # line 3
-    intervened_graph = graph.intervene(treatments)
+    intervened_graph = graph.remove_in_edges(treatments)
     no_effect_on_outcome = (vertices - treatments) - intervened_graph.ancestors_inclusive(outcomes)
     if no_effect_on_outcome:
         return identify(line_3(identification))
@@ -142,7 +142,7 @@ def line_3(identification: Identification) -> Identification:
     graph = identification.graph
     vertices = set(graph.nodes())
 
-    intervened_graph = graph.intervene(treatments)
+    intervened_graph = graph.remove_in_edges(treatments)
     no_effect_on_outcome = (vertices - treatments) - intervened_graph.ancestors_inclusive(outcomes)
     if not no_effect_on_outcome:
         raise ValueError(

src/y0/graph.py

@@ -354,39 +354,39 @@ def subgraph(self, vertices: Collection[Variable]) -> NxMixedGraph:
             undirected=_include_adjacent(self.undirected, vertices),
         )
 
-    def intervene(self, vertices: Collection[Variable]) -> NxMixedGraph:
+    def remove_in_edges(self, vertices: Union[Variable, Iterable[Variable]]) -> NxMixedGraph:
         """Return a mutilated graph given a set of interventions.
 
         :param vertices: a subset of nodes from which to remove incoming edges
         :returns: A NxMixedGraph subgraph
         """
-        vertices = set(vertices)
+        vertices = _ensure_set(vertices)
         return self.from_edges(
             nodes=vertices,
             directed=_exclude_target(self.directed, vertices),
             undirected=_exclude_adjacent(self.undirected, vertices),
         )
 
-    def remove_nodes_from(self, vertices: Collection[Variable]) -> NxMixedGraph:
+    def remove_nodes_from(self, vertices: Union[Variable, Iterable[Variable]]) -> NxMixedGraph:
         """Return a subgraph that does not contain any of the specified vertices.
 
         :param vertices: a set of nodes to remove from graph
         :returns: A NxMixedGraph subgraph
         """
-        vertices = set(vertices)
+        vertices = _ensure_set(vertices)
         return self.from_edges(
             nodes=self.nodes() - vertices,
             directed=_exclude_adjacent(self.directed, vertices),
             undirected=_exclude_adjacent(self.undirected, vertices),
         )
 
-    def remove_outgoing_edges_from(self, vertices: Collection[Variable]) -> NxMixedGraph:
+    def remove_out_edges(self, vertices: Union[Variable, Iterable[Variable]]) -> NxMixedGraph:
         """Return a subgraph that does not have any outgoing edges from any of the given vertices.
 
         :param vertices: a set of nodes whose outgoing edges get removed from the graph
         :returns: NxMixedGraph subgraph
         """
-        vertices = set(vertices)
+        vertices = _ensure_set(vertices)
         return self.from_edges(
             nodes=self.nodes(),
             directed=_exclude_source(self.directed, vertices),
@@ -552,5 +552,9 @@ def _get_latex(node) -> str:
     raise TypeError
 
 
+def _ensure_set(vertices: Union[Variable, Iterable[Variable]]) -> set[Variable]:
+    return {vertices} if isinstance(vertices, Variable) else set(vertices)
+
+
 class NoAnankeError(TypeError):
     """Thrown when an :mod:`ananke` graph was used but a y0 NxMixedGraph should have been used."""

tests/test_graph.py

@@ -143,13 +143,13 @@ def test_intervention(self):
             directed=[("X", "Y"), ("Z", "X")],
             undirected=[("X", "Z"), ("X", "Y"), ("Y", "Z")],
         )
-        self.assertEqual(graph, graph.intervene(set()))
+        self.assertEqual(graph, graph.remove_in_edges(set()))
 
         intervened_graph = NxMixedGraph.from_str_edges(
             directed=[("X", "Y")],
             undirected=[("Z", "Y")],
         )
-        self.assertEqual(intervened_graph, graph.intervene({X}))
+        self.assertEqual(intervened_graph, graph.remove_in_edges({X}))
 
     def test_remove_nodes_from(self):
         """Test generating a new graph without the given nodes."""
@@ -165,14 +165,14 @@ def test_remove_nodes_from(self):
     def test_remove_outgoing_edges_from(self):
         """Test generating a new graph without the outgoing edgs from the given nodes."""
         graph = NxMixedGraph.from_str_edges(directed=[("X", "Y")])
-        self.assertEqual(graph, graph.remove_outgoing_edges_from(set()))
+        self.assertEqual(graph, graph.remove_out_edges(set()))
 
         graph = NxMixedGraph.from_str_edges(undirected=[("X", "Y")])
-        self.assertEqual(graph, graph.remove_outgoing_edges_from(set()))
+        self.assertEqual(graph, graph.remove_out_edges(set()))
 
         graph = NxMixedGraph.from_str_edges(directed=[("W", "X"), ("X", "Y"), ("Y", "Z")])
         expected = NxMixedGraph.from_str_edges(directed=[("W", "X"), ("Y", "Z")])
-        self.assertEqual(expected, graph.remove_outgoing_edges_from({X}))
+        self.assertEqual(expected, graph.remove_out_edges({X}))
 
     def test_ancestors_inclusive(self):
         """Test getting ancestors, inclusive."""