From 89a81542ee570ef895c9122cba0369940efe426d Mon Sep 17 00:00:00 2001
From: vinitha-balachandran
 <balachan@asp2a-login-astar01.head.cm.asp2a.nscc.sg>
Date: Thu, 28 Mar 2024 15:09:11 +0800
Subject: [PATCH] Added changes for expectation calculation in qutensornet

---
 src/qibotn/eval_qu.py | 81 +++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 81 insertions(+)

diff --git a/src/qibotn/eval_qu.py b/src/qibotn/eval_qu.py
index 6cb9f395..4d06034c 100644
--- a/src/qibotn/eval_qu.py
+++ b/src/qibotn/eval_qu.py
@@ -44,3 +44,84 @@ def dense_vector_tn_qu(qasm: str, initial_state, mps_opts, backend="numpy"):
     amplitudes = interim.to_dense(backend=backend)
 
     return amplitudes
+
+
+
+def expectation_qu(qasm: str, pauli_string_pattern, initial_state, mps_opts, backend="numpy"):
+
+    qasm_mod, nqubits = modify_qasm(qasm)
+    if initial_state is not None:
+        initial_state = init_state_tn(nqubits, initial_state)
+
+    circ_cls = qtn.circuit.CircuitMPS if mps_opts else qtn.circuit.Circuit
+    circ_quimb = circ_cls.from_openqasm2_str(
+        qasm, psi0=initial_state, gate_opts=mps_opts
+    )
+
+    obs = pauli_string_gen(nqubits-1, pauli_string_pattern)
+
+    opt = ctg.ReusableHyperOptimizer(
+    # just do a few runs
+        max_repeats=32,
+    # only use the basic greedy optimizer ...
+        methods=['greedy'],
+    # ... but pair it with reconfiguration
+        reconf_opts={},
+    # just uniformly sample the space
+        optlib='random',
+    # terminate search if contraction is cheap
+        max_time='rate:1e6',
+    # account for both flops and write - usually wise for practical performance
+        minimize='combo-64',
+    # persist paths found in here
+        directory='cotengra_cache_eco',
+    )
+    expectation = circ_quimb.local_expectation(obs,  where=list(range(nqubits-1)), optimize=opt, 
+        simplify_sequence ="DRC")
+  
+    return expectation
+
+def modify_qasm(qasm_circ):
+    lines =qasm_circ.split("\n")
+    qasm_circ_mod =[]
+    while lines:
+        line = lines.pop(0).strip()
+        sta = re.compile(r"qreg\s+(\w+)\s*\[(\d+)\];")
+        match = sta.match(line)
+        if match:
+            name, nqubits = match.groups()
+            qasm_circ_mod.append(f"qreg q[{int(nqubits)+1}];")
+        else:
+            qasm_circ_mod.append(line)
+    qasm_circ_mod = "\n".join(qasm_circ_mod)
+    return qasm_circ_mod, int(nqubits)+1
+
+def pauli_string_gen(nqubits, pauli_string_pattern):
+    """Used internally to generate the string based on given pattern and number
+    of qubit.
+
+    Parameters:
+        nqubits(int): Number of qubits of Quantum Circuit
+        pauli_string_pattern(str): Strings representing sequence of pauli gates.
+
+    Returns:
+        String representation of the actual pauli string from the pattern.
+
+    Example: pattern: "XZ", number of qubit: 7, output = XZXZXZX
+    """
+    if nqubits <= 0:
+        return "Invalid input. N should be a positive integer."
+
+    result = ""
+
+    for i in range(nqubits):
+        char_to_add = pauli_string_pattern[i % len(pauli_string_pattern)]
+        result += char_to_add
+    
+
+    for i, c in enumerate(result):
+        if (i==0):
+            obs = qu.pauli(c)
+        else:
+            obs = obs&qu.pauli(c)
+    return obs