From 250c41bf4dbe72508ec32eafe2cd231942188d8c Mon Sep 17 00:00:00 2001
From: Alessandro Candido <candido.ale@gmail.com>
Date: Wed, 15 Feb 2023 12:37:50 +0100
Subject: [PATCH] Start reworking qasm parser

---
 src/qibotn/{qasm_quimb.py => quimb.py} | 48 +++++++-------------------
 tests/test_qasm_quimb_backend.py       | 11 +++---
 2 files changed, 17 insertions(+), 42 deletions(-)
 rename src/qibotn/{qasm_quimb.py => quimb.py} (55%)

diff --git a/src/qibotn/qasm_quimb.py b/src/qibotn/quimb.py
similarity index 55%
rename from src/qibotn/qasm_quimb.py
rename to src/qibotn/quimb.py
index f7546eb..93457f4 100644
--- a/src/qibotn/qasm_quimb.py
+++ b/src/qibotn/quimb.py
@@ -1,6 +1,6 @@
 import re
 
-import quimb as qu
+import qibo
 import quimb.tensor as qtn
 import numpy as np
 
@@ -22,7 +22,7 @@ def get_gate_functions(qasm_str, start_idx):
         idx_inc += 1
 
 
-def qasm_QFT(nqubits: int, qasm_str: str, with_swaps: bool = True, psi0=None):
+def convert(nqubits: int, qasm_str: str, with_swaps: bool = True, psi0=None):
     circ = qtn.Circuit(nqubits, psi0=psi0)
 
     qasm_str = qasm_str.split("\n")
@@ -53,44 +53,22 @@ def qasm_QFT(nqubits: int, qasm_str: str, with_swaps: bool = True, psi0=None):
     return circ
 
 
-def init_state_tn(nqubits, init_state_sv, tn_lib="quimb"):
+def init_state_tn(nqubits, init_state_sv):
     dims = tuple(2 * np.ones(nqubits, dtype=int))
 
-    if tn_lib == "quimb":
-        init_state_MPS = qtn.tensor_1d.MatrixProductState.from_dense(
-            init_state_sv, dims
-        )
-    else:
-        # TODO: Add cuquantum later
-        assert False, "Unsupported tensor network backend in initilization"
-
-    return init_state_MPS
+    return qtn.tensor_1d.MatrixProductState.from_dense(init_state_sv, dims)
 
 
-def tn_circ_eval(
-    nqubits, qasm_circ, init_state, swaps=True, tn_lib="quimb", backend="numpy"
-):
-    if tn_lib == "quimb":
+def eval(qasm: str, init_state, backend="numpy", swaps=True):
+    """Evaluate QASM with Quimb
 
-        circ_quimb = qasm_QFT(nqubits, qasm_circ, swaps, psi0=init_state)
-        interim = circ_quimb.psi.full_simplify(seq="DRC")
-        result = interim.to_dense(backend=backend).flatten()
-        return result
-    else:
-        # TODO: Change assert or value. Add cuquantum later
-        assert False, "Unsupported tensor network library"
+    backend (quimb): numpy, cupy, jax. Passed to ``opt_einsum``.
 
+    """
+    circuit = qibo.models.Circuit.from_qasm(qasm)
+    init_state_mps = init_state_tn(circuit.nqubits, init_state)
+    circ_quimb = convert(circuit, swaps=swaps, psi0=init_state_mps)
+    interim = circ_quimb.psi.full_simplify(seq="DRC")
+    amplitudes = interim.to_dense(backend=backend).flatten()
 
-def eval_QI_qft(nqubits, qasm_circ, init_state, backend="numpy", swaps=True):
-    # backend (quimb): numpy, cupy, jax. Passed to ``opt_einsum``.
-
-    # Quimb circuit
-    init_state_mps = init_state_tn(nqubits=nqubits, init_state_sv=init_state)
-    amplitudes = tn_circ_eval(
-        nqubits=nqubits,
-        qasm_circ=qasm_circ,
-        init_state=init_state_mps,
-        swaps=swaps,
-        tn_lib="quimb",
-    )
     return amplitudes
diff --git a/tests/test_qasm_quimb_backend.py b/tests/test_qasm_quimb_backend.py
index edb99a2..9f4c854 100644
--- a/tests/test_qasm_quimb_backend.py
+++ b/tests/test_qasm_quimb_backend.py
@@ -14,8 +14,7 @@ def init_state_sv(nqubits):
     init_state = np.random.random(2**nqubits) + 1j * np.random.random(2**nqubits)
     init_state = init_state / np.sqrt((np.abs(init_state) ** 2).sum())
     # An unmodified init_state has to be converted to tn format
-    init_state_for_tn = copy.deepcopy(init_state)
-    return init_state, init_state_for_tn
+    return init_state
 
 
 def qibo_qft(nqubits, init_state, swaps):
@@ -29,14 +28,12 @@ def test_eval(nqubits: int):
     os.environ["QUIMB_NUM_PROCS"] = str(os.cpu_count())
     from qibotn import qasm_quimb
 
-    init_state_qibo, init_state_for_tn = init_state_sv(nqubits=nqubits)
+    init_state = init_state_sv(nqubits=nqubits)
 
     # Test qibo
     qibo.set_backend(backend=config.qibo["backend"], platform=config.qibo["platform"])
     start_time = timer()
-    qibo_circ, result_sv = qibo_qft(
-        nqubits, init_state=init_state_qibo, swaps=config.qibo["swaps"]
-    )
+    qibo_circ, result_sv = qibo_qft(nqubits, init_state, swaps=config.qibo["swaps"])
     end_time = timer()
     qibo_time = end_time - start_time
 
@@ -49,7 +46,7 @@ def test_eval(nqubits: int):
     result_tn = qasm_quimb.eval_QI_qft(
         nqubits=nqubits,
         qasm_circ=qasm_circ,
-        init_state=init_state_for_tn,
+        init_state=copy.deepcopy(init_state),
         backend=config.quimb["backend"],
         swaps=config.quimb["swaps"],
     )