Add support for different measurement levels

qiskit/result/utils.py

@@ -127,10 +127,11 @@ def _adjust_creg_sizes(creg_sizes, indices):
 
 
 def marginal_memory(
-    memory: List[str],
+    memory: Union[List[str], np.ndarray],
     indices: Optional[List[int]] = None,
     int_return: bool = False,
     hex_return: bool = False,
+    avg_data: bool = False,
     parallel_threshold: int = 1000,
 ) -> Union[List[str], np.ndarray]:
     """Marginalize memory
@@ -141,7 +142,9 @@ def marginal_memory(
     to 4 threads.
 
     Args:
-        memory: The input memory list, this is a list of hexadecimal strings to be marginalized
+        memory: The input memory list, this is either a list of hexadecimal strings to be marginalized
+            representing measure level 2 memory or a numpy array representing level 0 measurement
+            memory (avg or single) or level 1 measurement memory (single or avg).
         indices: The bit positions of interest to marginalize over. If
             ``None`` (default), do not marginalize at all.
         int_return: If set to ``True`` the output will be a list of integers.
@@ -151,6 +154,9 @@ def marginal_memory(
             strings. By default the return type is a bit string. This and
             ``int_return`` are mutually exclusive and can not be specified
             at the same time.
+        avg_data: If a 2 dimensional numpy array is passed in for ``memory`` this can be set to
+            ``True`` to indicate it's a avg level 0 data instead of level 1
+            single data.
         parallel_threshold: The number of elements in ``memory`` to start running in multiple
             threads. If ``len(memory)`` is >= this value this function will run in multiple
             threads. By default this is set to 50.
@@ -163,6 +169,24 @@ def marginal_memory(
     """
     if int_return and hex_return:
         raise ValueError("Either int_return or hex_return can be specified but not both")
+
+    if isinstance(memory, np.ndarray):
+        if int_return:
+            raise ValueError("int_return option only works with memory list input")
+        if hex_return:
+            raise ValueError("hex_return option only works with memory list input")
+        if indices is None:
+            return memory.copy()
+        if memory.ndim == 1:
+            return results_rs.marginal_measure_level_1_avg(memory, indices)
+        if memory.ndim == 2:
+            if avg_data:
+                return results_rs.marginal_measure_level_0_avg(memory, indices)
+            else:
+                return results_rs.marginal_measure_level_1(memory, indices)
+        if memory.ndim == 3:
+            return results_rs.marginal_measure_level_0(memory, indices)
+        raise ValueError("Invalid input memory array")
     return results_rs.marginal_memory(
         memory,
         indices,

src/results/marginalization.rs

@@ -13,7 +13,11 @@
 use super::converters::hex_to_bin;
 use crate::getenv_use_multiple_threads;
 use hashbrown::HashMap;
+use ndarray::prelude::*;
 use num_bigint::BigUint;
+use num_complex::Complex64;
+use numpy::IntoPyArray;
+use numpy::{PyReadonlyArray1, PyReadonlyArray2, PyReadonlyArray3};
 use pyo3::prelude::*;
 use rayon::prelude::*;
 
@@ -110,7 +114,11 @@ fn map_memory(
     }
 }
 
-#[pyfunction(return_int = "false", return_hex = "false", parallel_threshold = "1000")]
+#[pyfunction(
+    return_int = "false",
+    return_hex = "false",
+    parallel_threshold = "1000"
+)]
 pub fn marginal_memory(
     py: Python,
     memory: Vec<String>,
@@ -157,3 +165,56 @@ pub fn marginal_memory(
         Ok(out_mem.to_object(py))
     }
 }
