Ensure LRE compatibility with all supported frontends

mitiq/interface/conversions.py

@@ -6,7 +6,18 @@
 """Functions for converting to/from Mitiq's internal circuit representation."""
 
 from functools import wraps
-from typing import Any, Callable, Dict, Iterable, Optional, Tuple, cast
+from typing import (
+    Any,
+    Callable,
+    Collection,
+    Concatenate,
+    Dict,
+    Optional,
+    ParamSpec,
+    Tuple,
+    TypeVar,
+    cast,
+)
 
 import cirq
 
@@ -199,13 +210,21 @@ def conversion_function(circ: cirq.Circuit) -> cirq.Circuit:
     return converted_circuit
 
 
+P = ParamSpec("P")
+R = TypeVar("R")
+
+
 def accept_any_qprogram_as_input(
-    accept_cirq_circuit_function: Callable[[cirq.Circuit], Any],
-) -> Callable[[QPROGRAM], Any]:
+    accept_cirq_circuit_function: Callable[Concatenate[cirq.Circuit, P], R],
+) -> Callable[Concatenate[QPROGRAM, P], R]:
+    """Converts functions which take as input cirq.Circuit object (and return
+    anything), to function which can accept any QPROGRAM.
+    """
+
     @wraps(accept_cirq_circuit_function)
     def accept_any_qprogram_function(
-        circuit: QPROGRAM, *args: Any, **kwargs: Any
-    ) -> Any:
+        circuit: QPROGRAM, *args: P.args, **kwargs: P.kwargs
+    ) -> R:
         cirq_circuit, _ = convert_to_mitiq(circuit)
         return accept_cirq_circuit_function(cirq_circuit, *args, **kwargs)
 
@@ -245,15 +264,19 @@ def qprogram_modifier(
 
 
 def atomic_one_to_many_converter(
-    cirq_circuit_modifier: Callable[..., Iterable[cirq.Circuit]],
-) -> Callable[..., Iterable[QPROGRAM]]:
+    cirq_circuit_modifier: Callable[..., Collection[cirq.Circuit]],
+) -> Callable[..., Collection[QPROGRAM]]:
+    """Convert function which returns multiple cirq.Circuits into a function
+    which returns multiple QPROGRAM instances.
+    """
+
     @wraps(cirq_circuit_modifier)
     def qprogram_modifier(
         circuit: QPROGRAM, *args: Any, **kwargs: Any
-    ) -> Iterable[QPROGRAM]:
+    ) -> Collection[QPROGRAM]:
         mitiq_circuit, input_circuit_type = convert_to_mitiq(circuit)
 
