Repository: unitaryfoundation/mitiq

Issue:
Make sure LRE executors work with observables
All the top-level functions in LRE should work with `observable`.  This issue is to make sure of this as well as add unit tests for the same.

Discussion:
@Shivansh20128: Hi! I can work on this issue if it is available. I am listing all the top-level functions I could find in the code (after comparing them to functions in ZNE).
1. `execute_with_lre`
2. `mitigate_executor`
3. `lre_decorator`

I hope these are the only functions when you say that all top-level functions in LRE should work with `observable`. Let me know if I missed any function.@natestemen: Yup those are correct! The main code changes will be in `execute_with_lre`, and then they should be easily propagated to `mitigate_executor` and `lre_decorator` through the addition of a new `observable` argument.@Shivansh20128: Great! You can assign this to me then.
Thanks@Shivansh20128: To work with observable, I think I need to change the `executor` as well ([here](https://github.com/unitaryfund/mitiq/blob/9e11cbce4348cc91e56a305966a18ee9f97ed4fd/mitiq/lre/lre.py#L85)).

So in the line:
```
    for scaled_circuit in noise_scaled_circuits:
        circ_exp_val = executor(scaled_circuit)
        lre_exp_values.append(circ_exp_val)
```
Since I need to pass the observables here, I need call the function `executor.evaluate()`.
```
    for scaled_circuit in noise_scaled_circuits:
        circ_exp_val = executor.evaluate(scaled_circuit, observables)
        lre_exp_values.append(circ_exp_val)
```
And for that to work, I also need to change the executor from the definition:
```
def execute_with_lre(
    input_circuit: QPROGRAM,
    executor: Callable[[QPROGRAM], float],
    ...
```
to:
```
def execute_with_lre(
    input_circuit: QPROGRAM,
    executor: Union[Executor, Callable[[QPROGRAM], float]],
    ...
```

This is further leading to more changes in the program and tests in the test file. Am I heading in the right direction?
Thanks@natestemen: Good find, that's exactly right! Sorry for not pointing this out earlier, but the first step is exactly what @bdg221 is working on in https://github.com/unitaryfund/mitiq/pull/2676. I'll give that PR a review today, and hopefully it will be merged soon. If you want to start this issue in the meantime you can start working on top of that branch.@Shivansh20128: Oh alright. This helps. But there is another problem. I am getting the below error when I tried running the `execute_with_lre` along with the parameter `observable`.  Note: this error is coming after I specified the return type of the executor as `float`.

`
ValueError: When using an executor that returns a float-like result, measurements should be added before the circuit is executed instead of with an observable.`

I tried running the below test:

```
@pytest.mark.parametrize(
    "degree, fold_multiplier, observable",
    [(2, 2, Observable(PauliString(spec="Z"))), (2, 3, Observable(PauliString(spec="Z"))), (3, 4, Observable(PauliString(spec="Z")))],
)
def test_lre_executor_with_observables(degree, fold_multiplier, observable):
    """Verify LRE executor with observable works as expected."""
    test_executor = Executor(execute)
    lre_exp_val = execute_with_lre(
        test_cirq,
        test_executor,
        degree=degree,
        fold_multiplier=fold_multiplier,
        observable=observable,
    )
    assert isinstance(lre_exp_val, float)

    assert test_executor.calls_to_executor == len(
        multivariate_layer_scaling(test_cirq, degree, fold_multiplier)
    )
```

I think the solution to this requires changing the return type of `executor` from `float`. Or we need to add measurements before the circuit is executed (as the error suggests), but I am unsure what that means and how to do that. @bdg221, I think you might have thought about this while working on the previous issue. Please suggest.@natestemen: In order to make this work, you'll need an executor function which returns either bitstrings or a density matrix. Some documentation can be found [here](https://mitiq.readthedocs.io/en/stable/guide/observables.html#using-observables-in-error-mitigation-techniques) and an example of how this is done in ZNE here:

https://github.com/unitaryfund/mitiq/blob/34eb6cedda1cc218f2f097704a9278bf15650d1f/mitiq/zne/tests/test_zne.py#L91-L94

Solution patch:
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index c2a55450..7e97fb56 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -111,6 +111,7 @@ For any part of the template that you weren't able to complete please mention th
 ### List of accepted RFCs
 This is a list of accepted request-for-comments (RFC) documents by date of creation (reverse chronological order):
 
+- [Virtual Distillation RFC](https://docs.google.com/document/d/1833IaqzpJ7DbfasgXt59AA93fJk4_EUDfBV4yNHHyic/edit?usp=sharing) by Chris de Hart, Jari Egbers, Kajetan Knopp and Lara Boloni Jan 31, 2025
 - [Probabilistic Error Amplification RFC](https://docs.google.com/document/d/1l-74EFdMA0CSFUpHjqCyQYb3ZKCmY77seB1_mOZo5Co/edit?usp=sharing) by Misty Wahl (@Misty-W) Dec 20, 2024
 - [Layerwise Richardson Extrapolation RFC](https://docs.google.com/document/d/1oFRl4wMGMtn57V0c_1egaHh0WUUAbtgW-U_QxNL9_kY/edit?usp=sharing) by Purva Thakre (@purva-thakre) and Vincent Russo (@vprusso) Apr 24, 2024
 - [Robust Shadow Estimation](https://docs.google.com/document/d/1B5FnqQDvoRYap5fGPqzcbp-RXIrUFjbBcLiWIUrLmuA) by Min Li (@Min-Li) Jun 16, 2023
diff --git a/README.md b/README.md
index 25b4e18c..ec033c0a 100644
--- a/README.md
+++ b/README.md
@@ -100,6 +100,7 @@ mitiq.qem_methods()
 | Robust Shadow Estimation   🚧           | [RSE](https://mitiq.readthedocs.io/en/stable/guide/shadows.html)| [`mitiq.qse`](https://github.com/unitaryfund/mitiq/tree/main/mitiq/shadows) | [2011.09636](https://arxiv.org/abs/2011.09636) <br> [2002.08953](https://arxiv.org/abs/2002.08953)|
 | Layerwise Richardson Extrapolation  | [LRE](https://mitiq.readthedocs.io/en/stable/guide/lre.html) |  [`mitiq.lre`](https://github.com/unitaryfund/mitiq/tree/main/mitiq/lre) | [2402.04000](https://arxiv.org/abs/2402.04000) |
 | Probabilistic Error Amplification 🚧  | Coming soon |  `mitiq.pea` | [Nature](https://www.nature.com/articles/s41586-023-06096-3) |
+| Virtual Distillation 🚧  | Coming soon |  `mitiq.vd` | [APS](https://journals.aps.org/prx/abstract/10.1103/PhysRevX.11.041036) |
 
 
 In addition, we also have a noise tailoring technique currently available with limited functionality:
diff --git a/mitiq/lre/lre.py b/mitiq/lre/lre.py
index 18892622..aadb3380 100644
--- a/mitiq/lre/lre.py
+++ b/mitiq/lre/lre.py
@@ -10,7 +10,7 @@ from typing import Any, Callable, Optional, Union
 
 import numpy as np
 
-from mitiq import QPROGRAM, Executor
+from mitiq import QPROGRAM, Executor, Observable, QuantumResult
 from mitiq.lre.inference import (
     multivariate_richardson_coefficients,
 )
@@ -22,9 +22,10 @@ from mitiq.zne.scaling import fold_gates_at_random
 
 def execute_with_lre(
     input_circuit: QPROGRAM,
-    executor: Union[Executor, Callable[[QPROGRAM], float]],
+    executor: Union[Executor, Callable[[QPROGRAM], QuantumResult]],
     degree: int,
     fold_multiplier: int,
+    observable: Optional[Observable] = None,
     folding_method: Callable[
         [QPROGRAM, float], QPROGRAM
     ] = fold_gates_at_random,  # type: ignore [has-type]
@@ -46,10 +47,14 @@ def execute_with_lre(
 
     Args:
         input_circuit: Circuit to be scaled.
-        executor: Executes a circuit and returns a `float`
+        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.
+        observable: Observable to compute the expectation value of. If
+            ``None``, the ``executor`` must return an expectation value.
+            Otherwise, the ``DensityMatrix`` or ``Bitstrings`` returned by
+            ``executor`` is used to compute the expectation of the observable.
         folding_method: Unitary folding method. Default is
             :func:`fold_gates_at_random`.
         num_chunks: Number of desired approximately equal chunks. When the
@@ -83,15 +88,16 @@ def execute_with_lre(
             + "multivariate extrapolation."
         )
 
-    lre_exp_values = executor.evaluate(noise_scaled_circuits)
+    lre_exp_values = executor.evaluate(noise_scaled_circuits, observable)
 
     return np.dot(lre_exp_values, linear_combination_coeffs)
 
 
 def mitigate_executor(
-    executor: Callable[[QPROGRAM], float],
+    executor: Callable[[QPROGRAM], QuantumResult],
     degree: int,
     fold_multiplier: int,
+    observable: Optional[Observable] = None,
     folding_method: Callable[
         [Union[Any], float], Union[Any]
     ] = fold_gates_at_random,
@@ -102,10 +108,14 @@ def mitigate_executor(
 
     Args:
         input_circuit: Circuit to be scaled.
-        executor: Executes a circuit and returns a `float`
+        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.
+        observable: Observable to compute the expectation value of. If
+            ``None``, the ``executor`` must return an expectation value.
+            Otherwise, the ``DensityMatrix`` or ``Bitstrings`` returned by
+            ``executor`` is used to compute the expectation of the observable.
         folding_method: Unitary folding method. Default is
             :func:`fold_gates_at_random`.
         num_chunks: Number of desired approximately equal chunks. When the
@@ -127,6 +137,7 @@ def mitigate_executor(
                 executor,
                 degree,
                 fold_multiplier,
+                observable,
                 folding_method,
                 num_chunks,
             )
@@ -140,6 +151,7 @@ def mitigate_executor(
                     executor,
                     degree,
                     fold_multiplier,
+                    observable,
                     folding_method,
                     num_chunks,
                 )
@@ -152,20 +164,25 @@ def mitigate_executor(
 def lre_decorator(
     degree: int,
     fold_multiplier: int,
+    observable: Optional[Observable] = None,
     folding_method: Callable[
         [QPROGRAM, float], QPROGRAM
     ] = fold_gates_at_random,
     num_chunks: Optional[int] = None,
-) -> Callable[[Callable[[QPROGRAM], float]], Callable[[QPROGRAM], float]]:
+) -> Callable[
+    [Callable[[QPROGRAM], QuantumResult]], Callable[[QPROGRAM], 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.
+        observable: Observable to compute the expectation value of. If
+            ``None``, the ``executor`` must return an expectation value.
+            Otherwise, the ``DensityMatrix`` or ``Bitstrings`` returned by
+            ``executor`` is used to compute the expectation of the observable.
         folding_method: Unitary folding method. Default is
             :func:`fold_gates_at_random`.
         num_chunks: Number of desired approximately equal chunks. When the
@@ -173,17 +190,19 @@ def lre_decorator(
             the input circuit is unchanged.
 
 
+
     Returns:
         Error-mitigated decorator.
     """
 
     def decorator(
-        executor: Callable[[QPROGRAM], float],
+        executor: Callable[[QPROGRAM], QuantumResult],
     ) -> Callable[[QPROGRAM], float]:
         return mitigate_executor(
             executor,
             degree,
             fold_multiplier,
+            observable,
             folding_method,
             num_chunks,
         )
diff --git a/requirements/requirements-qiskit.txt b/requirements/requirements-qiskit.txt
index e95c3632..7f043474 100644
--- a/requirements/requirements-qiskit.txt
+++ b/requirements/requirements-qiskit.txt
@@ -1,4 +1,4 @@
-qiskit~=1.4.1
+qiskit~=1.4.2
 qiskit-aer~=0.15.1
 qiskit-ibm-runtime~=0.36.1
 ply==3.11
\ No newline at end of file


Test patch:
diff --git a/mitiq/lre/tests/test_lre.py b/mitiq/lre/tests/test_lre.py
index 8a72b36c..701f7d74 100644
--- a/mitiq/lre/tests/test_lre.py
+++ b/mitiq/lre/tests/test_lre.py
@@ -6,10 +6,13 @@ import re
 from typing import List
 from unittest.mock import Mock
 
+import numpy as np
 import pytest
 from cirq import DensityMatrixSimulator, depolarize
 
 from mitiq import SUPPORTED_PROGRAM_TYPES, Executor, benchmarks
+from mitiq.interface import mitiq_cirq
+from mitiq.interface.mitiq_cirq import compute_density_matrix
 from mitiq.lre import (
     execute_with_lre,
     lre_decorator,
@@ -18,6 +21,7 @@ from mitiq.lre.multivariate_scaling.layerwise_folding import (
     _get_chunks,
     multivariate_layer_scaling,
 )
+from mitiq.observable import Observable, PauliString
 from mitiq.zne.scaling import fold_all, fold_global
 
 # default circuit for all unit tests
@@ -26,15 +30,33 @@ test_cirq = benchmarks.generate_rb_circuits(
     num_cliffords=2,
 )[0]
 
+test_cirq_two_qubits = benchmarks.generate_rb_circuits(
+    n_qubits=2,
+    num_cliffords=2,
+)[0]
+
+obs_x = Observable(PauliString(spec="X"))
+obs_y = Observable(PauliString(spec="Y"))
+obs_z = Observable(PauliString(spec="Z"))
+obs_zz = Observable(PauliString(spec="ZZ"))
+
 
 # default execute function for all unit tests
-def execute(circuit, noise_level=0.025):
+def execute(circuit, noise_level=0.025) -> float:
     """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
 
 
+def executor_density_matrix_typed(circuit) -> np.ndarray:
+    return compute_density_matrix(circuit, noise_level=(0,))
+
+
+def executor_density_matrix_batched(circuits) -> List[np.ndarray]:
+    return [executor_density_matrix_typed(circuit) for circuit in circuits]
+
+
 def batched_executor(circuits) -> List[float]:
     return [execute(circuit) for circuit in circuits]
 
@@ -45,7 +67,7 @@ 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."""
+    """Verify LRE executors works as expected."""
     lre_exp_val = execute_with_lre(
         test_cirq,
         execute,
@@ -55,6 +77,33 @@ def test_lre_exp_value(degree, fold_multiplier):
     assert abs(lre_exp_val - ideal_val) <= abs(noisy_val - ideal_val)
 
 
+@pytest.mark.parametrize(
+    "circuit, degree, fold_multiplier, observable",
+    [
+        (test_cirq, 2, 3, obs_z),
+        (test_cirq_two_qubits, 1, 2, obs_zz),
+    ],
+)
+def test_lre_exp_value_with_observable(
+    circuit, degree, fold_multiplier, observable
+):
+    """Verify LRE can be used with observables."""
+    test_executor = Executor(mitiq_cirq.compute_density_matrix)
+    lre_exp_val = execute_with_lre(
+        circuit,
+        test_executor,
+        degree=degree,
+        fold_multiplier=fold_multiplier,
+        observable=observable,
+    )
+
+    assert isinstance(lre_exp_val, float)
+
+    assert test_executor.calls_to_executor == len(
+        multivariate_layer_scaling(circuit, degree, fold_multiplier)
+    )
+
+
 @pytest.mark.parametrize(
     "degree, fold_multiplier",
     [(2, 2), (2, 3), (3, 4)],
@@ -97,6 +146,33 @@ def test_lre_batched_executor(degree, fold_multiplier):
     )
 
 
+@pytest.mark.parametrize(
+    "degree, fold_multiplier, observable",
+    [(2, 2, obs_x), (2, 3, obs_y), (3, 4, obs_z)],
+)
+def test_lre_batched_executor_with_observable(
+    degree, fold_multiplier, observable
+):
+    """Verify LRE batch executor with observable works as expected."""
+    test_batched_executor = Executor(
+        executor_density_matrix_batched, max_batch_size=200
+    )
+    lre_exp_val_batched = execute_with_lre(
+        test_cirq,
+        test_batched_executor,
+        degree=degree,
+        fold_multiplier=fold_multiplier,
+        observable=observable,
+    )
+    assert isinstance(lre_exp_val_batched, float)
+
+    assert test_batched_executor.calls_to_executor == 1
+    assert (
+        test_batched_executor.executed_circuits
+        == multivariate_layer_scaling(test_cirq, degree, fold_multiplier)
+    )
+
+
 @pytest.mark.parametrize("circuit_type", SUPPORTED_PROGRAM_TYPES.keys())
 def test_lre_all_qprogram(circuit_type):
     """Verify LRE works with all supported frontends."""