Simplify StatePreparationAliasSampling tests by relying on Cirq simulators to simulate operations allocating ancillas (#6204)

* Simplify StatePreparationAliasSampling tests by relying on Cirq simulators capability to simulate operations allocating ancilla qubits

* Update unary iteration implementation to proactively allocate all required ancillas

* Fix pylint errors

Author: tanujkhattar

cirq-ft/cirq_ft/algos/selected_majorana_fermion_test.py

@@ -22,34 +22,33 @@
 @pytest.mark.parametrize("selection_bitsize, target_bitsize", [(2, 4), (3, 8), (4, 9)])
 @pytest.mark.parametrize("target_gate", [cirq.X, cirq.Y])
 def test_selected_majorana_fermion_gate(selection_bitsize, target_bitsize, target_gate):
-    greedy_mm = cirq_ft.GreedyQubitManager(prefix="_a", maximize_reuse=True)
     gate = cirq_ft.SelectedMajoranaFermionGate(
         cirq_ft.SelectionRegisters(
             [cirq_ft.SelectionRegister('selection', selection_bitsize, target_bitsize)]
         ),
         target_gate=target_gate,
     )
-    g = cirq_ft.testing.GateHelper(gate, context=cirq.DecompositionContext(greedy_mm))
+    g = cirq_ft.testing.GateHelper(gate)
     assert len(g.all_qubits) <= gate.registers.total_bits() + selection_bitsize + 1
 
     sim = cirq.Simulator(dtype=np.complex128)
     for n in range(target_bitsize):
         # Initial qubit values
-        qubit_vals = {q: 0 for q in g.all_qubits}
+        qubit_vals = {q: 0 for q in g.operation.qubits}
         # All controls 'on' to activate circuit
         qubit_vals.update({c: 1 for c in g.quregs['control']})
         # Set selection according to `n`
         qubit_vals.update(zip(g.quregs['selection'], iter_bits(n, selection_bitsize)))
 
-        initial_state = [qubit_vals[x] for x in g.all_qubits]
+        initial_state = [qubit_vals[x] for x in g.operation.qubits]
 
         result = sim.simulate(
-            g.decomposed_circuit, initial_state=initial_state, qubit_order=g.all_qubits
+            g.circuit, initial_state=initial_state, qubit_order=g.operation.qubits
         )
 
         final_target_state = cirq.sub_state_vector(
             result.final_state_vector,
-            keep_indices=[g.all_qubits.index(q) for q in g.quregs['target']],
+            keep_indices=[g.operation.qubits.index(q) for q in g.quregs['target']],
         )
 
         expected_target_state = cirq.Circuit(

cirq-ft/cirq_ft/algos/state_preparation_test.py

@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import itertools
-
 import cirq
 import cirq_ft
 import numpy as np
@@ -22,49 +20,18 @@
 from cirq_ft.infra.jupyter_tools import execute_notebook
 
 
-def construct_gate_helper_and_qubit_order(data, eps):
-    gate = cirq_ft.StatePreparationAliasSampling.from_lcu_probs(
-        lcu_probabilities=data, probability_epsilon=eps
-    )
-    g = cirq_ft.testing.GateHelper(gate)
-    context = cirq.DecompositionContext(cirq.ops.SimpleQubitManager())
-
-    def map_func(op: cirq.Operation, _):
-        gateset = cirq.Gateset(cirq_ft.And, cirq_ft.LessThanEqualGate, cirq_ft.LessThanGate)
-        return cirq.Circuit(
-            cirq.decompose(op, on_stuck_raise=None, keep=gateset.validate, context=context)
-        )
-
-    # TODO: Do not decompose {cq.And, cq.LessThanEqualGate, cq.LessThanGate} because the
-    # `cq.map_clean_and_borrowable_qubits` currently gets confused and is not able to re-map qubits
-    # optimally; which results in a higher number of ancillas thus the tests fails due to OOO.
-    decomposed_circuit = cirq.map_operations_and_unroll(
-        g.circuit, map_func, raise_if_add_qubits=False
-    )
-    greedy_mm = cirq_ft.GreedyQubitManager(prefix="_a", size=25, maximize_reuse=True)
-    decomposed_circuit = cirq_ft.map_clean_and_borrowable_qubits(decomposed_circuit, qm=greedy_mm)
-    # We are fine decomposing the `cq.And` gates once the qubit re-mapping is complete. Ideally,
-    # we shouldn't require this two step process.
-    arithmetic_gateset = cirq.Gateset(cirq_ft.LessThanEqualGate, cirq_ft.LessThanGate)
-    decomposed_circuit = cirq.Circuit(
-        cirq.decompose(decomposed_circuit, keep=arithmetic_gateset.validate, on_stuck_raise=None)
-    )
-    ordered_input = list(itertools.chain(*g.quregs.values()))
-    qubit_order = cirq.QubitOrder.explicit(ordered_input, fallback=cirq.QubitOrder.DEFAULT)
-    return g, qubit_order, decomposed_circuit
-
-
 @pytest.mark.parametrize("num_sites, epsilon", [[2, 3e-3], [3, 3.0e-3], [4, 5.0e-3], [7, 8.0e-3]])
 def test_state_preparation_via_coherent_alias_sampling(num_sites, epsilon):
     lcu_coefficients = get_1d_Ising_lcu_coeffs(num_sites)
-    g, qubit_order, decomposed_circuit = construct_gate_helper_and_qubit_order(
-        lcu_coefficients, epsilon
-    )
-    # assertion to ensure that simulating the `decomposed_circuit` doesn't run out of memory.
-    assert len(decomposed_circuit.all_qubits()) < 25
-    result = cirq.Simulator(dtype=np.complex128).simulate(
-        decomposed_circuit, qubit_order=qubit_order
+    gate = cirq_ft.StatePreparationAliasSampling.from_lcu_probs(
+        lcu_probabilities=lcu_coefficients.tolist(), probability_epsilon=epsilon
     )
+    g = cirq_ft.testing.GateHelper(gate)
+    qubit_order = g.operation.qubits
+
+    # Assertion to ensure that simulating the `decomposed_circuit` doesn't run out of memory.
+    assert len(g.circuit.all_qubits()) < 20
+    result = cirq.Simulator(dtype=np.complex128).simulate(g.circuit, qubit_order=qubit_order)
     state_vector = result.final_state_vector
     # State vector is of the form |l>|temp_{l}>. We trace out the |temp_{l}> part to
     # get the coefficients corresponding to |l>.
@@ -82,7 +49,12 @@ def test_state_preparation_via_coherent_alias_sampling(num_sites, epsilon):
 
 def test_state_preparation_via_coherent_alias_sampling_diagram():
     data = np.asarray(range(1, 5)) / np.sum(range(1, 5))
-    g, qubit_order, _ = construct_gate_helper_and_qubit_order(data, 0.05)
+    gate = cirq_ft.StatePreparationAliasSampling.from_lcu_probs(
+        lcu_probabilities=data.tolist(), probability_epsilon=0.05
+    )
+    g = cirq_ft.testing.GateHelper(gate)
+    qubit_order = g.operation.qubits
+
     circuit = cirq.Circuit(cirq.decompose_once(g.operation))
     cirq.testing.assert_has_diagram(
         circuit,

cirq-ft/cirq_ft/algos/unary_iteration_gate.py

@@ -28,12 +28,12 @@ def _unary_iteration_segtree(
     ops: List[cirq.Operation],
     control: cirq.Qid,
     selection: Sequence[cirq.Qid],
+    ancilla: Sequence[cirq.Qid],
     sl: int,
     l: int,
     r: int,
     l_iter: int,
     r_iter: int,
-    qm: cirq.QubitManager,
 ) -> Iterator[Tuple[cirq.OP_TREE, cirq.Qid, int]]:
     """Constructs a unary iteration circuit by iterating over nodes of an implicit Segment Tree.
 
@@ -45,6 +45,8 @@ def _unary_iteration_segtree(
         selection: Sequence of selection qubits. The i'th qubit in the list corresponds to the i'th
             level in the segment tree.Thus, a total of O(logN) selection qubits are required for a
             tree on range `N = (r_iter - l_iter)`.
+        ancilla: Pre-allocated ancilla qubits to be used for constructing the unary iteration
+            circuit.
         sl: Current depth of the tree. `selection[sl]` gives the selection qubit corresponding to
             the current depth.
         l: Left index of the range represented by current node of the segment tree.
@@ -76,73 +78,73 @@ def _unary_iteration_segtree(
     if r_iter <= m:
         # Yield only left sub-tree.
         yield from _unary_iteration_segtree(
-            ops, control, selection, sl + 1, l, m, l_iter, r_iter, qm
+            ops, control, selection, ancilla, sl + 1, l, m, l_iter, r_iter
         )
         return
     if l_iter >= m:
         # Yield only right sub-tree
         yield from _unary_iteration_segtree(
-            ops, control, selection, sl + 1, m, r, l_iter, r_iter, qm
+            ops, control, selection, ancilla, sl + 1, m, r, l_iter, r_iter
         )
         return
-    anc, sq = qm.qalloc(1)[0], selection[sl]
+    anc, sq = ancilla[sl], selection[sl]
     ops.append(and_gate.And((1, 0)).on(control, sq, anc))
-    yield from _unary_iteration_segtree(ops, anc, selection, sl + 1, l, m, l_iter, r_iter, qm)
+    yield from _unary_iteration_segtree(ops, anc, selection, ancilla, sl + 1, l, m, l_iter, r_iter)
     ops.append(cirq.CNOT(control, anc))
-    yield from _unary_iteration_segtree(ops, anc, selection, sl + 1, m, r, l_iter, r_iter, qm)
+    yield from _unary_iteration_segtree(ops, anc, selection, ancilla, sl + 1, m, r, l_iter, r_iter)
     ops.append(and_gate.And(adjoint=True).on(control, sq, anc))
-    qm.qfree([anc])
 
 
 def _unary_iteration_zero_control(
     ops: List[cirq.Operation],
     selection: Sequence[cirq.Qid],
+    ancilla: Sequence[cirq.Qid],
     l_iter: int,
     r_iter: int,
-    qm: cirq.QubitManager,
 ) -> Iterator[Tuple[cirq.OP_TREE, cirq.Qid, int]]:
     sl, l, r = 0, 0, 2 ** len(selection)
     m = (l + r) >> 1
     ops.append(cirq.X(selection[0]))
     yield from _unary_iteration_segtree(
-        ops, selection[0], selection[1:], sl, l, m, l_iter, r_iter, qm
+        ops, selection[0], selection[1:], ancilla, sl, l, m, l_iter, r_iter
     )
     ops.append(cirq.X(selection[0]))
     yield from _unary_iteration_segtree(
-        ops, selection[0], selection[1:], sl, m, r, l_iter, r_iter, qm
+        ops, selection[0], selection[1:], ancilla, sl, m, r, l_iter, r_iter
     )
 
 
 def _unary_iteration_single_control(
     ops: List[cirq.Operation],
     control: cirq.Qid,
     selection: Sequence[cirq.Qid],
+    ancilla: Sequence[cirq.Qid],
     l_iter: int,
     r_iter: int,
-    qm: cirq.QubitManager,
 ) -> Iterator[Tuple[cirq.OP_TREE, cirq.Qid, int]]:
     sl, l, r = 0, 0, 2 ** len(selection)
-    yield from _unary_iteration_segtree(ops, control, selection, sl, l, r, l_iter, r_iter, qm)
+    yield from _unary_iteration_segtree(ops, control, selection, ancilla, sl, l, r, l_iter, r_iter)
 
 
 def _unary_iteration_multi_controls(
     ops: List[cirq.Operation],
     controls: Sequence[cirq.Qid],
     selection: Sequence[cirq.Qid],
+    ancilla: Sequence[cirq.Qid],
     l_iter: int,
     r_iter: int,
-    qm: cirq.QubitManager,
 ) -> Iterator[Tuple[cirq.OP_TREE, cirq.Qid, int]]:
     num_controls = len(controls)
-    and_ancilla = qm.qalloc(num_controls - 2)
-    and_target = qm.qalloc(1)[0]
+    and_ancilla = ancilla[: num_controls - 2]
+    and_target = ancilla[num_controls - 2]
     multi_controlled_and = and_gate.And((1,) * len(controls)).on_registers(
         control=np.array(controls), ancilla=np.array(and_ancilla), target=and_target
     )
     ops.append(multi_controlled_and)
-    yield from _unary_iteration_single_control(ops, and_target, selection, l_iter, r_iter, qm)
+    yield from _unary_iteration_single_control(
+        ops, and_target, selection, ancilla[num_controls - 1 :], l_iter, r_iter
+    )
     ops.append(cirq.inverse(multi_controlled_and))
-    qm.qfree(and_ancilla + [and_target])
 
 
 def unary_iteration(
@@ -203,18 +205,18 @@ def unary_iteration(
     """
     assert 2 ** len(selection) >= r_iter - l_iter
     assert len(selection) > 0
+    ancilla = qubit_manager.qalloc(max(0, len(controls) + len(selection) - 1))
     if len(controls) == 0:
-        yield from _unary_iteration_zero_control(
-            flanking_ops, selection, l_iter, r_iter, qubit_manager
-        )
+        yield from _unary_iteration_zero_control(flanking_ops, selection, ancilla, l_iter, r_iter)
     elif len(controls) == 1:
         yield from _unary_iteration_single_control(
-            flanking_ops, controls[0], selection, l_iter, r_iter, qubit_manager
+            flanking_ops, controls[0], selection, ancilla, l_iter, r_iter
         )
     else:
         yield from _unary_iteration_multi_controls(
-            flanking_ops, controls, selection, l_iter, r_iter, qubit_manager
+            flanking_ops, controls, selection, ancilla, l_iter, r_iter
         )
+    qubit_manager.qfree(ancilla)
 
 
 class UnaryIterationGate(infra.GateWithRegisters):