Supporting counts from raw samples

doc/introduction/measurements.rst

@@ -102,7 +102,6 @@ the measurement results always coincide, and the lists are therefore equal:
 >>> np.all(result[0] == result[1])
 True
 
-
 Tensor observables
 ------------------
 
@@ -127,6 +126,66 @@ accept observables as arguments,
 including :func:`~.pennylane.expval`, :func:`~.pennylane.var`,
 and :func:`~.pennylane.sample`.
 
+Counts
+------
+
+To avoid dealing with long arrays for the larger numbers of shots, one can pass an argument counts=True
+to :func:`~pennylane.sample`. In this case, the result will be a dictionary containing the number of occurrences for each
+unique sample. The previous example will be modified as follows:
+
+.. code-block:: python
+
+    dev = qml.device("default.qubit", wires=2, shots=1000)
+
+    @qml.qnode(dev)
+    def circuit():
+        qml.Hadamard(wires=0)
+        qml.CNOT(wires=[0, 1])
+        # passing the counts flag
+        return qml.sample(qml.PauliZ(0), counts=True), qml.sample(qml.PauliZ(1), counts=True)
+
+After executing the circuit, we can directly see how many times each measurement outcome occurred:
+
+>>> result = circuit()
+>>> print(result)
+[{-1: 526, 1: 474} {-1: 526, 1: 474}]
+
+Similarly, if the observable is not provided, the count of each computational basis state is returned.
+
+.. code-block:: python
+
+    dev = qml.device("default.qubit", wires=2, shots=1000)
+
+    @qml.qnode(dev)
+    def circuit():
+        qml.Hadamard(wires=0)
+        qml.CNOT(wires=[0, 1])
+        # passing the counts flag
+        return qml.sample(counts=True)
+
+And the result is:
+
+>>> result = circuit()
+>>> print(result)
+{'00': 495, '11': 505}
+
+If counts are obtained along with a measurement function other than :func:`~.pennylane.sample`,
+a tensor of tensors is returned to provide differentiability for the outputs of QNodes.
+
+.. code-block:: python
+
+    @qml.qnode(dev)
+    def circuit():
+        qml.Hadamard(wires=0)
+        qml.CNOT(wires=[0,1])
+        qml.PauliX(wires=2)
+        return qml.expval(qml.PauliZ(0)),qml.expval(qml.PauliZ(1)), qml.sample(counts=True)
+
+>>> result = circuit()
+>>> print(result)
+[tensor(0.026, requires_grad=True) tensor(0.026, requires_grad=True)
+ tensor({'001': 513, '111': 487}, dtype=object, requires_grad=True)]
+
 Probability
 -----------
 

doc/releases/changelog-dev.md

@@ -47,6 +47,43 @@
   [(#2699)](https://github.com/PennyLaneAI/pennylane/pull/2699)  
 
 <h3>Improvements</h3>
+
+* Samples can be grouped into counts by passing the `counts=True` flag to `qml.sample`.
+  [(#2686)](https://github.com/PennyLaneAI/pennylane/pull/2686)
+
+  Note that the change included creating a new `Counts` measurement type in `measurements.py`.
+
+  `counts=True` can be set when obtaining raw samples in the computational basis:
+
+  ```pycon
+  >>> dev = qml.device("default.qubit", wires=2, shots=1000)
+  >>>
+  >>> @qml.qnode(dev)
+  >>> def circuit():
+  ...     qml.Hadamard(wires=0)
+  ...     qml.CNOT(wires=[0, 1])
+  ...     # passing the counts flag
+  ...     return qml.sample(counts=True)   
+  >>> result = circuit()
+  >>> print(result)
+  {'00': 495, '11': 505}
+  ```
+
+  Counts can also be obtained when sampling the eigenstates of an observable:
+
+  ```pycon
+  >>> dev = qml.device("default.qubit", wires=2, shots=1000)
+  >>>
+  >>> @qml.qnode(dev)
+  >>> def circuit():
+  ...   qml.Hadamard(wires=0)
+  ...   qml.CNOT(wires=[0, 1])
+  ...   return qml.sample(qml.PauliZ(0), counts=True), qml.sample(qml.PauliZ(1), counts=True)
+  >>> result = circuit()
+  >>> print(result)
+  [tensor({-1: 526, 1: 474}, dtype=object, requires_grad=True)
+   tensor({-1: 526, 1: 474}, dtype=object, requires_grad=True)]
+  ```
 
 * The `qml.state` and `qml.density_matrix` measurements now support custom wire
   labels.
@@ -105,5 +142,6 @@
 
 This release contains contributions from (in alphabetical order):
 
-David Ittah, Edward Jiang, Ankit Khandelwal, Christina Lee, Ixchel Meza Chavez, Mudit Pandey,
+
+David Ittah, Edward Jiang, Ankit Khandelwal, Christina Lee, Ixchel Meza Chavez, Bogdan Reznychenko, Mudit Pandey,
 Antal Száva, Moritz Willmann

pennylane/_qubit_device.py

@@ -29,13 +29,15 @@
 from pennylane.operation import operation_derivative
 from pennylane.measurements import (
     Sample,
+    Counts,
     Variance,
     Expectation,
     Probability,
     State,
     VnEntropy,
     MutualInfo,
 )
+
 from pennylane import Device
 from pennylane.math import sum as qmlsum
 from pennylane.math import multiply as qmlmul
@@ -276,10 +278,13 @@ def execute(self, circuit, **kwargs):
 
                 if qml.math._multi_dispatch(r) == "jax":  # pylint: disable=protected-access
                     r = r[0]
-                else:
+                elif not isinstance(r[0], dict):
+                    # Measurement types except for Counts
                     r = qml.math.squeeze(r)
-
-                if shot_tuple.copies > 1:
+                if isinstance(r, (np.ndarray, list)) and r.shape and isinstance(r[0], dict):
+                    # This happens when measurement type is Counts
+                    results.append(r)
+                elif shot_tuple.copies > 1:
                     results.extend(r.T)
                 else:
                     results.append(r.T)
@@ -301,7 +306,7 @@ def execute(self, circuit, **kwargs):
                 if circuit.measurements[0].return_type is qml.measurements.State:
                     # State: assumed to only be allowed if it's the only measurement
                     results = self._asarray(results, dtype=self.C_DTYPE)
-                else:
+                elif circuit.measurements[0].return_type is not qml.measurements.Counts:
                     # Measurements with expval, var or probs
                     results = self._asarray(results, dtype=self.R_DTYPE)
 
@@ -311,7 +316,8 @@ def execute(self, circuit, **kwargs):
             ):
                 # Measurements with expval or var
                 results = self._asarray(results, dtype=self.R_DTYPE)
-            else:
+            elif any(ret is not qml.measurements.Counts for ret in ret_types):
+                # all the other cases except all counts
                 results = self._asarray(results)
 
         elif circuit.all_sampled and not self._has_partitioned_shots():
@@ -475,7 +481,14 @@ def statistics(self, observables, shot_range=None, bin_size=None):
                 results.append(self.var(obs, shot_range=shot_range, bin_size=bin_size))
 
             elif obs.return_type is Sample:
-                results.append(self.sample(obs, shot_range=shot_range, bin_size=bin_size))
+                results.append(
+                    self.sample(obs, shot_range=shot_range, bin_size=bin_size, counts=False)
+                )
+
+            elif obs.return_type is Counts:
+                results.append(
+                    self.sample(obs, shot_range=shot_range, bin_size=bin_size, counts=True)
+                )
 
             elif obs.return_type is Probability:
                 results.append(
@@ -958,20 +971,39 @@ def var(self, observable, shot_range=None, bin_size=None):
         samples = self.sample(observable, shot_range=shot_range, bin_size=bin_size)
         return np.squeeze(np.var(samples, axis=0))
 
-    def sample(self, observable, shot_range=None, bin_size=None):
+    def sample(self, observable, shot_range=None, bin_size=None, counts=False):
+        def _samples_to_counts(samples, no_observable_provided):
+            """Group the obtained samples into a dictionary.
+
+            **Example**
+
+                >>> samples
+                tensor([[0, 0, 1],
+                        [0, 0, 1],
+                        [1, 1, 1]], requires_grad=True)
+                >>> self._samples_to_counts(samples)
+                {'111':1, '001':2}
+            """
+            if no_observable_provided:
+                # If we describe a state vector, we need to convert its list representation
+                # into string (it's hashable and good-looking).
+                # Before converting to str, we need to extract elements from arrays
+                # to satisfy the case of jax interface, as jax arrays do not support str.
+                samples = ["".join([str(s.item()) for s in sample]) for sample in samples]
+            states, counts = np.unique(samples, return_counts=True)
+            return dict(zip(states, counts))
 
         # translate to wire labels used by device
         device_wires = self.map_wires(observable.wires)
         name = observable.name
         sample_slice = Ellipsis if shot_range is None else slice(*shot_range)
+        no_observable_provided = isinstance(observable, MeasurementProcess)
 
         if isinstance(name, str) and name in {"PauliX", "PauliY", "PauliZ", "Hadamard"}:
             # Process samples for observables with eigenvalues {1, -1}
             samples = 1 - 2 * self._samples[sample_slice, device_wires[0]]
 
-        elif isinstance(
-            observable, MeasurementProcess
-        ):  # if no observable was provided then return the raw samples
+        elif no_observable_provided:  # if no observable was provided then return the raw samples
             if (
                 len(observable.wires) != 0
             ):  # if wires are provided, then we only return samples from those wires
@@ -998,9 +1030,20 @@ def sample(self, observable, shot_range=None, bin_size=None):
                 ) from e
 
         if bin_size is None:
+            if counts:
+                return _samples_to_counts(samples, no_observable_provided)
             return samples
-
-        return samples.reshape((bin_size, -1))
+        if counts:
+            shape = (-1, bin_size, 3) if no_observable_provided else (-1, bin_size)
+            return [
+                _samples_to_counts(bin_sample, no_observable_provided)
+                for bin_sample in samples.reshape(shape)
+            ]
+        return (
+            samples.reshape((3, bin_size, -1))
+            if no_observable_provided
+            else samples.reshape((bin_size, -1))
+        )
 
     def adjoint_jacobian(self, tape, starting_state=None, use_device_state=False):
         """Implements the adjoint method outlined in

pennylane/interfaces/autograd.py

@@ -111,10 +111,16 @@ def _execute(
 
     for i, r in enumerate(res):
 
-        if isinstance(res[i], np.ndarray):
+        if isinstance(r, np.ndarray):
             # For backwards compatibility, we flatten ragged tape outputs
             # when there is no sampling
-            r = np.hstack(res[i]) if res[i].dtype == np.dtype("object") else res[i]
+            try:
+                if isinstance(r[0][0], dict):
+                    # This happens when measurement type is Counts and shot vector is passed
+                    continue
+            except (IndexError, KeyError):
+                pass
+            r = np.hstack(r) if r.dtype == np.dtype("object") else r
             res[i] = np.tensor(r)
 
         elif isinstance(res[i], tuple):

pennylane/measurements.py

@@ -38,6 +38,7 @@ class ObservableReturnTypes(Enum):
     """Enumeration class to represent the return types of an observable."""
 
     Sample = "sample"
+    Counts = "counts"
     Variance = "var"
     Expectation = "expval"
     Probability = "probs"
@@ -54,6 +55,10 @@ def __repr__(self):
 Sample = ObservableReturnTypes.Sample
 """Enum: An enumeration which represents sampling an observable."""
 
+Counts = ObservableReturnTypes.Counts
+"""Enum: An enumeration which represents returning the number of times
+ each sample was obtained."""
+
 Variance = ObservableReturnTypes.Variance
 """Enum: An enumeration which represents returning the variance of
 an observable on specified wires."""
@@ -578,9 +583,10 @@ def circuit(x):
     return MeasurementProcess(Variance, obs=op)
 
 
-def sample(op=None, wires=None):
+def sample(op=None, wires=None, counts=False):
     r"""Sample from the supplied observable, with the number of shots
-    determined from the ``dev.shots`` attribute of the corresponding device.
+    determined from the ``dev.shots`` attribute of the corresponding device,
+    returning raw samples (counts=False) or the number of counts for each sample (counts=True).
     If no observable is provided then basis state samples are returned directly
     from the device.
 
@@ -590,6 +596,7 @@ def sample(op=None, wires=None):
     Args:
         op (Observable or None): a quantum observable object
         wires (Sequence[int] or int or None): the wires we wish to sample from, ONLY set wires if op is None
+        counts (bool): return the result as number of counts for each sample
 
     Raises:
         QuantumFunctionError: `op` is not an instance of :class:`~.Observable`
@@ -642,6 +649,27 @@ def circuit(x):
            [1, 1],
            [0, 0]])
 
+    If specified counts=True, the function returns number of counts for each sample,
+    both for observables eigenvalues or the system eigenstates.
+
+    .. code-block:: python3
+
+        dev = qml.device('default.qubit', wires=3, shots=10)
+
+        @qml.qnode(dev)
+        def my_circ():
+            qml.Hadamard(wires=0)
+            qml.CNOT(wires=[0,1])
+            qml.PauliX(wires=2)
+            return qml.sample(qml.PauliZ(0), counts = True), qml.sample(counts=True)
+
+    Executing this QNode:
+
+    >>> my_circ()
+    tensor([tensor({-1: 5, 1: 5}, dtype=object, requires_grad=True),
+        tensor({'001': 5, '111': 5}, dtype=object, requires_grad=True)],
+       dtype=object, requires_grad=True)
+
     .. note::
 
         QNodes that return samples cannot, in general, be differentiated, since the derivative
@@ -655,16 +683,17 @@ def circuit(x):
             f"{op.name} is not an observable: cannot be used with sample"
         )
 
+    sample_or_counts = Counts if counts else Sample
+
     if wires is not None:
         if op is not None:
             raise ValueError(
                 "Cannot specify the wires to sample if an observable is "
                 "provided. The wires to sample will be determined directly from the observable."
             )
+        return MeasurementProcess(sample_or_counts, obs=op, wires=qml.wires.Wires(wires))
 
-        return MeasurementProcess(Sample, obs=op, wires=qml.wires.Wires(wires))
-
-    return MeasurementProcess(Sample, obs=op)
+    return MeasurementProcess(sample_or_counts, obs=op)
 
 
 def probs(wires=None, op=None):

pennylane/qnode.py

@@ -650,6 +650,9 @@ def __call__(self, *args, **kwargs):
             self.tape.is_sampled and self.device._has_partitioned_shots()
         ):
             return res
+        if self._qfunc_output.return_type is qml.measurements.Counts:
+            # return a dictionary with counts not as a single-element array
+            return res[0]
 
         return qml.math.squeeze(res)