Change signature of cirq_ft.And gate to use directional registers like Qualtran (#6302)

* Change signature of  gate to use directional registers like Qualtran

* Remove debug prints

Author: tanujkhattar

cirq-ft/cirq_ft/algos/and_gate.py

@@ -61,7 +61,16 @@ def _validate_cv(self, attribute, value):
 
     @cached_property
     def signature(self) -> infra.Signature:
-        return infra.Signature.build(control=len(self.cv), ancilla=len(self.cv) - 2, target=1)
+        one_side = infra.Side.RIGHT if not self.adjoint else infra.Side.LEFT
+        n_cv = len(self.cv)
+        junk_reg = [infra.Register('junk', 1, shape=n_cv - 2, side=one_side)] if n_cv > 2 else []
+        return infra.Signature(
+            [
+                infra.Register('ctrl', 1, shape=n_cv),
+                *junk_reg,
+                infra.Register('target', 1, side=one_side),
+            ]
+        )
 
     def __pow__(self, power: int) -> "And":
         if power == 1:
@@ -142,9 +151,9 @@ def decompose_from_registers(
         self, *, context: cirq.DecompositionContext, **quregs: NDArray[cirq.Qid]
     ) -> cirq.OP_TREE:
         control, ancilla, target = (
-            quregs['control'],
-            quregs.get('ancilla', np.array([])),
-            quregs['target'],
+            quregs['ctrl'].flatten(),
+            quregs.get('junk', np.array([])).flatten(),
+            quregs['target'].flatten(),
         )
         if len(self.cv) == 2:
             yield self._decompose_single_and(

cirq-ft/cirq_ft/algos/and_gate_test.py

@@ -46,7 +46,7 @@ def random_cv(n: int) -> List[int]:
 def test_multi_controlled_and_gate(cv: List[int]):
     gate = cirq_ft.And(cv)
     r = gate.signature
-    assert r.get_left('ancilla').total_bits() == r.get_left('control').total_bits() - 2
+    assert r.get_right('junk').total_bits() == r.get_left('ctrl').total_bits() - 2
     quregs = infra.get_named_qubits(r)
     and_op = gate.on_registers(**quregs)
     circuit = cirq.Circuit(and_op)
@@ -55,7 +55,7 @@ def test_multi_controlled_and_gate(cv: List[int]):
     qubit_order = infra.merge_qubits(gate.signature, **quregs)
 
     for input_control in input_controls:
-        initial_state = input_control + [0] * (r.get_left('ancilla').total_bits() + 1)
+        initial_state = input_control + [0] * (r.get_right('junk').total_bits() + 1)
         result = cirq.Simulator(dtype=np.complex128).simulate(
             circuit, initial_state=initial_state, qubit_order=qubit_order
         )
@@ -80,64 +80,67 @@ def test_and_gate_diagram():
     gate = cirq_ft.And((1, 0, 1, 0, 1, 0))
     qubit_regs = infra.get_named_qubits(gate.signature)
     op = gate.on_registers(**qubit_regs)
-    # Qubit order should be alternating (control, ancilla) pairs.
-    c_and_a = sum(zip(qubit_regs["control"][1:], qubit_regs["ancilla"]), ()) + (
-        qubit_regs["control"][-1],
+    ctrl, junk, target = (
+        qubit_regs["ctrl"].flatten(),
+        qubit_regs["junk"].flatten(),
+        qubit_regs['target'].flatten(),
     )
-    qubit_order = np.concatenate([qubit_regs["control"][0:1], c_and_a, qubit_regs["target"]])
+    # Qubit order should be alternating (control, ancilla) pairs.
+    c_and_a = sum(zip(ctrl[1:], junk), ()) + (ctrl[-1],)
+    qubit_order = np.concatenate([ctrl[0:1], c_and_a, target])
     # Test diagrams.
     cirq.testing.assert_has_diagram(
         cirq.Circuit(op),
         """
-control0: ───@─────
-             │
-control1: ───(0)───
-             │
-ancilla0: ───Anc───
-             │
-control2: ───@─────
-             │
-ancilla1: ───Anc───
-             │
-control3: ───(0)───
-             │
-ancilla2: ───Anc───
-             │
-control4: ───@─────
-             │
-ancilla3: ───Anc───
-             │
-control5: ───(0)───
-             │
-target: ─────And───
+ctrl[0]: ───@─────
+            │
+ctrl[1]: ───(0)───
+            │
+junk[0]: ───Anc───
+            │
+ctrl[2]: ───@─────
+            │
+junk[1]: ───Anc───
+            │
+ctrl[3]: ───(0)───
+            │
+junk[2]: ───Anc───
+            │
+ctrl[4]: ───@─────
+            │
+junk[3]: ───Anc───
+            │
+ctrl[5]: ───(0)───
+            │
+target: ────And───
 """,
         qubit_order=qubit_order,
     )
     cirq.testing.assert_has_diagram(
         cirq.Circuit(op**-1),
         """
-control0: ───@──────
-             │
-control1: ───(0)────
-             │
-ancilla0: ───Anc────
-             │
-control2: ───@──────
-             │
-ancilla1: ───Anc────
-             │
-control3: ───(0)────
-             │
-ancilla2: ───Anc────
-             │
-control4: ───@──────
-             │
-ancilla3: ───Anc────
-             │
-control5: ───(0)────
-             │
-target: ─────And†───
-    """,
+ctrl[0]: ───@──────
+            │
+ctrl[1]: ───(0)────
+            │
+junk[0]: ───Anc────
+            │
+ctrl[2]: ───@──────
+            │
+junk[1]: ───Anc────
+            │
+ctrl[3]: ───(0)────
+            │
+junk[2]: ───Anc────
+            │
+ctrl[4]: ───@──────
+            │
+junk[3]: ───Anc────
+            │
+ctrl[5]: ───(0)────
+            │
+target: ────And†───
+""",
         qubit_order=qubit_order,
     )
     # Test diagram of decomposed 3-qubit and ladder.
@@ -147,28 +150,28 @@ def test_and_gate_diagram():
     cirq.testing.assert_has_diagram(
         decomposed_circuit,
         """
-control0: ───@─────────────────────────────────────────────────────────@──────
-             │                                                         │
-control1: ───(0)───────────────────────────────────────────────────────(0)────
-             │                                                         │
-ancilla0: ───And───@────────────────────────────────────────────@──────And†───
-                   │                                            │
-control2: ─────────@────────────────────────────────────────────@─────────────
-                   │                                            │
-ancilla1: ─────────And───@───────────────────────────────@──────And†──────────
-                         │                               │
-control3: ───────────────(0)─────────────────────────────(0)──────────────────
-                         │                               │
-ancilla2: ───────────────And───@──────────────────@──────And†─────────────────
-                               │                  │
-control4: ─────────────────────@──────────────────@───────────────────────────
-                               │                  │
-ancilla3: ─────────────────────And───@─────@──────And†────────────────────────
-                                     │     │
-control5: ───────────────────────────(0)───(0)────────────────────────────────
-                                     │     │
-target: ─────────────────────────────And───And†───────────────────────────────
-    """,
+ctrl[0]: ───@─────────────────────────────────────────────────────────@──────
+            │                                                         │
+ctrl[1]: ───(0)───────────────────────────────────────────────────────(0)────
+            │                                                         │
+junk[0]: ───And───@────────────────────────────────────────────@──────And†───
+                  │                                            │
+ctrl[2]: ─────────@────────────────────────────────────────────@─────────────
+                  │                                            │
+junk[1]: ─────────And───@───────────────────────────────@──────And†──────────
+                        │                               │
+ctrl[3]: ───────────────(0)─────────────────────────────(0)──────────────────
+                        │                               │
+junk[2]: ───────────────And───@──────────────────@──────And†─────────────────
+                              │                  │
+ctrl[4]: ─────────────────────@──────────────────@───────────────────────────
+                              │                  │
+junk[3]: ─────────────────────And───@─────@──────And†────────────────────────
+                                    │     │
+ctrl[5]: ───────────────────────────(0)───(0)────────────────────────────────
+                                    │     │
+target: ────────────────────────────And───And†───────────────────────────────
+""",
         qubit_order=qubit_order,
     )
 

cirq-ft/cirq_ft/algos/hubbard_model.py

@@ -335,13 +335,13 @@ def decompose_from_registers(
         and_target = context.qubit_manager.qalloc(1)
         and_anc = context.qubit_manager.qalloc(1)
         yield and_gate.And(cv=(0, 0, 1)).on_registers(
-            control=[*U, *V, temp[-1]], ancilla=and_anc, target=and_target
+            ctrl=np.array([U, V, temp[-1:]]), junk=np.array([and_anc]), target=and_target
         )
         yield swap_network.MultiTargetCSwap.make_on(
             control=and_target, target_x=[*p_x, *p_y, *alpha], target_y=[*q_x, *q_y, *beta]
         )
         yield and_gate.And(cv=(0, 0, 1), adjoint=True).on_registers(
-            control=[*U, *V, temp[-1]], ancilla=and_anc, target=and_target
+            ctrl=np.array([U, V, temp[-1:]]), junk=np.array([and_anc]), target=and_target
         )
         context.qubit_manager.qfree([*and_anc, *and_target])
 

cirq-ft/cirq_ft/algos/multi_control_multi_target_pauli.py

@@ -85,15 +85,15 @@ def decompose_from_registers(
     ) -> cirq.OP_TREE:
         controls, target = quregs['controls'], quregs['target']
         qm = context.qubit_manager
-        and_ancilla, and_target = qm.qalloc(len(self.cvs) - 2), qm.qalloc(1)
+        and_ancilla, and_target = np.array(qm.qalloc(len(self.cvs) - 2)), qm.qalloc(1)
         yield and_gate.And(self.cvs).on_registers(
-            control=controls, ancilla=and_ancilla, target=and_target
+            ctrl=controls[:, np.newaxis], junk=and_ancilla[:, np.newaxis], target=and_target
         )
         yield self.target_gate.on(*target).controlled_by(*and_target)
         yield and_gate.And(self.cvs, adjoint=True).on_registers(
-            control=controls, ancilla=and_ancilla, target=and_target
+            ctrl=controls[:, np.newaxis], junk=and_ancilla[:, np.newaxis], target=and_target
         )
-        qm.qfree(and_ancilla + and_target)
+        qm.qfree([*and_ancilla, *and_target])
 
     def _circuit_diagram_info_(self, _) -> cirq.CircuitDiagramInfo:
         wire_symbols = ["@" if b else "@(0)" for b in self.cvs]

cirq-ft/cirq_ft/algos/prepare_uniform_superposition.py

@@ -94,7 +94,9 @@ def decompose_from_registers(
 
         and_ancilla = context.qubit_manager.qalloc(len(self.cv) + logL - 2)
         and_op = and_gate.And((0,) * logL + self.cv).on_registers(
-            control=[*logL_qubits, *controls], ancilla=and_ancilla, target=ancilla
+            ctrl=np.asarray([*logL_qubits, *controls])[:, np.newaxis],
+            junk=np.asarray(and_ancilla)[:, np.newaxis],
+            target=ancilla,
         )
         yield and_op
         yield cirq.Rz(rads=theta)(*ancilla)

cirq-ft/cirq_ft/algos/qrom.py

@@ -150,7 +150,9 @@ def decompose_zero_selection(
             and_ancilla = context.qubit_manager.qalloc(len(controls) - 2)
             and_target = context.qubit_manager.qalloc(1)[0]
             multi_controlled_and = and_gate.And((1,) * len(controls)).on_registers(
-                control=controls, ancilla=and_ancilla, target=and_target
+                ctrl=np.array(controls)[:, np.newaxis],
+                junk=np.array(and_ancilla)[:, np.newaxis],
+                target=and_target,
             )
             yield multi_controlled_and
             yield self._load_nth_data(zero_indx, lambda q: cirq.CNOT(and_target, q), **target_regs)

cirq-ft/cirq_ft/algos/unary_iteration_gate.py

@@ -153,7 +153,9 @@ def _unary_iteration_multi_controls(
     and_ancilla = ancilla[: num_controls - 2]
     and_target = ancilla[num_controls - 2]
     multi_controlled_and = and_gate.And((1,) * len(controls)).on_registers(
-        control=np.array(controls), ancilla=np.array(and_ancilla), target=and_target
+        ctrl=np.array(controls).reshape(len(controls), 1),
+        junk=np.array(and_ancilla).reshape(len(and_ancilla), 1),
+        target=and_target,
     )
     ops.append(multi_controlled_and)
     yield from _unary_iteration_single_control(

cirq-ft/cirq_ft/infra/jupyter_tools_test.py

@@ -27,7 +27,7 @@ def test_svg_circuit():
     svg_str = svg.data
 
     # check that the order is respected in the svg data.
-    assert svg_str.find('control') < svg_str.find('ancilla') < svg_str.find('target')
+    assert svg_str.find('ctrl') < svg_str.find('junk') < svg_str.find('target')
 
     # Check svg_circuit raises.
     with pytest.raises(ValueError):

cirq-ft/cirq_ft/infra/testing_test.py

@@ -34,10 +34,16 @@ def test_assert_circuit_inp_out_cirqsim():
 def test_gate_helper():
     g = cirq_ft.testing.GateHelper(cirq_ft.And(cv=(1, 0, 1, 0)))
     assert g.gate == cirq_ft.And(cv=(1, 0, 1, 0))
-    assert g.r == cirq_ft.Signature.build(control=4, ancilla=2, target=1)
+    assert g.r == cirq_ft.Signature(
+        [
+            cirq_ft.Register('ctrl', bitsize=1, shape=4),
+            cirq_ft.Register('junk', bitsize=1, shape=2, side=cirq_ft.infra.Side.RIGHT),
+            cirq_ft.Register('target', bitsize=1, side=cirq_ft.infra.Side.RIGHT),
+        ]
+    )
     expected_quregs = {
-        'control': cirq.NamedQubit.range(4, prefix='control'),
-        'ancilla': cirq.NamedQubit.range(2, prefix='ancilla'),
+        'ctrl': np.array([[cirq.q(f'ctrl[{i}]')] for i in range(4)]),
+        'junk': np.array([[cirq.q(f'junk[{i}]')] for i in range(2)]),
         'target': [cirq.NamedQubit('target')],
     }
     for key in expected_quregs: