Change wire_symbol to use Register and idx (#919)

* Change wire_symbol to use Register and idx

- Change wire_symbol to use Register and idx since the
wire_symbol should never need to use the binst of a Soquet.

Author: dstrain115

qualtran/_infra/adjoint.py

@@ -15,7 +15,6 @@
 from functools import cached_property
 from typing import Dict, List, Set, Tuple, TYPE_CHECKING
 
-import attrs
 import cirq
 from attrs import frozen
 from numpy.typing import NDArray
@@ -26,7 +25,7 @@
 from .registers import Signature
 
 if TYPE_CHECKING:
-    from qualtran import Bloq, CompositeBloq, Signature, Soquet, SoquetT
+    from qualtran import Bloq, CompositeBloq, Register, Signature, Soquet, SoquetT
     from qualtran.drawing import WireSymbol
     from qualtran.resource_counting import BloqCountT, SympySymbolAllocator
 
@@ -183,11 +182,11 @@ def __str__(self) -> str:
         """Delegate to subbloq's `__str__` method."""
         return f'Adjoint(subbloq={str(self.subbloq)})'
 
-    def wire_symbol(self, soq: 'Soquet') -> 'WireSymbol':
+    def wire_symbol(self, reg: 'Register', idx: Tuple[int, ...] = tuple()) -> 'WireSymbol':
         # Note: since we pass are passed a soquet which has the 'new' side, we flip it before
         # delegating and then flip back. Subbloqs only have to answer this protocol
         # if the provided soquet is facing the correct direction.
-        return self.subbloq.wire_symbol(attrs.evolve(soq, reg=soq.reg.adjoint())).adjoint()
+        return self.subbloq.wire_symbol(reg=reg.adjoint(), idx=idx).adjoint()
 
     def _t_complexity_(self):
         """The cirq-style `_t_complexity_` delegates to the subbloq's method with a special shim.

qualtran/_infra/adjoint_test.py

@@ -18,7 +18,7 @@
 import sympy
 
 import qualtran.testing as qlt_testing
-from qualtran import Adjoint, Bloq, BloqInstance, CompositeBloq, Side, Signature, Soquet
+from qualtran import Adjoint, Bloq, CompositeBloq, Side, Signature
 from qualtran._infra.adjoint import _adjoint_cbloq
 from qualtran.bloqs.basic_gates import CNOT, CSwap, ZeroState
 from qualtran.bloqs.for_testing.atom import TestAtom
@@ -162,11 +162,8 @@ def test_wire_symbol():
     (reg,) = zero.signature
     adj = Adjoint(zero)  # specifically use the Adjoint wrapper for testing
 
-    # TODO: Remove binst variable.  These BloqInstances are for typing only
-    # and are not really used by the function.
-    # See https://github.com/quantumlib/Qualtran/issues/608
-    ws = zero.wire_symbol(Soquet(BloqInstance(CNOT(), 1), reg))
-    adj_ws = adj.wire_symbol(Soquet(BloqInstance(CNOT(), 2), reg.adjoint()))
+    ws = zero.wire_symbol(reg)
+    adj_ws = adj.wire_symbol(reg.adjoint())
     assert isinstance(ws, LarrowTextBox)
     assert isinstance(adj_ws, RarrowTextBox)
 

qualtran/_infra/bloq.py

@@ -30,6 +30,7 @@
         BloqBuilder,
         CompositeBloq,
         CtrlSpec,
+        Register,
         Signature,
         Soquet,
         SoquetT,
@@ -503,7 +504,7 @@ def on_registers(
 
         return self.on(*merge_qubits(self.signature, **qubit_regs))
 
-    def wire_symbol(self, soq: 'Soquet') -> 'WireSymbol':
+    def wire_symbol(self, reg: 'Register', idx: Tuple[int, ...] = tuple()) -> 'WireSymbol':
         """On a musical score visualization, use this `WireSymbol` to represent `soq`.
 
         By default, we use a "directional text box", which is a text box that is either
@@ -516,4 +517,10 @@ def wire_symbol(self, soq: 'Soquet') -> 'WireSymbol':
         """
         from qualtran.drawing import directional_text_box
 
-        return directional_text_box(text=soq.pretty(), side=soq.reg.side)
+        label = reg.name
+        if len(idx) > 0:
+            pretty_str = f'{label}[{", ".join(str(i) for i in idx)}]'
+        else:
+            pretty_str = label
+
+        return directional_text_box(text=pretty_str, side=reg.side)

qualtran/_infra/controlled.py

@@ -166,15 +166,15 @@ def is_active(self, *vals: 'ClassicalValT') -> bool:
                 return False
         return True
 
-    def wire_symbol(self, i: int, soq: 'Soquet') -> 'WireSymbol':
+    def wire_symbol(self, i: int, reg: Register, idx: Tuple[int, ...] = tuple()) -> 'WireSymbol':
         # Return a circle for bits; a box otherwise.
         from qualtran.drawing import Circle, TextBox
 
-        if soq.reg.bitsize == 1:
-            cv = self.cvs[i][soq.idx]
+        if reg.bitsize == 1:
+            cv = self.cvs[i][idx]
             return Circle(filled=(cv == 1))
 
-        cv = self.cvs[i][soq.idx]
+        cv = self.cvs[i][idx]
         return TextBox(f'{cv}')
 
     @cached_property
@@ -431,14 +431,16 @@ def _unitary_(self):
         # Unable to determine the unitary effect.
         return NotImplemented
 
-    def wire_symbol(self, soq: 'Soquet') -> 'WireSymbol':
-        if soq.reg.name not in self.ctrl_reg_names:
+    def wire_symbol(self, reg: Register, idx: Tuple[int, ...] = tuple()) -> 'WireSymbol':
+        if reg.name not in self.ctrl_reg_names:
             # Delegate to subbloq
-            return self.subbloq.wire_symbol(soq)
+            print(self.subbloq)
+            print(type(self.subbloq))
+            return self.subbloq.wire_symbol(reg, idx)
 
         # Otherwise, it's part of the control register.
-        i = self.ctrl_reg_names.index(soq.reg.name)
-        return self.ctrl_spec.wire_symbol(i, soq)
+        i = self.ctrl_reg_names.index(reg.name)
+        return self.ctrl_spec.wire_symbol(i, reg, idx)
 
     def adjoint(self) -> 'Bloq':
         return self.subbloq.adjoint().controlled(self.ctrl_spec)

qualtran/_infra/gate_with_registers.py

@@ -33,7 +33,6 @@
 
 from qualtran._infra.bloq import Bloq, DecomposeNotImplementedError, DecomposeTypeError
 from qualtran._infra.composite_bloq import CompositeBloq
-from qualtran._infra.quantum_graph import Soquet
 from qualtran._infra.registers import Register, Side
 
 if TYPE_CHECKING:
@@ -309,10 +308,10 @@ def as_cirq_op(
         )
         return self.on_registers(**all_quregs), out_quregs
 
-    def wire_symbol(self, soq: 'Soquet') -> 'WireSymbol':
+    def wire_symbol(self, reg: Register, idx: Tuple[int, ...] = tuple()) -> 'WireSymbol':
         from qualtran.cirq_interop._cirq_to_bloq import _wire_symbol_from_gate
 
-        return _wire_symbol_from_gate(self, self.signature, soq)
+        return _wire_symbol_from_gate(self, self.signature, reg, idx)
 
     # Part-2: Cirq-FT style interface can be used to implemented algorithms by Bloq authors.
 

qualtran/bloqs/arithmetic/addition.py

@@ -150,13 +150,13 @@ def _circuit_diagram_info_(self, _) -> cirq.CircuitDiagramInfo:
         wire_symbols += ["In(y)/Out(x+y)"] * int(self.b_dtype.bitsize)
         return cirq.CircuitDiagramInfo(wire_symbols=wire_symbols)
 
-    def wire_symbol(self, soq: 'Soquet') -> 'WireSymbol':
+    def wire_symbol(self, reg: Register, idx: Tuple[int, ...] = tuple()) -> 'WireSymbol':
         from qualtran.drawing import directional_text_box
 
-        if soq.reg.name == 'a':
-            return directional_text_box('a', side=soq.reg.side)
-        elif soq.reg.name == 'b':
-            return directional_text_box('a+b', side=soq.reg.side)
+        if reg.name == 'a':
+            return directional_text_box('a', side=reg.side)
+        elif reg.name == 'b':
+            return directional_text_box('a+b', side=reg.side)
         else:
             raise ValueError()
 

qualtran/bloqs/arithmetic/comparison.py

@@ -13,7 +13,7 @@
 #  limitations under the License.
 
 from functools import cached_property
-from typing import Dict, Iterable, Iterator, List, Sequence, Set, TYPE_CHECKING, Union
+from typing import Dict, Iterable, Iterator, List, Sequence, Set, Tuple, TYPE_CHECKING, Union
 
 import attrs
 import cirq
@@ -599,14 +599,14 @@ def _t_complexity_(self) -> 'TComplexity':
         # See: https://github.com/quantumlib/Qualtran/issues/217
         return t_complexity(LessThanEqual(self.a_bitsize, self.b_bitsize))
 
-    def wire_symbol(self, soq: Soquet) -> WireSymbol:
-        if soq.reg.name == 'a':
+    def wire_symbol(self, reg: Register, idx: Tuple[int, ...] = tuple()) -> WireSymbol:
+        if reg.name == 'a':
             return TextBox("In(a)")
-        if soq.reg.name == 'b':
+        if reg.name == 'b':
             return TextBox("In(b)")
-        elif soq.reg.name == 'target':
+        elif reg.name == 'target':
             return TextBox("⨁(a > b)")
-        raise ValueError(f'Unknown register name {soq.reg.name}')
+        raise ValueError(f'Unknown register name {reg.name}')
 
     def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']:
         # TODO Determine precise clifford count and/or ignore.
@@ -677,7 +677,6 @@ def on_classical_vals(
     def build_composite_bloq(
         self, bb: 'BloqBuilder', a: Soquet, b: Soquet, target: SoquetT
     ) -> Dict[str, 'SoquetT']:
-
         # Base Case: Comparing two qubits.
         # Signed doesn't matter because we can't represent signed integers with 1 qubit.
         if self.bitsize == 1:
@@ -831,12 +830,12 @@ def _t_complexity_(self) -> TComplexity:
     def short_name(self) -> str:
         return f"x > {self.val}"
 
-    def wire_symbol(self, soq: Soquet) -> WireSymbol:
-        if soq.reg.name == 'x':
+    def wire_symbol(self, reg: Register, idx: Tuple[int, ...] = tuple()) -> WireSymbol:
+        if reg.name == 'x':
             return TextBox("In(x)")
-        elif soq.reg.name == 'target':
+        elif reg.name == 'target':
             return TextBox(f"⨁(x > {self.val})")
-        raise ValueError(f'Unknown register symbol {soq.reg.name}')
+        raise ValueError(f'Unknown register symbol {reg.name}')
 
     def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']:
         # TODO Determine precise clifford count and/or ignore.
@@ -884,12 +883,12 @@ def _t_complexity_(self) -> 'TComplexity':
     def short_name(self) -> str:
         return f"x == {self.val}"
 
-    def wire_symbol(self, soq: Soquet) -> WireSymbol:
-        if soq.reg.name == 'x':
+    def wire_symbol(self, reg: Register, idx: Tuple[int, ...] = tuple()) -> WireSymbol:
+        if reg.name == 'x':
             return TextBox("In(x)")
-        elif soq.reg.name == 'target':
+        elif reg.name == 'target':
             return TextBox(f"⨁(x = {self.val})")
-        raise ValueError(f'Unknown register symbol {soq.reg.name}')
+        raise ValueError(f'Unknown register symbol {reg.name}')
 
     def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']:
         # See: https://github.com/quantumlib/Qualtran/issues/219

qualtran/bloqs/arithmetic/conversions.py

@@ -13,7 +13,7 @@
 #  limitations under the License.
 
 from functools import cached_property
-from typing import Dict, Set, TYPE_CHECKING
+from typing import Dict, Set, Tuple, TYPE_CHECKING
 
 from attrs import frozen
 
@@ -28,7 +28,6 @@
     Side,
     Signature,
 )
-from qualtran._infra.quantum_graph import Soquet
 from qualtran.bloqs.basic_gates import Toffoli
 from qualtran.cirq_interop.t_complexity_protocol import TComplexity
 from qualtran.drawing import WireSymbol
@@ -89,10 +88,10 @@ def _t_complexity_(self) -> 'TComplexity':
         num_toffoli = self.bitsize**2 + self.bitsize - 1
         return TComplexity(t=4 * num_toffoli)
 
-    def wire_symbol(self, soq: Soquet) -> WireSymbol:
-        if soq.reg.name == 'mu':
+    def wire_symbol(self, reg: Register, idx: Tuple[int, ...] = tuple()) -> WireSymbol:
+        if reg.name == 'mu':
             return TextBox(r'$\mu$')
-        elif soq.reg.name == 'nu':
+        elif reg.name == 'nu':
             return TextBox(r'$\mu$')
         else:
             text = r'$\oplus\nu(\nu-1)/2+\mu$'

qualtran/bloqs/basic_gates/cnot.py

@@ -29,8 +29,8 @@
     CompositeBloq,
     CtrlSpec,
     DecomposeTypeError,
+    Register,
     Signature,
-    Soquet,
     SoquetT,
 )
 from qualtran.cirq_interop.t_complexity_protocol import TComplexity
@@ -131,12 +131,12 @@ def as_cirq_op(
         (target,) = target
         return cirq.CNOT(ctrl, target), {'ctrl': np.array([ctrl]), 'target': np.array([target])}
 
-    def wire_symbol(self, soq: 'Soquet') -> 'WireSymbol':
-        if soq.reg.name == 'ctrl':
+    def wire_symbol(self, reg: Register, idx: Tuple[int, ...] = tuple()) -> 'WireSymbol':
+        if reg.name == 'ctrl':
             return Circle(filled=True)
-        elif soq.reg.name == 'target':
+        elif reg.name == 'target':
             return ModPlus()
-        raise ValueError(f'Bad wire symbol soquet: {soq}')
+        raise ValueError(f'Unknown wire symbol register name: {reg.name}')
 
     def _t_complexity_(self) -> 'TComplexity':
         return TComplexity(clifford=1)

qualtran/bloqs/basic_gates/hadamard.py

@@ -24,8 +24,8 @@
     BloqDocSpec,
     CompositeBloq,
     DecomposeTypeError,
+    Register,
     Signature,
-    Soquet,
     SoquetT,
 )
 from qualtran.cirq_interop.t_complexity_protocol import TComplexity
@@ -95,7 +95,7 @@ def _t_complexity_(self):
     def short_name(self) -> 'str':
         return 'H'
 
-    def wire_symbol(self, soq: 'Soquet') -> 'WireSymbol':
+    def wire_symbol(self, reg: Register, idx: Tuple[int, ...] = tuple()) -> 'WireSymbol':
         return TextBox('H')