From b043e6a0173dc508a3cb9cbef10ca863be478064 Mon Sep 17 00:00:00 2001
From: Liwei Yang <yang0345@e.ntu.edu.sg>
Date: Mon, 24 Jul 2023 17:29:29 +0800
Subject: [PATCH] Add the pytest function for MPS in cuquantum.

---
 src/qibotn/QiboCircuitToMPS.py           | 12 ++++----
 src/qibotn/cutn.py                       | 23 +++++++-------
 tests/test_cuquantum_cutensor_backend.py | 38 ++++++++++++++++++++++++
 3 files changed, 57 insertions(+), 16 deletions(-)

diff --git a/src/qibotn/QiboCircuitToMPS.py b/src/qibotn/QiboCircuitToMPS.py
index 49aa125..4c18aac 100644
--- a/src/qibotn/QiboCircuitToMPS.py
+++ b/src/qibotn/QiboCircuitToMPS.py
@@ -2,14 +2,15 @@ import cupy as cp
 import numpy as np
 
 from cuquantum import cutensornet as cutn
-from QiboCircuitConvertor import QiboCircuitToEinsum
-from MPSUtils import get_initial_mps, apply_gate
+from qibotn.QiboCircuitConvertor import QiboCircuitToEinsum
+from qibotn.MPSUtils import get_initial_mps, apply_gate
+
 
 class QiboCircuitToMPS:
-    def __init__(self,circ_qibo, gate_algo, dtype = 'complex128',rand_seed=0,):
+    def __init__(self, circ_qibo, gate_algo, dtype='complex128', rand_seed=0,):
         np.random.seed(rand_seed)
         cp.random.seed(rand_seed)
-        
+
         self.num_qubits = circ_qibo.nqubits
         self.handle = cutn.create()
         self.options = {'handle': self.handle}
@@ -20,7 +21,8 @@ class QiboCircuitToMPS:
         for (gate, qubits) in circuitconvertor.gate_tensors:
             # mapping from qubits to qubit indices
             # apply the gate in-place
-            apply_gate(self.mps_tensors, gate, qubits, algorithm=gate_algo, options=self.options)
+            apply_gate(self.mps_tensors, gate, qubits,
+                       algorithm=gate_algo, options=self.options)
 
     def __del__(self):
         cutn.destroy(self.handle)
diff --git a/src/qibotn/cutn.py b/src/qibotn/cutn.py
index f29a62e..21c316f 100644
--- a/src/qibotn/cutn.py
+++ b/src/qibotn/cutn.py
@@ -1,33 +1,34 @@
 from qibotn.QiboCircuitConvertor import QiboCircuitToEinsum
 from cuquantum import contract
-from MpsHelper import MPSHelper
-import cuquantum 
 from cuquantum import cutensornet as cutn
 import cupy as cp
 import numpy as np
 from qibo.models import QFT
-from QiboCircuitToMPS import QiboCircuitToMPS
-from MPSContractionHelper import MPSContractionHelper
+from qibotn.QiboCircuitToMPS import QiboCircuitToMPS
+from qibotn.MPSContractionHelper import MPSContractionHelper
 
 
 def eval(qibo_circ, datatype):
     myconvertor = QiboCircuitToEinsum(qibo_circ, dtype=datatype)
     return contract(*myconvertor.state_vector_operands())
 
+
 def eval_mps(qibo_circ, gate_algo, datatype):
-    myconvertor = QiboCircuitToMPS(qibo_circ, gate_algo, datatype)  
+    myconvertor = QiboCircuitToMPS(qibo_circ, gate_algo, dtype=datatype)
     mps_helper = MPSContractionHelper(myconvertor.num_qubits)
-    sv_mps = mps_helper.contract_state_vector(myconvertor.mps_tensors,myconvertor.options)
+    sv_mps = mps_helper.contract_state_vector(
+        myconvertor.mps_tensors, myconvertor.options)
     return sv_mps
 
+
 if __name__ == "__main__":
-    num_qubits  = 25
+    num_qubits = 25
     swaps = True
     circ_qibo = QFT(num_qubits, swaps)
-    
-    exact_gate_algorithm = {'qr_method': False, 
-                            'svd_method':{'partition': 'UV', 'abs_cutoff':1e-12}}
-    dtype="complex128"
+
+    exact_gate_algorithm = {'qr_method': False,
+                            'svd_method': {'partition': 'UV', 'abs_cutoff': 1e-12}}
+    dtype = "complex128"
     sv_mps = eval_mps(circ_qibo, exact_gate_algorithm, dtype)
     sv_reference = eval(circ_qibo, dtype)
     state_vec = np.array(circ_qibo())
diff --git a/tests/test_cuquantum_cutensor_backend.py b/tests/test_cuquantum_cutensor_backend.py
index e7f2804..1df3d0f 100644
--- a/tests/test_cuquantum_cutensor_backend.py
+++ b/tests/test_cuquantum_cutensor_backend.py
@@ -2,6 +2,7 @@ from timeit import default_timer as timer
 
 import config
 import numpy as np
+import cupy as cp
 import pytest
 import qibo
 from qibo.models import QFT
@@ -46,3 +47,40 @@ def test_eval(nqubits: int, dtype="complex128"):
     assert 1e-2 * qibo_time < cutn_time < 1e2 * qibo_time
     assert np.allclose(
         result_sv, result_tn), "Resulting dense vectors do not match"
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("nqubits", [2, 5, 10])
+def test_mps(nqubits: int, dtype="complex128"):
+    """Evaluate MPS with cuQuantum.
+
+    Args:
+        nqubits (int): Total number of qubits in the system.
+        dtype (str): The data type for precision, 'complex64' for single,
+            'complex128' for double.
+    """
+    import qibotn.cutn
+
+    # Test qibo
+    qibo.set_backend(backend=config.qibo.backend,
+                     platform=config.qibo.platform)
+
+    qibo_time, (circ_qibo, result_sv) = time(
+        lambda: qibo_qft(nqubits, swaps=True))
+
+    result_sv_cp = cp.asarray(result_sv)
+
+    # Test of MPS
+    gate_algo = {'qr_method': False,
+                 'svd_method': {
+                     'partition': 'UV',
+                     'abs_cutoff': 1e-12,
+                 }}
+
+    cutn_time, result_tn = time(
+        lambda: qibotn.cutn.eval_mps(circ_qibo, gate_algo, dtype).flatten())
+
+    print(
+        f"State vector difference: {abs(result_tn - result_sv_cp).max():0.3e}")
+
+    assert cp.allclose(result_tn, result_sv_cp)