Repository: PennyLaneAI/pennylane
Pull request #6855 (2025-01-17 20:44:52)

Improve approximate set generation in `qml.ops.sk_decomposition`
**Context:** The approximate set for the Solovay-Kitaev decomposition is currently prepared by populating a trie-based data structure (TbDS) with a BFS traversal.  This is efficient but misses out on a few future gate sequences because some parent sequences are caught in the equality check condition. 

**Description of the Change:**
1. To allow traversal over some potential missed-out candidate sequences, we relax the equality check to include information regarding the tracked global phases. This unlocks.
2. While the previous step unlocks the previously forbidden parent sequences, it also adds redundant sequences in the set, which is undesirable. So we prune our TbDS by finding the redundant pool sequences and ranking them based on their T-gate count.  
3. To perform the T-gate ranking efficiently, we also build and store this information for each sequence during the traversal itself.
4. To unlock higher `max_length` for broader TbDS, we now separate checking for equivalence with a local and global `KDTree`. The local one is updated as the TbDS is being constructed for a given depth, whereas the global one is updated after the construction for a given depth. 

**Benefits:** This makes the approximate set decently dense than before, unlocking efficient decomposition for some previously missing edge cases. 

**Possible Drawbacks:** This adds a constant overhead incurred from traversing additional parent nodes and post-build pruning. However, this is negligible; as the approximate set once built is cached. 

**Related GitHub Issues:** [sc-82463] #6131 


