From 301e3cf259dcbb95ebf29e97fe53c5139233104a Mon Sep 17 00:00:00 2001 From: Anurudh Peduri <7265746+anurudhp@users.noreply.github.com> Date: Wed, 17 Jul 2024 06:54:22 -0700 Subject: [PATCH] Sparse state preparation via alias sampling (#1067) * WIP sparse SP alias sampling * remove cirq.value equality * types * example * add tests * mypy * fix decomp * fix validity test * docstring * one more test * add tolerance parameter for nonzero check * update docstring and parameter names * call preprocess only on nonzero values * cleanup * support sparse state input as map * move examples * symbolic example * notebook * test symbolic t complexity --------- Co-authored-by: Tanuj Khattar --- dev_tools/autogenerate-bloqs-notebooks-v2.py | 3 +- qualtran/bloqs/state_preparation/__init__.py | 1 + .../state_preparation_alias_sampling.ipynb | 170 +++++++++++ .../state_preparation_alias_sampling.py | 288 +++++++++++++++++- .../state_preparation_alias_sampling_test.py | 67 +++- qualtran/conftest.py | 1 + qualtran/serialization/resolver_dict.py | 1 + 7 files changed, 506 insertions(+), 25 deletions(-) diff --git a/dev_tools/autogenerate-bloqs-notebooks-v2.py b/dev_tools/autogenerate-bloqs-notebooks-v2.py index 277adc13e..36c055b48 100644 --- a/dev_tools/autogenerate-bloqs-notebooks-v2.py +++ b/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( diff --git a/qualtran/bloqs/state_preparation/__init__.py b/qualtran/bloqs/state_preparation/__init__.py index 9c781cf80..024a57c22 100644 --- a/qualtran/bloqs/state_preparation/__init__.py +++ b/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 ( diff --git a/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.ipynb b/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.ipynb index a7a66bade..fb06cda64 100644 --- a/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.ipynb +++ b/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": { diff --git a/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py b/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py index 89562b0b9..200155568 100644 --- a/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py +++ b/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py @@ -30,7 +30,7 @@ from qualtran import bloq_example, BloqDocSpec, BoundedQUInt, Register, Signature from qualtran._infra.gate_with_registers import total_bits from qualtran.bloqs.arithmetic import LessThanEqual -from qualtran.bloqs.basic_gates import CSwap, Hadamard +from qualtran.bloqs.basic_gates import CSwap, Hadamard, OnEach from qualtran.bloqs.block_encoding.lcu_select_and_prepare import PrepareOracle from qualtran.bloqs.data_loading.qrom import QROM from qualtran.bloqs.state_preparation.prepare_uniform_superposition import ( @@ -45,13 +45,21 @@ ignore_cliffords, ignore_split_join, ) -from qualtran.symbolics import bit_length, Shaped, SymbolicFloat, SymbolicInt +from qualtran.symbolics import bit_length, is_symbolic, Shaped, slen, SymbolicFloat, SymbolicInt if TYPE_CHECKING: + from qualtran import BloqBuilder, SoquetT from qualtran.resource_counting import BloqCountT, SympySymbolAllocator -@cirq.value_equality() +def _data_or_shape_to_tuple(data_or_shape: Union[NDArray, Shaped]) -> Tuple: + return ( + tuple(data_or_shape.flatten()) + if isinstance(data_or_shape, np.ndarray) + else (data_or_shape,) + ) + + @attrs.frozen class StatePreparationAliasSampling(PrepareOracle): r"""Initialize a state with $L$ coefficients using coherent alias sampling. @@ -110,8 +118,8 @@ class StatePreparationAliasSampling(PrepareOracle): selection_registers: Tuple[Register, ...] = attrs.field( converter=lambda v: (v,) if isinstance(v, Register) else tuple(v) ) - alt: Union[Shaped, NDArray[np.int_]] - keep: Union[Shaped, NDArray[np.int_]] + alt: Union[Shaped, NDArray[np.int_]] = attrs.field(eq=_data_or_shape_to_tuple) + keep: Union[Shaped, NDArray[np.int_]] = attrs.field(eq=_data_or_shape_to_tuple) mu: SymbolicInt sum_of_unnormalized_probabilities: SymbolicFloat @@ -219,14 +227,6 @@ def junk_registers(self) -> Tuple[Register, ...]: ) ) - def _value_equality_values_(self): - return ( - self.selection_registers, - self.alt if isinstance(self.alt, Shaped) else tuple(self.alt.ravel()), - self.keep if isinstance(self.keep, Shaped) else tuple(self.keep.ravel()), - self.mu, - ) - @cached_property def qrom_bloq(self) -> QROM: return QROM( @@ -287,3 +287,265 @@ def _state_prep_alias_symb() -> StatePreparationAliasSampling: import_line='from qualtran.bloqs.state_preparation import StatePreparationAliasSampling', examples=(_state_prep_alias, _state_prep_alias_symb), ) + + +@attrs.frozen +class SparseStatePreparationAliasSampling(PrepareOracle): + r"""Initialize a $d$-sparse state over $L$ indices using coherent alias sampling. + + In particular, we take the zero state to: + + $$ + \sum_{j=0}^{d-1} \sqrt{p_{\mathrm{ind}_j}} |\mathrm{ind}_j\rangle |\mathrm{temp}_j\rangle + $$ + + where $\mathrm{ind}_j \in [0, L)$ is the index of the $j$-th non-zero coefficient, + and the probabilities $p_l$ are $\mu$-bit binary approximations to the true values, + and the register $|\mathrm{temp}_j\rangle$ may be entangled with the index register. + + This bloq is nearly identical to :class:`StatePreparationByAliasSampling`, except + that it loads the non-zero indices from the QROM and prepares a dense state on them. + In comparison, this uses $\lceil \log d \rceil$ extra ancilla qubits, and reduces + the iteration length to $d$ from $L$. + + See :class:`StatePreparationAliasSampling` for an exposition on alias sampling. + + + Registers: + selection: The input/output register $|\mathrm{ind}_l\rangle$ of size lg(L) where the desired + coefficient state is prepared. + sigma_mu: A mu-sized register containing uniform probabilities for comparison against `keep`. + sparse_index: A lg(d)-sized register storing the sparse index $j \in [0, d)$. + alt: A lg(L)-sized register of alternate indices + keep: a mu-sized register of probabilities of keeping the initially sampled index. + less_than_equal: one bit for the result of the comparison. + + This gate corresponds to the following operations: + - UNIFORM_d on the `sparse_index` register. + - H^mu on the `sigma` register. + - QROM addressed by the `sparse_index` register into the `selection`, `alt`, and `keep` signature. + - LessThanEqualGate comparing the `keep` and `sigma` registers. + - Coherent swap between the `selection` and `alt` registers if the comparison returns True. + + Total space will be $(2 \log(L) + \log(d) + 2 \mu + 1)$ work qubits + $log(L)$ ancillas for QROM. + + References: + [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. + """ + selection_registers: Tuple[Register, ...] = attrs.field( + converter=lambda v: (v,) if isinstance(v, Register) else tuple(v) + ) + index: Union[Shaped, NDArray[np.int_]] = attrs.field(eq=_data_or_shape_to_tuple) + alt: Union[Shaped, NDArray[np.int_]] = attrs.field(eq=_data_or_shape_to_tuple) + keep: Union[Shaped, NDArray[np.int_]] = attrs.field(eq=_data_or_shape_to_tuple) + mu: SymbolicInt + sum_of_unnormalized_probabilities: SymbolicFloat + + def __attrs_post_init__(self): + if not is_symbolic(self.mu) and self.mu <= 0: + raise ValueError(f"{self.mu=} must be at least 1") + + @cached_property + def junk_registers(self) -> Tuple[Register, ...]: + return tuple( + Signature.build( + sigma_mu=self.mu, + sparse_index=self.sparse_index_bitsize, + alt=self.selection_bitsize, + keep=self.mu, + less_than_equal=1, + ) + ) + + @classmethod + def from_sparse_dict( + cls, unnormalized_probabilities: dict[int, float], N: int, *, precision: float = 1.0e-5 + ) -> 'SparseStatePreparationAliasSampling': + """Construct the state preparation gate for a given dictionary of non-zero probabilities. + + Args: + unnormalized_probabilities: a dictionary mapping indices to non-zero probabilities. + N: the maximum index (i.e. prepares a state with basis in [0, N)) + precision: The desired accuracy to represent each probability + (which sets mu size and keep/alt integers). + See `qualtran.linalg.lcu_util.preprocess_probabilities_for_reversible_sampling` + for more information. + """ + if not all(x >= 0 for x in unnormalized_probabilities.values()): + raise ValueError(f"{cls} expects only non-negative probabilities") + if not all(0 <= ix < N for ix in unnormalized_probabilities.keys()): + raise ValueError( + f"Sparse indices not in range [0, {N}): {unnormalized_probabilities.keys()}" + ) + + alt_compressed, keep, mu = preprocess_probabilities_for_reversible_sampling( + unnormalized_probabilities=list(unnormalized_probabilities.values()), epsilon=precision + ) + + index = list(unnormalized_probabilities.keys()) + alt = [index[idx] for idx in alt_compressed] + + return cls( + selection_registers=Register('selection', BoundedQUInt((N - 1).bit_length(), N)), + index=np.array(index), + alt=np.array(alt), + keep=np.array(keep), + mu=mu, + sum_of_unnormalized_probabilities=sum(unnormalized_probabilities.values()), + ) + + @classmethod + def from_dense_probabilities( + cls, + unnormalized_probabilities: Sequence[float], + *, + precision: float = 1.0e-5, + nonzero_threshold: float = 1e-6, + ) -> 'SparseStatePreparationAliasSampling': + """Factory to construct the state preparation gate for a given set of probability coefficients. + + Args: + unnormalized_probabilities: A dense list of all probabilities (i.e. including 0s) + precision: The desired accuracy to represent each probability + nonzero_threshold: minimum value for a probability entry to be considered non-zero. + """ + nonzero_value_map: dict[int, float] = { + ix: prob + for ix, prob in enumerate(unnormalized_probabilities) + if not np.isclose(prob, 0, atol=nonzero_threshold) + } + + return cls.from_sparse_dict( + nonzero_value_map, len(unnormalized_probabilities), precision=precision + ) + + @classmethod + def from_n_coeff( + cls, + n_coeff: SymbolicInt, + n_nonzero_coeff: SymbolicInt, + sum_of_terms: SymbolicFloat, + *, + precision: SymbolicFloat = 1.0e-5, + ) -> 'SparseStatePreparationAliasSampling': + """Factory to construct sparse state preparation for symbolic number of input probabilities. + + Args: + n_coeff: Symbolic number of LCU coefficients in the prepared state. + n_nonzero_coeff: Symbolic number of non-zero LCU coefficients in the prepared state. + sum_of_terms: Sum of absolute values of the input probabilities. + precision: The desired accuracy to represent each probability + (which sets mu size and keep/alt integers). + See `qualtran.linalg.lcu_util.preprocess_lcu_coefficients_for_reversible_sampling` + for more information. + """ + mu = sub_bit_prec_from_epsilon(n_coeff, precision) + selection_bitsize = bit_length(n_coeff - 1) + return cls( + selection_registers=Register('selection', BoundedQUInt(selection_bitsize, n_coeff)), + index=Shaped((n_nonzero_coeff,)), + alt=Shaped((n_nonzero_coeff,)), + keep=Shaped((n_nonzero_coeff,)), + mu=mu, + sum_of_unnormalized_probabilities=sum_of_terms, + ) + + @property + def n_coeff(self) -> SymbolicInt: + return self.selection_registers[0].dtype.iteration_length_or_zero() + + @property + def n_nonzero_coeff(self) -> SymbolicInt: + return slen(self.index) + + @cached_property + def l1_norm_of_coeffs(self) -> 'SymbolicFloat': + return self.sum_of_unnormalized_probabilities + + @cached_property + def selection_bitsize(self) -> SymbolicInt: + return total_bits(self.selection_registers) + + @cached_property + def sparse_index_bitsize(self) -> SymbolicInt: + return bit_length(self.n_nonzero_coeff - 1) + + @cached_property + def qrom_bloq(self) -> QROM: + return QROM( + (self.index, self.alt, self.keep), + (self.sparse_index_bitsize,), + (self.selection_bitsize, self.selection_bitsize, self.mu), + ) + + def build_composite_bloq(self, bb: 'BloqBuilder', **soqs: 'SoquetT') -> dict[str, 'SoquetT']: + soqs['sparse_index'] = bb.add( + PrepareUniformSuperposition(self.n_nonzero_coeff), target=soqs['sparse_index'] + ) + soqs['sigma_mu'] = bb.add(OnEach(self.mu, Hadamard()), q=soqs['sigma_mu']) + soqs['sparse_index'], soqs['selection'], soqs['alt'], soqs['keep'] = bb.add_t( + self.qrom_bloq, + selection=soqs['sparse_index'], + target0_=soqs['selection'], + target1_=soqs['alt'], + target2_=soqs['keep'], + ) + soqs['keep'], soqs['sigma_mu'], soqs['less_than_equal'] = bb.add_t( + LessThanEqual(self.mu, self.mu), + x=soqs['keep'], + y=soqs['sigma_mu'], + target=soqs['less_than_equal'], + ) + soqs['less_than_equal'], soqs['alt'], soqs['selection'] = bb.add_t( + CSwap(self.selection_bitsize), + ctrl=soqs['less_than_equal'], + x=soqs['alt'], + y=soqs['selection'], + ) + + return soqs + + +@bloq_example(generalizer=[cirq_to_bloqs, ignore_split_join, ignore_cliffords]) +def _sparse_state_prep_alias() -> SparseStatePreparationAliasSampling: + coeff_map = {0: 1.0, 3: 1.0, 5: 3.0, 7: 2.0} + N = 9 + mu = 3 + sparse_state_prep_alias = SparseStatePreparationAliasSampling.from_sparse_dict( + coeff_map, N, precision=2**-mu / len(coeff_map) + ) + return sparse_state_prep_alias + + +@bloq_example(generalizer=[cirq_to_bloqs, ignore_split_join, ignore_cliffords]) +def _sparse_state_prep_alias_from_list() -> SparseStatePreparationAliasSampling: + coeffs = [1.0, 0, 0, 1, 0, 3, 0, 2, 0] + mu = 3 + sparse_state_prep_alias_from_list = ( + SparseStatePreparationAliasSampling.from_dense_probabilities(coeffs, precision=2**-mu / 4) + ) + return sparse_state_prep_alias_from_list + + +@bloq_example(generalizer=[cirq_to_bloqs, ignore_split_join, ignore_cliffords]) +def _sparse_state_prep_alias_symb() -> SparseStatePreparationAliasSampling: + import sympy + + n_coeffs, n_nonzero_coeffs, sum_coeff, eps = sympy.symbols(r"L d \lambda \epsilon") + sparse_state_prep_alias_symb = SparseStatePreparationAliasSampling.from_n_coeff( + n_coeffs, n_nonzero_coeffs, sum_coeff, precision=eps + ) + return sparse_state_prep_alias_symb + + +_SPARSE_STATE_PREP_ALIAS_DOC = BloqDocSpec( + bloq_cls=SparseStatePreparationAliasSampling, + examples=( + _sparse_state_prep_alias, + _sparse_state_prep_alias_from_list, + _sparse_state_prep_alias_symb, + ), +) diff --git a/qualtran/bloqs/state_preparation/state_preparation_alias_sampling_test.py b/qualtran/bloqs/state_preparation/state_preparation_alias_sampling_test.py index 108e45cdf..e1cbdbb6e 100644 --- a/qualtran/bloqs/state_preparation/state_preparation_alias_sampling_test.py +++ b/qualtran/bloqs/state_preparation/state_preparation_alias_sampling_test.py @@ -17,13 +17,19 @@ import pytest import sympy +from qualtran import Bloq from qualtran.bloqs.chemistry.ising import get_1d_ising_lcu_coeffs from qualtran.bloqs.state_preparation.state_preparation_alias_sampling import ( + _sparse_state_prep_alias, + _sparse_state_prep_alias_from_list, + _sparse_state_prep_alias_symb, _state_prep_alias, _state_prep_alias_symb, + SparseStatePreparationAliasSampling, StatePreparationAliasSampling, ) from qualtran.cirq_interop.testing import GateHelper +from qualtran.symbolics import ceil, log2 from qualtran.testing import assert_valid_bloq_decomposition, execute_notebook @@ -31,6 +37,12 @@ def test_state_prep_alias_sampling_autotest(bloq_autotester): bloq_autotester(_state_prep_alias) +def test_sparse_state_prep_alias_sampling_autotest(bloq_autotester): + bloq_autotester(_sparse_state_prep_alias) + bloq_autotester(_sparse_state_prep_alias_from_list) + bloq_autotester(_sparse_state_prep_alias_symb) + + def test_mu_from_precision(): coeffs = [1.0, 1, 3, 2] mu = 3 @@ -79,11 +91,21 @@ def test_state_prep_alias_sampling_symb(): np.testing.assert_allclose(concrete_t_counts, symb_t_counts, rtol=5e-4) -def assert_state_preparation_valid_for_coefficient(lcu_coefficients: np.ndarray, epsilon: float): - gate = StatePreparationAliasSampling.from_probabilities( - unnormalized_probabilities=lcu_coefficients.tolist(), - precision=epsilon * np.sum(np.abs(lcu_coefficients)), - ) +def assert_state_preparation_valid_for_coefficient( + lcu_coefficients: np.ndarray, epsilon: float, *, sparse: bool = False, atol: float = 1e-6 +): + gate: Bloq + coeff_precision = epsilon * np.sum(np.abs(lcu_coefficients)) + if sparse: + gate = SparseStatePreparationAliasSampling.from_dense_probabilities( + unnormalized_probabilities=lcu_coefficients.tolist(), + precision=coeff_precision, + nonzero_threshold=atol, + ) + else: + gate = StatePreparationAliasSampling.from_probabilities( + unnormalized_probabilities=lcu_coefficients.tolist(), precision=coeff_precision + ) assert_valid_bloq_decomposition(gate) _ = gate.call_graph() @@ -100,14 +122,14 @@ def assert_state_preparation_valid_for_coefficient(lcu_coefficients: np.ndarray, L, logL = len(lcu_coefficients), len(g.quregs['selection']) qlambda = sum(abs(lcu_coefficients)) state_vector = state_vector.reshape(2**logL, len(state_vector) // 2**logL) - num_non_zero = (abs(state_vector) > 1e-6).sum(axis=1) - prepared_state = state_vector.sum(axis=1) - assert all(num_non_zero[:L] > 0) and all(num_non_zero[L:] == 0) - assert all(np.abs(prepared_state[:L]) > 1e-6) and all(np.abs(prepared_state[L:]) <= 1e-6) - prepared_state = prepared_state[:L] / np.sqrt(num_non_zero[:L]) + prepared_state = np.linalg.norm(state_vector, axis=1) + # Assert that the absolute square of prepared state (probabilities instead of amplitudes) is # same as `lcu_coefficients` upto `epsilon`. - np.testing.assert_allclose(lcu_coefficients / qlambda, abs(prepared_state) ** 2, atol=epsilon) + np.testing.assert_allclose( + abs(prepared_state[:L]) ** 2, lcu_coefficients / qlambda, atol=epsilon + ) + np.testing.assert_allclose(np.abs(prepared_state[L:]), 0, atol=epsilon) def test_state_preparation_via_coherent_alias_sampling_quick(): @@ -166,6 +188,29 @@ def test_state_preparation_via_coherent_alias_sampling_diagram(): ) +def test_sparse_state_preparation_via_coherent_alias(): + lcu_coefficients = np.array([1 / 8 if j < 8 else 0.0 for j in range(16)]) + assert_state_preparation_valid_for_coefficient(lcu_coefficients, 2e-1, sparse=True) + + lcu_coefficients = np.array([1 if j < 6 else 0.0 for j in range(10)]) + assert_state_preparation_valid_for_coefficient(lcu_coefficients, 2e-1, sparse=True) + + +def test_symbolic_sparse_state_prep_t_complexity(): + from qualtran.cirq_interop.t_complexity_protocol import TComplexity + + N, d, qlambda, eps = sympy.symbols(r"N d \lambda \epsilon") + logN = ceil(log2(N - 1)) + logd = ceil(log2(d - 1)) + mu = ceil(log2(1 / (N * eps))) + bloq = SparseStatePreparationAliasSampling.from_n_coeff(N, d, qlambda, precision=eps) + assert bloq.t_complexity() == TComplexity( + t=4 * d + 8 * mu + 7 * logN + 12 * logd - 8, + clifford=d * mu * logN**2 + 13 * d + 47 * mu + 10 * logN + 52 * logd - 12, + rotations=2, + ) + + @pytest.mark.notebook def test_notebook(): execute_notebook('state_preparation_alias_sampling') diff --git a/qualtran/conftest.py b/qualtran/conftest.py index 7c5fcd56c..f631b83b9 100644 --- a/qualtran/conftest.py +++ b/qualtran/conftest.py @@ -103,6 +103,7 @@ def assert_bloq_example_serializes_for_pytest(bloq_ex: BloqExample): 'apply_lth_bloq', 'linear_combination_block_encoding', 'phase_block_encoding', + 'sparse_state_prep_alias_symb', # cannot serialize Shaped 'sparse_permutation', 'permutation_cycle_symb', ]: diff --git a/qualtran/serialization/resolver_dict.py b/qualtran/serialization/resolver_dict.py index 41c65c744..2107d4baa 100644 --- a/qualtran/serialization/resolver_dict.py +++ b/qualtran/serialization/resolver_dict.py @@ -345,6 +345,7 @@ "qualtran.bloqs.block_encoding.lcu_select_and_prepare.SelectOracle": qualtran.bloqs.block_encoding.lcu_select_and_prepare.SelectOracle, "qualtran.bloqs.state_preparation.prepare_uniform_superposition.PrepareUniformSuperposition": qualtran.bloqs.state_preparation.prepare_uniform_superposition.PrepareUniformSuperposition, "qualtran.bloqs.state_preparation.state_preparation_alias_sampling.StatePreparationAliasSampling": qualtran.bloqs.state_preparation.state_preparation_alias_sampling.StatePreparationAliasSampling, + "qualtran.bloqs.state_preparation.state_preparation_alias_sampling.SparseStatePreparationAliasSampling": qualtran.bloqs.state_preparation.state_preparation_alias_sampling.SparseStatePreparationAliasSampling, "qualtran.bloqs.state_preparation.state_preparation_via_rotation.PRGAViaPhaseGradient": qualtran.bloqs.state_preparation.state_preparation_via_rotation.PRGAViaPhaseGradient, "qualtran.bloqs.state_preparation.state_preparation_via_rotation.StatePreparationViaRotations": qualtran.bloqs.state_preparation.state_preparation_via_rotation.StatePreparationViaRotations, "qualtran.bloqs.swap_network.cswap_approx.CSwapApprox": qualtran.bloqs.swap_network.cswap_approx.CSwapApprox,