Make marginal_counts act on Result memory

qiskit/result/utils.py

@@ -12,28 +12,38 @@
 
 """Utility functions for working with Results."""
 
+from typing import List, Union, Optional, Dict
 from collections import Counter
 from copy import deepcopy
 
 from qiskit.exceptions import QiskitError
 from qiskit.result.result import Result
-from qiskit.result.postprocess import _bin_to_hex
+from qiskit.result.postprocess import _bin_to_hex, _hex_to_bin
 
 
-def marginal_counts(result, indices=None, inplace=False, format_marginal=False):
+def marginal_counts(
+    result: Union[dict, Result],
+    indices: Optional[List[int]] = None,
+    inplace: bool = False,
+    format_marginal: bool = False,
+    marginalize_memory: Optional[bool] = True,
+) -> Union[Dict[str, int], Result]:
     """Marginalize counts from an experiment over some indices of interest.
 
     Args:
-        result (dict or Result): result to be marginalized
+        result: result to be marginalized
             (a Result object or a dict(str, int) of counts).
-        indices (list(int) or None): The bit positions of interest
+        indices: The bit positions of interest
             to marginalize over. If ``None`` (default), do not marginalize at all.
-        inplace (bool): Default: False. Operates on the original Result
+        inplace: Default: False. Operates on the original Result
             argument if True, leading to loss of original Job Result.
             It has no effect if ``result`` is a dict.
-        format_marginal (bool): Default: False. If True, takes the output of
+        format_marginal: Default: False. If True, takes the output of
             marginalize and formats it with placeholders between cregs and
             for non-indices.
+        marginalize_memory: If True, then also marginalize the memory field (if present).
+            If False, remove the memory field from the result.
+            If None, leave the memory field as is.
 
     Returns:
         Result or dict(str, int): A Result object or a dictionary with
@@ -53,9 +63,27 @@ def marginal_counts(result, indices=None, inplace=False, format_marginal=False):
             for k, v in new_counts.items():
                 new_counts_hex[_bin_to_hex(k)] = v
             experiment_result.data.counts = new_counts_hex
-            experiment_result.header.memory_slots = len(indices)
-            csize = experiment_result.header.creg_sizes
-            experiment_result.header.creg_sizes = _adjust_creg_sizes(csize, indices)
+
+            if indices is not None:
+                experiment_result.header.memory_slots = len(indices)
+                csize = experiment_result.header.creg_sizes
+                experiment_result.header.creg_sizes = _adjust_creg_sizes(csize, indices)
+
+            if getattr(experiment_result.data, "memory", None) is not None and indices is not None:
+                if marginalize_memory is False:
+                    delattr(experiment_result.data, "memory")
+                elif marginalize_memory is None:
+                    pass  # leave as is
+                else:
+                    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)])
+                        for s in bit_strings
+                    ]
+                    experiment_result.data.memory = [_bin_to_hex(s) for s in marginal_bit_strings]
         return result
     else:
         marg_counts = _marginalize(result, indices)

releasenotes/notes/marginal_counts_act_on_memory-0a9b58d0b95046dd.yaml

@@ -0,0 +1,22 @@
+---
+upgrade:
+  - |
+    The behaviour of :func:`~.marginal_counts` method on the `memory` field of experiment data has changed.
+    Previously, the `memory` field was not marginalized. In this release the `memory` field is 
+    marginalized by default. The old behaviour can be restored by setting `marginalize_memory=None`
+    as an argument to :func:`~.marginal_counts`.
+
+features:
+  - |
+    Update the :func:`~.marginal_counts` method to marginalize the memory field of the data.
+    For example, if `result`  is a qiskit :class:`~.Result` object obtained from a 4-qubit measurement we can
+    marginalize onto the first qubit with as::
+
+        print(result.results[0].data.memory)
+        marginal_result = marginal_counts(result, [0])
+        print(marginal_result.results[0].data.memory)
+
+    The output is as::
+
+        ['0x0', '0x1', '0x2', '0x3', '0x4', '0x5', '0x6', '0x7']
+        ['0x0', '0x1', '0x0', '0x1', '0x0', '0x1', '0x0', '0x1']

test/python/result/test_result.py

@@ -19,6 +19,7 @@
 from qiskit.result import Result
 from qiskit.qobj import QobjExperimentHeader
 from qiskit.test import QiskitTestCase
+from qiskit.exceptions import QiskitError
 
 
 class TestResultOperations(QiskitTestCase):
@@ -35,6 +36,19 @@ def setUp(self):
 
         super().setUp()
 
+    def generate_qiskit_result(self):
+        """Generate standard Result for testing"""
+        memory = [hex(ii) for ii in range(8)]
+        counts = {m: 1 for m in memory}
+        data_1 = models.ExperimentResultData(counts=counts, memory=memory)
+        exp_result_header_1 = QobjExperimentHeader(creg_sizes=[["c0", 4]], memory_slots=4)
+        exp_result_1 = models.ExperimentResult(
+            shots=8, success=True, data=data_1, header=exp_result_header_1
+        )
+
+        result = Result(results=[exp_result_1], **self.base_result_args)
+        return result
+
     def test_counts_no_header(self):
         """Test that counts are extracted properly without header."""
         raw_counts = {"0x0": 4, "0x2": 10}
@@ -170,14 +184,14 @@ def test_marginal_counts(self):
     def test_marginal_counts_result(self):
         """Test that a Result object containing counts marginalizes correctly."""
         raw_counts_1 = {"0x0": 4, "0x1": 7, "0x2": 10, "0x6": 5, "0x9": 11, "0xD": 9, "0xE": 8}
-        data_1 = models.ExperimentResultData(counts=dict(**raw_counts_1))
+        data_1 = models.ExperimentResultData(counts=raw_counts_1)
         exp_result_header_1 = QobjExperimentHeader(creg_sizes=[["c0", 4]], memory_slots=4)
         exp_result_1 = models.ExperimentResult(
             shots=54, success=True, data=data_1, header=exp_result_header_1
         )
 
         raw_counts_2 = {"0x2": 5, "0x3": 8}
-        data_2 = models.ExperimentResultData(counts=dict(**raw_counts_2))
+        data_2 = models.ExperimentResultData(counts=raw_counts_2)
         exp_result_header_2 = QobjExperimentHeader(creg_sizes=[["c0", 2]], memory_slots=2)
         exp_result_2 = models.ExperimentResult(
             shots=13, success=True, data=data_2, header=exp_result_header_2
@@ -187,9 +201,84 @@ def test_marginal_counts_result(self):
 
         expected_marginal_counts_1 = {"00": 4, "01": 27, "10": 23}
         expected_marginal_counts_2 = {"0": 5, "1": 8}
+        expected_marginal_counts_none = {
+            "0000": 4,
+            "0001": 7,
+            "0010": 10,
+            "0110": 5,
+            "1001": 11,
+            "1101": 9,
+            "1110": 8,
+        }
 
         self.assertEqual(marginal_counts(result, [0, 1]).get_counts(0), expected_marginal_counts_1)
         self.assertEqual(marginal_counts(result, [0]).get_counts(1), expected_marginal_counts_2)
+        self.assertEqual(marginal_counts(result, None).get_counts(0), expected_marginal_counts_none)
+
+    def test_marginal_counts_result_memory(self):
+        """Test that a Result object containing memory marginalizes correctly."""
+        result = self.generate_qiskit_result()
+        marginal_result = marginal_counts(result, indices=[0])
+        marginal_memory = marginal_result.results[0].data.memory
+        self.assertEqual(marginal_memory, [hex(ii % 2) for ii in range(8)])
+
+    def test_marginal_counts_result_memory_indices_None(self):
+        """Test that a Result object containing memory marginalizes correctly."""
+        result = self.generate_qiskit_result()
+        memory = "should not be touched"
+        result.results[0].data.memory = memory
+        marginal_result = marginal_counts(result, indices=None)
+        marginal_memory = marginal_result.results[0].data.memory
+        self.assertEqual(marginal_memory, memory)
+
+    def test_marginal_counts_result_invalid_indices(self):
+        """Test that a Result object containing memory marginalizes correctly inplace."""
+
+        result = self.generate_qiskit_result()
+        with self.assertRaises(QiskitError):
+            _ = marginal_counts(result, indices=[0, 1, 100], inplace=True)
+
+    def test_marginal_counts_result_marginalize_memory(self):
+        """Test that a Result object containing memory marginalizes correctly inplace."""
+
+        result = self.generate_qiskit_result()
+        marginal_result = marginal_counts(
+            result, indices=[0], inplace=True, marginalize_memory=False
+        )
+        self.assertFalse(hasattr(marginal_result.results[0].data, "memory"))
+        result = self.generate_qiskit_result()
+        marginal_result = marginal_counts(
+            result, indices=[0], inplace=True, marginalize_memory=None
+        )
+        self.assertTrue(hasattr(marginal_result.results[0].data, "memory"))
+        result = self.generate_qiskit_result()
+        marginal_result = marginal_counts(
+            result, indices=[0], inplace=True, marginalize_memory=True
+        )
+        self.assertTrue(hasattr(marginal_result.results[0].data, "memory"))
+
+        result = self.generate_qiskit_result()
+        marginal_result = marginal_counts(
+            result, indices=[0], inplace=False, marginalize_memory=False
+        )
+        self.assertFalse(hasattr(marginal_result.results[0].data, "memory"))
+        marginal_result = marginal_counts(
+            result, indices=[0], inplace=False, marginalize_memory=None
+        )
+        self.assertTrue(hasattr(marginal_result.results[0].data, "memory"))
+        marginal_result = marginal_counts(
+            result, indices=[0], inplace=False, marginalize_memory=True
+        )
+        self.assertTrue(hasattr(marginal_result.results[0].data, "memory"))
+
+    def test_marginal_counts_result_inplace(self):
+        """Test that a Result object containing memory marginalizes correctly inplace."""
+        result = self.generate_qiskit_result()
+
+        marginal_result = marginal_counts(result, indices=[0], inplace=True)
+        self.assertEqual(id(result), id(marginal_result))
+        marginal_memory = marginal_result.results[0].data.memory
+        self.assertEqual(marginal_memory, [hex(ii % 2) for ii in range(8)])
 
     def test_marginal_counts_result_creg_sizes(self):
         """Test that marginal_counts with Result input properly changes creg_sizes."""