Repository: PennyLaneAI/pennylane (python)
Issue: Hamiltonians created via addition produce incorrect results with `ExpvalCost(optimize=True)`
Good afternoon, 
I was having problem with vqe when I set the flag: optmize = True . I noticed that there were errors when entering Hamiltonians of the form Z1@Z2 + Z2@Z1. We know it is the same as 2.Z1.Z2 but the optimizer seems to have some kind of problem with this kind of Hamiltonians (vqe with optimize = False works correctly in this situation).
Although by preprocessing the Hamiltonian we avoid the problem, I was writing to take the problem into account.
Greetings!
Pull request description: Pop observables and coeffs that have been grouped already
**Context**

The `ExpvalCost` class uses the grouping module when the `optimize=True` argument is set. Using this option, the terms in the Hamiltonian (that are Pauli words) are grouped based on a relation selected (e.g., qubit-wise commutativity) and then their corresponding coefficients are also grouped accordingly.

That is, grouping involves first the actual grouping of Hamiltonian terms and then grouping the coefficients of the terms.

When there are multiple equivalent terms in the Hamiltonian, with the current logic, the coefficient of the first term that matches a Pauli word will be used when grouping the coefficients. This is so because the grouping module assumes that the terms in the Hamiltonian are unique.

Therefore, equal terms of the Hamiltonian might get the same coefficient incorrectly which will result in incorrect results for the expectation value of the entire Hamiltonian.
```python
import pennylane as qml

a = qml.PauliZ(wires=[0]) @ qml.PauliZ(wires=[4])
b = qml.PauliZ(wires=[4]) @ qml.PauliZ(wires=[0])

coeffs, observables = [0.5, 1], [a, b]
obs_groupings, coeffs_groupings = qml.grouping.group_observables(observables, coeffs)
print(coeffs_groupings)
```
```
[[0.5, 0.5]]   # <---- second coeff should be 1
```
*Note*: the logic of `ExpvalCost` does not hinder such a case from being valid and it doesn't explicitly perform the simplification of the Hamiltonian because that might incur considerable overhead for larger Hamiltonians. It is advisable for users to simplify their Hamiltonians in advance.

**Changes**

Updates the grouping mechanism to remove already processed observables and coefficients.

**Benefits**

The grouping module can function even if Hamiltonian terms are not unique.

**Related issues**

Closes #1389

*Further note*: although it should be fair to assume that the terms are unique in the Hamiltonian, making sure that grouping happens even with terms appearing multiple times can ensure that there is no unexpected behaviour. Users should be encouraged to simplify their Hamiltonians at all times.

Patch:
diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md
index a04d061aa..249848161 100644
--- a/.github/CHANGELOG.md
+++ b/.github/CHANGELOG.md
@@ -503,6 +503,10 @@ random_mat2 = rng.standard_normal(3, requires_grad=False)
 
 <h3>Bug fixes</h3>
 
