Repository: PennyLaneAI/pennylane (python)

Issue:
Implement `SampleMP.proces_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=10)
qml.sample(wires=0).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.

# Implementation Details

The task is to implement [SampleMP.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 to look like raw samples.

```pycon
>>> counts = {'00': 2, '10': 3}
>>> qml.sample(wires=(0,1)).process_counts(counts, wire_order=qml.wires.Wires((0,1)))
array([[0, 0],
       [0, 0],
       [1, 0],
       [1, 0],
       [1, 0]])
>>> qml.sample(wires=0).process_counts(counts, wire_order=qml.wires.Wires((0,1)))
array([0, 0, 1, 1, 1])
```

# Requirements

A completed PR should:

* Add the `process_counts` method to `SampleMP`
* Add unit tests to [pennylane/tests/measurements/test_sample.py](https://github.com/PennyLaneAI/pennylane/blob/master/tests/measurements/test_sample.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 :)

Pull request:
Implement process_counts in SampleMP
**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 `SampleMP`.
* Make `SampleMeasurement.process_counts` an abstract method

**Benefits:**
The method `SampleMP.process_counts` won't throw `NotImplementedError` exception

**Possible Drawbacks:**

**Related GitHub Issues:**
Fixes #5371 

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 1e579a934..54456bc79 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -179,8 +179,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`.
+* Implemented the method `process_counts` in `ExpectationMP`, `VarianceMP`, `CountsMP`, and `SampleMP`
   [(#5256)](https://github.com/PennyLaneAI/pennylane/pull/5256)
+  [(#5395)](https://github.com/PennyLaneAI/pennylane/pull/5395)
 
 <h3>Breaking changes 💔</h3>
 
diff --git a/pennylane/measurements/measurements.py b/pennylane/measurements/measurements.py
index b8c76522c..1ac4a4c32 100644
--- a/pennylane/measurements/measurements.py
+++ b/pennylane/measurements/measurements.py
@@ -542,6 +542,7 @@ class SampleMeasurement(MeasurementProcess):
                 provided, the entire shot range is treated as a single bin.
         """
 
+    @abstractmethod
     def process_counts(self, counts: dict, wire_order: Wires):
         """Calculate the measurement given a counts histogram dictionary.
 
@@ -551,7 +552,6 @@ class SampleMeasurement(MeasurementProcess):
 
         Note that the input dictionary may only contain states with non-zero entries (``all_outcomes=False``).
         """
-        raise NotImplementedError
 
 
 class StateMeasurement(MeasurementProcess):
diff --git a/pennylane/measurements/sample.py b/pennylane/measurements/sample.py
index eea7089ee..39f7c2fab 100644
--- a/pennylane/measurements/sample.py
+++ b/pennylane/measurements/sample.py
@@ -18,6 +18,8 @@ import functools
 import warnings
 from typing import Sequence, Tuple, Optional, Union
 
+import numpy as np
+
 import pennylane as qml
 from pennylane.operation import Operator
 from pennylane.wires import Wires
@@ -272,3 +274,45 @@ class SampleMP(SampleMeasurement):
                 ) from e
 
         return samples if bin_size is None else samples.reshape((bin_size, -1))
+
+    def process_counts(self, counts: dict, wire_order: Wires):
+        samples = []
+        mapped_counts = self._map_counts(counts, wire_order)
+        for outcome, count in mapped_counts.items():
+            outcome_sample = self._compute_outcome_sample(outcome)
+            if len(self.wires) == 1:
+                # If only one wire is sampled, flatten the list
+                outcome_sample = outcome_sample[0]
+            samples.extend([outcome_sample] * count)
+
+        return np.array(samples)
+
+    def _map_counts(self, counts_to_map, wire_order) -> 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
+        """
+        with qml.QueuingManager.stop_recording():
+            helper_counts = qml.counts(wires=self.wires, all_outcomes=False)
+        return helper_counts.process_counts(counts_to_map, wire_order)
+
+    def _compute_outcome_sample(self, outcome) -> list:
+        """
+        Args:
+            outcome (str): The binary string representation of the measurement outcome.
+
+        Returns:
+            list: A list of outcome samples for given binary string.
+                If eigenvalues exist, the binary outcomes are mapped to their corresponding eigenvalues.
+        """
+        outcome_samples = [int(bit) for bit in outcome]
+
+        if self.eigvals() is not None:
+            eigvals = self.eigvals()
+            outcome_samples = [eigvals[outcome] for outcome in outcome_samples]
+
+        return outcome_samples


Test patch:
diff --git a/pennylane/devices/tests/test_measurements.py b/pennylane/devices/tests/test_measurements.py
index 2a2b8b07d..f0488cfa4 100644
--- a/pennylane/devices/tests/test_measurements.py
+++ b/pennylane/devices/tests/test_measurements.py
@@ -27,6 +27,7 @@ from pennylane.measurements import (
     StateMeasurement,
     StateMP,
 )
+from pennylane.wires import Wires
 
 pytestmark = pytest.mark.skip_unsupported
 
@@ -1630,6 +1631,9 @@ class TestSampleMeasurement:
             def process_samples(self, samples, wire_order, shot_range=None, bin_size=None):
                 return 1
 
+            def process_counts(self, counts: dict, wire_order: Wires):
+                return 1
+
         @qml.qnode(dev)
         def circuit():
             qml.X(0)
@@ -1651,6 +1655,9 @@ class TestSampleMeasurement:
             def process_samples(self, samples, wire_order, shot_range=None, bin_size=None):
                 return 1
 
+            def process_counts(self, counts: dict, wire_order: Wires):
+                return 1
+
         @qml.qnode(dev)
         def circuit():
             qml.X(0)
diff --git a/tests/measurements/legacy/test_measurements_legacy.py b/tests/measurements/legacy/test_measurements_legacy.py
index 8290ab5a8..d524095f7 100644
--- a/tests/measurements/legacy/test_measurements_legacy.py
+++ b/tests/measurements/legacy/test_measurements_legacy.py
@@ -28,6 +28,9 @@ from pennylane.measurements import (
     var,
 )
 
+from pennylane.wires import Wires
+
+
 # pylint: disable=too-few-public-methods, unused-argument
 
 
@@ -90,6 +93,9 @@ class TestSampleMeasurement:
             def process_samples(self, samples, wire_order, shot_range, bin_size):
                 return qml.math.sum(samples[..., self.wires])
 
+            def process_counts(self, counts: dict, wire_order: Wires):
+                return counts
+
         dev = qml.device("default.qubit.legacy", wires=2, shots=1000)
 
         @qml.qnode(dev)
@@ -107,6 +113,9 @@ class TestSampleMeasurement:
             def process_samples(self, samples, wire_order, shot_range, bin_size):
                 return qml.math.sum(samples[..., self.wires])
 
+            def process_counts(self, counts: dict, wire_order: Wires):
+                return counts
+
         dev = qml.device("default.qubit.legacy", wires=2)
 
         @qml.qnode(dev)
diff --git a/tests/measurements/test_measurements.py b/tests/measurements/test_measurements.py
index e251a359a..37ae04233 100644
--- a/tests/measurements/test_measurements.py
+++ b/tests/measurements/test_measurements.py
@@ -583,6 +583,9 @@ class TestSampleMeasurement:
             def process_samples(self, samples, wire_order, shot_range=None, bin_size=None):
                 return qml.math.sum(samples[..., self.wires])
 
+            def process_counts(self, counts: dict, wire_order: Wires):
+                return counts
+
         dev = qml.device("default.qubit", wires=2, shots=1000)
 
         @qml.qnode(dev)
@@ -600,6 +603,9 @@ class TestSampleMeasurement:
             def process_samples(self, samples, wire_order, shot_range, bin_size):
                 return qml.math.sum(samples[..., self.wires])
 
+            def process_counts(self, counts: dict, wire_order: Wires):
+                return counts
+
             @property
             def return_type(self):
                 return Sample
@@ -617,22 +623,6 @@ class TestSampleMeasurement:
         ):
             circuit()
 
-    def test_process_counts_not_implemented(self):
-        """Test that process_counts is not implemented by default."""
-
-        class MyMeasurement(SampleMeasurement):
-            # pylint: disable=signature-differs
-            def process_samples(self, samples, wire_order, shot_range, bin_size):
-                return qml.math.sum(samples[..., self.wires])
-
-            @property
-            def return_type(self):
-                return Sample
-
-        m = MyMeasurement(wires=[0])
-        with pytest.raises(NotImplementedError):
-            m.process_counts({"0": 10}, wire_order=qml.wires.Wires(0))
-
 
 class TestStateMeasurement:
     """Tests for the SampleMeasurement class."""
@@ -690,6 +680,9 @@ class TestMeasurementTransform:
             def process_samples(self, samples, wire_order, shot_range=None, bin_size=None):
                 return [True]
 
+            def process_counts(self, counts: dict, wire_order: Wires):
+                return counts
+
         dev = qml.device("default.qubit", wires=2, shots=1000)
 
         @qml.qnode(dev)
diff --git a/tests/measurements/test_sample.py b/tests/measurements/test_sample.py
index fcb1c8430..e2920b274 100644
--- a/tests/measurements/test_sample.py
+++ b/tests/measurements/test_sample.py
@@ -506,3 +506,57 @@ def test_jitting_with_sampling_on_subset_of_wires(samples):
     assert (
         circuit._qfunc_output.shape(dev, Shots(samples)) == (samples, 2) if samples != 1 else (2,)
     )
+
+
+class TestSampleProcessCounts:
+    """Tests for the process_counts method in the SampleMP class."""
+
+    def test_process_counts_multiple_wires(self):
+        """Test process_counts method with multiple wires."""
+        sample_mp = qml.sample(wires=[0, 1])
+        counts = {"00": 2, "10": 3}
+        wire_order = qml.wires.Wires((0, 1))
+
+        result = sample_mp.process_counts(counts, wire_order)
+
+        assert np.array_equal(result, np.array([[0, 0], [0, 0], [1, 0], [1, 0], [1, 0]]))
+
+    def test_process_counts_single_wire(self):
+        """Test process_counts method with a single wire."""
+        sample_mp = qml.sample(wires=[0])
+        counts = {"00": 2, "10": 3}
+        wire_order = qml.wires.Wires((0, 1))
+
+        result = sample_mp.process_counts(counts, wire_order)
+
+        assert np.array_equal(result, np.array([0, 0, 1, 1, 1]))
+
+    def test_process_counts_with_eigen_values(self):
+        """Test process_counts method with eigen values."""
+        sample_mp = qml.sample(qml.Z(0))
+        counts = {"00": 2, "10": 3}
+        wire_order = qml.wires.Wires((0, 1))
+
+        result = sample_mp.process_counts(counts, wire_order)
+
+        assert np.array_equal(result, np.array([1, 1, -1, -1, -1]))
+
+    def test_process_counts_with_inverted_wire_order(self):
+        """Test process_counts method with inverted wire order."""
+        sample_mp = qml.sample(wires=[0, 1])
+        counts = {"00": 2, "01": 3}
+        wire_order = qml.wires.Wires((1, 0))
+
+        result = sample_mp.process_counts(counts, wire_order)
+
+        assert np.array_equal(result, np.array([[0, 0], [0, 0], [1, 0], [1, 0], [1, 0]]))
+
+    def test_process_counts_with_second_single_wire(self):
+        """Test process_counts method with the second single wire."""
+        sample_mp = qml.sample(wires=[1])
+        counts = {"00": 2, "10": 3}
+        wire_order = qml.wires.Wires((0, 1))
+
+        result = sample_mp.process_counts(counts, wire_order)
+
+        assert np.array_equal(result, np.array([0, 0, 0, 0, 0]))