Allowing complex expectation value

CQCL · Jun 7, 2024 · 44f7d6b · 44f7d6b
1 parent 4aeb976
commit 44f7d6b
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 7 deletions.
diff --git a/pytket/extensions/cutensornet/general_state/tensor_network_state.py b/pytket/extensions/cutensornet/general_state/tensor_network_state.py
@@ -212,7 +212,7 @@ def expectation_value(
         operator: QubitPauliOperator,
         attributes: Optional[dict] = None,
         scratch_fraction: float = 0.5,
-    ) -> float:
+    ) -> complex:
         """Calculates the expectation value of the given operator.
 
         Args:
@@ -374,9 +374,8 @@ def expectation_value(
             # Note: we can also return `state_norm.item()`, but this should be 1 since
             # we are always running unitary circuits
             assert np.isclose(state_norm.item(), 1.0)
-            # The expectation value is a real number
-            assert np.isclose(expectation_value.item().imag, 0.0)
-            return expectation_value.item().real  # type: ignore
+
+            return expectation_value.item()  # type: ignore
 
         finally:
             #####################################################

diff --git a/tests/test_general_state.py b/tests/test_general_state.py
@@ -181,14 +181,14 @@ def to_bool_tuple(n_qubits: int, x: int) -> tuple:
         ),
         QubitPauliOperator(
             {
-                QubitPauliString({Qubit(0): Pauli.X, Qubit(1): Pauli.Y}): 3.5,
+                QubitPauliString({Qubit(0): Pauli.X, Qubit(1): Pauli.Y}): 3.5 + 0.3j,
             }
         ),
         QubitPauliOperator(
             {
                 QubitPauliString({Qubit(0): Pauli.Z}): 0.25,
-                QubitPauliString({Qubit(1): Pauli.Y}): 0.33,
-                QubitPauliString({Qubit(0): Pauli.X, Qubit(1): Pauli.X}): 0.42,
+                QubitPauliString({Qubit(1): Pauli.Y}): 0.33j,
+                QubitPauliString({Qubit(0): Pauli.X, Qubit(1): Pauli.X}): 0.42 + 0.1j,
             }
         ),
     ],