+
+#[pyfunction]
+pub fn marginal_measure_level_0(
+    py: Python,
+    memory: PyReadonlyArray3<Complex64>,
+    indices: Vec<usize>,
+) -> PyObject {
+    let mem_arr: ArrayView3<Complex64> = memory.as_array();
+    let input_shape = mem_arr.shape();
+    let new_shape = [input_shape[0], indices.len(), input_shape[2]];
+    let out_arr: Array3<Complex64> =
+        Array3::from_shape_fn(new_shape, |(i, j, k)| mem_arr[[i, indices[j], k]]);
+    out_arr.into_pyarray(py).into()
+}
+
+#[pyfunction]
+pub fn marginal_measure_level_0_avg(
+    py: Python,
+    memory: PyReadonlyArray2<Complex64>,
+    indices: Vec<usize>,
+) -> PyObject {
+    let mem_arr: ArrayView2<Complex64> = memory.as_array();
+    let input_shape = mem_arr.shape();
+    let new_shape = [indices.len(), input_shape[1]];
+    let out_arr: Array2<Complex64> =
+        Array2::from_shape_fn(new_shape, |(i, j)| mem_arr[[indices[i], j]]);
+    out_arr.into_pyarray(py).into()
+}
+
+#[pyfunction]
+pub fn marginal_measure_level_1(
+    py: Python,
+    memory: PyReadonlyArray2<Complex64>,
+    indices: Vec<usize>,
+) -> PyObject {
+    let mem_arr: ArrayView2<Complex64> = memory.as_array();
+    let input_shape = mem_arr.shape();
+    let new_shape = [input_shape[0], indices.len()];
+    let out_arr: Array2<Complex64> =
+        Array2::from_shape_fn(new_shape, |(i, j)| mem_arr[[i, indices[j]]]);
+    out_arr.into_pyarray(py).into()
+}
+
+#[pyfunction]
+pub fn marginal_measure_level_1_avg(
+    py: Python,
+    memory: PyReadonlyArray1<Complex64>,
+    indices: Vec<usize>,
+) -> PyResult<PyObject> {
+    let mem_arr: &[Complex64] = memory.as_slice()?;
+    let out_arr: Vec<Complex64> = indices.into_iter().map(|idx| mem_arr[idx]).collect();
+    Ok(out_arr.into_pyarray(py).into())
+}

src/results/mod.rs

@@ -21,5 +21,10 @@ pub fn results(_py: Python, m: &PyModule) -> PyResult<()> {
     m.add_wrapped(wrap_pyfunction!(marginalization::marginal_counts))?;
     m.add_wrapped(wrap_pyfunction!(marginalization::marginal_distribution))?;
     m.add_wrapped(wrap_pyfunction!(marginalization::marginal_memory))?;
+    m.add_wrapped(wrap_pyfunction!(marginalization::marginal_measure_level_0))?;
+    m.add_wrapped(wrap_pyfunction!(
+        marginalization::marginal_measure_level_0_avg
+    ))?;
+    m.add_wrapped(wrap_pyfunction!(marginalization::marginal_measure_level_1))?;
     Ok(())
 }

test/python/result/test_memory_marginalization.py

@@ -12,6 +12,8 @@
 
 """Test marginal_memory() function."""
 
+import numpy as np
+
 from qiskit.test import QiskitTestCase
 from qiskit.result import marginal_memory
 
@@ -53,3 +55,53 @@ def test_error_on_multiple_return_types(self):
         """Test that ValueError raised if multiple return types are requested."""
         with self.assertRaises(ValueError):
             marginal_memory([], int_return=True, hex_return=True)
+
+    def test_memory_level_0(self):
+        """Test that a single level 0 measurement data is correctly marginalized."""
+        memory = np.asarray(
+            [
+                # qubit 0                    qubit 1                     qubit 2
+                [
+                    [-12974255.0, -28106672.0],
+                    [15848939.0, -53271096.0],
+                    [-18731048.0, -56490604.0],
+                ],  # shot 1
+                [
+                    [-18346508.0, -26587824.0],
+                    [-12065728.0, -44948360.0],
+                    [14035275.0, -65373000.0],
+                ],  # shot 2
+                [
+                    [12802274.0, -20436864.0],
+                    [-15967512.0, -37575556.0],
+                    [15201290.0, -65182832.0],
+                ],  # ...
+                [[-9187660.0, -22197716.0], [-17028016.0, -49578552.0], [13526576.0, -61017756.0]],
+                [[7006214.0, -32555228.0], [16144743.0, -33563124.0], [-23524160.0, -66919196.0]],
+            ],
+            dtype=complex,
+        )
+        result = marginal_memory(memory, [0, 2])
+        expected = np.asarray(
+            [
+                [[-12974255.0, -28106672.0], [-18731048.0, -56490604.0]],  # shot 1
+                [[-18346508.0, -26587824.0], [14035275.0, -65373000.0]],  # shot 2
+                [[12802274.0, -20436864.0], [15201290.0, -65182832.0]],  # ...
+                [[-9187660.0, -22197716.0], [13526576.0, -61017756.0]],
+                [[7006214.0, -32555228.0], [-23524160.0, -66919196.0]],
+            ],
+            dtype=complex,
+        )
+        np.testing.assert_array_equal(result, expected)
+
+    def test_memory_level_0_avg(self):
+        """Test that avg level 0 measurement data is correctly marginalized."""
+        memory = np.asarray(
+            [[-1059254.375, -26266612.0], [-9012669.0, -41877468.0], [6027076.0, -54875060.0]],
+            dtype=complex,
+        )
+        result = marginal_memory(memory, [0, 2], avg_data=True)
+        expected = np.asarray(
+            [[-1059254.375, -26266612.0], [6027076.0, -54875060.0]], dtype=complex
+        )
+        np.testing.assert_array_equal(result, expected)