Repository: Quantum-TII/qibo (python)

Issue:
Hamiltonian expectation for density matrices
**Describe the bug**
Hamiltonian expectation value calculation is wrong for density matrices, as they are treated as state vectors.

**To Reproduce**
```Python
import numpy as np
from qibo import hamiltonians

def random_density_matrix(nqubits):
    shape = 2 * (2 ** nqubits,)
    rho = np.random.random(shape) + 1j * np.random.random(shape)
    rho = rho + rho.T.conj()
    # Normalize
    ids = np.arange(2 ** nqubits)
    rho[ids, ids] = rho[ids, ids] / np.trace(rho)
    return rho


nqubits = 5
ham = hamiltonians.XXZ(nqubits)
state = random_density_matrix(nqubits)
print(ham.expectation(state))

matrix = np.array(ham.matrix)
# EV calculation for density matrices
print(np.trace(matrix @ state))
# EV calculation for state vectors
print(np.sum(np.conj(state) * (matrix @ state)))
```

**Expected behavior**
The calculated expectation value (first printed number) should agree with the `np.trace` calculation (second printed number) but instead it agrees with the third printed value, which is the formula for state vectors.


Pull request description:
Hamiltonian expectation for density matrices
Fixes #339.

Patch:
diff --git a/src/qibo/abstractions/states.py b/src/qibo/abstractions/states.py
index cd3035500..0ed416a20 100644
--- a/src/qibo/abstractions/states.py
+++ b/src/qibo/abstractions/states.py
@@ -262,3 +262,15 @@ class AbstractState(ABC):
                 ``p0`` will be used for both bitflips.
         """
         raise_error(NotImplementedError)
+
+    @abstractmethod
+    def expectation(self, hamiltonian, normalize=False): # pragma: no cover
+        """Calculates Hamiltonian expectation value with respect to the state.
+
+        Args:
+            hamiltonian (`qibo.abstractions.hamiltonians.Hamiltonian`):
+                Hamiltonian object to calculate the expectation value of.
+            normalize (bool): Normalize the result by dividing with the norm of
+                the state. Default is ``False``.
+        """
+        raise_error(NotImplementedError)
diff --git a/src/qibo/core/hamiltonians.py b/src/qibo/core/hamiltonians.py
index 03f788d87..47399495d 100644
--- a/src/qibo/core/hamiltonians.py
+++ b/src/qibo/core/hamiltonians.py
@@ -52,13 +52,24 @@ class Hamiltonian(hamiltonians.Hamiltonian):
 
     def expectation(self, state, normalize=False):
         if isinstance(state, states.AbstractState):
-            state = state.tensor
-        statec = self.K.conj(state)
-        hstate = self @ state
-        ev = self.K.real(self.K.sum(statec * hstate))
-        if normalize:
-            norm = self.K.sum(self.K.square(self.K.abs(state)))
-            return ev / norm
+            ev = state.expectation(self, normalize)
+        elif isinstance(state, K.tensor_types):
+            shape = tuple(state.shape)
+            if len(shape) == 1: # state vector
+                from qibo.core.states import VectorState
+                sobj = VectorState.from_tensor(state)
+            elif len(shape) == 2: # density matrix
+                from qibo.core.states import MatrixState
+                sobj = MatrixState.from_tensor(state)
+            else:
+                raise_error(ValueError, "Cannot calculate Hamiltonian "
+                                        "expectation value for state of shape "
+                                        "{}.".format(shape))
+            ev = sobj.expectation(self, normalize)
+        else:
+            raise_error(TypeError, "Cannot calculate Hamiltonian expectation "
+                                   "value for state of type {}."
+                                   "".format(type(state)))
         return ev
 
     def eye(self, n=None):
diff --git a/src/qibo/core/states.py b/src/qibo/core/states.py
index d57a2c52e..bc4405033 100644
--- a/src/qibo/core/states.py
+++ b/src/qibo/core/states.py
@@ -127,6 +127,15 @@ class VectorState(AbstractState):
         self.measurements = self.measurements.apply_bitflips(p0, p1)
         return self
 
+    def expectation(self, hamiltonian, normalize=False):
+        statec = K.conj(self.tensor)
+        hstate = hamiltonian @ self.tensor
+        ev = K.real(K.sum(statec * hstate))
+        if normalize:
+            norm = K.sum(K.square(K.abs(self.tensor)))
+            ev = ev / norm
+        return ev
+
 
 class MatrixState(VectorState):
 
@@ -148,6 +157,18 @@ class MatrixState(VectorState):
         state = K.einsum(traceout, K.reshape(self.tensor, shape))
         return K.cast(state, dtype='DTYPE')
 
+    def expectation(self, hamiltonian, normalize=False):
+        from qibo.abstractions.hamiltonians import TrotterHamiltonian
+        if isinstance(hamiltonian, TrotterHamiltonian):
+            # use dense form of Trotter Hamiltonians because their
+            # multiplication to rank-2 tensors is not implemented
+            hamiltonian = hamiltonian.dense
+        ev = K.real(K.trace(hamiltonian @ self.tensor))
+        if normalize:
+            norm = K.real(K.trace(self.tensor))
+            ev = ev / norm
+        return ev
+
 
 class DistributedState(VectorState):
     """Data structure that holds the pieces of a state vector.


