adds support to PassManager drawer to display stages

Signed-off-by: Paul S. Schweigert <paul@paulschweigert.com>

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)

qiskit/visualization/__init__.py

@@ -200,6 +200,7 @@
    :toctree: ../stubs/
 
    pass_manager_drawer
+   staged_pass_manager_drawer
 
 Pulse Visualizations
 ====================
@@ -267,6 +268,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

qiskit/visualization/pass_manager_visualization.py

@@ -193,3 +193,160 @@ def _get_node_color(pss, style):
             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.
+
+    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
+
+    # 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="", fontname="helvetica", labeljust="l")
+            for controller_group in passes:
+
+                # 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",
+                    )
+
+                    subgraph.add_node(node)
+                    component_id += 1
+
+                    # 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))
+
+                    prev_node = node
+
+                stagegraph.add_subgraph(subgraph)
+                idx += 1
+
+            graph.add_subgraph(stagegraph)
+
+    if raw:
+        if filename:
+            graph.write(filename, format="raw")
+            return None
+        else:
+            raise VisualizationError("if format=raw, then a filename is required.")
+
+    if not _optionals.HAS_PIL and filename:
+        # pylint says this isn't a method - it is
+        graph.write_png(filename)  # pylint: disable=no-member
+        return None
+
+    _optionals.HAS_PIL.require_now("pass manager drawer")
+
+    with tempfile.TemporaryDirectory() as tmpdirname:
+        from PIL import Image
+
+        tmppath = os.path.join(tmpdirname, "pass_manager.png")
+
+        # pylint says this isn't a method - it is
+        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

releasenotes/notes/staged-pass-manager-drawer-f1da0308895a30d9.yaml

@@ -0,0 +1,4 @@
+---
+features:
+  - |
+    Modifies the drawer for the `StagedPassManager` to add an outer box around the passes for each stage.