From 6614641f5fd67b702ccb99c47fc68dd04b94327f Mon Sep 17 00:00:00 2001
From: Purva Thakre <purva.athakre@gmail.com>
Date: Wed, 11 Sep 2024 09:11:02 -0500
Subject: [PATCH 01/16] main functions without corrected mypy errors, unit
 tests and docstrings

---
 docs/source/apidoc.md       |   5 ++
 mitiq/lre/__init__.py       |   4 +-
 mitiq/lre/lre.py            | 112 ++++++++++++++++++++++++++++++++++++
 mitiq/lre/tests/test_lre.py |  30 ++++++++++
 4 files changed, 150 insertions(+), 1 deletion(-)
 create mode 100644 mitiq/lre/lre.py
 create mode 100644 mitiq/lre/tests/test_lre.py

diff --git a/docs/source/apidoc.md b/docs/source/apidoc.md
index 9a3fbf10f..e9e5fe65a 100644
--- a/docs/source/apidoc.md
+++ b/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:
diff --git a/mitiq/lre/__init__.py b/mitiq/lre/__init__.py
index 71e51c78e..353568cdd 100644
--- a/mitiq/lre/__init__.py
+++ b/mitiq/lre/__init__.py
@@ -10,4 +10,6 @@
 from mitiq.lre.inference.multivariate_richardson import (
     multivariate_richardson_coefficients,
     sample_matrix,
-)
\ No newline at end of file
+)
+
+from mitiq.lre.lre import execute_with_lre, mitigate_executor, lre_decorator
\ No newline at end of file
diff --git a/mitiq/lre/lre.py b/mitiq/lre/lre.py
new file mode 100644
index 000000000..dd0354897
--- /dev/null
+++ b/mitiq/lre/lre.py
@@ -0,0 +1,112 @@
+# 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 Callable, Optional, Union
+
+import numpy as np
+from cirq import Circuit
+
+from mitiq import Executor, Observable, QuantumResult
+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: Union[Executor, Callable[[Circuit], QuantumResult]],
+    shots: int,
+    degree: int,
+    fold_multiplier: int,
+    folding_method: Callable[[Circuit, float], Circuit] = fold_gates_at_random,
+    num_chunks: Optional[int] = None,
+    observable: Optional[Observable] = None,
+)-> float:
+    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
+    )
+    normalized_shots_list = shots // len(linear_combination_coeffs)
+    rescaled_shots_list = [normalized_shots_list] * len(
+        linear_combination_coeffs
+    )
+
+    lre_exp_values = []
+    for circuit_shots, scaled_circuit in zip(
+        rescaled_shots_list, noise_scaled_circuits
+    ):
+        circ_exp_val = executor(scaled_circuit, shots=circuit_shots)
+        lre_exp_values.append(circ_exp_val)
+
+    # verify the linear combination coefficients and the calculated expectation
+    # values have the same length
+    assert len(lre_exp_values) == len(linear_combination_coeffs)
+
+    return np.dot(lre_exp_values, linear_combination_coeffs)
+
+
+def mitigate_executor(
+    executor: Union[Executor, Callable[[Circuit], QuantumResult]],
+    shots: int,
+    degree: int,
+    fold_multiplier: int,
+    folding_method: Callable[[Circuit, float], Circuit] = fold_gates_at_random,
+    num_chunks: Optional[int] = None,
+    observable: Optional[Observable] = None,
+) -> Callable[[Circuit], float]:
+    @wraps(executor)
+    def new_executor(input_circuit: Circuit) -> float:
+        return execute_with_lre(
+            input_circuit,
+            executor,
+            shots,
+            degree,
+            fold_multiplier,
+            folding_method,
+            num_chunks,
+            observable,
+        )
+
+    return new_executor
+
+
+def lre_decorator(
+    shots: int,
+    degree: int,
+    fold_multiplier: int,
+    folding_method: Callable[[Circuit, float], Circuit] = fold_gates_at_random,
+    num_chunks: Optional[int] = None,
+    observable: Optional[Observable] = None,
+) -> Callable[
+    [Callable[[Circuit], QuantumResult]], Callable[[Circuit], float]
+]:
+    # Raise an error if the decorator is used without parenthesis
+    if callable(observable):
+        raise TypeError(
+            "Decorator must be used with parentheses (i.e., @lre_decorator()) "
+            "with explicit arguments for shots, degree and fold_multiplier."
+        )
+
+    def decorator(
+        executor: Callable[[Circuit], QuantumResult],
+    ) -> Callable[[Circuit], float]:
+        return mitigate_executor(
+            executor,
+            shots,
+            degree,
+            fold_multiplier,
+            folding_method,
+            num_chunks,
+            observable,
+        )
+
+    return decorator
diff --git a/mitiq/lre/tests/test_lre.py b/mitiq/lre/tests/test_lre.py
new file mode 100644
index 000000000..e4ef5c4e9
--- /dev/null
+++ b/mitiq/lre/tests/test_lre.py
@@ -0,0 +1,30 @@
+"""Unit tests for the LRE extrapolation methods."""
+
+from cirq import DensityMatrixSimulator, depolarize
+
+from mitiq import benchmarks
+from mitiq.lre import execute_with_lre
+
+test_cirq = benchmarks.generate_rb_circuits(
+    n_qubits=1,
+    num_cliffords=2,
+)[0]
+
+
+def execute(circuit, noise_level=0.025, shots=1000):
+    """Returns Tr[ρ |0⟩⟨0|] where ρ is the state prepared by the circuit
+    executed with depolarizing noise.
+    """
+    # Replace with code based on your frontend and backend.
+    mitiq_circuit = circuit
+    noisy_circuit = mitiq_circuit.with_noise(depolarize(p=noise_level))
+    rho = DensityMatrixSimulator().simulate(noisy_circuit).final_density_matrix
+    return rho[0, 0].real
+
+
+def test_lre_exp_value():
+    noisy_val = execute(test_cirq)
+    ideal_val = execute(test_cirq, noise_level=0, shots=1000)
+    assert abs(ideal_val - noisy_val) > 0
+    lre_exp_val = execute_with_lre(test_cirq, execute, 1000, 2, 2)
+    assert lre_exp_val > noisy_val

