Add QROAMClean and QROAMCleanAdjoint bloqs for data loading (#1291)

* Add QROAMClean and QROAMCleanAdjoint bloqs for data loading

* Address nits, add lots of docstrings

* Generate notebook for QROAM bloq

* Add notebook, add ceil to cost so number of toffoli is not a fraction

dev_tools/autogenerate-bloqs-notebooks-v2.py

@@ -745,6 +745,14 @@
             qualtran.bloqs.data_loading.select_swap_qrom._SELECT_SWAP_QROM_DOC,
         ],
     ),
+    NotebookSpecV2(
+        title='Advanced QROM (aka QROAM) using clean ancilla',
+        module=qualtran.bloqs.data_loading.qroam_clean,
+        bloq_specs=[
+            qualtran.bloqs.data_loading.qrom_base._QROM_BASE_DOC,
+            qualtran.bloqs.data_loading.qroam_clean._QROAM_CLEAN_DOC,
+        ],
+    ),
     NotebookSpecV2(
         title='Reflections',
         module=qualtran.bloqs.reflections,

docs/bloqs/index.rst

@@ -129,6 +129,7 @@ Bloqs Library
     multiplexers/apply_lth_bloq.ipynb
     data_loading/qrom.ipynb
     data_loading/select_swap_qrom.ipynb
+    data_loading/qroam_clean.ipynb
     reflections/reflections.ipynb
     mcmt/multi_control_multi_target_pauli.ipynb
     multiplexers/select_pauli_lcu.ipynb

qualtran/bloqs/data_loading/__init__.py

@@ -14,5 +14,6 @@
 
 """Bloqs to load classical data in a quantum register"""
 
+from qualtran.bloqs.data_loading.qroam_clean import QROAMClean, QROAMCleanAdjoint
 from qualtran.bloqs.data_loading.qrom import QROM
 from qualtran.bloqs.data_loading.select_swap_qrom import SelectSwapQROM

qualtran/bloqs/data_loading/qroam_clean.ipynb

@@ -0,0 +1,342 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "f81ea157",
+   "metadata": {
+    "cq.autogen": "title_cell"
+   },
+   "source": [
+    "# Advanced QROM (aka QROAM) using clean ancilla"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "65f7e061",
+   "metadata": {
+    "cq.autogen": "top_imports"
+   },
+   "outputs": [],
+   "source": [
+    "from qualtran import Bloq, CompositeBloq, BloqBuilder, Signature, Register\n",
+    "from qualtran import QBit, QInt, QUInt, QAny\n",
+    "from qualtran.drawing import show_bloq, show_call_graph, show_counts_sigma\n",
+    "from typing import *\n",
+    "import numpy as np\n",
+    "import sympy\n",
+    "import cirq"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "de80d822",
+   "metadata": {
+    "cq.autogen": "QROMBase.bloq_doc.md"
+   },
+   "source": [
+    "## `QROMBase`\n",
+    "Interface for Bloqs to load `data[l]` when the selection register stores index `l`.\n",
+    "\n",
+    "## Overview\n",
+    "The action of a QROM can be described as\n",
+    "$$\n",
+    "        \\text{QROM}_{s_1, s_2, \\dots, s_K}^{d_1, d_2, \\dots, d_L}\n",
+    "        |s_1\\rangle |s_2\\rangle \\dots |s_K\\rangle\n",
+    "        |0\\rangle^{\\otimes b_1} |0\\rangle^{\\otimes b_2} \\dots |0\\rangle^{\\otimes b_L}\n",
+    "        \\rightarrow\n",
+    "        |s_1\\rangle |s_2\\rangle \\dots |s_K\\rangle\n",
+    "        |d_1[s_1, s_2, \\dots, s_k]\\rangle\n",
+    "        |d_2[s_1, s_2, \\dots, s_k]\\rangle \\dots\n",
+    "        |d_L[s_1, s_2, \\dots, s_k]\\rangle\n",
+    "$$\n",
+    "\n",
+    "A behavior of a QROM can be understood in terms of its classical analogue, where a for-loop\n",
+    "over one or more (selection) indices can be used to load one or more classical datasets, where\n",
+    "each of the classical dataset can be multidimensional.\n",
+    "\n",
+    "```\n",
+    ">>> # N, M, P, Q, R, S, T are pre-initialized integer parameters.\n",
+    ">>> output = [np.zeros((P, Q)), np.zeros((R, S, T))]\n",
+    ">>> # Load two different classical datasets; each of different shape.\n",
+    ">>> data = [np.random.rand(N, M, P, Q), np.random.rand(N, M, R, S, T)]\n",
+    ">>> for i in range(N): # For loop over two selection indices i and j.\n",
+    ">>>     for j in range(M):\n",
+    ">>>        # Load two multidimensional classical datasets data[0] and data[1] s.t.\n",
+    ">>>        # |i, j⟩|0⟩  -> |i, j⟩|data[0][i, j, :]⟩|data[1][i, j, :]⟩\n",
+    ">>>        output[0] = data[0][i, j, :]\n",
+    ">>>        output[1] = data[1][i, j, :]\n",
+    "```\n",
+    "\n",
+    "The parameters that control the behavior and costs of a QROM are -\n",
+    "\n",
+    "1. Number of selection registers (eg: $i$, $j$) and their iteration lengths (eg: $N$, $M$).\n",
+    "2. Number of target registers, their quantum datatype and shape.\n",
+    "    - Number of target registers: One for each classical dataset to load (eg: $\\text{data}[0]$\n",
+    "        and $\\text{data}[1]$)\n",
+    "    - QDType of target registers: Depends on `dtype` of the $i$'th classical dataset\n",
+    "    - Shape of target registers: Depends on shape of classical data (eg: $(P, Q)$ and\n",
+    "        $(R, S, T)$ above)\n",
+    "\n",
+    "### Specification of classical data via `data_or_shape`\n",
+    "Users can specify the classical data to load via QROM by passing in an appropriate value\n",
+    "for `data_or_shape` attribute. This is a list of numpy arrays or `Shaped` objects, where\n",
+    "each item of the list corresponds to a classical dataset to load.\n",
+    "\n",
+    "Each classical dataset to load can be specified as a numpy array (or a `Shaped` object for\n",
+    "symbolic bloqs). The shape of the dataset is a union of the selection shape and target shape,\n",
+    "s.t.\n",
+    "$$\n",
+    "    \\text{data[i].shape} = \\text{selection\\_shape} + \\text{target\\_shape[i]}\n",
+    "$$\n",
+    "\n",
+    "Note that the $\\text{selection\\_shape}$ should be same across all classical datasets to be\n",
+    "loaded and correspond to a tuple of iteration lengths of selection indices (i.e. $(N, M)$\n",
+    "in the example above).\n",
+    "\n",
+    "The target shape of each classical dataset can be different and parameterizes the size of\n",
+    "the desired output that should be loaded in a target register.\n",
+    "\n",
+    "### Number of selection registers and their iteration lengths\n",
+    "As describe in the previous section, the number of selection registers and their iteration\n",
+    "lengths can be inferred from the shape of the classical dataset. All classical datasets\n",
+    "to be loaded must have the same $\\text{selection\\_shape}$, which is a tuple of iteration\n",
+    "lengths over each dimension of the dataset (i.e. the range for each nested for-loop).\n",
+    "\n",
+    "In order to load a data set with $\\text{selection\\_shape} == (P, Q, R, S)$ the QROM bloq\n",
+    "needs four selection registers with bitsizes $(p, q, r, s)$ where each of\n",
+    "$p,q,r,s \\geq \\log_2{P}, \\log_2{Q}, \\log_2{R}, \\log_2{S}$.\n",
+    "\n",
+    "In general, to load $K$ dimensional data, we use $K$ named selection registers\n",
+    "$(\\text{selection}_0, \\text{selection}_1, ..., \\text{selection}_k)$ to index and\n",
+    "load the data. For the $i$'th selection register, its size is configured using\n",
+    "attribute $\\text{selection\\_bitsizes[i]}$ and the iteration range is configued\n",
+    "using $\\text{data\\_or\\_shape[0].shape[i]}$.\n",
+    "\n",
+    "### Number of target registers, their quantum datatype and shape\n",
+    "QROM bloq uses one target register for each entry corresponding to classical dataset in the\n",
+    "tuple `data_or_shape`. Thus, to load $L$ classical datsets, we use $L$ names target registers\n",
+    "$(\\text{target}_0, \\text{target}_1, ..., \\text{target}_L)$\n",
+    "\n",
+    "Each named target register has a bitsize $b_{i}=\\text{target\\_bitsizes[i]}$ that represents\n",
+    "the size of the register and depends upon the maximum value of individual elements in the\n",
+    "$i$'th classical dataset.\n",
+    "\n",
+    "Each named target register has a shape that can be configured using attribute\n",
+    "$\\text{target\\_shape[i]}$ that represents the number of target registers if the output to load\n",
+    "is multidimensional.\n",
+    "\n",
+    "#### Parameters\n",
+    " - `data_or_shape`: List of numpy ndarrays specifying the data to load. If the length of this list ($L$) is greater than one then we use the same selection indices to load each dataset. The shape of a classical dataset is a concatenation of selection_shape and target_shape[i]; i.e. `data_or_shape[i].shape = selection_shape + target_shape[i]`. Thus, each data set is required to have the same selection shape $(S_1, S_2, ..., S_K)$ and can have a different target shape given by `target_shapes[i]`. For symbolic QROMs, pass a list of `Shaped` objects instead with shape $(S_1, S_2, ..., S_K) + target_shape[i]$.\n",
+    " - `selection_bitsizes`: The number of bits used to represent each selection register corresponding to the size of each dimension of the selection_shape $(S_1, S_2, ..., S_K)$. Should be the same length as the selection shape of each of the datasets and $2**\\text{selection\\_bitsizes[i]} >= S_i$\n",
+    " - `target_shapes`: Shape of target registers for each classical dataset to be loaded. Must be consistent with `data_or_shape` s.t. `len(data_or_shape) == len(target_shapes)` and `data_or_shape[-len(target_shapes[i]):] == target_shapes[i]`.\n",
+    " - `target_bitsizes`: Bitsize (or qdtype) of the target registers for each classical dataset to be loaded. This can be deduced from the maximum element of each of the datasets. Must be consistent with `data_or_shape` s.t. `len(data_or_shape) == len(target_bitsizes)` and `target_bitsizes[i] >= max(data[i]).bitsize`.\n",
+    " - `num_controls`: The number of controls to instanstiate a controlled version of this bloq.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7853bb2e",
+   "metadata": {
+    "cq.autogen": "QROMBase.bloq_doc.py"
+   },
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "markdown",
+   "id": "5d91b534",
+   "metadata": {
+    "cq.autogen": "QROAMClean.bloq_doc.md"
+   },
+   "source": [
+    "## `QROAMClean`\n",
+    "Lower cost variant of SelectSwapQROM. Assumes target register is initially in |0> state.\n",
+    "\n",
+    "To load a classical dataset of $N$ elements, each of bitsize $b$, into a target register initialized\n",
+    "in the $|0\\rangle$ state, this construction uses:\n",
+    "    - $\\frac{N}{K} + (K - 1) \\times b$ Toffoli gates.\n",
+    "    - $(K - 1)$ ancilla registers, each of bitsize $b$, left in a junk state and should be kept\n",
+    "     around to get uncomputed by the adjoint bloq - `QROAMCleanAdjoint`.\n",
+    "\n",
+    "Here $K=2^k$ is a configurable constant and should be set to $\\sqrt{\\frac{N}{b}}$ for optimal cost.\n",
+    "\n",
+    "Similar to SelectSwapQROM, this bloq also supports loading multiple classical datasets,\n",
+    "each of which can be multidimensional. Factory methods `QROAMClean.build_from_data` and\n",
+    "`QROAMClean.build_from_bitsize` should be used to construct the bloq.\n",
+    "\n",
+    "The adjoint of the bloq is performed via `QROAMCleanAdjoint`, and reduces to a problem of\n",
+    "uncomputing a table lookup with $N$ elements, each of target bitsize $K \\times b$. The data to\n",
+    "be loaded for uncomputation is computed by this bloq in the `self.batched_data_permuted`\n",
+    "property.\n",
+    "\n",
+    "`QROAMCleanAdjoint` uses measurement based uncomputation to uncompute a table lookup of $N$\n",
+    "elements and target bitsize $b$ using only $\\frac{N}{K} + (K - 1)$ Toffoli gates\n",
+    "(instead of $\\frac{N}{K} + (K - 1) \\times b$ used by the original lookup). Thus, increasing the\n",
+    "target bitsize for uncomputation is preferred since complexity of uncomputation does not depend\n",
+    "upon the target bitsize of elements to be loaded.\n",
+    "\n",
+    "#### Registers\n",
+    " - `- control_registers`: If control is specified, a THRU register to denote the control qubits. Empty by default for uncontrolled version of the Bloq.\n",
+    " - `- selection_registers`: $N$ THRU registers, each with shape (), to load $N$ dimensional classical datasets.\n",
+    " - `- target_registers`: $M$ RIGHT registers to load $M$ different classical datasets. Each target register is of bitsize $b$ and shape described by a tuple of length $N$.\n",
+    " - `- junk_registers`: $K - 1$ RIGHT registers, each of bitsize $b$ used to load batches of size $K$ \n",
+    "\n",
+    "#### References\n",
+    " - [Qubitization of Arbitrary Basis Quantum Chemistry Leveraging Sparsity and Low Rank Factorization](https://arxiv.org/abs/1902.02134).     Berry et. al. (2019). Appendix A. and B.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b781be03",
+   "metadata": {
+    "cq.autogen": "QROAMClean.bloq_doc.py"
+   },
+   "outputs": [],
+   "source": [
+    "from qualtran.bloqs.data_loading.qroam_clean import QROAMClean"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "78f192d4",
+   "metadata": {
+    "cq.autogen": "QROAMClean.example_instances.md"
+   },
+   "source": [
+    "### Example Instances"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a4fa5027",
+   "metadata": {
+    "cq.autogen": "QROAMClean.qroam_clean_multi_data"
+   },
+   "outputs": [],
+   "source": [
+    "data1 = np.arange(5, dtype=int)\n",
+    "data2 = np.arange(5, dtype=int) + 1\n",
+    "qroam_clean_multi_data = QROAMClean.build_from_data(data1, data2, log_block_sizes=(1,))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6f19e239",
+   "metadata": {
+    "cq.autogen": "QROAMClean.qroam_clean_multi_dim"
+   },
+   "outputs": [],
+   "source": [
+    "data1 = np.arange(25, dtype=int).reshape((5, 5))\n",
+    "data2 = (np.arange(25, dtype=int) + 1).reshape((5, 5))\n",
+    "qroam_clean_multi_dim = QROAMClean.build_from_data(data1, data2, log_block_sizes=(1, 1))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "0ab833fb",
+   "metadata": {
+    "cq.autogen": "QROAMClean.qroam_clean_symb_1d"
+   },
+   "outputs": [],
+   "source": [
+    "N, b, k, c = sympy.symbols('N b k c')\n",
+    "qroam_clean_symb_1d = QROAMClean.build_from_bitsize(\n",
+    "    (N,), (b,), log_block_sizes=(k,), num_controls=c\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "4e97cfe8",
+   "metadata": {
+    "cq.autogen": "QROAMClean.qroam_clean_symb_2d"
+   },
+   "outputs": [],
+   "source": [
+    "N, M, b1, b2, k1, k2, c = sympy.symbols('N M b1 b2 k1 k2 c')\n",
+    "log_block_sizes = (k1, k2)\n",
+    "qroam_clean_symb_2d = QROAMClean.build_from_bitsize(\n",
+    "    (N, M), (b1, b2), log_block_sizes=log_block_sizes, num_controls=c\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "eefaa8d8",
+   "metadata": {
+    "cq.autogen": "QROAMClean.graphical_signature.md"
+   },
+   "source": [
+    "#### Graphical Signature"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "bbaa7ea7",
+   "metadata": {
+    "cq.autogen": "QROAMClean.graphical_signature.py"
+   },
+   "outputs": [],
+   "source": [
+    "from qualtran.drawing import show_bloqs\n",
+    "show_bloqs([qroam_clean_multi_data, qroam_clean_multi_dim],\n",
+    "           ['`qroam_clean_multi_data`', '`qroam_clean_multi_dim`'])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d9f0a049",
+   "metadata": {
+    "cq.autogen": "QROAMClean.call_graph.md"
+   },
+   "source": [
+    "### Call Graph"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7ddfc52d",
+   "metadata": {
+    "cq.autogen": "QROAMClean.call_graph.py"
+   },
+   "outputs": [],
+   "source": [
+    "from qualtran.resource_counting.generalizers import ignore_split_join\n",
+    "qroam_clean_multi_data_g, qroam_clean_multi_data_sigma = qroam_clean_multi_data.call_graph(max_depth=1, generalizer=ignore_split_join)\n",
+    "show_call_graph(qroam_clean_multi_data_g)\n",
+    "show_counts_sigma(qroam_clean_multi_data_sigma)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.8"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}