feat: visualize static operators and measurements as nodes

doc/releases/changelog-dev.md

@@ -4,6 +4,9 @@
 
 * Compiled programs can be visualized.
   [(#2213)](https://github.com/PennyLaneAI/catalyst/pull/2213)
+  [(#2229)](https://github.com/PennyLaneAI/catalyst/pull/2229)
+  [(#2214)](https://github.com/PennyLaneAI/catalyst/pull/2214)
+  [(#2118)](https://github.com/PennyLaneAI/catalyst/pull/2218)
 
 * Added ``catalyst.switch``, a qjit compatible, index-switch style control flow decorator.
   [(#2171)](https://github.com/PennyLaneAI/catalyst/pull/2171)

frontend/catalyst/python_interface/visualization/construct_circuit_dag.py

@@ -0,0 +1,153 @@
+# Copyright 2025 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Contains the ConstructCircuitDAG tool for constructing a DAG from an xDSL module."""
+
+from functools import singledispatchmethod
+from typing import Any
+
+from xdsl.dialects import builtin
+from xdsl.ir import Block, Operation, Region
+
+from catalyst.python_interface.dialects import quantum
+from catalyst.python_interface.visualization.dag_builder import DAGBuilder
+from catalyst.python_interface.visualization.xdsl_conversion import (
+    xdsl_to_qml_measurement,
+    xdsl_to_qml_op,
+)
+
+
+class ConstructCircuitDAG:
+    """A tool that analyzes an xDSL module and constructs a Directed Acyclic Graph (DAG)
+    using an injected DAGBuilder instance. This tool does not mutate the xDSL module."""
+
+    def __init__(self, dag_builder: DAGBuilder) -> None:
+        """Initialize the analysis pass by injecting the DAG builder dependency.
+
+        Args:
+            dag_builder (DAGBuilder): The concrete builder instance used for graph construction.
+        """
+        self.dag_builder: DAGBuilder = dag_builder
+
+        # Record clusters seen as a stack
+        # beginning with the base graph (None)
+        self._cluster_stack: list[str | None] = [None]
+
+    def _reset(self) -> None:
+        """Resets the instance."""
+        self._cluster_stack: list[str | None] = [None]
+
+    # =================================
+    # 1. CORE DISPATCH AND ENTRY POINT
+    # =================================
+
+    @singledispatchmethod
+    def _visit(self, op: Any) -> None:
+        """Central dispatch method (Visitor Pattern). Routes the operation 'op'
+        to the specialized handler registered for its type."""
+
+    def construct(self, module: builtin.ModuleOp) -> None:
+        """Constructs the DAG from the module."""
+        self._reset()
+
+        for op in module.ops:
+            self._visit(op)
+
+    # =======================
+    # 2. IR TRAVERSAL
+    # =======================
+
+    @_visit.register
+    def _operation(self, operation: Operation) -> None:
+        """Visit an xDSL Operation."""
+        for region in operation.regions:
+            self._visit(region)
+
+    @_visit.register
+    def _region(self, region: Region) -> None:
+        """Visit an xDSL Region operation."""
+        for block in region.blocks:
+            self._visit(block)
+
+    @_visit.register
+    def _block(self, block: Block) -> None:
+        """Visit an xDSL Block operation, dispatching handling for each contained Operation."""
+        for op in block.ops:
+            self._visit(op)
+
+    # ======================
+    # 3. QUANTUM OPERATIONS
+    # ======================
+
+    @_visit.register
+    def _unitary(
+        self,
+        op: quantum.CustomOp
+        | quantum.GlobalPhaseOp
+        | quantum.QubitUnitaryOp
+        | quantum.MultiRZOp,
+    ) -> None:
+        """Generic handler for unitary gates and quantum state preparation operations."""
+
+        qml_op = xdsl_to_qml_op(op)
+        # Build node on graph
+        self.dag_builder.add_node(
+            id=f"node_{id(op)}",
+            label=str(qml_op),
+            cluster_id=self._cluster_stack[-1],
+        )
+
+    # ========================
+    # 4. QUANTUM MEASUREMENTS
+    # ========================
+
+    @_visit.register
+    def _state_op(self, op: quantum.StateOp) -> None:
+        """Handler for the terminal state measurement operation."""
+
+        meas = xdsl_to_qml_measurement(op)
+        # Build node on graph
+        self.dag_builder.add_node(
+            id=f"node_{id(op)}",
+            label=str(meas),
+            cluster_id=self._cluster_stack[-1],
+        )
+
+    @_visit.register
+    def _statistical_measurement_ops(
+        self,
+        op: quantum.ExpvalOp | quantum.VarianceOp | quantum.ProbsOp | quantum.SampleOp,
+    ) -> None:
+        """Handler for statistical measurement operations."""
+
+        obs_op = op.obs.owner
+        meas = xdsl_to_qml_measurement(op, xdsl_to_qml_measurement(obs_op))
+        # Build node on graph
+        self.dag_builder.add_node(
+            id=f"node_{id(op)}",
+            label=str(meas),
+            cluster_id=self._cluster_stack[-1],
+        )
+
+    @_visit.register
+    def _projective_measure_op(self, op: quantum.MeasureOp) -> None:
+        """Handler for the single-qubit projective measurement operation."""
+
+        meas = xdsl_to_qml_measurement(op)
+        # Build node on graph
+        self.dag_builder.add_node(
+            id=f"node_{id(op)}",
+            label=str(meas),
+            cluster_id=self._cluster_stack[-1],
+        )

frontend/catalyst/python_interface/visualization/dag_builder.py

@@ -14,7 +14,10 @@
 """File that defines the DAGBuilder abstract base class."""
 
 from abc import ABC, abstractmethod
-from typing import Any
+from typing import Any, TypeAlias
+
+ClusterID: TypeAlias = str
+NodeID: TypeAlias = str
 
 
 class DAGBuilder(ABC):
@@ -28,50 +31,83 @@ class DAGBuilder(ABC):
 
     @abstractmethod
     def add_node(
-        self, node_id: str, node_label: str, parent_graph_id: str | None = None, **node_attrs: Any
+        self,
+        id: NodeID,
+        label: str,
+        cluster_id: ClusterID | None = None,
+        **attrs: Any,
     ) -> None:
         """Add a single node to the graph.
 
         Args:
-            node_id (str): Unique node ID to identify this node.
-            node_label (str): The text to display on the node when rendered.
-            parent_graph_id (str | None): Optional ID of the cluster this node belongs to.
-            **node_attrs (Any): Any additional styling keyword arguments.
+            id (str): Unique node ID to identify this node.
+            label (str): The text to display on the node when rendered.
+            cluster_id (str | None): Optional ID of the cluster this node belongs to. If `None`, this node gets
+                added on the base graph.
+            **attrs (Any): Any additional styling keyword arguments.
 
         """
         raise NotImplementedError
 
     @abstractmethod
-    def add_edge(self, from_node_id: str, to_node_id: str, **edge_attrs: Any) -> None:
+    def add_edge(self, from_id: NodeID, to_id: NodeID, **attrs: Any) -> None:
         """Add a single directed edge between nodes in the graph.
 
         Args:
-            from_node_id (str): The unique ID of the source node.
-            to_node_id (str): The unique ID of the destination node.
-            **edge_attrs (Any): Any additional styling keyword arguments.
+            from_id (str): The unique ID of the source node.
+            to_id (str): The unique ID of the destination node.
+            **attrs (Any): Any additional styling keyword arguments.
 
         """
         raise NotImplementedError
 
     @abstractmethod
     def add_cluster(
         self,
-        cluster_id: str,
+        id: ClusterID,
         node_label: str | None = None,
-        parent_graph_id: str | None = None,
-        **cluster_attrs: Any,
+        cluster_id: ClusterID | None = None,
+        **attrs: Any,
     ) -> None:
         """Add a single cluster to the graph.
 
         A cluster is a specific type of subgraph where the nodes and edges contained
         within it are visually and logically grouped.
 
         Args:
-            cluster_id (str): Unique cluster ID to identify this cluster.
+            id (str): Unique cluster ID to identify this cluster.
             node_label (str | None): The text to display on an information node within the cluster when rendered.
-            parent_graph_id (str | None): Optional ID of the cluster this cluster belongs to.
-            **cluster_attrs (Any): Any additional styling keyword arguments.
+            cluster_id (str | None): Optional ID of the cluster this cluster belongs to. If `None`, the cluster will be
+                placed on the base graph.
+            **attrs (Any): Any additional styling keyword arguments.
+
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_nodes(self) -> dict[NodeID, dict[str, Any]]:
+        """Retrieve the current set of nodes in the graph.
+
+        Returns:
+            nodes (dict[str, dict[str, Any]]): A dictionary that maps the node's ID to it's node information.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_edges(self) -> list[dict[str, Any]]:
+        """Retrieve the current set of edges in the graph.
+
+        Returns:
+            edges (list[dict[str, Any]]): A list of edges where each element in the list contains a dictionary of edge information.
+        """
+        raise NotImplementedError
 
+    @abstractmethod
+    def get_clusters(self) -> dict[ClusterID, dict[str, Any]]:
+        """Retrieve the current set of clusters in the graph.
+
+        Returns:
+            clusters (dict[str, dict[str, Any]]): A dictionary that maps the cluster's ID to it's cluster information.
         """
         raise NotImplementedError