From 3db28178e8c23d5ed08de193aec4a2abcfd8cc3c Mon Sep 17 00:00:00 2001
From: Purva Thakre <purva.athakre@gmail.com>
Date: Wed, 11 Sep 2024 14:36:11 -0500
Subject: [PATCH 02/16] add unit tests for the decorator - check for test
 coverage

---
 mitiq/lre/lre.py            |  9 +------
 mitiq/lre/tests/test_lre.py | 51 +++++++++++++++++++++++++++++++++----
 2 files changed, 47 insertions(+), 13 deletions(-)

diff --git a/mitiq/lre/lre.py b/mitiq/lre/lre.py
index dd0354897..32ce2cf86 100644
--- a/mitiq/lre/lre.py
+++ b/mitiq/lre/lre.py
@@ -28,7 +28,7 @@ def execute_with_lre(
     folding_method: Callable[[Circuit, float], Circuit] = fold_gates_at_random,
     num_chunks: Optional[int] = None,
     observable: Optional[Observable] = None,
-)-> float:
+) -> float:
     noise_scaled_circuits = multivariate_layer_scaling(
         input_circuit, degree, fold_multiplier, num_chunks, folding_method
     )
@@ -89,13 +89,6 @@ def lre_decorator(
 ) -> Callable[
     [Callable[[Circuit], QuantumResult]], Callable[[Circuit], float]
 ]:
-    # Raise an error if the decorator is used without parenthesis
-    if callable(observable):
-        raise TypeError(
-            "Decorator must be used with parentheses (i.e., @lre_decorator()) "
-            "with explicit arguments for shots, degree and fold_multiplier."
-        )
-
     def decorator(
         executor: Callable[[Circuit], QuantumResult],
     ) -> Callable[[Circuit], float]:
