Merge pull request #322 from djarecka/enh/dotfile

creating dotfiles for graphs (closes #197)

Author: djarecka

pydra/engine/core.py

@@ -737,7 +737,7 @@ def __init__(
             rerun=rerun,
         )
 
-        self.graph = DiGraph()
+        self.graph = DiGraph(name=name)
         self.name2obj = {}
 
         # store output connections
@@ -828,15 +828,17 @@ def add(self, task):
         logger.debug(f"Added {task}")
         return self
 
-    def create_connections(self, task):
+    def create_connections(self, task, detailed=False):
         """
         Add and connect a particular task to existing nodes in the workflow.
 
         Parameters
         ----------
         task : :class:`TaskBase`
             The task to be added.
-
+        detailed : :obj:`bool`
+            If True, `add_edges_description` is run for self.graph to add
+            a detailed descriptions of the connections (input/output fields names)
         """
         other_states = {}
         for field in attr_fields(task.inputs):
@@ -847,9 +849,12 @@ def create_connections(self, task):
                 # adding an edge to the graph if task id expecting output from a different task
                 if val.name != self.name:
                     # checking if the connection is already in the graph
-                    if (getattr(self, val.name), task) in self.graph.edges:
-                        continue
-                    self.graph.add_edges((getattr(self, val.name), task))
+                    if (getattr(self, val.name), task) not in self.graph.edges:
+                        self.graph.add_edges((getattr(self, val.name), task))
+                    if detailed:
+                        self.graph.add_edges_description(
+                            (task.name, field.name, val.name, val.field)
+                        )
                     logger.debug("Connecting %s to %s", val.name, task.name)
 
                     if (
@@ -861,6 +866,13 @@ def create_connections(self, task):
                             getattr(self, val.name).state,
                             field.name,
                         )
+                else:  # LazyField with the wf input
+                    # connections with wf input should be added to the detailed graph description
+                    if detailed:
+                        self.graph.add_edges_description(
+                            (task.name, field.name, val.name, val.field)
+                        )
+
         # if task has connections state has to be recalculated
         if other_states:
             if hasattr(task, "fut_combiner"):
@@ -1000,6 +1012,49 @@ def _collect_outputs(self):
                     raise ValueError(f"Task {val.name} raised an error")
         return attr.evolve(output, **output_wf)
 
+    def create_dotfile(self, type="simple", export=None, name=None):
+        """creating a graph - dotfile and optionally exporting to other formats"""
+        if not name:
+            name = f"graph_{self.name}"
+        if type == "simple":
+            for task in self.graph.nodes:
+                self.create_connections(task)
+            dotfile = self.graph.create_dotfile_simple(
+                outdir=self.output_dir, name=name
+            )
+        elif type == "nested":
+            for task in self.graph.nodes:
+                self.create_connections(task)
+            dotfile = self.graph.create_dotfile_nested(
+                outdir=self.output_dir, name=name
+            )
+        elif type == "detailed":
+            # create connections with detailed=True
+            for task in self.graph.nodes:
+                self.create_connections(task, detailed=True)
+            # adding wf outputs
+            for (wf_out, lf) in self._connections:
+                self.graph.add_edges_description((self.name, wf_out, lf.name, lf.field))
+            dotfile = self.graph.create_dotfile_detailed(
+                outdir=self.output_dir, name=name
+            )
+        else:
+            raise Exception(
+                f"type of the graph can be simple, detailed or nested, "
+                f"but {type} provided"
+            )
+        if not export:
+            return dotfile
+        else:
+            if export is True:
+                export = ["png"]
+            elif isinstance(export, str):
+                export = [export]
+            formatted_dot = []
+            for ext in export:
+                formatted_dot.append(self.graph.export_graph(dotfile=dotfile, ext=ext))
+            return dotfile, formatted_dot
+
 
 def is_task(obj):
     """Check whether an object looks like a task."""

pydra/engine/graph.py

@@ -1,12 +1,15 @@
 """Data structure to support :class:`~pydra.engine.core.Workflow` tasks."""
 from copy import copy
+from pathlib import Path
+import subprocess as sp
+
 from .helpers import ensure_list
 
 
 class DiGraph:
     """A simple Directed Graph object."""
 
-    def __init__(self, nodes=None, edges=None):
+    def __init__(self, name=None, nodes=None, edges=None):
         """
         Initialize a directed graph.
 
@@ -19,13 +22,15 @@ def __init__(self, nodes=None, edges=None):
             the graph.
 
         """
+        self.name = name
         self._nodes = []
         self.nodes = nodes
         self._edges = []
         self.edges = edges
         self._create_connections()
         self._sorted_nodes = None
         self._node_wip = []
+        self._nodes_details = {}
 
     def copy(self):
         """
@@ -92,6 +97,20 @@ def edges_names(self):
         """Get edges as pairs of the nodes they connect."""
         return [(edg[0].name, edg[1].name) for edg in self._edges]
 
+    @property
+    def nodes_details(self):
+        """ dictionary with details of the nodes
+            for each task, there are inputs/outputs and connections
+            (with input/output fields names)
+        """
+        # removing repeated fields from inputs and outputs
+        for el in self._nodes_details.values():
+            el["inputs"] = list(set(el["inputs"]))
+            el["inputs"].sort()
+            el["outputs"] = list(set(el["outputs"]))
+            el["outputs"].sort()
+        return self._nodes_details
+
     @property
     def sorted_nodes(self):
         """Return sorted nodes (runs sorting if needed)."""
@@ -136,6 +155,19 @@ def add_edges(self, new_edges):
             # starting from the previous sorted list, so it's faster
             self.sorting(presorted=self.sorted_nodes + [])
 
+    def add_edges_description(self, new_edge_details):
+        """ adding detailed description of the connections, filling _nodes_details"""
+        in_nd, in_fld, out_nd, out_fld = new_edge_details
+        for key in [in_nd, out_nd]:
+            self._nodes_details.setdefault(
+                key, {"inputs": [], "outputs": [], "connections": []}
+            )
+
+        if (in_fld, out_nd, out_fld) not in self._nodes_details[in_nd]["connections"]:
+            self._nodes_details[in_nd]["connections"].append((in_fld, out_nd, out_fld))
+            self._nodes_details[in_nd]["inputs"].append(in_fld)
+            self._nodes_details[out_nd]["outputs"].append(out_fld)
+
     def sorting(self, presorted=None):
         """
         Sort this graph.
@@ -318,3 +350,200 @@ def calculate_max_paths(self):
         for nm in first_nodes:
             self.max_paths[nm] = {}
             self._checking_path(node_name=nm, first_name=nm)
+
+    def create_dotfile_simple(self, outdir, name="graph"):
+        """ creates a simple dotfile (no nested structure)"""
+        from .core import is_workflow
+
+        dotstr = "digraph G {\n"
+        for nd in self.nodes:
+            # breakpoint()
+            if is_workflow(nd):
+                if nd.state:
+                    # adding color for wf with a state
+                    dotstr += f"{nd.name} [shape=box, color=blue]\n"
+                else:
+                    dotstr += f"{nd.name} [shape=box]\n"
+            else:
+                if nd.state:
+                    # adding color for nd with a state
+                    dotstr += f"{nd.name} [color=blue]\n"
+                else:
+                    dotstr += f"{nd.name}\n"
+        for ed in self.edges_names:
+            dotstr += f"{ed[0]} -> {ed[1]}\n"
+
+        dotstr += "}"
+        Path(outdir).mkdir(parents=True, exist_ok=True)
+        dotfile = Path(outdir) / f"{name}.dot"
+        dotfile.write_text(dotstr)
+        return dotfile
+
+    def create_dotfile_detailed(self, outdir, name="graph_det"):
+        """ creates a detailed dotfile (detailed connections - input/output fields,
+        but no nested structure)
+        """
+        dotstr = "digraph structs {\n"
+        dotstr += "node [shape=record];\n"
+        if not self._nodes_details:
+            raise Exception("node_details is empty, detailed dotfile can't be created")
+        for nd_nm, nd_det in self.nodes_details.items():
+            if nd_nm == self.name:  # the main workflow itself
+                # wf inputs
+                wf_inputs_str = f'{{<{nd_det["outputs"][0]}> {nd_det["outputs"][0]}'
+                for el in nd_det["outputs"][1:]:
+                    wf_inputs_str += f" | <{el}> {el}"
+                wf_inputs_str += "}"
+                dotstr += f'struct_{nd_nm} [color=red, label="{{WORKFLOW INPUT: | {wf_inputs_str}}}"];\n'
+                # wf outputs
+                wf_outputs_str = f'{{<{nd_det["inputs"][0]}> {nd_det["inputs"][0]}'
+                for el in nd_det["inputs"][1:]:
+                    wf_outputs_str += f" | <{el}> {el}"
+                wf_outputs_str += "}"
+                dotstr += f'struct_{nd_nm}_out [color=red, label="{{WORKFLOW OUTPUT: | {wf_outputs_str}}}"];\n'
+                # connections to the wf outputs
+                for con in nd_det["connections"]:
+                    dotstr += (
+                        f"struct_{con[1]}:{con[2]} -> struct_{nd_nm}_out:{con[0]};\n"
+                    )
+            else:  # elements of the main workflow
+                inputs_str = "{INPUT:"
+                for inp in nd_det["inputs"]:
+                    inputs_str += f" | <{inp}> {inp}"
+                inputs_str += "}"
+                outputs_str = "{OUTPUT:"
+                for out in nd_det["outputs"]:
+                    outputs_str += f" | <{out}> {out}"
+                outputs_str += "}"
+                dotstr += f'struct_{nd_nm} [shape=record, label="{inputs_str} | {nd_nm} | {outputs_str}"];\n'
+                # connections between elements
+                for con in nd_det["connections"]:
+                    dotstr += f"struct_{con[1]}:{con[2]} -> struct_{nd_nm}:{con[0]};\n"
+        dotstr += "}"
+        Path(outdir).mkdir(parents=True, exist_ok=True)
+        dotfile = Path(outdir) / f"{name}.dot"
+        dotfile.write_text(dotstr)
+        return dotfile
+
+    def create_dotfile_nested(self, outdir, name="graph"):
+        """dotfile that includes the nested structures for workflows"""
+        dotstr = "digraph G {\ncompound=true \n"
+        dotstr += self._create_dotfile_single_graph(nodes=self.nodes, edges=self.edges)
+        dotstr += "}"
+        Path(outdir).mkdir(parents=True, exist_ok=True)
+        dotfile = Path(outdir) / f"{name}.dot"
+        dotfile.write_text(dotstr)
+        return dotfile
+
+    def _create_dotfile_single_graph(self, nodes, edges):
+        from .core import is_workflow
+
+        wf_asnd = []
+        dotstr = ""
+        for nd in nodes:
+            if is_workflow(nd):
+                wf_asnd.append(nd.name)
+                for task in nd.graph.nodes:
+                    nd.create_connections(task)
+                dotstr += f"subgraph cluster_{nd.name} {{\n" f"label = {nd.name} \n"
+                dotstr += self._create_dotfile_single_graph(
+                    nodes=nd.graph.nodes, edges=nd.graph.edges
+                )
+                if nd.state:
+                    dotstr += "color=blue\n"
+                dotstr += "}\n"
+            else:
+                if nd.state:
+                    dotstr += f"{nd.name} [color=blue]\n"
+                else:
+                    dotstr += f"{nd.name}\n"
+
+        dotstr_edg = ""
+        for ed in edges:
+            if ed[0].name in wf_asnd and ed[1].name in wf_asnd:
+                head_nd = list(ed[1].nodes)[0].name
+                tail_nd = list(ed[0].nodes)[-1].name
+                dotstr_edg += (
+                    f"{tail_nd} -> {head_nd} "
+                    f"[ltail=cluster_{ed[0].name}, "
+                    f"lhead=cluster_{ed[1].name}]\n"
+                )
+            elif ed[0].name in wf_asnd:
+                tail_nd = list(ed[0].nodes)[-1].name
+                dotstr_edg += (
+                    f"{tail_nd} -> {ed[1].name} [ltail=cluster_{ed[0].name}]\n"
+                )
+            elif ed[1].name in wf_asnd:
+                head_nd = list(ed[1].nodes)[0].name
+                dotstr_edg += (
+                    f"{ed[0].name} -> {head_nd} [lhead=cluster_{ed[1].name}]\n"
+                )
+            else:
+                dotstr_edg += f"{ed[0].name} -> {ed[1].name}\n"
+        dotstr = dotstr + dotstr_edg
+        return dotstr
+
+    def export_graph(self, dotfile, ext="png"):
+        """ exporting dotfile to other format, equires the dot command"""
+        available_ext = [
+            "bmp",
+            "canon",
+            "cgimage",
+            "cmap",
+            "cmapx",
+            "cmapx_np",
+            "dot",
+            "dot_json",
+            "eps",
+            "exr",
+            "fig",
+            "gif",
+            "gv",
+            "icns",
+            "ico",
+            "imap",
+            "imap_np",
+            "ismap",
+            "jp2",
+            "jpe",
+            "jpeg",
+            "jpg",
+            "json",
+            "json0",
+            "mp",
+            "pct",
+            "pdf",
+            "pic",
+            "pict",
+            "plain",
+            "plain-ext",
+            "png",
+            "pov",
+            "ps",
+            "ps2",
+            "psd",
+            "sgi",
+            "svg",
+            "svgz",
+            "tga",
+            "tif",
+            "tiff",
+            "tk",
+            "vml",
+            "vmlz",
+            "xdot",
+            "xdot1.2",
+            "xdot1.4",
+            "xdot_json",
+        ]
+        if ext not in available_ext:
+            raise Exception(f"unvalid extension - {ext}, chose from {available_ext}")
+
+        dot_check = sp.run(["which", "dot"], stdout=sp.PIPE, stderr=sp.PIPE)
+        if not dot_check.stdout:
+            raise Exception(f"dot command not available, can't create a {ext} file")
+
+        formatted_dot = dotfile.with_suffix(f".{ext}")
+        cmd = f"dot -T{ext} -o {formatted_dot} {dotfile}"
+        sp.run(cmd.split(), stdout=sp.PIPE, stderr=sp.PIPE)
+        return formatted_dot