Remove pretty_name part 2 (#1400)

Change the default implementation of Bloq.pretty_name to go via __str__. Removing pretty_name implementations and moving them to __str__ or wire_symbols as necessary. Co-authored-by: Matthew Harrigan <[email protected]>
quantumlib · Sep 9, 2024 · 64fd963 · 64fd963
1 parent 83b7a82
commit 64fd963
Show file tree

Hide file tree

Showing 18 changed files with 81 additions and 50 deletions.
diff --git a/qualtran/bloqs/arithmetic/multiplication.py b/qualtran/bloqs/arithmetic/multiplication.py
@@ -12,7 +12,7 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
-from typing import Dict, Iterable, List, Sequence, TYPE_CHECKING, Union
+from typing import Dict, Iterable, List, Optional, Sequence, Tuple, TYPE_CHECKING, Union
 
 import cirq
 import numpy as np
@@ -34,6 +34,7 @@
 from qualtran.bloqs.arithmetic.subtraction import Subtract
 from qualtran.bloqs.basic_gates import CNOT, TGate, Toffoli, XGate
 from qualtran.bloqs.mcmt import MultiControlX
+from qualtran.drawing import Text, WireSymbol
 from qualtran.symbolics import ceil, HasLength, is_symbolic, log2, smax, SymbolicInt
 
 if TYPE_CHECKING:
@@ -72,8 +73,10 @@ def __attrs_post_init__(self):
                 f"bitsizes {self.a_bitsize} + {self.b_bitsize}"
             )
 
-    def pretty_name(self) -> str:
-        return "result -= a*b" if self.is_adjoint else "result += a*b"
+    def wire_symbol(self, reg: Optional[Register], idx: Tuple[int, ...] = tuple()) -> 'WireSymbol':
+        if reg is None:
+            return Text("result -= a*b") if self.is_adjoint else Text("result += a*b")
+        return super().wire_symbol(reg, idx)
 
     @property
     def signature(self) -> 'Signature':
@@ -186,8 +189,10 @@ def on_classical_vals(self, **vals: int) -> Dict[str, 'ClassicalValT']:
         a = vals["a"]
         return {'a': a, 'result': a**2}
 
-    def pretty_name(self) -> str:
-        return "a^2"
+    def wire_symbol(self, reg: Optional[Register], idx: Tuple[int, ...] = tuple()) -> 'WireSymbol':
+        if reg is None:
+            return Text("a^2")
+        return super().wire_symbol(reg, idx)
 
     def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT':
         # TODO Determine precise clifford count and/or ignore.
@@ -272,8 +277,10 @@ def signature(self):
             ]
         )
 
-    def pretty_name(self) -> str:
-        return "SOS"
+    def wire_symbol(self, reg: Optional[Register], idx: Tuple[int, ...] = tuple()) -> 'WireSymbol':
+        if reg is None:
+            return Text('SOS')
+        return super().wire_symbol(reg, idx)
 
     def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT':
         num_toff = self.k * self.bitsize**2 - self.bitsize
@@ -326,8 +333,10 @@ def signature(self):
             ]
         )
 
-    def pretty_name(self) -> str:
-        return "a*b"
+    def wire_symbol(self, reg: Optional[Register], idx: Tuple[int, ...] = tuple()) -> 'WireSymbol':
+        if reg is None:
+            return Text('a*b')
+        return super().wire_symbol(reg, idx)
 
     def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT':
         # TODO Determine precise clifford count and/or ignore.
@@ -387,8 +396,10 @@ def signature(self):
             ]
         )
 
-    def pretty_name(self) -> str:
-        return "r*i"
+    def wire_symbol(self, reg: Optional[Register], idx: Tuple[int, ...] = tuple()) -> 'WireSymbol':
+        if reg is None:
+            return Text('r*i')
+        return super().wire_symbol(reg, idx)
 
     def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT':
         # Eq. D8, we are assuming dA(r_bitsize) and dB(i_bitsize) are inputs and
@@ -444,8 +455,10 @@ def signature(self):
             ]
         )
 
