diff --git a/ceci/pipeline/graph.py b/ceci/pipeline/graph.py index 0aa2dce..0b1fb2b 100644 --- a/ceci/pipeline/graph.py +++ b/ceci/pipeline/graph.py @@ -207,13 +207,13 @@ def combine_output_nodes(graph): Parameters ---------- - graph : networkx.DiGraph + graph : pygraphviz.AGraph The pipeline graph to modify """ groups = collections.defaultdict(list) # find groups of nodes that have no outgoing edges and all # share the same incoming edge. - for node in graph.nodes_iter(): + for node in graph.nodes(): node_data = graph.get_node(node).attr if node_data["type"] != "output" or len(graph.successors(node)) != 0: continue @@ -238,3 +238,40 @@ def combine_output_nodes(graph): graph.add_edge(parent, new_node_name) # remove the old nodes graph.remove_nodes_from(group) + + +def combine_input_nodes(graph): + """ + Combine input nodes that share the same successor and have no + predecessors into a single node. + + This is useful for plotting flow charts. + + Parameters + ---------- + graph : pygraphviz.AGraph + The pipeline graph to modify + """ + groups = collections.defaultdict(list) + for node in graph.nodes(): + node_data = graph.get_node(node).attr + if node_data["type"] != "input" or len(graph.predecessors(node)) != 0: + continue + successors = list(graph.successors(node)) + + # This shouldn't happen in the default case, + # but might if things change later, so let's deal with it + # just in case. + if len(successors) != 1: + continue + + successor = successors[0] + groups[successor].append(node) + + for successor, group in groups.items(): + if len(group) < 2: + continue + new_node_name = "\n".join(group) + graph.add_node(new_node_name, type="input") + graph.add_edge(new_node_name, successor) + graph.remove_nodes_from(group) \ No newline at end of file diff --git a/ceci/pipeline/pipeline.py b/ceci/pipeline/pipeline.py index b3a66cc..73ec850 100644 --- a/ceci/pipeline/pipeline.py +++ b/ceci/pipeline/pipeline.py @@ -9,7 +9,7 @@ from ..stage import PipelineStage from ..sites import load, set_default_site, get_default_site from ..utils import extra_paths -from .graph import build_graph, get_static_ordering, trim_pipeline_graph, combine_output_nodes +from .graph import build_graph, get_static_ordering, trim_pipeline_graph, combine_output_nodes, combine_input_nodes from .file_manager import FileManager from .sec import StageExecutionConfig from .templates import read_and_apply_template @@ -910,9 +910,10 @@ def make_flow_chart(self, filename=None): # are final outputs not used in the rest of the pipeline. We # can combine these to make the graph more readable. combine_output_nodes(graph) + combine_input_nodes(graph) # set the colours and styles for the boxes - for node in graph.nodes_iter(): + for node in graph.nodes(): node_type = node.attr['type'] if node_type == "input": node.attr.update(shape="box", color="gold", style="filled") diff --git a/tests/test_graph.py b/tests/test_graph.py index 2e345d0..d043b09 100644 --- a/tests/test_graph.py +++ b/tests/test_graph.py @@ -1,5 +1,6 @@ -from ceci.pipeline.graph import trim_pipeline_graph +from ceci.pipeline.graph import trim_pipeline_graph, combine_input_nodes import networkx +import pygraphviz def check_membership(graph, everything, expected): @@ -97,3 +98,23 @@ def test_trim_pipeline_graph(): } check_membership(sub4, everything, expected4) + + +def test_combine_input_nodes(): + graph = pygraphviz.AGraph(directed=True) + graph.add_node("a", type="input") + graph.add_node("b", type="input") + graph.add_node("stage", type="stage") + graph.add_edge("a", "stage") + graph.add_edge("b", "stage") + + combine_input_nodes(graph) + + input_nodes = [ + str(node) + for node in graph.nodes() + if graph.get_node(node).attr["type"] == "input" + ] + assert len(input_nodes) == 1 + assert input_nodes[0] == "a\nb" + assert [str(node) for node in graph.predecessors("stage")] == ["a\nb"] \ No newline at end of file