Repository: PennyLaneAI/pennylane
Pull request #3287 (2022-11-08 08:52:48)

CMP #7: Add `Probability` MP class
1. Create `_Probability ` class (name will be changed to `Probability ` after the whole refactor).
2. Return `_Probability ` instance in `qml.probs`.
3. Copy code from `QubitDevice` into `_Probability` + improve code style.
4. Add `teardown_method` to the `test_probs.py` tests that checks that `QubitDevice.probability` and `_Probability.process` / `_Probability.process_state` return the same values for all tests.

github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [`doc/releases/changelog-dev.md`](/PennyLaneAI/pennylane/blob/master/doc/releases/changelog-dev.md) with:
* A one-to-two sentence description of the change. You may include a small working example for new features.
* A link back to this PR.
* Your name (or GitHub username) in the contributors section.

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index d72f80ad2..e99dfa660 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -12,6 +12,9 @@
   * Add `_Sample` class.
     [#3288](https://github.com/PennyLaneAI/pennylane/pull/3288)
 
+  * Add `_Probs` class.
+    [#3287](https://github.com/PennyLaneAI/pennylane/pull/3287)
+
   * Add `_Counts` class.
     [#3292](https://github.com/PennyLaneAI/pennylane/pull/3292)
 
diff --git a/pennylane/measurements/probs.py b/pennylane/measurements/probs.py
index 713d6e632..9ba4d1053 100644
--- a/pennylane/measurements/probs.py
+++ b/pennylane/measurements/probs.py
@@ -15,9 +15,12 @@
 """
 This module contains the qml.probs measurement.
 """
+from typing import Sequence, Tuple
+
 import pennylane as qml
+from pennylane.wires import Wires
 
-from .measurements import MeasurementProcess, Probability
+from .measurements import Probability, SampleMeasurement, StateMeasurement
 
 
 def probs(wires=None, op=None):
@@ -108,5 +111,161 @@ def probs(wires=None, op=None):
                 "Cannot specify the wires to probs if an observable is "
                 "provided. The wires for probs will be determined directly from the observable."
             )
-        return MeasurementProcess(Probability, wires=qml.wires.Wires(wires))
-    return MeasurementProcess(Probability, obs=op)
+        return _Probability(Probability, wires=qml.wires.Wires(wires))
+    return _Probability(Probability, obs=op)
+
+
+# TODO: Make public when removing the ObservableReturnTypes enum
+class _Probability(SampleMeasurement, StateMeasurement):
+    """Measurement process that computes the probability of each computational basis state."""
+
+    def process_samples(
+        self,
+        samples: Sequence[complex],
+        wire_order: Wires,
+        shot_range: Tuple[int] = None,
+        bin_size: int = None,
+    ):
+        wire_map = dict(zip(wire_order, range(len(wire_order))))
+        mapped_wires = [wire_map[w] for w in self.wires]
+        if shot_range is not None:
+            # Indexing corresponds to: (potential broadcasting, shots, wires). Note that the last
+            # colon (:) is required because shots is the second-to-last axis and the
+            # Ellipsis (...) otherwise would take up broadcasting and shots axes.
+            samples = samples[..., slice(*shot_range), :]
+
+        if mapped_wires:
+            # if wires are provided, then we only return samples from those wires
+            samples = samples[..., mapped_wires]
+
+        num_wires = qml.math.shape(samples)[-1]
+        # convert samples from a list of 0, 1 integers, to base 10 representation
+        powers_of_two = 2 ** qml.math.arange(num_wires)[::-1]
+        indices = samples @ powers_of_two
+
+        # `samples` typically has two axes ((shots, wires)) but can also have three with
+        # broadcasting ((batch_size, shots, wires)) so that we simply read out the batch_size.
+        batch_size = samples.shape[0] if qml.math.ndim(samples) == 3 else None
+        dim = 2**num_wires
+        # count the basis state occurrences, and construct the probability vector
+        new_bin_size = bin_size or samples.shape[-2]
+        new_shape = (-1, new_bin_size) if batch_size is None else (batch_size, -1, new_bin_size)
+        indices = indices.reshape(new_shape)
+        prob = self._count_samples(indices, batch_size, dim)
+        return qml.math.squeeze(prob) if bin_size is None else prob
+
+    def process_state(self, state: Sequence[complex], wire_order: Wires):
+        num_wires = len(wire_order)
+        dim = 2**num_wires
+        # Compute batch_size
+        expected_shape = [2] * num_wires
+        expected_size = dim
+        size = qml.math.size(state)
+        batch_size = (
+            size // expected_size
+            if qml.math.ndim(state) > len(expected_shape) or size > expected_size
+            else None
+        )
+        flat_state = qml.math.reshape(
+            state, (batch_size, dim) if batch_size is not None else (dim,)
+        )
+        real_state = qml.math.real(flat_state)
+        imag_state = qml.math.imag(flat_state)
+        return self.marginal_prob(real_state**2 + imag_state**2, wire_order, batch_size)
+
+    @staticmethod
+    def _count_samples(indices, batch_size, dim):
+        """Count the occurences of sampled indices and convert them to relative
+        counts in order to estimate their occurence probability."""
+        num_bins, bin_size = indices.shape[-2:]
+        if batch_size is None:
+            prob = qml.math.zeros((dim, num_bins), dtype="float64")
+            # count the basis state occurrences, and construct the probability vector for each bin
+            for b, idx in enumerate(indices):
+                basis_states, counts = qml.math.unique(idx, return_counts=True)
+                prob[basis_states, b] = counts / bin_size
+
+            return prob
+
+        prob = qml.math.zeros((batch_size, dim, num_bins), dtype="float64")
+        indices = indices.reshape((batch_size, num_bins, bin_size))
+
+        # count the basis state occurrences, and construct the probability vector
+        # for each bin and broadcasting index
+        for i, _indices in enumerate(indices):  # First iterate over broadcasting dimension
+            for b, idx in enumerate(_indices):  # Then iterate over bins dimension
+                basis_states, counts = qml.math.unique(idx, return_counts=True)
+                prob[i, basis_states, b] = counts / bin_size
+
+        return prob
+
+    def marginal_prob(self, prob, wire_order, batch_size):
+        r"""Return the marginal probability of the computational basis
+        states by summing the probabiliites on the non-specified wires.
+
+        If no wires are specified, then all the basis states representable by
+        the device are considered and no marginalization takes place.
+
+        .. note::
+
+            If the provided wires are not in the order as they appear on the device,
+            the returned marginal probabilities take this permutation into account.
+
+            For example, if the addressable wires on this device are ``Wires([0, 1, 2])`` and
+            this function gets passed ``wires=[2, 0]``, then the returned marginal
+            probability vector will take this 'reversal' of the two wires
+            into account:
+
+            .. math::
+
+                \mathbb{P}^{(2, 0)}
+                            = \left[
+                               |00\rangle, |10\rangle, |01\rangle, |11\rangle
+                              \right]
+
+        Args:
+            prob: The probabilities to return the marginal probabilities
+                for
+            wire_order (Iterable[Number, str], Number, str, Wires): wires to return
+                marginal probabilities for. Wires not provided
+                are traced out of the system.
+
+        Returns:
+            array[float]: array of the resulting marginal probabilities.
+        """
+        # TODO: Add when ``qml.probs()`` is supported
+        # if self.wires == Wires([]):
+        #     # no need to marginalize
+        #     return prob
+
+        # determine which subsystems are to be summed over
+        inactive_wires = Wires.unique_wires([wire_order, self.wires])
+
+        # translate to wire labels used by device
+        wire_map = dict(zip(wire_order, range(len(wire_order))))
+        mapped_wires = [wire_map[w] for w in self.wires]
+        inactive_wires = [wire_map[w] for w in inactive_wires]
+
+        # reshape the probability so that each axis corresponds to a wire
+        num_device_wires = len(wire_order)
+        shape = [2] * num_device_wires
+        if batch_size is not None:
+            shape.insert(0, batch_size)
+            inactive_wires = [idx + 1 for idx in inactive_wires]
+        # prob now is reshaped to have self.num_wires+1 axes in the case of broadcasting
+        prob = qml.math.reshape(prob, shape)
+
+        # sum over all inactive wires
+        flat_shape = (-1,) if batch_size is None else (batch_size, -1)
+        prob = qml.math.reshape(qml.math.sum(prob, axis=tuple(inactive_wires)), flat_shape)
+
+        # The wires provided might not be in consecutive order (i.e., wires might be [2, 0]).
+        # If this is the case, we must permute the marginalized probability so that
+        # it corresponds to the orders of the wires passed.
+        num_wires = len(mapped_wires)
+        basis_states = qml.QubitDevice.generate_basis_states(num_wires)
+        basis_states = basis_states[:, qml.math.argsort(qml.math.argsort(mapped_wires))]
+
+        powers_of_two = 2 ** qml.math.arange(num_wires)[::-1]
+        perm = basis_states @ powers_of_two
+        return prob[:, perm] if batch_size is not None else prob[perm]


Test patch:
diff --git a/tests/math/test_entropies_math.py b/tests/math/test_entropies_math.py
index 39764ab53..9e837c277 100644
--- a/tests/math/test_entropies_math.py
+++ b/tests/math/test_entropies_math.py
@@ -18,7 +18,6 @@ import pytest
 
 import pennylane as qml
 from pennylane import numpy as np
-from tests.math.test_density_matrices import single_wires_list
 
 pytestmark = pytest.mark.all_interfaces
 
diff --git a/tests/measurements/test_probs.py b/tests/measurements/test_probs.py
index 55321db0d..febba6611 100644
--- a/tests/measurements/test_probs.py
+++ b/tests/measurements/test_probs.py
@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Unit tests for the probs module"""
-
 import numpy as np
 import pytest
 
@@ -21,12 +20,41 @@ from pennylane import numpy as pnp
 from pennylane.measurements import MeasurementProcess, MeasurementShapeError, Probability
 from pennylane.queuing import AnnotatedQueue
 
+
+# TODO: Remove this when new CustomMP are the default
+def custom_measurement_process(device, spy):
+    assert len(spy.call_args_list) > 0  # make sure method is mocked properly
+
+    samples = device._samples  # pylint: disable=protected-access
+    state = device._state  # pylint: disable=protected-access
+    call_args_list = list(spy.call_args_list)
+    for call_args in call_args_list:
+        wires, shot_range, bin_size = (
+            call_args.kwargs["wires"],
+            call_args.kwargs["shot_range"],
+            call_args.kwargs["bin_size"],
+        )
+        # no need to use op, because the observable has already been applied to ``dev._state``
+        meas = qml.probs(wires=wires)
+        old_res = device.probability(wires=wires, shot_range=shot_range, bin_size=bin_size)
+        if device.shots is None:
+            new_res = meas.process_state(state=state, wire_order=device.wires)
+        else:
+            new_res = meas.process_samples(
+                samples=samples,
+                wire_order=device.wires,
+                shot_range=shot_range,
+                bin_size=bin_size,
+            )
+        assert qml.math.allequal(old_res, new_res)
+
+
 # make the test deterministic
 np.random.seed(42)
 
 
 @pytest.fixture
-def init_state(scope="session"):
+def init_state():
     """Fixture that creates an initial state"""
 
     def _init_state(n):
@@ -38,6 +66,7 @@ def init_state(scope="session"):
     return _init_state
 
 
+# pylint: disable=redefined-outer-name
 class TestProbs:
     """Tests for the probs function"""
 
@@ -88,9 +117,11 @@ class TestProbs:
         assert isinstance(meas_proc, MeasurementProcess)
         assert meas_proc.return_type == Probability
 
-    def test_full_prob(self, init_state, tol):
+    def test_full_prob(self, init_state, tol, mocker):
         """Test that the correct probability is returned."""
         dev = qml.device("default.qubit", wires=4)
+        spy = mocker.spy(qml.QubitDevice, "probability")
+
         state = init_state(4)
 
         @qml.qnode(dev)
@@ -102,9 +133,13 @@ class TestProbs:
         expected = np.abs(state) ** 2
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
-    def test_marginal_prob(self, init_state, tol):
+        custom_measurement_process(dev, spy)
+
+    def test_marginal_prob(self, init_state, tol, mocker):
         """Test that the correct marginal probability is returned."""
         dev = qml.device("default.qubit", wires=4)
+        spy = mocker.spy(qml.QubitDevice, "probability")
+
         state = init_state(4)
 
         @qml.qnode(dev)
@@ -117,10 +152,13 @@ class TestProbs:
         expected = np.einsum("ijkl->jl", expected).flatten()
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
+        custom_measurement_process(dev, spy)
+
     def test_marginal_prob_more_wires(self, init_state, mocker, tol):
         """Test that the correct marginal probability is returned, when the
         states_to_binary method is used for probability computations."""
         dev = qml.device("default.qubit", wires=4)
+        spy_probs = mocker.spy(qml.QubitDevice, "probability")
         state = init_state(4)
 
         spy = mocker.spy(qml.QubitDevice, "states_to_binary")
@@ -138,9 +176,12 @@ class TestProbs:
 
         spy.assert_called_once()
 
-    def test_integration(self, tol):
+        custom_measurement_process(dev, spy_probs)
+
+    def test_integration(self, tol, mocker):
         """Test the probability is correct for a known state preparation."""
         dev = qml.device("default.qubit", wires=2)
+        spy = mocker.spy(qml.QubitDevice, "probability")
 
         @qml.qnode(dev)
         def circuit():
@@ -155,27 +196,51 @@ class TestProbs:
         expected = np.array([0.5, 0.5, 0, 0])
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
-    def test_integration_analytic_false(self, tol):
+        custom_measurement_process(dev, spy)
+
+    @pytest.mark.parametrize("shots", [100, [1, 10, 100]])
+    def test_integration_analytic_false(self, tol, mocker, shots):
         """Test the probability is correct for a known state preparation when the
         analytic attribute is set to False."""
-        dev = qml.device("default.qubit", wires=3, shots=1000)
+        dev = qml.device("default.qubit", wires=3, shots=shots)
+        spy = mocker.spy(qml.QubitDevice, "probability")
 
         @qml.qnode(dev)
         def circuit():
             qml.PauliX(0)
-            return qml.probs(wires=[0])
+            return qml.probs(wires=dev.wires)
 
         res = circuit()
-        expected = np.array([0, 1])
+        expected = np.array([0, 0, 0, 0, 1, 0, 0, 0])
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
+        custom_measurement_process(dev, spy)
+
+    @pytest.mark.parametrize("shots", [None, 100])
+    def test_batch_size(self, mocker, shots):
+        """Test the probability is correct for a batched input."""
+        dev = qml.device("default.qubit", wires=1, shots=shots)
+        spy = mocker.spy(qml.QubitDevice, "probability")
+
+        @qml.qnode(dev)
+        def circuit(x):
+            qml.RX(x, 0)
+            return qml.probs(wires=dev.wires)  # TODO: Use ``qml.probs()`` when supported
+
+        x = np.array([0, np.pi / 2])
+        res = circuit(x)
+        expected = [[1.0, 0.0], [0.5, 0.5]]
+        assert np.allclose(res, expected, atol=0.1, rtol=0.1)
+
+        custom_measurement_process(dev, spy)
+
     @pytest.mark.autograd
-    def test_numerical_analytic_diff_agree(self, init_state, tol):
+    def test_numerical_analytic_diff_agree(self, tol, mocker):
         """Test that the finite difference and parameter shift rule
         provide the same Jacobian."""
         w = 4
         dev = qml.device("default.qubit", wires=w)
-        state = init_state(w)
+        spy = mocker.spy(qml.QubitDevice, "probability")
 
         def circuit(x, y, z):
             for i in range(w):
@@ -205,10 +270,13 @@ class TestProbs:
         # Check that they agree up to numeric tolerance
         assert all(np.allclose(_rF, _rA, atol=tol, rtol=0) for _rF, _rA in zip(res_F, res_A))
 
+        custom_measurement_process(dev, spy)
+
     @pytest.mark.parametrize("hermitian", [1 / np.sqrt(2) * np.array([[1, 1], [1, -1]])])
-    def test_prob_generalize_param_one_qubit(self, hermitian, init_state, tol):
+    def test_prob_generalize_param_one_qubit(self, hermitian, tol, mocker):
         """Test that the correct probability is returned."""
         dev = qml.device("default.qubit", wires=1)
+        spy = mocker.spy(qml.QubitDevice, "probability")
 
         @qml.qnode(dev)
         def circuit(x):
@@ -229,10 +297,13 @@ class TestProbs:
 
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
+        custom_measurement_process(dev, spy)
+
     @pytest.mark.parametrize("hermitian", [1 / np.sqrt(2) * np.array([[1, 1], [1, -1]])])
-    def test_prob_generalize_param(self, hermitian, init_state, tol):
+    def test_prob_generalize_param(self, hermitian, tol, mocker):
         """Test that the correct probability is returned."""
         dev = qml.device("default.qubit", wires=3)
+        spy = mocker.spy(qml.QubitDevice, "probability")
 
         @qml.qnode(dev)
         def circuit(x, y):
@@ -258,10 +329,13 @@ class TestProbs:
         expected = np.einsum("ijk->i", expected).flatten()
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
+        custom_measurement_process(dev, spy)
+
     @pytest.mark.parametrize("hermitian", [1 / np.sqrt(2) * np.array([[1, 1], [1, -1]])])
-    def test_prob_generalize_param_multiple(self, hermitian, init_state, tol):
+    def test_prob_generalize_param_multiple(self, hermitian, tol, mocker):
         """Test that the correct probability is returned."""
         dev = qml.device("default.qubit", wires=3)
+        spy = mocker.spy(qml.QubitDevice, "probability")
 
         @qml.qnode(dev)
         def circuit(x, y):
@@ -298,11 +372,14 @@ class TestProbs:
         assert np.allclose(res[1], expected_1, atol=tol, rtol=0)
         assert np.allclose(res[2], expected_2, atol=tol, rtol=0)
 
+        custom_measurement_process(dev, spy)
+
     @pytest.mark.parametrize("hermitian", [1 / np.sqrt(2) * np.array([[1, 1], [1, -1]])])
     @pytest.mark.parametrize("wire", [0, 1, 2, 3])
-    def test_prob_generalize_initial_state(self, hermitian, wire, init_state, tol):
+    def test_prob_generalize_initial_state(self, hermitian, wire, init_state, tol, mocker):
         """Test that the correct probability is returned."""
         dev = qml.device("default.qubit", wires=4)
+        spy = mocker.spy(qml.QubitDevice, "probability")
         state = init_state(4)
 
         @qml.qnode(dev)
@@ -338,11 +415,14 @@ class TestProbs:
 
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
+        custom_measurement_process(dev, spy)
+
     @pytest.mark.parametrize("operation", [qml.PauliX, qml.PauliY, qml.Hadamard])
     @pytest.mark.parametrize("wire", [0, 1, 2, 3])
-    def test_operation_prob(self, operation, wire, init_state, tol):
+    def test_operation_prob(self, operation, wire, init_state, tol, mocker):
         "Test the rotated probability with different wires and rotating operations."
         dev = qml.device("default.qubit", wires=4)
+        spy = mocker.spy(qml.QubitDevice, "probability")
         state = init_state(4)
 
         @qml.qnode(dev)
@@ -378,10 +458,13 @@ class TestProbs:
 
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
+        custom_measurement_process(dev, spy)
+
     @pytest.mark.parametrize("observable", [(qml.PauliX, qml.PauliY)])
-    def test_observable_tensor_prob(self, observable, init_state, tol):
+    def test_observable_tensor_prob(self, observable, init_state, tol, mocker):
         "Test the rotated probability with a tensor observable."
         dev = qml.device("default.qubit", wires=4)
+        spy = mocker.spy(qml.QubitDevice, "probability")
         state = init_state(4)
 
         @qml.qnode(dev)
@@ -411,8 +494,10 @@ class TestProbs:
 
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
+        custom_measurement_process(dev, spy)
+
     @pytest.mark.parametrize("coeffs, obs", [([1, 1], [qml.PauliX(wires=0), qml.PauliX(wires=1)])])
-    def test_hamiltonian_error(self, coeffs, obs, init_state, tol):
+    def test_hamiltonian_error(self, coeffs, obs, init_state):
         "Test that an error is returned for hamiltonians."
         H = qml.Hamiltonian(coeffs, obs)
 
@@ -490,3 +575,40 @@ class TestProbs:
             "provided. The wires for probs will be determined directly from the observable.",
         ):
             circuit()
+
+    @pytest.mark.parametrize(
+        "wires, expected",
+        [
+            (
+                [0],
+                [
+                    [[0, 0, 0.5], [1, 1, 0.5]],
+                    [[0.5, 0.5, 0], [0.5, 0.5, 1]],
+                    [[0, 0.5, 1], [1, 0.5, 0]],
+                ],
+            ),
+            (
+                [0, 1],
+                [
+                    [[0, 0, 0], [0, 0, 0.5], [0.5, 0, 0], [0.5, 1, 0.5]],
+                    [[0.5, 0.5, 0], [0, 0, 0], [0, 0, 0], [0.5, 0.5, 1]],
+                    [[0, 0.5, 0.5], [0, 0, 0.5], [0.5, 0, 0], [0.5, 0.5, 0]],
+                ],
+            ),
+        ],
+    )
+    def test_estimate_probability_with_binsize_with_broadcasting(self, wires, expected):
+        """Tests the estimate_probability method with a bin size and parameter broadcasting"""
+        samples = np.array(
+            [
+                [[1, 0], [1, 1], [1, 1], [1, 1], [1, 1], [0, 1]],
+                [[0, 0], [1, 1], [1, 1], [0, 0], [1, 1], [1, 1]],
+                [[1, 0], [1, 1], [1, 1], [0, 0], [0, 1], [0, 0]],
+            ]
+        )
+
+        res = qml.probs(wires=wires).process_samples(
+            samples=samples, wire_order=wires, shot_range=None, bin_size=2
+        )
+
+        assert np.allclose(res, expected)