Add zero count outcomes

doc/introduction/measurements.rst

@@ -129,9 +129,11 @@ 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:
+To avoid dealing with long arrays for the larger numbers of shots, one can use :func:`~pennylane.counts` rather than
+:func:`~pennylane.sample`. This performs the same measurement as sampling, but returns a dictionary containing the 
+possible measurement outcomes and the number of occurrences for each, rather than a list of all outcomes. 
+
+The previous example will be modified as follows:
 
 .. code-block:: python
 
@@ -141,16 +143,15 @@ unique sample. The previous example will be modified as follows:
     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)
+        return qml.counts(qml.PauliZ(0)), qml.counts(qml.PauliZ(1))
 
 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}]
+({1: 475, -1: 525}, {1: 475, -1: 525})
 
-Similarly, if the observable is not provided, the count of each computational basis state is returned.
+Similarly, if the observable is not provided, the count of the observed computational basis state is returned.
 
 .. code-block:: python
 
@@ -160,15 +161,37 @@ Similarly, if the observable is not provided, the count of each computational ba
     def circuit():
         qml.Hadamard(wires=0)
         qml.CNOT(wires=[0, 1])
-        # passing the counts flag
-        return qml.sample(counts=True)
+        return qml.counts()
 
 And the result is:
 
 >>> result = circuit()
 >>> print(result)
 {'00': 495, '11': 505}
 
+Per default, only observed outcomes are included in the dictionary. The kwarg ``all_outcomes=True`` can 
+be used to display all possible outcomes, including those that were observed ``0`` times in sampling.
+
+For example, we could run the previous circuit with ``all_outcomes=True``:
+
+.. 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])
+    return qml.counts(all_outcomes=True)
+
+>>> result = circuit()
+>>> print(result)
+{'00': 518, '01': 0, '10': 0, '11': 482}
+
+Note: For complicated Hamiltonians, this can add considerable overhead time (due to the cost of calculating 
+eigenvalues to determine possible outcomes), and as the number of qubits increases, the length of the output 
+dictionary showing possible computational basis states grows rapidly. 
+
 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.
 
@@ -178,8 +201,8 @@ a tensor of tensors is returned to provide differentiability for the outputs of
     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)
+        qml.PauliX(wires=1)
+        return qml.expval(qml.PauliZ(0)),qml.expval(qml.PauliZ(1)), qml.counts()
 
 >>> result = circuit()
 >>> print(result)

doc/releases/changelog-dev.md

@@ -404,6 +404,7 @@ of operators. [(#2622)](https://github.com/PennyLaneAI/pennylane/pull/2622)
   [(#2686)](https://github.com/PennyLaneAI/pennylane/pull/2686)
   [(#2839)](https://github.com/PennyLaneAI/pennylane/pull/2839)
   [(#2876)](https://github.com/PennyLaneAI/pennylane/pull/2876)
+  [(#2889)](https://github.com/PennyLaneAI/pennylane/pull/2889)
 
   Note that the change included creating a new `Counts` measurement type in `measurements.py`.
 
@@ -436,6 +437,21 @@ of operators. [(#2622)](https://github.com/PennyLaneAI/pennylane/pull/2622)
   >>> print(result)
   ({-1: 470, 1: 530}, {-1: 470, 1: 530})
   ```
+
+  Per default, counts returns only the outcomes observed in sampling. Optionally, specifying   `qml.counts(all_outcomes=True)` will return a dictionary containing all possible outcomes:
+
+    ```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.counts()
+  >>> result = circuit()
+  >>> print(result)
+  {'00': 495, '01': 0, '10': 0,  '11': 505}
+  ```
 
 <h3>Improvements</h3>
 
@@ -563,7 +579,8 @@ of operators. [(#2622)](https://github.com/PennyLaneAI/pennylane/pull/2622)
 
 This release contains contributions from (in alphabetical order):
 
-Samuel Banning, Juan Miguel Arrazola, Utkarsh Azad, David Ittah, Soran Jahangiri, Edward Jiang,
-Ankit Khandelwal, Christina Lee, Sergio Martínez-Losa, Albert Mitjans Coma, Ixchel Meza Chavez,
-Romain Moyard, Lee James O'Riordan, Mudit Pandey, Bogdan Reznychenko, Shuli Shu, Jay Soni,
-Modjtaba Shokrian-Zini, Antal Száva, David Wierichs, Moritz Willmann
+Samuel Banning, Juan Miguel Arrazola, Utkarsh Azad, Lillian Marie Austin Frederiksen, David Ittah, 
+Soran Jahangiri, Edward Jiang, Ankit Khandelwal, Christina Lee, Sergio Martínez-Losa, 
+Albert Mitjans Coma, Ixchel Meza Chavez,Romain Moyard, Lee James O'Riordan, Mudit Pandey, 
+Bogdan Reznychenko, Shuli Shu, Jay Soni, Modjtaba Shokrian-Zini, Antal Száva, David Wierichs, 
+Moritz Willmann

pennylane/_qubit_device.py

@@ -1436,6 +1436,84 @@ def var(self, observable, shot_range=None, bin_size=None):
         # TODO: do we need to squeeze here? Maybe remove with new return types
         return np.squeeze(np.var(samples, axis=axis))
 
+    def _samples_to_counts(self, samples, obs, num_wires):
+        """Groups the samples into a dictionary showing number of occurences for
+        each possible outcome.
+
+        The format of the dictionary depends on obs.return_type.all_outcomes (bool),
+        which is set when calling measurements.counts. Per default, the dictionary
+        will only contain the observed outcomes. If obs.return_type.all_outcomes=True,
+        the dictionary will instead contain all possible outcomes, with a count of 0
+        for those not observed. See example.
+
+
+        Args:
+            samples: samples in an array of dimension ``(shots,len(wires))``
+            obs (Observable): the observable sampled
+            num_wires (int): number of wires the sampled observable was performed on
+
+        Returns:
+            dict: dictionary with format ``{'outcome': num_occurences}``, including all
+                outcomes for the sampled observable
+
+        **Example**
+
+            >>> samples
+            tensor([[0, 0],
+                    [0, 0],
+                    [1, 0]], requires_grad=True)
+
+            Per default, this will return:
+            >>> self._samples_to_counts(samples, obs, num_wires)
+            {'00': 2, '10': 1}
+
+            However, if obs.return_type.all_outcomes is set to True, this will return:
+            >>> self._samples_to_counts(samples, obs, num_wires)
+            {'00': 2, '01': 0, '10': 1, '11': 0}
+
+            The variable all_outcomes can be set when running measurements.counts, i.e.:
+
+             .. code-block:: python3
+
+            dev = qml.device("default.qubit", wires=2, shots=4)
+
+            @qml.qnode(dev)
+            def circuit(x):
+                qml.RX(x, wires=0)
+                return qml.counts(all_outcomes=True)
+
+
+        """
+
+        outcomes = []
+
+        if isinstance(obs, MeasurementProcess):
+            # convert samples and outcomes (if using) from arrays to str for dict keys
+            samples = ["".join([str(s.item()) for s in sample]) for sample in samples]
+
+            if obs.return_type.all_outcomes:
+                outcomes = self.generate_basis_states(num_wires)
+                outcomes = ["".join([str(o.item()) for o in outcome]) for outcome in outcomes]
+        else:
+            if obs.return_type.all_outcomes:
+                try:
+                    outcomes = qml.eigvals(obs)
+                #  if observable has no info on eigenvalues, we cannot return this measurement
+                except qml.operation.EigvalsUndefinedError:
+                    warnings.warn(
+                        f"Warning: Error occurred in attempting to generate all outcomes."
+                        f"Cannot find outcomes for {obs.name} (EigvalsUndefinedError). "
+                        f"Only observed outcomes will be returned."
+                    )
+
+        # generate empty outcome dict, populate values with state counts
+        outcome_dict = {k: np.int64(0) for k in outcomes}
+        states, counts = np.unique(samples, return_counts=True)
+        for s, c in zip(states, counts):
+            outcome_dict[s] = c
+
+        return outcome_dict
+
     def sample(self, observable, shot_range=None, bin_size=None, counts=False):
         """Return samples of an observable.
 
@@ -1458,28 +1536,6 @@ def sample(self, observable, shot_range=None, bin_size=None, counts=False):
             dimension ``(shots,)`` or counts
         """
 
-        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
@@ -1522,16 +1578,16 @@ def _samples_to_counts(samples, no_observable_provided):
                     f"Cannot compute samples of {observable.name}."
                 ) from e
 
+        num_wires = len(device_wires) if len(device_wires) > 0 else self.num_wires
         if bin_size is None:
             if counts:
-                return _samples_to_counts(samples, no_observable_provided)
+                return self._samples_to_counts(samples, observable, num_wires)
             return samples
 
-        num_wires = len(device_wires) if len(device_wires) > 0 else self.num_wires
         if counts:
             shape = (-1, bin_size, num_wires) if no_observable_provided else (-1, bin_size)
             return [
-                _samples_to_counts(bin_sample, no_observable_provided)
+                self._samples_to_counts(bin_sample, observable, num_wires)
                 for bin_sample in samples.reshape(shape)
             ]
 

pennylane/measurements.py

@@ -678,7 +678,7 @@ def circuit(x):
     return MeasurementProcess(Sample, obs=op, wires=wires)
 
 
-def counts(op=None, wires=None):
+def counts(op=None, wires=None, all_outcomes=False):
     r"""Sample from the supplied observable, with the number of shots
     determined from the ``dev.shots`` attribute of the corresponding device,
     returning the number of counts for each sample. If no observable is provided then basis state
@@ -690,7 +690,9 @@ def counts(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
+            op is None
+        all_outcomes(bool): determines whether the returned dict will contain only the observed
+            outcomes (default), or whether it will display all possible outcomes for the system
 
     Raises:
         QuantumFunctionError: `op` is not an instance of :class:`~.Observable`
@@ -740,6 +742,34 @@ def circuit(x):
     >>> circuit(0.5)
     {'00': 3, '01': 1}
 
+    Per default, outcomes that were not observed will not be included in the dictionary.
+    Passing all_outcomes=True will create a dictionary that displays all possible outcomes.
+
+    .. code-block:: python3
+
+    dev = qml.device("default.qubit", wires=2, shots=4)
+
+    @qml.qnode(dev)
+    def circuit():
+        qml.PauliX(wires=0)
+        return qml.counts()
+
+    Executing this QNode shows only the observed outcomes:
+
+    >>> circuit()
+    {'01': 4}
+
+
+    @qml.qnode(dev)
+    def circuit(x):
+        qml.PauliX(wires=0)
+        return qml.counts(all_outcomes=True)
+
+    Executing this QNode shows counts for all states:
+
+    >>> circuit()
+    {'00': 0, '01': 0, '10': 4, '11': 0}
+
     .. note::
 
         QNodes that return samples cannot, in general, be differentiated, since the derivative
@@ -764,7 +794,13 @@ def circuit(x):
             )
         wires = qml.wires.Wires(wires)
 
-    return MeasurementProcess(Counts, obs=op, wires=wires)
+    meas_process = MeasurementProcess(Counts, obs=op, wires=wires)
+    if all_outcomes:
+        meas_process.return_type.all_outcomes = True
+    else:
+        meas_process.return_type.all_outcomes = False
+
+    return meas_process
 
 
 def probs(wires=None, op=None):