Add zero count outcomes (#2889)

* generate list of possible sampling outcomes

* add all outcomes as dict keys

* update _samples_to_counts docstring

* move _samples_to_counts up a layer

* pre-commit checks

* pre-commit checks

* minor doc string edit

* update measurements.rst doc

* update measurements.rst doc

* fix bug in wire numbering in example code block

* update tests accept additional dict keys in outcome

* update changelog-dev

* revert pre-commit config file

* make 0 count int64 data type

* update test_measurements.py for new dict format

* add EigvalsUndefinedError exception to _samples_to_counts

* add unit tests

* add pytest xfail to test_scalar_counts_with_obs

* Update doc/introduction/measurements.rst

Co-authored-by: antalszava <[email protected]>

* Update tests/test_new_return_types.py

Co-authored-by: antalszava <[email protected]>

* Update tests/test_measurements.py

Co-authored-by: antalszava <[email protected]>

* Update docstring _samples_to_counts

Co-authored-by: antalszava <[email protected]>

* Update tests/test_measurements.py

Co-authored-by: antalszava <[email protected]>

* remove ToDo

* qml.eigvals(obs) instead of obs.compute_eigvals()

Co-authored-by: antalszava <[email protected]>

* add all_outcomes=False kwarg to counts fn

* update docstrings to reflect new kwarg

* update measurements.rst

* update changelog

* revert tests for standard (all_outcomes=False) counts fn to previous format

* display default (all_outcomes=False) behaviour with warning for EigvalsUndefinedError

* tests default (all_outcomes=False) counts behaviour, test passes

* update test for counts(all_outcomes=True)

* fix typo

* add missing all_outcomes=True kwarg in test

* modify ppre-commit file to circumvent error on pc

* small formatting changes

* revert change to pre-commit config

* Update doc/introduction/measurements.rst

Co-authored-by: antalszava <[email protected]>

* Update doc/introduction/measurements.rst

Co-authored-by: antalszava <[email protected]>

* Update doc/introduction/measurements.rst

Co-authored-by: antalszava <[email protected]>

* Update doc/introduction/measurements.rst

Co-authored-by: antalszava <[email protected]>

* Update pennylane/measurements.py

Co-authored-by: antalszava <[email protected]>

* Update pennylane/measurements.py

Co-authored-by: antalszava <[email protected]>

* Update doc/introduction/measurements.rst

Co-authored-by: antalszava <[email protected]>

* Update doc/introduction/measurements.rst

Co-authored-by: antalszava <[email protected]>

* update counts docstring

* Update doc/releases/changelog-dev.md

Co-authored-by: Albert Mitjans <[email protected]>

* temp modification to pre-commit yaml

* Switch to AllCounts implementation

* Update measurements.rst

* Add test for multiple measurement types with counts(all_outcomes=True)

* remove completed todo

* Revert "temp modification to pre-commit yaml"

This reverts commit 0ff4690.

* indent codeblock in measurements.rst

* Add AllCounts to updates from master

* Remove irrelevant exception handling in _samples_to_counts

* Fix too-many-boolean-expressions CodeFactor issue

* indent code block

* tidy up

* Update tests/test_measurements.py

Co-authored-by: Albert Mitjans <[email protected]>

* indent docstring code-block

* fix indentation

* black reformatting

Co-authored-by: Lillian Frederiksen <[email protected]>
Co-authored-by: antalszava <[email protected]>
Co-authored-by: Albert Mitjans <[email protected]>

Author: 4 people

doc/introduction/measurements.rst

@@ -171,8 +171,31 @@ And the result is:
 >>> circuit()
 {'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.
+a tuple is returned to provide differentiability for the outputs of QNodes.
 
 .. code-block:: python
 
@@ -184,7 +207,7 @@ a tensor of tensors is returned to provide differentiability for the outputs of
         return qml.expval(qml.PauliZ(0)),qml.expval(qml.PauliZ(1)), qml.counts()
 
 >>> circuit()
-(-0.036, 0.036, {'01': 482, '10': 518})
+(-0.036, 0.036, {'01': 482, '10': 518}) 
 
 Probability
 -----------

doc/releases/changelog-dev.md

@@ -210,6 +210,22 @@
   False
   ```
 
+* 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. [(#2889)](https://github.com/PennyLaneAI/pennylane/pull/2889)
+  
+  ```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(all_outcomes=True)
+  >>> result = circuit()
+  >>> print(result)
+  {'00': 495, '01': 0, '10': 0,  '11': 505}
+  ```
+  
 * Internal use of in-place inversion is eliminated in preparation for its deprecation.
   [(#2965)](https://github.com/PennyLaneAI/pennylane/pull/2965)
 
@@ -332,6 +348,7 @@ Ankit Khandelwal,
 Korbinian Kottmann,
 Christina Lee,
 Meenu Kumari,
+Lillian Marie Austin Frederiksen,
 Albert Mitjans Coma,
 Rashid N H M,
 Zeyue Niu,

pennylane/_device.py

@@ -744,9 +744,15 @@ def batch_transform(self, circuit):
         elif (
             len(circuit._obs_sharing_wires) > 0
             and not hamiltonian_in_obs
-            and qml.measurements.Sample not in return_types
-            and qml.measurements.Probability not in return_types
-            and qml.measurements.Counts not in return_types
+            and all(
+                t not in return_types
+                for t in [
+                    qml.measurements.Sample,
+                    qml.measurements.Probability,
+                    qml.measurements.Counts,
+                    qml.measurements.AllCounts,
+                ]
+            )
         ):
             # Check for case of non-commuting terms and that there are no Hamiltonians
             # TODO: allow for Hamiltonians in list of observables as well.

pennylane/_qubit_device.py

@@ -32,6 +32,7 @@
 
 from pennylane.measurements import (
     Counts,
+    AllCounts,
     Expectation,
     MeasurementProcess,
     MutualInfo,
@@ -319,7 +320,9 @@ def execute(self, circuit, **kwargs):
             self._samples = self.generate_samples()
 
         ret_types = [m.return_type for m in circuit.measurements]
-        counts_exist = any(ret is qml.measurements.Counts for ret in ret_types)
+        counts_exist = any(
+            ret in (qml.measurements.Counts, qml.measurements.AllCounts) for ret in ret_types
+        )
 
         # compute the required statistics
         if not self.analytic and self._shot_vector is not None:
@@ -448,7 +451,9 @@ def shot_vec_statistics(self, circuit):
         s1 = 0
 
         ret_types = [m.return_type for m in circuit.measurements]
-        counts_exist = any(ret is qml.measurements.Counts for ret in ret_types)
+        counts_exist = any(
+            ret in (qml.measurements.Counts, qml.measurements.AllCounts) for ret in ret_types
+        )
         single_measurement = len(circuit.measurements) == 1
 
         for shot_tuple in self._shot_vector:
@@ -545,7 +550,7 @@ def _multi_meas_with_counts_shot_vec(self, circuit, shot_tuple, r):
                 else:
                     result = r_[idx]
 
-                if not circuit.observables[idx2].return_type is Counts:
+                if not circuit.observables[idx2].return_type in (Counts, AllCounts):
                     result = self._asarray(result.T)
 
                 result_group.append(result)
@@ -745,7 +750,7 @@ def statistics(self, observables, shot_range=None, bin_size=None, circuit=None):
                     self.sample(obs, shot_range=shot_range, bin_size=bin_size, counts=False)
                 )
 
-            elif obs.return_type is Counts:
+            elif obs.return_type in (Counts, AllCounts):
                 results.append(
                     self.sample(obs, shot_range=shot_range, bin_size=bin_size, counts=True)
                 )
@@ -896,7 +901,7 @@ def statistics_new(self, observables, shot_range=None, bin_size=None):
                 samples = self.sample(obs, shot_range=shot_range, bin_size=bin_size, counts=False)
                 result = qml.math.squeeze(samples)
 
-            elif obs.return_type is Counts:
+            elif obs.return_type in (Counts, AllCounts):
                 result = self.sample(obs, shot_range=shot_range, bin_size=bin_size, counts=True)
 
             elif obs.return_type is Probability:
@@ -1598,6 +1603,75 @@ 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, which is set when
+        calling measurements.counts by setting the kwarg all_outcomes (bool). Per default,
+        the dictionary will only contain the observed outcomes. Optionally (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 is AllCounts, 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 is AllCounts:
+                outcomes = self.generate_basis_states(num_wires)
+                outcomes = ["".join([str(o.item()) for o in outcome]) for outcome in outcomes]
+        elif obs.return_type is AllCounts:
+            outcomes = qml.eigvals(obs)
+
+        # 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.
 
@@ -1620,28 +1694,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
@@ -1684,16 +1736,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/devices/default_mixed.py

@@ -34,7 +34,7 @@
     Snapshot,
 )
 from pennylane import numpy as np
-from pennylane.measurements import Counts, MutualInfo, Sample, State, VnEntropy
+from pennylane.measurements import Counts, AllCounts, MutualInfo, Sample, State, VnEntropy
 from pennylane.operation import Channel
 from pennylane.ops.qubit.attributes import diagonal_in_z_basis
 from pennylane.wires import Wires
@@ -579,7 +579,7 @@ def execute(self, circuit, **kwargs):
                     # Assumed to only be allowed if it's the only measurement.
                     self.measured_wires = []
                     return super().execute(circuit, **kwargs)
-                if obs.return_type in (Sample, Counts):
+                if obs.return_type in (Sample, Counts, AllCounts):
                     if obs.name == "Identity" and obs.wires in (qml.wires.Wires([]), self.wires):
                         # Sample, Counts: Readout error applied to all device wires when wires
                         # not specified or all wires specified.

pennylane/interfaces/autograd.py

@@ -111,7 +111,10 @@ def _execute(
 
     for i, r in enumerate(res):
 
-        if any(m.return_type is qml.measurements.Counts for m in tapes[i].measurements):
+        if any(
+            m.return_type in (qml.measurements.Counts, qml.measurements.AllCounts)
+            for m in tapes[i].measurements
+        ):
             continue
 
         if isinstance(r, np.ndarray):

pennylane/interfaces/jax.py

@@ -175,7 +175,10 @@ def array_if_not_counts(tape, r):
         dictionaries. JAX NumPy arrays don't support dictionaries."""
         return (
             jnp.array(r)
-            if not any(m.return_type is qml.measurements.Counts for m in tape.measurements)
+            if not any(
+                m.return_type in (qml.measurements.Counts, qml.measurements.AllCounts)
+                for m in tape.measurements
+            )
             else r
         )
 

pennylane/interfaces/tensorflow.py

@@ -91,7 +91,10 @@ def execute(tapes, device, execute_fn, gradient_fn, gradient_kwargs, _n=1, max_d
     for i, tape in enumerate(tapes):
         # convert output to TensorFlow tensors
 
-        if any(m.return_type is qml.measurements.Counts for m in tape.measurements):
+        if any(
+            m.return_type in (qml.measurements.Counts, qml.measurements.AllCounts)
+            for m in tape.measurements
+        ):
             continue
 
         if isinstance(res[i], np.ndarray):

pennylane/interfaces/torch.py

@@ -97,7 +97,10 @@ def forward(ctx, kwargs, *parameters):  # pylint: disable=arguments-differ
                 # For backwards compatibility, we flatten ragged tape outputs
                 r = np.hstack(r)
 
-            if any(m.return_type is qml.measurements.Counts for m in ctx.tapes[i].measurements):
+            if any(
+                m.return_type in (qml.measurements.Counts, qml.measurements.AllCounts)
+                for m in ctx.tapes[i].measurements
+            ):
                 continue
 
             if isinstance(r, (list, tuple)):