Qubitization reformat (#6182)

This PR updates `Qubitization` to use `qml.PrepSelPrep`. It also fixes
this [bug](#6175)
Note that we are changing from decompose Qubitization in PrepSelPrep .
Reflection to Reflection.PrepSelPrep

---------

Co-authored-by: Utkarsh <utkarshazad98@gmail.com>

doc/releases/changelog-dev.md

@@ -13,6 +13,9 @@
   `from pennylane.capture.primitives import *`.
   [(#6129)](https://github.com/PennyLaneAI/pennylane/pull/6129)
 
+* Improve `qml.Qubitization` decomposition.
+  [(#6182)](https://github.com/PennyLaneAI/pennylane/pull/6182)
+
 * The `__repr__` methods for `FermiWord` and `FermiSentence` now returns a
   unique representation of the object.
   [(#6167)](https://github.com/PennyLaneAI/pennylane/pull/6167)
@@ -31,6 +34,9 @@
 * Fix `qml.PrepSelPrep` template to work with `torch`:
   [(#6191)](https://github.com/PennyLaneAI/pennylane/pull/6191)
 
+* Now `qml.equal` compares correctly `qml.PrepSelPrep` operators.
+  [(#6182)](https://github.com/PennyLaneAI/pennylane/pull/6182)
+
 * The ``qml.QSVT`` template now orders the ``projector`` wires first and the ``UA`` wires second, which is the expected order of the decomposition.
   [(#6212)](https://github.com/PennyLaneAI/pennylane/pull/6212)
 

pennylane/ops/functions/equal.py

@@ -41,7 +41,7 @@
 )
 from pennylane.pulse.parametrized_evolution import ParametrizedEvolution
 from pennylane.tape import QuantumScript
-from pennylane.templates.subroutines import ControlledSequence
+from pennylane.templates.subroutines import ControlledSequence, PrepSelPrep
 
 OPERANDS_MISMATCH_ERROR_MESSAGE = "op1 and op2 have different operands because "
 
@@ -807,3 +807,17 @@ def _equal_hilbert_schmidt(
         return False
 
     return True
+
+
+@_equal_dispatch.register
+def _equal_prep_sel_prep(
+    op1: PrepSelPrep, op2: PrepSelPrep, **kwargs
+):  # pylint: disable=unused-argument
+    """Determine whether two PrepSelPrep are equal"""
+    if op1.control != op2.control:
+        return f"op1 and op2 have different control wires. Got {op1.control} and {op2.control}."
+    if op1.wires != op2.wires:
+        return f"op1 and op2 have different wires. Got {op1.wires} and {op2.wires}."
+    if not qml.equal(op1.lcu, op2.lcu):
+        return f"op1 and op2 have different lcu. Got {op1.lcu} and {op2.lcu}"
+    return True

pennylane/templates/subroutines/qubitization.py

@@ -18,32 +18,10 @@
 import copy
 
 import pennylane as qml
-from pennylane import numpy as np
 from pennylane.operation import Operation
 from pennylane.wires import Wires
 
 
-def _positive_coeffs_hamiltonian(hamiltonian):
-    """Transforms a Hamiltonian to ensure that the coefficients are positive.
-
-    Args:
-        hamiltonian (Union[.Hamiltonian, .Sum, .Prod, .SProd, .LinearCombination]): The Hamiltonian written as a linear combination of unitaries.
-
-    Returns:
-        list(float), list(.Operation): The coefficients and unitaries of the transformed Hamiltonian.
-    """
-
-    new_unitaries = []
-
-    coeffs, ops = hamiltonian.terms()
-
-    for op, coeff in zip(ops, coeffs):
-        angle = np.pi * (0.5 * (1 - qml.math.sign(coeff)))
-        new_unitaries.append(op @ qml.GlobalPhase(angle, wires=op.wires))
-
-    return qml.math.abs(coeffs), new_unitaries
-
-
 class Qubitization(Operation):
     r"""Applies the `Qubitization <https://arxiv.org/abs/2204.11890>`__ operator.
 
@@ -163,31 +141,16 @@ def compute_decomposition(*_, **kwargs):  # pylint: disable=arguments-differ
 
         **Example:**
 
-        >>> print(qml.Qubitization.compute_decomposition(hamiltonian = 0.1 * qml.Z(0), control = 1))
-        [AmplitudeEmbedding(array([1., 0.]), wires=[1]), Select(ops=(Z(0),), control=Wires([1])), Adjoint(AmplitudeEmbedding(array([1., 0.]), wires=[1])), Reflection(, wires=[0])]
-
+        >>> print(qml.Qubitization.compute_decomposition(hamiltonian = 0.1 * qml.Z(0), control = 1)
+        [Reflection(3.141592653589793, wires=[1]), PrepSelPrep(coeffs=(0.1,), ops=(Z(0),), control=Wires([1]))]
         """
 
         hamiltonian = kwargs["hamiltonian"]
         control = kwargs["control"]
 
-        coeffs, unitaries = _positive_coeffs_hamiltonian(hamiltonian)
-
         decomp_ops = []
 
-        decomp_ops.append(
-            qml.AmplitudeEmbedding(qml.math.sqrt(coeffs), normalize=True, pad_with=0, wires=control)
-        )
-
-        decomp_ops.append(qml.Select(unitaries, control=control))
-        decomp_ops.append(
-            qml.adjoint(
-                qml.AmplitudeEmbedding(
-                    qml.math.sqrt(coeffs), normalize=True, pad_with=0, wires=control
-                )
-            )
-        )
-
         decomp_ops.append(qml.Reflection(qml.Identity(control)))
+        decomp_ops.append(qml.PrepSelPrep(hamiltonian, control=control))
 
         return decomp_ops

tests/ops/functions/test_equal.py

@@ -2801,3 +2801,29 @@ def test_ops_with_abstract_parameters_not_equal():
     assert not jax.jit(qml.equal)(qml.RX(0.1, 0), qml.RX(0.1, 0))
     with pytest.raises(AssertionError, match="Data contains a tracer"):
         jax.jit(assert_equal)(qml.RX(0.1, 0), qml.RX(0.1, 0))
+
+
+@pytest.mark.parametrize(
+    "op, other_op",
+    [
+        (
+            qml.PrepSelPrep(qml.dot([1.0, 2.0], [qml.Z(0), qml.X(0)]), control=1),
+            qml.PrepSelPrep(qml.dot([1.0, 2.0], [qml.Z(0), qml.X(0)]), control=2),
+        ),
+        (
+            qml.PrepSelPrep(qml.dot([1.0, 2.0], [qml.Z(2), qml.X(2)]), control=1),
+            qml.PrepSelPrep(qml.dot([1.0, 2.0], [qml.Z(0), qml.X(0)]), control=1),
+        ),
+        (
+            qml.PrepSelPrep(qml.dot([1.0, -2.0], [qml.Z(0), qml.X(0)]), control=1),
+            qml.PrepSelPrep(qml.dot([1.0, 2.0], [qml.Z(0), qml.X(0)]), control=1),
+        ),
+        (
+            qml.PrepSelPrep(qml.dot([1.0, 2.0], [qml.Z(0), qml.X(0)]), control=1),
+            qml.PrepSelPrep(qml.dot([1.0, 2.0], [qml.Y(0), qml.X(0)]), control=1),
+        ),
+    ],
+)
+def test_not_equal_prep_sel_prep(op, other_op):
+    """Test that two PrepSelPrep operators with different Hamiltonian are not equal."""
+    assert not qml.equal(op, other_op)

tests/templates/test_subroutines/test_qubitization.py

@@ -21,53 +21,6 @@
 
 import pennylane as qml
 from pennylane import numpy as np
-from pennylane.templates.subroutines.qubitization import _positive_coeffs_hamiltonian
-
-
-@pytest.mark.parametrize(
-    "hamiltonian, expected_unitaries",
-    (
-        (
-            qml.ops.LinearCombination(
-                np.array([1, -1, 2]), [qml.PauliX(0), qml.PauliY(0), qml.PauliZ(0)]
-            ),
-            [
-                qml.PauliX(0) @ qml.GlobalPhase(np.array([0.0]), wires=0),
-                qml.PauliY(0) @ qml.GlobalPhase(np.array(np.pi), wires=0),
-                qml.PauliZ(0) @ qml.GlobalPhase(np.array([0.0]), wires=0),
-            ],
-        ),
-        (
-            qml.ops.LinearCombination(
-                np.array([1.0, 1.0, 2.0]), [qml.PauliX(0), qml.PauliY(0), qml.PauliZ(0)]
-            ),
-            [
-                qml.PauliX(0) @ qml.GlobalPhase(np.array([0.0]), wires=0),
-                qml.PauliY(0) @ qml.GlobalPhase(np.array([0.0]), wires=0),
-                qml.PauliZ(0) @ qml.GlobalPhase(np.array([0.0]), wires=0),
-            ],
-        ),
-        (
-            qml.ops.LinearCombination(
-                np.array([-0.2, -0.6, 2.1]), [qml.PauliX(0), qml.PauliY(0), qml.PauliZ(0)]
-            ),
-            [
-                qml.PauliX(0) @ qml.GlobalPhase(np.array(np.pi), wires=0),
-                qml.PauliY(0) @ qml.GlobalPhase(np.array(np.pi), wires=0),
-                qml.PauliZ(0) @ qml.GlobalPhase(np.array(0), wires=0),
-            ],
-        ),
-    ),
-)
-def test_positive_coeffs_hamiltonian(hamiltonian, expected_unitaries):
-    """Tests that the function _positive_coeffs_hamiltonian correctly transforms the Hamiltonian"""
-
-    new_coeffs, new_unitaries = _positive_coeffs_hamiltonian(hamiltonian)
-
-    assert np.allclose(new_coeffs, np.abs(hamiltonian.terms()[0]))
-
-    for i, unitary in enumerate(new_unitaries):
-        qml.assert_equal(expected_unitaries[i], unitary)
 
 
 @pytest.mark.parametrize(
@@ -138,43 +91,23 @@ def test_legacy_new_opmath():
         (
             qml.ops.LinearCombination(np.array([1.0, 1.0]), [qml.PauliX(0), qml.PauliZ(0)]),
             [
-                qml.AmplitudeEmbedding(
-                    np.array([1.0, 1.0]) / np.sqrt(2), wires=[1], pad_with=0, normalize=True
-                ),
-                qml.Select(
-                    ops=(
-                        qml.PauliX(0) @ qml.GlobalPhase(np.array(0.0), wires=0),
-                        qml.PauliZ(0) @ qml.GlobalPhase(np.array(0.0), wires=0),
-                    ),
+                qml.Reflection(qml.I([1]), 3.141592653589793),
+                qml.PrepSelPrep(
+                    qml.ops.LinearCombination(np.array([1.0, 1.0]), [qml.PauliX(0), qml.PauliZ(0)]),
                     control=[1],
                 ),
-                qml.adjoint(
-                    qml.AmplitudeEmbedding(
-                        np.array([1.0, 1.0]) / np.sqrt(2), wires=[1], pad_with=0, normalize=True
-                    )
-                ),
-                qml.Reflection(qml.Identity(wires=[1])),
             ],
         ),
         (
             qml.ops.LinearCombination(np.array([-1.0, 1.0]), [qml.PauliX(0), qml.PauliZ(0)]),
             [
-                qml.AmplitudeEmbedding(
-                    np.array([1.0, 1.0]) / np.sqrt(2), wires=[1], pad_with=0, normalize=True
-                ),
-                qml.Select(
-                    ops=(
-                        qml.PauliX(0) @ qml.GlobalPhase(np.array(np.pi), wires=0),
-                        qml.PauliZ(0) @ qml.GlobalPhase(np.array(0.0), wires=0),
+                qml.Reflection(qml.I(1), 3.141592653589793),
+                qml.PrepSelPrep(
+                    qml.ops.LinearCombination(
+                        np.array([-1.0, 1.0]), [qml.PauliX(0), qml.PauliZ(0)]
                     ),
                     control=[1],
                 ),
-                qml.adjoint(
-                    qml.AmplitudeEmbedding(
-                        np.array([1.0, 1.0]) / np.sqrt(2), wires=[1], pad_with=0, normalize=True
-                    )
-                ),
-                qml.Reflection(qml.Identity(wires=[1])),
             ],
         ),
     ),
@@ -183,7 +116,6 @@ def test_decomposition(hamiltonian, expected_decomposition):
     """Tests that the Qubitization template is correctly decomposed."""
 
     decomposition = qml.Qubitization.compute_decomposition(hamiltonian=hamiltonian, control=[1])
-
     for i, op in enumerate(decomposition):
         qml.assert_equal(op, expected_decomposition[i])