github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [`doc/releases/changelog-dev.md`](/PennyLaneAI/pennylane/blob/master/doc/releases/changelog-dev.md) with:
* A one-to-two sentence description of the change. You may include a small working example for new features.
* A link back to this PR.
* Your name (or GitHub username) in the contributors section.
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/6855?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 99.47%. Comparing base [(`cb883b1`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/cb883b1dc47828ed1c251d60cb4ea1fdff92a2a0?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`a05ecca`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/a05eccac9968b17181c16cc126f1bdfbc9313de4?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
> Report is 1 commits behind head on master.

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master    #6855      +/-   ##
==========================================
- Coverage   99.60%   99.47%   -0.13%     
==========================================
  Files         476      476              
  Lines       45194    45359     +165     
==========================================
+ Hits        45015    45122     +107     
- Misses        179      237      +58     
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/6855?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).


Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 0dffc758d..d5bde5112 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -111,6 +111,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.
+  [(#6855)](https://github.com/PennyLaneAI/pennylane/pull/6855)
+
 * `qml.gradients.finite_diff_jvp` has been added to compute the jvp of an arbitrary numeric
   function.
   [(#6853)](https://github.com/PennyLaneAI/pennylane/pull/6853)
diff --git a/pennylane/ops/op_math/decompositions/solovay_kitaev.py b/pennylane/ops/op_math/decompositions/solovay_kitaev.py
index 4f55c09e7..7695b0ae9 100644
--- a/pennylane/ops/op_math/decompositions/solovay_kitaev.py
+++ b/pennylane/ops/op_math/decompositions/solovay_kitaev.py
@@ -63,25 +63,60 @@ 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.
-        tol (float): Tolerance for the match to be considered ``True``.
+        op_vecs (list(TensorLike)): list of quaternion for the operations that makes the search space.
+        kd_tree (scipy.spatial.KDTree): kd-tree 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 tuple including `True`/`False` for whether an operation similar to the
+        given operations was found, the quaternion representation of the searched operations, and its index in
+        the `op_vecs` or `kd_tree`.
     """
-    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])
+
+
+def _prune_approximate_set(
+    approx_set_ids, approx_set_mat, approx_set_gph, approx_set_qat, approx_set_sum
+):
+    """Prune the approximate set for equivalent gate sequences with higher T-gate counts.
+
+    Args:
+        approx_set_ids (list[list[~pennylane.operation.Operation]]): list of gate sequences
+        approx_set_mat (list[TensorLike]): list of SU(2) matrices
+        approx_set_gph (list[float]): list of global phases
+        approx_set_qat (list[TensorLike]): list of quaternion representations
+        approx_set_sum (list[int]): list of numbers of the T and Adjoint(T) gates in the sequences
+
+    Returns:
+        Tuple[list[list[~pennylane.operation.Operation]], list[TensorLike], list[float], list[TensorLike]]:
+        A tuple containing the pruned approximate set.
+    """
+    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
 
 
 @lru_cache()
@@ -89,9 +124,8 @@ 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>`_.
 
     Args:
-        basis_gates (list(str)): Basis set to be used for Solovay-Kitaev decomposition build using
-            following terms, ``['X', 'Y', 'Z', 'H', 'T', 'T*', 'S', 'S*']``, where `*` refers
-            to the gate adjoint.
+        basis_gates (tuple[str]): Basis set to be used for Solovay-Kitaev decomposition build using the following
+        terms, ``('X', 'Y', 'Z', 'H', 'T', 'T*', 'S', 'S*')``, where `*` refers to the gate adjoint.
         max_length (int): Maximum expansion length of Clifford+T sequences in the approximation set. Default is `10`
 
     Returns:
@@ -113,6 +147,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 = {}, {}
@@ -121,55 +156,91 @@ def _approximate_set(basis_gates, max_length=10):
         basis_mat.update({gate: su2_mat})
         basis_gph.update({gate: su2_gph})
 
-    # Maintains a trie-like structure for each depth
+    # Maintain a trie-like structure that consists of -
+    # gtrie_<ids / mat / gph / sum> stores <gates / SU2s / global phase / T-gate sum>
+    # each of them are list of lists, where each inner list stores the data at a depth D.
     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
+    # Maintains the approximate set for gates, SU2s, global phases, T-gate sums and quaternions,
+    # where each of the approx_set_<name> is the corresponding flattened verison of gtrie_<name>.
+    # We store the quaternion representations for the SU2 matrices to build a KDTree. This allows us to
+    # query for nearest neighbours of any newly built gate sequence and test for its prior existence.
     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
+    # We will perform a breadth-first search (BFS)-style trie building, starting from basis gates:
+    # We attempt to extend every gate sequence at previous depth (defined by a node) with all
+    # basis gates. We add the extended sequence and its corresponding data to the next depth by
+    # comparing its quaternion representation with the gate sequences already added to the trie.
     for depth in range(max_length - 1):
-        # 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]):
+        # Build a KDTree for the quaternions stored up to the current depth for querying.
+        kdtree = KDTree(qml.math.array(approx_set_qat))
+
+        # Add the local containers for extending the trie to the next depth while traversing current one.
+        ltrie_id, ltrie_mt, ltrie_gp, ltrie_sm, ltrie_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
 
-            # Now attempt extending the current node for each basis gate
+            # Now attempt extending the current node with each gate in the basis set.
             for op in basis:
-                # If basis gate is adjoint of last op in the node, skip.
+                # If the op is the adjoint of last operation in the node, skip.
                 if qml.equal(op, last_op):
                     continue
 
-                # Extend and check if the node already exists in the approximate set.
+                # Extend and check if the node already exists in the approximate set in two steps:
+                # 1. (local check) => within the gate sequences built in the current iteration.
+                # 2. (global check) => within the gate sequences built in the previous iterations.
                 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:
-                    approx_set_ids.append(node + [op])
-                    approx_set_mat.append(su2_op)
-                    approx_set_qat.append(quaternion)
 
-                    # Add to the containers for next depth
-                    gtrie_id.append(node + [op])
-                    gtrie_mt.append(su2_op)
+                exists, quaternion, global_index, local_index = False, None, -1, -1
+                if ltrie_qt:  # local check
+                    exists, quaternion, local_index = _contains_SU2(su2_op, ops_vecs=ltrie_qt)
 
-                    # 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)
+                if exists:  # get the global index from the local index
+                    global_index = local_index + len(approx_set_qat) - len(ltrie_qt)
+                else:  # global check
+                    exists, quaternion, global_index = _contains_SU2(su2_op, kd_tree=kdtree)
 
-        # Add to the next depth for next iteration
-        gtrie_ids.append(gtrie_id)
-        gtrie_mat.append(gtrie_mt)
-        gtrie_gph.append(gtrie_gp)
+                # Add the sequence if it is unique, i.e., new SU(2) representation or global phase.
+                global_phase = qml.math.mod(su2_gp, math.pi)  # Get the global phase in [0, \pi)
+                if not exists or global_phase != approx_set_gph[global_index]:
+                    # Add the data to the approximate set
+                    approx_set_ids.append(node + [op])
+                    approx_set_mat.append(su2_op)
+                    approx_set_gph.append(global_phase)
+                    approx_set_qat.append(quaternion)
 
-    return approx_set_ids, approx_set_mat, approx_set_gph, approx_set_qat
+                    # Add the data to the containers for next depth
+                    ltrie_id.append(node + [op])
+                    ltrie_mt.append(su2_op)
+                    ltrie_gp.append(global_phase)
+                    ltrie_qt.append(quaternion)
+
+                    # Add the T gate sum data
+                    tbool = int(op in t_set)
+                    approx_set_sum.append(tgsum + tbool)
+                    ltrie_sm.append(tgsum + tbool)
+
+        # Add to the next depth for new iteration
+        gtrie_ids.append(ltrie_id)
+        gtrie_mat.append(ltrie_mt)
+        gtrie_gph.append(ltrie_gp)
+        gtrie_sum.append(ltrie_sm)
+
+    # Prune the approximate set for equivalent operations with higher T-gate counts and return
+    return _prune_approximate_set(
+        approx_set_ids, approx_set_mat, approx_set_gph, approx_set_qat, approx_set_sum
+    )
 
 
 def _group_commutator_decompose(matrix, tol=1e-5):
@@ -221,9 +292,9 @@ def sk_decomposition(op, epsilon, *, max_depth=5, basis_set=("T", "T*", "H"), ba
     Keyword Args:
         max_depth (int): The maximum number of approximation passes. A smaller :math:`\epsilon` would generally require
             a greater number of passes. Default is ``5``.
-        basis_set (list[str]): Basis set to be used for the decomposition and building an approximate set internally.
-            It accepts the following gate terms: ``['X', 'Y', 'Z', 'H', 'T', 'T*', 'S', 'S*']``, where ``*`` refers
-            to the gate adjoint. Default value is ``['T', 'T*', 'H']``.
+        basis_set (tuple[str]): Basis set to be used for the decomposition and building an approximate set internally.
+            It accepts the following gate terms: ``('X', 'Y', 'Z', 'H', 'T', 'T*', 'S', 'S*')``, where ``*`` refers
+            to the gate adjoint. Default value is ``('T', 'T*', 'H')``.
         basis_length (int): Maximum expansion length of Clifford+T sequences in the internally-built approximate set.
             Default is ``10``.
 
@@ -267,7 +338,7 @@ def sk_decomposition(op, epsilon, *, max_depth=5, basis_set=("T", "T*", "H"), ba
     with QueuingManager.stop_recording():
         # Build the approximate set with caching
         approx_set_ids, approx_set_mat, approx_set_gph, approx_set_qat = _approximate_set(
-            basis_set, max_length=basis_length
+            tuple(basis_set), max_length=basis_length
         )
 
         # Build the k-d tree with the current approximation set for querying in the base case
@@ -340,7 +411,7 @@ def sk_decomposition(op, epsilon, *, max_depth=5, basis_set=("T", "T*", "H"), ba
     [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


Test patch:
diff --git a/tests/ops/op_math/test_solovay_kitaev.py b/tests/ops/op_math/test_solovay_kitaev.py
index 8e0952b61..fbd292c0f 100644
--- a/tests/ops/op_math/test_solovay_kitaev.py
+++ b/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))
diff --git a/tests/transforms/test_cliffordt_transform.py b/tests/transforms/test_cliffordt_transform.py
index eb6fcb5f3..eb7019556 100644
--- a/tests/transforms/test_cliffordt_transform.py
+++ b/tests/transforms/test_cliffordt_transform.py
@@ -193,7 +193,7 @@ class TestCliffordCompile:
         assert qml.equal(compiled_ops[2], qml.adjoint(qml.T(2)))
         assert qml.equal(compiled_ops[3], qml.T(3))
 
-    @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."""