remove failing test; simplify layerwise ZNE tests (#2545)

* move  param test to module; add  test

* remove flaky test

* use numbered indices

mitiq/lre/tests/test_lre.py

@@ -101,20 +101,6 @@ def test_lre_executor_with_chunking():
     assert abs(lre_exp_val - ideal_val) <= abs(noisy_val - ideal_val)
 
 
-@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=num_chunks
-    )
-    assert abs(lre_exp_val - ideal_val) >= abs(noisy_val - ideal_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

mitiq/zne/scaling/tests/test_layer_scaling.py

@@ -5,9 +5,16 @@
 
 """Unit tests for scaling by layer."""
 
+import random
+from itertools import product
+from unittest.mock import patch
+
+import pytest
 from cirq import Circuit, LineQubit, ops
 
-from mitiq.zne.scaling import layer_folding
+from mitiq import SUPPORTED_PROGRAM_TYPES
+from mitiq.interface import convert_from_mitiq, convert_to_mitiq
+from mitiq.zne.scaling import get_layer_folding, layer_folding
 
 
 def test_layer_folding_with_measurements():
@@ -53,24 +60,59 @@ def test_layer_folding():
     # Iterate over every possible combination of layerwise folds for a maximum
     # number of 5-folds.
     total_folds = 5
-    for i1 in range(total_folds):
-        for i2 in range(total_folds):
-            for i3 in range(total_folds):
-                layers_to_fold = [i1, i2, i3]
-
-                folded_circuit = layer_folding(circ, layers_to_fold)
-
-                # For a given layer, the number of copies on a layer will be
-                # 2n + 1 where "n" is the number of folds to perform.
-                qreg = LineQubit.range(3)
-                correct = Circuit(
-                    # Layer-1
-                    [ops.H.on_each(*qreg)] * (2 * (layers_to_fold[0]) + 1),
-                    # Layer-2
-                    [ops.CNOT.on(qreg[0], qreg[1])]
-                    * (2 * (layers_to_fold[1]) + 1),
-                    [ops.X.on(qreg[2])] * (2 * (layers_to_fold[1]) + 1),
-                    # Layer-3
-                    [ops.TOFFOLI.on(*qreg)] * (2 * (layers_to_fold[2]) + 1),
-                )
-                assert folded_circuit == correct
+    for i1, i2, i3 in product(range(total_folds), repeat=3):
+        folded_circuit = layer_folding(circ, [i1, i2, i3])
+
+        # For a given layer, the number of copies on a layer will be
+        # 2n + 1 where "n" is the number of folds to perform.
+        a, b, c = LineQubit.range(3)
+        correct = Circuit(
+            # Layer-1
+            [ops.H.on_each(*(a, b, c))] * (2 * i1 + 1),
+            # Layer-2
+            [ops.CNOT.on(a, b)] * (2 * i2 + 1),
+            [ops.X.on(c)] * (2 * i2 + 1),
+            # Layer-3
+            [ops.TOFFOLI.on(*(a, b, c))] * (2 * i3 + 1),
+        )
+        assert folded_circuit == correct
+
+
+@pytest.mark.parametrize("circuit_type", SUPPORTED_PROGRAM_TYPES.keys())
+def test_layer_folding_all_qprograms(circuit_type):
+    """This test only ensures proper depth of layer-folded non-cirq circuits
+    as the mitiq conversion functions alter structure/gate composition."""
+    qreg = LineQubit.range(3)
+    circuit = Circuit(
+        [ops.H.on_each(*qreg)],
+        [ops.CNOT.on(qreg[0], qreg[1])],
+        [ops.X.on(qreg[2])],
+        [ops.TOFFOLI.on(*qreg)],
+    )
+    num_layers = len(circuit)
+    circuit = convert_from_mitiq(circuit, circuit_type)
+    layers_to_fold = random.choices(range(5), k=num_layers)
+    folded_circuit = layer_folding(circuit, layers_to_fold)
+    folded_mitiq_circuit = convert_to_mitiq(folded_circuit)[0]
+    num_ideal_layers = sum(2 * n + 1 for n in layers_to_fold)
+    if circuit_type == "pyquil":
+        # this block is needed for pyquil because of some quirks that pop up
+        # when converting to and from pyquil that does not make exact equality.
+        assert len(folded_mitiq_circuit) >= num_ideal_layers
+    else:
+        assert len(folded_mitiq_circuit) == num_ideal_layers
+
+
+@patch("mitiq.zne.scaling.layer_scaling.layer_folding")
+def test_get_layer_folding(mock_layer_folding):
+    a, b = LineQubit.range(2)
+    circuit = Circuit(ops.X(a), ops.CNOT(a, b), ops.Y(b))
+    layer_index = 1
+    scale_factor = 3
+
+    folding_func = get_layer_folding(layer_index)
+    folding_func(circuit, scale_factor)
+
+    mock_layer_folding.assert_called_once_with(
+        circuit, layers_to_fold=[0, 1, 0]
+    )

mitiq/zne/tests/test_zne.py

@@ -6,7 +6,9 @@
 """Unit tests for zero-noise extrapolation."""
 
 import functools
+import random
 from typing import List
+from unittest.mock import Mock
 
 import cirq
 import numpy as np
@@ -493,32 +495,34 @@ def test_execute_with_zne_with_supported_circuits(circuit_type):
     assert abs(unmitigated - expected) > abs(zne_value - expected)
 
 
-@pytest.mark.parametrize("circuit_type", SUPPORTED_PROGRAM_TYPES.keys())
-def test_layerwise_execute_with_zne_with_supported_circuits(circuit_type):
+def test_layerwise_folding_with_zne():
     # Define a circuit equivalent to the identity
     qreg = cirq.LineQubit.range(2)
-    cirq_circuit = cirq.Circuit(
+    circuit = cirq.Circuit(
         cirq.H.on_each(qreg),
         cirq.CNOT(*qreg),
         cirq.CNOT(*qreg),
         cirq.H.on_each(qreg),
     )
-    # Convert to one of the supported program types
-    circuit = convert_from_mitiq(cirq_circuit, circuit_type)
-    expected = generic_executor(circuit, noise_level=0.0)
-    unmitigated = generic_executor(circuit)
-    # Use odd scale factors for deterministic results
-    fac = RichardsonFactory([1, 3, 5])
-    # Layerwise-fold
+    circuit_depth = len(circuit)
+    mock_executor = Mock(side_effect=lambda _: random.random())
     layer_to_fold = 0
     fold_layer_func = get_layer_folding(layer_to_fold)
+    scale_factors = [1, 3, 5]
+    factory = RichardsonFactory(scale_factors)
 
-    zne_value = execute_with_zne(
-        circuit, generic_executor, factory=fac, scale_noise=fold_layer_func
+    execute_with_zne(
+        circuit, mock_executor, factory=factory, scale_noise=fold_layer_func
     )
-
-    # Test zero noise limit is better than unmitigated expectation value
-    assert abs(unmitigated - expected) > abs(zne_value - expected)
+    assert mock_executor.call_count == len(scale_factors)
+    circuit_depths = [
+        len(args[0]) for args, kwargs in mock_executor.call_args_list
+    ]
+    assert circuit_depths == [
+        circuit_depth,
+        circuit_depth + 2,
+        circuit_depth + 4,
+    ]
 
 
 def test_execute_with_zne_transpiled_qiskit_circuit():