Repository: PennyLaneAI/pennylane
Pull request #5978 (2024-07-10 13:51:14)

BugFix - Returning sums of observables on the same wires provides incorrect results 
**Context:**

These are the codes that detected this bug in the first place. I report the current (correct) result and the previous (incorrect) result:

```
import numpy as np
import pennylane as qml

dev = qml.device("default.qubit", wires=2, shots=1000)

@qml.qnode(dev)
def circuit():
    qml.RX(np.pi / 2, 0)
    qml.RX(-np.pi / 2, 1)
    return qml.expval(qml.Y(0) + qml.Y(0)), qml.expval(qml.Y(1))

>>> circuit()
(2.0, 1.0)  # previous result: (2.0,)
```

```
@qml.qnode(dev)
def circuit():
    qml.RX(np.pi / 2, 0)
    qml.RX(-np.pi / 2, 1)
    return qml.expval(qml.Y(1)), qml.expval(qml.Y(0) + qml.Y(0))

>>> circuit()
(1.0, -2.0)  # previous result: (1.0, 2.0)
```

The bug is caused by the current logic in the `qml.devices.qubit.measure_with_samples` function. 
This function internally calls the `_group_measurements` function and selects a different measure function for each group returned by the former.
 
From now on, we assume that `qml.devices.qubit` is the current directory. 

The grouping strategy in the `_group_measurements` function is such that `_measure_sum_with_samples` is the selected measurement function when a sum is involved in the group. Unfortunately, the logic in the grouping process and `_measure_sum_with_samples` cannot correctly handle the case in which there are multiple measurement processes and one involves a sum of operators acting on the same wire. 

I tried several different solutions, noticing that:

1) Many functions in the source code depend on the `measure_with_samples` function. Any 'major change' to this function results in test failures across PL. Re-designing the entire logic of this function is beyond the scope of this PR.

2) Changing just the grouping logic in `_group_measurements` _or_ the workflow in `_measure_sum_with_samples` cannot really solve this issue, since the logic of the entire function is 'entangled' among all its components.

There is a way to save goats and cabbages. 

**Description of the Change:**

In the `_group_measurements` function, we simplify the measurement processes involving a sum. This leaves untouched the cases for which the function was originally designed, but converts the sum of operators acting on the same wires into a `SProd`. Therefore, the logic in `_group_measurements` becomes compatible with the workflow of the `measure_with_sample` function.

That is, `_measure_with_samples_diagonalizing_gates` is the chosen measurement function in place of `_measure_sum_with_samples`. The latter is still the selected measurement function for sums of different operators and/or sums of the same operator acting on the different wires (I suspect these were the cases taken into account by the original author in the first place).

**Benefits:** The function now returns the correct result.

**Possible Drawbacks:** We are simply re-routing an internal possible workflow of the `measure_with_samples` function to make it compatible with the logic intended by the author. Although everything is possible, it is unlikely that any major problem is caused directly by this PR.

**Related GitHub Issues:** #5976 

