Add error mitigation batch transform

.github/workflows/tests.yml

@@ -59,6 +59,15 @@ jobs:
         if: contains(matrix.config.interfaces, 'jax')
         run: pip3 install jax jaxlib
 
+      # Mitiq requires numpy~=1.20.1 but current versions of TF are only compatible with
+      # numpy<=1.19. Installing Mitiq and tensorflow together causes some tests to fail. We also
+      # install pennylane-qiskit which is needed for Mitiq integration.
+      - name: Conditionally install mitiq if TensorFlow is not installed
+        if: false == ${{contains(matrix.config.device, 'tf')}}
+        run: |
+          pip3 install git+git://github.com/unitaryfund/mitiq@master#egg=mitiq
+          pip3 install pennylane-qiskit
+
       - name: Install PennyLane
         run: |
           pip install -r requirements-ci.txt

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.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.SimplifiedTwoDesign(w1, w2, wires=range(2))
+      return qml.expval(qml.PauliZ(0))
+  ```
+
+  Now, executing `circuit` will be mitigated:
+
+  ```pycon
+  >>> circuit(w1, w2)
+  0.19113067083636542
+  ```
+
 * 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
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -119,6 +120,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,261 @@
+# 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, Tuple
+
+from pennylane import apply
+from pennylane.math import mean
+from pennylane.measure import MeasurementProcess
+from pennylane.operation import Operation
+from pennylane.tape import QuantumTape
+from pennylane.tape.tape import STATE_PREP_OPS
+from pennylane.transforms import batch_transform, single_tape_transform
+
+
+# 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 noise 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.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.SimplifiedTwoDesign(w1, w2, wires=range(2))
+            return qml.expval(qml.PauliZ(0))
+
+    Executions of ``circuit`` will now be mitigated:
+
+    >>> circuit(w1, w2)
+    0.19113067083636542
+
+    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.
+
+    .. 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 amount of noise present when a circuit is run on a
+        noisy device is enumerated by a parameter :math:`\gamma`. Suppose we have an input circuit
+        that experiences an amount of noise equal to :math:`\gamma = \gamma_{0}` when executed.
+        Ideally, we would like to evaluate the result of the circuit in the :math:`\gamma = 0`
+        noise-free setting.
+
+        To do this, we create a family of equivalent circuits whose ideal noise-free value is the
+        same as our input circuit. However, when run on a noisy device, each circuit experiences
+        a noise equal to :math:`\gamma = s \gamma_{0}` for some scale factor :math:`s`. By
+        evaluating the noisy outputs of each circuit, we can extrapolate to :math:`s=0` to estimate
+        the result of running a noise-free circuit.
+
+        A key element of ZNE is the ability to run equivalent circuits for a range of scale factors
+        :math:`s`. When the noise present in a circuit scales with the number of gates, :math:`s`
+        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. When no folding occurs, the scale factor is
+        :math:`s=1` and we are running our input circuit. On the other hand, when each gate has been
+        folded once, we have tripled the amount of noise in the circuit so that :math:`s=3`. For
+        :math:`s \geq 3`, each gate in the circuit will be folded more than once. An typical choice
+        of scale parameters is :math:`(1, 2, 3)`.
+
+        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 {}
+
+    tape_expanded = tape.expand(stop_at=lambda op: not isinstance(op, QuantumTape))
+    tape_removed, in_preps = _remove_preps(tape_expanded)
+    tape_removed = _remove_measurements(tape_removed)
+
+    tapes = [
+        [folding(tape_removed, 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
+    tapes = [_add_measurements(tape_, tape.measurements) for tape_ in tapes]
+    tapes = [_add_preps(tape_, in_preps) for tape_ in tapes]
+
+    def processing_fn(results):
+        """Maps from input tape executions to an error-mitigated estimate"""
+        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
+
+
+def _remove_preps(tape: QuantumTape) -> Tuple[QuantumTape, Tuple[Operation]]:
+    """Removes state preparations from an input tape.
+
+    Args:
+        tape (QuantumTape): the input quantum tape
+
+    Returns:
+        Tuple[QuantumTape, Tuple[Operation]]: the transformed tape with the state preparations
+        removed and an ordered tuple detailing the removed preparations
+    """
+    num_preps = sum(isinstance(o, STATE_PREP_OPS) for o in tape.operations)
+
+    with QuantumTape() as new_tape:
+        for op in tape.operations[num_preps:]:
+            apply(op)
+        for op in tape.measurements:
+            apply(op)
+
+    return new_tape, tape.operations[:num_preps]
+
+
+@single_tape_transform
+def _add_preps(tape: QuantumTape, preps: Tuple[Operation]) -> QuantumTape:
+    """Add state preparations to an input tape.
+
+    Args:
+        tape (QuantumTape): the input quantum tape
+        preps (tuple[Operation]): the state preparations to be added
+
+    Returns:
+        QuantumTape: the transformed tape with the state preparations added
+    """
+    for prep in preps:
+        apply(prep)
+    for op in tape.operations:
+        apply(op)
+    for m in tape.measurements:
+        apply(m)
+
+
+@single_tape_transform
+def _remove_measurements(tape: QuantumTape) -> QuantumTape:
+    """Removes measurements from an input tape.
+
+    Args:
+        tape (QuantumTape): the input quantum tape
+
+    Returns:
+        QuantumTape: the transformed tape with the measurements removed
+    """
+    for op in tape.operations:
+        apply(op)
+
+
+@single_tape_transform
+def _add_measurements(tape: QuantumTape, measurements: Tuple[MeasurementProcess]) -> QuantumTape:
+    """Add measurements to an input tape.
+
+    Args:
+        tape (QuantumTape): the input quantum tape
+        measurements (tuple[MeasurementProcess]): the measurements to be added
+
+    Returns:
+        QuantumTape: the transformed tape with the measurements added
+    """
+    for op in tape.operations:
+        apply(op)
+    for m in measurements:
+        apply(m)