Add SelectionRegisters and multi-indexed versions of SelectedMajoranaFermionGate and ApplyGateToLthQubit

cirq_qubitization/__init__.py

@@ -35,6 +35,7 @@
     qfree,
     Register,
     Registers,
+    SelectionRegisters,
     SimpleQubitManager,
 )
 from cirq_qubitization.generic_select import GenericSelect

cirq_qubitization/cirq_algos/apply_gate_to_lth_target.ipynb

@@ -47,10 +47,9 @@
     "`selection`-th qubit of `target` all controlled by the `control` register.\n",
     "\n",
     "#### Parameters\n",
-    " - `selection_bitsize`: The size of the indexing `select` register. This should be at most `log2(target_bitsize)`\n",
-    " - `target_bitsize`: The size of the `target` register. This also serves as the iteration length.\n",
-    " - `nth_gate`: A function mapping the selection index to a single-qubit gate.\n",
-    " - `control_bitsize`: The size of the control register.\n"
+    " - `selection_regs`: Indexing `select` registers of type `SelectionRegisters`. It also contains information about the iteration length of each selection register.\n",
+    " - `nth_gate`: A function mapping the composite selection index to a single-qubit gate.\n",
+    " - `control_regs`: Control registers for constructing a controlled version of the gate.\n"
    ]
   },
   {
@@ -63,14 +62,17 @@
    "outputs": [],
    "source": [
     "from cirq_qubitization.cirq_algos.apply_gate_to_lth_target import ApplyGateToLthQubit\n",
+    "from cirq_qubitization.cirq_infra.gate_with_registers import Registers, SelectionRegisters\n",
     "\n",
     "def _z_to_odd(n: int):\n",
     "    if n % 2 == 1:\n",
     "        return cirq.Z\n",
     "    return cirq.I\n",
     "\n",
     "apply_z_to_odd = ApplyGateToLthQubit(\n",
-    "    selection_bitsize=3, target_bitsize=4, nth_gate=_z_to_odd, control_bitsize=2\n",
+    "    SelectionRegisters.build(selection=(3, 4)),\n",
+    "    nth_gate=_z_to_odd,\n",
+    "    control_regs=Registers.build(control=2),\n",
     ")\n",
     "\n",
     "g = cq_testing.GateHelper(\n",

cirq_qubitization/cirq_algos/apply_gate_to_lth_target.py

@@ -1,12 +1,15 @@
+import itertools
 from functools import cached_property
 from typing import Callable, Sequence, Tuple
 
 import cirq
+from attrs import frozen
 
 from cirq_qubitization.cirq_algos.unary_iteration import UnaryIterationGate
-from cirq_qubitization.cirq_infra.gate_with_registers import Registers
+from cirq_qubitization.cirq_infra.gate_with_registers import Registers, SelectionRegisters
 
 
+@frozen
 class ApplyGateToLthQubit(UnaryIterationGate):
     r"""A controlled SELECT operation for single-qubit gates.
 
@@ -20,62 +23,52 @@ class ApplyGateToLthQubit(UnaryIterationGate):
     `selection`-th qubit of `target` all controlled by the `control` register.
 
     Args:
-        selection_bitsize: The size of the indexing `select` register. This should be at most
-            `log2(target_bitsize)`
-        target_bitsize: The size of the `target` register. This also serves as the iteration
-            length.
-        nth_gate: A function mapping the selection index to a single-qubit gate.
-        control_bitsize: The size of the control register.
+        selection_regs: Indexing `select` registers of type `SelectionRegisters`. It also contains
+            information about the iteration length of each selection register.
+        nth_gate: A function mapping the composite selection index to a single-qubit gate.
+        control_regs: Control registers for constructing a controlled version of the gate.
     """
-
-    def __init__(
-        self,
-        selection_bitsize: int,
-        target_bitsize: int,
-        nth_gate: Callable[[int], cirq.Gate],
-        *,
-        control_bitsize: int = 1,
-    ):
-        self._selection_bitsize = selection_bitsize
-        self._target_bitsize = target_bitsize
-        self._nth_gate = nth_gate
-        self._control_bitsize = control_bitsize
+    selection_regs: SelectionRegisters
+    nth_gate: Callable[[int, ...], cirq.Gate]
+    control_regs: Registers = Registers.build(control=1)
 
     @classmethod
     def make_on(
-        cls, *, nth_gate: Callable[[int], cirq.Gate], **quregs: Sequence[cirq.Qid]
+        cls, *, nth_gate: Callable[[int, ...], cirq.Gate], **quregs: Sequence[cirq.Qid]
     ) -> cirq.Operation:
         """Helper constructor to automatically deduce bitsize attributes."""
         return cls(
-            selection_bitsize=len(quregs['selection']),
-            target_bitsize=len(quregs['target']),
+            SelectionRegisters.build(selection=(len(quregs['selection']), len(quregs['target']))),
             nth_gate=nth_gate,
-            control_bitsize=len(quregs['control']),
+            control_regs=Registers.build(control=len(quregs['control'])),
         ).on_registers(**quregs)
 
     @cached_property
     def control_registers(self) -> Registers:
-        return Registers.build(control=self._control_bitsize)
+        return self.control_regs
 
     @cached_property
-    def selection_registers(self) -> Registers:
-        return Registers.build(selection=self._selection_bitsize)
+    def selection_registers(self) -> SelectionRegisters:
+        return self.selection_regs
 
     @cached_property
     def target_registers(self) -> Registers:
-        return Registers.build(target=self._target_bitsize)
+        return Registers.build(target=self.selection_registers.total_iteration_size)
 
     @cached_property
     def iteration_lengths(self) -> Tuple[int, ...]:
-        return (self._target_bitsize,)
+        return self.selection_registers.iteration_lengths
 
     def _circuit_diagram_info_(self, args: cirq.CircuitDiagramInfoArgs) -> cirq.CircuitDiagramInfo:
         wire_symbols = ["@"] * self.control_registers.bitsize
         wire_symbols += ["In"] * self.selection_registers.bitsize
-        wire_symbols += [str(self._nth_gate(i)) for i in range(self._target_bitsize)]
+        for it in itertools.product(*[range(x) for x in self.selection_regs.iteration_lengths]):
+            wire_symbols += [str(self.nth_gate(*it))]
         return cirq.CircuitDiagramInfo(wire_symbols=wire_symbols)
 
     def nth_operation(
-        self, selection: int, control: cirq.Qid, target: Sequence[cirq.Qid]
+        self, control: cirq.Qid, target: Sequence[cirq.Qid], **selection_indices: int
     ) -> cirq.OP_TREE:
-        return self._nth_gate(selection).on(target[-(selection + 1)]).controlled_by(control)
+        selection_idx = tuple(selection_indices[reg.name] for reg in self.selection_regs)
+        target_idx = self.selection_registers.to_flat_idx(*selection_idx)
+        return self.nth_gate(*selection_idx).on(target[target_idx]).controlled_by(control)

cirq_qubitization/cirq_algos/apply_gate_to_lth_target_test.py

@@ -11,7 +11,10 @@
 def test_apply_gate_to_lth_qubit(selection_bitsize, target_bitsize):
     greedy_mm = cq.cirq_infra.GreedyQubitManager(prefix="_a", maximize_reuse=True)
     with cq.cirq_infra.memory_management_context(greedy_mm):
-        gate = cq.ApplyGateToLthQubit(selection_bitsize, target_bitsize, lambda _: cirq.X)
+        gate = cq.ApplyGateToLthQubit(
+            cq.SelectionRegisters.build(selection=(selection_bitsize, target_bitsize)),
+            lambda _: cirq.X,
+        )
         g = cq_testing.GateHelper(gate)
         # Upper bounded because not all ancillas may be used as part of unary iteration.
         assert (
@@ -28,7 +31,7 @@ def test_apply_gate_to_lth_qubit(selection_bitsize, target_bitsize):
         qubit_vals |= zip(g.quregs['selection'], iter_bits(n, selection_bitsize))
 
         initial_state = [qubit_vals[x] for x in g.all_qubits]
-        qubit_vals[g.quregs['target'][-(n + 1)]] = 1
+        qubit_vals[g.quregs['target'][n]] = 1
         final_state = [qubit_vals[x] for x in g.all_qubits]
         cq_testing.assert_circuit_inp_out_cirqsim(
             g.decomposed_circuit, g.all_qubits, initial_state, final_state
@@ -37,7 +40,11 @@ def test_apply_gate_to_lth_qubit(selection_bitsize, target_bitsize):
 
 def test_apply_gate_to_lth_qubit_diagram():
     # Apply Z gate to all odd targets and Identity to even targets.
-    gate = cq.ApplyGateToLthQubit(3, 5, lambda n: cirq.Z if n & 1 else cirq.I, control_bitsize=2)
+    gate = cq.ApplyGateToLthQubit(
+        cq.SelectionRegisters.build(selection=(3, 5)),
+        lambda n: cirq.Z if n & 1 else cirq.I,
+        control_regs=cq.Registers.build(control=2),
+    )
     circuit = cirq.Circuit(gate.on_registers(**gate.registers.get_named_qubits()))
     qubits = list(q for v in gate.registers.get_named_qubits().values() for q in v)
     cirq.testing.assert_has_diagram(
@@ -68,14 +75,19 @@ def test_apply_gate_to_lth_qubit_diagram():
 
 
 def test_apply_gate_to_lth_qubit_make_on():
-    gate = cq.ApplyGateToLthQubit(3, 5, lambda n: cirq.Z if n & 1 else cirq.I, control_bitsize=2)
+    gate = cq.ApplyGateToLthQubit(
+        cq.SelectionRegisters.build(selection=(3, 5)),
+        lambda n: cirq.Z if n & 1 else cirq.I,
+        control_regs=cq.Registers.build(control=2),
+    )
     op = gate.on_registers(**gate.registers.get_named_qubits())
     op2 = cq.ApplyGateToLthQubit.make_on(
         nth_gate=lambda n: cirq.Z if n & 1 else cirq.I, **gate.registers.get_named_qubits()
     )
     # Note: ApplyGateToLthQubit doesn't support value equality.
     assert op.qubits == op2.qubits
-    assert op.gate._selection_bitsize == op2.gate._selection_bitsize
+    assert op.gate.selection_regs == op2.gate.selection_regs
+    assert op.gate.control_regs == op2.gate.control_regs
 
 
 def test_notebook():

cirq_qubitization/cirq_algos/selected_majorana_fermion.py

@@ -2,81 +2,87 @@
 from typing import Sequence, Tuple
 
 import cirq
+from attrs import frozen
 
+from cirq_qubitization import cirq_infra
 from cirq_qubitization.cirq_algos import unary_iteration
-from cirq_qubitization.cirq_infra.gate_with_registers import Registers
+from cirq_qubitization.cirq_infra.gate_with_registers import Registers, SelectionRegisters
 
 
-@cirq.value_equality()
+@frozen
 class SelectedMajoranaFermionGate(unary_iteration.UnaryIterationGate):
     """Implements U s.t. U|l>|Psi> -> |l> T_{l} . Z_{l - 1} ... Z_{0} |Psi>
 
     where T = single qubit target gate. Defaults to pauli Y.
 
-    Uses:
-    * 1 Control qubit.
-    * 1 Accumulator qubit.
-    * `selection_bitsize` number of selection qubits.
-    * `target_bitsize` number of target qubits.
 
-    See Fig 9 of https://arxiv.org/abs/1805.03662 for more details.
+    Args:
+        selection_regs: Indexing `select` registers of type `SelectionRegisters`. It also contains
+            information about the iteration length of each selection register.
+        control_regs: Control registers for constructing a controlled version of the gate.
+        target_gate: Single qubit gate to be applied to the target qubits.
+
+    References:
+        See Fig 9 of https://arxiv.org/abs/1805.03662 for more details.
     """
 
-    def __init__(self, selection_bitsize: int, target_bitsize: int, target_gate=cirq.Y):
-        self._selection_bitsize = selection_bitsize
-        self._target_bitsize = target_bitsize
-        self._target_gate = target_gate
+    selection_regs: SelectionRegisters
+    control_regs: Registers = Registers.build(control=1)
+    target_gate: cirq.Gate = cirq.Y
 
     @classmethod
     def make_on(cls, *, target_gate=cirq.Y, **quregs: Sequence[cirq.Qid]) -> cirq.Operation:
-        """Helper constructor to automatically deduce bitsize attributes."""
+        """Helper constructor to automatically deduce selection_regs attribute."""
         return cls(
-            selection_bitsize=len(quregs['selection']),
-            target_bitsize=len(quregs['target']),
+            selection_regs=SelectionRegisters.build(
+                selection=(len(quregs['selection']), len(quregs['target']))
+            ),
             target_gate=target_gate,
         ).on_registers(**quregs)
 
-    def _value_equality_values_(self):
-        return self._selection_bitsize, self._target_bitsize, self._target_gate
-
     @cached_property
     def control_registers(self) -> Registers:
-        return Registers.build(control=1)
+        return self.control_regs
 
     @cached_property
-    def selection_registers(self) -> Registers:
-        return Registers.build(selection=self._selection_bitsize)
+    def selection_registers(self) -> SelectionRegisters:
+        return self.selection_regs
 
     @cached_property
     def target_registers(self) -> Registers:
-        return Registers.build(target=self._target_bitsize)
+        return Registers.build(target=self.selection_regs.total_iteration_size)
 
     @cached_property
     def iteration_lengths(self) -> Tuple[int, ...]:
-        return (self._target_bitsize,)
+        return self.selection_registers.iteration_lengths
 
     @cached_property
     def extra_registers(self) -> Registers:
         return Registers.build(accumulator=1)
 
     def decompose_from_registers(self, **qubit_regs: Sequence[cirq.Qid]) -> cirq.OP_TREE:
-        yield cirq.CNOT(*qubit_regs['control'], *qubit_regs['accumulator'])
+        qubit_regs['accumulator'] = cirq_infra.qalloc(1)
+        yield cirq.X(*qubit_regs['accumulator']).controlled_by(
+            *qubit_regs[self.control_regs[0].name]
+        )
         yield super().decompose_from_registers(**qubit_regs)
+        cirq_infra.qfree(qubit_regs['accumulator'])
 
     def _circuit_diagram_info_(self, args: cirq.CircuitDiagramInfoArgs) -> cirq.CircuitDiagramInfo:
         wire_symbols = ["@"] * self.control_registers.bitsize
         wire_symbols += ["In"] * self.selection_registers.bitsize
-        wire_symbols += [f"Z{self._target_gate}"] * self.target_registers.bitsize
-        wire_symbols += ["Acc"]
+        wire_symbols += [f"Z{self.target_gate}"] * self.target_registers.bitsize
         return cirq.CircuitDiagramInfo(wire_symbols=wire_symbols)
 
     def nth_operation(
         self,
-        selection: int,
         control: cirq.Qid,
         target: Sequence[cirq.Qid],
         accumulator: Sequence[cirq.Qid],
+        **selection_indices: int,
     ) -> cirq.OP_TREE:
-        yield cirq.CNOT(control, accumulator[0])
-        yield self._target_gate(target[selection]).controlled_by(control)
-        yield cirq.Z(target[selection]).controlled_by(accumulator[0])
+        selection_idx = tuple(selection_indices[reg.name] for reg in self.selection_regs)
+        target_idx = self.selection_registers.to_flat_idx(*selection_idx)
+        yield cirq.CNOT(control, *accumulator)
+        yield self.target_gate(target[target_idx]).controlled_by(control)
+        yield cirq.CZ(*accumulator, target[target_idx])