From c1be865c71c53884a2a29f92114ecc2b4374131a Mon Sep 17 00:00:00 2001
From: Purva Thakre <66048318+purva-thakre@users.noreply.github.com>
Date: Thu, 13 Feb 2025 18:07:03 -0600
Subject: [PATCH] Default scaling method in modularized ZNE function (#2666)

* add default option and copy-paste unit test

* mypy failure

* scaled_circuits default option assert

* nate's suggestion: change to_frontend

* separate test
---
 mitiq/zne/tests/test_zne.py | 25 +++++++++++++++++++++----
 mitiq/zne/zne.py            |  2 +-
 2 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/mitiq/zne/tests/test_zne.py b/mitiq/zne/tests/test_zne.py
index 80c6c3efc..114741daf 100644
--- a/mitiq/zne/tests/test_zne.py
+++ b/mitiq/zne/tests/test_zne.py
@@ -584,11 +584,11 @@ def qs_noisy_simulation(
     "extrapolation_factory", [RichardsonFactory, LinearFactory]
 )
 @pytest.mark.parametrize(
-    "to_frontend",
+    "conversion_func",
     [None, to_qiskit, to_braket, to_pennylane, to_pyquil, to_qibo],
 )
 def test_two_stage_zne(
-    noise_scaling_method, extrapolation_factory, to_frontend
+    noise_scaling_method, extrapolation_factory, conversion_func
 ):
     qreg = cirq.LineQubit.range(2)
     cirq_circuit = cirq.Circuit(
@@ -597,8 +597,8 @@ def test_two_stage_zne(
         cirq.CNOT(*qreg),
         cirq.H.on_each(qreg),
     )
-    if to_frontend is not None:
-        frontend_circuit = to_frontend(cirq_circuit)
+    if conversion_func is not None:
+        frontend_circuit = conversion_func(cirq_circuit)
     else:
         frontend_circuit = cirq_circuit
 
@@ -630,3 +630,20 @@ def executor(circuit):
         scale_noise=noise_scaling_method,
     )
     assert np.isclose(zne_res, two_stage_zne_res)
+
+
+def test_default_scaling_option_two_stage_zne():
+    qreg = cirq.LineQubit.range(2)
+    cirq_circuit = cirq.Circuit(
+        cirq.H.on_each(qreg),
+        cirq.CNOT(*qreg),
+        cirq.CNOT(*qreg),
+        cirq.H.on_each(qreg),
+    )
+
+    scale_factors = [3, 5, 6, 8]
+
+    circs_default_scaling_method = scaled_circuits(cirq_circuit, scale_factors)
+
+    for i in range(len(scale_factors)):
+        assert len(circs_default_scaling_method[i]) > len(cirq_circuit)
diff --git a/mitiq/zne/zne.py b/mitiq/zne/zne.py
index b23013748..cc4d56589 100644
--- a/mitiq/zne/zne.py
+++ b/mitiq/zne/zne.py
@@ -16,7 +16,7 @@
 def scaled_circuits(
     circuit: QPROGRAM,
     scale_factors: list[float],
-    scale_method: Callable[[QPROGRAM, float], QPROGRAM],
+    scale_method: Callable[[QPROGRAM, float], QPROGRAM] = fold_gates_at_random,  # type:ignore [has-type]
 ) -> list[QPROGRAM]:
     """Given a circuit, scale_factors and a scale_method, outputs a list
        of circuits that will be used in ZNE.