Improve approximate set generation in qml.ops.sk_decomposition

doc/releases/changelog-dev.md

@@ -44,6 +44,10 @@
 * An informative error is raised when a `QNode` with `diff_method=None` is differentiated.
   [(#6770)](https://github.com/PennyLaneAI/pennylane/pull/6770)
 
+* `qml.ops.sk_decomposition` has been improved to produce less gates for certain edge cases. This greatly impacts
+  the performance of `qml.clifford_t_decomposition`, which should now give less extraneous `qml.T` gates for them.
+  [(#6855)](https://github.com/PennyLaneAI/pennylane/pull/6855)
+
 * The requested `diff_method` is now validated when program capture is enabled.
   [(#6852)](https://github.com/PennyLaneAI/pennylane/pull/6852)
 
@@ -124,6 +128,7 @@
 
 This release contains contributions from (in alphabetical order):
 
+Utkarsh Azad,
 Yushao Chen,
 Isaac De Vlugt,
 Diksha Dhawan,

pennylane/ops/op_math/decompositions/solovay_kitaev.py

@@ -63,27 +63,28 @@ def _quaternion_transform(matrix):
     )
 
 
-def _contains_SU2(op_mat, ops_vecs, tol=1e-8):
+def _contains_SU2(op_mat, ops_vecs=None, kd_tree=None, tol=1e-8):
     r"""Checks if a given SU(2) matrix is contained in a list of quaternions for a given tolerance.
 
     Args:
         op_mat (TensorLike): SU(2) matrix for the operation to be searched
         op_vecs (list(TensorLike)): List of quaternion for the operations that makes the search space.
+        kd_tree (KDTree): KDTree object built from the list of quaternions. Default is ``None``.
         tol (float): Tolerance for the match to be considered ``True``.
 
     Returns:
-        Tuple(bool, TensorLike): A bool that shows whether an operation similar to the given operations
-        was found, and the quaternion representation of the searched operation.
+        Tuple(bool, TensorLike, int): A bool that shows whether an operation similar to the given operations
+        was found, the quaternion representation of the searched operation and its index in the list.
     """
-    node_points = qml.math.array(ops_vecs)
     gate_points = qml.math.array([_quaternion_transform(op_mat)])
 
-    tree = KDTree(node_points)
-    dist = tree.query(gate_points, workers=-1)[0][0]
+    tree = kd_tree or KDTree(qml.math.array(ops_vecs))
+    dist, indx = tree.query(gate_points, workers=-1)
 
-    return (dist < tol, gate_points[0])
+    return (dist[0] < tol, gate_points[0], indx[0])
 
 
+# pylint: disable=too-many-statements
 @lru_cache()
 def _approximate_set(basis_gates, max_length=10):
     r"""Builds an approximate unitary set required for the `Solovay-Kitaev algorithm <https://arxiv.org/abs/quant-ph/0505030>`_.
@@ -113,6 +114,7 @@ def _approximate_set(basis_gates, max_length=10):
     }
     # Maintains the basis gates
     basis = [_CLIFFORD_T_BASIS[gate.upper()] for gate in basis_gates]
+    t_set = {qml.T(0), qml.adjoint(qml.T(0))}
 
     # Get the SU(2) data for the gates in basis set
     basis_mat, basis_gph = {}, {}
@@ -125,18 +127,23 @@ def _approximate_set(basis_gates, max_length=10):
     gtrie_ids = [[[gate] for gate in basis]]
     gtrie_mat = [list(basis_mat.values())]
     gtrie_gph = [list(basis_gph.values())]
+    gtrie_sum = [[int(gate in t_set) for gate in basis]]
 
     # Maintains the approximate set for gates' names, SU(2)s, global phases and quaternions
     approx_set_ids = list(gtrie_ids[0])
     approx_set_mat = list(gtrie_mat[0])
     approx_set_gph = list(gtrie_gph[0])
+    approx_set_sum = list(gtrie_sum[0])
     approx_set_qat = [_quaternion_transform(mat) for mat in approx_set_mat]
 
     # We will perform a breadth-first search (BFS) style set building for the set
     for depth in range(max_length - 1):
+        kdtree = KDTree(qml.math.array(approx_set_qat))
         # Add the containers for next depth while we explore the current
-        gtrie_id, gtrie_mt, gtrie_gp = [], [], []
-        for node, su2m, gphase in zip(gtrie_ids[depth], gtrie_mat[depth], gtrie_gph[depth]):
+        gtrie_id, gtrie_mt, gtrie_gp, gtrie_sm, gtrie_qt = [], [], [], [], []
+        for node, su2m, gphase, tgsum in zip(
+            gtrie_ids[depth], gtrie_mat[depth], gtrie_gph[depth], gtrie_sum[depth]
+        ):
             # Get the last operation in the current node
             last_op = qml.adjoint(node[-1], lazy=False) if node else None
 
@@ -149,25 +156,59 @@ def _approximate_set(basis_gates, max_length=10):
                 # Extend and check if the node already exists in the approximate set.
                 su2_gp = basis_gph[op] + gphase
                 su2_op = (-1.0) ** bool(su2_gp >= math.pi) * (basis_mat[op] @ su2m)
-                exists, quaternion = _contains_SU2(su2_op, approx_set_qat)
-                if not exists:
+
+                exists, quaternion, index = False, None, -1
+                if gtrie_qt:
+                    exists, quaternion, index = _contains_SU2(su2_op, ops_vecs=gtrie_qt)
+
+                if exists:
+                    index += len(approx_set_qat) - len(gtrie_qt)
+                else:
+                    exists, quaternion, index = _contains_SU2(su2_op, kd_tree=kdtree)
+
+                global_phase = qml.math.mod(su2_gp, math.pi)
+                if not exists or global_phase != approx_set_gph[index]:
                     approx_set_ids.append(node + [op])
                     approx_set_mat.append(su2_op)
+
+                    # Add the quaternion data
                     approx_set_qat.append(quaternion)
+                    gtrie_qt.append(quaternion)
 
                     # Add to the containers for next depth
                     gtrie_id.append(node + [op])
                     gtrie_mt.append(su2_op)
 
                     # Add the global phase data
-                    global_phase = qml.math.mod(su2_gp, math.pi)
                     approx_set_gph.append(global_phase)
                     gtrie_gp.append(global_phase)
 
+                    # Add the T gate sum data
+                    tbool = int(op in t_set)
+                    approx_set_sum.append(tgsum + tbool)
+                    gtrie_sm.append(tgsum + tbool)
+
         # Add to the next depth for next iteration
         gtrie_ids.append(gtrie_id)
         gtrie_mat.append(gtrie_mt)
         gtrie_gph.append(gtrie_gp)
+        gtrie_sum.append(gtrie_sm)
+
+    # Prune the approximate set for equivalent operations with higher T-gate counts
+    if approx_set_qat:
+        tree, tsum = KDTree(approx_set_qat), qml.math.array(approx_set_sum)
+        dists, indxs = tree.query(approx_set_qat, workers=-1, k=10)
+
+        prune_ixs = []
+        for dist, indx in zip(dists, indxs):
+            eq_idx = qml.math.sort(indx[qml.math.where(dist.round(8) == 0.0)])
+            prune_ixs.extend(eq_idx[qml.math.argsort(tsum[eq_idx])][1:])
+
+        for ix in sorted(set(prune_ixs), reverse=True):
+            del approx_set_ids[ix]
+            del approx_set_mat[ix]
+            del approx_set_gph[ix]
+            del approx_set_qat[ix]
 
     return approx_set_ids, approx_set_mat, approx_set_gph, approx_set_qat
 
@@ -340,7 +381,7 @@ def _solovay_kitaev(umat, n, u_n1_ids, u_n1_mat):
     [map_tape], _ = qml.map_wires(new_tape, wire_map={0: op.wires[0]}, queue=True)
 
     # Get phase information based on the decomposition effort
-    phase = approx_set_gph[index] - gate_gph if depth or qml.math.allclose(gate_gph, 0.0) else 0.0
+    phase = approx_set_gph[index] - gate_gph
     global_phase = qml.GlobalPhase(qml.math.array(phase, like=interface))
 
     # Return the gates from the mapped tape and global phase

tests/ops/op_math/test_solovay_kitaev.py

@@ -85,7 +85,7 @@ def test_contains_SU2():
 
     approx_ids, _, _, approx_vec = _approximate_set(("T", "T*", "H"), max_length=3)
 
-    exists, quaternion = _contains_SU2(target, approx_vec)
+    exists, quaternion, _ = _contains_SU2(target, approx_vec)
     assert exists
 
     result = [qml.adjoint(qml.T(0)), qml.adjoint(qml.T(0)), qml.adjoint(qml.T(0))]
@@ -138,19 +138,23 @@ def test_group_commutator_decompose(op):
 
 
 @pytest.mark.parametrize(
-    ("op"),
+    ("op", "max_depth"),
     [
-        qml.RX(math.pi / 42, wires=[1]),
-        qml.RY(math.pi / 7, wires=["a"]),
-        qml.prod(*[qml.RX(1.0, "a"), qml.T("a")]),
-        qml.prod(*[qml.T(0), qml.Hadamard(0)] * 5),
+        (qml.RX(math.pi / 42, wires=[1]), 5),
+        (qml.RY(math.pi / 7, wires=["a"]), 5),
+        (qml.prod(*[qml.RX(1.0, "a"), qml.T("a")]), 5),
+        (qml.prod(*[qml.T(0), qml.Hadamard(0)] * 5), 5),
+        (qml.RZ(-math.pi / 2, wires=[1]), 1),
+        (qml.adjoint(qml.S(wires=["a"])), 1),
+        (qml.PhaseShift(5 * math.pi / 2, wires=[0]), 1),
+        (qml.PhaseShift(-3 * math.pi / 4, wires=["b"]), 1),
     ],
 )
-def test_solovay_kitaev(op):
-    """Test Solovay-Kitaev decomposition method"""
+def test_solovay_kitaev(op, max_depth):
+    """Test Solovay-Kitaev decomposition method with specified max-depth"""
 
     with qml.queuing.AnnotatedQueue() as q:
-        gates = sk_decomposition(op, epsilon=1e-4, max_depth=5, basis_set=("T", "T*", "H"))
+        gates = sk_decomposition(op, epsilon=1e-4, max_depth=max_depth, basis_set=("T", "T*", "H"))
     assert q.queue == gates
 
     matrix_sk = qml.matrix(qml.tape.QuantumScript(gates))

tests/transforms/test_cliffordt_transform.py

@@ -171,7 +171,7 @@ def qfunc():
             for op in tape.operations
         )
 
-    @pytest.mark.parametrize("epsilon", [2e-2, 5e-2, 9e-2])
+    @pytest.mark.parametrize("epsilon", [2e-2, 5e-2, 7e-2])
     @pytest.mark.parametrize("circuit", [circuit_3, circuit_4, circuit_5])
     def test_total_error(self, epsilon, circuit):
         """Ensure that given a certain epsilon, the total operator error is below the threshold."""
@@ -411,7 +411,7 @@ def test_zero_global_phase(self):
 
         [tape], _ = qml.clifford_t_decomposition(tape)
 
-        assert not sum([isinstance(op, qml.GlobalPhase) for op in tape.operations])
+        assert not sum(isinstance(op, qml.GlobalPhase) for op in tape.operations)
 
     def test_raise_with_decomposition_method(self):
         """Test that exception is correctly raise when using incorrect decomposing method"""