nate's suggestions

mitiq/lre/lre.py

@@ -11,7 +11,7 @@
 import numpy as np
 from cirq import Circuit
 
-from mitiq import QPROGRAM, QuantumResult
+from mitiq import QPROGRAM
 from mitiq.lre import (
     multivariate_layer_scaling,
     multivariate_richardson_coefficients,
@@ -21,7 +21,7 @@
 
 def execute_with_lre(
     input_circuit: Circuit,
-    executor: Callable[[Circuit], QuantumResult],
+    executor: Callable[[Circuit], float],
     degree: int,
     fold_multiplier: int,
     folding_method: Callable[
@@ -45,7 +45,7 @@ def execute_with_lre(
 
     Args:
         input_circuit: Circuit to be scaled.
-        executor: Executes a circuit and returns a `QuantumResult`
+        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.
@@ -63,18 +63,14 @@ def execute_with_lre(
     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
     )
 
-    lre_exp_values = []
-    for scaled_circuit in noise_scaled_circuits:
-        circ_exp_val = executor(scaled_circuit)
-        lre_exp_values.append(circ_exp_val)
-
     # verify the linear combination coefficients and the calculated expectation
     # values have the same length
-    if not len(lre_exp_values) == len(  # pragma: no cover
+    if len(noise_scaled_circuits) != len(  # pragma: no cover
         linear_combination_coeffs
     ):
         raise AssertionError(
@@ -83,11 +79,16 @@ def execute_with_lre(
             + "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], QuantumResult],
+    executor: Callable[[Circuit], float],
     degree: int,
     fold_multiplier: int,
     folding_method: Callable[
@@ -100,7 +101,7 @@ def mitigate_executor(
 
     Args:
         input_circuit: Circuit to be scaled.
-        executor: Executes a circuit and returns a `QuantumResult`
+        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.
@@ -134,15 +135,13 @@ def lre_decorator(
     fold_multiplier: int,
     folding_method: Callable[[Circuit, float], Circuit] = fold_gates_at_random,
     num_chunks: Optional[int] = None,
-) -> Callable[
-    [Callable[[Circuit], QuantumResult]], Callable[[Circuit], float]
-]:
+) -> 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 `QuantumResult`
+        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.
@@ -158,7 +157,7 @@ def lre_decorator(
     """
 
     def decorator(
-        executor: Callable[[Circuit], QuantumResult],
+        executor: Callable[[Circuit], float],
     ) -> Callable[[Circuit], float]:
         return mitigate_executor(
             executor,

mitiq/lre/tests/test_lre.py

@@ -28,24 +28,20 @@ def execute(circuit, noise_level=0.025):
 ideal_val = execute(test_cirq, noise_level=0)
 
 
-@pytest.mark.parametrize(
-    "input_degree, input_fold_multiplier", [(2, 2), (2, 3), (3, 4)]
-)
-def test_lre_exp_value(input_degree, input_fold_multiplier):
+@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=input_degree,
-        fold_multiplier=input_fold_multiplier,
+        degree=degree,
+        fold_multiplier=fold_multiplier,
     )
     assert abs(lre_exp_val - ideal_val) <= abs(noisy_val - ideal_val)
 
 
-@pytest.mark.parametrize(
-    "input_degree, input_fold_multiplier", [(2, 2), (2, 3), (3, 4)]
-)
-def test_lre_exp_value_decorator(input_degree, input_fold_multiplier):
+@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
@@ -105,21 +101,18 @@ def test_lre_executor_with_chunking():
     assert abs(lre_exp_val - ideal_val) <= abs(noisy_val - ideal_val)
 
 
-@pytest.mark.parametrize(
-    "test_input", [(1), (2), (3), (4), (5), (6), (7), (8), (9)]
-)
-@pytest.mark.xfail
-def test_lre_executor_with_chunking_failures(test_input):
-    """Verify chunking fails when a large number of layers are chunked into a
-    smaller number of circuit chunks."""
+@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=test_input
+        test_cirq, execute, degree=2, fold_multiplier=2, num_chunks=num_chunks
     )
-    assert abs(lre_exp_val - ideal_val) <= abs(noisy_val - ideal_val)
+    assert abs(lre_exp_val - ideal_val) >= abs(noisy_val - ideal_val)
 
 
 @pytest.mark.parametrize("input_method", [(fold_global), (fold_all)])