Add ApplyGateToLthQubit Bloq from Cirq-FT

qualtran/bloqs/apply_gate_to_lth_target.ipynb

@@ -0,0 +1,127 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "60432dd0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#  Copyright 2023 Google LLC\n",
+    "#\n",
+    "#  Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+    "#  you may not use this file except in compliance with the License.\n",
+    "#  You may obtain a copy of the License at\n",
+    "#\n",
+    "#      https://www.apache.org/licenses/LICENSE-2.0\n",
+    "#\n",
+    "#  Unless required by applicable law or agreed to in writing, software\n",
+    "#  distributed under the License is distributed on an \"AS IS\" BASIS,\n",
+    "#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
+    "#  See the License for the specific language governing permissions and\n",
+    "#  limitations under the License."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ac3bfb05",
+   "metadata": {
+    "cq.autogen": "title_cell"
+   },
+   "source": [
+    "# Apply to L-th Target"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4e214a27",
+   "metadata": {
+    "cq.autogen": "top_imports"
+   },
+   "outputs": [],
+   "source": [
+    "import cirq\n",
+    "import numpy as np\n",
+    "from qualtran import Signature, SelectionRegister\n",
+    "from qualtran.bloqs.apply_gate_to_lth_target import ApplyGateToLthQubit\n",
+    "import qualtran.cirq_interop.testing as cq_testing\n",
+    "from qualtran.cirq_interop.jupyter_tools import display_gate_and_compilation\n",
+    "from typing import *"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "249829b0",
+   "metadata": {
+    "cq.autogen": "_make_ApplyGateToLthQubit.md"
+   },
+   "source": [
+    "## `ApplyGateToLthQubit`\n",
+    "A controlled SELECT operation for single-qubit gates.\n",
+    "\n",
+    "$$\n",
+    "\\mathrm{SELECT} = \\sum_{l}|l \\rangle \\langle l| \\otimes [G(l)]_l\n",
+    "$$\n",
+    "\n",
+    "Where $G$ is a function that maps an index to a single-qubit gate.\n",
+    "\n",
+    "This gate uses the unary iteration scheme to apply `nth_gate(selection)` to the\n",
+    "`selection`-th qubit of `target` all controlled by the `control` register.\n",
+    "\n",
+    "#### Parameters\n",
+    " - `selection_regs`: Indexing `select` signature of type Tuple[`SelectionRegister`, ...]. 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 signature for constructing a controlled version of the gate.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b8d2a7bf",
+   "metadata": {
+    "cq.autogen": "_make_ApplyGateToLthQubit.py"
+   },
+   "outputs": [],
+   "source": [
+    "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",
+    "    SelectionRegister('selection', 3, 4),\n",
+    "    nth_gate=_z_to_odd,\n",
+    "    control_regs=Signature.build(control=2),\n",
+    ")\n",
+    "\n",
+    "g = cq_testing.GateHelper(\n",
+    "    apply_z_to_odd\n",
+    ")\n",
+    "\n",
+    "display_gate_and_compilation(g)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.8.16"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

qualtran/bloqs/apply_gate_to_lth_target.py

@@ -0,0 +1,104 @@
+#  Copyright 2023 Google LLC
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+import itertools
+from typing import Callable, Sequence, Tuple
+
+import attr
+import cirq
+import numpy as np
+from cirq._compat import cached_property
+
+from qualtran import Register, SelectionRegister
+from qualtran._infra.gate_with_registers import total_bits
+from qualtran.bloqs.unary_iteration_bloq import UnaryIterationGate
+
+
+@attr.frozen
+class ApplyGateToLthQubit(UnaryIterationGate):
+    r"""A controlled SELECT operation for single-qubit gates.
+
+    $$
+    \mathrm{SELECT} = \sum_{l}|l \rangle \langle l| \otimes [G(l)]_l
+    $$
+
+    Where $G$ is a function that maps an index to a single-qubit gate.
+
+    This gate uses the unary iteration scheme to apply `nth_gate(selection)` to the
+    `selection`-th qubit of `target` all controlled by the `control` register.
+
+    Args:
+        selection_regs: Indexing `select` signature of type Tuple[`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 signature for constructing a controlled version of the gate.
+
+    References:
+            [Encoding Electronic Spectra in Quantum Circuits with Linear T Complexity]
+        (https://arxiv.org/abs/1805.03662).
+        Babbush et. al. (2018). Section III.A. and Figure 7.
+    """
+    selection_regs: Tuple[SelectionRegister, ...] = attr.field(
+        converter=lambda v: (v,) if isinstance(v, SelectionRegister) else tuple(v)
+    )
+    nth_gate: Callable[..., cirq.Gate]
+    control_regs: Tuple[Register, ...] = attr.field(
+        converter=lambda v: (v,) if isinstance(v, Register) else tuple(v),
+        default=(Register('control', 1),),
+    )
+
+    @classmethod
+    def make_on(
+        cls, *, nth_gate: Callable[..., cirq.Gate], **quregs: Sequence[cirq.Qid]
+    ) -> cirq.Operation:
+        """Helper constructor to automatically deduce bitsize attributes."""
+        return ApplyGateToLthQubit(
+            SelectionRegister('selection', len(quregs['selection']), len(quregs['target'])),
+            nth_gate=nth_gate,
+            control_regs=Register('control', len(quregs['control'])),
+        ).on_registers(**quregs)
+
+    @cached_property
+    def control_registers(self) -> Tuple[Register, ...]:
+        return self.control_regs
+
+    @cached_property
+    def selection_registers(self) -> Tuple[SelectionRegister, ...]:
+        return self.selection_regs
+
+    @cached_property
+    def target_registers(self) -> Tuple[Register, ...]:
+        total_iteration_size = np.prod(
+            tuple(reg.iteration_length for reg in self.selection_registers)
+        )
+        return (Register('target', int(total_iteration_size)),)
+
+    def _circuit_diagram_info_(self, args: cirq.CircuitDiagramInfoArgs) -> cirq.CircuitDiagramInfo:
+        wire_symbols = ["@"] * total_bits(self.control_registers)
+        wire_symbols += ["In"] * total_bits(self.selection_registers)
+        for it in itertools.product(*[range(reg.iteration_length) for reg in self.selection_regs]):
+            wire_symbols += [str(self.nth_gate(*it))]
+        return cirq.CircuitDiagramInfo(wire_symbols=wire_symbols)
+
+    def nth_operation(  # type: ignore[override]
+        self,
+        context: cirq.DecompositionContext,
+        control: cirq.Qid,
+        target: Sequence[cirq.Qid],
+        **selection_indices: int,
+    ) -> cirq.OP_TREE:
+        selection_shape = tuple(reg.iteration_length for reg in self.selection_regs)
+        selection_idx = tuple(selection_indices[reg.name] for reg in self.selection_regs)
+        target_idx = int(np.ravel_multi_index(selection_idx, selection_shape))
+        return self.nth_gate(*selection_idx).on(target[target_idx]).controlled_by(control)

qualtran/bloqs/apply_gate_to_lth_target_test.py

@@ -0,0 +1,118 @@
+#  Copyright 2023 Google LLC
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+import cirq
+import pytest
+
+from qualtran import SelectionRegister, Signature
+from qualtran._infra.gate_with_registers import get_named_qubits, total_bits
+from qualtran.bloqs.apply_gate_to_lth_target import ApplyGateToLthQubit
+from qualtran.cirq_interop.bit_tools import iter_bits
+from qualtran.cirq_interop.testing import assert_circuit_inp_out_cirqsim, GateHelper
+from qualtran.testing import assert_valid_bloq_decomposition, execute_notebook
+
+
+@pytest.mark.parametrize("selection_bitsize,target_bitsize", [[3, 5], [3, 7], [4, 5]])
+def test_apply_gate_to_lth_qubit(selection_bitsize, target_bitsize):
+    greedy_mm = cirq.GreedyQubitManager(prefix="_a", maximize_reuse=True)
+    gate = ApplyGateToLthQubit(
+        SelectionRegister('selection', selection_bitsize, target_bitsize), lambda _: cirq.X
+    )
+    g = GateHelper(gate, context=cirq.DecompositionContext(greedy_mm))
+    # Upper bounded because not all ancillas may be used as part of unary iteration.
+    assert (
+        len(g.all_qubits)
+        <= target_bitsize + 2 * (selection_bitsize + total_bits(gate.control_registers)) - 1
+    )
+
+    for n in range(target_bitsize):
+        # Initial qubit values
+        qubit_vals = {q: 0 for q in g.all_qubits}
+        # All controls 'on' to activate circuit
+        qubit_vals.update({c: 1 for c in g.quregs['control']})
+        # Set selection according to `n`
+        qubit_vals.update(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
+        final_state = [qubit_vals[x] for x in g.all_qubits]
+        assert_circuit_inp_out_cirqsim(
+            g.decomposed_circuit, g.all_qubits, initial_state, final_state
+        )
+
+
+def test_apply_gate_to_lth_qubit_diagram():
+    # Apply Z gate to all odd targets and Identity to even targets.
+    gate = ApplyGateToLthQubit(
+        SelectionRegister('selection', 3, 5),
+        lambda n: cirq.Z if n & 1 else cirq.I,
+        control_regs=Signature.build(control=2),
+    )
+    circuit = cirq.Circuit(gate.on_registers(**get_named_qubits(gate.signature)))
+    qubits = list(q for v in get_named_qubits(gate.signature).values() for q in v)
+    cirq.testing.assert_has_diagram(
+        circuit,
+        """
+control0: ─────@────
+               │
+control1: ─────@────
+               │
+selection0: ───In───
+               │
+selection1: ───In───
+               │
+selection2: ───In───
+               │
+target0: ──────I────
+               │
+target1: ──────Z────
+               │
+target2: ──────I────
+               │
+target3: ──────Z────
+               │
+target4: ──────I────
+""",
+        qubit_order=qubits,
+    )
+
+
+def test_apply_gate_to_lth_qubit_make_on():
+    gate = ApplyGateToLthQubit(
+        SelectionRegister('selection', 3, 5),
+        lambda n: cirq.Z if n & 1 else cirq.I,
+        control_regs=Signature.build(control=2),
+    )
+    op = gate.on_registers(**get_named_qubits(gate.signature))
+    op2 = ApplyGateToLthQubit.make_on(
+        nth_gate=lambda n: cirq.Z if n & 1 else cirq.I, **get_named_qubits(gate.signature)
+    )
+    # Note: ApplyGateToLthQubit doesn't support value equality.
+    assert op.qubits == op2.qubits
+    assert op.gate.selection_regs == op2.gate.selection_regs
+    assert op.gate.control_regs == op2.gate.control_regs
+
+
+@pytest.mark.parametrize("selection_bitsize,target_bitsize", [[3, 5], [3, 7], [4, 5]])
+def test_bloq_has_consistent_decomposition(selection_bitsize, target_bitsize):
+    bloq = ApplyGateToLthQubit(
+        SelectionRegister('selection', selection_bitsize, target_bitsize),
+        lambda n: cirq.Z if n & 1 else cirq.I,
+        control_regs=Signature.build(control=2),
+    )
+    assert_valid_bloq_decomposition(bloq)
+
+
+def test_notebook():
+    execute_notebook('apply_gate_to_lth_target')