Add PrepareUniformSuperposition Bloq from Cirq-FT (#405)

* Add all arithmetic gates from Cirq-FT as Bloqs * Use qualtran bloqs in comparision_gates.ipynb * Change pretty_name for LessThanConstant to avoid pydot error * Add PrepreUniformSuperposition Bloq from Cirq-FT --------- Co-authored-by: Fionn Malone <[email protected]>
quantumlib · Oct 16, 2023 · d4612af · d4612af
1 parent 8260b16
commit d4612af
Show file tree

Hide file tree

Showing 2 changed files with 190 additions and 0 deletions.
diff --git a/qualtran/bloqs/prepare_uniform_superposition.py b/qualtran/bloqs/prepare_uniform_superposition.py
@@ -0,0 +1,110 @@
+#  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.
+
+from typing import Tuple
+
+import attrs
+import cirq
+import numpy as np
+from cirq._compat import cached_property
+from numpy.typing import NDArray
+
+from qualtran import GateWithRegisters, Signature
+from qualtran.bloqs.and_bloq import And, MultiAnd
+from qualtran.bloqs.arithmetic import LessThanConstant
+
+
+@attrs.frozen
+class PrepareUniformSuperposition(GateWithRegisters):
+    r"""Prepares a uniform superposition over first $n$ basis states using $O(log(n))$ T-gates.
+
+    Performs a single round of amplitude amplification and prepares a uniform superposition over
+    the first $n$ basis states $|0>, |1>, ..., |n - 1>$. The expected T-complexity should be
+    $10 * log(L) + 2 * K$ T-gates and $2$ single qubit rotation gates, where $n = L * 2^K$.
+
+    However, the current T-complexity is $12 * log(L)$ T-gates and $2 + 2 * (K + log(L))$ rotations
+    because of two open issues:
+        - https://github.com/quantumlib/cirq-qubitization/issues/233 and
+        - https://github.com/quantumlib/cirq-qubitization/issues/235
+
+    Args:
+        n: The gate prepares a uniform superposition over first $n$ basis states.
+        cvs: Control values for each control qubit. If specified, a controlled version
+            of the gate is constructed.
+
+    References:
+        See Fig 12 of https://arxiv.org/abs/1805.03662 for more details.
+    """
+
+    n: int
+    cvs: Tuple[int, ...] = attrs.field(
+        converter=lambda v: (v,) if isinstance(v, int) else tuple(v), default=()
+    )
+
+    @cached_property
+    def signature(self) -> Signature:
+        return Signature.build(ctrl=len(self.cvs), target=(self.n - 1).bit_length())
+
+    def _circuit_diagram_info_(self, args: cirq.CircuitDiagramInfoArgs) -> cirq.CircuitDiagramInfo:
+        control_symbols = ["@" if cv else "@(0)" for cv in self.cvs]
+        target_symbols = ['target'] * self.signature.get_left('target').total_bits()
+        target_symbols[0] = f"UNIFORM({self.n})"
+        return cirq.CircuitDiagramInfo(wire_symbols=control_symbols + target_symbols)
+
+    def decompose_from_registers(
+        self,
+        *,
+        context: cirq.DecompositionContext,
+        **quregs: NDArray[cirq.Qid],  # type:ignore[type-var]
+    ) -> cirq.OP_TREE:
+        controls, target = quregs.get('ctrl', ()), quregs['target']
+        # Find K and L as per https://arxiv.org/abs/1805.03662 Fig 12.
+        n, k = self.n, 0
+        while n > 1 and n % 2 == 0:
+            k += 1
+            n = n // 2
+        l, logL = int(n), self.signature.get_left('target').total_bits() - k
+        logL_qubits = target[:logL]
+
+        yield [
+            op.controlled_by(*controls, control_values=self.cvs) for op in cirq.H.on_each(*target)
+        ]
+        if not len(logL_qubits):
+            return
+
+        ancilla = context.qubit_manager.qalloc(1)
+        theta = np.arccos(1 - (2 ** np.floor(np.log2(l))) / l)
+        yield LessThanConstant(logL, l).on_registers(x=logL_qubits, target=ancilla)
+        yield cirq.Rz(rads=theta)(*ancilla)
+        yield LessThanConstant(logL, l).on_registers(x=logL_qubits, target=ancilla) ** -1
+        context.qubit_manager.qfree(ancilla)
+
+        yield cirq.H.on_each(*logL_qubits)
+
+        and_ancilla = context.qubit_manager.qalloc(len(self.cvs) + logL - 2)
+        and_target = context.qubit_manager.qalloc(1)
+        and_cv = (0,) * logL + self.cvs
+        ctrl = np.asarray([*logL_qubits, *controls])[:, np.newaxis]
+        junk = np.asarray(and_ancilla)[:, np.newaxis]
+        if len(and_cv) <= 2:
+            and_op = And(*and_cv).on_registers(ctrl=ctrl, target=and_target)
+        else:
+            and_op = MultiAnd(cvs=and_cv).on_registers(ctrl=ctrl, junk=junk, target=and_target)
+
+        yield and_op
+        yield cirq.Rz(rads=theta)(*and_target)
+        yield cirq.inverse(and_op)
+
+        yield cirq.H.on_each(*logL_qubits)
+        context.qubit_manager.qfree([*and_target, *and_ancilla])
diff --git a/qualtran/bloqs/prepare_uniform_superposition_test.py b/qualtran/bloqs/prepare_uniform_superposition_test.py
@@ -0,0 +1,80 @@
+#  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 numpy as np
+import pytest
+
+from qualtran._infra.gate_with_registers import total_bits
+from qualtran.bloqs.prepare_uniform_superposition import PrepareUniformSuperposition
+from qualtran.cirq_interop.t_complexity_protocol import t_complexity
+from qualtran.testing import assert_valid_bloq_decomposition
+
+
+@pytest.mark.parametrize("n", [*range(3, 20), 25, 41])
+@pytest.mark.parametrize("num_controls", [0, 1])
+def test_prepare_uniform_superposition(n, num_controls):
+    gate = PrepareUniformSuperposition(n, cvs=[1] * num_controls)
+    all_qubits = cirq.LineQubit.range(cirq.num_qubits(gate))
+    control, target = (all_qubits[:num_controls], all_qubits[num_controls:])
+    turn_on_controls = [cirq.X(c) for c in control]
+    prepare_uniform_op = gate.on(*control, *target)
+    circuit = cirq.Circuit(turn_on_controls, prepare_uniform_op)
+    result = cirq.Simulator(dtype=np.complex128).simulate(circuit, qubit_order=all_qubits)
+    final_target_state = cirq.sub_state_vector(
+        result.final_state_vector,
+        keep_indices=list(range(num_controls, num_controls + len(target))),
+    )
+    expected_target_state = np.asarray([np.sqrt(1.0 / n)] * n + [0] * (2 ** len(target) - n))
+    cirq.testing.assert_allclose_up_to_global_phase(
+        expected_target_state, final_target_state, atol=1e-6
+    )
+    assert_valid_bloq_decomposition(gate)
+
+
+@pytest.mark.parametrize("n", [*range(3, 41, 3)])
+def test_prepare_uniform_superposition_t_complexity(n: int):
+    gate = PrepareUniformSuperposition(n)
+    result = t_complexity(gate)
+    assert result.rotations <= 2
+    # TODO(#235): Uncomputing `LessThanGate` should take 0 T-gates instead of 4 * n
+    # and therefore the total complexity should come down to `8 * logN`
+    assert result.t <= 12 * (n - 1).bit_length()
+
+    gate = PrepareUniformSuperposition(n, cvs=(1,))
+    result = t_complexity(gate)
+    # TODO(#233): Controlled-H is currently counted as a separate rotation, but it can be
+    # implemented using 2 T-gates.
+    assert result.rotations <= 2 + 2 * total_bits(gate.signature)
+    assert result.t <= 12 * (n - 1).bit_length()
+
+
+def test_prepare_uniform_superposition_consistent_protocols():
+    gate = PrepareUniformSuperposition(5, cvs=(1, 0))
+    # Diagrams
+    expected_symbols = ('@', '@(0)', 'UNIFORM(5)', 'target', 'target')
+    assert cirq.circuit_diagram_info(gate).wire_symbols == expected_symbols
+    # Equality
+    equals_tester = cirq.testing.EqualsTester()
+    equals_tester.add_equality_group(
+        PrepareUniformSuperposition(5, cvs=(1, 0)), PrepareUniformSuperposition(5, cvs=[1, 0])
+    )
+    equals_tester.add_equality_group(
+        PrepareUniformSuperposition(5, cvs=(0, 1)), PrepareUniformSuperposition(5, cvs=[0, 1])
+    )
+    equals_tester.add_equality_group(
+        PrepareUniformSuperposition(5),
+        PrepareUniformSuperposition(5, cvs=()),
+        PrepareUniformSuperposition(5, cvs=[]),
+    )