Merge branch 'main' into 2024/08/06-t-compl-use-rotation-eps

Author: anurudhp

qualtran/bloqs/for_testing/atom.py

@@ -28,6 +28,7 @@
     Signature,
 )
 from qualtran.cirq_interop.t_complexity_protocol import TComplexity
+from qualtran.resource_counting import CostKey, GateCounts, QECGatesCost
 
 if TYPE_CHECKING:
     import quimb.tensor as qtn
@@ -66,6 +67,11 @@ def my_tensors(
             )
         ]
 
+    def my_static_costs(self, cost_key: 'CostKey'):
+        if cost_key == QECGatesCost():
+            return GateCounts(t=100)
+        return NotImplemented
+
     def _t_complexity_(self) -> 'TComplexity':
         return TComplexity(100)
 

qualtran/drawing/_show_funcs.py

@@ -15,12 +15,14 @@
 """Convenience functions for showing rich displays in Jupyter notebook."""
 
 import os
-from typing import Dict, Optional, Sequence, TYPE_CHECKING, Union
+from typing import Dict, Optional, overload, Sequence, TYPE_CHECKING, Union
 
 import IPython.display
 import ipywidgets
 
-from .bloq_counts_graph import format_counts_sigma, GraphvizCounts
+from qualtran import Bloq
+
+from .bloq_counts_graph import format_counts_sigma, GraphvizCallGraph, GraphvizCounts
 from .flame_graph import get_flame_graph_svg_data
 from .graphviz import PrettyGraphDrawer, TypedGraphDrawer
 from .musical_score import draw_musical_score, get_musical_score_data
@@ -30,8 +32,6 @@
     import networkx as nx
     import sympy
 
