Two stage ZNE (#2452)

* started with the two-stage ZNE approach together with Nate

* added docstrings

* edit docstring for scaled_circuits() function in zne.py

* Update mitiq/zne/tests/test_zne.py

remove comment in test_zne.py

Co-authored-by: nate stemen <[email protected]>

* Update mitiq/zne/zne.py

use Sequence instead of list in typehinting for scale_factors

Co-authored-by: nate stemen <[email protected]>

* import Sequence in zne.py

* add tests for different frontends

* format fix

* remove redundant tests in test_two_stage_zne()

* use Sequence instead of list in combine_results() in zne.py

---------

Co-authored-by: nate stemen <[email protected]>

mitiq/zne/tests/test_zne.py

@@ -18,13 +18,18 @@
 from mitiq import QPROGRAM, SUPPORTED_PROGRAM_TYPES
 from mitiq.benchmarks.randomized_benchmarking import generate_rb_circuits
 from mitiq.interface import accept_any_qprogram_as_input, convert_from_mitiq
+from mitiq.interface.mitiq_braket import to_braket
 from mitiq.interface.mitiq_cirq import (
     compute_density_matrix,
     sample_bitstrings,
 )
+from mitiq.interface.mitiq_pennylane import to_pennylane
+from mitiq.interface.mitiq_pyquil import to_pyquil
+from mitiq.interface.mitiq_qibo import to_qibo
 from mitiq.interface.mitiq_qiskit import (
     execute_with_shots_and_noise,
     initialized_depolarizing_noise,
+    to_qiskit,
 )
 from mitiq.observable import Observable, PauliString
 from mitiq.zne import (
@@ -41,10 +46,13 @@
     RichardsonFactory,
 )
 from mitiq.zne.scaling import (
+    fold_all,
     fold_gates_at_random,
+    fold_global,
     get_layer_folding,
     insert_id_layers,
 )
+from mitiq.zne.zne import combine_results, scaled_circuits
 
 BASE_NOISE = 0.007
 TEST_DEPTH = 30
@@ -560,3 +568,59 @@ def qs_noisy_simulation(
 
     mitigated = execute_with_zne(circuit, qs_noisy_simulation)
     assert abs(mitigated) < 1000
+
+
+@pytest.mark.parametrize(
+    "noise_scaling_method",
+    [fold_gates_at_random, insert_id_layers, fold_global, fold_all],
+)
+@pytest.mark.parametrize(
+    "extrapolation_factory", [RichardsonFactory, LinearFactory]
+)
+@pytest.mark.parametrize(
+    "to_frontend",
+    [None, to_qiskit, to_braket, to_pennylane, to_pyquil, to_qibo],
+)
+def test_two_stage_zne(
+    noise_scaling_method, extrapolation_factory, to_frontend
+):
+    qreg = cirq.LineQubit.range(2)
+    cirq_circuit = cirq.Circuit(
+        cirq.H.on_each(qreg),
+        cirq.CNOT(*qreg),
+        cirq.CNOT(*qreg),
+        cirq.H.on_each(qreg),
+    )
+    if to_frontend is not None:
+        frontend_circuit = to_frontend(cirq_circuit)
+    else:
+        frontend_circuit = cirq_circuit
+
+    scale_factors = [1, 3, 5]
+    circs = scaled_circuits(
+        frontend_circuit, scale_factors, noise_scaling_method
+    )
+
+    assert len(circs) == len(scale_factors)
+
+    np.random.seed(42)
+
+    def executor(circuit):
+        return np.random.random()
+
+    results = [executor(cirq_circuit) for _ in range(3)]
+    extrapolation_method = extrapolation_factory.extrapolate
+    two_stage_zne_res = combine_results(
+        scale_factors, results, extrapolation_method
+    )
+
+    assert isinstance(two_stage_zne_res, float)
+
+    np.random.seed(42)
+    zne_res = execute_with_zne(
+        cirq_circuit,
+        executor,
+        factory=extrapolation_factory(scale_factors),
+        scale_noise=noise_scaling_method,
+    )
+    assert np.isclose(zne_res, two_stage_zne_res)

mitiq/zne/zne.py

@@ -6,13 +6,62 @@
 """High-level zero-noise extrapolation tools."""
 
 from functools import wraps
-from typing import Callable, List, Optional, Union
+from typing import Callable, List, Optional, Sequence, Union
 
 from mitiq import QPROGRAM, Executor, Observable, QuantumResult
 from mitiq.zne.inference import Factory, RichardsonFactory
 from mitiq.zne.scaling import fold_gates_at_random
 
 
+def scaled_circuits(
+    circuit: QPROGRAM,
+    scale_factors: list[float],
+    scale_method: Callable[[QPROGRAM, float], QPROGRAM],
+) -> list[QPROGRAM]:
+    """Given a circuit, scale_factors and a scale_method, outputs a list
+       of circuits that will be used in ZNE.
+
+    Args:
+        circuit: The input circuit to execute with ZNE.
+        scale_factors: A list of ``float``s to scale the circuit with.
+        scale_method: The function for scaling the noise of a quantum circuit.
+            A list of built-in functions can be found in ``mitiq.zne.scaling``.
+
+    Returns:
+        The scaled circuits using the scale_method.
+    """
+    circuits = []
+    for scale_factor in scale_factors:
+        circuits.append(scale_method(circuit, scale_factor))
+
+    return circuits
+
+
+def combine_results(
+    scale_factors: Sequence[float],
+    results: Sequence[float],
+    extrapolation_method: Callable[[Sequence[float], Sequence[float]], float],
+) -> float:
+    """Computes the error-mitigated expectation value associated to the
+    input results from executing the scaled circuits, via the application
+    of zero-noise extrapolation (ZNE).
+
+    Args:
+        scale_factors: A list of ``float``s to scale the circuit with.
+        results: A list of ``float``s that is the result of applying an
+            executor to the scaled circuits.
+        extrapolation_method: The function for scaling the noise of a
+            quantum circuit. A list of built-in functions can be found
+            in ``mitiq.zne.scaling``.
+
+    Returns:
+        The expectation value estimated with ZNE.
+    """
+    res = extrapolation_method(scale_factors, results)
+
+    return res
+
+
 def execute_with_zne(
     circuit: QPROGRAM,
     executor: Union[Executor, Callable[[QPROGRAM], QuantumResult]],