-        modified_circuits: Iterable[cirq.Circuit] = cirq_circuit_modifier(
+        modified_circuits = cirq_circuit_modifier(
             mitiq_circuit, *args, **kwargs
         )
 

mitiq/lre/inference/multivariate_richardson.py

@@ -15,6 +15,7 @@
 from cirq import Circuit
 from numpy.typing import NDArray
 
+from mitiq.interface import accept_any_qprogram_as_input
 from mitiq.lre.multivariate_scaling.layerwise_folding import (
     _get_scale_factor_vectors,
 )
@@ -120,6 +121,7 @@ def sample_matrix(
     return sample_matrix
 
 
+@accept_any_qprogram_as_input
 def multivariate_richardson_coefficients(
     input_circuit: Circuit,
     degree: int,

mitiq/lre/lre.py

@@ -9,7 +9,6 @@
 from typing import Any, Callable, Optional, Union
 
 import numpy as np
-from cirq import Circuit
 
 from mitiq import QPROGRAM
 from mitiq.lre.inference import (
@@ -22,8 +21,8 @@
 
 
 def execute_with_lre(
-    input_circuit: Circuit,
-    executor: Callable[[Circuit], float],
+    input_circuit: QPROGRAM,
+    executor: Callable[[QPROGRAM], float],
     degree: int,
     fold_multiplier: int,
     folding_method: Callable[
@@ -90,14 +89,14 @@ def execute_with_lre(
 
 
 def mitigate_executor(
-    executor: Callable[[Circuit], float],
+    executor: Callable[[QPROGRAM], 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]:
+) -> Callable[[QPROGRAM], float]:
     """Returns a modified version of the input `executor` which is
     error-mitigated with layerwise richardson extrapolation (LRE).
 
@@ -119,7 +118,7 @@ def mitigate_executor(
     """
 
     @wraps(executor)
-    def new_executor(input_circuit: Circuit) -> float:
+    def new_executor(input_circuit: QPROGRAM) -> float:
         return execute_with_lre(
             input_circuit,
             executor,
@@ -135,9 +134,11 @@ def new_executor(input_circuit: Circuit) -> float:
 def lre_decorator(
     degree: int,
     fold_multiplier: int,
-    folding_method: Callable[[Circuit, float], Circuit] = fold_gates_at_random,
+    folding_method: Callable[
+        [QPROGRAM, float], QPROGRAM
+    ] = fold_gates_at_random,
     num_chunks: Optional[int] = None,
-) -> Callable[[Callable[[Circuit], float]], Callable[[Circuit], float]]:
+) -> Callable[[Callable[[QPROGRAM], float]], Callable[[QPROGRAM], float]]:
     """Decorator which adds an error-mitigation layer based on
     layerwise richardson extrapolation (LRE).
 
@@ -159,8 +160,8 @@ def lre_decorator(
     """
 
     def decorator(
-        executor: Callable[[Circuit], float],
-    ) -> Callable[[Circuit], float]:
+        executor: Callable[[QPROGRAM], float],
+    ) -> Callable[[QPROGRAM], float]:
         return mitigate_executor(
             executor,
             degree,

mitiq/lre/multivariate_scaling/layerwise_folding.py

@@ -15,6 +15,7 @@
 from cirq import Circuit
 
 from mitiq import QPROGRAM
+from mitiq.interface import accept_qprogram_and_validate
 from mitiq.utils import _append_measurements, _pop_measurements
 from mitiq.zne.scaling import fold_gates_at_random
 from mitiq.zne.scaling.folding import _check_foldable
@@ -134,7 +135,7 @@ def _get_scale_factor_vectors(
     ]
 
 
-def multivariate_layer_scaling(
+def _multivariate_layer_scaling(
     input_circuit: Circuit,
     degree: int,
     fold_multiplier: int,
@@ -208,3 +209,8 @@ def multivariate_layer_scaling(
         multiple_folded_circuits.append(folded_circuit)
 
     return multiple_folded_circuits
+
+
+multivariate_layer_scaling = accept_qprogram_and_validate(
+    _multivariate_layer_scaling, one_to_many=True
+)

mitiq/lre/tests/test_layerwise_folding.py

@@ -11,6 +11,8 @@
 import pytest
 from cirq import Circuit, LineQubit, ops
 
+from mitiq import SUPPORTED_PROGRAM_TYPES
+from mitiq.interface import convert_from_mitiq
 from mitiq.lre.multivariate_scaling.layerwise_folding import (
     _get_chunks,
     _get_num_layers_without_measurements,
@@ -30,36 +32,41 @@
 test_circuit1_with_measurements.append(ops.measure_each(*qreg1))
 
 
-def test_multivariate_layerwise_scaling():
+@pytest.mark.parametrize("circuit_type", SUPPORTED_PROGRAM_TYPES.keys())
+def test_multivariate_layerwise_scaling(circuit_type):
     """Checks if multiple scaled circuits are returned to fit the required
     folding pattern for multivariate extrapolation."""
-    multiple_scaled_circuits = multivariate_layer_scaling(
-        test_circuit1, 2, 2, 3
-    )
+    circuit = convert_from_mitiq(test_circuit1, circuit_type)
+    scaled_circuits = multivariate_layer_scaling(circuit, 2, 2, 3)
 
-    assert len(multiple_scaled_circuits) == 10
-    folding_pattern = [
-        (1, 1, 1),
-        (5, 1, 1),
-        (1, 5, 1),
-        (1, 1, 5),
-        (9, 1, 1),
-        (5, 5, 1),
-        (5, 1, 5),
-        (1, 9, 1),
-        (1, 5, 5),
-        (1, 1, 9),
-    ]
+    assert len(scaled_circuits) == 10
+    assert all(
+        isinstance(circuit, SUPPORTED_PROGRAM_TYPES._python_type(circuit_type))
+        for circuit in scaled_circuits
+    )
 
-    for i, scale_factor_vector in enumerate(folding_pattern):
-        scale_layer1, scale_layer2, scale_layer3 = scale_factor_vector
-        expected_circuit = Circuit(
-            [ops.H.on_each(*qreg1)] * scale_layer1,
-            [ops.CNOT.on(qreg1[0], qreg1[1]), ops.X.on(qreg1[2])]
-            * scale_layer2,
-            [ops.TOFFOLI.on(*qreg1)] * scale_layer3,
-        )
-        assert expected_circuit == multiple_scaled_circuits[i]
+    if circuit_type == "cirq":
+        folding_pattern = [
+            (1, 1, 1),
+            (5, 1, 1),
+            (1, 5, 1),
+            (1, 1, 5),
+            (9, 1, 1),
+            (5, 5, 1),
+            (5, 1, 5),
+            (1, 9, 1),
+            (1, 5, 5),
+            (1, 1, 9),
+        ]
+        for scale_factors, circuit in zip(folding_pattern, scaled_circuits):
+            scale_layer1, scale_layer2, scale_layer3 = scale_factors
+            expected_circuit = Circuit(
+                [ops.H.on_each(*qreg1)] * scale_layer1,
+                [ops.CNOT.on(qreg1[0], qreg1[1]), ops.X.on(qreg1[2])]
+                * scale_layer2,
+                [ops.TOFFOLI.on(*qreg1)] * scale_layer3,
+            )
+            assert expected_circuit == circuit
 
 
 @pytest.mark.parametrize(

mitiq/lre/tests/test_lre.py

@@ -1,11 +1,13 @@
 """Unit tests for the LRE extrapolation methods."""
 
+import random
 import re
+from unittest.mock import Mock
 
 import pytest
 from cirq import DensityMatrixSimulator, depolarize
 
-from mitiq import benchmarks
+from mitiq import SUPPORTED_PROGRAM_TYPES, benchmarks
 from mitiq.lre import execute_with_lre, lre_decorator, mitigate_executor
 from mitiq.zne.scaling import fold_all, fold_global
 
@@ -40,8 +42,27 @@ def test_lre_exp_value(degree, fold_multiplier):
     assert abs(lre_exp_val - ideal_val) <= abs(noisy_val - ideal_val)
 
 
+@pytest.mark.parametrize("circuit_type", SUPPORTED_PROGRAM_TYPES.keys())
+def test_lre_all_qprogram(circuit_type):
+    """Verify LRE works with all supported frontends."""
+    degree, fold_multiplier = 2, 3
+    circuit = benchmarks.generate_ghz_circuit(3, circuit_type)
+
+    mock_executor = Mock(side_effect=lambda _: random.random())
+
+    lre_exp_val = execute_with_lre(
+        circuit,
+        mock_executor,
+        degree=degree,
+        fold_multiplier=fold_multiplier,
+    )
+
+    assert isinstance(lre_exp_val, float)
+    assert mock_executor.call_count == 10
+
+
 @pytest.mark.parametrize("degree, fold_multiplier", [(2, 2), (2, 3), (3, 4)])
-def test_lre_exp_value_decorator(degree, fold_multiplier):
+def test_lre_mitigate_executor(degree, fold_multiplier):
     """Verify LRE mitigated executor work as expected."""
     mitigated_executor = mitigate_executor(
         execute, degree=2, fold_multiplier=2

mitiq/typing.py

@@ -48,16 +48,6 @@
         return [member.value for member in cls]
 
 
-# Supported quantum programs.
-class SUPPORTED_PROGRAM_TYPES(EnhancedEnum):
-    BRAKET = "braket"
-    CIRQ = "cirq"
-    PENNYLANE = "pennylane"
-    PYQUIL = "pyquil"
-    QIBO = "qibo"
-    QISKIT = "qiskit"
-
-
 try:
     from pyquil import Program as _Program
 except ImportError:  # pragma: no cover
@@ -90,6 +80,33 @@
 ]
 
 
+# Supported quantum programs.
+class SUPPORTED_PROGRAM_TYPES(EnhancedEnum):
+    BRAKET = "braket"
+    CIRQ = "cirq"
+    PENNYLANE = "pennylane"
+    PYQUIL = "pyquil"
+    QIBO = "qibo"
+    QISKIT = "qiskit"
+
+    @classmethod
+    def _python_type(cls, identifier: str) -> QPROGRAM:
+        if identifier == "braket":
+            return _BKCircuit
+        elif identifier == "cirq":
+            return _Circuit
+        elif identifier == "pennylane":
+            return _QuantumTape
+        elif identifier == "pyquil":
+            return _Program
+        elif identifier == "qibo":
+            return _QiboCircuit
+        elif identifier == "qiskit":
+            return _QuantumCircuit
+        else:
+            raise ValueError(f"Invalid identifier: {identifier}")
+
+
 # Define MeasurementResult, a result obtained by measuring qubits on a quantum
 # computer.
 Bitstring = Union[str, List[int]]