Repository: Qiskit/qiskit-terra (python)

Issue:
Result.utils.marginal_counts returns inconsistent creg structure for multiple cregs
<!-- ⚠️ If you do not respect this template, your issue will be closed -->
<!-- ⚠️ Make sure to browse the opened and closed issues -->

### Information

- **Qiskit Terra version**: master @ d26f273d
- **Python version**: 3.6
- **Operating system**: osx

### Steps to reproduce the problem

```python
from qiskit import *

qr_data = QuantumRegister(3, 'qr_data')
qr_anc = QuantumRegister(2, 'qr_anc')
cr_data = ClassicalRegister(3, 'cr_data')
cr_anc = ClassicalRegister(2, 'cr_anc')

qc = QuantumCircuit(qr_data, qr_anc, cr_data, cr_anc)
qc.x(qr_data)
qc.measure(qr_data, cr_data)
qc.measure(qr_anc, cr_anc)
qc.draw('mpl')
```
![image](https://user-images.githubusercontent.com/766693/101812675-aaaa2380-3ae9-11eb-9d8a-cfe8c0afc6d1.png)

```python
result = execute(qc, Aer.get_backend('qasm_simulator')).result()
result.get_counts()
```
```
{'00 111': 1024}
```

```python
from qiskit.result.utils import marginal_counts
marginal_counts(result, [0, 1, 2]).get_counts()
```
```
{'11 1': 1024}
```
### What is the expected behavior?



### Suggested solutions




Pull request description:
Fix marginal_counts incorrect creg_sizes with Result object and add format_marginal kwarg
<!--
⚠️ If you do not respect this template, your pull request will be closed.
⚠️ Your pull request title should be short detailed and understandable for all.
⚠️ Also, please add a release note file using reno if the change needs to be
  documented in the release notes.
⚠️ If your pull request fixes an open issue, please link to the issue.

- [ ] I have added the tests to cover my changes.
- [ ] I have updated the documentation accordingly.
- [ ] I have read the CONTRIBUTING document.
-->

### Summary

Fixes #5424

### Details and comments

When a Result object was passed to `result.utils.marginal_counts`, the `memory_slots` value was decremented to match the number of indices passed to `marginal_counts`. The `creg_sizes` though were not decremented, which resulted in incorrect formatting if `get_counts` was called on the result returned from `marginal_counts`. This PR properly decrements the `creg_sizes` to match the indices.

In addition there is a problem with `marginal_counts` in that information as to which indices are used is lost. For example, if there were 2 cregs of size 3 and 2 respectively, this count, `{11 111: 1024}` might be returned from `get_counts`. If `marginal_counts` is called with indices `[0, 2, 4]`, the return value is `{111: 1024}`. This value would be the same for `[0, 1, 2]`, `[1, 3, 4]`, etc. 

An optional `kwarg`, `format_marginal`, has been added to `marginal_counts `which defaults to False. If set to True, underscores are inserted as placeholders in the non-indices positions. For the above example, the output would be `{1_ 1_1: 1024}`. For `[0, 1, 2]`, this would be `{__ 111: 1024}`, and for `[1, 3, 4]` it would be `{11 _1_: 1024}`, allowing differentiation between the different indices inputs. Since these are strings, they will display this way when visualizations like `plot_histogram` are called.

Finally code from @aeddins-ibm in #5241 was used to replace complex and slow code in the `marginalize` function. In basic tests with 20 cregs, the new code was 7 times faster than the old and produced the same output in all the tests.

Patch:
diff --git a/qiskit/result/counts.py b/qiskit/result/counts.py
index a5956c45a..707c6574c 100644
--- a/qiskit/result/counts.py
+++ b/qiskit/result/counts.py
@@ -97,7 +97,7 @@ class Counts(dict):
                                     'Counts objects with dit strings do not '
                                     'currently support dit string formatting parameters '
                                     'creg_sizes or memory_slots')
-                            int_key = int(bitstring.replace(" ", ""), 2)
+                            int_key = self._remove_space_underscore(bitstring)
                             int_dict[int_key] = value
                             hex_dict[hex(int_key)] = value
                         self.hex_raw = hex_dict
@@ -155,7 +155,7 @@ class Counts(dict):
                     raise exceptions.QiskitError(
                         'Counts objects with dit strings do not '
                         'currently support conversion to hexadecimal')
-                int_key = int(bitstring.replace(" ", ""), 2)
+                int_key = self._remove_space_underscore(bitstring)
                 out_dict[hex(int_key)] = value
             return out_dict
 
@@ -176,6 +176,11 @@ class Counts(dict):
                     raise exceptions.QiskitError(
                         'Counts objects with dit strings do not '
                         'currently support conversion to integer')
-                int_key = int(bitstring.replace(" ", ""), 2)
+                int_key = self._remove_space_underscore(bitstring)
                 out_dict[int_key] = value
             return out_dict
+
+    @staticmethod
+    def _remove_space_underscore(bitstring):
+        """Removes all spaces and underscores from bitstring"""
+        return int(bitstring.replace(" ", "").replace("_", ""), 2)
diff --git a/qiskit/result/utils.py b/qiskit/result/utils.py
index b50528370..b121dd265 100644
--- a/qiskit/result/utils.py
+++ b/qiskit/result/utils.py
@@ -12,8 +12,7 @@
 
 """Utility functions for working with Results."""
 
-from functools import reduce
-from re import match
+from collections import Counter
 from copy import deepcopy
 
 from qiskit.exceptions import QiskitError
@@ -21,22 +20,25 @@ from qiskit.result.result import Result
 from qiskit.result.postprocess import _bin_to_hex
 
 
