Repository: PennyLaneAI/pennylane
Pull request #3312 (2022-11-10 14:07:25)

CMP #9: Add `Var` MP class
1. Create `_Variance` class (name will be changed to `Variance` after the whole refactor).
2. Return `_Variance` instance in `qml.var`.
3. Copy code from `QubitDevice` into `_Variance.process` + improve code style.
4. Add `teardown_method` to the `test_var.py` tests that checks that `QubitDevice.var()` and `_Variance.process` 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.
codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3312?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#3312](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3312?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (b9d3202) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/4dd6088e83c5eed49abd031808ded21c4f841adb?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (4dd6088) will **decrease** coverage by `0.12%`.
> The diff coverage is `80.68%`.

```diff
@@            Coverage Diff             @@
##           master    #3312      +/-   ##
==========================================
- Coverage   99.70%   99.57%   -0.13%     
==========================================
  Files         299      299              
  Lines       26122    26280     +158     
==========================================
+ Hits        26045    26169     +124     
- Misses         77      111      +34     
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3312?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/measurements/probs.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3312/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21lYXN1cmVtZW50cy9wcm9icy5weQ==) | `74.00% <69.76%> (-26.00%)` | :arrow_down: |
| [pennylane/measurements/var.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3312/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21lYXN1cmVtZW50cy92YXIucHk=) | `83.33% <81.39%> (-16.67%)` | :arrow_down: |
| [pennylane/\_qubit\_device.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3312/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL19xdWJpdF9kZXZpY2UucHk=) | `99.66% <100.00%> (+<0.01%)` | :arrow_up: |
| [pennylane/gradients/parameter\_shift.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3312/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2dyYWRpZW50cy9wYXJhbWV0ZXJfc2hpZnQucHk=) | `100.00% <100.00%> (ø)` | |
| [pennylane/measurements/\_\_init\_\_.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3312/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21lYXN1cmVtZW50cy9fX2luaXRfXy5weQ==) | `100.00% <100.00%> (ø)` | |
| [pennylane/measurements/measurements.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3312/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21lYXN1cmVtZW50cy9tZWFzdXJlbWVudHMucHk=) | `100.00% <100.00%> (ø)` | |
| [pennylane/measurements/sample.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3312/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21lYXN1cmVtZW50cy9zYW1wbGUucHk=) | `100.00% <100.00%> (ø)` | |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)


Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 89f600807..fd196be83 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -15,7 +15,10 @@
   * Add `_Sample` class.
     [#3288](https://github.com/PennyLaneAI/pennylane/pull/3288)
 
-  * Add `_Probs` class.
+  * Add `_Var` class.
+    [#3312](https://github.com/PennyLaneAI/pennylane/pull/3312)
+
+  * Add `_Probability` class.
     [#3287](https://github.com/PennyLaneAI/pennylane/pull/3287)
 
   * Add `_Counts` class.
diff --git a/pennylane/measurements/expval.py b/pennylane/measurements/expval.py
index b1e7f1d27..03ed7ef98 100644
--- a/pennylane/measurements/expval.py
+++ b/pennylane/measurements/expval.py
@@ -104,10 +104,3 @@ class _Expectation(SampleMeasurement, StateMeasurement):
         prob = qml.probs(wires=permuted_wires).process_state(state=state, wire_order=wire_order)
         # In case of broadcasting, `prob` has two axes and this is a matrix-vector product
         return qml.math.dot(prob, eigvals)
-
-    def _permute_wires(self, wires: Wires):
-        wire_map = dict(zip(wires, range(len(wires))))
-        ordered_obs_wire_lst = sorted(self.wires.tolist(), key=lambda label: wire_map[label])
-        mapped_wires = [wire_map[w] for w in self.wires]
-        permutation = qml.math.argsort(mapped_wires)  # extract permutation via argsort
-        return Wires([ordered_obs_wire_lst[index] for index in permutation])
diff --git a/pennylane/measurements/measurements.py b/pennylane/measurements/measurements.py
index 727564bd4..7d3dba20c 100644
--- a/pennylane/measurements/measurements.py
+++ b/pennylane/measurements/measurements.py
@@ -693,6 +693,60 @@ class MeasurementProcess:
             new_measurement._wires = Wires([wire_map.get(wire, wire) for wire in self.wires])
         return new_measurement
 
+    def _permute_wires(self, wires: Wires):
+        r"""Given an observable which acts on multiple wires, permute the wires to
+          be consistent with the device wire order.
+
+          Suppose we are given an observable :math:`\hat{O} = \Identity \otimes \Identity \otimes \hat{Z}`.
+          This observable can be represented in many ways:
+
+        .. code-block:: python
+
+              O_1 = qml.Identity(wires=0) @ qml.Identity(wires=1) @ qml.PauliZ(wires=2)
+              O_2 = qml.PauliZ(wires=2) @ qml.Identity(wires=0) @ qml.Identity(wires=1)
+
+          Notice that while the explicit tensor product matrix representation of :code:`O_1` and :code:`O_2` is
+          different, the underlying operator is identical due to the wire labelling (assuming the labels in
+          ascending order are {0,1,2}). If we wish to compute the expectation value of such an observable, we must
+          ensure it is identical in both cases. To facilitate this, we permute the wires in our state vector such
+          that they are consistent with this swapping of order in the tensor observable.
+
+        .. code-block:: python
+
+              >>> print(O_1.wires)
+              <Wires = [0, 1, 2]>
+              >>> print(O_2.wires)
+              <Wires = [2, 0, 1]>
+
+          We might naively think that we must permute our state vector to match the wire order of our tensor observable.
+          We must be careful and realize that the wire order of the terms in the tensor observable DOES NOT match the
+          permutation of the terms themselves. As an example we directly compare :code:`O_1` and :code:`O_2`:
+
+          The first term in :code:`O_1` (:code:`qml.Identity(wires=0)`) became the second term in :code:`O_2`.
+          By similar comparison we see that each term in the tensor product was shifted one position forward
+          (i.e 0 --> 1, 1 --> 2, 2 --> 0). The wires in our permuted quantum state should follow their respective
+          terms in the tensor product observable.
+
+          Thus, the correct wire ordering should be :code:`permuted_wires = <Wires = [1, 2, 0]>`. But if we had
+          taken the naive approach we would have permuted our state according to
+          :code:`permuted_wires = <Wires = [2, 0, 1]>` which is NOT correct.
+
+          This function uses the observable wires and the global device wire ordering in order to determine the
+          permutation of the wires in the observable required such that if our quantum state vector is
+          permuted accordingly then the amplitudes of the state will match the matrix representation of the observable.
+
+          Args:
+              observable (Observable): the observable whose wires are to be permuted.
+
+          Returns:
+              permuted_wires (Wires): permuted wires object
+        """
+        wire_map = dict(zip(wires, range(len(wires))))
+        ordered_obs_wire_lst = sorted(self.wires.tolist(), key=lambda label: wire_map[label])
+        mapped_wires = [wire_map[w] for w in self.wires]
+        permutation = qml.math.argsort(mapped_wires)  # extract permutation via argsort
+        return Wires([ordered_obs_wire_lst[index] for index in permutation])
+
 
 class SampleMeasurement(MeasurementProcess, ABC):
     """Sample-based measurement process."""
diff --git a/pennylane/measurements/var.py b/pennylane/measurements/var.py
index 1001e5585..9bc973a8a 100644
--- a/pennylane/measurements/var.py
+++ b/pennylane/measurements/var.py
@@ -16,10 +16,14 @@
 This module contains the qml.var measurement.
 """
 import warnings
+from typing import Sequence, Tuple
 
+import pennylane as qml
 from pennylane.operation import Operator
+from pennylane.ops import Projector
+from pennylane.wires import Wires
 
-from .measurements import MeasurementProcess, Variance
+from .measurements import SampleMeasurement, StateMeasurement, Variance
 
 
 def var(op: Operator):
@@ -51,4 +55,54 @@ def var(op: Operator):
     """
     if not op.is_hermitian:
         warnings.warn(f"{op.name} might not be hermitian.")
-    return MeasurementProcess(Variance, obs=op)
+    return _Variance(Variance, obs=op)
+
+
+class _Variance(SampleMeasurement, StateMeasurement):
+    """Measurement process that computes the variance of the supplied observable."""
+
+    def process_samples(
+        self,
+        samples: Sequence[complex],
+        wire_order: Wires,
+        shot_range: Tuple[int] = None,
+        bin_size: int = None,
+    ):
+        if isinstance(self.obs, Projector):
+            # branch specifically to handle the projector observable
+            idx = int("".join(str(i) for i in self.obs.parameters[0]), 2)
+            # we use ``self.wires`` instead of ``self.obs`` because the observable was
+            # already applied before the sampling
+            probs = qml.probs(wires=self.wires).process_samples(
+                samples=samples, wire_order=wire_order, shot_range=shot_range, bin_size=bin_size
+            )
+            return probs[idx] - probs[idx] ** 2
+
+        # estimate the variance
+        samples = qml.sample(op=self.obs).process_samples(
+            samples=samples, wire_order=wire_order, shot_range=shot_range, bin_size=bin_size
+        )
+        # With broadcasting, we want to take the variance over axis 1, which is the -1st/-2nd with/
+        # without bin_size. Without broadcasting, axis 0 is the -1st/-2nd with/without bin_size
+        axis = -1 if bin_size is None else -2
+        # TODO: do we need to squeeze here? Maybe remove with new return types
+        return qml.math.squeeze(qml.math.var(samples, axis=axis))
+
+    def process_state(self, state: Sequence[complex], wire_order: Wires):
+        if isinstance(self.obs, Projector):
+            # branch specifically to handle the projector observable
+            idx = int("".join(str(i) for i in self.obs.parameters[0]), 2)
+            # we use ``self.wires`` instead of ``self.obs`` because the observable was
+            # already applied to the state
+            probs = qml.probs(wires=self.wires).process_state(state=state, wire_order=wire_order)
+            return probs[idx] - probs[idx] ** 2
+
+        eigvals = qml.math.asarray(self.obs.eigvals(), dtype=float)
+
+        # the probability vector must be permuted to account for the permuted wire order of the observable
+        new_obs_wires = self._permute_wires(self.obs.wires)
+        # we use ``wires`` instead of ``op`` because the observable was
+        # already applied to the state
+        prob = qml.probs(wires=new_obs_wires).process_state(state=state, wire_order=wire_order)
+        # In case of broadcasting, `prob` has two axes and these are a matrix-vector products
+        return qml.math.dot(prob, (eigvals**2)) - qml.math.dot(prob, eigvals) ** 2


Test patch:
diff --git a/tests/measurements/test_var.py b/tests/measurements/test_var.py
index 3e33cc3ff..cbf6919a1 100644
--- a/tests/measurements/test_var.py
+++ b/tests/measurements/test_var.py
@@ -19,13 +19,39 @@ import pennylane as qml
 from pennylane.measurements import Variance
 
 
+# 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
+    state = device._state
+    call_args_list = list(spy.call_args_list)
+    for call_args in call_args_list:
+        obs = call_args.args[1]
+        shot_range, bin_size = (
+            call_args.kwargs["shot_range"],
+            call_args.kwargs["bin_size"],
+        )
+        meas = qml.var(op=obs)
+        old_res = device.var(obs, shot_range, bin_size)
+        if samples is not None:
+            new_res = meas.process_samples(
+                samples=samples, wire_order=device.wires, shot_range=shot_range, bin_size=bin_size
+            )
+        else:
+            new_res = meas.process_state(state=state, wire_order=device.wires)
+        assert qml.math.allequal(old_res, new_res)
+
+
 class TestVar:
     """Tests for the var function"""
 
+    @pytest.mark.parametrize("shots", [None, 1000, [1000, 10000]])
     @pytest.mark.parametrize("r_dtype", [np.float32, np.float64])
-    def test_value(self, tol, r_dtype):
+    def test_value(self, tol, r_dtype, mocker, shots):
         """Test that the var function works"""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit", wires=2, shots=shots)
+        spy = mocker.spy(qml.QubitDevice, "var")
         dev.R_DTYPE = r_dtype
 
         @qml.qnode(dev, diff_method="parameter-shift")
@@ -35,15 +61,20 @@ class TestVar:
 
         x = 0.54
         res = circuit(x)
-        expected = np.sin(x) ** 2
+        expected = [np.sin(x) ** 2, np.sin(x) ** 2] if isinstance(shots, list) else np.sin(x) ** 2
+        atol = tol if shots is None else 0.05
+        rtol = 0 if shots is None else 0.05
 
-        assert np.allclose(res, expected, atol=tol, rtol=0)
+        assert np.allclose(res, expected, atol=atol, rtol=rtol)
         assert res.dtype == r_dtype
 
-    def test_not_an_observable(self):
+        custom_measurement_process(dev, spy)
+
+    def test_not_an_observable(self, mocker):
         """Test that a UserWarning is raised if the provided
         argument might not be hermitian."""
         dev = qml.device("default.qubit", wires=2)
+        spy = mocker.spy(qml.QubitDevice, "var")
 
         @qml.qnode(dev)
         def circuit():
@@ -53,9 +84,12 @@ class TestVar:
         with pytest.warns(UserWarning, match="Prod might not be hermitian."):
             _ = circuit()
 
-    def test_observable_return_type_is_variance(self):
+        custom_measurement_process(dev, spy)
+
+    def test_observable_return_type_is_variance(self, mocker):
         """Test that the return type of the observable is :attr:`ObservableReturnTypes.Variance`"""
         dev = qml.device("default.qubit", wires=2)
+        spy = mocker.spy(qml.QubitDevice, "var")
 
         @qml.qnode(dev)
         def circuit():
@@ -65,6 +99,8 @@ class TestVar:
 
         circuit()
 
+        custom_measurement_process(dev, spy)
+
     @pytest.mark.parametrize(
         "obs",
         [qml.PauliZ(0), qml.Hermitian(np.diag([1, 2]), 0), qml.Hermitian(np.diag([1.0, 2.0]), 0)],
@@ -93,3 +129,23 @@ class TestVar:
         shot_vector = (1, 2, 3)
         dev = qml.device("default.qubit", wires=3, shots=shot_vector)
         assert res.shape(dev) == (len(shot_vector),)
+
+    @pytest.mark.parametrize("shots", [None, 1000, [1000, 10000]])
+    def test_projector_var(self, shots, mocker):
+        """Tests that the variance of a ``Projector`` object is computed correctly."""
+        dev = qml.device("default.qubit", wires=3, shots=shots)
+        spy = mocker.spy(qml.QubitDevice, "var")
+
+        basis_state = np.array([0, 0, 0])
+
+        @qml.qnode(dev)
+        def circuit():
+            qml.Hadamard(0)
+            return qml.var(qml.Projector(basis_state, wires=range(3)))
+
+        res = circuit()
+        expected = [0.25, 0.25] if isinstance(shots, list) else 0.25
+
+        assert np.allclose(res, expected, atol=0.02, rtol=0.02)
+
+        custom_measurement_process(dev, spy)