-    from qualtran import Bloq
-
 
 def show_bloq(bloq: 'Bloq', type: str = 'graph'):  # pylint: disable=redefined-builtin
     """Display a visual representation of the bloq in IPython.
@@ -75,9 +75,51 @@ def show_bloqs(bloqs: Sequence['Bloq'], labels: Optional[Sequence[Optional[str]]
     IPython.display.display(box)
 
 
-def show_call_graph(g: 'nx.DiGraph') -> None:
-    """Display a graph representation of the counts graph `g`."""
-    IPython.display.display(GraphvizCounts(g).get_svg())
+@overload
+def show_call_graph(
+    item: 'Bloq', /, *, max_depth: Optional[int] = None, agg_gate_counts: Optional[str] = None
+) -> None:
+    ...
+
+
+@overload
+def show_call_graph(
+    item: 'nx.Graph', /, *, max_depth: Optional[int] = None, agg_gate_counts: Optional[str] = None
+) -> None:
+    ...
+
+
+def show_call_graph(
+    item: Union['Bloq', 'nx.Graph'],
+    /,
+    *,
+    max_depth: Optional[int] = None,
+    agg_gate_counts: Optional[str] = None,
+) -> None:
+    """Display a graph representation of the call graph.
+
+    Args:
+        item: Either a networkx graph or a bloq. If a networkx graph, it should be a "call graph"
+            which is passed verbatim to the graph drawer and the additional arguments to this
+            function are ignored. If it is a bloq, the factory
+            method `GraphvizCallGraph.from_bloq()` is used to construct the call graph, compute
+            relevant costs, and display the call graph annotated with the costs.
+        max_depth: The maximum depth (from the root bloq) of the call graph to draw. Note
+            that the cost computations will walk the whole call graph, but only the nodes
+            within this depth will be drawn.
+        agg_gate_counts: One of 'factored', 'total_t', 't_and_ccz', or 'beverland' to
+            (optionally) aggregate the gate counts. If not specified, the 'factored'
+            approach is used where each type of gate is counted individually.
+
+    """
+    if isinstance(item, Bloq):
+        IPython.display.display(
+            GraphvizCallGraph.from_bloq(
+                item, max_depth=max_depth, agg_gate_counts=agg_gate_counts
+            ).get_svg()
+        )
+    else:
+        IPython.display.display(GraphvizCounts(item).get_svg())
 
 
 def show_counts_sigma(sigma: Dict['Bloq', Union[int, 'sympy.Expr']]):

qualtran/drawing/bloq_counts_graph.py

@@ -15,7 +15,8 @@
 """Classes for drawing bloq counts graphs with Graphviz."""
 import abc
 import html
-from typing import Any, Dict, Iterable, Optional, Tuple, Union
+import warnings
+from typing import Any, cast, Dict, Iterable, Mapping, Optional, Tuple, TYPE_CHECKING, Union
 
 import attrs
 import IPython.display
@@ -24,6 +25,10 @@
 import sympy
 
 from qualtran import Bloq, CompositeBloq
+from qualtran.symbolics import SymbolicInt
+
+if TYPE_CHECKING:
+    from qualtran.resource_counting import CostKey, CostValT, GateCounts
 
 
 class _CallGraphDrawerBase(metaclass=abc.ABCMeta):
@@ -176,7 +181,15 @@ class GraphvizCallGraph(_CallGraphDrawerBase):
     Each edge is labeled with the number of times the "caller" (predecessor) bloq calls the
     "callee" (successor) bloq.
 
-    This class follows the behavior described in https://github.com/quantumlib/Qualtran/issues/791
+    The constructor of this class assumes you have already generated the call graph as a networkx
+    graph and constructed any associated data. See the factory method
+    `GraphvizCallGraph.from_bloq()` to set up a call graph diagram from a bloq with sensible
+    defaults.
+
+    This class uses a bloq's `__str__` string to title the bloq. Arbitrary additional tabular
+    data can be provided with `bloq_data`.
+
+    This graph drawer is the successor to the `GraphvizCounts` existing drawer,
     and will replace `GraphvizCounts` when all bloqs have been migrated to use `__str__()`.
 
     Args:
@@ -193,6 +206,133 @@ def __init__(self, g: nx.DiGraph, bloq_data: Optional[Dict['Bloq', Dict[Any, Any
 
         self.bloq_data = bloq_data
 
+    @classmethod
+    def format_qubit_count(cls, val: SymbolicInt) -> Dict[str, str]:
+        """Format `QubitCount` cost values as a string.
+
+        Args:
+            val: The qubit count value, which should be an integer
+
+        Returns:
+            A dictionary mapping a string cost name to a string cost value.
+        """
+        return {'Qubits': f'{val}'}
+
+    @classmethod
+    def format_qec_gates_cost(cls, val: 'GateCounts', agg: Optional[str] = None) -> Dict[str, str]:
+        """Format `QECGatesCost` cost values as a string.
+
+        Args:
+            val: The qec gate costs value, which should be a `GateCounts` dataclass.
+            agg: One of 'factored', 'total_t', 't_and_ccz', or 'beverland' to
+                (optionally) aggregate the gate counts. If not specified, the 'factored'
+                approach is used where each type of gate is counted individually. See the
+                methods on `GateCounts` for more information.
+
+        Returns:
+            A dictionary mapping string cost names to string cost values.
+        """
+        labels = {
+            't': 'Ts',
+            'n_t': 'Ts',
+            'toffoli': 'Toffolis',
+            'n_ccz': 'CCZs',
+            'cswap': 'CSwaps',
+            'and_bloq': 'Ands',
+            'clifford': 'Cliffords',
+            'rotation': 'Rotations',
+            'measurement': 'Measurements',
+        }
+        counts_dict: Mapping[str, SymbolicInt]
+        if agg is None or agg == 'factored':
+            counts_dict = val.asdict()
+        elif agg == 'total_t':
+            counts_dict = {'t': val.total_t_count()}
+        elif agg == 't_and_ccz':
+            counts_dict = val.total_t_and_ccz_count()
+        elif agg == 'beverland':
+            counts_dict = val.total_beverland_count()
+        else:
+            raise ValueError(f"Unknown aggregation mode {agg}.")
+
+        return {labels.get(gate_k, gate_k): f'{gate_v}' for gate_k, gate_v in counts_dict.items()}
+
+    @classmethod
+    def format_cost_data(
+        cls,
+        cost_data: Dict['Bloq', Dict['CostKey', 'CostValT']],
+        agg_gate_counts: Optional[str] = None,
+    ) -> Dict['Bloq', Dict[str, str]]:
+        """Format `cost_data` as human-readable strings.
+
+        Args:
+            cost_data: The cost data, likely returned from a call to `query_costs()`. This
+                class method will delegate to `format_qubit_count` and `format_qec_gates_cost`
+                for `QubitCount` and `QECGatesCost` cost keys, respectively.
+            agg_gate_counts: One of 'factored', 'total_t', 't_and_ccz', or 'beverland' to
+                (optionally) aggregate the gate counts. If not specified, the 'factored'
+                approach is used where each type of gate is counted individually. See the
+                methods on `GateCounts` for more information.
+
+        Returns:
+            For each bloq key, a table of label/value pairs consisting of
+            human-readable labels and formatted values.
+        """
+        from qualtran.resource_counting import GateCounts, QECGatesCost, QubitCount
+
+        disp_data: Dict['Bloq', Dict[str, str]] = {}
+        for bloq in cost_data.keys():
+            bloq_disp: Dict[str, str] = {}
+            for cost_key, cost_val in cost_data[bloq].items():
+                if isinstance(cost_key, QubitCount):
+                    bloq_disp |= cls.format_qubit_count(cast(SymbolicInt, cost_val))
+                elif isinstance(cost_key, QECGatesCost):
+                    assert isinstance(cost_val, GateCounts)
+                    bloq_disp |= cls.format_qec_gates_cost(cost_val, agg=agg_gate_counts)
+                else:
+                    warnings.warn(f"Unknown cost key {cost_key}")
+                    bloq_disp[str(cost_key)] = str(cost_val)
+
+            disp_data[bloq] = bloq_disp
+        return disp_data
+
+    @classmethod
+    def from_bloq(
+        cls, bloq: Bloq, *, max_depth: Optional[int] = None, agg_gate_counts: Optional[str] = None
+    ) -> 'GraphvizCallGraph':
+        """Draw a bloq call graph.
+
+        This factory method will generate a call graph from the bloq, query the `QECGatesCost`
+        and `QubitCount` costs, format the cost data, and merge it with the call graph
+        to create a call graph diagram with annotated costs.
+
+        For additional customization, users can construct the call graph and bloq data themselves
+        and use the normal constructor, or provide minor display customizations by
+        overriding the `format_xxx` class methods.
+
+        Args:
+            bloq: The bloq from which we construct the call graph and query the costs.
+            max_depth: The maximum depth (from the root bloq) of the call graph to draw. Note
+                that the cost computations will walk the whole call graph, but only the nodes
+                within this depth will be drawn.
+            agg_gate_counts: One of 'factored', 'total_t', 't_and_ccz', or 'beverland' to
+                (optionally) aggregate the gate counts. If not specified, the 'factored'
+                approach is used where each type of gate is counted individually. See the
+                methods on `GateCounts` for more information.
+
+        Returns:
+            A `GraphvizCallGraph` diagram-drawer, whose methods can be used to generate
+            graphviz inputs or SVG diagrams.
+        """
+        from qualtran.resource_counting import QECGatesCost, QubitCount, query_costs
+
+        call_graph, _ = bloq.call_graph(max_depth=max_depth)
+        cost_data: Dict['Bloq', Dict[CostKey, Any]] = query_costs(
+            bloq, [QubitCount(), QECGatesCost()]
+        )
+        formatted_cost_data = cls.format_cost_data(cost_data, agg_gate_counts=agg_gate_counts)
+        return cls(g=call_graph, bloq_data=formatted_cost_data)
+
     def get_node_title(self, b: Bloq):
         return str(b)
 

qualtran/drawing/bloq_counts_graph_test.py

@@ -15,6 +15,8 @@
 import re
 from typing import List
 
+import networkx as nx
+
 from qualtran.bloqs.for_testing import TestBloqWithCallGraph
 from qualtran.bloqs.mcmt.and_bloq import MultiAnd
 from qualtran.drawing import (
@@ -128,3 +130,47 @@ def test_graphviz_call_graph_with_data():
                 '<tr><td colspan="2"><font point-size="10">TestBloqWithCallGraph</font></td></tr>\n'
                 '<tr><td>T count</td><td>100*_n0 + 600</td></tr><tr><td>clifford</td><td>0</td></tr><tr><td>rot</td><td>0</td></tr></table></font>>'
             )
+
+
+def test_graphviz_call_graph_from_bloq():
+    bloq = TestBloqWithCallGraph()
+    drawer = GraphvizCallGraph.from_bloq(bloq)
+
+    node_labels = _get_node_labels_from_pydot_graph(drawer)
+    for nl in node_labels:
+        # Spot check one of the nodes
+        if 'TestBloqWithCallGraph' in nl:
+            assert nl == (
+                '<<font point-size="10"><table border="0" cellborder="1" cellspacing="0" cellpadding="5">\n'
+                '<tr><td colspan="2"><font point-size="10">TestBloqWithCallGraph</font></td></tr>\n'
+                '<tr><td>Qubits</td><td>3</td></tr>'
+                '<tr><td>Ts</td><td>100*_n0 + 600</td></tr>'
+                '</table></font>>'
+            )
+
+
+def test_graphviz_call_graph_from_bloq_agg():
+    bloq = TestBloqWithCallGraph()
+    drawer = GraphvizCallGraph.from_bloq(bloq, agg_gate_counts='t_and_ccz')
+
+    node_labels = _get_node_labels_from_pydot_graph(drawer)
+    for nl in node_labels:
+        # Spot check one of the nodes
+        # Note the additional cell.
+        if 'TestBloqWithCallGraph' in nl:
+            assert nl == (
+                '<<font point-size="10"><table border="0" cellborder="1" cellspacing="0" cellpadding="5">\n'
+                '<tr><td colspan="2"><font point-size="10">TestBloqWithCallGraph</font></td></tr>\n'
+                '<tr><td>Qubits</td><td>3</td></tr>'
+                '<tr><td>Ts</td><td>100*_n0 + 600</td></tr>'
+                '<tr><td>CCZs</td><td>0</td></tr>'
+                '</table></font>>'
+            )
+
+
+def test_graphviz_call_graph_from_bloq_max_depth():
+    bloq = TestBloqWithCallGraph()
+    drawer = GraphvizCallGraph.from_bloq(bloq)
+    assert len(list(nx.topological_generations(drawer.g))) == 3
+    drawer2 = GraphvizCallGraph.from_bloq(bloq, max_depth=1)
+    assert len(list(nx.topological_generations(drawer2.g))) == 2