-    def pretty_name(self) -> str:
-        return "a*b"
+    def wire_symbol(self, reg: Optional[Register], idx: Tuple[int, ...] = tuple()) -> 'WireSymbol':
+        if reg is None:
+            return Text('a*b')
+        return super().wire_symbol(reg, idx)
 
     def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT':
         # Eq. D13, there it is suggested keeping both registers the same size is optimal.
@@ -503,8 +516,10 @@ def signature(self):
             ]
         )
 
-    def pretty_name(self) -> str:
-        return "a^2"
+    def wire_symbol(self, reg: Optional[Register], idx: Tuple[int, ...] = tuple()) -> 'WireSymbol':
+        if reg is None:
+            return Text('a^2')
+        return super().wire_symbol(reg, idx)
 
     def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT':
         # Bottom of page 74
@@ -559,8 +574,10 @@ def signature(self):
             ]
         )
 
-    def pretty_name(self) -> str:
-        return "1/a"
+    def wire_symbol(self, reg: Optional[Register], idx: Tuple[int, ...] = tuple()) -> 'WireSymbol':
+        if reg is None:
+            return Text('1/a')
+        return super().wire_symbol(reg, idx)
 
     def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT':
         # initial approximation: Figure 4

diff --git a/qualtran/bloqs/block_encoding/chebyshev_polynomial.py b/qualtran/bloqs/block_encoding/chebyshev_polynomial.py
@@ -94,7 +94,7 @@ def signature(self) -> Signature:
         )
 
     def pretty_name(self) -> str:
-        return f"B[T_{self.order}({self.block_encoding.pretty_name()})]"
+        return f"B[T_{self.order}({self.block_encoding})]"
 
     @property
     def system_bitsize(self) -> SymbolicInt:
