From 76f2f21e91712c6bcf6fa2f72a4857113de4f0a5 Mon Sep 17 00:00:00 2001
From: Wong Zi Cheng <70616433+chmwzc@users.noreply.github.com>
Date: Mon, 16 Dec 2024 07:01:35 +0000
Subject: [PATCH] Small bug fix to expectation_from_samples

Added new test for the bug too
---
 src/qibo/hamiltonians/hamiltonians.py |  4 ++--
 tests/test_hamiltonians_symbolic.py   | 22 +++++++++++++---------
 2 files changed, 15 insertions(+), 11 deletions(-)

diff --git a/src/qibo/hamiltonians/hamiltonians.py b/src/qibo/hamiltonians/hamiltonians.py
index 0016ddb05b..498e73c96d 100644
--- a/src/qibo/hamiltonians/hamiltonians.py
+++ b/src/qibo/hamiltonians/hamiltonians.py
@@ -635,7 +635,7 @@ def expectation_from_samples(self, freq: dict, qubit_map: dict = None) -> float:
 
         Args:
             freq (dict): input frequencies of the samples.
-            qubit_map (dict): qubit map.
+            qubit_map (list): qubit map.
 
         Returns:
             (float): the calculated expectation value.
@@ -661,7 +661,7 @@ def expectation_from_samples(self, freq: dict, qubit_map: dict = None) -> float:
             expvals.extend(
                 [
                     term.coefficient.real
-                    * (-1) ** [state[qubit_map[q]] for q in qubits].count("1")
+                    * (-1) ** [state[qubit_map.index(q)] for q in qubits].count("1")
                     for state in keys
                 ]
             )
diff --git a/tests/test_hamiltonians_symbolic.py b/tests/test_hamiltonians_symbolic.py
index 0a0f225728..fbc998e061 100644
--- a/tests/test_hamiltonians_symbolic.py
+++ b/tests/test_hamiltonians_symbolic.py
@@ -286,25 +286,29 @@ def test_symbolic_hamiltonian_state_expectation_different_nqubits(
             local_ev = local_ham.expectation(state)
 
 
-def test_hamiltonian_expectation_from_samples(backend):
+@pytest.mark.parametrize(
+    "observable,qubit_map",
+    [
+        (2 * Z(0) * Z(3) + Z(0) * Z(2), [0, 1, 2, 3]),
+        (Z(1) + Z(3), [0, 1, 3]),
+    ]
+)
+def test_hamiltonian_expectation_from_samples(backend, observable, qubit_map):
     """Test Hamiltonian expectation value calculation."""
     backend.set_seed(0)
-    obs0 = 2 * Z(0) * Z(1) + Z(0) * Z(2)
-    obs1 = 2 * Z(0) * Z(1) + Z(0) * Z(2) * I(3)
-    h0 = SymbolicHamiltonian(obs0, backend=backend)
-    h1 = SymbolicHamiltonian(obs1, backend=backend)
+    hamiltonian = SymbolicHamiltonian(observable, backend=backend)
     c = Circuit(4)
     c.add(gates.RX(0, np.random.rand()))
     c.add(gates.RX(1, np.random.rand()))
     c.add(gates.RX(2, np.random.rand()))
     c.add(gates.RX(3, np.random.rand()))
-    c.add(gates.M(0, 1, 2, 3))
+    c.add(gates.M(*qubit_map)
     nshots = 10**5
     result = backend.execute_circuit(c, nshots=nshots)
     freq = result.frequencies(binary=True)
-    ev0 = h0.expectation_from_samples(freq, qubit_map=None)
-    ev1 = h1.expectation(result.state())
-    backend.assert_allclose(ev0, ev1, atol=20 / np.sqrt(nshots))
+    from_samples = hamiltonian.expectation_from_samples(freq, qubit_map=qubit_map)
+    from_state  = hamiltonian.expectation(result.state())
+    backend.assert_allclose(from_samples, from_state, atol=20 / np.sqrt(nshots))
 
 
 @pytest.mark.parametrize("density_matrix", [False, True])