**Related ShortCut Stories** [sc-68218]


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/5978?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.67%. Comparing base [(`e63362d`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/e63362d0a83ae412436cebcb6ca7d13e5c6f3579?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`40c1e34`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/40c1e34a390f70ada8076d2c5b021414b5456b58?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).

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


```diff
@@            Coverage Diff             @@
##           master    #5978      +/-   ##
==========================================
- Coverage   99.67%   99.67%   -0.01%     
==========================================
  Files         426      426              
  Lines       40921    40625     -296     
==========================================
- Hits        40788    40491     -297     
- Misses        133      134       +1     
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5978?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).

astralcai:  Potentially we can re-use `split_non_commuting` to replace this completely at some point.
PietropaoloFrisoni: > Potentially we can re-use split_non_commuting to replace this completely at some point.

At some point, I agree. It seems to me that this would require changing the grouping strategy _and_ the current logic of `measure_with_sampling`. This is a simple bugfix to return the correct result, so it's is beyond the purpose of the PR 👍 

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index c95127ac1..84103a037 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -45,6 +45,10 @@
 
 <h3>Bug fixes 🐛</h3>
 
+* `qml.devices.qubit.measure_with_samples` now returns the correct result if the provided measurements
+  contain sum of operators acting on the same wire.
+  [(#5978)](https://github.com/PennyLaneAI/pennylane/pull/5978)
+
 * `qml.AmplitudeEmbedding` has better support for features using low precision integer data types.
 [(#5969)](https://github.com/PennyLaneAI/pennylane/pull/5969)
 
@@ -55,6 +59,7 @@ This release contains contributions from (in alphabetical order):
 Ahmed Darwish
 Astral Cai,
 Yushao Chen,
+Pietropaolo Frisoni,
 Christina Lee,
 William Maxwell,
 Vincent Michaud-Rioux,
diff --git a/pennylane/devices/qubit/sampling.py b/pennylane/devices/qubit/sampling.py
index 2a585655c..be2a2ae8b 100644
--- a/pennylane/devices/qubit/sampling.py
+++ b/pennylane/devices/qubit/sampling.py
@@ -25,7 +25,7 @@ from pennylane.measurements import (
     ShadowExpvalMP,
     Shots,
 )
-from pennylane.ops import Hamiltonian, LinearCombination, Sum
+from pennylane.ops import Hamiltonian, LinearCombination, Prod, SProd, Sum
 from pennylane.typing import TensorLike
 
 from .apply_operation import apply_operation
@@ -65,6 +65,8 @@ def _group_measurements(mps: List[Union[SampleMeasurement, ClassicalShadowMP, Sh
     mp_no_obs_indices = []
 
     for i, mp in enumerate(mps):
+        if isinstance(mp.obs, (Sum, SProd, Prod)):
+            mps[i].obs = qml.simplify(mp.obs)
         if isinstance(mp, (ClassicalShadowMP, ShadowExpvalMP)):
             mp_other_obs.append([mp])
             mp_other_obs_indices.append([i])


Test patch:
diff --git a/tests/devices/qubit/test_sampling.py b/tests/devices/qubit/test_sampling.py
index 550691478..cb4938414 100644
--- a/tests/devices/qubit/test_sampling.py
+++ b/tests/devices/qubit/test_sampling.py
@@ -189,6 +189,7 @@ class TestSampleState:
         assert np.allclose(reordered_probs, random_probs, atol=APPROX_ATOL)
 
 
+# pylint: disable=too-many-public-methods
 class TestMeasureSamples:
     """Test that the measure_with_samples function works as expected"""
 
@@ -542,6 +543,44 @@ class TestMeasureSamples:
         [result] = measure_with_samples([mp], state, shots=qml.measurements.Shots(1))
         assert qml.math.allclose(result, 1)  # -1 + 2
 
+    @pytest.mark.parametrize(
+        "state, measurements, expected_results",
+        [
+            [
+                np.array([[0.5, 0.5j], [-0.5j, 0.5]]),
+                [qml.expval(qml.Y(0) + qml.Y(0)), qml.expval(qml.Y(1))],
+                (-2.0, 1.0),
+            ],
+            [
+                np.array([[0.5, -0.5j], [0.5j, 0.5]]),
+                [qml.expval(qml.Y(0) + qml.Y(0)), qml.expval(qml.Y(1))],
+                (2.0, -1.0),
+            ],
+            [
+                np.array([[0.5, 0.5j], [-0.5j, 0.5]]),
+                [qml.expval(qml.Y(1)), qml.expval(qml.Y(0) + qml.Y(0))],
+                (1.0, -2.0),
+            ],
+            [
+                np.array([[0.5, 0.5j], [-0.5j, 0.5]]),
+                [
+                    qml.expval(qml.Y(1) - qml.Y(1)),
+                    qml.expval(2 * (qml.Y(0) + qml.Y(0) - 5 * (qml.Y(0) + qml.Y(0)))),
+                    qml.expval(
+                        (2 * (qml.Y(0) + qml.Y(0)))
+                        @ ((5 * (qml.Y(0) + qml.Y(0)) + 3 * (qml.Y(0) + qml.Y(0))))
+                    ),
+                ],
+                (0.0, 16.0, 64.0),
+            ],
+        ],
+    )
+    def test_sum_same_wires(self, state, measurements, expected_results):
+        """Test that the sum of observables acting on the same wires works as expected."""
+
+        results = measure_with_samples(measurements, state, qml.measurements.Shots(1000))
+        assert qml.math.allclose(results, expected_results)
+
 
 class TestInvalidStateSamples:
     """Tests for state vectors containing nan values or shot vectors with zero shots."""