[GQSP] support negative powers (#810)

* [GQSP] support negative powers

* fix notebook

* call graph: only add non-zero counts, add test

* fix `RandomGate.adjoint`

qualtran/bloqs/generalized_qsp.ipynb

@@ -16,9 +16,25 @@
    "outputs": [],
    "source": [
     "from qualtran.bloqs.qubitization_walk_operator_test import get_walk_operator_for_1d_ising_model\n",
-    "from qualtran.drawing import show_bloq\n",
+    "from qualtran.bloqs.generalized_qsp import GeneralizedQSP\n",
+    "from qualtran.drawing import show_bloq"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "6963c30f339d42de",
+   "metadata": {},
+   "source": [
+    "`GeneralizedQSP` implements the Quantum Eigenvalue Transform on a unitary $U$ using QSP. Given a complex GQSP polynomial $P$ (and its complement $Q$), it implements the unitary:\n",
+    "$$U' = \\begin{bmatrix} P(U) & \\cdot \\\\ Q(U) & \\cdot \\end{bmatrix}$$\n",
+    "\n",
+    "Here, the polynomials $P, Q$ must satisfy the following constraint:\n",
     "\n",
-    "from qualtran.bloqs.generalized_qsp import GeneralizedQSP"
+    "$$\\left| P(e^{i\\theta}) \\right|^2 + \\left| Q(e^{i\\theta}) \\right|^2 = 1 ~~\\text{for every}~ \\theta \\in [0, 2\\pi]$$\n",
+    "\n",
+    "A polynomial $P$ is said to be a GQSP polynomial iff it satisfies $\\left| P(e^{i\\theta}) \\right|^2 \\le 1$ for every $\\theta \\in [0, 2\\pi]$. \n",
+    "\n",
+    "Reference: https://doi.org/10.48550/arXiv.2308.01501"
    ]
   },
   {
@@ -32,42 +48,56 @@
     "show_bloq(U.decompose_bloq())"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "dc7100fd94d58c6b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "pU = GeneralizedQSP(U, (0.5, 0.5), (-0.5, 0.5))\n",
+    "show_bloq(pU.decompose_bloq())"
+   ]
+  },
   {
    "cell_type": "markdown",
-   "id": "6963c30f339d42de",
+   "id": "935a03f7-5843-4b11-abe6-5eb9048c0ab5",
    "metadata": {},
    "source": [
-    "`GeneralizedQSP` implements the Quantum Eigenvalue Transform on a unitary $U$ using QSP. Given a complex QSP polynomial $P$ (and its complement $Q$), it implements the unitary:\n",
-    "$$U' = \\begin{bmatrix} P(U) & \\cdot \\\\ Q(U) & \\cdot \\end{bmatrix}$$\n",
-    "\n",
-    "Here, the polynomials $P, Q$ must satisfy the following constraint:\n",
-    "\n",
-    "$$\\left\\mid P(e^{i\\theta}) \\right\\mid^2 + \\left\\mid Q(e^{i\\theta}) \\right\\mid^2 = 1 ~~\\text{for every}~ \\theta \\in [0, 2\\pi]$$\n",
-    "\n",
-    "\n",
-    "Reference: https://arxiv.org/abs/2308.01501"
+    "There is also a method that directly computes $Q$ from $P$:"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "dc7100fd94d58c6b",
+   "id": "78cd3857297f092b",
    "metadata": {},
    "outputs": [],
    "source": [
-    "pU = GeneralizedQSP(U, (0.5, 0.5))\n",
+    "pU = GeneralizedQSP.from_qsp_polynomial(U, (0.5, 0, 0.5))\n",
     "show_bloq(pU.decompose_bloq())"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "a58f06ba-9287-435d-92a3-256f747024c2",
+   "metadata": {},
+   "source": [
+    "### Negative degree terms\n",
+    "\n",
+    "To apply GQSP for a polynomial $P'(z) = z^{-k} P(z)$, we can just pass the polynomial $P$ along with negative power $k$.\n",
+    "The QSP angle sequence is the same for both, and $P'$ can be achieved by running $(U^\\dagger)^k$ at any point in the circuit."
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "78cd3857297f092b",
+   "id": "ee60e95f-979f-420d-91b1-b6955b9c5a3b",
    "metadata": {},
    "outputs": [],
    "source": [
-    "pU = GeneralizedQSP(U, (0.5, 0, 0.5))\n",
-    "show_bloq(pU.decompose_bloq())"
+    "pU = GeneralizedQSP.from_qsp_polynomial(U, (0.5, 0, 0.5), negative_power=1)\n",
+    "show_bloq(pU)"
    ]
   }
  ],

qualtran/bloqs/generalized_qsp.py

@@ -12,11 +12,11 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 from functools import cached_property
-from typing import Sequence, Tuple
+from typing import Sequence, Set, Tuple
 
 import cirq
 import numpy as np
-from attrs import frozen
+from attrs import field, frozen
 from numpy.polynomial import Polynomial
 from numpy.typing import NDArray
 
@@ -239,9 +239,11 @@ def safe_angle(x):
 class GeneralizedQSP(GateWithRegisters):
     r"""Applies a QSP polynomial $P$ to a unitary $U$ to obtain a block-encoding of $P(U)$.
 
+    Can optionally provide a negative power offset $k$ (defaults to 0),
+    to obtain $U^{-k} P(U)$. (Theorem 6)
     This gate represents the following unitary:
 
-        $$ \begin{bmatrix} P(U) & \cdot \\ \cdot & \cdot \end{bmatrix} $$
+        $$ \begin{bmatrix} U^{-k} P(U) & \cdot \\ Q(U) & \cdot \end{bmatrix} $$
 
     The polynomial $P$ must satisfy:
     $\abs{P(e^{i \theta})}^2 \le 1$ for every $\theta \in \mathbb{R}$.
@@ -251,22 +253,28 @@ class GeneralizedQSP(GateWithRegisters):
     Args:
         U: Unitary operation.
         P: Co-efficients of a complex polynomial.
+        negative_power: value of $k$, which effectively applies $z^{-k} P(z)$. defaults to 0.
 
     References:
         [Generalized Quantum Signal Processing](https://arxiv.org/abs/2308.01501)
-            Motlagh and Wiebe. (2023). Theorem 3; Figure 2.
+            Motlagh and Wiebe. (2023). Theorem 3; Figure 2; Theorem 6.
     """
 
     U: GateWithRegisters
-    P: Sequence[complex]
+    P: Tuple[complex, ...] = field(converter=tuple)
+    Q: Tuple[complex, ...] = field(converter=tuple)
+    negative_power: int = field(default=0, kw_only=True)
 
     @cached_property
     def signature(self) -> Signature:
         return Signature([Register('signal', QBit()), *self.U.signature])
 
-    @cached_property
-    def Q(self):
-        return qsp_complementary_polynomial(self.P)
+    @classmethod
+    def from_qsp_polynomial(
+        cls, U: GateWithRegisters, P: Sequence[complex], *, negative_power: int = 0
+    ) -> 'GeneralizedQSP':
+        Q = qsp_complementary_polynomial(P)
+        return GeneralizedQSP(U, P, Q, negative_power=negative_power)
 
     @cached_property
     def _qsp_phases(self) -> Tuple[Sequence[float], Sequence[float], float]:
@@ -284,13 +292,48 @@ def _phi(self) -> Sequence[float]:
     def _lambda(self) -> float:
         return self._qsp_phases[2]
 
+    @cached_property
+    def signal_rotations(self) -> NDArray[SU2RotationGate]:
+        return np.array(
+            [
+                SU2RotationGate(theta, phi, self._lambda if i == 0 else 0)
+                for i, (theta, phi) in enumerate(zip(self._theta, self._phi))
+            ]
+        )
+
     def decompose_from_registers(
         self, *, context: cirq.DecompositionContext, signal, **quregs: NDArray[cirq.Qid]
     ) -> cirq.OP_TREE:
         assert len(signal) == 1
         signal_qubit = signal[0]
 
-        yield SU2RotationGate(self._theta[0], self._phi[0], self._lambda).on(signal_qubit)
-        for theta, phi in zip(self._theta[1:], self._phi[1:]):
-            yield self.U.on_registers(**quregs).controlled_by(signal_qubit, control_values=[0])
-            yield SU2RotationGate(theta, phi, 0).on(signal_qubit)
+        num_inverse_applications = self.negative_power
+
+        yield self.signal_rotations[0].on(signal_qubit)
+        for signal_rotation in self.signal_rotations[1:]:
+            if num_inverse_applications > 0:
+                # apply C-U^\dagger
+                yield self.U.adjoint().on_registers(**quregs).controlled_by(signal_qubit)
+                num_inverse_applications -= 1
+            else:
+                # apply C[0]-U
+                yield self.U.on_registers(**quregs).controlled_by(signal_qubit, control_values=[0])
+            yield signal_rotation.on(signal_qubit)
+
+        for _ in range(num_inverse_applications):
+            yield self.U.adjoint().on_registers(**quregs)
+
+    def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']:
+        degree = len(self.P)
+
+        counts = {(rotation, 1) for rotation in self.signal_rotations}
+
+        if degree > self.negative_power:
+            counts.add((self.U.controlled(control_values=[0]), degree - self.negative_power))
+        elif self.negative_power > degree:
+            counts.add((self.U.adjoint(), self.negative_power - degree))
+
+        if self.negative_power > 0:
+            counts.add((self.U.adjoint().controlled(), min(degree, self.negative_power)))
+
+        return counts

qualtran/bloqs/generalized_qsp_test.py

@@ -12,24 +12,25 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 from functools import cached_property
-from typing import Sequence, Tuple, Union
+from typing import Optional, Sequence, Tuple, Union
 
 import cirq
 import numpy as np
 import pytest
 import sympy
-from attrs import define, frozen
+from attrs import define, field, frozen
 from cirq.testing import random_unitary
 from numpy.polynomial import Polynomial
 from numpy.typing import NDArray
 
-from qualtran import GateWithRegisters, Signature
+from qualtran import Bloq, GateWithRegisters, Signature
 from qualtran.bloqs.generalized_qsp import (
     GeneralizedQSP,
     qsp_complementary_polynomial,
     qsp_phase_factors,
     SU2RotationGate,
 )
+from qualtran.resource_counting import SympySymbolAllocator
 
 
 def assert_angles_almost_equal(
@@ -99,6 +100,7 @@ def test_complementary_polynomial(degree: int):
         check_polynomial_pair_on_random_points_on_unit_circle(P, Q, random_state=random_state)
 
 
+@pytest.mark.slow
 @pytest.mark.parametrize("degree", [3, 4, 5, 10, 20, 30, 100])
 def test_real_polynomial_has_real_complementary_polynomial(degree: int):
     random_state = np.random.RandomState(42)
@@ -113,12 +115,14 @@ def test_real_polynomial_has_real_complementary_polynomial(degree: int):
 @frozen
 class RandomGate(GateWithRegisters):
     bitsize: int
-    matrix: Tuple[Tuple[int, ...], ...]
+    matrix: Tuple[Tuple[complex, ...], ...] = field(
+        converter=lambda mat: tuple(tuple(row) for row in mat)
+    )
 
     @staticmethod
     def create(bitsize: int, *, random_state=None) -> 'RandomGate':
         matrix = random_unitary(2**bitsize, random_state=random_state)
-        return RandomGate(bitsize, tuple(tuple(x) for x in matrix.tolist()))
+        return RandomGate(bitsize, matrix)
 
     @property
     def signature(self) -> Signature:
@@ -127,11 +131,21 @@ def signature(self) -> Signature:
     def _unitary_(self):
         return np.array(self.matrix)
 
+    def adjoint(self) -> 'RandomGate':
+        return RandomGate(self.bitsize, np.conj(self.matrix).T)
+
+    def __pow__(self, power):
+        if power == -1:
+            return self.adjoint()
+        return NotImplemented
 
-def evaluate_polynomial_of_matrix(P: Sequence[complex], U: NDArray) -> NDArray:
+
+def evaluate_polynomial_of_matrix(
+    P: Sequence[complex], U: NDArray, *, negative_power: int = 0
+) -> NDArray:
     assert U.ndim == 2 and U.shape[0] == U.shape[1]
 
-    pow_U = np.identity(U.shape[0], dtype=U.dtype)
+    pow_U = np.linalg.matrix_power(U.conj().T, negative_power)
     result = np.zeros(U.shape, dtype=U.dtype)
 
     for c in P:
@@ -146,17 +160,30 @@ def assert_matrices_almost_equal(A: NDArray, B: NDArray):
     assert np.linalg.norm(A - B) <= 1e-5
 
 
-def verify_generalized_qsp(U: GateWithRegisters, P: Sequence[complex]):
+def verify_generalized_qsp(
+    U: GateWithRegisters,
+    P: Sequence[complex],
+    Q: Optional[Sequence[complex]] = None,
+    *,
+    negative_power: int = 0,
+):
     input_unitary = cirq.unitary(U)
     N = input_unitary.shape[0]
-    gqsp_U = GeneralizedQSP(U, P)
+    if Q is None:
+        gqsp_U = GeneralizedQSP.from_qsp_polynomial(U, P, negative_power=negative_power)
+    else:
+        gqsp_U = GeneralizedQSP(U, P, Q, negative_power=negative_power)
     result_unitary = cirq.unitary(gqsp_U)
 
-    expected_top_left = evaluate_polynomial_of_matrix(P, input_unitary)
+    expected_top_left = evaluate_polynomial_of_matrix(
+        P, input_unitary, negative_power=negative_power
+    )
     actual_top_left = result_unitary[:N, :N]
     assert_matrices_almost_equal(expected_top_left, actual_top_left)
 
-    expected_bottom_left = evaluate_polynomial_of_matrix(gqsp_U.Q, input_unitary)
+    expected_bottom_left = evaluate_polynomial_of_matrix(
+        gqsp_U.Q, input_unitary, negative_power=negative_power
+    )
     actual_bottom_left = result_unitary[N:, :N]
     assert_matrices_almost_equal(expected_bottom_left, actual_bottom_left)
 
@@ -176,13 +203,44 @@ def test_generalized_qsp_with_real_poly_on_random_unitaries(bitsize: int, degree
 @pytest.mark.slow
 @pytest.mark.parametrize("bitsize", [1, 2, 3])
 @pytest.mark.parametrize("degree", [2, 3, 4, 5, 50, 100, 120])
-def test_generalized_qsp_with_complex_poly_on_random_unitaries(bitsize: int, degree: int):
+@pytest.mark.parametrize("negative_power", [0, 1, 2])
+def test_generalized_qsp_with_complex_poly_on_random_unitaries(
+    bitsize: int, degree: int, negative_power: int
+):
     random_state = np.random.RandomState(42)
 
     for _ in range(10):
         U = RandomGate.create(bitsize, random_state=random_state)
         P = random_qsp_polynomial(degree, random_state=random_state)
-        verify_generalized_qsp(U, P)
+        verify_generalized_qsp(U, P, negative_power=negative_power)
+
+
+@pytest.mark.parametrize("negative_power", [0, 1, 2])
+def test_call_graph(negative_power: int):
+    random_state = np.random.RandomState(42)
+
+    ssa = SympySymbolAllocator()
+    theta = ssa.new_symbol("theta")
+    phi = ssa.new_symbol("phi")
+    lambd = ssa.new_symbol("lambda")
+    arbitrary_rotation = SU2RotationGate(theta, phi, lambd)
+
+    def catch_rotations(bloq: Bloq) -> Bloq:
+        if isinstance(bloq, SU2RotationGate):
+            return arbitrary_rotation
+        return bloq
+
+    U = RandomGate.create(1, random_state=random_state)
+    P = (0.5, 0, 0.5)
+    gsqp_U = GeneralizedQSP.from_qsp_polynomial(U, P, negative_power=negative_power)
+
+    g, sigma = gsqp_U.call_graph(max_depth=1, generalizer=catch_rotations)
+
+    expected_counts = {U.controlled(control_values=[0]): 3 - negative_power, arbitrary_rotation: 3}
+    if negative_power > 0:
+        expected_counts[U.adjoint().controlled()] = negative_power
+
+    assert sigma == expected_counts
 
 
 @define(slots=False)