Implement process_counts in ExpectationMP, VarianceMP and CountsMP (#…

…5256)

**Context:**
Under #4941 abstract method `process_counts` was added to
`SampleMeasurement`. This method provides support to process counts
dictionary produced by external systems.

**Description of the Change:**
* Implement `process_counts` in `ExpectationMP`, `VarianceMP` and
`CountsMP`.
* While implementing `process_counts` in `CountsMP` some common logic
was extracted from `ProbabilityMP` to `CountsMP`

**Benefits:**
The below methods won't throw `NotImplementedError` exception
* `ExpectationMP.process_counts`
* `VarianceMP.process_counts`
* `CountsMP.process_counts`

**Possible Drawbacks:**
* All classes which inherit from `SampleMeasurement` implement
`process_counts` except `SampleMP`.
Should I add an implementation for that ?
* After implementing `process_counts` in all child classes can we make
the method abstract similar to `process_samples` ? This might break some
tests where some classes inherit from `SampleMeasurement` as there
`process_counts` is not implemented
* It is assumed that `counts` dictionary is valid and caller is
responsible to call with valid dictionary. It has already been discussed
in this
[conversation](https://github.com/PennyLaneAI/pennylane/pull/4952/files/eb9c1ee81ea87c274a5ec094f90a0065d05fdc36#r1433117236)
and this is a design choice.

**Related GitHub Issues:**
This PR
* fixes #5249 
* fixes #5244 
* fixes #5241

***

**Further details**
I'm excited to be making my first open-source contribution with this PR
😄 .
As a newcomer to the community, I'm eager to learn and improve. 
Any feedback for enhancements would be greatly appreciated!

**Internal Shortcut Stories**
- [sc-57166]
- [sc-57274]
- [sc-57307]

---------

Co-authored-by: Christina Lee <[email protected]>
Co-authored-by: Christina Lee <[email protected]>

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.

pennylane/measurements/counts.py

@@ -341,3 +341,62 @@ def convert(x):
                 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]

pennylane/measurements/expval.py

@@ -20,7 +20,6 @@
 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 @@ def process_samples(
     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)

pennylane/measurements/probs.py

@@ -17,9 +17,9 @@
 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 @@ def process_state(self, state: Sequence[complex], wire_order: Wires):
         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

pennylane/measurements/var.py

@@ -21,7 +21,6 @@
 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 @@ def process_samples(
     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

tests/measurements/test_counts.py

@@ -867,3 +867,74 @@ def circuit():
     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

tests/measurements/test_expval.py

@@ -319,3 +319,20 @@ def cost_circuit(params):
         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)

tests/measurements/test_var.py

@@ -215,3 +215,20 @@ def circuit2():
             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)