SelectSwapQROM revamp and upgrades (#986)

* SwapWithZero learns how to swap multidimensional selection index.

* Clear output

* Fix mypy and use wire_symbol instead of cirq diagram. Fix a bug in _wire_symbol_to_cirq_diagram_info

* SelectSwapQROM revamp and upgrades

* Fix pylint and typos

* Fix mypy errors and address comments

* Add type ignores, fix notebook and nits

* Update docstrings and regenerate ipynb files

dev_tools/autogenerate-bloqs-notebooks-v2.py

@@ -75,6 +75,8 @@
 import qualtran.bloqs.chemistry.trotter.ising.unitaries
 import qualtran.bloqs.chemistry.trotter.trotterized_unitary
 import qualtran.bloqs.data_loading.qrom
+import qualtran.bloqs.data_loading.qrom_base
+import qualtran.bloqs.data_loading.select_swap_qrom
 import qualtran.bloqs.factoring.ecc
 import qualtran.bloqs.factoring.mod_exp
 import qualtran.bloqs.hamiltonian_simulation.hamiltonian_simulation_by_gqsp
@@ -496,7 +498,18 @@
     NotebookSpecV2(
         title='QROM',
         module=qualtran.bloqs.data_loading.qrom,
-        bloq_specs=[qualtran.bloqs.data_loading.qrom._QROM_DOC],
+        bloq_specs=[
+            qualtran.bloqs.data_loading.qrom_base._QROM_BASE_DOC,
+            qualtran.bloqs.data_loading.qrom._QROM_DOC,
+        ],
+    ),
+    NotebookSpecV2(
+        title='SelectSwapQROM',
+        module=qualtran.bloqs.data_loading.select_swap_qrom,
+        bloq_specs=[
+            qualtran.bloqs.data_loading.qrom_base._QROM_BASE_DOC,
+            qualtran.bloqs.data_loading.select_swap_qrom._SELECT_SWAP_QROM_DOC,
+        ],
     ),
     NotebookSpecV2(
         title='Block Encoding',

docs/bloqs/index.rst

@@ -95,6 +95,7 @@ Bloqs Library
     hubbard_model.ipynb
     multiplexers/apply_gate_to_lth_target.ipynb
     data_loading/qrom.ipynb
+    data_loading/select_swap_qrom.ipynb
     block_encoding.ipynb
     reflection.ipynb
     mcmt/multi_control_multi_target_pauli.ipynb

qualtran/bloqs/chemistry/sparse/prepare.py

@@ -43,6 +43,7 @@
     PrepareUniformSuperposition,
 )
 from qualtran.linalg.lcu_util import preprocess_lcu_coefficients_for_reversible_sampling
+from qualtran.symbolics.math_funcs import ceil, log2
 
 if TYPE_CHECKING:
     from qualtran import Bloq
@@ -313,10 +314,10 @@ def build_qrom_bloq(self) -> 'Bloq':
             (n_n,) * 4 + (1,) * 2 + (n_n,) * 4 + (1,) * 2 + (self.num_bits_state_prep,)
         )
         if self.qroam_block_size is None:
-            block_size = 2 ** find_optimal_log_block_size(self.num_non_zero, sum(target_bitsizes))
+            log_block_sizes = find_optimal_log_block_size(self.num_non_zero, sum(target_bitsizes))
         else:
-            block_size = self.qroam_block_size
-        qrom = SelectSwapQROM(
+            log_block_sizes = ceil(log2(self.qroam_block_size))
+        qrom = SelectSwapQROM.build_from_data(
             self.ind_pqrs[0],
             self.ind_pqrs[1],
             self.ind_pqrs[2],
@@ -331,7 +332,7 @@ def build_qrom_bloq(self) -> 'Bloq':
             self.alt_one_body,
             self.keep,
             target_bitsizes=target_bitsizes,
-            block_size=block_size,
+            log_block_sizes=log_block_sizes,
         )
         return qrom
 

qualtran/bloqs/chemistry/thc/notebook_utils.py

@@ -48,8 +48,8 @@
 
 
 def custom_qroam_repr(self) -> str:
-    target_repr = repr(self._target_bitsizes)
-    return f"SelectSwapQROM(target_bitsizes={target_repr}, block_size={self.block_size})"
+    target_repr = repr(self.target_bitsizes)
+    return f"SelectSwapQROM(target_bitsizes={target_repr}, block_sizes={self.block_sizes})"
 
 
 # TODO: better way of customizing label

qualtran/bloqs/chemistry/thc/prepare.py

@@ -382,7 +382,7 @@ def build_composite_bloq(
         # 2. Make contiguous register from mu and nu and store in register `s`.
         mu, nu, s = bb.add(ToContiguousIndex(log_mu, log_d), mu=mu, nu=nu, s=s)
         # 3. Load alt / keep values
-        qroam = SelectSwapQROM(
+        qroam = SelectSwapQROM.build_from_data(
             *(self.theta, self.alt_theta, self.alt_mu, self.alt_nu, self.keep),
             target_bitsizes=(1, 1, log_mu, log_mu, self.keep_bitsize),
         )
@@ -444,7 +444,7 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']:
         data_size = self.num_spin_orb // 2 + self.num_mu * (self.num_mu + 1) // 2
         nd = (data_size - 1).bit_length()
         cost_2 = (ToContiguousIndex(nmu, nd), 1)
-        qroam = SelectSwapQROM(
+        qroam = SelectSwapQROM.build_from_data(
             *(self.theta, self.alt_theta, self.alt_mu, self.alt_nu, self.keep),
             target_bitsizes=(1, 1, nmu, nmu, self.keep_bitsize),
         )

qualtran/bloqs/chemistry/writing_algorithms.ipynb

@@ -375,15 +375,15 @@
     "        target_bitsizes = (n_n,) * 4 + (self.num_bits_state_prep,)\n",
     "        ns = self.num_spin_orb // 2\n",
     "        data_size = ns ** 2 + ns**4\n",
-    "        block_size = 2 ** find_optimal_log_block_size(data_size, sum(target_bitsizes))\n",
-    "        qroam = SelectSwapQROM(\n",
+    "        log_block_size = find_optimal_log_block_size(data_size, sum(target_bitsizes))\n",
+    "        qroam = SelectSwapQROM.build_from_data(\n",
     "            self.alt_pqrs[0],\n",
     "            self.alt_pqrs[1],\n",
     "            self.alt_pqrs[2],\n",
     "            self.alt_pqrs[3],\n",
     "            self.keep,\n",
     "            target_bitsizes=target_bitsizes,\n",
-    "            block_size=block_size,\n",
+    "            log_block_sizes=[log_block_size],\n",
     "        )\n",
     "        (l, alt_pqrs[0], alt_pqrs[1], alt_pqrs[2], alt_pqrs[3], keep) = bb.add(\n",
     "            qroam,\n",