Ensure LRE compatibility with all supported frontends (#2547)

* add non-cirq support in LRE

* clean up some utility conversions

* use  directly for parametrization

* replace magic number with computed

* split test into two

mitiq/calibration/settings.py

@@ -95,7 +95,7 @@ def converted_circuit(
         """
         circuit = self.circuit.copy()
         circuit.append(cirq.measure(circuit.all_qubits()))
-        return convert_from_mitiq(circuit, circuit_type.value)
+        return convert_from_mitiq(circuit, circuit_type.name)
 
     @property
     def num_qubits(self) -> int:

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
 
@@ -150,6 +161,7 @@ def convert_from_mitiq(
         circuit: Mitiq circuit to convert.
         conversion_type: String specifier for the converted circuit type.
     """
+    conversion_type = conversion_type.lower()
     conversion_function: Callable[[cirq.Circuit], QPROGRAM]
     if conversion_type == "qiskit":
         from mitiq.interface.mitiq_qiskit.conversions import to_qiskit
@@ -199,13 +211,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 +265,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

@@ -6,11 +6,14 @@
 """Unit tests for scaling noise by unitary folding of layers in the input
 circuit to allow for multivariate extrapolation."""
 
+import math
 from copy import deepcopy
 
 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,14 +33,35 @@
 test_circuit1_with_measurements.append(ops.measure_each(*qreg1))
 
 
-def test_multivariate_layerwise_scaling():
-    """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
+def test_multivariate_layerwise_scaling_num_circuits():
+    """Ensure the correct number of circuits are generated."""
+    degree, fold_multiplier, num_chunks = 2, 2, 3
+    scaled_circuits = multivariate_layer_scaling(
+        test_circuit1, degree, fold_multiplier, num_chunks
     )
 
-    assert len(multiple_scaled_circuits) == 10
+    depth = len(test_circuit1)
+    # number of circuit is `degree` + `depth` choose `degree`
+    assert len(scaled_circuits) == math.comb(degree + depth, degree)
+
+
+@pytest.mark.parametrize("circuit_type", SUPPORTED_PROGRAM_TYPES)
+def test_multivariate_layerwise_scaling_types(circuit_type):
+    """Ensure layer scaling returns circuits of the correct type."""
+    circuit = convert_from_mitiq(test_circuit1, circuit_type.name)
+    degree, fold_multiplier, num_chunks = 2, 2, 3
+    scaled_circuits = multivariate_layer_scaling(
+        circuit, degree, fold_multiplier, num_chunks
+    )
+
+    assert all(
+        isinstance(circuit, circuit_type.value) for circuit in scaled_circuits
+    )
+
+
+def test_multivariate_layerwise_scaling_cirq():
+    """Ensure the folding pattern is correct for a nontrivial case."""
+    scaled_circuits = multivariate_layer_scaling(test_circuit1, 2, 2, 3)
     folding_pattern = [
         (1, 1, 1),
         (5, 1, 1),
@@ -50,16 +74,15 @@ def test_multivariate_layerwise_scaling():
         (1, 5, 5),
         (1, 1, 9),
     ]
-
-    for i, scale_factor_vector in enumerate(folding_pattern):
-        scale_layer1, scale_layer2, scale_layer3 = scale_factor_vector
+    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 == multiple_scaled_circuits[i]
+        assert expected_circuit == circuit
 
 
 @pytest.mark.parametrize(

mitiq/lre/tests/test_lre.py

@@ -1,11 +1,14 @@
 """Unit tests for the LRE extrapolation methods."""
 
+import math
+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 +43,28 @@ 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)
+    depth = 3  # not all circuit types have a simple way to compute depth
+
+    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 == math.comb(degree + depth, degree)
+
+
 @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

@@ -20,7 +20,6 @@
 from typing import (
     Any,
     Dict,
-    Iterable,
     List,
     Optional,
     Sequence,
@@ -37,25 +36,17 @@
 
 class EnhancedEnumMeta(EnumMeta):
     def __str__(cls) -> str:
-        return ", ".join([member.value for member in cast(Type[Enum], cls)])
+        return ", ".join(
+            [member.name.lower() for member in cast(Type[Enum], cls)]
+        )
 
 
 class EnhancedEnum(Enum, metaclass=EnhancedEnumMeta):
     # This is for backwards compatibility with the old representation
     # of SUPPORTED_PROGRAM_TYPES, which was a dictionary
     @classmethod
-    def keys(cls) -> Iterable[str]:
-        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"
+    def keys(cls) -> list[str]:
+        return [member.name.lower() for member in cls]
 
 
 try:
@@ -90,6 +81,16 @@ class SUPPORTED_PROGRAM_TYPES(EnhancedEnum):
 ]
 
 
+# Supported quantum programs.
+class SUPPORTED_PROGRAM_TYPES(EnhancedEnum):
+    BRAKET = _BKCircuit
+    CIRQ = _Circuit
+    PENNYLANE = _QuantumTape
+    PYQUIL = _Program
+    QIBO = _QiboCircuit
+    QISKIT = _QuantumCircuit
+
+
 # Define MeasurementResult, a result obtained by measuring qubits on a quantum
 # computer.
 Bitstring = Union[str, List[int]]