Add dedicated functions for memory marginalization

Cargo.toml

@@ -18,10 +18,12 @@ rand_distr = "0.4.3"
 indexmap = "1.9.1"
 ahash = "0.7.6"
 num-complex = "0.4"
+num-bigint = "0.4"
+lazy_static = "1.4.0"
 
 [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,13 +18,15 @@
 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
 from qiskit.result.distributions.probability import ProbDistribution
 from qiskit.result.distributions.quasi import QuasiDistribution
 
-from qiskit.result.postprocess import _bin_to_hex, _hex_to_bin
+from qiskit.result.postprocess import _bin_to_hex
 
 # pylint: disable=import-error, no-name-in-module
 from qiskit._accelerate import results as results_rs
@@ -88,12 +90,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)
@@ -128,14 +127,85 @@ def _adjust_creg_sizes(creg_sizes, indices):
     return new_creg_sizes
 
 
+def marginal_memory(
+    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 shot 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 either a list of hexadecimal strings to be marginalized
+            representing measure level 2 memory or a numpy array representing level 0 measurement
+            memory (single or avg) 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.
+            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. This option only has an
+            effect with memory level 2.
+        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. This option only has an effect with memory level 2.
+        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, the function will run in multiple
+            threads. By default this is set to 1000.
+
+    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")
+
+    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,
+        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:

releasenotes/notes/marginal-memory-29d9d6586ae78590.yaml

@@ -0,0 +1,9 @@
+---
+features:
+  - |
+    Added a new function :func:`~.marginal_memory` which is used to marginalize
+    shot memory arrays. Provided with the shot memory array and the indices
+    of interest the function will return a maginized shot memory array. This
+    function differs from the memory support in the :func:`~.marginal_counts`
+    method which only works on the ``memory`` field in a :class:`~.Results`
+    object.

src/lib.rs

@@ -10,6 +10,9 @@
 // copyright notice, and modified files need to carry a notice indicating
 // that they have been altered from the originals.
 
+#[macro_use]
+extern crate lazy_static;
+
 use std::env;
 
 use pyo3::prelude::*;

src/results/converters.rs

@@ -0,0 +1,48 @@
+// 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.
+
+lazy_static! {
+    static ref HEX_TO_BIN_LUT: [&'static str; 256] = {
+        let mut lookup = [""; 256];
+        lookup[b'0' as usize] = "0000";
+        lookup[b'1' as usize] = "0001";
+        lookup[b'2' as usize] = "0010";
+        lookup[b'3' as usize] = "0011";
+        lookup[b'4' as usize] = "0100";
+        lookup[b'5' as usize] = "0101";
+        lookup[b'6' as usize] = "0110";
+        lookup[b'7' as usize] = "0111";
+        lookup[b'8' as usize] = "1000";
+        lookup[b'9' as usize] = "1001";
+        lookup[b'A' as usize] = "1010";
+        lookup[b'B' as usize] = "1011";
+        lookup[b'C' as usize] = "1100";
+        lookup[b'D' as usize] = "1101";
+        lookup[b'E' as usize] = "1110";
+        lookup[b'F' as usize] = "1111";
+        lookup[b'a' as usize] = "1010";
+        lookup[b'b' as usize] = "1011";
+        lookup[b'c' as usize] = "1100";
+        lookup[b'd' as usize] = "1101";
+        lookup[b'e' as usize] = "1110";
+        lookup[b'f' as usize] = "1111";
+        lookup
+    };
+}
+
+#[inline]
+pub fn hex_to_bin(hex: &str) -> String {
+    hex[2..]
+        .chars()
+        .map(|c| HEX_TO_BIN_LUT[c as usize])
+        .collect()
+}