From c05fe12e46b7e368b19fb4dadb1dbeda1e09f9cd Mon Sep 17 00:00:00 2001
From: tankya2 <tankya2@ihpc.a-star.edu.sg>
Date: Tue, 21 Mar 2023 15:17:29 +0800
Subject: [PATCH] Updated with pytest script for  testing

---
 setup.py                                 |  5 ++++
 src/qibotn/cutn.py                       | 10 +++++++
 tests/test_cuquantum_cutensor_backend.py | 38 ++++++++++++++++++++++++
 3 files changed, 53 insertions(+)
 create mode 100644 src/qibotn/cutn.py
 create mode 100644 tests/test_cuquantum_cutensor_backend.py

diff --git a/setup.py b/setup.py
index aefac89..147e58b 100644
--- a/setup.py
+++ b/setup.py
@@ -42,6 +42,11 @@ setup(
         "qibo>=0.1.10",
         "qibojit>=0.0.7",
         "quimb[tensor]>=1.4.0",
+        "cupy>=11.6.0",
+        "cuquantum-python-cu11>=22.11.0.1",
+        "custatevec-cu11>=1.2.0",
+        "cutensor-cu11>=1.6.2",
+        "cutensornet-cu11>=2.0.0",
     ],
     extras_require={
         "docs": [],
diff --git a/src/qibotn/cutn.py b/src/qibotn/cutn.py
new file mode 100644
index 0000000..fe0ae38
--- /dev/null
+++ b/src/qibotn/cutn.py
@@ -0,0 +1,10 @@
+# from qibotn import quimb as qiboquimb
+from QiboCircuitConvertor import QiboCircuitToEinsum
+from cuquantum import contract
+
+
+def eval(qibo_circ):
+    myconvertor = QiboCircuitToEinsum(qibo_circ, dtype="complex128")
+    operands_expression = myconvertor.state_vector()
+    results = contract(*operands_expression)
+    return results.flatten()
diff --git a/tests/test_cuquantum_cutensor_backend.py b/tests/test_cuquantum_cutensor_backend.py
new file mode 100644
index 0000000..9c989a3
--- /dev/null
+++ b/tests/test_cuquantum_cutensor_backend.py
@@ -0,0 +1,38 @@
+import copy
+import os
+from timeit import default_timer as timer
+
+import config
+import numpy as np
+import pytest
+import qibo
+from qibo.models import QFT
+
+
+def qibo_qft(nqubits, swaps):
+    circ_qibo = QFT(nqubits, swaps)
+    state_vec = np.array(circ_qibo())
+    return circ_qibo, state_vec
+
+
+def time(func):
+    start = timer()
+    res = func()
+    end = timer()
+    time = end - start
+    return time, res
+
+
+@pytest.mark.parametrize("nqubits", [1, 2, 5, 10])
+def test_eval(nqubits: int):
+    import qibotn.cutn
+
+    # Test qibo
+    qibo.set_backend(backend=config.qibo.backend, platform=config.qibo.platform)
+    qibo_time, (qibo_circ, result_sv) = time(lambda: qibo_qft(nqubits, swaps=True))
+
+    # Test Cuquantum
+    cutn_time, result_tn = time(lambda: qibotn.cutn.eval(qibo_circ))
+
+    assert 1e-2 * qibo_time < cutn_time < 1e2 * qibo_time
+    assert np.allclose(result_sv, result_tn), "Resulting dense vectors do not match"