LRE Executors (#2499)

* main functions without corrected mypy errors, unit tests and docstrings

* add unit tests for the decorator - check for test coverage

* fix #2500 (comment)

* mypy - remove shots

* more unit tests + docstrings

* dosctring args formatting

* fix #2499 (comment)

* weird chunking failure

* try chunking to 2

* num_chunks = 5 with test_cirq * 200

* push to check for test coverage

* chunking failures

* split decorator unit test

* cleanup

* chunking failure again + docker failure

* nate's suggestions

.gitignore

@@ -15,7 +15,7 @@ mitiq.egg-info/
 dist/
 build/
 jupyter_execute/
-
+.mypy_cache/
 # Coverage reports
 coverage.xml
 .coverage

docs/source/apidoc.md

@@ -89,6 +89,11 @@ See Ref. {cite}`Czarnik_2021_Quantum` for more details on these methods.
 
 ### Layerwise Richardson Extrapolation
 
+```{eval-rst}
+.. automodule:: mitiq.lre.lre
+   :members:
+```
+
 ```{eval-rst}
 .. automodule:: mitiq.lre.multivariate_scaling.layerwise_folding
    :members:

mitiq/lre/__init__.py

@@ -10,4 +10,6 @@
 from mitiq.lre.inference.multivariate_richardson import (
     multivariate_richardson_coefficients,
     sample_matrix,
-)
+)
+
+from mitiq.lre.lre import execute_with_lre, mitigate_executor, lre_decorator

mitiq/lre/lre.py

@@ -0,0 +1,170 @@
+# Copyright (C) Unitary Fund
+#
+# This source code is licensed under the GPL license (v3) found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Extrapolation methods for Layerwise Richardson Extrapolation (LRE)"""
+
+from functools import wraps
+from typing import Any, Callable, Optional, Union
+
+import numpy as np
+from cirq import Circuit
+
+from mitiq import QPROGRAM
+from mitiq.lre import (
+    multivariate_layer_scaling,
+    multivariate_richardson_coefficients,
+)
+from mitiq.zne.scaling import fold_gates_at_random
+
+
+def execute_with_lre(
+    input_circuit: Circuit,
+    executor: Callable[[Circuit], float],
+    degree: int,
+    fold_multiplier: int,
+    folding_method: Callable[
+        [QPROGRAM, float], QPROGRAM
+    ] = fold_gates_at_random,  # type: ignore [has-type]
+    num_chunks: Optional[int] = None,
+) -> float:
+    r"""
+    Defines the executor required for Layerwise Richardson
+    Extrapolation as defined in :cite:`Russo_2024_LRE`.
+
+    Note that this method only works for the multivariate extrapolation
+    methods. It does not allows a user to choose which layers in the input
+    circuit will be scaled.
+
+    .. seealso::
+
+        If you would prefer to choose the layers for unitary
+        folding, use :func:`mitiq.zne.scaling.layer_scaling.get_layer_folding`
+        instead.
+
+    Args:
+        input_circuit: Circuit to be scaled.
+        executor: Executes a circuit and returns a `float`
+        degree: Degree of the multivariate polynomial.
+        fold_multiplier: Scaling gap value required for unitary folding which
+            is used to generate the scale factor vectors.
+        folding_method: Unitary folding method. Default is
+            :func:`fold_gates_at_random`.
+        num_chunks: Number of desired approximately equal chunks. When the
+            number of chunks is the same as the layers in the input circuit,
+            the input circuit is unchanged.
+
+
+    Returns:
+        Error-mitigated expectation value
+
+    """
+    noise_scaled_circuits = multivariate_layer_scaling(
+        input_circuit, degree, fold_multiplier, num_chunks, folding_method
+    )
+
+    linear_combination_coeffs = multivariate_richardson_coefficients(
+        input_circuit, degree, fold_multiplier, num_chunks
+    )
+
+    # verify the linear combination coefficients and the calculated expectation
+    # values have the same length
+    if len(noise_scaled_circuits) != len(  # pragma: no cover
+        linear_combination_coeffs
+    ):
+        raise AssertionError(
+            "The number of expectation values are not equal "
+            + "to the number of coefficients required for "
+            + "multivariate extrapolation."
+        )
+
+    lre_exp_values = []
+    for scaled_circuit in noise_scaled_circuits:
+        circ_exp_val = executor(scaled_circuit)
+        lre_exp_values.append(circ_exp_val)
+
+    return np.dot(lre_exp_values, linear_combination_coeffs)
+
+
+def mitigate_executor(
+    executor: Callable[[Circuit], float],
+    degree: int,
+    fold_multiplier: int,
+    folding_method: Callable[
+        [Union[Any], float], Union[Any]
+    ] = fold_gates_at_random,
+    num_chunks: Optional[int] = None,
+) -> Callable[[Circuit], float]:
+    """Returns a modified version of the input `executor` which is
+    error-mitigated with layerwise richardson extrapolation (LRE).
+
+    Args:
+        input_circuit: Circuit to be scaled.
+        executor: Executes a circuit and returns a `float`
+        degree: Degree of the multivariate polynomial.
+        fold_multiplier Scaling gap value required for unitary folding which
+            is used to generate the scale factor vectors.
+        folding_method: Unitary folding method. Default is
+            :func:`fold_gates_at_random`.
+        num_chunks: Number of desired approximately equal chunks. When the
+            number of chunks is the same as the layers in the input circuit,
+            the input circuit is unchanged.
+
+
+    Returns:
+        Error-mitigated version of the circuit executor.
+    """
+
+    @wraps(executor)
+    def new_executor(input_circuit: Circuit) -> float:
+        return execute_with_lre(
+            input_circuit,
+            executor,
+            degree,
+            fold_multiplier,
+            folding_method,
+            num_chunks,
+        )
+
+    return new_executor
+
+
+def lre_decorator(
+    degree: int,
+    fold_multiplier: int,
+    folding_method: Callable[[Circuit, float], Circuit] = fold_gates_at_random,
+    num_chunks: Optional[int] = None,
+) -> Callable[[Callable[[Circuit], float]], Callable[[Circuit], float]]:
+    """Decorator which adds an error-mitigation layer based on
+    layerwise richardson extrapolation (LRE).
+
+    Args:
+        input_circuit: Circuit to be scaled.
+        executor: Executes a circuit and returns a `float`
+        degree: Degree of the multivariate polynomial.
+        fold_multiplier Scaling gap value required for unitary folding which
+            is used to generate the scale factor vectors.
+        folding_method: Unitary folding method. Default is
+            :func:`fold_gates_at_random`.
+        num_chunks: Number of desired approximately equal chunks. When the
+            number of chunks is the same as the layers in the input circuit,
+            the input circuit is unchanged.
+
+
+    Returns:
+        Error-mitigated decorator.
+    """
+
+    def decorator(
+        executor: Callable[[Circuit], float],
+    ) -> Callable[[Circuit], float]:
+        return mitigate_executor(
+            executor,
+            degree,
+            fold_multiplier,
+            folding_method,
+            num_chunks,
+        )
+
+    return decorator

mitiq/lre/tests/test_lre.py

@@ -0,0 +1,129 @@
+"""Unit tests for the LRE extrapolation methods."""
+
+import re
+
+import pytest
+from cirq import DensityMatrixSimulator, depolarize
+
+from mitiq import benchmarks
+from mitiq.lre import execute_with_lre, lre_decorator, mitigate_executor
+from mitiq.zne.scaling import fold_all, fold_global
+
+# default circuit for all unit tests
+test_cirq = benchmarks.generate_rb_circuits(
+    n_qubits=1,
+    num_cliffords=2,
+)[0]
+
+
+# default execute function for all unit tests
+def execute(circuit, noise_level=0.025):
+    """Default executor for all unit tests."""
+    noisy_circuit = circuit.with_noise(depolarize(p=noise_level))
+    rho = DensityMatrixSimulator().simulate(noisy_circuit).final_density_matrix
+    return rho[0, 0].real
+
+
+noisy_val = execute(test_cirq)
+ideal_val = execute(test_cirq, noise_level=0)
+
+
+@pytest.mark.parametrize("degree, fold_multiplier", [(2, 2), (2, 3), (3, 4)])
+def test_lre_exp_value(degree, fold_multiplier):
+    """Verify LRE executors work as expected."""
+    lre_exp_val = execute_with_lre(
+        test_cirq,
+        execute,
+        degree=degree,
+        fold_multiplier=fold_multiplier,
+    )
+    assert abs(lre_exp_val - ideal_val) <= abs(noisy_val - ideal_val)
+
+
+@pytest.mark.parametrize("degree, fold_multiplier", [(2, 2), (2, 3), (3, 4)])
+def test_lre_exp_value_decorator(degree, fold_multiplier):
+    """Verify LRE mitigated executor work as expected."""
+    mitigated_executor = mitigate_executor(
+        execute, degree=2, fold_multiplier=2
+    )
+    exp_val_from_mitigate_executor = mitigated_executor(test_cirq)
+    assert abs(exp_val_from_mitigate_executor - ideal_val) <= abs(
+        noisy_val - ideal_val
+    )
+
+
+def test_lre_decorator():
+    """Verify LRE decorators work as expected."""
+
+    @lre_decorator(degree=2, fold_multiplier=2)
+    def execute(circuit, noise_level=0.025):
+        noisy_circuit = circuit.with_noise(depolarize(p=noise_level))
+        rho = (
+            DensityMatrixSimulator()
+            .simulate(noisy_circuit)
+            .final_density_matrix
+        )
+        return rho[0, 0].real
+
+    assert abs(execute(test_cirq) - ideal_val) <= abs(noisy_val - ideal_val)
+
+
+def test_lre_decorator_raised_error():
+    """Verify an error is raised when the required parameters for the decorator
+    are not specified."""
+    with pytest.raises(TypeError, match=re.escape("lre_decorator() missing")):
+
+        @lre_decorator()
+        def execute(circuit, noise_level=0.025):
+            noisy_circuit = circuit.with_noise(depolarize(p=noise_level))
+            rho = (
+                DensityMatrixSimulator()
+                .simulate(noisy_circuit)
+                .final_density_matrix
+            )
+            return rho[0, 0].real
+
+        assert abs(execute(test_cirq) - ideal_val) <= abs(
+            noisy_val - ideal_val
+        )
+
+
+def test_lre_executor_with_chunking():
+    """Verify the executor works as expected for chunking a large circuit into
+    a smaller circuit."""
+    # define a larger circuit
+    test_cirq = benchmarks.generate_rb_circuits(n_qubits=1, num_cliffords=12)[
+        0
+    ]
+    lre_exp_val = execute_with_lre(
+        test_cirq, execute, degree=2, fold_multiplier=2, num_chunks=14
+    )
+    assert abs(lre_exp_val - ideal_val) <= abs(noisy_val - ideal_val)
+
+
+@pytest.mark.parametrize("num_chunks", [(1), (2), (3), (4), (5), (6), (7)])
+def test_large_circuit_with_small_chunks_poor_performance(num_chunks):
+    """Verify chunking performs poorly when a large number of layers are
+    chunked into a smaller number of circuit chunks."""
+    # define a larger circuit
+    test_cirq = benchmarks.generate_rb_circuits(n_qubits=1, num_cliffords=15)[
+        0
+    ]
+    lre_exp_val = execute_with_lre(
+        test_cirq, execute, degree=2, fold_multiplier=2, num_chunks=num_chunks
+    )
+    assert abs(lre_exp_val - ideal_val) >= abs(noisy_val - ideal_val)
+
+
+@pytest.mark.parametrize("input_method", [(fold_global), (fold_all)])
+def test_lre_executor_with_different_folding_methods(input_method):
+    """Verify the executor works as expected for using non-default unitary
+    folding methods."""
+    lre_exp_val = execute_with_lre(
+        test_cirq,
+        execute,
+        degree=2,
+        fold_multiplier=2,
+        folding_method=input_method,
+    )
+    assert abs(lre_exp_val - ideal_val) <= abs(noisy_val - ideal_val)