From 1c0ed631866bd3b9d6b387a6e599a28fea3fb27f Mon Sep 17 00:00:00 2001
From: vinitha-balachandran
 <balachan@x1001c0s0b0n0.hostmgmt2001.cm.asp2a.nscc.sg>
Date: Wed, 24 Apr 2024 16:38:38 +0800
Subject: [PATCH] Fixed import errors

---
 src/qibotn/backends/quimb.py | 20 +++++++++++++++-----
 src/qibotn/eval_qu.py        | 36 ++++++++++++++++++++++++++++++++++++
 2 files changed, 51 insertions(+), 5 deletions(-)

diff --git a/src/qibotn/backends/quimb.py b/src/qibotn/backends/quimb.py
index 35c55f58..f73dd7bf 100644
--- a/src/qibotn/backends/quimb.py
+++ b/src/qibotn/backends/quimb.py
@@ -12,11 +12,21 @@ def __init__(self, runcard):
         if runcard is not None:
             self.MPI_enabled = runcard.get("MPI_enabled", False)
             self.NCCL_enabled = runcard.get("NCCL_enabled", False)
-            self.expectation_enabled = True
-            expectation_enabled_dict = runcard.get("expectation_enabled", {})
-            self.pauli_string_pattern = expectation_enabled_dict.get(
-                "pauli_string_pattern", None
-            )
+
+            expectation_enabled_value = runcard.get("expectation_enabled")
+            if expectation_enabled_value is True:
+                self.expectation_enabled = True
+                self.pauli_string_pattern = "ZZZZ"
+            elif expectation_enabled_value is False:
+                self.expectation_enabled = False
+            elif isinstance(expectation_enabled_value, dict):
+                self.expectation_enabled = True
+                expectation_enabled_dict = runcard.get("expectation_enabled", {})
+                self.pauli_string_pattern = expectation_enabled_dict.get(
+                    "pauli_string_pattern", None
+                )
+            else:
+                raise TypeError("expectation_enabled has an unexpected type")
 
             mps_enabled_value = runcard.get("MPS_enabled")
             if mps_enabled_value is True:
diff --git a/src/qibotn/eval_qu.py b/src/qibotn/eval_qu.py
index df9425f6..302a62c0 100644
--- a/src/qibotn/eval_qu.py
+++ b/src/qibotn/eval_qu.py
@@ -49,8 +49,25 @@ def dense_vector_tn_qu(qasm: str, initial_state, mps_opts, backend="numpy"):
 def expectation_qu(
     qasm: str, pauli_string_pattern, initial_state, mps_opts, backend="numpy"
 ):
+    """Calculates the expectation value of a given global observable as pauli
+    string pattern.
 
+    Args:
+        qasm (str): QASM program.
+        pauli_string_pattern (str): Pattern for Pauli string that defines the global observable.
+        initial_state (list): Initial state in the dense vector form. If ``None`` the default ``|00...0>`` state is used.
+        mps_opts (dict): Parameters to tune the gate_opts for mps settings in ``class quimb.tensor.circuit.CircuitMPS``.
+        backend (str):  Backend to perform the contraction with, e.g. ``numpy``, ``cupy``, ``jax``. Passed to ``opt_einsum``.
+
+    Returns:
+        float: Expectation value of the global observable for the final state after the simulation of the circuit.
+    """
+    # use cotengra package for tensor contractions
+    import cotengra as ctg
+
+    # used to add one more qubit to the circuit to make the observable local
     qasm_mod, nqubits = modify_qasm(qasm)
+
     if initial_state is not None:
         initial_state = init_state_tn(nqubits, initial_state)
 
@@ -59,8 +76,10 @@ def expectation_qu(
         qasm, psi0=initial_state, gate_opts=mps_opts
     )
 
+    # generates the global observable
     obs = pauli_string_gen(nqubits - 1, pauli_string_pattern)
 
+    # parameters to find the contraction path using cotengra
     opt = ctg.ReusableHyperOptimizer(
         # just do a few runs
         max_repeats=32,
@@ -77,6 +96,8 @@ def expectation_qu(
         # persist paths found in here
         directory="cotengra_cache_eco",
     )
+
+    # expectation value
     expectation = circ_quimb.local_expectation(
         obs, where=list(range(nqubits - 1)), optimize=opt, simplify_sequence="DRC"
     )
@@ -85,7 +106,19 @@ def expectation_qu(
 
 
 def modify_qasm(qasm_circ):
+    """Generate a modified qasm string.
+
+    Args:
+        qasm (str): QASM program.
+
+    Returns:
+        string: QASM program with an additional auxillary qubit for the calculation of expectation
+    """
+
+    import re
+
     lines = qasm_circ.split("\n")
+
     qasm_circ_mod = []
     while lines:
         line = lines.pop(0).strip()
@@ -97,6 +130,7 @@ def modify_qasm(qasm_circ):
         else:
             qasm_circ_mod.append(line)
     qasm_circ_mod = "\n".join(qasm_circ_mod)
+
     return qasm_circ_mod, int(nqubits) + 1
 
 
@@ -113,6 +147,8 @@ def pauli_string_gen(nqubits, pauli_string_pattern):
 
     Example: pattern: "XZ", number of qubit: 7, output = XZXZXZX
     """
+    import quimb as qu
+
     if nqubits <= 0:
         return "Invalid input. N should be a positive integer."