-def marginal_counts(result, indices=None, inplace=False):
+def marginal_counts(result, indices=None, inplace=False, format_marginal=False):
     """Marginalize counts from an experiment over some indices of interest.
 
     Args:
         result (dict or Result): result to be marginalized
-            (a Result object or a dict of counts).
+            (a Result object or a dict(str, int) of counts).
         indices (list(int) or None): 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
             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
+            marginalize and formats it with placeholders between cregs and
+            for non-indices.
 
     Returns:
-        Result or dict[str, int]: a dictionary with the observed counts,
-        marginalized to only account for frequency of observations
-        of bits of interest.
+        Result or dict(str, int): A Result object or a dictionary with
+            the observed counts, marginalized to only account for frequency
+            of observations of bits of interest.
 
     Raises:
         QiskitError: in case of invalid indices to marginalize over.
@@ -52,63 +54,89 @@ def marginal_counts(result, indices=None, inplace=False):
                 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)
+        return result
     else:
-        result = _marginalize(result, indices)
-
-    return result
-
-
-def count_keys(num_clbits):
-    """Return ordered count keys."""
-    return [bin(j)[2:].zfill(num_clbits) for j in range(2 ** num_clbits)]
+        marg_counts = _marginalize(result, indices)
+        if format_marginal and indices is not None:
+            marg_counts = _format_marginal(result, marg_counts, indices)
+        return marg_counts
+
+
+def _adjust_creg_sizes(creg_sizes, indices):
+    """Helper to reduce creg_sizes to match indices"""
+
+    # Zero out creg_sizes list
+    new_creg_sizes = [[creg[0], 0] for creg in creg_sizes]
+    indices_sort = sorted(indices)
+
+    # Get creg num values and then convert to the cumulative last index per creg.
+    # e.g. [2, 1, 3] => [1, 2, 5]
+    creg_nums = [x for _, x in creg_sizes]
+    creg_limits = [sum(creg_nums[0:x:1])-1 for x in range(0, len(creg_nums)+1)][1:]
+
+    # Now iterate over indices and find which creg that index is in.
+    # When found increment the creg size
+    creg_idx = 0
+    for ind in indices_sort:
+        for idx in range(creg_idx, len(creg_limits)):
+            if ind <= creg_limits[idx]:
+                creg_idx = idx
+                new_creg_sizes[idx][1] += 1
+                break
+    # Throw away any cregs with 0 size
+    new_creg_sizes = [creg for creg in new_creg_sizes if creg[1] != 0]
+    return new_creg_sizes
 
 
 def _marginalize(counts, indices=None):
-    # Extract total number of clbits from first count key
-    # We trim the whitespace separating classical registers
-    # and count the number of digits
+    """Get the marginal counts for the given set of indices"""
     num_clbits = len(next(iter(counts)).replace(' ', ''))
 
-    # Check if we do not need to marginalize. In this case we just trim
-    # whitespace from count keys
+    # Check if we do not need to marginalize and if so, trim
+    # whitespace and '_' and return
     if (indices is None) or set(range(num_clbits)) == set(indices):
         ret = {}
         for key, val in counts.items():
-            key = key.replace(' ', '')
+            key = _remove_space_underscore(key)
             ret[key] = val
         return ret
 
-    if not set(indices).issubset(set(range(num_clbits))):
+    if not indices or not set(indices).issubset(set(range(num_clbits))):
         raise QiskitError('indices must be in range [0, {}].'.format(num_clbits-1))
 
     # Sort the indices to keep in decending order
     # Since bitstrings have qubit-0 as least significant bit
     indices = sorted(indices, reverse=True)
 
-    # Generate bitstring keys for indices to keep
-    meas_keys = count_keys(len(indices))
-
-    # Get regex match strings for suming outcomes of other qubits
-    rgx = []
-    for key in meas_keys:
-        def _helper(x, y):
-            if y in indices:
-                return key[indices.index(y)] + x
-            return '\\d' + x
-        rgx.append(reduce(_helper, range(num_clbits), ''))
-
     # Build the return list
-    meas_counts = []
-    for m in rgx:
-        c = 0
-        for key, val in counts.items():
-            if match(m, key.replace(' ', '')):
-                c += val
-        meas_counts.append(c)
-
-    # Return as counts dict on desired indices only
-    ret = {}
-    for key, val in zip(meas_keys, meas_counts):
-        if val != 0:
-            ret[key] = val
-    return ret
+    new_counts = Counter()
+    for key, val in counts.items():
+        new_key = ''.join([_remove_space_underscore(key)[-idx-1] for idx in indices])
+        new_counts[new_key] += val
+    return dict(new_counts)
+
+
+def _format_marginal(counts, marg_counts, indices):
+    """Take the output of marginalize and add placeholders for
+    multiple cregs and non-indices."""
+    format_counts = {}
+    counts_template = next(iter(counts))
+    counts_len = len(counts_template.replace(' ', ''))
+    indices_rev = sorted(indices, reverse=True)
+
+    for count in marg_counts:
+        index_dict = dict(zip(indices_rev, count))
+        count_bits = ''.join([index_dict[index] if index in index_dict else '_'
+                              for index in range(counts_len)])[::-1]
+        for index, bit in enumerate(counts_template):
+            if bit == ' ':
+                count_bits = count_bits[:index] + ' ' + count_bits[index:]
+        format_counts[count_bits] = marg_counts[count]
+    return format_counts
+
+
+def _remove_space_underscore(bitstring):
+    """Removes all spaces and underscores from bitstring"""
+    return bitstring.replace(" ", "").replace("_", "")
diff --git a/releasenotes/notes/marginal_counts_format-dc20b15d6c04ec94.yaml b/releasenotes/notes/marginal_counts_format-dc20b15d6c04ec94.yaml
new file mode 100644
index 000000000..35c701f69
--- /dev/null
+++ b/releasenotes/notes/marginal_counts_format-dc20b15d6c04ec94.yaml
@@ -0,0 +1,25 @@
+---
+features:
+  - |
+    A new kwarg, ``format_marginal`` has been added to
+    :func:`~qiskit.result.utils.marginal_counts` which formats the counts
+    output according to the cregs in the circuit. The missing indices are
+    now represented with a ``_``. For example, this code produces
+    {'0_0_1': 1024}
+
+    .. jupyter-execute::
+
+        from qiskit import QuantumCircuit, execute, BasicAer, result
+        from qiskit.result.utils import marginal_counts
+        qc = QuantumCircuit(5, 5)
+        qc.x(0)
+        qc.measure(0, 0)
+
+        result = execute(qc, BasicAer.get_backend('qasm_simulator')).result()
+        print(marginal_counts(result.get_counts(), [0, 2, 4], format_marginal=True))
+
+fixes:
+  - |
+    This fixes an issue when calling :meth:`~qiskit.result.get_counts` using the
+    result returned from :func:`~qiskit.result.utils.marginal_counts` that would
+    produce an improperly formatted ``counts`` dict with certain inputs.


Test patch:
diff --git a/test/python/result/test_result.py b/test/python/result/test_result.py
index a4da93493..f95899cd9 100644
--- a/test/python/result/test_result.py
+++ b/test/python/result/test_result.py
@@ -174,6 +174,45 @@ class TestResultOperations(QiskitTestCase):
         self.assertEqual(marginal_counts(result, [0]).get_counts(1),
                          expected_marginal_counts_2)
 
+    def test_marginal_counts_result_creg_sizes(self):
+        """Test that marginal_counts with Result input properly changes creg_sizes."""
+        raw_counts = {'0x0': 4, '0x1': 7, '0x2': 10, '0x6': 5, '0x9': 11, '0xD': 9, '0xE': 8}
+        data = models.ExperimentResultData(counts=dict(**raw_counts))
+        exp_result_header = QobjExperimentHeader(creg_sizes=[['c0', 1], ['c1', 3]],
+                                                 memory_slots=4)
+        exp_result = models.ExperimentResult(shots=54, success=True, data=data,
+                                             header=exp_result_header)
+
+        result = Result(results=[exp_result], **self.base_result_args)
+
+        expected_marginal_counts = {'0 0': 14, '0 1': 18, '1 0': 13, '1 1': 9}
+        expected_creg_sizes = [['c0', 1], ['c1', 1]]
+        expected_memory_slots = 2
+        marginal_counts_result = marginal_counts(result, [0, 2])
+        self.assertEqual(marginal_counts_result.results[0].header.creg_sizes,
+                         expected_creg_sizes)
+        self.assertEqual(marginal_counts_result.results[0].header.memory_slots,
+                         expected_memory_slots)
+        self.assertEqual(marginal_counts_result.get_counts(0),
+                         expected_marginal_counts)
+
+    def test_marginal_counts_result_format(self):
+        """Test that marginal_counts with format_marginal true properly formats output."""
+        raw_counts_1 = {'0x0': 4, '0x1': 7, '0x2': 10, '0x6': 5, '0x9': 11, '0xD': 9, '0x12': 8}
+        data_1 = models.ExperimentResultData(counts=dict(**raw_counts_1))
+        exp_result_header_1 = QobjExperimentHeader(creg_sizes=[['c0', 2], ['c1', 3]],
+                                                   memory_slots=5)
+        exp_result_1 = models.ExperimentResult(shots=54, success=True, data=data_1,
+                                               header=exp_result_header_1)
+
+        result = Result(results=[exp_result_1], **self.base_result_args)
+
+        expected_marginal_counts_1 = {'0_0 _0': 14, '0_0 _1': 18,
+                                      '0_1 _0': 5, '0_1 _1': 9, '1_0 _0': 8}
+        marginal_counts_result = marginal_counts(result.get_counts(), [0, 2, 4],
+                                                 format_marginal=True)
+        self.assertEqual(marginal_counts_result, expected_marginal_counts_1)
+
     def test_marginal_counts_inplace_true(self):
         """Test marginal_counts(Result, inplace = True)
         """