Modularize DDD functionality (#2610)

* separated generate circuits and ddd value into separate functions

* Add docstrings and test

* Cast a float for mypy

* Use combine_results to match ZNE and PEC

* increase num_trials testing

---------

Co-authored-by: nate stemen <nate@unitary.fund>

mitiq/ddd/__init__.py

@@ -14,4 +14,4 @@
     insert_ddd_sequences,
 )
 
-from mitiq.ddd.ddd import execute_with_ddd, mitigate_executor, ddd_decorator
+from mitiq.ddd.ddd import execute_with_ddd, mitigate_executor, ddd_decorator, generate_circuits_with_ddd, combine_results

mitiq/ddd/ddd.py

@@ -59,13 +59,11 @@ def execute_with_ddd(
     if not isinstance(executor, Executor):
         executor = Executor(executor)
 
-    rule_partial: Callable[[int], QPROGRAM]
-    rule_partial = partial(rule, **rule_args)
-
     # Insert DDD sequences in (a copy of) the input circuit
-    circuits_with_ddd = [
-        insert_ddd_sequences(circuit, rule_partial) for _ in range(num_trials)
-    ]
+    circuits_with_ddd = generate_circuits_with_ddd(
+        circuit, rule, rule_args, num_trials
+    )
+
     results = executor.evaluate(
         circuits_with_ddd,
         observable,
@@ -74,7 +72,8 @@ def execute_with_ddd(
 
     assert len(results) == num_trials
 
-    ddd_value = np.sum(results) / num_trials
+    ddd_value = combine_results(results)
+
     if not full_output:
         return ddd_value
 
@@ -86,6 +85,50 @@ def execute_with_ddd(
     return ddd_value, ddd_data
 
 
+def combine_results(results: list[float]) -> float:
+    """Averages over the DDD results to get the expectation value from using
+    DDD.
+
+    Args:
+        results: Results as obtained from running circuits.
+
+    Returns:
+        The expectation value estimated with DDD.
+    """
+    return float(np.average(results))
+
+
+def generate_circuits_with_ddd(
+    circuit: QPROGRAM,
+    rule: Callable[[int], QPROGRAM],
+    rule_args: Dict[str, Any] = {},
+    num_trials: int = 1,
+) -> list[QPROGRAM]:
+    """Generates a list of circuits with DDD sequences inserted.
+
+    Args:
+        circuit: The quantum circuit to be modified with DD.
+        rule: A function that takes as main argument a slack length (i.e. the
+            number of idle moments) of a slack window (i.e. a single-qubit idle
+            window in a circuit) and returns the DDD sequence of gates to be
+            applied in that window.
+        rule_args: An optional dictionary of keyword arguments for ``rule``.
+        num_trials: The number of circuits to generate with DDD insertions.
+
+    Returns:
+        A list of circuits with DDD inserted.
+    """
+    rule_partial: Callable[[int], QPROGRAM]
+    rule_partial = partial(rule, **rule_args)
+
+    # Insert DDD sequences in (a copy of) the input circuit
+    circuits_with_ddd = [
+        insert_ddd_sequences(circuit, rule_partial) for _ in range(num_trials)
+    ]
+
+    return circuits_with_ddd
+
+
 def mitigate_executor(
     executor: Callable[[QPROGRAM], QuantumResult],
     observable: Optional[Observable] = None,

mitiq/ddd/tests/test_ddd.py

@@ -12,7 +12,12 @@
 from pytest import mark
 
 from mitiq import QPROGRAM, SUPPORTED_PROGRAM_TYPES, Executor
-from mitiq.ddd import ddd_decorator, execute_with_ddd, mitigate_executor
+from mitiq.ddd import (
+    ddd_decorator,
+    execute_with_ddd,
+    generate_circuits_with_ddd,
+    mitigate_executor,
+)
 from mitiq.ddd.rules import xx, xyxy, yy
 from mitiq.interface import convert_from_mitiq, convert_to_mitiq
 from mitiq.interface.mitiq_cirq import compute_density_matrix
@@ -226,3 +231,14 @@ def exec_xx_small_spacing(circuit):
     # What is important to test is getting different results.
     assert not np.isclose(unmitigated, mitigated_small_spacing)
     assert not np.isclose(mitigated_large_spacing, mitigated_small_spacing)
+
+
+@mark.parametrize("num_trials", [1, 10, 20, 30])
+def test_num_trials_generates_circuits(num_trials: int):
+    """Test that the number of generated circuits follows num_trials."""
+
+    circuits = generate_circuits_with_ddd(
+        circuit_cirq_a, rule=xx, num_trials=num_trials
+    )
+
+    assert num_trials == len(circuits)