Repository: PennyLaneAI/pennylane
Pull request #5552 (2024-04-22 17:59:02)

Improve support for `snapshot` on `default.mixed`
**Context:** Improves support for performing snapshots on `default.mixed`.

**Description of the Change:** Adds the ability to perform arbitrary state-based measurements. 

**Benefits:** Device will support `qml.Snapshot` fully.

**Possible Drawbacks:**

**Related GitHub Issues:**


obliviateandsurrender: [sc-60987]
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5552?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
Attention: Patch coverage is `97.29730%` with `1 lines` in your changes are missing coverage. Please review.
> Project coverage is 99.68%. Comparing base [(`d03216e`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/d03216e6ce5407a7f28489aa4d4b23a9c4345fac?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`f4023a2`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5552?dropdown=coverage&src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
> Report is 6 commits behind head on master.

| [Files](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5552?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Patch % | Lines |
|---|---|---|
| [pennylane/devices/default\_mixed.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5552?src=pr&el=tree&filepath=pennylane%2Fdevices%2Fdefault_mixed.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RldmljZXMvZGVmYXVsdF9taXhlZC5weQ==) | 97.29% | [1 Missing :warning: ](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5552?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master    #5552      +/-   ##
==========================================
+ Coverage   99.67%   99.68%   +0.01%     
==========================================
  Files         407      409       +2     
  Lines       38019    37939      -80     
==========================================
- Hits        37895    37820      -75     
+ Misses        124      119       -5     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5552?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).

astralcai: @obliviateandsurrender Hi, we have a feature freeze for the v0.36 release starting next Monday. When you feel that this PR is ready to be merged, please label it merge-ready. I will be manually managing the merge queue over the next few days to make sure everything gets merged in time for feature freeze without wasting too much time on repeated CI runs.

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index ecb9a8b46..ce323abc7 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -373,6 +373,9 @@
   [(#5514)](https://github.com/PennyLaneAI/pennylane/pull/5514)
   [(#5530)](https://github.com/PennyLaneAI/pennylane/pull/5530)
 
+* `default.mixed` now supports arbitrary state-based measurements with `qml.Snapshot`.
+  [(#5552)](https://github.com/PennyLaneAI/pennylane/pull/5552)
+
 * Replaced `cache_execute` with an alternate implementation based on `@transform`.
   [(#5318)](https://github.com/PennyLaneAI/pennylane/pull/5318)
 
diff --git a/pennylane/devices/default_mixed.py b/pennylane/devices/default_mixed.py
index 55376174f..add4c7f7a 100644
--- a/pennylane/devices/default_mixed.py
+++ b/pennylane/devices/default_mixed.py
@@ -35,7 +35,18 @@ from pennylane import (
     StatePrep,
     Snapshot,
 )
-from pennylane.measurements import CountsMP, MutualInfoMP, SampleMP, StateMP, VnEntropyMP, PurityMP
+from pennylane.measurements import (
+    CountsMP,
+    MutualInfoMP,
+    SampleMP,
+    StateMP,
+    VnEntropyMP,
+    PurityMP,
+    DensityMatrixMP,
+    ExpectationMP,
+    VarianceMP,
+    ProbabilityMP,
+)
 from pennylane.operation import Channel
 from pennylane.ops.qubit.attributes import diagonal_in_z_basis
 from pennylane.wires import Wires
@@ -582,18 +593,82 @@ class DefaultMixed(QubitDevice):
             self._state = qnp.asarray(rho, dtype=self.C_DTYPE)
             self._pre_rotated_state = self._state
 
+    def _snapshot_measurements(self, density_matrix, measurement):
+        """Perform state-based snapshot measurement"""
+        meas_wires = measurement.wires
+
+        pre_rotated_state = self._state
+        if isinstance(measurement, (ProbabilityMP, ExpectationMP, VarianceMP)):
+            for diag_gate in measurement.diagonalizing_gates():
+                self._apply_operation(diag_gate)
+
+        if isinstance(measurement, (StateMP, DensityMatrixMP)):
+            map_wires = self.map_wires(meas_wires)
+            snap_result = qml.math.reduce_dm(
+                density_matrix, indices=map_wires, c_dtype=self.C_DTYPE
+            )
+
+        elif isinstance(measurement, PurityMP):
+            map_wires = self.map_wires(meas_wires)
+            snap_result = qml.math.purity(density_matrix, indices=map_wires, c_dtype=self.C_DTYPE)
+
+        elif isinstance(measurement, ProbabilityMP):
+            snap_result = self.analytic_probability(wires=meas_wires)
+
+        elif isinstance(measurement, ExpectationMP):
+            eigvals = self._asarray(measurement.obs.eigvals(), dtype=self.R_DTYPE)
+            probs = self.analytic_probability(wires=meas_wires)
+            snap_result = self._dot(probs, eigvals)
+
+        elif isinstance(measurement, VarianceMP):
+            eigvals = self._asarray(measurement.obs.eigvals(), dtype=self.R_DTYPE)
+            probs = self.analytic_probability(wires=meas_wires)
+            snap_result = self._dot(probs, (eigvals**2)) - self._dot(probs, eigvals) ** 2
+
+        elif isinstance(measurement, VnEntropyMP):
+            base = measurement.log_base
+            map_wires = self.map_wires(meas_wires)
+            snap_result = qml.math.vn_entropy(
+                density_matrix, indices=map_wires, c_dtype=self.C_DTYPE, base=base
+            )
+
+        elif isinstance(measurement, MutualInfoMP):
+            base = measurement.log_base
+            wires0, wires1 = list(map(self.map_wires, measurement.raw_wires))
+            snap_result = qml.math.mutual_info(
+                density_matrix,
+                indices0=wires0,
+                indices1=wires1,
+                c_dtype=self.C_DTYPE,
+                base=base,
+            )
+
+        else:
+            raise DeviceError(
+                f"Snapshots of {type(measurement)} are not yet supported on default.mixed"
+            )
+
+        self._state = pre_rotated_state
+        self._pre_rotated_state = self._state
+
+        return snap_result
+
     def _apply_snapshot(self, operation):
         """Applies the snapshot operation"""
         measurement = operation.hyperparameters["measurement"]
-        if measurement:
-            raise DeviceError("Snapshots of measurements are not yet supported on default.mixed")
+
         if self._debugger and self._debugger.active:
             dim = 2**self.num_wires
             density_matrix = qnp.reshape(self._state, (dim, dim))
+            snap_result = density_matrix
+
+            if measurement:
+                snap_result = self._snapshot_measurements(density_matrix, measurement)
+
             if operation.tag:
-                self._debugger.snapshots[operation.tag] = density_matrix
+                self._debugger.snapshots[operation.tag] = snap_result
             else:
-                self._debugger.snapshots[len(self._debugger.snapshots)] = density_matrix
+                self._debugger.snapshots[len(self._debugger.snapshots)] = snap_result
 
     def _apply_operation(self, operation):
         """Applies operations to the internal device state.


Test patch:
diff --git a/tests/devices/test_default_mixed.py b/tests/devices/test_default_mixed.py
index cafdb89af..51f86edd0 100644
--- a/tests/devices/test_default_mixed.py
+++ b/tests/devices/test_default_mixed.py
@@ -898,15 +898,62 @@ class TestApplyOperation:
         spy_diagonal_unitary.assert_not_called()
         spy_apply_channel.assert_not_called()
 
+    @pytest.mark.parametrize(
+        "measurement",
+        [
+            qml.expval(op=qml.Z(1)),
+            qml.expval(op=qml.Y(0) @ qml.X(1)),
+            qml.var(op=qml.X(0)),
+            qml.var(op=qml.X(0) @ qml.Z(1)),
+            qml.density_matrix(wires=[1]),
+            qml.density_matrix(wires=[0, 1]),
+            qml.probs(op=qml.Y(0)),
+            qml.probs(op=qml.X(0) @ qml.Y(1)),
+            qml.vn_entropy(wires=[0]),
+            qml.mutual_info(wires0=[1], wires1=[0]),
+            qml.purity(wires=[1]),
+        ],
+    )
+    def test_snapshot_supported(self, measurement):
+        """Tests that applying snapshot of measurements is done correctly"""
+
+        def circuit():
+            """Snapshot circuit"""
+            qml.Hadamard(wires=0)
+            qml.Hadamard(wires=1)
+            qml.Snapshot(measurement=qml.expval(qml.Z(0) @ qml.Z(1)))
+            qml.RX(0.123, wires=[0])
+            qml.RY(0.123, wires=[0])
+            qml.CNOT(wires=[0, 1])
+            qml.Snapshot(measurement=measurement)
+            qml.RZ(0.467, wires=[0])
+            qml.RX(0.235, wires=[0])
+            qml.CZ(wires=[1, 0])
+            qml.Snapshot("meas2", measurement=measurement)
+            return qml.probs(op=qml.Y(1) @ qml.Z(0))
+
+        dev_qubit = qml.device("default.qubit", wires=2)
+        dev_mixed = qml.device("default.mixed", wires=2)
+
+        qnode_qubit = qml.QNode(circuit, device=dev_qubit)
+        qnode_mixed = qml.QNode(circuit, device=dev_mixed)
+
+        snaps_qubit = qml.snapshots(qnode_qubit)()
+        snaps_mixed = qml.snapshots(qnode_mixed)()
+
+        for key1, key2 in zip(snaps_qubit, snaps_mixed):
+            assert key1 == key2
+            assert qml.math.allclose(snaps_qubit[key1], snaps_mixed[key2])
+
     def test_snapshot_not_supported(self):
-        """Tests that an error is raised when applying snapshot of measurements"""
+        """Tests that an error is raised when applying snapshot of sample-based measurements"""
 
         dev = qml.device("default.mixed", wires=1)
-        with pytest.raises(DeviceError, match="Snapshots of measurements are not yet supported"):
-            dev._apply_operation(qml.Snapshot(measurement=qml.expval(qml.PauliZ(0))))
-
-        # assert that a snapshot still works without a measurement
-        _ = dev._apply_operation(qml.Snapshot())
+        measurement = qml.sample(op=qml.Z(0))
+        with pytest.raises(
+            DeviceError, match=f"Snapshots of {type(measurement)} are not yet supported"
+        ):
+            dev._snapshot_measurements(dev.state, measurement)
 
 
 class TestApply: