Repository: PennyLaneAI/pennylane
Pull request #2185 (2022-02-09 21:44:07)

Partition the Pauli group into qubit-wise commuting terms
**Context:**

As part of the QCut module, we need to efficiently measure over the full N-qubit Pauli group. Although there are `4^N` terms, we can measure only `3^N` groups.

**Description of the Change:**

Added the `partition_pauli_group` function to the `grouping` module.

**Possible Drawbacks:**

The function returns Pauli words expressed as simple strings. This is sufficient for the QCut module's needs but some users may prefer the Pauli words as PL observables/tensors. Since the current implementation is sufficient for now, I propose to leave returning PL observables/tensors as a follow up addition.

Example:
```pycon
>>> qml.grouping.partition_pauli_group(3)
[['III', 'IIZ', 'IZI', 'IZZ', 'ZII', 'ZIZ', 'ZZI', 'ZZZ'],
 ['IIX', 'IZX', 'ZIX', 'ZZX'],
 ['IIY', 'IZY', 'ZIY', 'ZZY'],
 ['IXI', 'IXZ', 'ZXI', 'ZXZ'],
 ['IXX', 'ZXX'],
 ['IXY', 'ZXY'],
 ['IYI', 'IYZ', 'ZYI', 'ZYZ'],
 ['IYX', 'ZYX'],
 ['IYY', 'ZYY'],
 ['XII', 'XIZ', 'XZI', 'XZZ'],
 ['XIX', 'XZX'],
 ['XIY', 'XZY'],
 ['XXI', 'XXZ'],
 ['XXX'],
 ['XXY'],
 ['XYI', 'XYZ'],
 ['XYX'],
 ['XYY'],
 ['YII', 'YIZ', 'YZI', 'YZZ'],
 ['YIX', 'YZX'],
 ['YIY', 'YZY'],
 ['YXI', 'YXZ'],
 ['YXX'],
 ['YXY'],
 ['YYI', 'YYZ'],
 ['YYX'],
 ['YYY']]
```

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.
trbromley: [sc-14640]
codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2185?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#2185](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2185?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (548f6dd) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/99dd676bf0647e0859614aa996c9da72332b3087?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (99dd676) will **decrease** coverage by `2.79%`.
> The diff coverage is `100.00%`.

