ZPow (known angle) using phase gradient (#1275)

* zpow(t) using phase grad

* str

* add unitary test

* docstring

* qubits

qualtran/bloqs/rotations/zpow_via_phase_gradient.py

@@ -0,0 +1,153 @@
+#  Copyright 2024 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 functools import cached_property
+from typing import cast
+
+import sympy
+from attrs import frozen
+
+from qualtran import (
+    Bloq,
+    bloq_example,
+    BloqBuilder,
+    DecomposeTypeError,
+    QBit,
+    QFxp,
+    Signature,
+    Soquet,
+    SoquetT,
+)
+from qualtran.bloqs.arithmetic import XorK
+from qualtran.bloqs.rotations.phase_gradient import AddIntoPhaseGrad
+from qualtran.resource_counting import BloqCountT, SympySymbolAllocator
+from qualtran.resource_counting.generalizers import ignore_alloc_free
+from qualtran.symbolics import ceil, is_symbolic, log2, pi, SymbolicFloat, SymbolicInt
+
+
+@frozen
+class ZPowConstViaPhaseGradient(Bloq):
+    r"""Apply an $Z**t$ on a qubit using a phase gradient state.
+
+    This bloq implements a `Z**t` by conditionally loading `t/2` into a quantum
+    register, conditioned on the qubit `q` (rotation target), and then adding
+    this value to the phase gradient to get a phase kickback, and uncomputes the load.
+    This controlled-load trick is taken from Ref. [2], Fig 2a.
+
+    See :class:`PhaseGradientState` for details on phase gradients.
+
+    It loads an approximation of `t/2` to `phase_grad_bitsize` bits,
+    which is loaded using `phase_grad_bitsize` clean ancilla.
+
+    The total Tofolli cost is `phase_grad_bitsize - 2`.
+
+
+    Args:
+        exponent: value of `t` to apply `Z**t`
+        phase_grad_bitsize: number of qubits of the phase gradient state.
+
+    Registers:
+        q: qubit to apply rotation on.
+        phase_grad: phase gradient state of type `QFxp` with `phase_grad_bitsize` fractional bits.
+
+    References:
+        [Improved quantum circuits for elliptic curve discrete logarithms](https://arxiv.org/abs/2001.09580).
+        Haner et. al. 2020. Section 3: Components. "Integer addition" and Fig 2a.
+    """
+    exponent: SymbolicFloat
+    phase_grad_bitsize: SymbolicInt
+
+    @property
+    def signature(self) -> 'Signature':
+        return Signature.build_from_dtypes(q=QBit(), phase_grad=self.phase_grad_dtype)
+
+    @classmethod
+    def from_precision(
+        cls, exponent: SymbolicFloat, *, eps: SymbolicFloat
+    ) -> 'ZPowConstViaPhaseGradient':
+        r"""Apply a ZPow(t) with precision `eps`.
+
+        Uses a phase gradient of size $\ceil(\log(2\pi / \epsilon)$.
+
+        Args:
+            exponent: value of `t` to apply `Z**t`
+            eps: precision to approximate the unitary to.
+        """
+        b_grad = ceil(log2(2 * pi(eps) / eps))
+        return cls(exponent, b_grad)
+
+    @cached_property
+    def phase_grad_dtype(self) -> QFxp:
+        return QFxp(self.phase_grad_bitsize, self.phase_grad_bitsize)
+
+    @property
+    def _load_bloq(self) -> XorK:
+        if is_symbolic(self.exponent) or is_symbolic(self.phase_grad_bitsize):
+            return XorK(self.phase_grad_dtype, cast(sympy.Expr, self.exponent / 2))
+
+        k_int = self.phase_grad_dtype.to_fixed_width_int(self.exponent / 2)
+        return XorK(self.phase_grad_dtype, k_int)
+
+    def build_composite_bloq(
+        self, bb: BloqBuilder, q: Soquet, phase_grad: Soquet
+    ) -> dict[str, SoquetT]:
+        if is_symbolic(self.exponent):
+            raise DecomposeTypeError(f"cannot decompose {self} with symbolic {self.exponent=}")
+
+        # load the angle
+        t = bb.allocate(dtype=self.phase_grad_dtype)
+        q, t = bb.add(self._load_bloq.controlled(), ctrl=q, x=t)
+
+        # add
+        t, phase_grad = bb.add(
+            AddIntoPhaseGrad(
+                x_bitsize=self.phase_grad_bitsize, phase_bitsize=self.phase_grad_bitsize
+            ),
+            x=t,
+            phase_grad=phase_grad,
+        )
+
+        # unload the angle
+        q, t = bb.add(self._load_bloq.controlled(), ctrl=q, x=t)
+        bb.free(t)
+
+        return {'q': q, 'phase_grad': phase_grad}
+
+    def build_call_graph(self, ssa: 'SympySymbolAllocator') -> set['BloqCountT']:
+        return {
+            (self._load_bloq.controlled(), 2),
+            (AddIntoPhaseGrad(self.phase_grad_bitsize, self.phase_grad_bitsize), 1),
+        }
+
+    def __str__(self) -> str:
+        return f'ZPow({self.exponent})'
+
+
+@bloq_example(generalizer=ignore_alloc_free)
+def _zpow_const_via_phase_grad() -> ZPowConstViaPhaseGradient:
+    zpow_const_via_phase_grad = ZPowConstViaPhaseGradient.from_precision(3 / 8, eps=1e-11)
+    return zpow_const_via_phase_grad
+
+
+@bloq_example(generalizer=ignore_alloc_free)
+def _zpow_const_via_phase_grad_symb_prec() -> ZPowConstViaPhaseGradient:
+    eps = sympy.symbols(r"\epsilon")
+    zpow_const_via_phase_grad_symb_prec = ZPowConstViaPhaseGradient.from_precision(3 / 8, eps=eps)
+    return zpow_const_via_phase_grad_symb_prec
+
+
+@bloq_example(generalizer=ignore_alloc_free)
+def _zpow_const_via_phase_grad_symb_angle() -> ZPowConstViaPhaseGradient:
+    t = sympy.symbols(r"t")
+    zpow_const_via_phase_grad_symb_angle = ZPowConstViaPhaseGradient.from_precision(t, eps=1e-11)
+    return zpow_const_via_phase_grad_symb_angle

qualtran/bloqs/rotations/zpow_via_phase_gradient_test.py

@@ -0,0 +1,97 @@
+#  Copyright 2024 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 pytest
+from attrs import frozen
+
+from qualtran import Bloq, BloqBuilder, Signature, Soquet, SoquetT
+from qualtran.bloqs.basic_gates import ZPowGate
+from qualtran.bloqs.rotations.phase_gradient import PhaseGradientState
+from qualtran.bloqs.rotations.zpow_via_phase_gradient import (
+    _zpow_const_via_phase_grad,
+    _zpow_const_via_phase_grad_symb_angle,
+    _zpow_const_via_phase_grad_symb_prec,
+    ZPowConstViaPhaseGradient,
+)
+from qualtran.linalg.testing import assert_unitaries_equivalent_upto_global_phase
+from qualtran.resource_counting import GateCounts, get_cost_value, QECGatesCost
+
+
+@pytest.mark.parametrize(
+    "bloq",
+    [
+        _zpow_const_via_phase_grad,
+        _zpow_const_via_phase_grad_symb_prec,
+        _zpow_const_via_phase_grad_symb_angle,
+    ],
+)
+def test_examples(bloq_autotester, bloq):
+    if bloq_autotester.check_name == 'serialize':
+        pytest.skip()
+
+    bloq_autotester(bloq)
+
+
+def test_cost():
+    bloq = _zpow_const_via_phase_grad()
+
+    costs = get_cost_value(bloq, QECGatesCost())
+    b_grad = bloq.phase_grad_bitsize
+    assert costs == GateCounts(toffoli=b_grad - 2, clifford=4)
+
+
+def test_cost_symbolic_angle():
+    bloq = _zpow_const_via_phase_grad_symb_angle()
+
+    costs = get_cost_value(bloq, QECGatesCost())
+    b_grad = bloq.phase_grad_bitsize
+    assert costs == GateCounts(toffoli=b_grad - 2, clifford=80)
+
+
+def test_cost_symbolic_prec():
+    bloq = _zpow_const_via_phase_grad_symb_prec()
+
+    costs = get_cost_value(bloq, QECGatesCost())
+    b_grad = bloq.phase_grad_bitsize
+    assert costs == GateCounts(toffoli=b_grad - 2, clifford=2 * b_grad)
+
+
+@frozen
+class TestZPowUsingPhaseGradient(Bloq):
+    exponent: float
+    phase_grad_bitsize: int
+
+    @property
+    def signature(self) -> 'Signature':
+        return Signature.build(q=1)
+
+    def build_composite_bloq(self, bb: 'BloqBuilder', q: 'Soquet') -> dict[str, 'SoquetT']:
+        phase_grad = bb.add(PhaseGradientState(self.phase_grad_bitsize))
+        q, phase_grad = bb.add(
+            ZPowConstViaPhaseGradient(self.exponent, self.phase_grad_bitsize),
+            q=q,
+            phase_grad=phase_grad,
+        )
+        bb.add(PhaseGradientState(self.phase_grad_bitsize).adjoint(), phase_grad=phase_grad)
+        return {'q': q}
+
+
+@pytest.mark.parametrize(
+    ("exponent", "phase_grad_bitsize"),
+    [(3 / 4, 3), (3 / 4, 4), pytest.param(0.175, 7, marks=pytest.mark.slow)],
+)
+def test_unitary(exponent: float, phase_grad_bitsize: int):
+    bloq = TestZPowUsingPhaseGradient(exponent, phase_grad_bitsize)
+    actual = bloq.tensor_contract()
+    expected = ZPowGate(exponent).tensor_contract()
+    assert_unitaries_equivalent_upto_global_phase(actual, expected, atol=2**-phase_grad_bitsize)

qualtran/linalg/matrix.py

@@ -21,3 +21,31 @@ def random_hermitian_matrix(dim: int, random_state: np.random.RandomState) -> ND
     """Return a random Hermitian matrix of given dimension and norm <= 1."""
     a = random_orthogonal(dim, random_state=random_state)
     return (a + a.conj().T) / 2
+
+
+def unitary_distance_ignoring_global_phase(U: NDArray, V: NDArray) -> float:
+    r"""Distance between two unitaries ignoring global phase.
+
+    Given two $N \times N$ unitaries $U, V$, their distance is given by (Eq. 1 of Ref. [1]):
+
+        $$
+            \sqrt{\frac{N - \abs{\tr(U^\dagger V)}}{N}}
+        $$
+
+    Args:
+        U: N x N unitary matrix
+        V: N x N unitary matrix
+
+    References:
+        [Constructing arbitrary Steane code single logical qubit fault-tolerant gates](https://arxiv.org/abs/quant-ph/0411206)
+        Fowler. 2010. Section 1, Eq 1.
+    """
+    if U.shape != V.shape:
+        raise ValueError(f"Inputs must have same dimension, got {U.shape}, {V.shape}")
+    if U.ndim != 2 or U.shape[0] != U.shape[1]:
+        raise ValueError(f"Inputs must be square matrices, got shape {U.shape}")
+
+    N = U.shape[0]
+    trace = np.trace(U.conj().T @ V)
+    d_squared = 1 - np.abs(trace) / N
+    return np.sqrt(max(0, d_squared))

qualtran/linalg/matrix_test.py

@@ -0,0 +1,28 @@
+#  Copyright 2024 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 numpy as np
+import pytest
+from cirq.testing import random_unitary
+
+from qualtran.linalg.matrix import unitary_distance_ignoring_global_phase
+
+
+@pytest.mark.parametrize("dim", [2, 3])
+def test_unitary_distance_zero(dim: int):
+    rs = np.random.RandomState(1234)
+    for _ in range(3):
+        U = random_unitary(dim, random_state=rs)
+        V = np.exp(1j * rs.random()) * U
+        d = unitary_distance_ignoring_global_phase(U, V)
+        np.testing.assert_almost_equal(d, 0)

qualtran/linalg/testing.py

@@ -14,6 +14,8 @@
 import numpy as np
 from numpy.typing import NDArray
 
+from qualtran.linalg.matrix import unitary_distance_ignoring_global_phase
+
 
 def assert_matrices_almost_equal(A: NDArray, B: NDArray, *, atol: float = 1e-5):
     r"""Asserts that two matrices are close to each other by bounding the matrix norm of their difference.
@@ -22,3 +24,8 @@ def assert_matrices_almost_equal(A: NDArray, B: NDArray, *, atol: float = 1e-5):
     """
     assert A.shape == B.shape
     assert np.linalg.norm(A - B) <= atol
+
+
+def assert_unitaries_equivalent_upto_global_phase(U: NDArray, V: NDArray, *, atol: float = 1e-5):
+    d = unitary_distance_ignoring_global_phase(U, V)
+    assert d <= atol, f"unitaries are not equivalent: distance={d} ({atol=})"