Add dedicated functions for memory marginalization

Cargo.toml

@@ -18,10 +18,11 @@ rand_distr = "0.4.3"
 indexmap = "1.8.1"
 ahash = "0.7.6"
 num-complex = "0.4"
+num-bigint = "0.4"
 
 [dependencies.pyo3]
 version = "0.16.5"
-features = ["extension-module", "hashbrown", "num-complex"]
+features = ["extension-module", "hashbrown", "num-complex", "num-bigint"]
 
 [dependencies.ndarray]
 version = "^0.15.0"

qiskit/result/__init__.py

@@ -25,6 +25,7 @@
    Counts
    marginal_counts
    marginal_distribution
+   marginal_memory
 
 Distributions
 =============
@@ -50,6 +51,7 @@
 from .exceptions import ResultError
 from .utils import marginal_counts
 from .utils import marginal_distribution
+from .utils import marginal_memory
 from .counts import Counts
 
 from .distributions.probability import ProbDistribution

qiskit/result/utils.py

@@ -18,6 +18,8 @@
 from collections import Counter
 from copy import deepcopy
 
+import numpy as np
+
 from qiskit.exceptions import QiskitError
 from qiskit.result.result import Result
 from qiskit.result.counts import Counts
@@ -87,12 +89,9 @@ def marginal_counts(
                     sorted_indices = sorted(
                         indices, reverse=True
                     )  # same convention as for the counts
-                    bit_strings = [_hex_to_bin(s) for s in experiment_result.data.memory]
-                    marginal_bit_strings = [
-                        "".join([s[-idx - 1] for idx in sorted_indices if idx < len(s)]) or "0"
-                        for s in bit_strings
-                    ]
-                    experiment_result.data.memory = [_bin_to_hex(s) for s in marginal_bit_strings]
+                    experiment_result.data.memory = results_rs.marginal_memory(
+                        experiment_result.data.memory, sorted_indices, return_hex=True
+                    )
         return result
     else:
         marg_counts = _marginalize(result, indices)
@@ -127,14 +126,60 @@ def _adjust_creg_sizes(creg_sizes, indices):
     return new_creg_sizes
 
 
+def marginal_memory(
+    memory: List[str],
+    indices: Optional[List[int]] = None,
+    int_return: bool = False,
+    hex_return: bool = False,
+    parallel_threshold: int = 50,
+) -> Union[List[str], np.ndarray]:
+    """Marginalize memory
+
+    This function is multithreaded and will launch a thread pool with threads equal to the number
+    of CPUs by default. You can tune the number of threads with the ``RAYON_NUM_THREADS``
+    environment variable. For example, setting ``RAYON_NUM_THREADS=4`` would limit the thread pool
+    to 4 threads.
+
+    Args:
+        memory: The input memory list, this is a list of hexadecimal strings to be marginalized
+        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.
+            By default the return type is a bit string. This and ``hex_return``
+            are mutually exclusive and can not be specified at the same time.
+        hex_return: If set to ``True`` the output will be a list of hexadecimal
+            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.
+        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.
+
+    Returns:
+        marginal_memory: The list of marginalized memory
+
+    Raises:
+        ValueError: if both ``int_return`` and ``hex_return`` are set to ``True``
+    """
+    if int_return and hex_return:
+        raise ValueError("Either int_return or hex_return can be specified but not both")
+    return results_rs.marginal_memory(
+        memory,
+        indices,
+        return_int=int_return,
+        return_hex=hex_return,
+        parallel_threshold=parallel_threshold,
+    )
+
+
 def marginal_distribution(
     counts: dict, indices: Optional[List[int]] = None, format_marginal: bool = False
 ) -> Dict[str, int]:
     """Marginalize counts from an experiment over some indices of interest.
 
     Unlike :func:`~.marginal_counts` this function respects the order of
-    the input ``indices``. If the input ``indices`` list is specified, the order
-    the bit indices will be the output order of the bitstrings
+    the input ``indices``. If the input ``indices`` list is specified then the order
+    the bit indices are specified will be the output order of the bitstrings
     in the marginalized output.
 
     Args:

src/results/converters.rs

@@ -0,0 +1,45 @@
+// This code is part of Qiskit.
+//
+// (C) Copyright IBM 2022
+//
+// This code is licensed under the Apache License, Version 2.0. You may
+// obtain a copy of this license in the LICENSE.txt file in the root directory
+// of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
+//
+// Any modifications or derivative works of this code must retain this
+// copyright notice, and modified files need to carry a notice indicating
+// that they have been altered from the originals.
+
+#[inline]
+pub fn hex_char_to_bin(c: char) -> &'static str {
+    match c {
+        '0' => "0000",
+        '1' => "0001",
+        '2' => "0010",
+        '3' => "0011",
+        '4' => "0100",
+        '5' => "0101",
+        '6' => "0110",
+        '7' => "0111",
+        '8' => "1000",
+        '9' => "1001",
+        'A' => "1010",
+        'B' => "1011",
+        'C' => "1100",
+        'D' => "1101",
+        'E' => "1110",
+        'F' => "1111",
+        'a' => "1010",
+        'b' => "1011",
+        'c' => "1100",
+        'd' => "1101",
+        'e' => "1110",
+        'f' => "1111",
+        _ => "",
+    }
+}
+
+#[inline]
+pub fn hex_to_bin(hex: &str) -> String {
+    hex[2..].chars().map(hex_char_to_bin).collect()
+}

src/results/marginalization.rs

@@ -10,8 +10,12 @@
 // copyright notice, and modified files need to carry a notice indicating
 // that they have been altered from the originals.
 
+use super::converters::hex_to_bin;
+use crate::getenv_use_multiple_threads;
 use hashbrown::HashMap;
+use num_bigint::BigUint;
 use pyo3::prelude::*;
+use rayon::prelude::*;
 
 fn marginalize<T: std::ops::AddAssign + Copy>(
     counts: HashMap<String, T>,
@@ -68,3 +72,83 @@ pub fn marginal_distribution(
 ) -> HashMap<String, f64> {
     marginalize(counts, indices)
 }
+
+#[inline]
+fn map_memory(
+    hexstring: &str,
+    indices: &Option<Vec<usize>>,
+    clbit_size: usize,
+    return_hex: bool,
+) -> String {
+    let out = match indices {
+        Some(indices) => {
+            let bitstring = hex_to_bin(hexstring);
+            let bit_array = bitstring.as_bytes();
+            indices
+                .iter()
+                .map(|bit| {
+                    let index = clbit_size - *bit - 1;
+                    match bit_array.get(index) {
+                        Some(bit) => *bit as char,
+                        None => '0',
+                    }
+                })
+                .rev()
+                .collect()
+        }
+        None => hex_to_bin(hexstring),
+    };
+    if return_hex {
+        format!("0x{:x}", BigUint::parse_bytes(out.as_bytes(), 2).unwrap())
+    } else {
+        out
+    }
+}
+
+#[pyfunction(return_int = "false", return_hex = "false", parallel_threshold = "50")]
+pub fn marginal_memory(
+    py: Python,
+    memory: Vec<String>,
+    indices: Option<Vec<usize>>,
+    return_int: bool,
+    return_hex: bool,
+    parallel_threshold: usize,
+) -> PyResult<PyObject> {
+    let run_in_parallel = getenv_use_multiple_threads();
+    let first_elem = memory.get(0);
+    if first_elem.is_none() {
+        let res: Vec<String> = Vec::new();
+        return Ok(res.to_object(py));
+    }
+
+    let clbit_size = hex_to_bin(first_elem.unwrap()).len();
+
+    let out_mem: Vec<String> = if memory.len() < parallel_threshold || !run_in_parallel {
+        memory
+            .iter()
+            .map(|x| map_memory(x, &indices, clbit_size, return_hex))
+            .collect()
+    } else {
+        memory
+            .par_iter()
+            .map(|x| map_memory(x, &indices, clbit_size, return_hex))
+            .collect()
+    };
+    if return_int {
+        if out_mem.len() < parallel_threshold || !run_in_parallel {
+            Ok(out_mem
+                .iter()
+                .map(|x| BigUint::parse_bytes(x.as_bytes(), 2).unwrap())
+                .collect::<Vec<BigUint>>()
+                .to_object(py))
+        } else {
+            Ok(out_mem
+                .par_iter()
+                .map(|x| BigUint::parse_bytes(x.as_bytes(), 2).unwrap())
+                .collect::<Vec<BigUint>>()
+                .to_object(py))
+        }
+    } else {
+        Ok(out_mem.to_object(py))
+    }
+}

src/results/mod.rs

@@ -10,6 +10,7 @@
 // copyright notice, and modified files need to carry a notice indicating
 // that they have been altered from the originals.
 
+pub mod converters;
 pub mod marginalization;
 
 use pyo3::prelude::*;
@@ -19,5 +20,6 @@ use pyo3::wrap_pyfunction;
 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))?;
     Ok(())
 }

test/python/result/test_memory_marginalization.py

@@ -0,0 +1,55 @@
+# This code is part of Qiskit.
+#
+# (C) Copyright IBM 2017, 2019.
+#
+# This code is licensed under the Apache License, Version 2.0. You may
+# obtain a copy of this license in the LICENSE.txt file in the root directory
+# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# Any modifications or derivative works of this code must retain this
+# copyright notice, and modified files need to carry a notice indicating
+# that they have been altered from the originals.
+
+"""Test marginal_memory() function."""
+
+from qiskit.test import QiskitTestCase
+from qiskit.result import marginal_memory
+
+
+class TestMarginalMemory(QiskitTestCase):
+    """Result operations methods."""
+
+    def test_marginalize_memory(self):
+        """Test that memory marginalizes correctly."""
+        memory = [hex(ii) for ii in range(8)]
+        res = marginal_memory(memory, indices=[0])
+        self.assertEqual(res, [bin(ii % 2)[2:] for ii in range(8)])
+
+    def test_marginalize_memory_int(self):
+        """Test that memory marginalizes correctly int output."""
+        memory = [hex(ii) for ii in range(8)]
+        res = marginal_memory(memory, indices=[0], int_return=True)
+        self.assertEqual(res, [ii % 2 for ii in range(8)])
+
+    def test_marginalize_memory_hex(self):
+        """Test that memory marginalizes correctly hex output."""
+        memory = [hex(ii) for ii in range(8)]
+        res = marginal_memory(memory, indices=[0], hex_return=True)
+        self.assertEqual(res, [hex(ii % 2) for ii in range(8)])
+
+    def test_marginal_counts_result_memory_indices_None(self):
+        """Test that a memory marginalizes correctly with indices=None."""
+        memory = [hex(ii) for ii in range(8)]
+        res = marginal_memory(memory, hex_return=True)
+        self.assertEqual(res, memory)
+
+    def test_marginalize_memory_in_parallel(self):
+        """Test that memory marginalizes correctly multithreaded."""
+        memory = [hex(ii) for ii in range(15)]
+        res = marginal_memory(memory, indices=[0], parallel_threshold=1)
+        self.assertEqual(res, [bin(ii % 2)[2:] for ii in range(15)])
+
+    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)