GQSP template and angle solver

doc/introduction/templates.rst

@@ -326,6 +326,10 @@ Other useful templates which do not belong to the previous categories can be fou
     :description: :doc:`QSVT<../code/api/pennylane.QSVT>`
     :figure: _static/templates/subroutines/qsvt.png
 
+.. gallery-item::
+    :description: :doc:`GQSP<../code/api/pennylane.GQSP>`
+    :figure: _static/templates/subroutines/gqsp.png
+
 .. gallery-item::
     :description: :doc:`Select<../code/api/pennylane.Select>`
     :figure: _static/templates/subroutines/select.png

doc/releases/changelog-dev.md

@@ -57,6 +57,11 @@
   [(#6426)](https://github.com/PennyLaneAI/pennylane/pull/6426)
   [(#6645)](https://github.com/PennyLaneAI/pennylane/pull/6645)
 
+  * New `qml.GQSP` template has been added to perform Generalized Quantum Signal Processing (GQSP).
+    The functionality `qml.poly_to_angles` has been also extended to support GQSP.
+    [(#6565)](https://github.com/PennyLaneAI/pennylane/pull/6565)
+
+
 <h4>New `labs` module `dla` for handling dynamical Lie algebras (DLAs)</h4>
 
 * Added a dense implementation of computing the Lie closure in a new function

pennylane/templates/subroutines/__init__.py

@@ -52,3 +52,4 @@
 from .out_adder import OutAdder
 from .mod_exp import ModExp
 from .out_poly import OutPoly
+from .gqsp import GQSP

pennylane/templates/subroutines/gqsp.py

@@ -0,0 +1,167 @@
+# Copyright 2024 Xanadu Quantum Technologies Inc.
+
+# 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
+
+#     http://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.
+"""
+Contains the GQSP template.
+"""
+# pylint: disable=too-many-arguments
+
+import copy
+
+import pennylane as qml
+from pennylane.operation import Operation
+from pennylane.queuing import QueuingManager
+from pennylane.wires import Wires
+
+
+class GQSP(Operation):
+    r"""
+    Implements the generalized quantum signal processing (GQSP) circuit.
+
+    This operation encodes a polynomial transformation of an input unitary operator following the algorithm
+    described in `arXiv:2308.01501 <https://arxiv.org/abs/2308.01501>`__ as:
+
+    .. math::
+         U
+         \xrightarrow{GQSP}
+         \begin{pmatrix}
+         \text{poly}(U) & * \\
+         * & * \\
+         \end{pmatrix}
+
+    The implementation requires one auxiliary qubit.
+
+    Args:
+
+        unitary (Operator): the operator to be encoded by the GQSP circuit
+        angles (tensor[float]): array of angles defining the polynomial transformation. The shape of the array must be `(3, d+1)`, where `d` is the degree of the polynomial.
+        control (Union[Wires, int, str]): control qubit used to encode the polynomial transformation
+
+    .. note::
+
+       The  :func:`~.poly_to_angles` function can be used to calculate the angles for a given polynomial.
+
+    Example:
+
+    .. code-block:: python
+
+        # P(x) = 0.1 + 0.2j x + 0.3 x^2
+        poly = [0.1, 0.2j, 0.3]
+
+        angles = qml.poly_to_angles(poly, "GQSP")
+
+        @qml.prod # transforms the qfunc into an Operator
+        def unitary(wires):
+            qml.RX(0.3, wires)
+
+        dev = qml.device("default.qubit")
+
+        @qml.qnode(dev)
+        def circuit(angles):
+            qml.GQSP(unitary(wires = 1), angles, control = 0)
+            return qml.state()
+
+        matrix = qml.matrix(circuit, wire_order=[0, 1])(angles)
+
+    .. code-block:: pycon
+
+        >>> print(np.round(matrix,3)[:2, :2])
+        [[0.387+0.198j 0.03 -0.089j]
+        [0.03 -0.089j 0.387+0.198j]]
+    """
+
+    grad_method = None
+
+    def __init__(self, unitary, angles, control, id=None):
+        total_wires = qml.wires.Wires(control) + unitary.wires
+
+        self._hyperparameters = {"unitary": unitary, "control": qml.wires.Wires(control)}
+
+        super().__init__(angles, *unitary.data, wires=total_wires, id=id)
+
+    def _flatten(self):
+        data = self.parameters
+        return data, (
+            self.hyperparameters["unitary"],
+            self.hyperparameters["control"],
+        )
+
+    @classmethod
+    def _unflatten(cls, data, metadata):
+        return cls(unitary=metadata[0], angles=data[0], control=metadata[1])
+
+    @classmethod
+    def _primitive_bind_call(cls, *args, **kwargs):
+        return cls._primitive.bind(*args, **kwargs)
+
+    def map_wires(self, wire_map: dict):
+        # pylint: disable=protected-access
+        new_op = copy.deepcopy(self)
+        new_op._wires = Wires([wire_map.get(wire, wire) for wire in self.wires])
+        new_op._hyperparameters["unitary"] = qml.map_wires(
+            new_op._hyperparameters["unitary"], wire_map
+        )
+        new_op._hyperparameters["control"] = tuple(
+            wire_map.get(w, w) for w in new_op._hyperparameters["control"]
+        )
+
+        return new_op
+
+    @staticmethod
+    def compute_decomposition(*parameters, **hyperparameters):  # pylint: disable=arguments-differ
+        r"""
+        Representation of the operator as a product of other operators (static method).
+
+        .. math:: O = O_1 O_2 \dots O_n.
+
+        .. seealso:: :meth:`~.Operator.decomposition`.
+
+        Args:
+            *parameters (list): trainable parameters of the operator, as stored in the ``parameters`` attribute
+            wires (Iterable[Any], Wires): wires that the operator acts on
+            **hyperparams (dict): non-trainable hyperparameters of the operator, as stored in the ``hyperparameters`` attribute
+
+        Returns:
+            list[Operator]: decomposition of the operator
+        """
+
+        unitary = hyperparameters["unitary"]
+        control = hyperparameters["control"]
+
+        angles = parameters[0]
+
+        thetas, phis, lambds = angles[0], angles[1], angles[2]
+
+        op_list = []
+
+        # These four gates adapt PennyLane's qml.U3 to the chosen U3 format in the GQSP paper.
+        op_list.append(qml.X(control))
+        op_list.append(qml.U3(2 * thetas[0], phis[0], lambds[0], wires=control))
+        op_list.append(qml.X(control))
+        op_list.append(qml.Z(control))
+
+        for theta, phi, lamb in zip(thetas[1:], phis[1:], lambds[1:]):
+
+            op_list.append(qml.ctrl(unitary, control=control, control_values=0))
+
+            op_list.append(qml.X(control))
+            op_list.append(qml.U3(2 * theta, phi, lamb, wires=control))
+            op_list.append(qml.X(control))
+            op_list.append(qml.Z(control))
+
+        return op_list
+
+    def queue(self, context=QueuingManager):
+        context.remove(self.hyperparameters["unitary"])
+        context.append(self)
+        return self

pennylane/templates/subroutines/qsvt.py

@@ -571,6 +571,70 @@ def _compute_qsp_angle(poly_coeffs):
     return rotation_angles
 
 
+def _gqsp_u3_gate(theta, phi, lambd):
+    r"""
+    Computes the U3 gate matrix for Generalized Quantum Signal Processing (GQSP) as described
+    in [`arXiv:2406.04246 <https://arxiv.org/abs/2406.04246>`_]
+    """
+
+    exp_phi = qml.math.exp(1j * phi)
+    exp_lambda = qml.math.exp(1j * lambd)
+    exp_lambda_phi = qml.math.exp(1j * (lambd + phi))
+
+    matrix = np.array(
+        [
+            [exp_lambda_phi * qml.math.cos(theta), exp_phi * qml.math.sin(theta)],
+            [exp_lambda * qml.math.sin(theta), -qml.math.cos(theta)],
+        ],
+        dtype=complex,
+    )
+
+    return matrix
+
+
+def _compute_gqsp_angles(poly_coeffs):
+    r"""
+    Computes the Generalized Quantum Signal Processing (GQSP) angles given the coefficients of a polynomial P.
+
+    The method for computing the GQSP angles is based on the algorithm described in [`arXiv:2406.04246 <https://arxiv.org/abs/2406.04246>`_].
+    The complementary polynomial is calculated using root-finding methods.
+
+    Args:
+        poly_coeffs (tensor-like): Coefficients of the input polynomial P.
+
+    Returns:
+        (tensor-like): QSP angles corresponding to the input polynomial P. The shape is (3, P-degree)
+    """
+
+    complementary = _complementary_poly(poly_coeffs)
+
+    # Algorithm 1 in [arXiv:2308.01501]
+    input_data = qml.math.array([poly_coeffs, complementary])
+    num_elements = input_data.shape[1]
+
+    angles_theta, angles_phi, angles_lambda = qml.math.zeros([3, num_elements])
+
+    for idx in range(num_elements - 1, -1, -1):
+
+        component_a, component_b = input_data[:, idx]
+        angles_theta[idx] = qml.math.arctan2(np.abs(component_b), qml.math.abs(component_a))
+        angles_phi[idx] = (
+            0
+            if qml.math.isclose(qml.math.abs(component_b), 0, atol=1e-10)
+            else qml.math.angle(component_a * qml.math.conj(component_b))
+        )
+
+        if idx == 0:
+            angles_lambda[0] = qml.math.angle(component_b)
+        else:
+            updated_matrix = (
+                _gqsp_u3_gate(angles_theta[idx], angles_phi[idx], 0).conj().T @ input_data
+            )
+            input_data = qml.math.array([updated_matrix[0][1 : idx + 1], updated_matrix[1][0:idx]])
+
+    return angles_theta, angles_phi, angles_lambda
+
+
 def transform_angles(angles, routine1, routine2):
     r"""
     Converts angles for quantum signal processing (QSP) and quantum singular value transformation (QSVT) routines.
@@ -766,4 +830,7 @@ def circuit_qsvt():
         if angle_solver == "root-finding":
             return _compute_qsp_angle(poly)
         raise AssertionError("Invalid angle solver method. Valid value is 'root-finding'")
+
+    if routine == "GQSP":
+        return _compute_gqsp_angles(poly)
     raise AssertionError("Invalid routine. Valid values are 'QSP' and 'QSVT'")

tests/capture/test_templates.py

@@ -253,9 +253,11 @@ def fn(*args):
     qml.OutAdder,
     qml.ModExp,
     qml.OutPoly,
+    qml.GQSP,
 ]
 
 
+# pylint: disable=too-many-public-methods
 class TestModifiedTemplates:
     """Test that templates with custom primitive binds are captured as expected."""
 
@@ -927,6 +929,39 @@ def qfunc():
         assert len(q) == 1
         qml.assert_equal(q.queue[0], qml.OutPoly(**kwargs))
 
+    def test_gqsp(self):
+        """Test the primitive bind call of GQSP."""
+
+        kwargs = {
+            "unitary": qml.RX(1, wires=1),
+            "control": 0,
+            "angles": np.ones([3, 3]),
+        }
+
+        def qfunc():
+            qml.GQSP(**kwargs)
+
+        # Validate inputs
+        qfunc()
+
+        # Actually test primitive bind
+        jaxpr = jax.make_jaxpr(qfunc)()
+
+        assert len(jaxpr.eqns) == 1
+
+        eqn = jaxpr.eqns[0]
+        assert eqn.primitive == qml.GQSP._primitive
+        assert eqn.invars == jaxpr.jaxpr.invars
+        assert eqn.params == kwargs
+        assert len(eqn.outvars) == 1
+        assert isinstance(eqn.outvars[0], jax.core.DropVar)
+
+        with qml.queuing.AnnotatedQueue() as q:
+            jax.core.eval_jaxpr(jaxpr.jaxpr, jaxpr.consts)
+
+        assert len(q) == 1
+        qml.assert_equal(q.queue[0], qml.GQSP(**kwargs))
+
     @pytest.mark.parametrize(
         "template, kwargs",
         [