Test patch:
diff --git a/src/qibo/tests/test_hamiltonians.py b/src/qibo/tests/test_hamiltonians.py
index c0b8a0bef..ee593a2b0 100644
--- a/src/qibo/tests/test_hamiltonians.py
+++ b/src/qibo/tests/test_hamiltonians.py
@@ -151,18 +151,34 @@ def test_hamiltonian_exponentiation(numpy, trotter):
 
 @pytest.mark.parametrize("numpy", [True, False])
 @pytest.mark.parametrize("trotter", [True, False])
-def test_hamiltonian_expectation(numpy, trotter):
+@pytest.mark.parametrize("density_matrix", [True, False])
+def test_hamiltonian_expectation(numpy, trotter, density_matrix):
     h = XXZ(nqubits=3, delta=0.5, numpy=numpy, trotter=trotter)
     matrix = np.array(h.matrix)
 
-    state = utils.random_numpy_complex(8)
-    norm = (np.abs(state) ** 2).sum()
-    target_ev = (state.conj() * matrix.dot(state)).sum().real
+    if density_matrix:
+        state = utils.random_numpy_complex((8, 8))
+        state = state + state.T.conj()
+        norm = np.trace(state)
+        target_ev = np.trace(matrix.dot(state)).real
+    else:
+        state = utils.random_numpy_complex(8)
+        norm = np.sum(np.abs(state) ** 2)
+        target_ev = np.sum(state.conj() * matrix.dot(state)).real
 
     np.testing.assert_allclose(h.expectation(state), target_ev)
     np.testing.assert_allclose(h.expectation(state, True), target_ev / norm)
 
 
+def test_hamiltonian_expectation_errors():
+    h = XXZ(nqubits=3, delta=0.5)
+    state = utils.random_numpy_complex((4, 4, 4))
+    with pytest.raises(ValueError):
+        h.expectation(state)
+    with pytest.raises(TypeError):
+        h.expectation("test")
+
+
 @pytest.mark.parametrize("numpy", [True, False])
 def test_hamiltonian_runtime_errors(numpy):
     """Testing hamiltonian runtime errors."""
diff --git a/src/qibo/tests_new/test_core_states.py b/src/qibo/tests_new/test_core_states.py
index 05f93453c..21367fd91 100644
--- a/src/qibo/tests_new/test_core_states.py
+++ b/src/qibo/tests_new/test_core_states.py
@@ -158,3 +158,40 @@ def test_state_apply_bitflips():
     with pytest.raises(RuntimeError):
         state.apply_bitflips(0.1)
     # Bitflips are tested in measurement tests
+
+
+@pytest.mark.parametrize("trotter", [True, False])
+def test_vector_state_expectation(backend, trotter):
+    original_backend = qibo.get_backend()
+    qibo.set_backend(backend)
+    from qibo.hamiltonians import XXZ
+    ham = XXZ(nqubits=5, delta=0.5, trotter=trotter)
+    matrix = np.array(ham.matrix)
+
+    state = np.random.random(32) + 1j * np.random.random(32)
+    norm = np.sum(np.abs(state) ** 2)
+    target_ev = np.sum(state.conj() * matrix.dot(state)).real
+    state = states.VectorState.from_tensor(state)
+
+    np.testing.assert_allclose(state.expectation(ham), target_ev)
+    np.testing.assert_allclose(state.expectation(ham, True), target_ev / norm)
+    qibo.set_backend(original_backend)
+
+
+@pytest.mark.parametrize("trotter", [True, False])
+def test_matrix_state_expectation(backend, trotter):
+    original_backend = qibo.get_backend()
+    qibo.set_backend(backend)
+    from qibo.hamiltonians import TFIM
+    ham = TFIM(nqubits=2, h=1.0, trotter=trotter)
+    matrix = np.array(ham.matrix)
+
+    state = np.random.random((4, 4)) + 1j * np.random.random((4, 4))
+    state = state + state.T.conj()
+    norm = np.trace(state)
+    target_ev = np.trace(matrix.dot(state)).real
+    state = states.MatrixState.from_tensor(state)
+
+    np.testing.assert_allclose(state.expectation(ham), target_ev)
+    np.testing.assert_allclose(state.expectation(ham, True), target_ev / norm)
+    qibo.set_backend(original_backend)