Improve testing of do, and fix various quirks in implementation (#81)

* improve do test

* don't add edges that include the do node

* restructure do code, make it remove predecessors of targets of do

* add test that edges are correct after do

* make code more readable, upset ruff

* check for label, not node object

* noqa

* test graph in #80

* further simplify do function

* pop not del

* Apply suggestions from code review

Co-authored-by: Will Graham <32364977+willGraham01@users.noreply.github.com>

* update docstring, variable names, ruff

* move test_algorithms.py to its own folder

* make graph functions return tuples

* make algorithm functions return tuples

* docstring

* tuple ... and other fixes

* mypy

* ruff

* tuples in tests

* more ruff

---------

Co-authored-by: Will Graham <32364977+willGraham01@users.noreply.github.com>

Author: mscroggs

src/causalprog/algorithms/do.py

@@ -2,7 +2,55 @@
 
 from copy import deepcopy
 
-from causalprog.graph import Graph
+from causalprog.graph import Graph, Node
+
+
+def get_included_excluded_successors(
+    graph: Graph, node_list: dict[str, Node], successors_of: str
+) -> tuple[tuple[str, ...], tuple[str, ...]]:
+    """
+    Split successors of a node into nodes included and not included in a list.
+
+    Split the successorts of a node into a list of nodes that are included in
+    the input node list and a list of nodes that are not in the list.
+
+    Args:
+        graph: The graph
+        node_list: A dictionary of nodes, indexed by label
+        successors_of: The node to check the successors of
+
+    Returns:
+        Lists of included and excluded nodes
+
+    """
+    included = []
+    excluded = []
+    for n in graph.successors[graph.get_node(successors_of)]:
+        if n.label in node_list:
+            included.append(n)
+        else:
+            excluded.append(n)
+    return tuple(included), tuple(excluded)
+
+
+def removable_nodes(graph: Graph, nodes: dict[str, Node]) -> tuple[str, ...]:
+    """
+    Generate list of nodes that can be removed from the graph.
+
+    Args:
+        graph: The graph
+        nodes: A dictionary of nodes, indexed by label
+
+    Returns:
+        List of labels of removable nodes
+
+    """
+    removable: list[str] = []
+    for n in nodes:
+        included, excluded = get_included_excluded_successors(graph, nodes, n)
+        if len(excluded) > 0 and len(included) == 0:
+            removable.append(n)
+    return tuple(removable)
 
 
 def do(graph: Graph, node: str, value: float, label: str | None = None) -> Graph:
@@ -22,46 +70,47 @@ def do(graph: Graph, node: str, value: float, label: str | None = None) -> Graph
     if label is None:
         label = f"{graph.label}|do({node}={value})"
 
-    old_g = graph._graph  # noqa: SLF001
-    g = deepcopy(old_g)
-
-    nodes_by_label = {n.label: n for n in g.nodes}
-    g.remove_node(nodes_by_label[node])
+    nodes = {n.label: deepcopy(n) for n in graph.nodes if n.label != node}
 
-    new_nodes = {}
     # Search through the old graph, identifying nodes that had parameters which were
     # defined by the node being fixed in the DO operation.
     # We recreate these nodes, but replace each such parameter we encounter with
     # a constant parameter equal that takes the fixed value given as an input.
-    for original_node in old_g.nodes:
-        new_n = None
-        for parameter_name, parameter_target_node in original_node.parameters.items():
-            if parameter_target_node == node:
-                # If this parameter in the original_node was determined by the node we
-                # are fixing with DO.
-                if new_n is None:
-                    new_n = deepcopy(original_node)
+    for n in nodes.values():
+        params = tuple(n.parameters.keys())
+        for parameter_name in params:
+            if n.parameters[parameter_name] == node:
                 # Swap the parameter to a constant parameter, giving it the fixed value
-                new_n.constant_parameters[parameter_name] = value
+                n.constant_parameters[parameter_name] = value
                 # Remove the parameter from the node's record of non-constant parameters
-                new_n.parameters.pop(parameter_name)
-        # If we had to recreate a new node, add it to the new (Di)Graph.
-        # Also record the name of the node that it is set to replace
-        if new_n is not None:
-            g.add_node(new_n)
-            new_nodes[original_node.label] = new_n
-
-    # Any new_nodes whose counterparts connect to other nodes in the network need
-    # to mimic these links.
-    for edge in old_g.edges:
-        if edge[0].label in new_nodes or edge[1].label in new_nodes:
-            g.add_edge(
-                new_nodes.get(edge[0].label, edge[0]),
-                new_nodes.get(edge[1].label, edge[1]),
+                n.parameters.pop(parameter_name)
+
+    # Recursively remove nodes that are predecessors of removed nodes
+    nodes_to_remove: tuple[str, ...] = (node,)
+    while len(nodes_to_remove) > 0:
+        nodes_to_remove = removable_nodes(graph, nodes)
+        for n in removable_nodes(graph, nodes):
+            nodes.pop(n)
+
+    # Check for nodes that are predecessors of both a removed node and a remaining node
+    # and throw an error if one of these is found
+    for n in nodes:
+        _, excluded = get_included_excluded_successors(graph, nodes, n)
+        if len(excluded) > 0:
+            msg = (
+                "Node that is predecessor of node set by do and "
+                f'nodes that are not removed found ("{n}")'
             )
-    # Now that the new_nodes are present in the graph, and correctly connected, remove
-    # their counterparts from the graph.
-    for original_node in new_nodes:
-        g.remove_node(nodes_by_label[original_node])
+            raise ValueError(msg)
+
+    g = Graph(label=f"{label}|do[{node}={value}]")
+    for n in nodes.values():
+        g.add_node(n)
+
+    # Any nodes whose counterparts connect to other nodes in the network need
+    # to mimic these links.
+    for edge in graph.edges:
+        if edge[0].label in nodes and edge[1].label in nodes:
+            g.add_edge(edge[0].label, edge[1].label)
 
-    return Graph(label=label, graph=g)
+    return g

src/causalprog/graph/graph.py

@@ -4,7 +4,7 @@
 import numpy.typing as npt
 
 from causalprog._abc.labelled import Labelled
-from causalprog.graph.node import DistributionNode, Node, ParameterNode
+from causalprog.graph.node import Node
 
 
 class Graph(Labelled):
@@ -89,7 +89,7 @@ def add_edge(self, start_node: Node | str, end_node: Node | str) -> None:
         self._graph.add_edge(start_node, end_node)
 
     @property
-    def parameter_nodes(self) -> tuple[ParameterNode, ...]:
+    def parameter_nodes(self) -> tuple[Node, ...]:
         """
         Returns all parameter nodes in the graph.
 
@@ -107,31 +107,31 @@ def parameter_nodes(self) -> tuple[ParameterNode, ...]:
         return tuple(node for node in self.ordered_nodes if node.is_parameter)
 
     @property
-    def predecessors(self) -> dict[Node, list[Node]]:
+    def predecessors(self) -> dict[Node, tuple[Node, ...]]:
         """
         Get predecessors of every node.
 
         Returns:
             Mapping of each Node to its predecessor Nodes
 
         """
-        return {node: list(self._graph.predecessors(node)) for node in self.nodes}
+        return {node: tuple(self._graph.predecessors(node)) for node in self.nodes}
 
     @property
-    def successors(self) -> dict[Node, list[Node]]:
+    def successors(self) -> dict[Node, tuple[Node, ...]]:
         """
         Get successors of every node.
 
         Returns:
             Mapping of each Node to its successor Nodes.
 
         """
-        return {node: list(self._graph.successors(node)) for node in self.nodes}
+        return {node: tuple(self._graph.successors(node)) for node in self.nodes}
 
     @property
-    def nodes(self) -> list[Node]:
+    def nodes(self) -> tuple[Node, ...]:
         """
-        Get the nodes of the graph, with no enforeced ordering.
+        Get the nodes of the graph, with no enforced ordering.
 
         Returns:
             A list of all the nodes in the graph.
@@ -140,10 +140,21 @@ def nodes(self) -> list[Node]:
             ordered_nodes: Fetch an ordered list of the nodes in the graph.
 
         """
-        return list(self._graph.nodes())
+        return tuple(self._graph.nodes())
 
     @property
-    def ordered_nodes(self) -> list[Node]:
+    def edges(self) -> tuple[tuple[Node, Node], ...]:
+        """
+        Get the edges of the graph.
+
+        Returns:
+            A tuple of all the edges in the graph.
+
+        """
+        return tuple(self._graph.edges())
+
+    @property
+    def ordered_nodes(self) -> tuple[Node, ...]:
         """
         Nodes ordered so that each node appears after its dependencies.
 
@@ -155,23 +166,23 @@ def ordered_nodes(self) -> list[Node]:
         if not nx.is_directed_acyclic_graph(self._graph):
             msg = "Graph is not acyclic."
             raise RuntimeError(msg)
-        return list(nx.topological_sort(self._graph))
+        return tuple(nx.topological_sort(self._graph))
 
     @property
-    def ordered_dist_nodes(self) -> list[DistributionNode]:
+    def ordered_dist_nodes(self) -> tuple[Node, ...]:
         """
         `DistributionNode`s in dependency order.
 
         Each `DistributionNode` in the returned list appears after all its
         dependencies. Order is derived from `self.ordered_nodes`, selecting
         only those nodes where `is_distribution` is `True`.
         """
-        return [node for node in self.ordered_nodes if node.is_distribution]
+        return tuple(node for node in self.ordered_nodes if node.is_distribution)
 
     def roots_down_to_outcome(
         self,
         outcome_node_label: str,
-    ) -> list[Node]:
+    ) -> tuple[Node, ...]:
         """
         Get ordered list of nodes that outcome depends on.
 
@@ -186,9 +197,9 @@ def roots_down_to_outcome(
         """
         outcome = self.get_node(outcome_node_label)
         ancestors = nx.ancestors(self._graph, outcome)
-        return [
+        return tuple(
             node for node in self.ordered_nodes if node == outcome or node in ancestors
-        ]
+        )
 
     def model(self, **parameter_values: npt.ArrayLike) -> dict[str, npt.ArrayLike]:
         """

src/causalprog/graph/node.py

@@ -201,7 +201,12 @@ def copy(self) -> Node:
 
     @override
     def __repr__(self) -> str:
-        return f'DistributionNode("{self.label}")'
+        r = f'DistributionNode({self._dist.__name__}, label="{self.label}"'
+        if len(self._parameters) > 0:
+            r += f", parameters={self._parameters}"
+        if len(self._constant_parameters) > 0:
+            r += f", constant_parameters={self._constant_parameters}"
+        return r
 
     @override
     @property
@@ -283,7 +288,7 @@ def copy(self) -> Node:
 
     @override
     def __repr__(self) -> str:
-        return f'ParameterNode("{self.label}")'
+        return f'ParameterNode(label="{self.label}")'
 
     @override
     @property

tests/test_algorithms/test_do.py

@@ -0,0 +1,76 @@
+"""Tests for the do algorithm."""
+
+from causalprog import algorithms
+from causalprog.graph import Graph, ParameterNode
+
+max_samples = 10**5
+
+
+def test_do(two_normal_graph, raises_context):
+    graph = two_normal_graph(5.0, 1.2, 0.8)
+    graph2 = algorithms.do(graph, "UX", 4.0)
+
+    assert "loc" in graph.get_node("X").parameters
+    assert "loc" not in graph.get_node("X").constant_parameters
+    assert "loc" not in graph2.get_node("X").parameters
+    assert "loc" in graph2.get_node("X").constant_parameters
+
+    graph.get_node("UX")
+    with raises_context(KeyError('Node not found with label "UX"')):
+        graph2.get_node("UX")
+
+
+def test_do_removes_dependencies(two_normal_graph, raises_context):
+    graph = two_normal_graph()
+    graph2 = algorithms.do(graph, "UX", 4.0)
+
+    for node in ["UX", "mean", "cov"]:
+        graph.get_node(node)
+        with raises_context(KeyError(f'Node not found with label "{node}"')):
+            graph2.get_node(node)
+
+
+def test_do_edges(two_normal_graph):
+    graph = two_normal_graph()
+    graph2 = algorithms.do(graph, "UX", 4.0)
+
+    edges = [(e[0].label, e[1].label) for e in graph.edges]
+    edges2 = [(e[0].label, e[1].label) for e in graph2.edges]
+
+    # Check that correct edges are removed
+    for e in [
+        ("UX", "X"),
+        ("mean", "UX"),
+        ("cov", "UX"),
+    ]:
+        assert e in edges
+        assert e not in edges2
+
+    # Check that correct edges remain
+    for e in [
+        ("cov2", "X"),
+    ]:
+        assert e in edges
+        assert e in edges2
+
+
+def test_do_error(raises_context):
+    graph = Graph(label="ABC")
+    graph.add_node(ParameterNode(label="A"))
+    graph.add_node(ParameterNode(label="B1"))
+    graph.add_node(ParameterNode(label="B2"))
+    graph.add_node(ParameterNode(label="C"))
+    graph.add_edge("A", "B1")
+    graph.add_edge("A", "B2")
+    graph.add_edge("B1", "C")
+    graph.add_edge("B2", "C")
+
+    # Currently, applying to do to a node that had predecessors that cannot be removed
+    # raises an error, see https://github.com/UCL/causalprog/issues/80
+    with raises_context(
+        ValueError(
+            "Node that is predecessor of node set by do and nodes that are not removed "
+            "found"
+        )
+    ):
+        algorithms.do(graph, "B1", 1.0)