diff --git a/mitiq/lre/tests/test_lre.py b/mitiq/lre/tests/test_lre.py
index e4ef5c4e9..93addfb91 100644
--- a/mitiq/lre/tests/test_lre.py
+++ b/mitiq/lre/tests/test_lre.py
@@ -1,9 +1,12 @@
 """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
+from mitiq.lre import execute_with_lre, lre_decorator, mitigate_executor
 
 test_cirq = benchmarks.generate_rb_circuits(
     n_qubits=1,
@@ -12,9 +15,6 @@
 
 
 def execute(circuit, noise_level=0.025, shots=1000):
-    """Returns Tr[ρ |0⟩⟨0|] where ρ is the state prepared by the circuit
-    executed with depolarizing noise.
-    """
     # Replace with code based on your frontend and backend.
     mitiq_circuit = circuit
     noisy_circuit = mitiq_circuit.with_noise(depolarize(p=noise_level))
@@ -22,9 +22,50 @@ def execute(circuit, noise_level=0.025, shots=1000):
     return rho[0, 0].real
 
 
+noisy_val = execute(test_cirq)
+
+
 def test_lre_exp_value():
-    noisy_val = execute(test_cirq)
     ideal_val = execute(test_cirq, noise_level=0, shots=1000)
     assert abs(ideal_val - noisy_val) > 0
     lre_exp_val = execute_with_lre(test_cirq, execute, 1000, 2, 2)
     assert lre_exp_val > noisy_val
+
+    # verify the mitigated decorator work as expected
+    mitigated_executor = mitigate_executor(execute, 1000, 2, 2)
+    exp_val_from_mitigate_executor = mitigated_executor(test_cirq)
+    assert exp_val_from_mitigate_executor > noisy_val
+
+
+def test_lre_decorator():
+    @lre_decorator(100, 2, 2)
+    def execute(circuit, noise_level=0.025, shots=100):
+        # Replace with code based on your frontend and backend.
+        mitiq_circuit = circuit
+        noisy_circuit = mitiq_circuit.with_noise(depolarize(p=noise_level))
+        rho = (
+            DensityMatrixSimulator()
+            .simulate(noisy_circuit)
+            .final_density_matrix
+        )
+        return rho[0, 0].real
+
+    assert noisy_val < execute(test_cirq)
+
+
+def test_lre_decorator_raised_error():
+    with pytest.raises(TypeError, match=re.escape("lre_decorator() missing")):
+
+        @lre_decorator()
+        def execute(circuit, noise_level=0.025, shots=100):
+            # Replace with code based on your frontend and backend.
+            mitiq_circuit = circuit
+            noisy_circuit = mitiq_circuit.with_noise(depolarize(p=noise_level))
+            rho = (
+                DensityMatrixSimulator()
+                .simulate(noisy_circuit)
+                .final_density_matrix
+            )
+            return rho[0, 0].real
+
+        assert noisy_val < execute(test_cirq)

From 759fbc5c577ce20544ec8af0867e35e8978f201b Mon Sep 17 00:00:00 2001
From: Purva Thakre <purva.athakre@gmail.com>
Date: Thu, 12 Sep 2024 07:42:38 -0500
Subject: [PATCH 03/16] fix
 https://github.com/unitaryfund/mitiq/issues/2500#issuecomment-2344872470

---
 .gitignore | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.gitignore b/.gitignore
index 523b7c4a7..1770047cd 100644
--- a/.gitignore
+++ b/.gitignore
@@ -15,7 +15,7 @@ mitiq.egg-info/
 dist/
 build/
 jupyter_execute/
-
+.mypy_cache/
 # Coverage reports
 coverage.xml
 .coverage

From 2c35e74ded2683901a2e2b0ca473a6b35048803a Mon Sep 17 00:00:00 2001
From: Purva Thakre <purva.athakre@gmail.com>
Date: Thu, 12 Sep 2024 09:39:58 -0500
Subject: [PATCH 04/16] mypy - remove shots

---
 mitiq/lre/lre.py            | 31 ++++++++++++-------------------
 mitiq/lre/tests/test_lre.py | 14 +++++++-------
 2 files changed, 19 insertions(+), 26 deletions(-)

diff --git a/mitiq/lre/lre.py b/mitiq/lre/lre.py
index 32ce2cf86..f44418932 100644
--- a/mitiq/lre/lre.py
+++ b/mitiq/lre/lre.py
@@ -6,12 +6,12 @@
 """Extrapolation methods for Layerwise Richardson Extrapolation (LRE)"""
 
 from functools import wraps
-from typing import Callable, Optional, Union
+from typing import Any, Callable, Optional, Union
 
 import numpy as np
 from cirq import Circuit
 
-from mitiq import Executor, Observable, QuantumResult
+from mitiq import Observable, QuantumResult
 from mitiq.lre import (
     multivariate_layer_scaling,
     multivariate_richardson_coefficients,
@@ -21,11 +21,12 @@
 
 def execute_with_lre(
     input_circuit: Circuit,
-    executor: Union[Executor, Callable[[Circuit], QuantumResult]],
-    shots: int,
+    executor: Callable[[Circuit], QuantumResult],
     degree: int,
     fold_multiplier: int,
-    folding_method: Callable[[Circuit, float], Circuit] = fold_gates_at_random,
+    folding_method: Callable[
+        [Union[Any], float], Union[Any]
+    ] = fold_gates_at_random,
     num_chunks: Optional[int] = None,
     observable: Optional[Observable] = None,
 ) -> float:
@@ -35,16 +36,10 @@ def execute_with_lre(
     linear_combination_coeffs = multivariate_richardson_coefficients(
         input_circuit, degree, fold_multiplier, num_chunks
     )
-    normalized_shots_list = shots // len(linear_combination_coeffs)
-    rescaled_shots_list = [normalized_shots_list] * len(
-        linear_combination_coeffs
-    )
 
     lre_exp_values = []
-    for circuit_shots, scaled_circuit in zip(
-        rescaled_shots_list, noise_scaled_circuits
-    ):
-        circ_exp_val = executor(scaled_circuit, shots=circuit_shots)
+    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
@@ -55,11 +50,12 @@ def execute_with_lre(
 
 
 def mitigate_executor(
-    executor: Union[Executor, Callable[[Circuit], QuantumResult]],
-    shots: int,
+    executor: Callable[[Circuit], QuantumResult],
     degree: int,
     fold_multiplier: int,
-    folding_method: Callable[[Circuit, float], Circuit] = fold_gates_at_random,
+    folding_method: Callable[
+        [Union[Any], float], Union[Any]
+    ] = fold_gates_at_random,
     num_chunks: Optional[int] = None,
     observable: Optional[Observable] = None,
 ) -> Callable[[Circuit], float]:
@@ -68,7 +64,6 @@ def new_executor(input_circuit: Circuit) -> float:
         return execute_with_lre(
             input_circuit,
             executor,
-            shots,
             degree,
             fold_multiplier,
             folding_method,
@@ -80,7 +75,6 @@ def new_executor(input_circuit: Circuit) -> float:
 
 
 def lre_decorator(
-    shots: int,
     degree: int,
     fold_multiplier: int,
     folding_method: Callable[[Circuit, float], Circuit] = fold_gates_at_random,
@@ -94,7 +88,6 @@ def decorator(
     ) -> Callable[[Circuit], float]:
         return mitigate_executor(
             executor,
-            shots,
             degree,
             fold_multiplier,
             folding_method,
diff --git a/mitiq/lre/tests/test_lre.py b/mitiq/lre/tests/test_lre.py
index 93addfb91..ce696a6b2 100644
--- a/mitiq/lre/tests/test_lre.py
+++ b/mitiq/lre/tests/test_lre.py
@@ -14,7 +14,7 @@
 )[0]
 
 
-def execute(circuit, noise_level=0.025, shots=1000):
+def execute(circuit, noise_level=0.025):
     # Replace with code based on your frontend and backend.
     mitiq_circuit = circuit
     noisy_circuit = mitiq_circuit.with_noise(depolarize(p=noise_level))
@@ -26,20 +26,20 @@ def execute(circuit, noise_level=0.025, shots=1000):
 
 
 def test_lre_exp_value():
-    ideal_val = execute(test_cirq, noise_level=0, shots=1000)
+    ideal_val = execute(test_cirq, noise_level=0)
     assert abs(ideal_val - noisy_val) > 0
-    lre_exp_val = execute_with_lre(test_cirq, execute, 1000, 2, 2)
+    lre_exp_val = execute_with_lre(test_cirq, execute, 2, 2)
     assert lre_exp_val > noisy_val
 
     # verify the mitigated decorator work as expected
-    mitigated_executor = mitigate_executor(execute, 1000, 2, 2)
+    mitigated_executor = mitigate_executor(execute, 2, 2)
     exp_val_from_mitigate_executor = mitigated_executor(test_cirq)
     assert exp_val_from_mitigate_executor > noisy_val
 
 
 def test_lre_decorator():
-    @lre_decorator(100, 2, 2)
-    def execute(circuit, noise_level=0.025, shots=100):
+    @lre_decorator(2, 2)
+    def execute(circuit, noise_level=0.025):
         # Replace with code based on your frontend and backend.
         mitiq_circuit = circuit
         noisy_circuit = mitiq_circuit.with_noise(depolarize(p=noise_level))
@@ -57,7 +57,7 @@ def test_lre_decorator_raised_error():
     with pytest.raises(TypeError, match=re.escape("lre_decorator() missing")):
 
         @lre_decorator()
-        def execute(circuit, noise_level=0.025, shots=100):
+        def execute(circuit, noise_level=0.025):
             # Replace with code based on your frontend and backend.
             mitiq_circuit = circuit
             noisy_circuit = mitiq_circuit.with_noise(depolarize(p=noise_level))

From d0f005293b5e0a40bcb65f0f0fb5ab4d335497a5 Mon Sep 17 00:00:00 2001
From: Purva Thakre <purva.athakre@gmail.com>
Date: Fri, 13 Sep 2024 04:41:27 -0500
Subject: [PATCH 05/16] more unit tests + docstrings

---
 mitiq/lre/lre.py            | 75 ++++++++++++++++++++++++++++++++++---
 mitiq/lre/tests/test_lre.py | 50 +++++++++++++++++++++----
 2 files changed, 112 insertions(+), 13 deletions(-)

diff --git a/mitiq/lre/lre.py b/mitiq/lre/lre.py
index f44418932..56e0e9a64 100644
--- a/mitiq/lre/lre.py
+++ b/mitiq/lre/lre.py
@@ -11,7 +11,7 @@
 import numpy as np
 from cirq import Circuit
 
-from mitiq import Observable, QuantumResult
+from mitiq import QuantumResult
 from mitiq.lre import (
     multivariate_layer_scaling,
     multivariate_richardson_coefficients,
@@ -28,8 +28,37 @@ def execute_with_lre(
         [Union[Any], float], Union[Any]
     ] = fold_gates_at_random,
     num_chunks: Optional[int] = None,
-    observable: Optional[Observable] = 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 `QuantumResult`
+        degree: Degree of the multivariate polynomial.
+        fold_multiplier: Scaling gap required by unitary folding.
+        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
     )
@@ -57,8 +86,26 @@ def mitigate_executor(
         [Union[Any], float], Union[Any]
     ] = fold_gates_at_random,
     num_chunks: Optional[int] = None,
-    observable: Optional[Observable] = 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 `QuantumResult`
+    degree: Degree of the multivariate polynomial.
+    fold_multiplier: Scaling gap required by unitary folding.
+    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(
@@ -68,7 +115,6 @@ def new_executor(input_circuit: Circuit) -> float:
             fold_multiplier,
             folding_method,
             num_chunks,
-            observable,
         )
 
     return new_executor
@@ -79,10 +125,28 @@ def lre_decorator(
     fold_multiplier: int,
     folding_method: Callable[[Circuit, float], Circuit] = fold_gates_at_random,
     num_chunks: Optional[int] = None,
-    observable: Optional[Observable] = None,
 ) -> Callable[
     [Callable[[Circuit], QuantumResult]], 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`
+    degree: Degree of the multivariate polynomial.
+    fold_multiplier: Scaling gap required by unitary folding.
+    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], QuantumResult],
     ) -> Callable[[Circuit], float]:
@@ -92,7 +156,6 @@ def decorator(
             fold_multiplier,
             folding_method,
             num_chunks,
-            observable,
         )
 
     return decorator
diff --git a/mitiq/lre/tests/test_lre.py b/mitiq/lre/tests/test_lre.py
index ce696a6b2..6ae39b222 100644
--- a/mitiq/lre/tests/test_lre.py
+++ b/mitiq/lre/tests/test_lre.py
@@ -7,7 +7,9 @@
 
 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,
@@ -15,9 +17,9 @@
 
 
 def execute(circuit, noise_level=0.025):
+    """Default executor for all unit tests."""
     # Replace with code based on your frontend and backend.
-    mitiq_circuit = circuit
-    noisy_circuit = mitiq_circuit.with_noise(depolarize(p=noise_level))
+    noisy_circuit = circuit.with_noise(depolarize(p=noise_level))
     rho = DensityMatrixSimulator().simulate(noisy_circuit).final_density_matrix
     return rho[0, 0].real
 
@@ -26,21 +28,27 @@ def execute(circuit, noise_level=0.025):
 
 
 def test_lre_exp_value():
+    """Verify LRE executors work as expected."""
     ideal_val = execute(test_cirq, noise_level=0)
     assert abs(ideal_val - noisy_val) > 0
-    lre_exp_val = execute_with_lre(test_cirq, execute, 2, 2)
+    lre_exp_val = execute_with_lre(
+        test_cirq, execute, degree=2, fold_multiplier=2
+    )
     assert lre_exp_val > noisy_val
 
     # verify the mitigated decorator work as expected
-    mitigated_executor = mitigate_executor(execute, 2, 2)
+    mitigated_executor = mitigate_executor(
+        execute, degree=2, fold_multiplier=2
+    )
     exp_val_from_mitigate_executor = mitigated_executor(test_cirq)
     assert exp_val_from_mitigate_executor > noisy_val
 
 
 def test_lre_decorator():
-    @lre_decorator(2, 2)
+    """Verify LRE decorators work as expected."""
+
+    @lre_decorator(degree=2, fold_multiplier=2)
     def execute(circuit, noise_level=0.025):
-        # Replace with code based on your frontend and backend.
         mitiq_circuit = circuit
         noisy_circuit = mitiq_circuit.with_noise(depolarize(p=noise_level))
         rho = (
@@ -54,11 +62,12 @@ def execute(circuit, noise_level=0.025):
 
 
 def test_lre_decorator_raised_error():
+    """Verify error is raised when the defualt 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):
-            # Replace with code based on your frontend and backend.
             mitiq_circuit = circuit
             noisy_circuit = mitiq_circuit.with_noise(depolarize(p=noise_level))
             rho = (
@@ -69,3 +78,30 @@ def execute(circuit, noise_level=0.025):
             return rho[0, 0].real
 
         assert noisy_val < execute(test_cirq)
+
+
+def test_lre_executor_with_chunking():
+    """Verify the executor works as expected for chunking a large circuit into
+    a smaller circuit."""
+    ideal_val = execute(test_cirq * 20, noise_level=0)
+    assert abs(ideal_val - noisy_val) > 0
+    lre_exp_val = execute_with_lre(
+        test_cirq, execute, degree=2, fold_multiplier=2, num_chunks=5
+    )
+    assert lre_exp_val > noisy_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."""
+    ideal_val = execute(test_cirq, noise_level=0)
+    assert abs(ideal_val - noisy_val) > 0
+    lre_exp_val = execute_with_lre(
+        test_cirq,
+        execute,
+        degree=2,
+        fold_multiplier=2,
+        folding_method=input_method,
+    )
+    assert lre_exp_val > noisy_val

From 7708f50c458fe655a37ea40581339e8bee91c6d7 Mon Sep 17 00:00:00 2001
From: Purva Thakre <purva.athakre@gmail.com>
Date: Fri, 13 Sep 2024 04:45:07 -0500
Subject: [PATCH 06/16] dosctring args formatting

---
 mitiq/lre/lre.py | 36 ++++++++++++++++++------------------
 1 file changed, 18 insertions(+), 18 deletions(-)

diff --git a/mitiq/lre/lre.py b/mitiq/lre/lre.py
index 56e0e9a64..3c7022538 100644
--- a/mitiq/lre/lre.py
+++ b/mitiq/lre/lre.py
@@ -91,15 +91,15 @@ def mitigate_executor(
     error-mitigated with layerwise richardson extrapolation (LRE).
 
     Args:
-    input_circuit: Circuit to be scaled.
-    executor: Executes a circuit and returns a `QuantumResult`
-    degree: Degree of the multivariate polynomial.
-    fold_multiplier: Scaling gap required by unitary folding.
-    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.
+        input_circuit: Circuit to be scaled.
+        executor: Executes a circuit and returns a `QuantumResult`
+        degree: Degree of the multivariate polynomial.
+        fold_multiplier: Scaling gap required by unitary folding.
+        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:
@@ -132,15 +132,15 @@ def lre_decorator(
     layerwise richardson extrapolation (LRE).
 
     Args:
-    input_circuit: Circuit to be scaled.
-    executor: Executes a circuit and returns a `QuantumResult`
-    degree: Degree of the multivariate polynomial.
-    fold_multiplier: Scaling gap required by unitary folding.
-    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.
+        input_circuit: Circuit to be scaled.
+        executor: Executes a circuit and returns a `QuantumResult`
+        degree: Degree of the multivariate polynomial.
+        fold_multiplier: Scaling gap required by unitary folding.
+        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:

From 003d3a116b10ef56227cfb4dbdf85b2e8826c69d Mon Sep 17 00:00:00 2001
From: Purva Thakre <purva.athakre@gmail.com>
Date: Fri, 13 Sep 2024 04:55:16 -0500
Subject: [PATCH 07/16] fix
 https://github.com/unitaryfund/mitiq/pull/2499#discussion_r1758547125

---
 mitiq/lre/tests/test_lre.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/mitiq/lre/tests/test_lre.py b/mitiq/lre/tests/test_lre.py
index 6ae39b222..c0ee76acc 100644
--- a/mitiq/lre/tests/test_lre.py
+++ b/mitiq/lre/tests/test_lre.py
@@ -27,12 +27,18 @@ def execute(circuit, noise_level=0.025):
 noisy_val = execute(test_cirq)
 
 
-def test_lre_exp_value():
+@pytest.mark.parametrize(
+    "input_degree, input_fold_multiplier", [(2, 2), (2, 3), (3, 4)]
+)
+def test_lre_exp_value(input_degree, input_fold_multiplier):
     """Verify LRE executors work as expected."""
     ideal_val = execute(test_cirq, noise_level=0)
     assert abs(ideal_val - noisy_val) > 0
     lre_exp_val = execute_with_lre(
-        test_cirq, execute, degree=2, fold_multiplier=2
+        test_cirq,
+        execute,
+        degree=input_degree,
+        fold_multiplier=input_fold_multiplier,
     )
     assert lre_exp_val > noisy_val
 

From c970474d247fa716918a120547bac0ace84850b0 Mon Sep 17 00:00:00 2001
From: Purva Thakre <purva.athakre@gmail.com>
Date: Fri, 13 Sep 2024 05:04:44 -0500
Subject: [PATCH 08/16] weird chunking failure

---
 mitiq/lre/tests/test_lre.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mitiq/lre/tests/test_lre.py b/mitiq/lre/tests/test_lre.py
index c0ee76acc..71da9e00f 100644
--- a/mitiq/lre/tests/test_lre.py
+++ b/mitiq/lre/tests/test_lre.py
@@ -89,7 +89,7 @@ def execute(circuit, noise_level=0.025):
 def test_lre_executor_with_chunking():
     """Verify the executor works as expected for chunking a large circuit into
     a smaller circuit."""
-    ideal_val = execute(test_cirq * 20, noise_level=0)
+    ideal_val = execute(test_cirq * 200, noise_level=0)
     assert abs(ideal_val - noisy_val) > 0
     lre_exp_val = execute_with_lre(
         test_cirq, execute, degree=2, fold_multiplier=2, num_chunks=5

From 9146116c6e210ec0016bbc566684df31f670e459 Mon Sep 17 00:00:00 2001
From: Purva Thakre <purva.athakre@gmail.com>
Date: Fri, 13 Sep 2024 05:16:02 -0500
Subject: [PATCH 09/16] try chunking to 2

---
 mitiq/lre/tests/test_lre.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/mitiq/lre/tests/test_lre.py b/mitiq/lre/tests/test_lre.py
index 71da9e00f..c1e4a819c 100644
--- a/mitiq/lre/tests/test_lre.py
+++ b/mitiq/lre/tests/test_lre.py
@@ -18,7 +18,6 @@
 
 def execute(circuit, noise_level=0.025):
     """Default executor for all unit tests."""
-    # Replace with code based on your frontend and backend.
     noisy_circuit = circuit.with_noise(depolarize(p=noise_level))
     rho = DensityMatrixSimulator().simulate(noisy_circuit).final_density_matrix
     return rho[0, 0].real
@@ -92,7 +91,7 @@ def test_lre_executor_with_chunking():
     ideal_val = execute(test_cirq * 200, noise_level=0)
     assert abs(ideal_val - noisy_val) > 0
     lre_exp_val = execute_with_lre(
-        test_cirq, execute, degree=2, fold_multiplier=2, num_chunks=5
+        test_cirq, execute, degree=2, fold_multiplier=2, num_chunks=2
     )
     assert lre_exp_val > noisy_val
 

From d926c5c593fb6a0096519ecafe00b9f8a0272d2c Mon Sep 17 00:00:00 2001
From: Purva Thakre <purva.athakre@gmail.com>
Date: Tue, 24 Sep 2024 11:13:49 -0500
Subject: [PATCH 10/16] num_chunks = 5 with test_cirq * 200

---
 mitiq/lre/tests/test_lre.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mitiq/lre/tests/test_lre.py b/mitiq/lre/tests/test_lre.py
index c1e4a819c..5625bf81f 100644
--- a/mitiq/lre/tests/test_lre.py
+++ b/mitiq/lre/tests/test_lre.py
@@ -91,7 +91,7 @@ def test_lre_executor_with_chunking():
     ideal_val = execute(test_cirq * 200, noise_level=0)
     assert abs(ideal_val - noisy_val) > 0
     lre_exp_val = execute_with_lre(
-        test_cirq, execute, degree=2, fold_multiplier=2, num_chunks=2
+        test_cirq * 200, execute, degree=2, fold_multiplier=2, num_chunks=5
     )
     assert lre_exp_val > noisy_val
 

From 595360a9c53b8d0397ac6f93a20d441608a8272d Mon Sep 17 00:00:00 2001
From: Purva Thakre <purva.athakre@gmail.com>
Date: Wed, 25 Sep 2024 17:51:51 -0500
Subject: [PATCH 11/16] push to check for test coverage

---
 mitiq/lre/lre.py            | 18 +++++++++++++-----
 mitiq/lre/tests/test_lre.py |  4 ++--
 2 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/mitiq/lre/lre.py b/mitiq/lre/lre.py
index 3c7022538..b76d3fdc8 100644
--- a/mitiq/lre/lre.py
+++ b/mitiq/lre/lre.py
@@ -11,7 +11,7 @@
 import numpy as np
 from cirq import Circuit
 
-from mitiq import QuantumResult
+from mitiq import QPROGRAM, QuantumResult
 from mitiq.lre import (
     multivariate_layer_scaling,
     multivariate_richardson_coefficients,
@@ -25,8 +25,8 @@ def execute_with_lre(
     degree: int,
     fold_multiplier: int,
     folding_method: Callable[
-        [Union[Any], float], Union[Any]
-    ] = fold_gates_at_random,
+        [QPROGRAM, float], QPROGRAM
+    ] = fold_gates_at_random,  # type: ignore [has-type]
     num_chunks: Optional[int] = None,
 ) -> float:
     r"""
@@ -47,7 +47,8 @@ def execute_with_lre(
         input_circuit: Circuit to be scaled.
         executor: Executes a circuit and returns a `QuantumResult`
         degree: Degree of the multivariate polynomial.
-        fold_multiplier: Scaling gap required by unitary folding.
+        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
@@ -73,7 +74,14 @@ def execute_with_lre(
 
     # verify the linear combination coefficients and the calculated expectation
     # values have the same length
-    assert len(lre_exp_values) == len(linear_combination_coeffs)
+    if not len(lre_exp_values) == 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."
+        )
 
     return np.dot(lre_exp_values, linear_combination_coeffs)
 
diff --git a/mitiq/lre/tests/test_lre.py b/mitiq/lre/tests/test_lre.py
index 5625bf81f..9fdb36266 100644
--- a/mitiq/lre/tests/test_lre.py
+++ b/mitiq/lre/tests/test_lre.py
@@ -88,10 +88,10 @@ def execute(circuit, noise_level=0.025):
 def test_lre_executor_with_chunking():
     """Verify the executor works as expected for chunking a large circuit into
     a smaller circuit."""
-    ideal_val = execute(test_cirq * 200, noise_level=0)
+    ideal_val = execute(test_cirq * 4, noise_level=0)
     assert abs(ideal_val - noisy_val) > 0
     lre_exp_val = execute_with_lre(
-        test_cirq * 200, execute, degree=2, fold_multiplier=2, num_chunks=5
+        test_cirq * 4, execute, degree=2, fold_multiplier=2, num_chunks=5
     )
     assert lre_exp_val > noisy_val
 

From 41d08f50cdee5e1a08df79d2db76457db609e017 Mon Sep 17 00:00:00 2001
From: Purva Thakre <purva.athakre@gmail.com>
Date: Wed, 25 Sep 2024 18:12:29 -0500
Subject: [PATCH 12/16] chunking failures

---
 mitiq/lre/tests/test_lre.py | 49 ++++++++++++++++++++++++++++---------
 1 file changed, 38 insertions(+), 11 deletions(-)

diff --git a/mitiq/lre/tests/test_lre.py b/mitiq/lre/tests/test_lre.py
index 9fdb36266..974280be2 100644
--- a/mitiq/lre/tests/test_lre.py
+++ b/mitiq/lre/tests/test_lre.py
@@ -24,6 +24,7 @@ def execute(circuit, noise_level=0.025):
 
 
 noisy_val = execute(test_cirq)
+ideal_val = execute(test_cirq, noise_level=0)
 
 
 @pytest.mark.parametrize(
@@ -31,7 +32,6 @@ def execute(circuit, noise_level=0.025):
 )
 def test_lre_exp_value(input_degree, input_fold_multiplier):
     """Verify LRE executors work as expected."""
-    ideal_val = execute(test_cirq, noise_level=0)
     assert abs(ideal_val - noisy_val) > 0
     lre_exp_val = execute_with_lre(
         test_cirq,
@@ -39,14 +39,16 @@ def test_lre_exp_value(input_degree, input_fold_multiplier):
         degree=input_degree,
         fold_multiplier=input_fold_multiplier,
     )
-    assert lre_exp_val > noisy_val
+    assert abs(lre_exp_val - ideal_val) <= abs(noisy_val - ideal_val)
 
     # verify the mitigated decorator work as expected
     mitigated_executor = mitigate_executor(
         execute, degree=2, fold_multiplier=2
     )
     exp_val_from_mitigate_executor = mitigated_executor(test_cirq)
-    assert exp_val_from_mitigate_executor > noisy_val
+    assert abs(exp_val_from_mitigate_executor - ideal_val) <= abs(
+        noisy_val - ideal_val
+    )
 
 
 def test_lre_decorator():
@@ -63,12 +65,12 @@ def execute(circuit, noise_level=0.025):
         )
         return rho[0, 0].real
 
-    assert noisy_val < execute(test_cirq)
+    assert abs(execute(test_cirq) - ideal_val) <= abs(noisy_val - ideal_val)
 
 
 def test_lre_decorator_raised_error():
-    """Verify error is raised when the defualt parameters for the decorator are
-    not specified."""
+    """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()
@@ -82,18 +84,43 @@ def execute(circuit, noise_level=0.025):
             )
             return rho[0, 0].real
 
-        assert noisy_val < execute(test_cirq)
+        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."""
-    ideal_val = execute(test_cirq * 4, noise_level=0)
+    # define a larger circuit
+    test_cirq = benchmarks.generate_rb_circuits(n_qubits=1, num_cliffords=12)[
+        0
+    ]
+    ideal_val = execute(test_cirq, noise_level=0)
+    assert abs(ideal_val - noisy_val) > 0
+    lre_exp_val = execute_with_lre(
+        test_cirq, execute, degree=2, fold_multiplier=2, num_chunks=10
+    )
+    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 layers."""
+    # define a larger circuit
+    test_cirq = benchmarks.generate_rb_circuits(n_qubits=1, num_cliffords=15)[
+        0
+    ]
+    ideal_val = execute(test_cirq, noise_level=0)
     assert abs(ideal_val - noisy_val) > 0
     lre_exp_val = execute_with_lre(
-        test_cirq * 4, execute, degree=2, fold_multiplier=2, num_chunks=5
+        test_cirq, execute, degree=2, fold_multiplier=2, num_chunks=test_input
     )
-    assert lre_exp_val > noisy_val
+    assert abs(lre_exp_val - ideal_val) <= abs(noisy_val - ideal_val)
 
 
 @pytest.mark.parametrize("input_method", [(fold_global), (fold_all)])
@@ -109,4 +136,4 @@ def test_lre_executor_with_different_folding_methods(input_method):
         fold_multiplier=2,
         folding_method=input_method,
     )
-    assert lre_exp_val > noisy_val
+    assert abs(lre_exp_val - ideal_val) <= abs(noisy_val - ideal_val)

From ce46c58a332d0380aaf4950a7bc322564342c6e4 Mon Sep 17 00:00:00 2001
From: Purva Thakre <purva.athakre@gmail.com>
Date: Fri, 27 Sep 2024 11:36:17 -0500
Subject: [PATCH 13/16] split decorator unit test

---
 mitiq/lre/tests/test_lre.py | 15 ++++++++-------
 1 file changed, 8 insertions(+), 7 deletions(-)

diff --git a/mitiq/lre/tests/test_lre.py b/mitiq/lre/tests/test_lre.py
index 974280be2..5c4229feb 100644
--- a/mitiq/lre/tests/test_lre.py
+++ b/mitiq/lre/tests/test_lre.py
@@ -32,7 +32,6 @@ def execute(circuit, noise_level=0.025):
 )
 def test_lre_exp_value(input_degree, input_fold_multiplier):
     """Verify LRE executors work as expected."""
-    assert abs(ideal_val - noisy_val) > 0
     lre_exp_val = execute_with_lre(
         test_cirq,
         execute,
@@ -41,6 +40,14 @@ def test_lre_exp_value(input_degree, input_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):
+    """Verify LRE executors work as expected."""
+
+    # move to a separate test
     # verify the mitigated decorator work as expected
     mitigated_executor = mitigate_executor(
         execute, degree=2, fold_multiplier=2
@@ -96,8 +103,6 @@ def test_lre_executor_with_chunking():
     test_cirq = benchmarks.generate_rb_circuits(n_qubits=1, num_cliffords=12)[
         0
     ]
-    ideal_val = execute(test_cirq, noise_level=0)
-    assert abs(ideal_val - noisy_val) > 0
     lre_exp_val = execute_with_lre(
         test_cirq, execute, degree=2, fold_multiplier=2, num_chunks=10
     )
@@ -115,8 +120,6 @@ def test_lre_executor_with_chunking_failures(test_input):
     test_cirq = benchmarks.generate_rb_circuits(n_qubits=1, num_cliffords=15)[
         0
     ]
-    ideal_val = execute(test_cirq, noise_level=0)
-    assert abs(ideal_val - noisy_val) > 0
     lre_exp_val = execute_with_lre(
         test_cirq, execute, degree=2, fold_multiplier=2, num_chunks=test_input
     )
@@ -127,8 +130,6 @@ def test_lre_executor_with_chunking_failures(test_input):
 def test_lre_executor_with_different_folding_methods(input_method):
     """Verify the executor works as expected for using non-default unitary
     folding methods."""
-    ideal_val = execute(test_cirq, noise_level=0)
-    assert abs(ideal_val - noisy_val) > 0
     lre_exp_val = execute_with_lre(
         test_cirq,
         execute,

From 4561acbcf4fa3a737b29110c2a88069e49d3b26a Mon Sep 17 00:00:00 2001
From: Purva Thakre <purva.athakre@gmail.com>
Date: Fri, 27 Sep 2024 14:54:15 -0500
Subject: [PATCH 14/16] cleanup

---
 mitiq/lre/lre.py            |  6 ++++--
 mitiq/lre/tests/test_lre.py | 14 +++++---------
 2 files changed, 9 insertions(+), 11 deletions(-)

diff --git a/mitiq/lre/lre.py b/mitiq/lre/lre.py
index b76d3fdc8..fb5159f7c 100644
--- a/mitiq/lre/lre.py
+++ b/mitiq/lre/lre.py
@@ -102,7 +102,8 @@ def mitigate_executor(
         input_circuit: Circuit to be scaled.
         executor: Executes a circuit and returns a `QuantumResult`
         degree: Degree of the multivariate polynomial.
-        fold_multiplier: Scaling gap required by unitary folding.
+        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
@@ -143,7 +144,8 @@ def lre_decorator(
         input_circuit: Circuit to be scaled.
         executor: Executes a circuit and returns a `QuantumResult`
         degree: Degree of the multivariate polynomial.
-        fold_multiplier: Scaling gap required by unitary folding.
+        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
diff --git a/mitiq/lre/tests/test_lre.py b/mitiq/lre/tests/test_lre.py
index 5c4229feb..894497028 100644
--- a/mitiq/lre/tests/test_lre.py
+++ b/mitiq/lre/tests/test_lre.py
@@ -16,6 +16,7 @@
 )[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))
@@ -45,10 +46,7 @@ def test_lre_exp_value(input_degree, input_fold_multiplier):
     "input_degree, input_fold_multiplier", [(2, 2), (2, 3), (3, 4)]
 )
 def test_lre_exp_value_decorator(input_degree, input_fold_multiplier):
-    """Verify LRE executors work as expected."""
-
-    # move to a separate test
-    # verify the mitigated decorator work as expected
+    """Verify LRE mitigated executor work as expected."""
     mitigated_executor = mitigate_executor(
         execute, degree=2, fold_multiplier=2
     )
@@ -63,8 +61,7 @@ def test_lre_decorator():
 
     @lre_decorator(degree=2, fold_multiplier=2)
     def execute(circuit, noise_level=0.025):
-        mitiq_circuit = circuit
-        noisy_circuit = mitiq_circuit.with_noise(depolarize(p=noise_level))
+        noisy_circuit = circuit.with_noise(depolarize(p=noise_level))
         rho = (
             DensityMatrixSimulator()
             .simulate(noisy_circuit)
@@ -82,8 +79,7 @@ def test_lre_decorator_raised_error():
 
         @lre_decorator()
         def execute(circuit, noise_level=0.025):
-            mitiq_circuit = circuit
-            noisy_circuit = mitiq_circuit.with_noise(depolarize(p=noise_level))
+            noisy_circuit = circuit.with_noise(depolarize(p=noise_level))
             rho = (
                 DensityMatrixSimulator()
                 .simulate(noisy_circuit)
@@ -115,7 +111,7 @@ def test_lre_executor_with_chunking():
 @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 layers."""
+    smaller number of circuit chunks."""
     # define a larger circuit
     test_cirq = benchmarks.generate_rb_circuits(n_qubits=1, num_cliffords=15)[
         0

From 3d25cd2d127505ad1220095d773ac7a333eaf227 Mon Sep 17 00:00:00 2001
From: Purva Thakre <purva.athakre@gmail.com>
Date: Fri, 27 Sep 2024 15:41:42 -0500
Subject: [PATCH 15/16] chunking failure again + docker failure

---
 mitiq/lre/tests/test_lre.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mitiq/lre/tests/test_lre.py b/mitiq/lre/tests/test_lre.py
index 894497028..85dc1a01f 100644
--- a/mitiq/lre/tests/test_lre.py
+++ b/mitiq/lre/tests/test_lre.py
@@ -100,7 +100,7 @@ def test_lre_executor_with_chunking():
         0
     ]
     lre_exp_val = execute_with_lre(
-        test_cirq, execute, degree=2, fold_multiplier=2, num_chunks=10
+        test_cirq, execute, degree=2, fold_multiplier=2, num_chunks=14
     )
     assert abs(lre_exp_val - ideal_val) <= abs(noisy_val - ideal_val)
 

From 21aeb4d3580565432c8e5dbe356e67f441188db4 Mon Sep 17 00:00:00 2001
From: Purva Thakre <purva.athakre@gmail.com>
Date: Tue, 1 Oct 2024 14:35:02 -0500
Subject: [PATCH 16/16] nate's suggestions

---
 mitiq/lre/lre.py            | 31 +++++++++++++++----------------
 mitiq/lre/tests/test_lre.py | 31 ++++++++++++-------------------
 2 files changed, 27 insertions(+), 35 deletions(-)

diff --git a/mitiq/lre/lre.py b/mitiq/lre/lre.py
index fb5159f7c..c2fd7b52a 100644
--- a/mitiq/lre/lre.py
+++ b/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,
diff --git a/mitiq/lre/tests/test_lre.py b/mitiq/lre/tests/test_lre.py
index 85dc1a01f..4679e746a 100644
--- a/mitiq/lre/tests/test_lre.py
+++ b/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)])