[![Impacted file tree graph](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2185/graphs/tree.svg?width=650&height=150&src=pr&token=9S3YzlUgGk&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2185?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)

```diff
@@            Coverage Diff             @@
##           master    #2185      +/-   ##
==========================================
- Coverage   99.20%   96.41%   -2.80%     
==========================================
  Files         231      231              
  Lines       17832    17853      +21     
==========================================
- Hits        17691    17213     -478     
- Misses        141      640     +499     
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2185?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/grouping/\_\_init\_\_.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2185/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2dyb3VwaW5nL19faW5pdF9fLnB5) | `100.00% <100.00%> (ø)` | |
| [pennylane/grouping/pauli.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2185/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2dyb3VwaW5nL3BhdWxpLnB5) | `100.00% <100.00%> (ø)` | |
| [pennylane/interfaces/batch/jax.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2185/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvYmF0Y2gvamF4LnB5) | `0.00% <0.00%> (-100.00%)` | :arrow_down: |
| [pennylane/interfaces/batch/jax\_jit.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2185/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvYmF0Y2gvamF4X2ppdC5weQ==) | `0.00% <0.00%> (-100.00%)` | :arrow_down: |
| [pennylane/interfaces/jax.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2185/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvamF4LnB5) | `0.00% <0.00%> (-97.44%)` | :arrow_down: |
| [pennylane/devices/default\_qubit\_jax.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2185/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RldmljZXMvZGVmYXVsdF9xdWJpdF9qYXgucHk=) | `0.00% <0.00%> (-96.63%)` | :arrow_down: |
| [pennylane/qnn/cost.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2185/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3Fubi9jb3N0LnB5) | `42.85% <0.00%> (-57.15%)` | :arrow_down: |
| [pennylane/transforms/classical\_jacobian.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2185/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvY2xhc3NpY2FsX2phY29iaWFuLnB5) | `77.08% <0.00%> (-22.92%)` | :arrow_down: |
| [pennylane/math/utils.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2185/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21hdGgvdXRpbHMucHk=) | `84.41% <0.00%> (-15.59%)` | :arrow_down: |
| [pennylane/math/multi\_dispatch.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2185/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21hdGgvbXVsdGlfZGlzcGF0Y2gucHk=) | `87.73% <0.00%> (-12.27%)` | :arrow_down: |
| ... and [15 more](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2185/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2185?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2185?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Last update [99dd676...548f6dd](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2185?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?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 c7ee27f1d..f85c4c305 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -35,6 +35,42 @@
 
 <h3>Improvements</h3>
 
+* Added a new `partition_pauli_group` function to the `grouping` module for
+  efficiently measuring the `N`-qubit Pauli group with `3 ** N`
+  qubit-wise commuting terms.
+  [(#2185)](https://github.com/PennyLaneAI/pennylane/pull/2185)
+  
+  ```pycon
+  >>> qml.grouping.partition_pauli_group(3)
+  [['III', 'IIZ', 'IZI', 'IZZ', 'ZII', 'ZIZ', 'ZZI', 'ZZZ'],
+   ['IIX', 'IZX', 'ZIX', 'ZZX'],
+   ['IIY', 'IZY', 'ZIY', 'ZZY'],
+   ['IXI', 'IXZ', 'ZXI', 'ZXZ'],
+   ['IXX', 'ZXX'],
+   ['IXY', 'ZXY'],
+   ['IYI', 'IYZ', 'ZYI', 'ZYZ'],
+   ['IYX', 'ZYX'],
+   ['IYY', 'ZYY'],
+   ['XII', 'XIZ', 'XZI', 'XZZ'],
+   ['XIX', 'XZX'],
+   ['XIY', 'XZY'],
+   ['XXI', 'XXZ'],
+   ['XXX'],
+   ['XXY'],
+   ['XYI', 'XYZ'],
+   ['XYX'],
+   ['XYY'],
+   ['YII', 'YIZ', 'YZI', 'YZZ'],
+   ['YIX', 'YZX'],
+   ['YIY', 'YZY'],
+   ['YXI', 'YXZ'],
+   ['YXX'],
+   ['YXY'],
+   ['YYI', 'YYZ'],
+   ['YYX'],
+   ['YYY']]
+  ```
+
 <h3>Breaking changes</h3>
 
 * The `MultiControlledX` operation now accepts a single `wires` keyword argument for both `control_wires` and `wires`.
diff --git a/pennylane/grouping/__init__.py b/pennylane/grouping/__init__.py
index 36cac90ab..7dadd45c3 100644
--- a/pennylane/grouping/__init__.py
+++ b/pennylane/grouping/__init__.py
@@ -38,4 +38,4 @@ from .utils import (
     observables_to_binary_matrix,
     qwc_complement_adj_matrix,
 )
-from .pauli import pauli_group, pauli_mult, pauli_mult_with_phase
+from .pauli import pauli_group, pauli_mult, pauli_mult_with_phase, partition_pauli_group
diff --git a/pennylane/grouping/pauli.py b/pennylane/grouping/pauli.py
index 06a3b392d..4786d2ce2 100644
--- a/pennylane/grouping/pauli.py
+++ b/pennylane/grouping/pauli.py
@@ -15,16 +15,19 @@
 Functions for constructing the :math:`n`-qubit Pauli group, and performing the
 group operation (multiplication).
 """
+import itertools
+from functools import lru_cache
+from typing import List
 
 import numpy as np
 
 from pennylane import Identity
 from pennylane.grouping.utils import (
+    _wire_map_from_pauli_pair,
+    are_identical_pauli_words,
     binary_to_pauli,
     pauli_to_binary,
     pauli_word_to_string,
-    are_identical_pauli_words,
-    _wire_map_from_pauli_pair,
 )
 
 
@@ -250,3 +253,92 @@ def pauli_mult_with_phase(pauli_1, pauli_2, wire_map=None):
             phase *= -1j
 
     return pauli_product, phase
+
+
+@lru_cache()
+def partition_pauli_group(n_qubits: int) -> List[List[str]]:
+    """Partitions the :math:`n`-qubit Pauli group into qubit-wise commuting terms.
+
+    The :math:`n`-qubit Pauli group is composed of :math:`4^{n}` terms that can be partitioned into
+    :math:`3^{n}` qubit-wise commuting groups.
+
+    Args:
+        n_qubits (int): number of qubits
+
+    Returns:
+        List[List[str]]: A collection of qubit-wise commuting groups containing Pauli words as
+        strings
+
+    **Example**
+
+    >>> qml.grouping.partition_pauli_group(3)
+    [['III', 'IIZ', 'IZI', 'IZZ', 'ZII', 'ZIZ', 'ZZI', 'ZZZ'],
+     ['IIX', 'IZX', 'ZIX', 'ZZX'],
+     ['IIY', 'IZY', 'ZIY', 'ZZY'],
+     ['IXI', 'IXZ', 'ZXI', 'ZXZ'],
+     ['IXX', 'ZXX'],
+     ['IXY', 'ZXY'],
+     ['IYI', 'IYZ', 'ZYI', 'ZYZ'],
+     ['IYX', 'ZYX'],
+     ['IYY', 'ZYY'],
+     ['XII', 'XIZ', 'XZI', 'XZZ'],
+     ['XIX', 'XZX'],
+     ['XIY', 'XZY'],
+     ['XXI', 'XXZ'],
+     ['XXX'],
+     ['XXY'],
+     ['XYI', 'XYZ'],
+     ['XYX'],
+     ['XYY'],
+     ['YII', 'YIZ', 'YZI', 'YZZ'],
+     ['YIX', 'YZX'],
+     ['YIY', 'YZY'],
+     ['YXI', 'YXZ'],
+     ['YXX'],
+     ['YXY'],
+     ['YYI', 'YYZ'],
+     ['YYX'],
+     ['YYY']]
+    """
+    # Cover the case where n_qubits may be passed as a float
+    if isinstance(n_qubits, float):
+        if n_qubits.is_integer():
+            n_qubits = int(n_qubits)
+
+    # If not an int, or a float representing a int, raise an error
+    if not isinstance(n_qubits, int):
+        raise TypeError("Must specify an integer number of qubits.")
+
+    if n_qubits <= 0:
+        raise ValueError("Number of qubits must be at least 1.")
+
+    strings = set()  # tracks all the strings that have already been grouped
+    groups = []
+
+    # We know that I and Z always commute on a given qubit. The following generates all product
+    # sequences of len(n_qubits) over "FXYZ", with F indicating a free slot that can be swapped for
+    # the product over I and Z, and all other terms fixed to the given X/Y/Z. For example, if
+    # ``n_qubits = 3`` our first value for ``string`` will be ``('F', 'F', 'F')``. We then expand
+    # the product of I and Z over the three free slots, giving
+    # ``['III', 'IIZ', 'IZI', 'IZZ', 'ZII', 'ZIZ', 'ZZI', 'ZZZ']``, which is our first group. The
+    # next element of ``string`` will be ``('F', 'F', 'X')`` which we use to generate our second
+    # group ``['IIX', 'IZX', 'ZIX', 'ZZX']``.
+    for string in itertools.product("FXYZ", repeat=n_qubits):
+        if string not in strings:
+            num_free_slots = string.count("F")
+
+            group = []
+            commuting = itertools.product("IZ", repeat=num_free_slots)
+
+            for commuting_string in commuting:
+                commuting_string = list(commuting_string)
+                new_string = tuple(commuting_string.pop(0) if s == "F" else s for s in string)
+
+                if new_string not in strings:  # only add if string has not already been grouped
+                    group.append("".join(new_string))
+                    strings |= {new_string}
+
+            if len(group) > 0:
+                groups.append(group)
+
+    return groups


Test patch:
diff --git a/tests/grouping/test_pauli_group.py b/tests/grouping/test_pauli_group.py
index e75221f15..90f6cfae5 100644
--- a/tests/grouping/test_pauli_group.py
+++ b/tests/grouping/test_pauli_group.py
@@ -15,9 +15,15 @@
 Unit tests for the :mod:`pauli_group`  functions in ``grouping/pauli.py``.
 """
 import pytest
-from pennylane import Identity, PauliX, PauliY, PauliZ
 
-from pennylane.grouping.pauli import pauli_group, pauli_mult, pauli_mult_with_phase
+from pennylane import Identity, PauliX, PauliY, PauliZ
+from pennylane.grouping import is_commuting, string_to_pauli_word
+from pennylane.grouping.pauli import (
+    partition_pauli_group,
+    pauli_group,
+    pauli_mult,
+    pauli_mult_with_phase,
+)
 
 
 class TestPauliGroup:
@@ -196,3 +202,66 @@ class TestPauliGroup:
         """Test that multiplication including phases works as expected."""
         _, obtained_phase = pauli_mult_with_phase(pauli_word_1, pauli_word_2, wire_map=wire_map)
         assert obtained_phase == expected_phase
+
+
+class TestPartitionPauliGroup:
+    """Tests for the partition_pauli_group function"""
+
+    def test_expected_answer(self):
+        """Test if we get the expected answer for 3 qubits"""
+        expected = [
+            ["III", "IIZ", "IZI", "IZZ", "ZII", "ZIZ", "ZZI", "ZZZ"],
+            ["IIX", "IZX", "ZIX", "ZZX"],
+            ["IIY", "IZY", "ZIY", "ZZY"],
+            ["IXI", "IXZ", "ZXI", "ZXZ"],
+            ["IXX", "ZXX"],
+            ["IXY", "ZXY"],
+            ["IYI", "IYZ", "ZYI", "ZYZ"],
+            ["IYX", "ZYX"],
+            ["IYY", "ZYY"],
+            ["XII", "XIZ", "XZI", "XZZ"],
+            ["XIX", "XZX"],
+            ["XIY", "XZY"],
+            ["XXI", "XXZ"],
+            ["XXX"],
+            ["XXY"],
+            ["XYI", "XYZ"],
+            ["XYX"],
+            ["XYY"],
+            ["YII", "YIZ", "YZI", "YZZ"],
+            ["YIX", "YZX"],
+            ["YIY", "YZY"],
+            ["YXI", "YXZ"],
+            ["YXX"],
+            ["YXY"],
+            ["YYI", "YYZ"],
+            ["YYX"],
+            ["YYY"],
+        ]
+        assert expected == partition_pauli_group(3.0)
+
+    @pytest.mark.parametrize("n", range(1, 9))
+    def test_scaling(self, n):
+        """Test if the number of groups is equal to 3**n"""
+        assert len(partition_pauli_group(n)) == 3**n
+
+    @pytest.mark.parametrize("n", range(1, 6))
+    def test_is_qwc(self, n):
+        """Test if each group contains only qubit-wise commuting terms"""
+        for group in partition_pauli_group(n):
+            size = len(group)
+            for i in range(size):
+                for j in range(i, size):
+                    s1 = group[i]
+                    s2 = group[j]
+                    w1 = string_to_pauli_word(s1)
+                    w2 = string_to_pauli_word(s2)
+                    assert is_commuting(w1, w2)
+
+    def test_invalid_input(self):
+        """Test that invalid inputs are handled correctly."""
+        with pytest.raises(TypeError, match="Must specify an integer number"):
+            partition_pauli_group("3")
+
+        with pytest.raises(ValueError, match="Number of qubits must be at least 1"):
+            partition_pauli_group(-1)