From a7a496d8fcad9041138a57e2f106b0b586718c03 Mon Sep 17 00:00:00 2001
From: simone bordoni <simone.bordoni@uniroma1.it>
Date: Tue, 12 Mar 2024 17:31:47 +0400
Subject: [PATCH] improve calculate_hamiltonian_state_product method

---
 src/qibo/backends/numpy.py   | 20 +++++++++-----------
 src/qibo/backends/pytorch.py | 16 ++++------------
 tests/test_hamiltonians.py   |  5 +++++
 3 files changed, 18 insertions(+), 23 deletions(-)

diff --git a/src/qibo/backends/numpy.py b/src/qibo/backends/numpy.py
index 6e0b4ad3ef..ae96a6eef5 100644
--- a/src/qibo/backends/numpy.py
+++ b/src/qibo/backends/numpy.py
@@ -718,9 +718,9 @@ def calculate_overlap(self, state1, state2):
         return self.np.abs(self.np.sum(np.conj(self.cast(state1)) * self.cast(state2)))
 
     def calculate_overlap_density_matrix(self, state1, state2):
-        state1 = self.cast(state1)
-        state2 = self.cast(state2)
-        return self.np.trace(self.np.transpose(np.conj(state1)) @ state2)
+        return self.np.trace(
+            self.np.matmul(self.np.conj(self.cast(state1)).T, self.cast(state2))
+        )
 
     def calculate_eigenvalues(self, matrix, k=6):
         if self.issparse(matrix):
@@ -767,20 +767,18 @@ def calculate_expectation_density_matrix(self, hamiltonian, state, normalize):
             ev = ev / norm
         return ev
 
+    # TODO: remove this method
     def calculate_hamiltonian_matrix_product(self, matrix1, matrix2):
-        return self.np.dot(matrix1, matrix2)
+        return matrix1 @ matrix2
 
+    # TODO: remove this method
     def calculate_hamiltonian_state_product(self, matrix, state):
-        rank = len(tuple(state.shape))
-        if rank == 1:  # vector
-            return matrix.dot(state[:, np.newaxis])[:, 0]
-        elif rank == 2:  # matrix
-            return matrix.dot(state)
-        else:
+        if len(tuple(state.shape)) > 2:
             raise_error(
                 ValueError,
-                f"Cannot multiply Hamiltonian with rank-{rank} tensor.",
+                f"Cannot multiply Hamiltonian with rank-{len(tuple(state.shape))} tensor.",
             )
+        return matrix @ state
 
     def assert_allclose(self, value, target, rtol=1e-7, atol=0.0):
         if isinstance(value, CircuitResult) or isinstance(value, QuantumState):
diff --git a/src/qibo/backends/pytorch.py b/src/qibo/backends/pytorch.py
index bfd45222a9..9a321364ee 100644
--- a/src/qibo/backends/pytorch.py
+++ b/src/qibo/backends/pytorch.py
@@ -148,10 +148,10 @@ def sample_shots(self, probabilities, nshots):
             self.cast(probabilities, dtype="float"), nshots, replacement=True
         )
 
-    def calculate_overlap_density_matrix(self, state1, state2):
-        return self.np.trace(
-            self.np.matmul(self.np.conj(self.cast(state1)).T, self.cast(state2))
-        )
+    # def calculate_overlap_density_matrix(self, state1, state2):
+    #     return self.np.trace(
+    #         self.np.matmul(self.np.conj(self.cast(state1)).T, self.cast(state2))
+    #     )
 
     def calculate_eigenvalues(self, matrix, k=6):
         return self.np.linalg.eigvalsh(matrix)  # pylint: disable=not-callable
@@ -168,14 +168,6 @@ def calculate_matrix_exp(self, a, matrix, eigenvectors=None, eigenvalues=None):
         ud = self.np.conj(eigenvectors).T
         return self.np.matmul(eigenvectors, self.np.matmul(expd, ud))
 
-    def calculate_hamiltonian_matrix_product(self, matrix1, matrix2):
-        if self.issparse(matrix1) or self.issparse(matrix2):  # pragma: no cover
-            return self.np.sparse.mm(matrix1, matrix2)  # pylint: disable=E1102
-        return self.np.matmul(matrix1, matrix2)
-
-    def calculate_hamiltonian_state_product(self, matrix, state):
-        return self.np.matmul(matrix, state)
-
     def test_regressions(self, name):
         if name == "test_measurementresult_apply_bitflips":
             return [
diff --git a/tests/test_hamiltonians.py b/tests/test_hamiltonians.py
index 1d1742e328..d388bffe39 100644
--- a/tests/test_hamiltonians.py
+++ b/tests/test_hamiltonians.py
@@ -177,6 +177,11 @@ def test_hamiltonian_matmul(backend, sparse_type):
         backend.assert_allclose((H1 @ H2).matrix, (m1 @ m2))
         backend.assert_allclose((H2 @ H1).matrix, (m2 @ m1))
 
+    try:
+        H1 @ np.zeros(3 * (2**nqubits,), dtype=m1.dtype)
+    except Exception as error:
+        print(error)
+
     with pytest.raises(ValueError):
         H1 @ np.zeros(3 * (2**nqubits,), dtype=m1.dtype)
     with pytest.raises(NotImplementedError):