Improve implementation of group_observables (#6043)

**Context:**

`qml.pauli.group_observables` has several areas of improvement that make
the calculation slow. This PR attempts to address some of them.

**Description of the Change:**

1. Solve the Graph colouring algorithm using `Rustworkx` instead of
custom implementation. This adds 2 additional methods to solve the
Minimum Clique cover problem: DSATUR (`'dsatur'`) and IndependentSet
(`'gis'`)
2. Use the symplectic representation of the Pauli observables to
construct the adjacency matrix of the complement graph.
3. Introduce a new function to obtain the indices of the partitions
directly: `qml.pauli.compute_partition_indices`

**Benefits:**

Improves of orders of magnitude compared to the current implementation
of the `qwc` grouping type strategy. The current implementation is based
on an $\mathcal{O}(m^2)$ computation, with $m$ the number of observables
to be grouped. Instead, taking advantage of the symplectic inner
product, and its relation with the commutator and anti commutator of
Pauli observables [1], we can obtain a significant improvement.


![image](https://github.com/user-attachments/assets/3f0ec70b-2082-44a3-acf9-9cdf4238da08)

In addition, by introducing `qml.pauli.compute_partition_indices`, we
now avoid the (rather often) inefficient computation of partitions of
observables, just to then calculate the indices from them - see
https://github.com/PennyLaneAI/pennylane/blob/2892a9a0aa6797912ef4a03f3c4759eaff01d8ef/pennylane/ops/qubit/hamiltonian.py#L37

Instead, we obtain the indices of the partitions directly from the graph
colouring algorithm. Putting all together, the improvement in
performance when calculating the partitions of large Hamiltonians is
evident.


![image](https://github.com/user-attachments/assets/7bd047a9-a628-46c8-9d3e-58818fd7acc0)


This graph was obtained by running 
``` python
%timeit hamiltonian.compute_grouping(grouping_type="qwc", method='lf')
```

**Notes:**

Profiling the implementation from `linear_combination` demonstrates that
the post processing to obtain the indices remains a significant burden
in the calculations, and therefore should be addressed for the remaining
functions - see reference code above.


![image](https://github.com/user-attachments/assets/09cf4ca6-601a-4e91-8865-8bf727dd3f1d)

[1] Andrew Jena (2019). Partitioning Pauli Operators: in Theory and in
Practice. UWSpace. http://hdl.handle.net/10012/15017

[sc-64594]

## Performance comparisons between colouring algorithms

Comparison for `LinearCombination.compute_grouping()`


![image](https://github.com/user-attachments/assets/042d040a-f75b-4642-b6e4-0e903c5d58c5)


![image](https://github.com/user-attachments/assets/a89197ee-9c3d-424d-9bba-36829ee33388)

Comparison for `colour_paui_graph` (only between `Rustworkx` algorithms)


![image](https://github.com/user-attachments/assets/52ae4d0f-acaa-4d45-aab5-d3c40c37863a)

---------

Co-authored-by: Ali Asadi <[email protected]>

Author: EmilianoG-byte

doc/releases/changelog-dev.md

@@ -52,6 +52,14 @@
 * A new method `to_mat` has been added to the `FermiWord` and `FermiSentence` classes, which allows
   computing the matrix representation of these Fermi operators.
   [(#5920)](https://github.com/PennyLaneAI/pennylane/pull/5920)
+  
+* `qml.pauli.group_observables` now uses `Rustworkx` colouring algorithms to solve the Minimum Clique Cover problem.
+  This adds two new options for the `method` argument: `dsatur` and `gis`. In addition, the creation of the adjancecy matrix 
+  now takes advantage of the symplectic representation of the Pauli observables. An additional function `qml.pauli.compute_partition_indices` 
+  is added to calculate the indices from the partitioned observables more efficiently. `qml.pauli.grouping.PauliGroupingStrategy.idx_partitions_from_graph` 
+  can be used to compute partitions of custom indices. These changes improve the wall time of `qml.LinearCombination.compute_grouping` 
+  and the `grouping_type='qwc'` by orders of magnitude. 
+  [(#6043)](https://github.com/PennyLaneAI/pennylane/pull/6043)
 
 <h4>Improvements to operators</h4>
 

doc/requirements.txt

@@ -31,7 +31,7 @@ torch==1.9.0; sys_platform == "darwin" and python_version < "3.10"
 torch==1.13.1; sys_platform != "darwin" and python_version == "3.10"
 torch==1.13.1; sys_platform == "darwin" and python_version == "3.10"
 jinja2==3.0.3
-rustworkx==0.12.1
+rustworkx>=0.14.0
 networkx==2.6
 requests~=2.28.1
 # we do not pin the sphinx theme, to allow the

pennylane/devices/qubit/sampling.py

@@ -63,7 +63,6 @@ def _group_measurements(mps: list[Union[SampleMeasurement, ClassicalShadowMP, Sh
     # measurements with no observables
     mp_no_obs = []
     mp_no_obs_indices = []
-
     for i, mp in enumerate(mps):
         if isinstance(mp.obs, (Sum, SProd, Prod)):
             mps[i].obs = qml.simplify(mp.obs)
@@ -78,13 +77,11 @@ def _group_measurements(mps: list[Union[SampleMeasurement, ClassicalShadowMP, Sh
         else:
             mp_other_obs.append([mp])
             mp_other_obs_indices.append([i])
-
     if mp_pauli_obs:
         i_to_pauli_mp = dict(mp_pauli_obs)
         _, group_indices = qml.pauli.group_observables(
             [mp.obs for mp in i_to_pauli_mp.values()], list(i_to_pauli_mp.keys())
         )
-
         mp_pauli_groups = []
         for indices in group_indices:
             mp_group = [i_to_pauli_mp[i] for i in indices]
@@ -94,7 +91,6 @@ def _group_measurements(mps: list[Union[SampleMeasurement, ClassicalShadowMP, Sh
 
     mp_no_obs_indices = [mp_no_obs_indices] if mp_no_obs else []
     mp_no_obs = [mp_no_obs] if mp_no_obs else []
-
     all_mp_groups = mp_pauli_groups + mp_no_obs + mp_other_obs
     all_indices = group_indices + mp_no_obs_indices + mp_other_obs_indices
 
@@ -240,7 +236,6 @@ def measure_with_samples(
     mps = measurements[0 : -len(mid_measurements)] if mid_measurements else measurements
 
     groups, indices = _group_measurements(mps)
-
     all_res = []
     for group in groups:
         if isinstance(group[0], ExpectationMP) and isinstance(

pennylane/ops/functions/dot.py

@@ -49,7 +49,7 @@ def dot(
         grouping_type (str): The type of binary relation between Pauli words used to compute
             the grouping. Can be ``'qwc'``, ``'commuting'``, or ``'anticommuting'``. Note that if
             ``pauli=True``, the grouping will be ignored.
-        method (str): The graph coloring heuristic to use in solving minimum clique cover for
+        method (str): The graph colouring heuristic to use in solving minimum clique cover for
             grouping, which can be ``'lf'`` (Largest First) or ``'rlf'`` (Recursive Largest
             First). This keyword argument is ignored if ``grouping_type`` is ``None``.
 

pennylane/ops/op_math/linear_combination.py

@@ -38,16 +38,20 @@ class LinearCombination(Sum):
         coeffs (tensor_like): coefficients of the ``LinearCombination`` expression
         observables (Iterable[Observable]): observables in the ``LinearCombination`` expression, of same length as ``coeffs``
         simplify (bool): Specifies whether the ``LinearCombination`` is simplified upon initialization
-                         (like-terms are combined). The default value is `False`. Note that ``coeffs`` cannot
-                         be differentiated when using the ``'torch'`` interface and ``simplify=True``. Use of this argument is deprecated.
+                        (like-terms are combined). The default value is `False`. Note that ``coeffs`` cannot
+                        be differentiated when using the ``'torch'`` interface and ``simplify=True``. Use of this argument is deprecated.
         grouping_type (str): If not ``None``, compute and store information on how to group commuting
             observables upon initialization. This information may be accessed when a :class:`~.QNode` containing this
             ``LinearCombination`` is executed on devices. The string refers to the type of binary relation between Pauli words.
             Can be ``'qwc'`` (qubit-wise commuting), ``'commuting'``, or ``'anticommuting'``.
-        method (str): The graph coloring heuristic to use in solving minimum clique cover for grouping, which
-            can be ``'lf'`` (Largest First) or ``'rlf'`` (Recursive Largest First). Ignored if ``grouping_type=None``.
+        method (str): The graph colouring heuristic to use in solving minimum clique cover for grouping, which
+            can be ``'lf'`` (Largest First), ``'rlf'`` (Recursive Largest First), ``'dsatur'`` (Degree of Saturation), or ``'gis'`` (IndependentSet).
+            Defaults to ``'lf'``. Ignored if ``grouping_type=None``.
         id (str): name to be assigned to this ``LinearCombination`` instance
 
+    .. seealso:: `rustworkx.ColoringStrategy <https://www.rustworkx.org/apiref/rustworkx.ColoringStrategy.html#coloringstrategy>`_
+        for more information on the ``('lf', 'dsatur', 'gis')`` strategies.
+
     .. warning::
         The ``simplify`` argument is deprecated and will be removed in a future release.
         Instead, you can call ``qml.simplify`` on the constructed operator.
@@ -122,7 +126,7 @@ def __init__(
         observables: list[Operator],
         simplify=False,
         grouping_type=None,
-        method="rlf",
+        method="lf",
         _grouping_indices=None,
         _pauli_rep=None,
         id=None,
@@ -229,7 +233,7 @@ def terms(self):
         """
         return self.coeffs, self.ops
 
-    def compute_grouping(self, grouping_type="qwc", method="rlf"):
+    def compute_grouping(self, grouping_type="qwc", method="lf"):
         """
         Compute groups of operators and coefficients corresponding to commuting
         observables of this ``LinearCombination``.
@@ -242,9 +246,10 @@ def compute_grouping(self, grouping_type="qwc", method="rlf"):
         Args:
             grouping_type (str): The type of binary relation between Pauli words used to compute
                 the grouping. Can be ``'qwc'``, ``'commuting'``, or ``'anticommuting'``.
-            method (str): The graph coloring heuristic to use in solving minimum clique cover for
+                Defaults to ``'qwc'``.
+            method (str): The graph colouring heuristic to use in solving minimum clique cover for
                 grouping, which can be ``'lf'`` (Largest First) or ``'rlf'`` (Recursive Largest
-                First).
+                First). Defaults to ``'lf'``.
 
         **Example**
 
@@ -271,27 +276,9 @@ def compute_grouping(self, grouping_type="qwc", method="rlf"):
 
         _, ops = self.terms()
 
-        with qml.QueuingManager.stop_recording():
-            op_groups = qml.pauli.group_observables(ops, grouping_type=grouping_type, method=method)
-
-        ops = copy(ops)
-
-        indices = []
-        available_indices = list(range(len(ops)))
-        for partition in op_groups:  # pylint:disable=too-many-nested-blocks
-            indices_this_group = []
-            for pauli_word in partition:
-                # find index of this pauli word in remaining original observables,
-                for ind, observable in enumerate(ops):
-                    if qml.pauli.are_identical_pauli_words(pauli_word, observable):
-                        indices_this_group.append(available_indices[ind])
-                        # delete this observable and its index, so it cannot be found again
-                        ops.pop(ind)
-                        available_indices.pop(ind)
-                        break
-            indices.append(tuple(indices_this_group))
-
-        self._grouping_indices = tuple(indices)
+        self._grouping_indices = qml.pauli.compute_partition_indices(
+            ops, grouping_type=grouping_type, method=method
+        )
 
     @property
     def wires(self):

pennylane/ops/op_math/sum.py

@@ -43,7 +43,7 @@ def sum(*summands, grouping_type=None, method="rlf", id=None, lazy=True):
             of the operators is already a sum operator, its operands (summands) will be used instead.
         grouping_type (str): The type of binary relation between Pauli words used to compute
             the grouping. Can be ``'qwc'``, ``'commuting'``, or ``'anticommuting'``.
-        method (str): The graph coloring heuristic to use in solving minimum clique cover for
+        method (str): The graph colouring heuristic to use in solving minimum clique cover for
             grouping, which can be ``'lf'`` (Largest First) or ``'rlf'`` (Recursive Largest
             First). This keyword argument is ignored if ``grouping_type`` is ``None``.
 
@@ -129,7 +129,7 @@ class Sum(CompositeOp):
     Keyword Args:
         grouping_type (str): The type of binary relation between Pauli words used to compute
             the grouping. Can be ``'qwc'``, ``'commuting'``, or ``'anticommuting'``.
-        method (str): The graph coloring heuristic to use in solving minimum clique cover for
+        method (str): The graph colouring heuristic to use in solving minimum clique cover for
             grouping, which can be ``'lf'`` (Largest First) or ``'rlf'`` (Recursive Largest
             First). This keyword argument is ignored if ``grouping_type`` is ``None``.
         id (str or None): id for the sum operator. Default is None.
@@ -480,7 +480,7 @@ def compute_grouping(self, grouping_type="qwc", method="rlf"):
         Args:
             grouping_type (str): The type of binary relation between Pauli words used to compute
                 the grouping. Can be ``'qwc'``, ``'commuting'``, or ``'anticommuting'``.
-            method (str): The graph coloring heuristic to use in solving minimum clique cover for
+            method (str): The graph colouring heuristic to use in solving minimum clique cover for
                 grouping, which can be ``'lf'`` (Largest First) or ``'rlf'`` (Recursive Largest
                 First).
 

pennylane/ops/qubit/hamiltonian.py

@@ -87,7 +87,7 @@ class Hamiltonian(Observable):
             observables upon initialization. This information may be accessed when QNodes containing this
             Hamiltonian are executed on devices. The string refers to the type of binary relation between Pauli words.
             Can be ``'qwc'`` (qubit-wise commuting), ``'commuting'``, or ``'anticommuting'``.
-        method (str): The graph coloring heuristic to use in solving minimum clique cover for grouping, which
+        method (str): The graph colouring heuristic to use in solving minimum clique cover for grouping, which
             can be ``'lf'`` (Largest First) or ``'rlf'`` (Recursive Largest First). Ignored if ``grouping_type=None``.
         id (str): name to be assigned to this Hamiltonian instance
 
@@ -476,7 +476,7 @@ def compute_grouping(
         Args:
             grouping_type (str): The type of binary relation between Pauli words used to compute the grouping.
                 Can be ``'qwc'``, ``'commuting'``, or ``'anticommuting'``.
-            method (str): The graph coloring heuristic to use in solving minimum clique cover for grouping, which
+            method (str): The graph colouring heuristic to use in solving minimum clique cover for grouping, which
                 can be ``'lf'`` (Largest First) or ``'rlf'`` (Recursive Largest First).
         """
 

pennylane/pauli/__init__.py

@@ -48,6 +48,7 @@
     PauliGroupingStrategy,
     optimize_measurements,
     graph_colouring,
+    compute_partition_indices,
 )
 
 from .dla import PauliVSpace, lie_closure, structure_constants, center

pennylane/pauli/grouping/__init__.py

@@ -17,5 +17,5 @@
 """
 
 from . import graph_colouring
-from .group_observables import PauliGroupingStrategy, group_observables
+from .group_observables import PauliGroupingStrategy, group_observables, compute_partition_indices
 from .optimize_measurements import optimize_measurements