Switch return type to a proper array

qiboteam · Jan 22, 2024 · 17d0a26 · 17d0a26
1 parent e144580
commit 17d0a26
Show file tree

Hide file tree

Showing 3 changed files with 89 additions and 4 deletions.
diff --git a/crate/Cargo.lock b/crate/Cargo.lock
diff --git a/crate/Cargo.toml b/crate/Cargo.toml
@@ -6,6 +6,8 @@ edition = "2021"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
+ndarray = "0.15.6"
+numpy = "0.20.0"
 pyo3 = { version = "0.20.2", features = ["auto-initialize"] }
 
 [dev-dependencies]

diff --git a/crate/src/lib.rs b/crate/src/lib.rs
@@ -1,19 +1,22 @@
+use ndarray::Array2;
+use numpy::PyArray2;
 use pyo3::prelude::*;
 use pyo3::types::PyDict;
 
-pub fn execute_qasm(circuit: String, platform: String, nshots: u32) -> PyResult<Vec<u32>> {
+pub fn execute_qasm(circuit: String, platform: String, nshots: u32) -> PyResult<Array2<i32>> {
     Python::with_gil(|py| {
         let kwargs = PyDict::new(py);
         kwargs.set_item("circuit", circuit)?;
         kwargs.set_item("platform", platform)?;
         kwargs.set_item("nshots", nshots)?;
 
         let qibolab = PyModule::import(py, "qibolab")?;
-        qibolab
+        let pyarray: &PyArray2<i32> = qibolab
             .getattr("execute_qasm")?
             .call((), Some(kwargs))?
             .call_method0("samples")?
-            .call_method0("ravel")?
-            .extract()
+            .extract()?;
+
+        Ok(pyarray.to_owned_array())
     })
 }