diff --git a/dev_tools/autogenerate-bloqs-notebooks-v2.py b/dev_tools/autogenerate-bloqs-notebooks-v2.py index cb675a426..3943eda03 100644 --- a/dev_tools/autogenerate-bloqs-notebooks-v2.py +++ b/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, diff --git a/docs/bloqs/index.rst b/docs/bloqs/index.rst index 256f60daf..7c9b559b1 100644 --- a/docs/bloqs/index.rst +++ b/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 diff --git a/qualtran/bloqs/data_loading/__init__.py b/qualtran/bloqs/data_loading/__init__.py index 3d40fd5b6..8be91c4bb 100644 --- a/qualtran/bloqs/data_loading/__init__.py +++ b/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 diff --git a/qualtran/bloqs/data_loading/qroam_clean.ipynb b/qualtran/bloqs/data_loading/qroam_clean.ipynb new file mode 100644 index 000000000..2d5652cfc --- /dev/null +++ b/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 +} diff --git a/qualtran/bloqs/data_loading/qroam_clean.py b/qualtran/bloqs/data_loading/qroam_clean.py new file mode 100644 index 000000000..bb70cb98f --- /dev/null +++ b/qualtran/bloqs/data_loading/qroam_clean.py @@ -0,0 +1,469 @@ +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import numbers +from collections import defaultdict +from functools import cached_property +from typing import cast, Dict, List, Optional, Set, Tuple, Type, TYPE_CHECKING, Union + +import attrs +import numpy as np +import sympy +from numpy.typing import ArrayLike + +from qualtran import bloq_example, BloqDocSpec, GateWithRegisters, Register, Side, Signature +from qualtran.bloqs.basic_gates import Toffoli +from qualtran.bloqs.data_loading.qrom_base import QROMBase +from qualtran.drawing import Circle, LarrowTextBox, RarrowTextBox, Text, TextBox, WireSymbol +from qualtran.symbolics import ceil, is_symbolic, log2, prod, SymbolicFloat, SymbolicInt + +if TYPE_CHECKING: + from qualtran import Bloq, BloqBuilder, SoquetT, QDType + from qualtran.simulation.classical_sim import ClassicalValT + from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + +from qualtran.bloqs.data_loading.select_swap_qrom import _alloc_anc_for_reg, SelectSwapQROM + + +def _alloc_anc_for_reg_except_first( + bb: 'BloqBuilder', dtype: 'QDType', shape: Tuple[int, ...], dirty: bool +) -> 'SoquetT': + if not shape: + return bb.allocate(dtype=dtype, dirty=dirty) + soqs = np.empty(shape, dtype=object) + ndindex_iter = np.ndindex(shape) + _ = next(ndindex_iter) # Skip the first element. + for idx in ndindex_iter: + soqs[idx] = bb.allocate(dtype=dtype, dirty=dirty) + return soqs + + +def qroam_cost(x, data_size: SymbolicInt, bitsize: SymbolicInt, adjoint: bool = False): + # See appendix B of https://arxiv.org/pdf/1902.02134 + if adjoint: + return data_size / x + x + else: + return data_size / x + bitsize * (x - 1) + + +def get_optimal_log_block_size_clean_ancilla( + data_size: SymbolicInt, + bitsize: SymbolicInt, + adjoint: bool = False, + qroam_block_size: Optional[SymbolicInt] = None, +) -> SymbolicInt: + if qroam_block_size is None: + if adjoint: + k: SymbolicFloat = 0.5 * log2(data_size) + else: + k = 0.5 * log2(data_size / bitsize) + else: + k = log2(qroam_block_size) + if is_symbolic(k): + return k + k_int = np.array([np.floor(k), np.ceil(k)]) + return int(k_int[np.argmin(qroam_cost(2**k_int, data_size, bitsize, adjoint))]) + + +@attrs.frozen +class QROAMCleanAdjoint(QROMBase, GateWithRegisters): # type: ignore[misc] + r"""Measurement based uncomputation of a table lookup. Assumes target is left in |0> state. + + Measurement based uncomputation of a table, as described in Ref[1]. Follows the usual API + of a QROM bloq. Assumes that the target register was initially in the $|0\rangle$ state before + the lookup was performed and must be deallocated and left in |0\rangle$ state after + uncomputation. + + Reduces the problem of uncomputating a data lookup with $N$ elements and target bitsize $b$ to + doing another data lookup with $\frac{N}{2}$ elements, each of bitsize 2. The reduction is + explained in Ref[1] and uses X-basis measurements + classically controlled Cliffords. A variant + of `QROAM` bloq can be used to perform the resulting lookup and resulting junk registers can be + cleaned up with measurement based uncomputation of SwapWithZero. + + Thus, cost of uncomputing a data lookup with $N$ elements and target bitsize $b$ is + $\frac{N}{K} + (K - 1)$ Toffoli gates (instead of $\frac{N}{K} + (K - 1) \times b$ used by the + original lookup). + + Registers: + - control_registers: If control is specified, a THRU register to denote the control qubits. + Empty by default for uncontrolled version of the Bloq. + - selection_registers: $N$ THRU registers, each with shape (), to load $N$ dimensional + classical datasets. + - target_registers: $M$ LEFT registers to load $M$ different classical datasets. Each target + register is of bitsize $b$ and shape described by a tuple of length $N + S$. Here $S$ is + a parameter that describes the shape of the output to be loaded for each selection index. + + + References: + [Qubitization of Arbitrary Basis Quantum Chemistry Leveraging Sparsity and Low Rank Factorization](https://arxiv.org/abs/1902.02134). + Berry et. al. (2019). Appendix C. + """ + log_block_sizes: Tuple[SymbolicInt, ...] = attrs.field( + converter=lambda x: tuple(x.tolist() if isinstance(x, np.ndarray) else x) + ) + + @cached_property + def _target_reg_side(self) -> Side: + return Side.LEFT + + @classmethod + def build_from_data( + cls: Type['QROAMCleanAdjoint'], + *data: ArrayLike, + target_bitsizes: Optional[Union[SymbolicInt, Tuple[SymbolicInt, ...]]] = None, + target_shapes: Tuple[Tuple[SymbolicInt, ...], ...] = (), + num_controls: SymbolicInt = 0, + log_block_sizes: Optional[Union[SymbolicInt, Tuple[SymbolicInt, ...]]] = None, + ) -> 'QROAMCleanAdjoint': + qroam: 'QROAMCleanAdjoint' = cls._build_from_data( + *data, + target_bitsizes=target_bitsizes, + target_shapes=target_shapes, + num_controls=num_controls, + ) + return qroam.with_log_block_sizes(log_block_sizes) + + @classmethod + def build_from_bitsize( + cls: Type['QROAMCleanAdjoint'], + data_len_or_shape: Union[SymbolicInt, Tuple[SymbolicInt, ...]], + target_bitsizes: Union[SymbolicInt, Tuple[SymbolicInt, ...]], + *, + target_shapes: Tuple[Tuple[SymbolicInt, ...], ...] = (), + selection_bitsizes: Tuple[SymbolicInt, ...] = (), + num_controls: SymbolicInt = 0, + log_block_sizes: Optional[Union[SymbolicInt, Tuple[SymbolicInt, ...]]] = None, + ) -> 'QROAMCleanAdjoint': + qroam: 'QROAMCleanAdjoint' = cls._build_from_bitsize( + data_len_or_shape, + target_bitsizes, + selection_bitsizes=selection_bitsizes, + target_shapes=target_shapes, + num_controls=num_controls, + ) + return qroam.with_log_block_sizes(log_block_sizes=log_block_sizes) + + @log_block_sizes.default + def _default_log_block_sizes(self) -> Tuple[SymbolicInt, ...]: + target_bitsize = sum(self.target_bitsizes) * sum( + prod(shape) for shape in self.target_shapes + ) + return tuple( + get_optimal_log_block_size_clean_ancilla(ilen, target_bitsize, adjoint=True) + for ilen in self.data_shape + ) + + def with_log_block_sizes( + self, log_block_sizes: Optional[Union[SymbolicInt, Tuple[SymbolicInt, ...]]] = None + ) -> 'QROAMCleanAdjoint': + if log_block_sizes is None: + return self + if isinstance(log_block_sizes, (int, sympy.Basic, numbers.Number)): + log_block_sizes = (log_block_sizes,) + if not is_symbolic(*log_block_sizes): + assert all(1 <= 2**bs <= ilen for bs, ilen in zip(log_block_sizes, self.data_shape)) + return attrs.evolve(self, log_block_sizes=log_block_sizes) + + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + block_sizes = prod([2**k for k in self.log_block_sizes]) + data_size = prod(self.data_shape) + n_toffoli = ceil(data_size / block_sizes) + block_sizes + return {(Toffoli(), n_toffoli)} + + @cached_property + def signature(self) -> Signature: + return Signature( + [*self.control_registers, *self.selection_registers, *self.target_registers] + ) + + def adjoint(self) -> 'QROAMClean': + return QROAMClean(**attrs.asdict(self)) + + def wire_symbol(self, reg: Optional[Register], idx: Tuple[int, ...] = tuple()) -> 'WireSymbol': + if reg is None: + return Text('QROAM').adjoint() + name = reg.name + if name == 'selection': + return TextBox('In').adjoint() + elif 'target' in name: + trg_indx = int(name.replace('target', '').replace('_', '')) + # match the sel index + subscript = chr(ord('a') + trg_indx) + return LarrowTextBox(f'QROAM_{subscript}') + elif name == 'control': + return Circle() + raise ValueError(f'Unknown register name {name}') + + +@attrs.frozen +class QROAMClean(SelectSwapQROM): + r"""Lower cost variant of SelectSwapQROM. Assumes target register is initially in |0> state. + + To load a classical dataset of $N$ elements, each of bitsize $b$, into a target register initialized + in the $|0\rangle$ state, this construction uses: + - $\frac{N}{K} + (K - 1) \times b$ Toffoli gates. + - $(K - 1)$ ancilla registers, each of bitsize $b$, left in a junk state and should be kept + around to get uncomputed by the adjoint bloq - `QROAMCleanAdjoint`. + + Here $K=2^k$ is a configurable constant and should be set to $\sqrt{\frac{N}{b}}$ for optimal cost. + + Similar to SelectSwapQROM, this bloq also supports loading multiple classical datasets, + each of which can be multidimensional. Factory methods `QROAMClean.build_from_data` and + `QROAMClean.build_from_bitsize` should be used to construct the bloq. + + The adjoint of the bloq is performed via `QROAMCleanAdjoint`, and reduces to a problem of + uncomputing a table lookup with $N$ elements, each of target bitsize $K \times b$. The data to + be loaded for uncomputation is computed by this bloq in the `self.batched_data_permuted` + property. + + `QROAMCleanAdjoint` uses measurement based uncomputation to uncompute a table lookup of $N$ + elements and target bitsize $b$ using only $\frac{N}{K} + (K - 1)$ Toffoli gates + (instead of $\frac{N}{K} + (K - 1) \times b$ used by the original lookup). Thus, increasing the + target bitsize for uncomputation is preferred since complexity of uncomputation does not depend + upon the target bitsize of elements to be loaded. + + Registers: + - control_registers: If control is specified, a THRU register to denote the control qubits. + Empty by default for uncontrolled version of the Bloq. + - selection_registers: $N$ THRU registers, each with shape (), to load $N$ dimensional + classical datasets. + - 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$. + - junk_registers: $K - 1$ RIGHT registers, each of bitsize $b$ used to load batches of size $K$ + + References: + [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. + """ + log_block_sizes: Tuple[SymbolicInt, ...] = attrs.field( + converter=lambda x: tuple(x.tolist() if isinstance(x, np.ndarray) else x) + ) + use_dirty_ancilla: bool = attrs.field(init=False, default=False, repr=False) + + @cached_property + def _target_reg_side(self) -> Side: + return Side.RIGHT + + @log_block_sizes.default + def _default_log_block_sizes(self) -> Tuple[SymbolicInt, ...]: + target_bitsize = sum(self.target_bitsizes) * sum( + prod(shape) for shape in self.target_shapes + ) + return tuple( + get_optimal_log_block_size_clean_ancilla(ilen, target_bitsize) + for ilen in self.data_shape + ) + + @classmethod + def build_from_data( + cls: Type['QROAMClean'], + *data: ArrayLike, + target_bitsizes: Optional[Union[SymbolicInt, Tuple[SymbolicInt, ...]]] = None, + num_controls: SymbolicInt = 0, + log_block_sizes: Optional[Union[SymbolicInt, Tuple[SymbolicInt, ...]]] = None, + ) -> 'QROAMClean': + qroam: 'QROAMClean' = cls._build_from_data( + *data, target_bitsizes=target_bitsizes, num_controls=num_controls + ) + return qroam.with_log_block_sizes(log_block_sizes=log_block_sizes) + + @classmethod + def build_from_bitsize( + cls: Type['QROAMClean'], + data_len_or_shape: Union[SymbolicInt, Tuple[SymbolicInt, ...]], + target_bitsizes: Union[SymbolicInt, Tuple[SymbolicInt, ...]], + *, + selection_bitsizes: Tuple[SymbolicInt, ...] = (), + num_controls: SymbolicInt = 0, + log_block_sizes: Optional[Union[SymbolicInt, Tuple[SymbolicInt, ...]]] = None, + ) -> 'QROAMClean': + qroam: 'QROAMClean' = cls._build_from_bitsize( + data_len_or_shape, + target_bitsizes, + selection_bitsizes=selection_bitsizes, + num_controls=num_controls, + ) + return qroam.with_log_block_sizes(log_block_sizes=log_block_sizes) + + @cached_property + def signature(self) -> Signature: + return Signature( + [ + *self.control_registers, + *self.selection_registers, + *self.target_registers, + *self.junk_registers, + ] + ) + + @cached_property + def batched_data_permuted(self) -> List[np.ndarray]: + if is_symbolic(*self.block_sizes): + raise ValueError( + f"Cannot decompose SelectSwapQROM bloq with symbolic block sizes. Found {self.block_sizes=}" + ) + block_sizes = cast(Tuple[int, ...], self.block_sizes) + ret = [] + for data, swz in zip(self.batched_data, self.swap_with_zero_bloqs): + permuted_batched_data = np.zeros(data.shape + block_sizes, dtype=data.dtype) + for sel_l in np.ndindex(cast(Tuple[int, ...], self.batched_qrom_shape)): + for sel_k in np.ndindex(block_sizes): + sel_kwargs = {reg.name: sel for reg, sel in zip(swz.selection_registers, sel_k)} + curr_data = swz.call_classically(**sel_kwargs, targets=np.copy(data[sel_l]))[-1] + idx = (*sel_l, *sel_k) + permuted_batched_data[idx][:] = curr_data + + n_blocks = len(self.block_sizes) + transpose_axes = [x for i in range(n_blocks) for x in [i, i + n_blocks]] + transpose_axes += [i + 2 * n_blocks for i in range(n_blocks)] + desired_shape = tuple( + NbyK * K for NbyK, K in zip(self.batched_qrom_shape, self.block_sizes) + ) + ret.append( + permuted_batched_data.transpose(transpose_axes).reshape( + desired_shape + self.block_sizes + ) + ) + return ret + + @cached_property + def junk_registers(self) -> Tuple[Register, ...]: + # The newly allocated registers should be kept around for measurement based uncomputation. + junk_regs = [] + block_size = prod(self.block_sizes) + for reg in self.target_registers: + assert reg.shape == () + if is_symbolic(block_size) or block_size > 1: + junk_regs += [attrs.evolve(reg, name='junk_' + reg.name, shape=(block_size - 1,))] + return tuple(junk_regs) + + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + ret: Dict[Bloq, SymbolicInt] = defaultdict(lambda: 0) + ret[self.qrom_bloq] += 1 + for swz in self.swap_with_zero_bloqs: + if any(is_symbolic(s) or s > 0 for s in swz.selection_bitsizes): + ret[swz] += 1 + return set(ret.items()) + + def _build_composite_bloq_with_swz_clean( + self, + bb: 'BloqBuilder', + ctrl: List['SoquetT'], + selection: List['SoquetT'], + qrom_targets: List['SoquetT'], + ) -> Tuple[List['SoquetT'], List['SoquetT'], List['SoquetT']]: + sel_l, sel_k = self._partition_sel_register(bb, selection) + ctrl, sel_l, qrom_targets = self._add_qrom_bloq(bb, ctrl, sel_l, qrom_targets) + sel_k, qrom_targets = self._add_swap_with_zero_bloq(bb, sel_k, qrom_targets) + selection = self._unpartition_sel_register(bb, sel_l, sel_k) + return ctrl, selection, qrom_targets + + def build_composite_bloq(self, bb: 'BloqBuilder', **soqs: 'SoquetT') -> Dict[str, 'SoquetT']: + # Get the ctrl and target register for the SelectSwapQROM. + ctrl = [soqs.pop(reg.name) for reg in self.control_registers] + selection = [soqs.pop(reg.name) for reg in self.selection_registers] + if is_symbolic(*self.block_sizes): + raise ValueError( + f"Cannot decompose QROAM bloq with symbolic block sizes. Found {self.block_sizes=}" + ) + block_sizes = cast(Tuple[int, ...], self.block_sizes) + # Allocate intermediate clean/dirty ancilla for the underlying QROM call. + qrom_targets = [] + for reg in self.target_registers: + qrom_target = _alloc_anc_for_reg_except_first( + bb, reg.dtype, block_sizes, self.use_dirty_ancilla + ) + qrom_target[np.unravel_index(0, block_sizes)] = _alloc_anc_for_reg( # type: ignore[index] + bb, reg.dtype, reg.shape, dirty=False + ) + qrom_targets.append(qrom_target) + # Assert that all registers have been used by now. + assert not soqs, f"All registers must have been used by now. Found: {soqs}" + # Add the bloq decomposition + if any(b > 1 for b in block_sizes): + ctrl, selection, qrom_targets = self._build_composite_bloq_with_swz_clean( + bb, ctrl, selection, qrom_targets + ) + else: + ctrl, selection, qrom_targets = self._add_qrom_bloq(bb, ctrl, selection, qrom_targets) + # Construct and return dictionary of final soquets. + soqs |= {reg.name: soq for reg, soq in zip(self.control_registers, ctrl)} + soqs |= {reg.name: soq for reg, soq in zip(self.selection_registers, selection)} + soqs |= {reg.name: soq.flat[1:] for reg, soq in zip(self.junk_registers, qrom_targets)} # type: ignore[union-attr] + soqs |= {reg.name: soq.flat[0] for reg, soq in zip(self.target_registers, qrom_targets)} # type: ignore[union-attr] + return soqs + + def on_classical_vals( + self, **vals: Union['sympy.Symbol', 'ClassicalValT'] + ) -> Dict[str, 'ClassicalValT']: + vals_without_junk = super().on_classical_vals(**vals) + selection = cast(Tuple[int, ...], tuple(vals[reg.name] for reg in self.selection_registers)) + for d, junk_reg in zip(self.batched_data_permuted, self.junk_registers): + vals_without_junk[junk_reg.name] = d[selection].flat[1:] + return vals_without_junk + + def adjoint(self) -> 'QROAMCleanAdjoint': + if self.has_data(): + return QROAMCleanAdjoint.build_from_data( + *self.batched_data_permuted, + target_shapes=(self.block_sizes,) * len(self.batched_data_permuted), + ) + else: + return QROAMCleanAdjoint.build_from_bitsize( + self.data_shape, + target_bitsizes=self.target_bitsizes, + target_shapes=(self.block_sizes,) * len(self.target_bitsizes), + ) + + def wire_symbol(self, reg: Optional[Register], idx: Tuple[int, ...] = tuple()) -> 'WireSymbol': + if reg is None: + return Text('QROAM') + name = reg.name + if name == 'selection': + return TextBox('In') + elif 'target' in name: + trg_indx = int(name.replace('target', '').replace('_', '')) + # match the sel index + subscript = chr(ord('a') + trg_indx) + return RarrowTextBox(f'QROAM_{subscript}') + elif 'junk' in name: + junk_idx = int(name.replace('junk_target', '').replace('_', '')) + # match the sel index + subscript = chr(ord('a') + junk_idx) + return RarrowTextBox(f'junk_{subscript}') + elif name == 'control': + return Circle() + raise ValueError(f'Unknown register name {name}') + + +@bloq_example +def _qroam_clean_multi_data() -> QROAMClean: + data1 = np.arange(5, dtype=int) + data2 = np.arange(5, dtype=int) + 1 + qroam_clean_multi_data = QROAMClean.build_from_data(data1, data2, log_block_sizes=(1,)) + return qroam_clean_multi_data + + +@bloq_example +def _qroam_clean_multi_dim() -> QROAMClean: + data1 = np.arange(25, dtype=int).reshape((5, 5)) + data2 = (np.arange(25, dtype=int) + 1).reshape((5, 5)) + qroam_clean_multi_dim = QROAMClean.build_from_data(data1, data2, log_block_sizes=(1, 1)) + return qroam_clean_multi_dim + + +_QROAM_CLEAN_DOC = BloqDocSpec( + bloq_cls=QROAMClean, + import_line='from qualtran.bloqs.data_loading.qroam_clean import QROAMClean', + examples=[_qroam_clean_multi_data, _qroam_clean_multi_dim], +) diff --git a/qualtran/bloqs/data_loading/qroam_clean_test.py b/qualtran/bloqs/data_loading/qroam_clean_test.py new file mode 100644 index 000000000..932e9848c --- /dev/null +++ b/qualtran/bloqs/data_loading/qroam_clean_test.py @@ -0,0 +1,141 @@ +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import numpy as np +import pytest +import sympy + +from qualtran.bloqs.data_loading.qroam_clean import ( + _qroam_clean_multi_data, + _qroam_clean_multi_dim, + QROAMClean, + QROAMCleanAdjoint, +) +from qualtran.symbolics import ceil + + +def test_bloq_examples(bloq_autotester): + bloq_autotester(_qroam_clean_multi_data) + bloq_autotester(_qroam_clean_multi_dim) + + +def test_t_complexity_1d_data_symbolic(): + # 1D data, 1 dataset + 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 + bloq_inv = bloq.adjoint() + assert isinstance(bloq_inv, QROAMCleanAdjoint) + inv_k = sympy.symbols('kinv') + inv_K = 2**inv_k + bloq_inv = bloq_inv.with_log_block_sizes(log_block_sizes=(inv_k,)) + expected_toffoli_inv = ceil(N / inv_K) + inv_K + assert bloq_inv.t_complexity().t == 4 * expected_toffoli_inv + + +def test_t_complexity_2d_data_symbolic(): + # 2D data, 1 dataset + 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 + bloq_inv = bloq.adjoint() + assert isinstance(bloq_inv, QROAMCleanAdjoint) + inv_k1, inv_k2 = sympy.symbols('kinv1, kinv2') + inv_K1, inv_K2 = 2**inv_k1, 2**inv_k2 + bloq_inv = bloq_inv.with_log_block_sizes(log_block_sizes=(inv_k1, inv_k2)) + expected_toffoli_inv = ceil(N1 * N2 / (inv_K1 * inv_K2)) + inv_K1 * inv_K2 + assert bloq_inv.t_complexity().t == 4 * expected_toffoli_inv + + +def test_qroam_clean_classical_sim(): + rng = np.random.default_rng(42) + # 1D data, 1 dataset + N, max_N, log_block_sizes = 25, 2**10, 3 + data = rng.integers(max_N, size=N) + bloq = QROAMClean.build_from_data(data, log_block_sizes=log_block_sizes) + cbloq = bloq.decompose_bloq() + bloq_inv = bloq.adjoint() + assert isinstance(bloq_inv, QROAMCleanAdjoint) + for x in range(N): + vals = bloq.call_classically(selection=x) + cvals = cbloq.call_classically(selection=x) + assert vals[0:2] == cvals[0:2] == (x, data[x]) + assert np.array_equal(vals[2], cvals[2]) + target_with_junk = np.array([vals[1], *vals[2]]) # type: ignore[misc] + assert bloq_inv.call_classically(selection=vals[0], target0_=target_with_junk) == (x,) + + # 2D data, 1 datasets + N, M, max_N, log_block_sizes = 7, 11, 2**5, (2, 3) + data = rng.integers(max_N, size=N * M).reshape(N, M) + bloq = QROAMClean.build_from_data(data, log_block_sizes=log_block_sizes) + cbloq = bloq.decompose_bloq() + bloq_inv = bloq.adjoint() + assert isinstance(bloq_inv, QROAMCleanAdjoint) + for x in range(N): + for y in range(M): + vals = bloq.call_classically(selection0=x, selection1=y) + cvals = cbloq.call_classically(selection0=x, selection1=y) + assert vals[0:3] == cvals[0:3] == (x, y, data[x][y]) + assert np.array_equal(vals[3], cvals[3]) + target_with_junk = np.array([vals[2], *vals[3]]).reshape(2 ** np.array(log_block_sizes)) # type: ignore[misc] + assert bloq_inv.call_classically( + selection0=x, selection1=y, target0_=target_with_junk + ) == (x, y) + + +@pytest.mark.slow +def test_qroam_clean_classical_sim_multi_dataset(): + rng = np.random.default_rng(42) + # 1D data, 2 datasets + N, max_N, log_block_sizes = 25, 2**20, 3 + data = [rng.integers(max_N, size=N), rng.integers(max_N, size=N)] + bloq = QROAMClean.build_from_data(*data, log_block_sizes=log_block_sizes) + cbloq = bloq.decompose_bloq() + bloq_inv = bloq.adjoint() + assert isinstance(bloq_inv, QROAMCleanAdjoint) + for x in range(N): + vals = bloq.call_classically(selection=x) + cvals = cbloq.call_classically(selection=x) + assert vals[0:3] == cvals[0:3] == (x, data[0][x], data[1][x]) + assert np.array_equal(vals[3], cvals[3]) and np.array_equal(vals[4], cvals[4]) + targets_with_junk0 = np.array([vals[1], *vals[3]]) # type: ignore[misc] + targets_with_junk1 = np.array([vals[2], *vals[4]]) # type: ignore[misc] + assert bloq_inv.call_classically( + selection=vals[0], target0_=targets_with_junk0, target1_=targets_with_junk1 + ) == (x,) + + # 2D data, 2 datasets + N, M, max_N, log_block_sizes = 7, 11, 2**5, np.array([2, 3]) # type: ignore[misc] + data = [ + rng.integers(max_N, size=N * M).reshape(N, M), + rng.integers(max_N, size=N * M).reshape(N, M), + ] + bloq = QROAMClean.build_from_data(*data, log_block_sizes=tuple(log_block_sizes.tolist())) + cbloq = bloq.decompose_bloq() + bloq_inv = bloq.adjoint() + assert isinstance(bloq_inv, QROAMCleanAdjoint) + for x in range(N): + for y in range(M): + vals = bloq.call_classically(selection0=x, selection1=y) + cvals = cbloq.call_classically(selection0=x, selection1=y) + assert vals[0:4] == cvals[0:4] == (x, y, data[0][x][y], data[1][x][y]) + assert np.array_equal(vals[4], cvals[4]) and np.array_equal(vals[5], cvals[5]) + targets_with_junk0 = np.array([vals[2], *vals[4]]).reshape(2**log_block_sizes) # type: ignore[misc] + targets_with_junk1 = np.array([vals[3], *vals[5]]).reshape(2**log_block_sizes) # type: ignore[misc] + assert bloq_inv.call_classically( + selection0=x, selection1=y, target0_=targets_with_junk0, target1_=targets_with_junk1 + ) == (x, y) diff --git a/qualtran/bloqs/data_loading/qrom.py b/qualtran/bloqs/data_loading/qrom.py index 6491caad7..8524be9d9 100644 --- a/qualtran/bloqs/data_loading/qrom.py +++ b/qualtran/bloqs/data_loading/qrom.py @@ -96,10 +96,14 @@ def build_from_data( cls, *data: ArrayLike, target_bitsizes: Optional[Union[SymbolicInt, Tuple[SymbolicInt, ...]]] = None, + target_shapes: Tuple[Tuple[SymbolicInt, ...], ...] = (), num_controls: SymbolicInt = 0, ) -> 'QROM': return cls._build_from_data( - *data, target_bitsizes=target_bitsizes, num_controls=num_controls + *data, + target_bitsizes=target_bitsizes, + target_shapes=target_shapes, + num_controls=num_controls, ) @classmethod diff --git a/qualtran/bloqs/data_loading/qrom_base.py b/qualtran/bloqs/data_loading/qrom_base.py index 7e69eb06f..dc751949c 100644 --- a/qualtran/bloqs/data_loading/qrom_base.py +++ b/qualtran/bloqs/data_loading/qrom_base.py @@ -23,7 +23,7 @@ import sympy from numpy.typing import ArrayLike, NDArray -from qualtran import BloqDocSpec, BoundedQUInt, QAny, Register +from qualtran import BloqDocSpec, BoundedQUInt, QAny, Register, Side from qualtran.simulation.classical_sim import ClassicalValT from qualtran.symbolics import bit_length, is_symbolic, shape, Shaped, SymbolicInt @@ -221,16 +221,22 @@ def _build_from_data( cls: Type[QROM_T], *data: ArrayLike, target_bitsizes: Optional[Union[SymbolicInt, Tuple[SymbolicInt, ...]]] = None, + target_shapes: Tuple[Tuple[SymbolicInt, ...], ...] = (), num_controls: SymbolicInt = 0, ) -> QROM_T: _data = [np.array(d, dtype=int) for d in data] - selection_bitsizes = tuple((s - 1).bit_length() for s in _data[0].shape) if target_bitsizes is None: target_bitsizes = tuple(max(int(np.max(d)).bit_length(), 1) for d in data) + assert isinstance(target_bitsizes, tuple) # Make mypy happy. + if target_shapes == (): + target_shapes = ((),) * len(target_bitsizes) + selection_len = len(_data[0].shape) - len(target_shapes[0]) + selection_bitsizes = tuple((s - 1).bit_length() for s in _data[0].shape[:selection_len]) return cls( data_or_shape=_data, selection_bitsizes=selection_bitsizes, target_bitsizes=target_bitsizes, + target_shapes=target_shapes, num_controls=num_controls, ) @@ -287,10 +293,14 @@ def selection_registers(self) -> Tuple[Register, ...]: return (Register('selection', types[0]),) return tuple(Register(f'selection{i}', qdtype) for i, qdtype in enumerate(types)) + @cached_property + def _target_reg_side(self) -> Side: + return Side.THRU + @cached_property def target_registers(self) -> Tuple[Register, ...]: return tuple( - Register(f'target{i}_', QAny(l), shape=sh) + Register(f'target{i}_', QAny(l), shape=sh, side=self._target_reg_side) for i, (l, sh) in enumerate(zip(self.target_bitsizes, self.target_shapes)) if is_symbolic(l) or l ) @@ -323,13 +333,15 @@ def on_classical_vals( # Retrieve the data; bitwise add them in to the input target values targets = {f'target{d_i}_': d[idx] for d_i, d in enumerate(self.data)} - targets = {k: v ^ vals[k] for k, v in targets.items()} + targets = {k: v ^ vals.get(k, 0) for k, v in targets.items()} + if not (self._target_reg_side & Side.RIGHT): + for reg_name, reg_val in targets.items(): + if np.any(reg_val): + raise ValueError( + f"Target register {reg_name} must be uncomputed before de-allocation. Found values {reg_val}" + ) + targets = {} return controls | selections | targets - def __pow__(self, power: int): - if power in [1, -1]: - return self - return NotImplemented # pragma: no cover - _QROM_BASE_DOC = BloqDocSpec(bloq_cls=QROMBase, import_line='', examples=[]) diff --git a/qualtran/bloqs/data_loading/select_swap_qrom.ipynb b/qualtran/bloqs/data_loading/select_swap_qrom.ipynb index b1d610ca4..4bd012a63 100644 --- a/qualtran/bloqs/data_loading/select_swap_qrom.ipynb +++ b/qualtran/bloqs/data_loading/select_swap_qrom.ipynb @@ -215,7 +215,7 @@ "source": [ "data1 = np.arange(5)\n", "data2 = np.arange(5) + 1\n", - "qroam_multi_data = SelectSwapQROM.build_from_data([data1, data2])" + "qroam_multi_data = SelectSwapQROM.build_from_data(data1, data2)" ] }, { @@ -229,7 +229,7 @@ "source": [ "data1 = np.arange(25).reshape((5, 5))\n", "data2 = (np.arange(25) + 1).reshape((5, 5))\n", - "qroam_multi_dim = SelectSwapQROM.build_from_data([data1, data2])" + "qroam_multi_dim = SelectSwapQROM.build_from_data(data1, data2)" ] }, { diff --git a/qualtran/bloqs/data_loading/select_swap_qrom.py b/qualtran/bloqs/data_loading/select_swap_qrom.py index 046c7d908..504a9a914 100644 --- a/qualtran/bloqs/data_loading/select_swap_qrom.py +++ b/qualtran/bloqs/data_loading/select_swap_qrom.py @@ -14,7 +14,7 @@ import numbers from collections import defaultdict from functools import cached_property -from typing import cast, Dict, List, Optional, Set, Tuple, Type, TYPE_CHECKING, Union +from typing import cast, Dict, List, Optional, Set, Tuple, Type, TYPE_CHECKING, TypeVar, Union import attrs import cirq @@ -35,6 +35,8 @@ from qualtran import Bloq, BloqBuilder, QDType, SoquetT from qualtran.resource_counting import BloqCountT, SympySymbolAllocator +SelSwapQROM_T = TypeVar('SelSwapQROM_T', bound='SelectSwapQROM') + def find_optimal_log_block_size( iteration_length: SymbolicInt, target_bitsize: SymbolicInt @@ -136,7 +138,7 @@ def signature(self) -> Signature: # Builder methods and helpers. @log_block_sizes.default - def _default_block_sizes(self) -> Tuple[SymbolicInt, ...]: + def _default_log_block_sizes(self) -> Tuple[SymbolicInt, ...]: target_bitsize = sum(self.target_bitsizes) * sum( prod(shape) for shape in self.target_shapes ) @@ -181,8 +183,9 @@ def build_from_bitsize( return qroam.with_log_block_sizes(log_block_sizes=log_block_sizes) def with_log_block_sizes( - self, log_block_sizes: Optional[Union[SymbolicInt, Tuple[SymbolicInt, ...]]] = None - ) -> 'SelectSwapQROM': + self: SelSwapQROM_T, + log_block_sizes: Optional[Union[SymbolicInt, Tuple[SymbolicInt, ...]]] = None, + ) -> 'SelSwapQROM_T': if log_block_sizes is None: return self if isinstance(log_block_sizes, (int, sympy.Basic, numbers.Number)): @@ -222,7 +225,7 @@ def batched_data(self) -> List[np.ndarray]: # # For data[N1][N2] with block sizes (k1, k2), you load batches of size `(k1, k2)` at once. # Thus, you load batch[N1/k1][N2/k2] where batch[i][j] = data[i*k1:(i + 1)*k1][j*k2:(j + 1)*k2] - batched_data = [np.empty(self.batched_data_shape) for _ in range(len(self.target_bitsizes))] + batched_data = [np.zeros(self.batched_data_shape, dtype=int) for _ in self.target_bitsizes] block_slices = [slice(0, k) for k in self.block_sizes] for i, data in enumerate(self.padded_data): for batch_idx in np.ndindex(cast(Tuple[int, ...], self.batched_qrom_shape)): @@ -451,7 +454,7 @@ def wire_symbol(self, reg: Optional[Register], idx: Tuple[int, ...] = tuple()) - def _qroam_multi_data() -> SelectSwapQROM: data1 = np.arange(5) data2 = np.arange(5) + 1 - qroam_multi_data = SelectSwapQROM.build_from_data([data1, data2]) + qroam_multi_data = SelectSwapQROM.build_from_data(data1, data2) return qroam_multi_data @@ -459,7 +462,7 @@ def _qroam_multi_data() -> SelectSwapQROM: def _qroam_multi_dim() -> SelectSwapQROM: data1 = np.arange(25).reshape((5, 5)) data2 = (np.arange(25) + 1).reshape((5, 5)) - qroam_multi_dim = SelectSwapQROM.build_from_data([data1, data2]) + qroam_multi_dim = SelectSwapQROM.build_from_data(data1, data2) return qroam_multi_dim diff --git a/qualtran/drawing/qpic_diagram.py b/qualtran/drawing/qpic_diagram.py index e8dae8f97..befd4f1f1 100644 --- a/qualtran/drawing/qpic_diagram.py +++ b/qualtran/drawing/qpic_diagram.py @@ -26,7 +26,7 @@ from collections import defaultdict from typing import Dict, List, Optional, Set, Tuple, TYPE_CHECKING, Union -from qualtran import DanglingT, LeftDangle, QBit, Side, Soquet +from qualtran import DanglingT, LeftDangle, QBit, RightDangle, Side, Soquet from qualtran.drawing.musical_score import ( _soq_to_symb, Circle, @@ -147,14 +147,27 @@ def __init__(self): # don't have any incoming or outgoing edges. self.empty_wire = None - def add_wires_for_signature(self, signature: 'Signature') -> None: - for reg in signature: + def add_left_wires_for_signature(self, signature: 'Signature') -> None: + for reg in signature.lefts(): for idx in reg.all_idxs(): self._alloc_wire_for_soq(Soquet(LeftDangle, reg, idx)) # Add horizontal blank space since left dangling wires would have annotations # corresponding to their QDType, which takes up horizontal space. self.wires += ['LABEL length=10'] + def add_right_wires_for_signature(self, signature: 'Signature') -> None: + add_space = False + for reg in signature.rights(): + if reg.side & Side.LEFT: + continue + for idx in reg.all_idxs(): + soq = Soquet(RightDangle, reg, idx) + wire_name = self.wire_manager.soq_to_wirename(self.soq_map[soq]) + self.gates += [f'{wire_name} / {_format_label_text(soq.pretty(), scale=0.5)} '] + add_space = True + if add_space: + self.gates += ['LABEL length=10'] + @property def data(self) -> List[str]: return self.wires + self.gates @@ -270,9 +283,10 @@ def get_qpic_data(bloq: 'Bloq', file_path: Union[None, pathlib.Path, str] = None """ cbloq = bloq.as_composite_bloq() qpic_circuit = QpicCircuit() - qpic_circuit.add_wires_for_signature(cbloq.signature) + qpic_circuit.add_left_wires_for_signature(cbloq.signature) for binst, pred, succ in cbloq.iter_bloqnections(): qpic_circuit.add_bloq(binst.bloq, pred, succ) + qpic_circuit.add_right_wires_for_signature(cbloq.signature) if file_path: with open(file_path, 'w') as f: f.write('\n'.join(qpic_circuit.data)) diff --git a/qualtran/serialization/resolver_dict.py b/qualtran/serialization/resolver_dict.py index b66e5b8ff..4f3782808 100644 --- a/qualtran/serialization/resolver_dict.py +++ b/qualtran/serialization/resolver_dict.py @@ -94,6 +94,7 @@ import qualtran.bloqs.chemistry.trotter.hubbard.interaction import qualtran.bloqs.chemistry.trotter.ising.unitaries import qualtran.bloqs.chemistry.trotter.trotterized_unitary +import qualtran.bloqs.data_loading.qroam_clean import qualtran.bloqs.data_loading.qrom import qualtran.bloqs.data_loading.select_swap_qrom import qualtran.bloqs.factoring.mod_exp @@ -317,6 +318,8 @@ "qualtran.bloqs.chemistry.trotter.hubbard.hopping.HoppingTileHWP": qualtran.bloqs.chemistry.trotter.hubbard.hopping.HoppingTileHWP, "qualtran.bloqs.chemistry.trotter.trotterized_unitary": qualtran.bloqs.chemistry.trotter.trotterized_unitary, "qualtran.bloqs.data_loading.qrom.QROM": qualtran.bloqs.data_loading.qrom.QROM, + "qualtran.bloqs.data_loading.qroam_clean.QROAMClean": qualtran.bloqs.data_loading.qroam_clean.QROAMClean, + "qualtran.bloqs.data_loading.qroam_clean.QROAMCleanAdjoint": qualtran.bloqs.data_loading.qroam_clean.QROAMCleanAdjoint, "qualtran.bloqs.data_loading.select_swap_qrom.SelectSwapQROM": qualtran.bloqs.data_loading.select_swap_qrom.SelectSwapQROM, "qualtran.bloqs.mod_arithmetic.CModAddK": qualtran.bloqs.mod_arithmetic.CModAddK, "qualtran.bloqs.mod_arithmetic.mod_addition.ModAddK": qualtran.bloqs.mod_arithmetic.mod_addition.ModAddK,