Sparse state preparation via alias sampling

dev_tools/autogenerate-bloqs-notebooks-v2.py

@@ -634,7 +634,8 @@
         title='State Preparation via Alias Sampling',
         module=qualtran.bloqs.state_preparation.state_preparation_alias_sampling,
         bloq_specs=[
-            qualtran.bloqs.state_preparation.state_preparation_alias_sampling._STATE_PREP_ALIAS_DOC
+            qualtran.bloqs.state_preparation.state_preparation_alias_sampling._STATE_PREP_ALIAS_DOC,
+            qualtran.bloqs.state_preparation.state_preparation_alias_sampling._SPARSE_STATE_PREP_ALIAS_DOC,
         ],
     ),
     NotebookSpecV2(

qualtran/bloqs/state_preparation/__init__.py

@@ -16,6 +16,7 @@
     PrepareUniformSuperposition,
 )
 from qualtran.bloqs.state_preparation.state_preparation_alias_sampling import (
+    SparseStatePreparationAliasSampling,
     StatePreparationAliasSampling,
 )
 from qualtran.bloqs.state_preparation.state_preparation_via_rotation import (

qualtran/bloqs/state_preparation/state_preparation_alias_sampling.ipynb

@@ -197,6 +197,176 @@
     "show_call_graph(state_prep_alias_g)\n",
     "show_counts_sigma(state_prep_alias_sigma)"
    ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f1559dde",
+   "metadata": {
+    "cq.autogen": "SparseStatePreparationAliasSampling.bloq_doc.md"
+   },
+   "source": [
+    "## `SparseStatePreparationAliasSampling`\n",
+    "Initialize a $d$-sparse state over $L$ indices using coherent alias sampling.\n",
+    "\n",
+    "In particular, we take the zero state to:\n",
+    "\n",
+    "$$\n",
+    "    \\sum_{j=0}^{d-1} \\sqrt{p_{\\mathrm{ind}_j}} |\\mathrm{ind}_j\\rangle |\\mathrm{temp}_j\\rangle\n",
+    "$$\n",
+    "\n",
+    "where $\\mathrm{ind}_j \\in [0, L)$ is the index of the $j$-th non-zero coefficient,\n",
+    "and the probabilities $p_l$ are $\\mu$-bit binary approximations to the true values,\n",
+    "and the register $|\\mathrm{temp}_j\\rangle$ may be entangled with the index register.\n",
+    "\n",
+    "This bloq is nearly identical to :class:`StatePreparationByAliasSampling`, except\n",
+    "that it loads the non-zero indices from the QROM and prepares a dense state on them.\n",
+    "In comparison, this uses $\\lceil \\log d \\rceil$ extra ancilla qubits, and reduces\n",
+    "the iteration length to $d$ from $L$.\n",
+    "\n",
+    "See :class:`StatePreparationAliasSampling` for an exposition on alias sampling.\n",
+    "\n",
+    "\n",
+    "#### Registers\n",
+    " - `selection`: The input/output register $|\\mathrm{ind}_l\\rangle$ of size lg(L) where the desired coefficient state is prepared.\n",
+    " - `sigma_mu`: A mu-sized register containing uniform probabilities for comparison against `keep`.\n",
+    " - `sparse_index`: A lg(d)-sized register storing the sparse index $j \\in [0, d)$.\n",
+    " - `alt`: A lg(L)-sized register of alternate indices\n",
+    " - `keep`: a mu-sized register of probabilities of keeping the initially sampled index.\n",
+    " - `less_than_equal`: one bit for the result of the comparison. \n",
+    "\n",
+    "This gate corresponds to the following operations:\n",
+    " - UNIFORM_d on the `sparse_index` register.\n",
+    " - H^mu on the `sigma` register.\n",
+    " - QROM addressed by the `sparse_index` register into the `selection`, `alt`, and `keep` signature.\n",
+    " - LessThanEqualGate comparing the `keep` and `sigma` registers.\n",
+    " - Coherent swap between the `selection` and `alt` registers if the comparison returns True.\n",
+    "\n",
+    "Total space will be $(2 \\log(L) + \\log(d) + 2 \\mu + 1)$ work qubits + $log(L)$ ancillas for QROM.\n",
+    "\n",
+    "#### References\n",
+    " - [1] [Qubitization of Arbitrary Basis Quantum Chemistry Leveraging Sparsity and Low Rank Factorization](https://arxiv.org/pdf/1902.02134#page=15.30) Berry et al. (2019). Section 5, Eqs. 43, 44. [2] [Encoding Electronic Spectra in Quantum Circuits with Linear T Complexity](https://arxiv.org/abs/1805.03662). Babbush et al. (2018). Section III.D. and Figure 11.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a5c09b86",
+   "metadata": {
+    "cq.autogen": "SparseStatePreparationAliasSampling.bloq_doc.py"
+   },
+   "outputs": [],
+   "source": [
+    "from qualtran.bloqs.state_preparation import SparseStatePreparationAliasSampling"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4cba7cd8",
+   "metadata": {
+    "cq.autogen": "SparseStatePreparationAliasSampling.example_instances.md"
+   },
+   "source": [
+    "### Example Instances"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "90068f1c",
+   "metadata": {
+    "cq.autogen": "SparseStatePreparationAliasSampling.sparse_state_prep_alias"
+   },
+   "outputs": [],
+   "source": [
+    "coeff_map = {0: 1.0, 3: 1.0, 5: 3.0, 7: 2.0}\n",
+    "N = 9\n",
+    "mu = 3\n",
+    "sparse_state_prep_alias = SparseStatePreparationAliasSampling.from_sparse_dict(\n",
+    "    coeff_map, N, precision=2**-mu / len(coeff_map)\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "beeea709",
+   "metadata": {
+    "cq.autogen": "SparseStatePreparationAliasSampling.sparse_state_prep_alias_from_list"
+   },
+   "outputs": [],
+   "source": [
+    "coeffs = [1.0, 0, 0, 1, 0, 3, 0, 2, 0]\n",
+    "mu = 3\n",
+    "sparse_state_prep_alias_from_list = (\n",
+    "    SparseStatePreparationAliasSampling.from_dense_probabilities(coeffs, precision=2**-mu / 4)\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1bf9d4f9",
+   "metadata": {
+    "cq.autogen": "SparseStatePreparationAliasSampling.sparse_state_prep_alias_symb"
+   },
+   "outputs": [],
+   "source": [
+    "import sympy\n",
+    "\n",
+    "n_coeffs, n_nonzero_coeffs, sum_coeff, eps = sympy.symbols(r\"L d \\lambda \\epsilon\")\n",
+    "sparse_state_prep_alias_symb = SparseStatePreparationAliasSampling.from_n_coeff(\n",
+    "    n_coeffs, n_nonzero_coeffs, sum_coeff, precision=eps\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d3ca3eec",
+   "metadata": {
+    "cq.autogen": "SparseStatePreparationAliasSampling.graphical_signature.md"
+   },
+   "source": [
+    "#### Graphical Signature"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7a294570",
+   "metadata": {
+    "cq.autogen": "SparseStatePreparationAliasSampling.graphical_signature.py"
+   },
+   "outputs": [],
+   "source": [
+    "from qualtran.drawing import show_bloqs\n",
+    "show_bloqs([sparse_state_prep_alias, sparse_state_prep_alias_from_list, sparse_state_prep_alias_symb],\n",
+    "           ['`sparse_state_prep_alias`', '`sparse_state_prep_alias_from_list`', '`sparse_state_prep_alias_symb`'])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4f02d080",
+   "metadata": {
+    "cq.autogen": "SparseStatePreparationAliasSampling.call_graph.md"
+   },
+   "source": [
+    "### Call Graph"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "56ac92d0",
+   "metadata": {
+    "cq.autogen": "SparseStatePreparationAliasSampling.call_graph.py"
+   },
+   "outputs": [],
+   "source": [
+    "from qualtran.resource_counting.generalizers import ignore_split_join\n",
+    "sparse_state_prep_alias_g, sparse_state_prep_alias_sigma = sparse_state_prep_alias.call_graph(max_depth=1, generalizer=ignore_split_join)\n",
+    "show_call_graph(sparse_state_prep_alias_g)\n",
+    "show_counts_sigma(sparse_state_prep_alias_sigma)"
+   ]
   }
  ],
  "metadata": {