+* Fixes a bug where multiple identical Hamiltonian terms will produce a
+  different result with ``optimize=True`` using ``ExpvalCost``.
+  [(#1405)](https://github.com/XanaduAI/pennylane/pull/1405)
+
 * Fixes bug where `shots=None` was not reset when changing shots temporarily in a QNode call
   like `circuit(0.1, shots=3)`.
   [(#1392)](https://github.com/XanaduAI/pennylane/pull/1392)
diff --git a/pennylane/grouping/group_observables.py b/pennylane/grouping/group_observables.py
index 28b2b2182..5491721e3 100644
--- a/pennylane/grouping/group_observables.py
+++ b/pennylane/grouping/group_observables.py
@@ -15,15 +15,18 @@
 This module contains the high-level Pauli-word-partitioning functionality used in measurement optimization.
 """
 
-from pennylane.wires import Wires
+from copy import copy
+
+import numpy as np
+
+from pennylane.grouping.graph_colouring import largest_first, recursive_largest_first
 from pennylane.grouping.utils import (
-    observables_to_binary_matrix,
-    binary_to_pauli,
     are_identical_pauli_words,
+    binary_to_pauli,
+    observables_to_binary_matrix,
     qwc_complement_adj_matrix,
 )
-from pennylane.grouping.graph_colouring import largest_first, recursive_largest_first
-import numpy as np
+from pennylane.wires import Wires
 
 GROUPING_TYPES = frozenset(["qwc", "commuting", "anticommuting"])
 GRAPH_COLOURING_METHODS = {"lf": largest_first, "rlf": recursive_largest_first}
@@ -236,11 +239,16 @@ def group_observables(observables, coefficients=None, grouping_type="qwc", metho
 
     partitioned_coeffs = [[0] * len(g) for g in partitioned_paulis]
 
+    observables = copy(observables)
+    coefficients = copy(coefficients)
     for i, partition in enumerate(partitioned_paulis):
         for j, pauli_word in enumerate(partition):
             for observable in observables:
                 if are_identical_pauli_words(pauli_word, observable):
-                    partitioned_coeffs[i][j] = coefficients[observables.index(observable)]
+                    ind = observables.index(observable)
+                    partitioned_coeffs[i][j] = coefficients[ind]
+                    observables.pop(ind)
+                    coefficients.pop(ind)
                     break
 
     return partitioned_paulis, partitioned_coeffs


Test patch:
diff --git a/tests/grouping/test_optimize_measurements.py b/tests/grouping/test_optimize_measurements.py
index 1f0184f82..a8ab748f1 100644
--- a/tests/grouping/test_optimize_measurements.py
+++ b/tests/grouping/test_optimize_measurements.py
@@ -111,6 +111,38 @@ class TestOptimizeMeasurements:
             grouped_coeffs[i] == grouped_coeffs_sol[i] for i in range(len(grouped_coeffs_sol))
         )
 
+    @pytest.mark.parametrize(
+        "obs",
+        [
+            [PauliZ(0), PauliZ(0)],
+            [PauliX(0) @ PauliX(1), PauliX(0) @ PauliX(1)],
+            [PauliX(0) @ PauliX(1), PauliX(1) @ PauliX(0)],
+        ],
+    )
+    def test_optimize_measurements_qwc_term_multiple_times(self, obs):
+        """Tests if coefficients are properly re-structured even if the same
+        terms appear multiple times.
+
+        Although it should be fair to assume that the terms are unique in the
+        Hamiltonian, making sure that grouping happens even with terms appearing
+        multiple times can ensure that there is no unexpected behaviour.
+        """
+        coefficients = [1.43, 4.21]
+
+        diagonalized_groupings_sol = [obs]
+
+        grouped_coeffs_sol = [[1.43, 4.21]]
+
+        grouped_coeffs = optimize_measurements(
+            obs, coefficients, grouping="qwc", colouring_method="rlf"
+        )[2]
+
+        assert len(grouped_coeffs) == len(grouped_coeffs)
+
+        assert all(
+            grouped_coeffs[i] == grouped_coeffs_sol[i] for i in range(len(grouped_coeffs_sol))
+        )
+
     def test_optimize_measurements_not_implemented_catch(self):
         """Tests that NotImplementedError is raised for methods other than ``'qwc'``."""
 
diff --git a/tests/test_vqe.py b/tests/test_vqe.py
index a8fc91906..bd0c7d079 100644
--- a/tests/test_vqe.py
+++ b/tests/test_vqe.py
@@ -990,6 +990,62 @@ class TestVQE:
 
         assert np.allclose(c1, c2)
 
+    @pytest.mark.parametrize("interface", ["tf", "torch", "autograd"])
+    def test_optimize_multiple_terms(self, interface, tf_support, torch_support):
+        """Test that an ExpvalCost with observable optimization gives the same
+        result as another ExpvalCost without observable optimization even when there
+        are non-unique Hamiltonian terms."""
+        if interface == "tf" and not tf_support:
+            pytest.skip("This test requires TensorFlow")
+        if interface == "torch" and not torch_support:
+            pytest.skip("This test requires Torch")
+
+        dev = qml.device("default.qubit", wires=5)
+        obs = [
+            qml.PauliZ(wires=[2]) @ qml.PauliZ(wires=[4]),  # <---- These two terms
+            qml.PauliZ(wires=[4]) @ qml.PauliZ(wires=[2]),  # <---- are equal
+            qml.PauliZ(wires=[1]),
+            qml.PauliZ(wires=[2]),
+            qml.PauliZ(wires=[1]) @ qml.PauliZ(wires=[2]),
+            qml.PauliZ(wires=[2]) @ qml.PauliZ(wires=[0]),
+            qml.PauliZ(wires=[3]) @ qml.PauliZ(wires=[1]),
+            qml.PauliZ(wires=[4]) @ qml.PauliZ(wires=[3]),
+        ]
+
+        coefs = (np.random.rand(len(obs)) - 0.5) * 2
+        hamiltonian = qml.Hamiltonian(coefs, obs)
+
+        cost = qml.ExpvalCost(
+            qml.templates.StronglyEntanglingLayers,
+            hamiltonian,
+            dev,
+            optimize=True,
+            interface=interface,
+            diff_method="parameter-shift",
+        )
+        cost2 = qml.ExpvalCost(
+            qml.templates.StronglyEntanglingLayers,
+            hamiltonian,
+            dev,
+            optimize=False,
+            interface=interface,
+            diff_method="parameter-shift",
+        )
+
+        w = qml.init.strong_ent_layers_uniform(2, 5, seed=1967)
+
+        c1 = cost(w)
+        exec_opt = dev.num_executions
+        dev._num_executions = 0
+
+        c2 = cost2(w)
+        exec_no_opt = dev.num_executions
+
+        assert exec_opt == 1  # Number of groups in the Hamiltonian
+        assert exec_no_opt == 8
+
+        assert np.allclose(c1, c2)
+
     def test_optimize_grad(self):
         """Test that the gradient of ExpvalCost is accessible and correct when using observable
         optimization and the autograd interface."""