Fix unwanted phase in QROAMClean

qualtran/bloqs/data_loading/qroam_clean.py

@@ -42,6 +42,7 @@
     from qualtran.simulation.classical_sim import ClassicalValT
 
 from qualtran.bloqs.data_loading.select_swap_qrom import _alloc_anc_for_reg, SelectSwapQROM
+from qualtran.bloqs.swap_network.swap_with_zero import SwapWithZero
 
 
 def _alloc_anc_for_reg_except_first(
@@ -362,6 +363,18 @@ class QROAMClean(SelectSwapQROM):
     )
     use_dirty_ancilla: bool = attrs.field(init=False, default=False, repr=False)
 
+    @cached_property
+    def swap_with_zero_bloqs(self) -> List[SwapWithZero]:
+        return [
+            SwapWithZero(
+                self.log_block_sizes,
+                target_bitsize=target_bitsize,
+                n_target_registers=self.block_sizes,
+                use_exact_cswap=True,
+            )
+            for target_bitsize in self.target_bitsizes
+        ]
+
     @cached_property
     def _target_reg_side(self) -> Side:
         return Side.RIGHT

qualtran/bloqs/data_loading/qroam_clean_test.py

@@ -15,8 +15,10 @@
 import pytest
 import sympy
 
+from qualtran import BloqBuilder
 from qualtran._infra.data_types import QAny
 from qualtran._infra.registers import Register
+from qualtran.bloqs.basic_gates import QUIntState
 from qualtran.bloqs.data_loading.qroam_clean import (
     _qroam_clean_multi_data,
     _qroam_clean_multi_dim,
@@ -88,8 +90,8 @@ def test_t_complexity_1d_data_symbolic():
     N, b, k = sympy.symbols('N b k')
     bloq = QROAMClean.build_from_bitsize((N,), (b,), log_block_sizes=(k,))
     K = 2**k
-    expected_toffoli = ceil(N / K) + (K - 1) * b - 2
-    assert bloq.t_complexity().t == 4 * expected_toffoli
+    expected_t = 4 * (ceil(N / K) - 2) + 7 * (K - 1) * b
+    assert sympy.simplify(bloq.t_complexity().t - expected_t) == 0
     bloq_inv = bloq.adjoint()
     assert isinstance(bloq_inv, QROAMCleanAdjointWrapper)
     inv_k = sympy.symbols('kinv')
@@ -104,8 +106,8 @@ def test_t_complexity_2d_data_symbolic():
     N1, N2, b, k1, k2 = sympy.symbols('N1 N2 b k1, k2')
     bloq = QROAMClean.build_from_bitsize((N1, N2), (b,), log_block_sizes=(k1, k2))
     K1, K2 = 2**k1, 2**k2
-    expected_toffoli = ceil(N1 / K1) * ceil(N2 / K2) + (K1 * K2 - 1) * b - 2
-    assert bloq.t_complexity().t == 4 * expected_toffoli
+    expected_t = 4 * (ceil(N1 / K1) * ceil(N2 / K2) - 2) + 7 * (K1 * K2 - 1) * b
+    assert sympy.simplify(bloq.t_complexity().t - expected_t) == 0
     bloq_inv = bloq.adjoint()
     assert isinstance(bloq_inv, QROAMCleanAdjointWrapper)
     inv_k1, inv_k2 = sympy.symbols('kinv1, kinv2')
@@ -227,3 +229,15 @@ def test_qroam_clean_small_bloq_counts():
         qroam_clean=qroam_clean, log_block_sizes=(0,)
     )
     assert qroam_clean_adj_wrapper.t_complexity().t == 0
+
+
+def test_qroam_clean_no_unwanted_phase():
+    qroam = QROAMClean.build_from_data(np.ones(4, dtype=int), log_block_sizes=(1,))
+    bb = BloqBuilder()
+    sel = bb.add(QUIntState(1, 2))
+    psi = bb.finalize(**bb.add_d(qroam, selection=sel)).tensor_contract()
+
+    # The expected state is |0111>, with a positive amplitude.
+    expected = np.zeros(16, dtype=complex)
+    expected[7] = 1
+    assert np.allclose(psi, expected)

qualtran/bloqs/swap_network/swap_with_zero.py

@@ -32,6 +32,7 @@
     Soquet,
     SoquetT,
 )
+from qualtran.bloqs.basic_gates import CSwap
 from qualtran.bloqs.swap_network.cswap_approx import CSwapApprox
 from qualtran.drawing import TextBox, WireSymbol
 from qualtran.resource_counting.generalizers import ignore_split_join
@@ -128,6 +129,7 @@ class SwapWithZero(GateWithRegisters):
     target_bitsize: SymbolicInt
     n_target_registers: Tuple[SymbolicInt, ...] = attrs.field(converter=_to_tuple)
     uncompute: bool = False
+    use_exact_cswap: bool = False
 
     def __attrs_post_init__(self):
         assert len(self.n_target_registers) == len(self.selection_bitsizes)
@@ -154,7 +156,9 @@ def signature(self) -> Signature:
         return Signature([*self.selection_registers, *self.target_registers])
 
     @cached_property
-    def cswap_n(self) -> 'CSwapApprox':
+    def cswap_n(self) -> Union['CSwapApprox', 'CSwap']:
+        if self.use_exact_cswap:
+            return CSwap(self.target_bitsize)
         return CSwapApprox(self.target_bitsize)
 
     @cached_property