Repository: PennyLaneAI/pennylane

Issue:
Implement `CountsMP.process_counts`
# Context

Currently, our sample-based measurement processes that inherit from `SampleMeasurement` implement a [process_samples](https://github.com/PennyLaneAI/pennylane/blob/a05ae286c5925e0c881099cea751193d368779fc/pennylane/measurements/measurements.py#L535) method. This method is capable of turning any samples (array of int64) into the required measurement result. This method is currently used internally by `default.qubit`.

```
@qml.qnode(qml.device('default.qubit'))
def circuit():
    qml.Hadamard(0)
    return qml.sample(wires=(0,1))

samples = circuit(shots=100)
qml.counts(wires=0, all_outcomes=True).process_samples(samples, wire_order=qml.wires.Wires((0,1)))
```

But external devices and plugins may want to store sampling information in the form of a counts dictionary instead, like `{'00': 45, '01': 55}`. This representation is much more condensed and memory efficient. Because of this, we have recently added a `SampleMeasurement.process_counts` method, and implemented it for `ProbabilityMP`.

# Implementation Details

The task is to implement [CountsMP.process_counts(counts, wire_order)](https://github.com/PennyLaneAI/pennylane/blob/3407a9e3fd76f0b112ef22e14ef8a4ba3d5a339f/pennylane/measurements/counts.py#L159) in such a way that a counts dictionary can be processed in such a way that agrees with the specification for the measurement process.

This includes:

* Adding in all outcomes if `all_outcomes=True`
* Removing entries with a value of `0` if `all_outcomes=False`
* Tracing over wires

```pycon
>>> counts = {'000': 100, '100': 100 }
>>> wire_order = qml.wires.Wires((0,1,2))
>>> qml.counts(wires=(0,1), all_outcomes=True).process_counts(counts, wire_order)
{'00': 100, '01': 0, '10': 100, '11': 0}
>>> qml.counts(wires=(0,2), all_outcomes=False).process_counts(counts, wire_order)
{'00': 100, '10': 10}
>>> counts2 = {'00': 0, '10': 0, '01': 0, '11': 50}
>>> qml.counts(wires=0, all_outcomes=False).process_counts(counts2, wire_order=qml.wires.Wires((0,1)))
{'1': 50}
```

Note that the counts are assumed to be taken in the basis for the provided observable already, with any diagonalizing gates already applied.

**Warning:** This processing should leave the input dictionary unchanged.

# Requirements

A completed PR should:

* Add the `process_counts` method to `CountsMP`
* Add unit tests to [pennylane/tests/measurements/test_counts.py](https://github.com/PennyLaneAI/pennylane/blob/master/tests/measurements/test_counts.py)
* Pass black and pylint checks
* Add a changelog entry to [pennylane/releases/changelog-dev.md](https://github.com/PennyLaneAI/pennylane/blob/master/doc/releases/changelog-dev.md) under Improvements. Don't forget to add your name to the list of contributors at the bottom :)

Discussion:
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5256?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
Attention: Patch coverage is `97.56098%` with `1 lines` in your changes are missing coverage. Please review.
> Project coverage is 99.67%. Comparing base [(`2ae851b`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/2ae851b0dc3676a13f36d4c78c244e1821c7cd82?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`f4ab5d5`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5256?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).

| [Files](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5256?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Patch % | Lines |
|---|---|---|
| [pennylane/measurements/counts.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5256?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21lYXN1cmVtZW50cy9jb3VudHMucHk=) | 95.65% | [1 Missing :warning: ](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5256?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master    #5256      +/-   ##
==========================================
- Coverage   99.68%   99.67%   -0.01%     
==========================================
  Files         399      399              
  Lines       36871    36611     -260     
==========================================
- Hits        36754    36492     -262     
- Misses        117      119       +2     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5256?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).

trbromley: Thanks a lot for opening this @Tarun-Kumar07! We'll give this a review soon and also help with the CI issue you mentioned. Sorry if things are a little delayed this week, we're busy preparing for the PennyLane 0.35 release next week :rocket:
Tarun-Kumar07: @trbromley Thank you for your response.
If there's anything I can do to help with the PennyLane 0.35 release preparations, please let me know. 
I'd be more than happy to assist in any way I can! 🚀
timmysilv: hi @Tarun-Kumar07 , thanks for bearing with us! I've pushed a fix for the strange "Documentation Check" issue to the master branch. If you update your fork with the latest master, it should disappear 🤞 give it a try, I'll keep an eye out to make sure it's resolved here.
Tarun-Kumar07: Hey @albi3ro,

I noticed that all classes which inherit from SampleMeasurement implement process_counts except SampleMP. Should I add an implementation for that?
Tarun-Kumar07: Hi @albi3ro and @trbromley,

Thanks for reviewing my PR. 
It's my first open-source contribution, and I'm eager to contribute even more 🎉.

Appreciate it!
Tarun-Kumar07: @albi3ro, regarding adding `SampleMP.process_counts`. 
Should I create a PR directly or make an issue about it? 
What's the preferred process?
albi3ro: > @albi3ro, regarding adding `SampleMP.process_counts`. Should I create a PR directly or make an issue about it? What's the preferred process?

I'll go ahead and create an issue right.

Since this PR now has two approvals 🎉 , I'll go ahead and enable auto-merge.  Congrats on the PR!

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 2a70356a3..445064a36 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -114,6 +114,9 @@
 * `qml.transforms.split_non_commuting` will now work with single-term operator arithmetic.
   [(#5314)](https://github.com/PennyLaneAI/pennylane/pull/5314)
 
+* Implemented the method `process_counts` in `ExpectationMP`, `VarianceMP`, and `CountsMP`.
+  [(#5256)](https://github.com/PennyLaneAI/pennylane/pull/5256)
+
 <h3>Breaking changes 💔</h3>
 
 * The private functions ``_pauli_mult``, ``_binary_matrix`` and ``_get_pauli_map`` from the ``pauli`` module have been removed. The same functionality can be achieved using newer features in the ``pauli`` module.
@@ -184,6 +187,7 @@
 
 This release contains contributions from (in alphabetical order):
 
+Tarun Kumar Allamsetty,
 Guillermo Alonso,
 Mikhail Andrenkov,
 Utkarsh Azad,
@@ -195,4 +199,4 @@ Soran Jahangiri,
 Korbinian Kottmann,
 Christina Lee,
 Mudit Pandey,
-Matthew Silverman.
+Matthew Silverman.
\ No newline at end of file
diff --git a/pennylane/measurements/counts.py b/pennylane/measurements/counts.py
index 554ce3e9e..f5a0acf36 100644
--- a/pennylane/measurements/counts.py
+++ b/pennylane/measurements/counts.py
@@ -341,3 +341,62 @@ class CountsMP(SampleMeasurement):
                 outcome_dict[state] = count
 
         return outcome_dicts if batched else outcome_dicts[0]
+
+    # pylint: disable=redefined-outer-name
+    def process_counts(self, counts: dict, wire_order: Wires) -> dict:
+        mapped_counts = self._map_counts(counts, wire_order)
+        if self.all_outcomes:
+            self._include_all_outcomes(mapped_counts)
+        else:
+            _remove_unobserved_outcomes(mapped_counts)
+        return mapped_counts
+
+    def _map_counts(self, counts_to_map: dict, wire_order: Wires) -> dict:
+        """
+        Args:
+            counts_to_map: Dictionary where key is binary representation of the outcome and value is its count
+            wire_order: Order of wires to which counts_to_map should be ordered in
+
+        Returns:
+            Dictionary where counts_to_map has been reordered according to wire_order
+        """
+        wire_map = dict(zip(wire_order, range(len(wire_order))))
+        mapped_wires = [wire_map[w] for w in self.wires]
+
+        mapped_counts = {}
+        for outcome, occurrence in counts_to_map.items():
+            mapped_outcome = "".join(outcome[i] for i in mapped_wires)
+            mapped_counts[mapped_outcome] = mapped_counts.get(mapped_outcome, 0) + occurrence
+
+        return mapped_counts
+
+    def _include_all_outcomes(self, outcome_counts: dict) -> None:
+        """
+        Includes missing outcomes in outcome_counts.
+        If an outcome is not present in outcome_counts, it's count is considered 0
+
+        Args:
+            outcome_counts(dict): Dictionary where key is binary representation of the outcome and value is its count
+        """
+        num_wires = len(self.wires)
+        num_outcomes = 2**num_wires
+        if num_outcomes == len(outcome_counts.keys()):
+            return
+
+        binary_pattern = "{0:0" + str(num_wires) + "b}"
+        for outcome in range(num_outcomes):
+            outcome_binary = binary_pattern.format(outcome)
+            if outcome_binary not in outcome_counts:
+                outcome_counts[outcome_binary] = 0
+
+
+def _remove_unobserved_outcomes(outcome_counts: dict):
+    """
+    Removes unobserved outcomes, i.e. whose count is 0 from the outcome_count dictionary.
+
+    Args:
+        outcome_counts(dict): Dictionary where key is binary representation of the outcome and value is its count
+    """
+    for outcome in list(outcome_counts.keys()):
+        if outcome_counts[outcome] == 0:
+            del outcome_counts[outcome]
diff --git a/pennylane/measurements/expval.py b/pennylane/measurements/expval.py
index 6e8f14649..1fffa7469 100644
--- a/pennylane/measurements/expval.py
+++ b/pennylane/measurements/expval.py
@@ -20,7 +20,6 @@ from typing import Sequence, Tuple, Union
 import pennylane as qml
 from pennylane.operation import Operator
 from pennylane.wires import Wires
-
 from .measurements import Expectation, SampleMeasurement, StateMeasurement
 from .mid_measure import MeasurementValue
 
@@ -128,11 +127,24 @@ class ExpectationMP(SampleMeasurement, StateMeasurement):
     def process_state(self, state: Sequence[complex], wire_order: Wires):
         # This also covers statistics for mid-circuit measurements manipulated using
         # arithmetic operators
-        eigvals = qml.math.asarray(self.eigvals(), dtype="float64")
-
         # we use ``self.wires`` instead of ``self.obs`` because the observable was
         # already applied to the state
         with qml.queuing.QueuingManager.stop_recording():
             prob = qml.probs(wires=self.wires).process_state(state=state, wire_order=wire_order)
         # In case of broadcasting, `prob` has two axes and this is a matrix-vector product
-        return qml.math.dot(prob, eigvals)
+        return self._calculate_expectation(prob)
+
+    def process_counts(self, counts: dict, wire_order: Wires):
+        with qml.QueuingManager.stop_recording():
+            probs = qml.probs(wires=self.wires).process_counts(counts=counts, wire_order=wire_order)
+        return self._calculate_expectation(probs)
+
+    def _calculate_expectation(self, probabilities):
+        """
+        Calculate the of expectation set of probabilities.
+
+        Args:
+            probabilities (array): the probabilities of collapsing to eigen states
+        """
+        eigvals = qml.math.asarray(self.eigvals(), dtype="float64")
+        return qml.math.dot(probabilities, eigvals)
diff --git a/pennylane/measurements/probs.py b/pennylane/measurements/probs.py
index 315bd4cf9..387580fe1 100644
--- a/pennylane/measurements/probs.py
+++ b/pennylane/measurements/probs.py
@@ -17,9 +17,9 @@ This module contains the qml.probs measurement.
 from typing import Sequence, Tuple
 
 import numpy as np
+
 import pennylane as qml
 from pennylane.wires import Wires
-
 from .measurements import Probability, SampleMeasurement, StateMeasurement
 from .mid_measure import MeasurementValue
 
@@ -237,26 +237,18 @@ class ProbabilityMP(SampleMeasurement, StateMeasurement):
         return qml.math.reshape(prob, flat_shape)
 
     def process_counts(self, counts: dict, wire_order: Wires) -> np.ndarray:
-        wire_map = dict(zip(wire_order, range(len(wire_order))))
-        mapped_wires = [wire_map[w] for w in self.wires]
+        with qml.QueuingManager.stop_recording():
+            helper_counts = qml.counts(wires=self.wires, all_outcomes=False)
+        mapped_counts = helper_counts.process_counts(counts, wire_order)
 
-        # when reducing wires, two keys may become equal
-        # the following structure was chosen to maintain compatibility with 'process_samples'
-        if mapped_wires:
-            mapped_counts = {}
-            for outcome, occurrence in counts.items():
-                mapped_outcome = "".join(outcome[i] for i in mapped_wires)
-                mapped_counts[mapped_outcome] = mapped_counts.get(mapped_outcome, 0) + occurrence
-            counts = mapped_counts
-
-        num_shots = sum(counts.values())
-        num_wires = len(next(iter(counts)))
+        num_shots = sum(mapped_counts.values())
+        num_wires = len(next(iter(mapped_counts)))
         dim = 2**num_wires
 
         # constructs the probability vector
         # converts outcomes from binary strings to integers (base 10 representation)
         prob_vector = qml.math.zeros((dim), dtype="float64")
-        for outcome, occurrence in counts.items():
+        for outcome, occurrence in mapped_counts.items():
             prob_vector[int(outcome, base=2)] = occurrence / num_shots
 
         return prob_vector
diff --git a/pennylane/measurements/var.py b/pennylane/measurements/var.py
index d62750dc7..a8b8f8384 100644
--- a/pennylane/measurements/var.py
+++ b/pennylane/measurements/var.py
@@ -21,7 +21,6 @@ from typing import Sequence, Tuple, Union
 import pennylane as qml
 from pennylane.operation import Operator
 from pennylane.wires import Wires
-
 from .measurements import SampleMeasurement, StateMeasurement, Variance
 from .mid_measure import MeasurementValue
 
@@ -122,11 +121,24 @@ class VarianceMP(SampleMeasurement, StateMeasurement):
     def process_state(self, state: Sequence[complex], wire_order: Wires):
         # This also covers statistics for mid-circuit measurements manipulated using
         # arithmetic operators
-        eigvals = qml.math.asarray(self.eigvals(), dtype="float64")
-
         # we use ``wires`` instead of ``op`` because the observable was
         # already applied to the state
         with qml.queuing.QueuingManager.stop_recording():
             prob = qml.probs(wires=self.wires).process_state(state=state, wire_order=wire_order)
         # In case of broadcasting, `prob` has two axes and these are a matrix-vector products
-        return qml.math.dot(prob, (eigvals**2)) - qml.math.dot(prob, eigvals) ** 2
+        return self._calculate_variance(prob)
+
+    def process_counts(self, counts: dict, wire_order: Wires):
+        with qml.QueuingManager.stop_recording():
+            probs = qml.probs(wires=self.wires).process_counts(counts=counts, wire_order=wire_order)
+        return self._calculate_variance(probs)
+
+    def _calculate_variance(self, probabilities):
+        """
+        Calculate the variance of a set of probabilities.
+
+        Args:
+            probabilities (array): the probabilities of collapsing to eigen states
+        """
+        eigvals = qml.math.asarray(self.eigvals(), dtype="float64")
+        return qml.math.dot(probabilities, (eigvals**2)) - qml.math.dot(probabilities, eigvals) ** 2


Test patch:
diff --git a/tests/measurements/test_counts.py b/tests/measurements/test_counts.py
index 87c013703..c364c5537 100644
--- a/tests/measurements/test_counts.py
+++ b/tests/measurements/test_counts.py
@@ -867,3 +867,74 @@ def test_batched_counts_and_expval_operator_finite_shots(interface):
     assert isinstance(res, tuple) and len(res) == 2
     assert res[0] == type(res[0])([{-1: n_shots}, {1: n_shots}])
     assert len(res[1]) == 2 and qml.math.allequal(res[1], [-1, 1])
+
+
+class TestProcessCounts:
+    """Unit tests for the counts.process_counts method"""
+
+    @pytest.mark.parametrize("all_outcomes", [True, False])
+    def test_should_not_modify_counts_dictionary(self, all_outcomes):
+        """Tests that count dictionary is not modified"""
+
+        counts = {"000": 100, "100": 100}
+        expected = counts.copy()
+        wire_order = qml.wires.Wires((0, 1, 2))
+
+        qml.counts(wires=(0, 1), all_outcomes=all_outcomes).process_counts(counts, wire_order)
+
+        assert counts == expected
+
+    def test_all_outcomes_is_true(self):
+        """When all_outcomes is True, 0 count should be added to missing outcomes in the counts dictionary"""
+
+        counts_to_process = {"00": 100, "10": 100}
+        wire_order = qml.wires.Wires((0, 1))
+
+        actual = qml.counts(wires=wire_order, all_outcomes=True).process_counts(
+            counts_to_process, wire_order=wire_order
+        )
+
+        expected_counts = {"00": 100, "01": 0, "10": 100, "11": 0}
+        assert actual == expected_counts
+
+    def test_all_outcomes_is_false(self):
+        """When all_outcomes is True, 0 count should be removed from the counts dictionary"""
+
+        counts_to_process = {"00": 0, "01": 0, "10": 0, "11": 100}
+        wire_order = qml.wires.Wires((0, 1))
+
+        actual = qml.counts(wires=wire_order, all_outcomes=False).process_counts(
+            counts_to_process, wire_order=wire_order
+        )
+
+        expected_counts = {"11": 100}
+        assert actual == expected_counts
+
+    @pytest.mark.parametrize(
+        "wires, expected",
+        [
+            ((0, 1), {"00": 100, "10": 100}),
+            ((1, 0), {"00": 100, "01": 100}),
+        ],
+    )
+    def test_wire_order(self, wires, expected):
+        """Test impact of wires in qml.counts"""
+        counts = {"000": 100, "100": 100}
+        wire_order = qml.wires.Wires((0, 1, 2))
+
+        actual = qml.counts(wires=wires, all_outcomes=False).process_counts(counts, wire_order)
+
+        assert actual == expected
+
+    @pytest.mark.parametrize("all_outcomes", [True, False])
+    def test_process_count_returns_same_count_dictionary(self, all_outcomes):
+        """
+        Test that process_count returns same dictionary when all outcomes are in count dictionary and wire_order is same
+        """
+
+        expected = {"0": 100, "1": 100}
+        wires = qml.wires.Wires(0)
+
+        actual = qml.counts(wires=wires, all_outcomes=all_outcomes).process_counts(expected, wires)
+
+        assert actual == expected
diff --git a/tests/measurements/test_expval.py b/tests/measurements/test_expval.py
index 57474643f..ed0bf60cc 100644
--- a/tests/measurements/test_expval.py
+++ b/tests/measurements/test_expval.py
@@ -319,3 +319,20 @@ class TestExpval:
         energy_batched = cost_circuit(params)
 
         assert qml.math.allequal(energy_batched, energy)
+
+    @pytest.mark.parametrize(
+        "wire, expected",
+        [
+            (0, 0.0),
+            (1, 1.0),
+        ],
+    )
+    def test_estimate_expectation_with_counts(self, wire, expected):
+        """Test that the expectation value of an observable is estimated correctly using counts"""
+        counts = {"000": 100, "100": 100}
+
+        wire_order = qml.wires.Wires((0, 1, 2))
+
+        res = qml.expval(qml.Z(wire)).process_counts(counts=counts, wire_order=wire_order)
+
+        assert np.allclose(res, expected)
diff --git a/tests/measurements/test_var.py b/tests/measurements/test_var.py
index b6a71c14d..88d9a8a96 100644
--- a/tests/measurements/test_var.py
+++ b/tests/measurements/test_var.py
@@ -215,3 +215,20 @@ class TestVar:
             return qml.var(obs_2)
 
         assert circuit() == circuit2()
+
+    @pytest.mark.parametrize(
+        "wire, expected",
+        [
+            (0, 1.0),
+            (1, 0.0),
+        ],
+    )
+    def test_estimate_variance_with_counts(self, wire, expected):
+        """Test that the variance of an observable is estimated correctly using counts."""
+        counts = {"000": 100, "100": 100}
+
+        wire_order = qml.wires.Wires((0, 1, 2))
+
+        res = qml.var(qml.Z(wire)).process_counts(counts=counts, wire_order=wire_order)
+
+        assert np.allclose(res, expected)