@@ -272,7 +272,7 @@ def build_composite_bloq(self, bb: BloqBuilder, **soqs: SoquetT) -> Dict[str, So
         return bb.add_d(self.linear_combination, **soqs)
 
     def __str__(self) -> str:
-        return f"B[T_{self.order}({self.block_encoding.pretty_name()})]"
+        return f"B[T_{self.order}({self.block_encoding})]"
 
 
 @bloq_example

diff --git a/qualtran/bloqs/block_encoding/linear_combination.py b/qualtran/bloqs/block_encoding/linear_combination.py
@@ -329,7 +329,7 @@ def build_composite_bloq(
         return out
 
     def __str__(self) -> str:
-        return f"B[{'+'.join(be.pretty_name()[2:-1] for be in self.signed_block_encodings)}]"
+        return f"B[{'+'.join(str(be)[2:-1] for be in self.signed_block_encodings)}]"
 
 
 @bloq_example

diff --git a/qualtran/bloqs/block_encoding/phase.py b/qualtran/bloqs/block_encoding/phase.py
@@ -88,7 +88,7 @@ def build_composite_bloq(self, bb: BloqBuilder, **soqs: SoquetT) -> Dict[str, So
         return bb.add_d(self.block_encoding, **soqs)
 
     def __str__(self) -> str:
-        return f"B[exp({self.phi}i){self.block_encoding.pretty_name()[2:-1]}]"
+        return f"B[exp({self.phi}i){str(self.block_encoding)[2:-1]}]"
 
 
 @bloq_example

diff --git a/qualtran/bloqs/block_encoding/product.py b/qualtran/bloqs/block_encoding/product.py
@@ -248,7 +248,7 @@ def build_composite_bloq(
         return out
 
     def __str__(self) -> str:
-        return f"B[{'*'.join(u.pretty_name()[2:-1] for u in self.block_encodings)}]"
+        return f"B[{'*'.join(str(u)[2:-1] for u in self.block_encodings)}]"
 
 
 @bloq_example(generalizer=ignore_split_join)

diff --git a/qualtran/bloqs/block_encoding/tensor_product.py b/qualtran/bloqs/block_encoding/tensor_product.py
@@ -176,7 +176,7 @@ def build_composite_bloq(
         return soqs_out
 
     def __str__(self) -> str:
-        return f"B[{'⊗'.join(u.pretty_name()[2:-1] for u in self.block_encodings)}]"
+        return f"B[{'⊗'.join(str(u)[2:-1] for u in self.block_encodings)}]"
 
 
 @bloq_example

diff --git a/qualtran/bloqs/block_encoding/unitary.py b/qualtran/bloqs/block_encoding/unitary.py
@@ -87,7 +87,7 @@ def build_composite_bloq(
         }
 
     def __str__(self) -> str:
-        return f"B[{self.U.pretty_name()}]"
+        return f"B[{self.U}]"
 
 
 @bloq_example

diff --git a/qualtran/bloqs/chemistry/thc/prepare.py b/qualtran/bloqs/chemistry/thc/prepare.py
@@ -106,7 +106,7 @@ def signature(self) -> Signature:
             ]
         )
 
-    def pretty_name(self) -> str:
+    def __str__(self) -> str:
         return r'$\sum_{\mu < \nu} |\mu\nu\rangle$'
 
     def wire_symbol(self, reg: Optional[Register], idx: Tuple[int, ...] = tuple()) -> 'WireSymbol':

diff --git a/qualtran/bloqs/chemistry/thc/select_bloq.py b/qualtran/bloqs/chemistry/thc/select_bloq.py
@@ -91,7 +91,7 @@ def signature(self) -> Signature:
     def adjoint(self) -> 'Bloq':
         return evolve(self, is_adjoint=not self.is_adjoint)
 
-    def pretty_name(self) -> str:
+    def __str__(self) -> str:
         dag = '†' if self.is_adjoint else ''
         return f"In_mu-R{dag}"
 

diff --git a/qualtran/bloqs/mcmt/multi_control_pauli.py b/qualtran/bloqs/mcmt/multi_control_pauli.py
@@ -109,7 +109,7 @@ def build_composite_bloq(self, bb: 'BloqBuilder', **soqs: 'SoquetT') -> Dict[str
 
         return {'controls': out_soqs[ctrl_reg_name], 'target': out_soqs[target_reg_name]}
 
-    def pretty_name(self) -> str:
+    def __str__(self) -> str:
         n = self.n_ctrls
         ctrl = f'C^{n}' if is_symbolic(n) or n > 2 else ['', 'C', 'CC'][int(n)]
         return f'{ctrl}{self.target_gate!s}'

diff --git a/qualtran/bloqs/multiplexers/black_box_select.py b/qualtran/bloqs/multiplexers/black_box_select.py
@@ -56,7 +56,7 @@ class BlackBoxSelect(Bloq):
 
     select: SelectOracle
 
-    def pretty_name(self) -> str:
+    def __str__(self) -> str:
         return 'SELECT'
 
     @cached_property

diff --git a/qualtran/bloqs/phase_estimation/lp_resource_state.py b/qualtran/bloqs/phase_estimation/lp_resource_state.py
@@ -15,16 +15,17 @@
 """Resource states proposed by A. Luis and J. Peřina (1996) for optimal phase measurements"""
 from collections import Counter
 from functools import cached_property
-from typing import Dict, TYPE_CHECKING
+from typing import Dict, Optional, Tuple, TYPE_CHECKING
 
 import attrs
 import numpy as np
 import sympy
 
-from qualtran import Bloq, bloq_example, BloqDocSpec, GateWithRegisters, QBit, Signature
+from qualtran import Bloq, bloq_example, BloqDocSpec, GateWithRegisters, QBit, Register, Signature
 from qualtran.bloqs.basic_gates import CZ, Hadamard, OnEach, Ry, Rz, XGate
 from qualtran.bloqs.phase_estimation.qpe_window_state import QPEWindowStateBase
 from qualtran.bloqs.reflections.reflection_using_prepare import ReflectionUsingPrepare
+from qualtran.drawing import Text, WireSymbol
 from qualtran.symbolics import acos, ceil, is_symbolic, log2, pi, SymbolicFloat, SymbolicInt
 
 if TYPE_CHECKING:
@@ -55,8 +56,10 @@ class LPRSInterimPrep(GateWithRegisters):
     def signature(self) -> 'Signature':
         return Signature.build(m=self.bitsize, anc=1)
 
-    def pretty_name(self) -> str:
-        return 'LPRS'
+    def wire_symbol(self, reg: Optional[Register], idx: Tuple[int, ...] = tuple()) -> 'WireSymbol':
+        if reg is None:
+            return Text('LPRS')
+        return super().wire_symbol(reg, idx)
 
     def build_composite_bloq(
         self, bb: 'BloqBuilder', *, m: 'SoquetT', anc: 'Soquet'

diff --git a/qualtran/bloqs/qft/two_bit_ffft.py b/qualtran/bloqs/qft/two_bit_ffft.py
@@ -12,7 +12,7 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 from functools import cached_property
-from typing import Dict, List, TYPE_CHECKING
+from typing import Dict, List, Optional, Tuple, TYPE_CHECKING
 
 import numpy as np
 from attrs import frozen
@@ -35,6 +35,7 @@
 from qualtran.bloqs.basic_gates.rotation import Rz
 from qualtran.bloqs.basic_gates.s_gate import SGate
 from qualtran.bloqs.basic_gates.t_gate import TGate
+from qualtran.drawing import Text, WireSymbol
 from qualtran.resource_counting import SympySymbolAllocator
 from qualtran.symbolics.types import SymbolicFloat
 
@@ -86,8 +87,10 @@ def __attrs_post_init__(self):
     def signature(self) -> Signature:
         return Signature([Register('x', QBit()), Register('y', QBit())])
 
-    def pretty_name(self) -> str:
-        return 'F(k, n)'
+    def wire_symbol(self, reg: Optional[Register], idx: Tuple[int, ...] = tuple()) -> 'WireSymbol':
+        if reg is None:
+            return Text('F(k, n)')
+        return super().wire_symbol(reg, idx)
 
     def my_tensors(
         self, incoming: Dict[str, 'ConnectionT'], outgoing: Dict[str, 'ConnectionT']

diff --git a/qualtran/bloqs/rotations/hamming_weight_phasing.py b/qualtran/bloqs/rotations/hamming_weight_phasing.py
@@ -13,7 +13,7 @@
 #  limitations under the License.
 
 from functools import cached_property
-from typing import Dict, TYPE_CHECKING
+from typing import Dict, Optional, Tuple, TYPE_CHECKING
 
 import attrs
 import numpy as np
@@ -22,6 +22,7 @@
 from qualtran.bloqs.arithmetic import HammingWeightCompute
 from qualtran.bloqs.basic_gates import ZPowGate
 from qualtran.bloqs.rotations.quantum_variable_rotation import QvrPhaseGradient
+from qualtran.drawing import Text, WireSymbol
 from qualtran.symbolics import SymbolicFloat, SymbolicInt
 
 if TYPE_CHECKING:
@@ -93,8 +94,10 @@ def build_composite_bloq(self, bb: 'BloqBuilder', **soqs: 'SoquetT') -> Dict[str
         )
         return soqs
 
-    def pretty_name(self) -> str:
-        return f'HWP_{self.bitsize}(Z^{self.exponent})'
+    def wire_symbol(self, reg: Optional[Register], idx: Tuple[int, ...] = tuple()) -> 'WireSymbol':
+        if reg is None:
+            return Text(f'HWP_{self.bitsize}(Z^{self.exponent})')
+        return super().wire_symbol(reg, idx)
 
     def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT':
         return {
@@ -190,8 +193,10 @@ def build_composite_bloq(
         x = bb.add(HammingWeightCompute(self.bitsize).adjoint(), x=x, junk=junk, out=out)
         return {'x': x, 'phase_grad': phase_grad}
 
-    def pretty_name(self) -> str:
-        return f'HWPG_{self.bitsize}(Z^{self.exponent})'
+    def wire_symbol(self, reg: Optional[Register], idx: Tuple[int, ...] = tuple()) -> 'WireSymbol':
+        if reg is None:
+            return Text(f'HWPG_{self.bitsize}(Z^{self.exponent})')
+        return super().wire_symbol(reg, idx)
 
 
 @bloq_example

diff --git a/qualtran/bloqs/rotations/phase_gradient.py b/qualtran/bloqs/rotations/phase_gradient.py
@@ -36,6 +36,7 @@
 from qualtran.bloqs.basic_gates import Hadamard, Toffoli
 from qualtran.bloqs.basic_gates.on_each import OnEach
 from qualtran.bloqs.basic_gates.rotation import CZPowGate, ZPowGate
+from qualtran.drawing import Text, WireSymbol
 from qualtran.resource_counting import SympySymbolAllocator
 from qualtran.symbolics.types import is_symbolic
 
@@ -259,9 +260,11 @@ class AddIntoPhaseGrad(GateWithRegisters, cirq.ArithmeticGate):  # type: ignore[
     sign: int = +1
     controlled_by: Optional[int] = None
 
-    def pretty_name(self) -> str:
+    def wire_symbol(self, reg: Optional[Register], idx: Tuple[int, ...] = tuple()) -> 'WireSymbol':
         sign = '+' if self.sign > 0 else '-'
-        return f'pg{sign}=x>>{self.right_shift}' if self.right_shift else f'pg{sign}=x'
+        if reg is None:
+            return Text(f'pg{sign}=x>>{self.right_shift}' if self.right_shift else f'pg{sign}=x')
+        return super().wire_symbol(reg, idx)
 
     @cached_property
     def signature(self) -> 'Signature':

diff --git a/qualtran/bloqs/state_preparation/black_box_prepare.py b/qualtran/bloqs/state_preparation/black_box_prepare.py
@@ -86,7 +86,7 @@ def build_composite_bloq(self, bb: BloqBuilder, **soqs: SoquetT) -> Dict[str, So
             )
         return bb.add_d(AutoPartition(self.prepare, partitions), **soqs)
 
-    def pretty_name(self) -> str:
+    def __str__(self) -> str:
         return 'Prep'
 
 

diff --git a/qualtran/bloqs/swap_network/cswap_approx.py b/qualtran/bloqs/swap_network/cswap_approx.py
@@ -13,16 +13,17 @@
 #  limitations under the License.
 
 from functools import cached_property
-from typing import Dict, Iterator, TYPE_CHECKING
+from typing import Dict, Iterator, Optional, Tuple, TYPE_CHECKING
 
 import cirq
 from attrs import frozen
 from numpy.typing import NDArray
 
-from qualtran import bloq_example, BloqDocSpec, GateWithRegisters, Signature
+from qualtran import bloq_example, BloqDocSpec, GateWithRegisters, Register, Signature
 from qualtran.bloqs.basic_gates import TGate
 from qualtran.bloqs.bookkeeping import ArbitraryClifford
 from qualtran.bloqs.mcmt.multi_target_cnot import MultiTargetCNOT
+from qualtran.drawing import Text, WireSymbol
 from qualtran.resource_counting.generalizers import (
     cirq_to_bloqs,
     generalize_rotation_angle,
@@ -93,8 +94,10 @@ def on_classical_vals(
             return {'ctrl': 1, 'x': y, 'y': x}
         raise ValueError("Bad control value for CSwap classical simulation.")
 
-    def pretty_name(self) -> str:
-        return '~swap'
+    def wire_symbol(self, reg: Optional[Register], idx: Tuple[int, ...] = tuple()) -> 'WireSymbol':
+        if reg is None:
+            return Text('~swap')
+        return super().wire_symbol(reg, idx)
 
     def _circuit_diagram_info_(self, args: cirq.CircuitDiagramInfoArgs) -> cirq.CircuitDiagramInfo:
         if not args.use_unicode_characters:

diff --git a/qualtran/bloqs/swap_network/swap_with_zero.py b/qualtran/bloqs/swap_network/swap_with_zero.py
@@ -213,9 +213,6 @@ def on_classical_vals(
     def adjoint(self) -> 'SwapWithZero':
         return attrs.evolve(self, uncompute=not self.uncompute)
 
-    def pretty_name(self) -> str:
-        return 'SwapWithZero†' if self.uncompute else 'SwapWithZero'
-
     def __str__(self) -> str:
         return 'SwapWithZero†' if self.uncompute else 'SwapWithZero'