Add error mitigation batch transform

doc/releases/changelog-dev.md

@@ -4,6 +4,54 @@
 
 <h3>New features since last release</h3>
 
+* Error mitigation using the zero-noise extrapolation method is now available through the
+  `transforms.mitigate_with_zne` transform. This transform can integrate with the
+  [Mitiq](https://mitiq.readthedocs.io/en/stable/) package for unitary folding and extrapolation
+  functionality.
+  [(#1813)](https://github.com/PennyLaneAI/pennylane/pull/1813)
+
+  Consider the following noisy device:
+
+  ```python
+  import pennylane as qml
+
+  noise_strength = 0.05
+
+  dev = qml.device("default.mixed", wires=2)
+  dev = qml.transforms.insert(qml.AmplitudeDamping, noise_strength)(dev)
+  ```
+
+  We can mitigate the effects of this noise for circuits run on this device by using the added
+  transform:
+
+  ```python
+  from pennylane import numpy as np
+  from pennylane.beta import qnode
+
+  from mitiq.zne.scaling import fold_global
+  from mitiq.zne.inference import RichardsonFactory
+
+  n_wires = 2
+  n_layers = 2
+
+  shapes = qml.templates.SimplifiedTwoDesign.shape(n_wires, n_layers)
+  np.random.seed(0)
+  w1, w2 = [np.random.random(s) for s in shapes]
+
+  @qml.transforms.mitigate_with_zne([1, 2, 3], fold_global, RichardsonFactory.extrapolate)
+  @qnode(dev)
+  def circuit(w1, w2):
+      qml.templates.SimplifiedTwoDesign(w1, w2, wires=range(2))
+      return qml.expval(qml.PauliZ(0))
+  ```
+
+  Now, executing `circuit` will be mitigated:
+
+  ```pycon
+  >>> circuit(w1, w2)
+  array([0.19113067])
+  ```
+
 * The `insert` transform has now been added, providing a way to insert single-qubit operations into
   a quantum circuit. The transform can apply to quantum functions, tapes, and devices.
   [(#1795)](https://github.com/PennyLaneAI/pennylane/pull/1795)

pennylane/transforms/__init__.py

@@ -35,6 +35,7 @@
     ~transforms.get_unitary_matrix
     ~metric_tensor
     ~specs
+    ~transforms.mitigate_with_zne
 
 Transforms that act on quantum functions
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -104,10 +105,16 @@
     ~transforms.expand_invalid_trainable
     ~transforms.expand_multipar
     ~transforms.expand_nonunitary_gen
+    ~transforms.support_preparations_and_measurements
 """
 # Import the decorators first to prevent circular imports when used in other transforms
 from .batch_transform import batch_transform
-from .qfunc_transforms import make_tape, single_tape_transform, qfunc_transform
+from .qfunc_transforms import (
+    make_tape,
+    single_tape_transform,
+    qfunc_transform,
+    support_preparations_and_measurements,
+)
 from .adjoint import adjoint
 from .batch_params import batch_params
 from .classical_jacobian import classical_jacobian
@@ -119,6 +126,7 @@
 from .measurement_grouping import measurement_grouping
 from .metric_tensor import metric_tensor
 from .insert_ops import insert
+from .mitigate import mitigate_with_zne
 from .optimization import (
     cancel_inverses,
     commute_controlled,

pennylane/transforms/insert_ops.py

@@ -184,6 +184,8 @@ def f(w, x, y, z):
         >>> qnode_noisy(0.9, 0.4, 0.5, 0.6)
         tensor(0.72945434, requires_grad=True)
     """
+    circuit = circuit.expand(stop_at=lambda op: not isinstance(op, QuantumTape))
+
     if not isinstance(op, FunctionType) and op.num_wires != 1:
         raise ValueError("Only single-qubit operations can be inserted into the circuit")
 

pennylane/transforms/mitigate.py

@@ -0,0 +1,168 @@
+# Copyright 2021 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Provides transforms for mitigating quantum circuits."""
+from typing import Any, Dict, Optional, Sequence
+
+from pennylane.math import mean
+from pennylane.tape import QuantumTape
+from pennylane.transforms import batch_transform, support_preparations_and_measurements
+
+
+# pylint: disable=too-many-arguments
+@batch_transform
+def mitigate_with_zne(
+    tape: QuantumTape,
+    scale_factors: Sequence[float],
+    folding: callable,
+    extrapolate: callable,
+    folding_kwargs: Optional[Dict[str, Any]] = None,
+    extrapolate_kwargs: Optional[Dict[str, Any]] = None,
+    reps_per_factor=1,
+) -> float:
+    r"""Mitigate an input circuit using zero-noise extrapolation.
+
+    Error mitigation is a precursor to error correction and is compatible with near-term quantum
+    devices. It aims to lower the impact of noise when evaluating a circuit on a quantum device by
+    evaluating multiple variations of the circuit and post-processing the results into a
+    noise-reduced estimate. This transform implements the zero-noise extrapolation (ZNE) method
+    originally introduced by
+    `Temme et al. <https://journals.aps.org/prl/abstract/10.1103/PhysRevLett.119.180509>`__ and
+    `Li et al. <https://journals.aps.org/prx/abstract/10.1103/PhysRevX.7.021050>`__.
+
+    Details on the functions passed to the ``folding`` and ``extrapolate`` arguments of this
+    transform can be found in the usage details. This transform is compatible with functionality
+    from the `Mitiq <https://mitiq.readthedocs.io/en/stable/>`__ package (version 0.11.0 and above),
+    see the example and usage details for further information.
+
+    Args:
+        tape (QuantumTape): the circuit to be error mitigated
+        scale_factors (Sequence[float]): the range of scale factors used
+        folding (callable): a function that returns a folded circuit for a specified scale factor
+        extrapolate (callable): a function that returns an extrapolated result when provided a
+            range of scale factors and corresponding results
+        folding_kwargs (dict): optional keyword arguments passed to the ``folding`` function
+        extrapolate_kwargs (dict): optional keyword arguments passed to the ``extrapolate`` function
+        reps_per_factor (int): Number of circuits generated for each scale factor. Useful when the
+            folding function is stochastic.
+
+    Returns:
+        float: the result of evaluating the circuit when mitigated using 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:
+
+    .. code-block:: python3
+
+        import pennylane as qml
+
+        noise_strength = 0.05
+
+        dev = qml.device("default.mixed", wires=2)
+        dev = qml.transforms.insert(qml.AmplitudeDamping, noise_strength)(dev)
+
+    We can now set up a mitigated QNode by harnessing functionality from the
+    `Mitiq <https://mitiq.readthedocs.io/en/stable/>`__ package:
+
+    .. code-block:: python3
+
+        from pennylane import numpy as np
+        from pennylane.beta import qnode
+
+        from mitiq.zne.scaling import fold_global
+        from mitiq.zne.inference import RichardsonFactory
+
+        n_wires = 2
+        n_layers = 2
+
+        shapes = qml.templates.SimplifiedTwoDesign.shape(n_wires, n_layers)
+        np.random.seed(0)
+        w1, w2 = [np.random.random(s) for s in shapes]
+
+        @qml.transforms.mitigate_with_zne([1, 2, 3], fold_global, RichardsonFactory.extrapolate)
+        @qnode(dev)
+        def circuit(w1, w2):
+            qml.templates.SimplifiedTwoDesign(w1, w2, wires=range(2))
+            return qml.expval(qml.PauliZ(0))
+
+    Executions of ``circuit`` will now be mitigated:
+
+    >>> circuit(w1, w2)
+    array([0.19113067])
+
+    .. UsageDetails::
+
+        A summary of ZNE can be found in `LaRose et al. <https://arxiv.org/abs/2009.04417>`__. The
+        method works by assuming that the circuit experiences a fixed amount of noise when executed
+        on a noisy device that is enumerated by the parameter :math:`\gamma`. If an equivalent
+        circuit can be run for a range of noise parameters :math:`\gamma`, then the results can be
+        extrapolated to the :math:`\gamma = 0` noiseless case.
+
+        A key element of ZNE is the ability to run equivalent circuits for a range of noise
+        parameters :math:`\gamma`. When :math:`\gamma` scales with the number of gates in the
+        circuit, it can be varied using
+        `unitary folding <https://ieeexplore.ieee.org/document/9259940>`__. Unitary folding
+        works by noticing that a unitary :math:`U` is equivalent to :math:`U U^{\dagger} U`. This
+        type of transform can be applied to individual gates in the circuit or to the whole circuit
+        and is controlled by a scale parameter :math:`s` which is calibrated so that :math:`s = 1`
+        corresponds to the (unfolded) input circuit and :math:`s = 3` is a folding of all gates in
+        the circuit once.
+
+        This transform applies ZNE to an input circuit using the unitary folding approach. It
+        requires a callable to be passed as the ``folding`` argument with signature
+        ``fn(circuit, scale_factor, **folding_kwargs)`` where ``circuit`` is a quantum tape,
+        ``scale_factor`` is a float, and ``folding_kwargs`` are optional arguments passed to the
+        folding function. The output of the function should be the folded circuit as a quantum tape.
+        Folding functionality is available from the
+        `Mitiq <https://mitiq.readthedocs.io/en/stable/>`__ package (version 0.11.0 and above)
+        in the
+        `zne.scaling.folding <https://mitiq.readthedocs.io/en/stable/apidoc.html#module-mitiq.zne.scaling.folding>`__
+        module.
+
+        This transform also requires a callable to be passed to the ``extrapolate`` argument that
+        returns the extrapolated value(s). Its function should be
+        ``fn(scale_factors, results, **extrapolate_kwargs)`` where ``scale_factors`` are the ZNE
+        scale factors, ``results`` are the execution results of the circuit at the specified scale
+        factors, and ``extrapolate_kwargs`` are optional keyword arguments. The shape of ``results``
+        will be ``(len(scale_factors), num_returns)``, where ``num_returns`` is the number of
+        returns in the mitigated QNode. The output of ``extrapolate`` should be a flat array of
+        length ``num_returns``. Extrapolation functionality is available using ``extrapolate``
+        methods of the factories in the
+        `mitiq.zne.inference <https://mitiq.readthedocs.io/en/stable/apidoc.html#module-mitiq.zne.inference>`__
+        module (version 0.11.0 and above).
+
+        .. warning::
+
+            Calculating the gradient of mitigated circuits is not supported when using the Mitiq
+            package as a backend for folding or extrapolation.
+    """
+    folding_kwargs = folding_kwargs or {}
+    extrapolate_kwargs = extrapolate_kwargs or {}
+    folding = support_preparations_and_measurements(folding)
+
+    tapes = [
+        [folding(tape, s, **folding_kwargs) for _ in range(reps_per_factor)] for s in scale_factors
+    ]
+    tapes = [tape_ for tapes_ in tapes for tape_ in tapes_]  # flattens nested list
+
+    def processing_fn(results):
+        results = [
+            results[i : i + reps_per_factor] for i in range(0, len(results), reps_per_factor)
+        ]  # creates nested list according to reps_per_factor
+        results = mean(results, axis=1)
+        extrapolated = extrapolate(scale_factors, results, **extrapolate_kwargs)
+        return extrapolated[0] if len(extrapolated) == 1 else extrapolated
+
+    return tapes, processing_fn