Repository: Qiskit/qiskit

Issue:
Add support to PassManager drawer to display stages
### What should we add?

In #6403 we added the concept of a staged passmanger where the different sections of a pass manager are organized into named stages. However, when calling the `draw()` method the additional metadata for the stages are missed. We should  modify the drawer for the `StagedPassManager` class to annotate the output diagram with the stages. The idea I had in mind was putting each stage as an outer box around the passes.


Discussion:
I'd like to work on this
Just a quick sanity check to make sure I'm not overcomplicating things:

`StagedPassManager` inherits the `draw()` method from `PassManager`. Which, if I'm understanding correctly, just iterates over the list of passes and draws the graph without considering any possible stage. Thus, to add an outer box around the passes for each stage, we'd thus need to implement a new  `draw()` method for `StagedPassManager` that:

* builds a subgraph for each stage
* links those subgraphs together (i.e. connects the last node in one stage to the first node in the next)
* returns the final graph

If there's something obvious I'm missing (and I very well may be), feel free to point me in another direction...


Yeah, that's basically the design I had in mind when I opened the issue. I think that plan makes sense.

Resolving patch:
diff --git a/qiskit/transpiler/passmanager.py b/qiskit/transpiler/passmanager.py
--- a/qiskit/transpiler/passmanager.py
+++ b/qiskit/transpiler/passmanager.py
@@ -526,3 +526,9 @@ def run(
     ) -> Union[QuantumCircuit, List[QuantumCircuit]]:
         self._update_passmanager()
         return super().run(circuits, output_name, callback)
+
+    def draw(self, filename=None, style=None, raw=False):
+        """Draw the staged pass manager."""
+        from qiskit.visualization import staged_pass_manager_drawer
+
+        return staged_pass_manager_drawer(self, filename=filename, style=style, raw=raw)
diff --git a/qiskit/visualization/__init__.py b/qiskit/visualization/__init__.py
--- a/qiskit/visualization/__init__.py
+++ b/qiskit/visualization/__init__.py
@@ -286,6 +286,7 @@
 from .dag_visualization import dag_drawer
 from .gate_map import plot_gate_map, plot_circuit_layout, plot_error_map, plot_coupling_map
 from .pass_manager_visualization import pass_manager_drawer
+from .pass_manager_visualization import staged_pass_manager_drawer
 
 from .pulse.interpolation import step_wise, linear, cubic_spline
 from .pulse.qcstyle import PulseStyle, SchedStyle
diff --git a/qiskit/visualization/pass_manager_visualization.py b/qiskit/visualization/pass_manager_visualization.py
--- a/qiskit/visualization/pass_manager_visualization.py
+++ b/qiskit/visualization/pass_manager_visualization.py
@@ -20,7 +20,6 @@
 
 from qiskit.utils import optionals as _optionals
 from qiskit.transpiler.basepasses import AnalysisPass, TransformationPass
-from . import utils
 from .exceptions import VisualizationError
 
 DEFAULT_STYLE = {AnalysisPass: "red", TransformationPass: "blue"}
@@ -91,66 +90,168 @@ def pass_manager_drawer(pass_manager, filename=None, style=None, raw=False):
     prev_node = None
 
     for index, controller_group in enumerate(passes):
+        subgraph, component_id, prev_node = draw_subgraph(
+            controller_group, component_id, style, prev_node, index
+        )
+        graph.add_subgraph(subgraph)
+
+    output = make_output(graph, raw, filename)
+    return output
+
+
+def _get_node_color(pss, style):
+    # look in the user provided dict first
+    for typ, color in style.items():
+        if isinstance(pss, typ):
+            return color
+
+    # failing that, look in the default
+    for typ, color in DEFAULT_STYLE.items():
+        if isinstance(pss, typ):
+            return color
+
+    return "black"
+
+
+@_optionals.HAS_GRAPHVIZ.require_in_call
+@_optionals.HAS_PYDOT.require_in_call
+def staged_pass_manager_drawer(pass_manager, filename=None, style=None, raw=False):
+    """
+    Draws the staged pass manager.
+
+        This function needs `pydot <https://github.com/erocarrera/pydot>`__, which in turn needs
+    `Graphviz <https://www.graphviz.org/>`__ to be installed.
 
-        # label is the name of the flow controller parameter
-        label = "[{}] {}".format(index, ", ".join(controller_group["flow_controllers"]))
+    Args:
+        pass_manager (StagedPassManager): the staged pass manager to be drawn
+        filename (str): file path to save image to
+        style (dict or OrderedDict): keys are the pass classes and the values are
+            the colors to make them. An example can be seen in the DEFAULT_STYLE. An ordered
+            dict can be used to ensure a priority coloring when pass falls into multiple
+            categories. Any values not included in the provided dict will be filled in from
+            the default dict
+        raw (Bool) : True if you want to save the raw Dot output not an image. The
+            default is False.
+    Returns:
+        PIL.Image or None: an in-memory representation of the pass manager. Or None if
+        no image was generated or PIL is not installed.
+    Raises:
+        MissingOptionalLibraryError: when nxpd or pydot not installed.
+        VisualizationError: If raw=True and filename=None.
+
+    Example:
+        .. code-block::
+
+            %matplotlib inline
+            from qiskit.providers.fake_provider import FakeLagosV2
+            from qiskit.transpiler.preset_passmanagers import generate_preset_pass_manager
+
+            pass_manager = generate_preset_pass_manager(3, FakeLagosV2())
+            pass_manager.draw()
+    """
+    import pydot
 
-        # create the subgraph for this controller
-        subgraph = pydot.Cluster(
-            str(component_id), label=label, fontname="helvetica", labeljust="l"
+    # only include stages that have passes
+    stages = list(filter(lambda s: s is not None, pass_manager.expanded_stages))
+
+    if not style:
+        style = DEFAULT_STYLE
+
+    # create the overall graph
+    graph = pydot.Dot()
+
+    # identifiers for nodes need to be unique, so assign an id
+    # can't just use python's id in case the exact same pass was
+    # appended more than once
+    component_id = 0
+
+    # keep a running count of indexes across stages
+    idx = 0
+
+    prev_node = None
+
+    for st in stages:
+        stage = getattr(pass_manager, st)
+
+        if stage is not None:
+            passes = stage.passes()
+            stagegraph = pydot.Cluster(str(st), label=str(st), fontname="helvetica", labeljust="l")
+            for controller_group in passes:
+                subgraph, component_id, prev_node = draw_subgraph(
+                    controller_group, component_id, style, prev_node, idx
+                )
+                stagegraph.add_subgraph(subgraph)
+                idx += 1
+            graph.add_subgraph(stagegraph)
+
+    output = make_output(graph, raw, filename)
+    return output
+
+
+def draw_subgraph(controller_group, component_id, style, prev_node, idx):
+    """Draw subgraph."""
+    import pydot
+
+    # label is the name of the flow controller parameter
+    label = "[{}] {}".format(idx, ", ".join(controller_group["flow_controllers"]))
+
+    # create the subgraph for this controller
+    subgraph = pydot.Cluster(str(component_id), label=label, fontname="helvetica", labeljust="l")
+    component_id += 1
+
+    for pass_ in controller_group["passes"]:
+
+        # label is the name of the pass
+        node = pydot.Node(
+            str(component_id),
+            label=str(type(pass_).__name__),
+            color=_get_node_color(pass_, style),
+            shape="rectangle",
+            fontname="helvetica",
         )
-        component_id += 1
 
-        for pass_ in controller_group["passes"]:
+        subgraph.add_node(node)
+        component_id += 1
 
-            # label is the name of the pass
-            node = pydot.Node(
-                str(component_id),
-                label=str(type(pass_).__name__),
-                color=_get_node_color(pass_, style),
-                shape="rectangle",
+        # the arguments that were provided to the pass when it was created
+        arg_spec = inspect.getfullargspec(pass_.__init__)
+        # 0 is the args, 1: to remove the self arg
+        args = arg_spec[0][1:]
+
+        num_optional = len(arg_spec[3]) if arg_spec[3] else 0
+
+        # add in the inputs to the pass
+        for arg_index, arg in enumerate(args):
+            nd_style = "solid"
+            # any optional args are dashed
+            # the num of optional counts from the end towards the start of the list
+            if arg_index >= (len(args) - num_optional):
+                nd_style = "dashed"
+
+            input_node = pydot.Node(
+                component_id,
+                label=arg,
+                color="black",
+                shape="ellipse",
+                fontsize=10,
+                style=nd_style,
                 fontname="helvetica",
             )
-
-            subgraph.add_node(node)
+            subgraph.add_node(input_node)
             component_id += 1
+            subgraph.add_edge(pydot.Edge(input_node, node))
 
-            # the arguments that were provided to the pass when it was created
-            arg_spec = inspect.getfullargspec(pass_.__init__)
-            # 0 is the args, 1: to remove the self arg
-            args = arg_spec[0][1:]
-
-            num_optional = len(arg_spec[3]) if arg_spec[3] else 0
-
-            # add in the inputs to the pass
-            for arg_index, arg in enumerate(args):
-                nd_style = "solid"
-                # any optional args are dashed
-                # the num of optional counts from the end towards the start of the list
-                if arg_index >= (len(args) - num_optional):
-                    nd_style = "dashed"
-
-                input_node = pydot.Node(
-                    component_id,
-                    label=arg,
-                    color="black",
-                    shape="ellipse",
-                    fontsize=10,
-                    style=nd_style,
-                    fontname="helvetica",
-                )
-                subgraph.add_node(input_node)
-                component_id += 1
-                subgraph.add_edge(pydot.Edge(input_node, node))
+        # if there is a previous node, add an edge between them
+        if prev_node:
+            subgraph.add_edge(pydot.Edge(prev_node, node))
 
-            # if there is a previous node, add an edge between them
-            if prev_node:
-                subgraph.add_edge(pydot.Edge(prev_node, node))
+        prev_node = node
 
-            prev_node = node
+    return subgraph, component_id, prev_node
 
-        graph.add_subgraph(subgraph)
 
+def make_output(graph, raw, filename):
+    """Produce output for pass_manager."""
     if raw:
         if filename:
             graph.write(filename, format="raw")
@@ -174,22 +275,7 @@ def pass_manager_drawer(pass_manager, filename=None, style=None, raw=False):
         graph.write_png(tmppath)  # pylint: disable=no-member
 
         image = Image.open(tmppath)
-        image = utils._trim(image)
         os.remove(tmppath)
         if filename:
             image.save(filename, "PNG")
         return image
-
-
-def _get_node_color(pss, style):
-    # look in the user provided dict first
-    for typ, color in style.items():
-        if isinstance(pss, typ):
-            return color
-
-    # failing that, look in the default
-    for typ, color in DEFAULT_STYLE.items():
-        if isinstance(pss, typ):
-            return color
-
-    return "black"


Test patch: