Adds elementary gate decomposition for single-qubit QubitUnitary

pennylane/ops/qubit.py

@@ -2157,6 +2157,43 @@ def _matrix(cls, *params):
 
         return U
 
+    @staticmethod
+    def decomposition(U, wires):
+        # Decompose arbitrary single-qubit unitaries as the form RZ RY RZ
+        if U.shape[0] == 2:
+            wires = Wires(wires)
+
+            # Remove the global phase if present; cannot just divide by square root of det
+            # because sometimes it is negative.
+            det = U[0, 0] * U[1, 1] - U[0, 1] * U[1, 0]
+            if not qml.math.isclose(det, 1):
+                U = U * (qml.math.exp(-1j * qml.math.angle(det) / 2))
+
+            # Compute the angle of the RY
+            theta = 2 * qml.math.arccos(qml.math.sqrt(U[0, 0] * U[1, 1]))
+
+            # If it's close to 0, the matrix is diagonal and we have just an RZ rotation
+            if qml.math.isclose(theta, 0, atol=1e-7):
+                omega = 2 * qml.math.angle(U[1, 1])
+                return [RZ(omega, wires=wires[0])]
+
+            # Otherwise recover the decomposition as a Rot, which can be further decomposed
+            # if desired. If the top left element is 0, can only use the off-diagonal elements
+            if qml.math.isclose(U[0, 0], 0):
+                phi = 1j * qml.math.log(U[0, 1] / U[1, 0])
+                omega = -phi - 2 * qml.math.angle(U[1, 0])
+            else:
+                omega = 1j * qml.math.log(qml.math.tan(theta / 2) * U[0, 0] / U[1, 0])
+                phi = -omega - 2 * qml.math.angle(U[0, 0])
+
+            return [
+                qml.Rot(
+                    qml.math.real(phi), qml.math.real(theta), qml.math.real(omega), wires=wires[0]
+                )
+            ]
+        else:
+            return NotImplementedError
+
     def adjoint(self):
         return QubitUnitary(qml.math.T(qml.math.conj(self.matrix)), wires=self.wires)
 

tests/ops/test_qubit_ops.py

@@ -1445,6 +1445,29 @@ def test_qubit_unitary_not_matrix_exception(self, U):
         with pytest.raises(ValueError, match="must be a square matrix"):
             qml.QubitUnitary(U, wires=0).matrix
 
+    @pytest.mark.parametrize(
+        "U,expected_gate,expected_params",
+        [  # First set of gates are diagonal and converted to RZ
+            (I, qml.RZ, [0]),
+            (Z, qml.RZ, [np.pi]),
+            (S, qml.RZ, [np.pi / 2]),
+            (T, qml.RZ, [np.pi / 4]),
+            (qml.RZ(0.3, wires=0).matrix, qml.RZ, [0.3]),
+            (qml.RZ(-0.5, wires=0).matrix, qml.RZ, [-0.5]),
+            # Next set of gates are non-diagonal and decomposed as Rots
+            (H, qml.Rot, [np.pi, np.pi / 2, 0]),
+            (X, qml.Rot, [0.0, np.pi, np.pi]),
+            (qml.Rot(0.2, 0.5, -0.3, wires=0).matrix, qml.Rot, [0.2, 0.5, -0.3]),
+            (np.exp(1j * 0.02) * qml.Rot(-1, 2, -3, wires=0).matrix, qml.Rot, [-1, 2, -3]),
+        ],
+    )
+    def test_qubit_unitary_decomposition(self, U, expected_gate, expected_params):
+        decomp = qml.QubitUnitary.decomposition(U, wires=0)
+
+        assert len(decomp) == 1
+        assert isinstance(decomp[0], expected_gate)
+        assert np.allclose(decomp[0].parameters, expected_params)
+
     def test_iswap_eigenval(self):
         """Tests that the ISWAP eigenvalue matches the numpy eigenvalues of the ISWAP matrix"""
         op = qml.ISWAP(wires=[0, 1])