NoiseModel example for mitigate_with_zne (#6346)

doc/releases/changelog-dev.md

@@ -90,6 +90,9 @@
   a sparse matrix.
   [(#6173)](https://github.com/PennyLaneAI/pennylane/pull/6173)
 
+* `mitigate_with_zne` now gives clearer error message when being used with circuits with channel noise.
+  [(#6346)](https://github.com/PennyLaneAI/pennylane/pull/6346)
+
 * The `make_plxpr` function is added, to take a function and create a `Callable` that,
   when called, will return a PLxPR representation of the input function.
   [(#6326)](https://github.com/PennyLaneAI/pennylane/pull/6326)

pennylane/transforms/mitigate.py

@@ -215,6 +215,12 @@ def qfunc(op):
     # Generate base_circuit without measurements
     # Treat all circuits as lists of operations, build new tape in the end
     base_ops = circuit.operations
+    if any((isinstance(op, qml.operation.Channel) for op in base_ops)):
+        raise ValueError(
+            "Circuits containing quantum channels cannot be folded with mitigate_with_zne. "
+            "To use zero-noise extrapolation on the circuit with channel noise, "
+            "please add the noise on the device rather than the circuit."
+        )
 
     num_global_folds, fraction_scale = _divmod(scale_factor - 1, 2)
 
@@ -414,16 +420,26 @@ def mitigate_with_zne(
     **Example:**
 
     We first create a noisy device using ``default.mixed`` by adding :class:`~.AmplitudeDamping` to
-    each gate of circuits executed on the device using the :func:`~.transforms.insert` transform:
+    each gate of circuits executed on the device using the :func:`~.transforms.add_noise` transform:
 
     .. code-block:: python3
 
         import pennylane as qml
 
-        noise_strength = 0.05
-
         dev = qml.device("default.mixed", wires=2)
-        dev = qml.transforms.insert(dev, qml.AmplitudeDamping, noise_strength)
+
+        fcond = qml.noise.wires_in(dev.wires)
+        noise = qml.noise.partial_wires(qml.AmplitudeDamping, 0.05)
+        noise_model = qml.NoiseModel({fcond: noise})
+
+        noisy_dev = qml.add_noise(dev, noise_model)
+
+    .. note ::
+
+        The :func:`~.transforms.add_noise` transform should be used on the device instead of
+        the circuit if the defined ``noise_model`` contains a :class:`~.operation.Channel`
+        instance. This is to prevent ``mitigate_with_zne`` from computing the adjoint of
+        the channel operation during `folding`, which is currently not supported.
 
     We can now set up a mitigated ``QNode`` by first decomposing it into a target gate set via :func:`~.pennylane.transforms.decompose`
     and then applying this transform by passing ``folding`` and ``extrapolate`` functions. PennyLane provides native
@@ -450,17 +466,18 @@ def mitigate_with_zne(
             scale_factors=[1., 2., 3.],
             folding=fold_global,
             extrapolate=poly_extrapolate,
-            extrapolate_kwargs={'order': 2})
+            extrapolate_kwargs={'order' : 2},
+        )
         @partial(qml.transforms.decompose, gate_set = ["RY", "CZ"])
-        @qnode(dev)
+        @qnode(noisy_dev)
         def circuit(w1, w2):
             qml.SimplifiedTwoDesign(w1, w2, wires=range(2))
             return qml.expval(qml.Z(0))
 
     Executions of ``circuit`` will now be mitigated:
 
     >>> circuit(w1, w2)
-    0.19113067083636542
+    0.19113067088978522
 
     The unmitigated circuit result is ``0.33652776`` while the ideal circuit result is
     ``0.23688169`` and we can hence see that mitigation has helped reduce our estimation error.

tests/transforms/test_mitigate.py

@@ -71,6 +71,7 @@ def same_tape(tape1, tape2):
 class TestMitigateWithZNE:
     """Tests for the mitigate_with_zne function"""
 
+    # pylint:disable = unnecessary-lambda-assignment
     folding = lambda *args, **kwargs: tape_base
     extrapolate = lambda *args, **kwargs: [3.141]
 
@@ -219,6 +220,52 @@ def original_qnode(inputs):
         result_expanded = mitigated_qnode_expanded(inputs)
         assert qml.math.allclose(result_orig, result_expanded)
 
+    # pylint:disable=not-callable
+    def test_zne_with_noise_models(self):
+        """Test that mitigate_with_zne transform works with noise models"""
+        fcond = qml.noise.wires_in([0, 1])
+        noise = qml.noise.partial_wires(qml.AmplitudeDamping, 0.05)
+        noise_model = qml.NoiseModel({fcond: noise})
+
+        def circuit():
+            qml.RX(1.23, wires=0)
+            qml.RZ(0.45, wires=1)
+            return qml.expval(qml.Z(0) @ qml.Z(1))
+
+        noise_qnode = qml.QNode(circuit, device=qml.add_noise(dev_ideal, noise_model))
+        zne_qnode = qml.transforms.mitigate_with_zne(
+            noise_qnode, [1, 2, 3], fold_global, richardson_extrapolate
+        )
+
+        # following result has been obtained manually and also by using
+        # qml.transforms.mitigate_with_zne(
+        #     noise_qnode, [1, 2, 3],
+        #     mitiq.zne.scaling.fold_global, mitiq.zne.inference.RichardsonFactory.extrapolate
+        # )()
+        mitigated_result = 0.39843788456
+        assert qml.math.allclose(zne_qnode(), mitigated_result, atol=1e-2)
+
+    # pylint:disable=not-callable
+    def test_zne_error_with_channels(self):
+        """Test that mitigate_with_zne transform raises correct error with channels"""
+        fcond = qml.noise.wires_in([0, 1])
+        noise = qml.noise.partial_wires(qml.AmplitudeDamping, 0.05)
+        noise_model = qml.NoiseModel({fcond: noise})
+
+        def circuit():
+            qml.RX(1.23, wires=0)
+            qml.RZ(0.45, wires=1)
+            return qml.expval(qml.Z(0) @ qml.Z(1))
+
+        with pytest.raises(
+            ValueError,
+            match="Circuits containing quantum channels cannot be folded with mitigate_with_zne.",
+        ):
+            noisy_qnode = qml.add_noise(qml.QNode(circuit, device=dev_ideal), noise_model)
+            qml.transforms.mitigate_with_zne(
+                noisy_qnode, [1, 2, 3], fold_global, richardson_extrapolate
+            )()
+
 
 @pytest.fixture
 def skip_if_no_mitiq_support():