From bc5aa0402e206ca79f868f2603b0d3745bfade96 Mon Sep 17 00:00:00 2001 From: Charles Yuan Date: Tue, 16 Jul 2024 19:44:12 -0700 Subject: [PATCH 1/3] Remove unnecessary `Union` (#1146) --- .../bloqs/chemistry/trotter/hubbard/hopping.py | 12 ++++++------ .../chemistry/trotter/hubbard/interaction.py | 18 +++++++++--------- .../chemistry/trotter/trotterized_unitary.py | 6 +++--- .../bloqs/rotations/hamming_weight_phasing.py | 4 ++-- .../rotations/quantum_variable_rotation.py | 14 +++++++------- 5 files changed, 27 insertions(+), 27 deletions(-) diff --git a/qualtran/bloqs/chemistry/trotter/hubbard/hopping.py b/qualtran/bloqs/chemistry/trotter/hubbard/hopping.py index b54672b4a..8db597c8d 100644 --- a/qualtran/bloqs/chemistry/trotter/hubbard/hopping.py +++ b/qualtran/bloqs/chemistry/trotter/hubbard/hopping.py @@ -13,7 +13,7 @@ # limitations under the License. """Bloqs implementing unitary evolution under the one-body hopping Hamiltonian in 2D.""" from functools import cached_property -from typing import Set, TYPE_CHECKING, Union +from typing import Set, TYPE_CHECKING from attrs import frozen @@ -67,8 +67,8 @@ class HoppingPlaquette(Bloq): page 13 Eq. E4 and E5 (Appendix E) """ - kappa: Union[SymbolicFloat] - eps: Union[SymbolicFloat] = 1e-9 + kappa: SymbolicFloat + eps: SymbolicFloat = 1e-9 @cached_property def signature(self) -> Signature: @@ -110,10 +110,10 @@ class HoppingTile(Bloq): see Eq. 21 and App E. """ - length: Union[SymbolicInt] - angle: Union[SymbolicFloat] + length: SymbolicInt + angle: SymbolicFloat tau: float = 1.0 - eps: Union[SymbolicFloat] = 1e-9 + eps: SymbolicFloat = 1e-9 pink: bool = True def __attrs_post_init__(self): diff --git a/qualtran/bloqs/chemistry/trotter/hubbard/interaction.py b/qualtran/bloqs/chemistry/trotter/hubbard/interaction.py index 3f33f6b60..a910cedcd 100644 --- a/qualtran/bloqs/chemistry/trotter/hubbard/interaction.py +++ b/qualtran/bloqs/chemistry/trotter/hubbard/interaction.py @@ -14,7 +14,7 @@ r"""Bloqs implementing unitary evolution under the interacting part of the Hubbard Hamiltonian.""" from functools import cached_property -from typing import Set, TYPE_CHECKING, Union +from typing import Set, TYPE_CHECKING from attrs import frozen @@ -53,10 +53,10 @@ class Interaction(Bloq): Eq. 6 page 2 and page 13 paragraph 1. """ - length: Union[SymbolicInt] - angle: Union[SymbolicFloat] - hubb_u: Union[SymbolicFloat] - eps: Union[SymbolicFloat] = 1e-9 + length: SymbolicInt + angle: SymbolicFloat + hubb_u: SymbolicFloat + eps: SymbolicFloat = 1e-9 @cached_property def signature(self) -> Signature: @@ -96,10 +96,10 @@ class InteractionHWP(Bloq): 14 paragraph 3 right column. The apply 2 batches of $L^2/2$ rotations. """ - length: Union[SymbolicInt] - angle: Union[SymbolicFloat] - hubb_u: Union[SymbolicFloat] - eps: Union[SymbolicFloat] = 1e-9 + length: SymbolicInt + angle: SymbolicFloat + hubb_u: SymbolicFloat + eps: SymbolicFloat = 1e-9 @cached_property def signature(self) -> Signature: diff --git a/qualtran/bloqs/chemistry/trotter/trotterized_unitary.py b/qualtran/bloqs/chemistry/trotter/trotterized_unitary.py index 78624a29a..91dd0736a 100644 --- a/qualtran/bloqs/chemistry/trotter/trotterized_unitary.py +++ b/qualtran/bloqs/chemistry/trotter/trotterized_unitary.py @@ -14,7 +14,7 @@ """Bloq for building a Trotterized unitary""" from functools import cached_property -from typing import Dict, Sequence, Union +from typing import Dict, Sequence import attrs @@ -83,8 +83,8 @@ class TrotterizedUnitary(Bloq): bloqs: Sequence[Bloq] indices: Sequence[int] - coeffs: Sequence[Union[SymbolicFloat]] - timestep: Union[SymbolicFloat] + coeffs: Sequence[SymbolicFloat] + timestep: SymbolicFloat def __attrs_post_init__(self): ref_sig = self.bloqs[0].signature diff --git a/qualtran/bloqs/rotations/hamming_weight_phasing.py b/qualtran/bloqs/rotations/hamming_weight_phasing.py index 84b863389..bd59b3ecc 100644 --- a/qualtran/bloqs/rotations/hamming_weight_phasing.py +++ b/qualtran/bloqs/rotations/hamming_weight_phasing.py @@ -13,7 +13,7 @@ # limitations under the License. from functools import cached_property -from typing import Dict, Set, TYPE_CHECKING, Union +from typing import Dict, Set, TYPE_CHECKING import attrs import numpy as np @@ -73,7 +73,7 @@ class HammingWeightPhasing(GateWithRegisters): bitsize: int exponent: float = 1 - eps: Union[SymbolicFloat] = 1e-10 + eps: SymbolicFloat = 1e-10 @cached_property def signature(self) -> 'Signature': diff --git a/qualtran/bloqs/rotations/quantum_variable_rotation.py b/qualtran/bloqs/rotations/quantum_variable_rotation.py index 5fb3dc24d..1559b7e94 100644 --- a/qualtran/bloqs/rotations/quantum_variable_rotation.py +++ b/qualtran/bloqs/rotations/quantum_variable_rotation.py @@ -56,7 +56,7 @@ import abc from functools import cached_property -from typing import cast, Dict, Sequence, Set, TYPE_CHECKING, Union +from typing import cast, Dict, Sequence, Set, TYPE_CHECKING import attrs import numpy as np @@ -146,12 +146,12 @@ class QvrZPow(QvrInterface): """ cost_reg: Register - gamma: Union[SymbolicFloat] = 1.0 - eps: Union[SymbolicFloat] = 1e-9 + gamma: SymbolicFloat = 1.0 + eps: SymbolicFloat = 1e-9 @classmethod def from_bitsize( - cls, bitsize: int, gamma: Union[SymbolicFloat] = 1.0, eps: Union[SymbolicFloat] = 1e-9 + cls, bitsize: int, gamma: SymbolicFloat = 1.0, eps: SymbolicFloat = 1e-9 ) -> 'QvrZPow': cost_reg = Register("x", QFxp(bitsize, bitsize, signed=False)) return QvrZPow(cost_reg, gamma=gamma, eps=eps) @@ -379,8 +379,8 @@ class QvrPhaseGradient(QvrInterface): """ cost_reg: Register - gamma: Union[SymbolicFloat] = 1.0 - eps: Union[SymbolicFloat] = 1e-9 + gamma: SymbolicFloat = 1.0 + eps: SymbolicFloat = 1e-9 def __attrs_post_init__(self): dtype = self.cost_reg.dtype @@ -389,7 +389,7 @@ def __attrs_post_init__(self): @classmethod def from_bitsize( - cls, bitsize: int, gamma: Union[SymbolicFloat] = 1.0, eps: Union[SymbolicFloat] = 1e-9 + cls, bitsize: int, gamma: SymbolicFloat = 1.0, eps: SymbolicFloat = 1e-9 ) -> 'QvrPhaseGradient': cost_reg = Register("x", QFxp(bitsize, bitsize, signed=False)) return QvrPhaseGradient(cost_reg, gamma=gamma, eps=eps) From d807b7299868a7969994390d3d2514b66ce3663b Mon Sep 17 00:00:00 2001 From: Anurudh Peduri <7265746+anurudhp@users.noreply.github.com> Date: Wed, 17 Jul 2024 06:37:36 -0700 Subject: [PATCH 2/3] Permutation Bloq (#1110) * Permutation Bloq * address comments: use `x` instead of `q` invert code in `build_call_graph` add note for neg-ctrl CNOT * move `XorK` to `bitwise.py` * tests * XorK bloq examples * undo trivial changes * mypy * cleanup EqualsAConstant * fix decomp * use dictionaries for sparse/partial permutations, simplify code * fix bug * regen notebooks * docstring * add notebook * docstring unitary --- dev_tools/autogenerate-bloqs-notebooks-v2.py | 9 + docs/bloqs/index.rst | 1 + qualtran/bloqs/arithmetic/bitwise.py | 120 ++++++ qualtran/bloqs/arithmetic/bitwise_test.py | 38 ++ qualtran/bloqs/arithmetic/comparison.ipynb | 2 +- qualtran/bloqs/arithmetic/comparison.py | 42 ++- qualtran/bloqs/arithmetic/comparison_test.py | 4 +- qualtran/bloqs/arithmetic/permutation.ipynb | 344 ++++++++++++++++++ qualtran/bloqs/arithmetic/permutation.py | 316 ++++++++++++++++ qualtran/bloqs/arithmetic/permutation_test.py | 138 +++++++ qualtran/bloqs/mcmt/and_bloq.py | 2 +- qualtran/conftest.py | 2 + qualtran/linalg/permutation.py | 58 +++ qualtran/linalg/permutation_test.py | 28 ++ qualtran/serialization/resolver_dict.py | 5 + 15 files changed, 1097 insertions(+), 12 deletions(-) create mode 100644 qualtran/bloqs/arithmetic/bitwise.py create mode 100644 qualtran/bloqs/arithmetic/bitwise_test.py create mode 100644 qualtran/bloqs/arithmetic/permutation.ipynb create mode 100644 qualtran/bloqs/arithmetic/permutation.py create mode 100644 qualtran/bloqs/arithmetic/permutation_test.py create mode 100644 qualtran/linalg/permutation.py create mode 100644 qualtran/linalg/permutation_test.py diff --git a/dev_tools/autogenerate-bloqs-notebooks-v2.py b/dev_tools/autogenerate-bloqs-notebooks-v2.py index 778409140..277adc13e 100644 --- a/dev_tools/autogenerate-bloqs-notebooks-v2.py +++ b/dev_tools/autogenerate-bloqs-notebooks-v2.py @@ -53,6 +53,7 @@ from qualtran_dev_tools.jupyter_autogen import NotebookSpecV2, render_notebook import qualtran.bloqs.arithmetic.addition +import qualtran.bloqs.arithmetic.permutation import qualtran.bloqs.arithmetic.sorting import qualtran.bloqs.arithmetic.subtraction import qualtran.bloqs.basic_gates.swap @@ -380,6 +381,14 @@ qualtran.bloqs.arithmetic.conversions._TO_CONTG_INDX, ], ), + NotebookSpecV2( + title='Permutations', + module=qualtran.bloqs.arithmetic.permutation, + bloq_specs=[ + qualtran.bloqs.arithmetic.permutation._PERMUTATION_DOC, + qualtran.bloqs.arithmetic.permutation._PERMUTATION_CYCLE_DOC, + ], + ), ] MOD_ARITHMETIC = [ diff --git a/docs/bloqs/index.rst b/docs/bloqs/index.rst index a039f2ae2..c7b1cc302 100644 --- a/docs/bloqs/index.rst +++ b/docs/bloqs/index.rst @@ -64,6 +64,7 @@ Bloqs Library arithmetic/comparison.ipynb arithmetic/sorting.ipynb arithmetic/conversions.ipynb + arithmetic/permutation.ipynb .. toctree:: :maxdepth: 2 diff --git a/qualtran/bloqs/arithmetic/bitwise.py b/qualtran/bloqs/arithmetic/bitwise.py new file mode 100644 index 000000000..3f80b0155 --- /dev/null +++ b/qualtran/bloqs/arithmetic/bitwise.py @@ -0,0 +1,120 @@ +# 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. +from functools import cached_property +from typing import cast, Dict, Optional, Set, Tuple, TYPE_CHECKING + +import numpy as np +from attrs import frozen + +from qualtran import ( + bloq_example, + BloqBuilder, + DecomposeTypeError, + QBit, + QDType, + QUInt, + Register, + Signature, + Soquet, + SoquetT, +) +from qualtran._infra.gate_with_registers import SpecializedSingleQubitControlledGate +from qualtran.bloqs.basic_gates import CNOT, XGate +from qualtran.resource_counting.generalizers import ignore_split_join +from qualtran.symbolics import is_symbolic, SymbolicInt + +if TYPE_CHECKING: + from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + + +def _cvs_converter(vv): + if isinstance(vv, (int, np.integer)): + return (int(vv),) + return tuple(int(v) for v in vv) + + +@frozen +class XorK(SpecializedSingleQubitControlledGate): + r"""Maps |x> to |x \oplus k> for a constant k. + + Args: + dtype: Data type of the input register `x`. + k: The classical integer value to be XOR-ed to x. + control_val: an optional single bit control, apply the operation when + the control qubit equals the `control_val`. + + Registers: + x: A quantum register of type `self.dtype` (see above). + ctrl: A sequence of control qubits (only when `control_val` is not None). + """ + dtype: QDType + k: SymbolicInt + control_val: Optional[int] = None + + @cached_property + def signature(self) -> 'Signature': + return Signature([*self.control_registers, Register('x', self.dtype)]) + + @cached_property + def control_registers(self) -> Tuple[Register, ...]: + if self.control_val is not None: + return (Register('ctrl', QBit()),) + return () + + @cached_property + def bitsize(self) -> SymbolicInt: + return self.dtype.num_qubits + + def is_symbolic(self): + return is_symbolic(self.k, self.dtype) + + def build_composite_bloq(self, bb: 'BloqBuilder', **soqs: 'SoquetT') -> Dict[str, 'SoquetT']: + if self.is_symbolic(): + raise DecomposeTypeError(f"cannot decompose symbolic {self}") + + # TODO clean this up once https://github.com/quantumlib/Qualtran/pull/1137 is merged + ctrl = soqs.pop('ctrl', None) + + xs = bb.split(cast(Soquet, soqs.pop('x'))) + + for i, bit in enumerate(self.dtype.to_bits(self.k)): + if bit == 1: + if ctrl is not None: + ctrl, xs[i] = bb.add(CNOT(), ctrl=ctrl, target=xs[i]) + else: + xs[i] = bb.add(XGate(), q=xs[i]) + + soqs['x'] = bb.join(xs) + + if ctrl is not None: + soqs['ctrl'] = ctrl + return soqs + + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + bit_flip_bloq = CNOT() if self.control_val is not None else XGate() + num_flips = self.bitsize if self.is_symbolic() else sum(self.dtype.to_bits(self.k)) + return {(bit_flip_bloq, num_flips)} + + +@bloq_example(generalizer=ignore_split_join) +def _xork() -> XorK: + xork = XorK(QUInt(8), 0b01010111) + return xork + + +@bloq_example(generalizer=ignore_split_join) +def _cxork() -> XorK: + cxork = XorK(QUInt(8), 0b01010111).controlled() + assert isinstance(cxork, XorK) + return cxork diff --git a/qualtran/bloqs/arithmetic/bitwise_test.py b/qualtran/bloqs/arithmetic/bitwise_test.py new file mode 100644 index 000000000..6dbd02473 --- /dev/null +++ b/qualtran/bloqs/arithmetic/bitwise_test.py @@ -0,0 +1,38 @@ +# 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. +from qualtran import QUInt +from qualtran.bloqs.arithmetic.bitwise import _cxork, _xork, XorK + + +def test_examples(bloq_autotester): + bloq_autotester(_cxork) + bloq_autotester(_xork) + + +def test_xork_classical_sim(): + k = 0b01101010 + bloq = XorK(QUInt(9), k) + cbloq = bloq.controlled() + + for x in bloq.dtype.get_classical_domain(): + (x_out,) = bloq.call_classically(x=x) + assert x_out == x ^ k + + ctrl_out, x_out = cbloq.call_classically(ctrl=0, x=x) + assert ctrl_out == 0 + assert x_out == x + + ctrl_out, x_out = cbloq.call_classically(ctrl=1, x=x) + assert ctrl_out == 1 + assert x_out == x ^ k diff --git a/qualtran/bloqs/arithmetic/comparison.ipynb b/qualtran/bloqs/arithmetic/comparison.ipynb index ecdf62c7a..d8dfa2bb9 100644 --- a/qualtran/bloqs/arithmetic/comparison.ipynb +++ b/qualtran/bloqs/arithmetic/comparison.ipynb @@ -256,7 +256,7 @@ }, "source": [ "## `EqualsAConstant`\n", - "Implements $U_a|x\\rangle = U_a|x\\rangle|z\\rangle = |x\\rangle |z \\land (x = a)\\rangle$\n", + "Implements $U_a|x\\rangle|z\\rangle = |x\\rangle |z \\oplus (x = a)\\rangle$\n", "\n", "The bloq_counts and t_complexity are derived from:\n", "https://qualtran.readthedocs.io/en/latest/bloqs/comparison_gates.html#equality-as-a-special-case\n", diff --git a/qualtran/bloqs/arithmetic/comparison.py b/qualtran/bloqs/arithmetic/comparison.py index 5cd7599e5..0408d0fac 100644 --- a/qualtran/bloqs/arithmetic/comparison.py +++ b/qualtran/bloqs/arithmetic/comparison.py @@ -38,6 +38,7 @@ Bloq, bloq_example, BloqDocSpec, + DecomposeTypeError, GateWithRegisters, QAny, QBit, @@ -48,12 +49,12 @@ Soquet, SoquetT, ) -from qualtran.bloqs.basic_gates import CNOT, TGate, XGate +from qualtran.bloqs.basic_gates import CNOT, XGate from qualtran.bloqs.mcmt.and_bloq import And, MultiAnd -from qualtran.bloqs.mcmt.multi_control_multi_target_pauli import MultiControlX +from qualtran.bloqs.mcmt.multi_control_multi_target_pauli import MultiControlPauli, MultiControlX from qualtran.drawing import WireSymbol from qualtran.drawing.musical_score import Text, TextBox -from qualtran.symbolics import is_symbolic, SymbolicInt +from qualtran.symbolics import HasLength, is_symbolic, SymbolicInt if TYPE_CHECKING: from qualtran import BloqBuilder @@ -926,7 +927,7 @@ def _gt_k() -> GreaterThanConstant: @frozen class EqualsAConstant(Bloq): - r"""Implements $U_a|x\rangle = U_a|x\rangle|z\rangle = |x\rangle |z \land (x = a)\rangle$ + r"""Implements $U_a|x\rangle|z\rangle = |x\rangle |z \oplus (x = a)\rangle$ The bloq_counts and t_complexity are derived from: https://qualtran.readthedocs.io/en/latest/bloqs/comparison_gates.html#equality-as-a-special-case @@ -940,8 +941,8 @@ class EqualsAConstant(Bloq): target: Register to hold result of comparison. """ - bitsize: int - val: int + bitsize: SymbolicInt + val: SymbolicInt @cached_property def signature(self) -> Signature: @@ -956,10 +957,33 @@ def wire_symbol(self, reg: Optional[Register], idx: Tuple[int, ...] = tuple()) - return TextBox(f"⨁(x = {self.val})") raise ValueError(f'Unknown register symbol {reg.name}') + def is_symbolic(self): + return is_symbolic(self.bitsize, self.val) + + @property + def bits_k(self) -> Union[tuple[int, ...], HasLength]: + if self.is_symbolic(): + return HasLength(self.bitsize) + + assert not isinstance(self.bitsize, sympy.Expr) + assert not isinstance(self.val, sympy.Expr) + return tuple(QUInt(self.bitsize).to_bits(self.val)) + + def build_composite_bloq( + self, bb: 'BloqBuilder', x: 'Soquet', target: 'Soquet' + ) -> Dict[str, 'SoquetT']: + if is_symbolic(self.bitsize): + raise DecomposeTypeError(f"Cannot decompose {self} with symbolic {self.bitsize=}") + + xs = bb.split(x) + xs, target = bb.add( + MultiControlPauli(self.bits_k, target_gate=cirq.X), controls=xs, target=target + ) + x = bb.join(xs) + return {'x': x, 'target': target} + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: - # See: https://github.com/quantumlib/Qualtran/issues/219 - # See: https://github.com/quantumlib/Qualtran/issues/217 - return {(TGate(), 4 * (self.bitsize - 1))} + return {(MultiControlPauli(self.bits_k, target_gate=cirq.X), 1)} def _make_equals_a_constant(): diff --git a/qualtran/bloqs/arithmetic/comparison_test.py b/qualtran/bloqs/arithmetic/comparison_test.py index c73bf0930..fa0076ccd 100644 --- a/qualtran/bloqs/arithmetic/comparison_test.py +++ b/qualtran/bloqs/arithmetic/comparison_test.py @@ -298,7 +298,9 @@ def test_equals_a_constant(): qlt_testing.assert_wire_symbols_match_expected( EqualsAConstant(bitsize, 17), ['In(x)', '⨁(x = 17)'] ) - assert t_complexity(EqualsAConstant(bitsize, 17)) == TComplexity(t=4 * (bitsize - 1)) + assert t_complexity(EqualsAConstant(bitsize, 17)) == TComplexity( + t=4 * (bitsize - 1), clifford=65 + ) @pytest.mark.notebook diff --git a/qualtran/bloqs/arithmetic/permutation.ipynb b/qualtran/bloqs/arithmetic/permutation.ipynb new file mode 100644 index 000000000..51b510d88 --- /dev/null +++ b/qualtran/bloqs/arithmetic/permutation.ipynb @@ -0,0 +1,344 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "a48878df", + "metadata": { + "cq.autogen": "title_cell" + }, + "source": [ + "# Permutations" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d083d346", + "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": "56539662", + "metadata": { + "cq.autogen": "Permutation.bloq_doc.md" + }, + "source": [ + "## `Permutation`\n", + "Apply a permutation of [0, N - 1] on the basis states.\n", + "\n", + "Given a permutation $P : [0, N - 1] \\to [0, N - 1]$, this bloq applies the unitary:\n", + "\n", + "$$\n", + " U|x\\rangle = |P(x)\\rangle\n", + "$$\n", + "\n", + "Decomposes a permutation into cycles and applies them in order.\n", + "See :meth:`from_dense_permutation` to construct this bloq from a permutation,\n", + "and :meth:`from_partial_permutation_map` to construct it from a mapping.\n", + "\n", + "#### Parameters\n", + " - `N`: the total size the permutation acts on.\n", + " - `cycles`: a sequence of permutation cycles that form the permutation. \n", + "\n", + "#### Registers\n", + " - `x`: integer register storing a value in [0, ..., N - 1] \n", + "\n", + "#### References\n", + " - [A simple quantum algorithm to efficiently prepare sparse states](https://arxiv.org/abs/2310.19309v1). Appendix B.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "df73124f", + "metadata": { + "cq.autogen": "Permutation.bloq_doc.py" + }, + "outputs": [], + "source": [ + "from qualtran.bloqs.arithmetic.permutation import Permutation" + ] + }, + { + "cell_type": "markdown", + "id": "31147b0d", + "metadata": { + "cq.autogen": "Permutation.example_instances.md" + }, + "source": [ + "### Example Instances" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ec6460b6", + "metadata": { + "cq.autogen": "Permutation.permutation" + }, + "outputs": [], + "source": [ + "permutation = Permutation.from_dense_permutation([1, 3, 0, 2])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0c8e04f0", + "metadata": { + "cq.autogen": "Permutation.permutation_symb" + }, + "outputs": [], + "source": [ + "import sympy\n", + "\n", + "from qualtran.symbolics import Shaped\n", + "\n", + "N, k = sympy.symbols(\"N k\")\n", + "permutation_symb = Permutation(N, Shaped((k,)))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8578d036", + "metadata": { + "cq.autogen": "Permutation.permutation_symb_with_cycles" + }, + "outputs": [], + "source": [ + "import sympy\n", + "\n", + "from qualtran.symbolics import Shaped\n", + "\n", + "N = sympy.symbols(\"N\")\n", + "n_cycles = 4\n", + "d = sympy.IndexedBase('d', shape=(n_cycles,))\n", + "permutation_symb_with_cycles = Permutation(N, tuple(Shaped((d[i],)) for i in range(n_cycles)))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "84b5606a", + "metadata": { + "cq.autogen": "Permutation.sparse_permutation" + }, + "outputs": [], + "source": [ + "sparse_permutation = Permutation.from_partial_permutation_map(\n", + " 16, {0: 1, 1: 3, 2: 8, 3: 15, 4: 12}\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "369df9cb", + "metadata": { + "cq.autogen": "Permutation.graphical_signature.md" + }, + "source": [ + "#### Graphical Signature" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "25b2dc2b", + "metadata": { + "cq.autogen": "Permutation.graphical_signature.py" + }, + "outputs": [], + "source": [ + "from qualtran.drawing import show_bloqs\n", + "show_bloqs([permutation, permutation_symb, permutation_symb_with_cycles, sparse_permutation],\n", + " ['`permutation`', '`permutation_symb`', '`permutation_symb_with_cycles`', '`sparse_permutation`'])" + ] + }, + { + "cell_type": "markdown", + "id": "ad4321ad", + "metadata": { + "cq.autogen": "Permutation.call_graph.md" + }, + "source": [ + "### Call Graph" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3bfbf3f7", + "metadata": { + "cq.autogen": "Permutation.call_graph.py" + }, + "outputs": [], + "source": [ + "from qualtran.resource_counting.generalizers import ignore_split_join\n", + "permutation_g, permutation_sigma = permutation.call_graph(max_depth=1, generalizer=ignore_split_join)\n", + "show_call_graph(permutation_g)\n", + "show_counts_sigma(permutation_sigma)" + ] + }, + { + "cell_type": "markdown", + "id": "b3e895db", + "metadata": { + "cq.autogen": "PermutationCycle.bloq_doc.md" + }, + "source": [ + "## `PermutationCycle`\n", + "Apply a single permutation cycle on the basis states.\n", + "\n", + "Given a permutation cycle $C = (v_0 v_2 \\ldots v_{k - 1})$, applies the following unitary:\n", + "\n", + " $$\n", + " U|v_i\\rangle \\mapsto |v_{(i + 1)\\mod k}\\rangle\n", + " $$\n", + "\n", + "for each $i \\in [0, k)$, and\n", + "\n", + " $$\n", + " U|x\\rangle \\mapsto |x\\rangle\n", + " $$\n", + "\n", + "and for every $x \\not\\in C$.\n", + "\n", + "#### Parameters\n", + " - `N`: the total size the permutation acts on.\n", + " - `cycle`: the permutation cycle to apply. \n", + "\n", + "#### Registers\n", + " - `x`: integer register storing a value in [0, ..., N - 1] \n", + "\n", + "#### References\n", + " - [A simple quantum algorithm to efficiently prepare sparse states](https://arxiv.org/abs/2310.19309v1). Appendix B, Algorithm 7.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d569cd2c", + "metadata": { + "cq.autogen": "PermutationCycle.bloq_doc.py" + }, + "outputs": [], + "source": [ + "from qualtran.bloqs.arithmetic.permutation import PermutationCycle" + ] + }, + { + "cell_type": "markdown", + "id": "a93c6e89", + "metadata": { + "cq.autogen": "PermutationCycle.example_instances.md" + }, + "source": [ + "### Example Instances" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "264ba946", + "metadata": { + "cq.autogen": "PermutationCycle.permutation_cycle" + }, + "outputs": [], + "source": [ + "permutation_cycle = PermutationCycle(4, (0, 1, 2))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "78572528", + "metadata": { + "cq.autogen": "PermutationCycle.permutation_cycle_symb" + }, + "outputs": [], + "source": [ + "import sympy\n", + "\n", + "from qualtran.symbolics import Shaped\n", + "\n", + "N, L = sympy.symbols(\"N L\")\n", + "cycle = Shaped((L,))\n", + "permutation_cycle_symb = PermutationCycle(N, cycle)" + ] + }, + { + "cell_type": "markdown", + "id": "87615783", + "metadata": { + "cq.autogen": "PermutationCycle.graphical_signature.md" + }, + "source": [ + "#### Graphical Signature" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f6961010", + "metadata": { + "cq.autogen": "PermutationCycle.graphical_signature.py" + }, + "outputs": [], + "source": [ + "from qualtran.drawing import show_bloqs\n", + "show_bloqs([permutation_cycle, permutation_cycle_symb],\n", + " ['`permutation_cycle`', '`permutation_cycle_symb`'])" + ] + }, + { + "cell_type": "markdown", + "id": "16e1c105", + "metadata": { + "cq.autogen": "PermutationCycle.call_graph.md" + }, + "source": [ + "### Call Graph" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "27628f10", + "metadata": { + "cq.autogen": "PermutationCycle.call_graph.py" + }, + "outputs": [], + "source": [ + "from qualtran.resource_counting.generalizers import ignore_split_join\n", + "permutation_cycle_g, permutation_cycle_sigma = permutation_cycle.call_graph(max_depth=1, generalizer=ignore_split_join)\n", + "show_call_graph(permutation_cycle_g)\n", + "show_counts_sigma(permutation_cycle_sigma)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/qualtran/bloqs/arithmetic/permutation.py b/qualtran/bloqs/arithmetic/permutation.py new file mode 100644 index 000000000..8738fab1c --- /dev/null +++ b/qualtran/bloqs/arithmetic/permutation.py @@ -0,0 +1,316 @@ +# 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. +# +# 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. +from functools import cached_property +from typing import cast, Sequence, Set, TYPE_CHECKING, TypeAlias, Union + +from attrs import field, frozen + +from qualtran import ( + Bloq, + bloq_example, + BloqDocSpec, + BoundedQUInt, + DecomposeTypeError, + QBit, + QUInt, + Signature, + Soquet, +) +from qualtran.bloqs.arithmetic.bitwise import XorK +from qualtran.bloqs.arithmetic.comparison import EqualsAConstant +from qualtran.linalg.permutation import ( + CycleT, + decompose_permutation_into_cycles, + decompose_permutation_map_into_cycles, +) +from qualtran.symbolics import bit_length, is_symbolic, Shaped, slen, SymbolicInt + +if TYPE_CHECKING: + import sympy + + from qualtran import BloqBuilder, SoquetT + from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + +SymbolicCycleT: TypeAlias = Union[CycleT, Shaped] + + +def _convert_cycle(cycle) -> Union[tuple[int, ...], Shaped]: + if isinstance(cycle, Shaped): + return cycle + return tuple(cycle) + + +@frozen +class PermutationCycle(Bloq): + r"""Apply a single permutation cycle on the basis states. + + Given a permutation cycle $C = (v_0 v_2 \ldots v_{k - 1})$, applies the following unitary: + + $$ + U|v_i\rangle \mapsto |v_{(i + 1)\mod k}\rangle + $$ + + for each $i \in [0, k)$, and + + $$ + U|x\rangle \mapsto |x\rangle + $$ + + and for every $x \not\in C$. + + Args: + N: the total size the permutation acts on. + cycle: the permutation cycle to apply. + + Registers: + x: integer register storing a value in [0, ..., N - 1] + + References: + [A simple quantum algorithm to efficiently prepare sparse states](https://arxiv.org/abs/2310.19309v1) + Appendix B, Algorithm 7. + """ + + N: SymbolicInt + cycle: Union[tuple[int, ...], Shaped] = field(converter=_convert_cycle) + + @cached_property + def signature(self) -> Signature: + return Signature.build_from_dtypes(x=BoundedQUInt(self.bitsize, self.N)) + + @cached_property + def bitsize(self): + return bit_length(self.N - 1) + + def is_symbolic(self): + return is_symbolic(self.N, self.cycle) + + def build_composite_bloq(self, bb: 'BloqBuilder', x: 'SoquetT') -> dict[str, 'SoquetT']: + if self.is_symbolic(): + raise DecomposeTypeError(f"cannot decompose symbolic {self}") + assert not isinstance(self.cycle, Shaped) + + a: 'SoquetT' = bb.allocate(dtype=QBit()) + + for k, x_k in enumerate(self.cycle): + x, a = bb.add_t(EqualsAConstant(self.bitsize, x_k), x=x, target=a) + + delta = x_k ^ self.cycle[(k + 1) % len(self.cycle)] + a, x = bb.add_t(XorK(QUInt(self.bitsize), delta).controlled(), ctrl=a, x=x) + + x, a = bb.add_t(EqualsAConstant(self.bitsize, self.cycle[0]), x=x, target=a) + + bb.free(cast(Soquet, a)) + + return {'x': x} + + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + if self.is_symbolic(): + x = ssa.new_symbol('x') + cycle_len = slen(self.cycle) + return { + (EqualsAConstant(self.bitsize, x), cycle_len + 1), + (XorK(QUInt(self.bitsize), x).controlled(), cycle_len), + } + + return super().build_call_graph(ssa) + + +@bloq_example +def _permutation_cycle() -> PermutationCycle: + permutation_cycle = PermutationCycle(4, (0, 1, 2)) + return permutation_cycle + + +@bloq_example +def _permutation_cycle_symb() -> PermutationCycle: + import sympy + + from qualtran.symbolics import Shaped + + N, L = sympy.symbols("N L") + cycle = Shaped((L,)) + permutation_cycle_symb = PermutationCycle(N, cycle) + return permutation_cycle_symb + + +_PERMUTATION_CYCLE_DOC = BloqDocSpec( + bloq_cls=PermutationCycle, + import_line='from qualtran.bloqs.arithmetic.permutation import PermutationCycle', + examples=[_permutation_cycle, _permutation_cycle_symb], +) + + +def _convert_cycles(cycles) -> Union[tuple[SymbolicCycleT, ...], Shaped]: + if isinstance(cycles, Shaped): + return cycles + return tuple(_convert_cycle(cycle) for cycle in cycles) + + +@frozen +class Permutation(Bloq): + r"""Apply a permutation of [0, N - 1] on the basis states. + + Given a permutation $P : [0, N - 1] \to [0, N - 1]$, this bloq applies the unitary: + + $$ + U|x\rangle = |P(x)\rangle + $$ + + Decomposes a permutation into cycles and applies them in order. + See :meth:`from_dense_permutation` to construct this bloq from a permutation, + and :meth:`from_partial_permutation_map` to construct it from a mapping. + + Args: + N: the total size the permutation acts on. + cycles: a sequence of permutation cycles that form the permutation. + + Registers: + x: integer register storing a value in [0, ..., N - 1] + + References: + [A simple quantum algorithm to efficiently prepare sparse states](https://arxiv.org/abs/2310.19309v1) + Appendix B. + """ + + N: SymbolicInt + cycles: Union[tuple[SymbolicCycleT, ...], Shaped] = field(converter=_convert_cycles) + + @cached_property + def signature(self) -> Signature: + return Signature.build_from_dtypes(x=BoundedQUInt(self.bitsize, self.N)) + + @cached_property + def bitsize(self): + return bit_length(self.N - 1) + + def is_symbolic(self): + return is_symbolic(self.N, self.cycles) or ( + isinstance(self.cycles, tuple) and is_symbolic(*self.cycles) + ) + + @classmethod + def from_dense_permutation(cls, permutation: Sequence[int]): + """Construct a permutation bloq from a dense permutation of size `N`. + + Args: + permutation: a sequence of length `N` containing a permutation of `[0, N)`. + """ + N = len(permutation) + cycles = tuple(decompose_permutation_into_cycles(permutation)) + return cls(N, cycles) + + @classmethod + def from_partial_permutation_map(cls, N: int, permutation_map: dict[int, int]): + """Construct a permutation bloq from a (partial) permutation mapping + + Constructs a permuation of `[0, N)` from a partial mapping. Any numbers that + do not occur in `permutation_map` (i.e. as keys or values) are treated as + mapping to themselves. + + Args: + N: the upper limit, i.e. permutation is on range `[0, N)` + permutation_map: a dictionary defining the permutation + """ + cycles = tuple(decompose_permutation_map_into_cycles(permutation_map)) + return cls(N, cycles) + + @classmethod + def from_cycle_lengths(cls, N: SymbolicInt, cycle_lengths: tuple[SymbolicInt, ...]): + """Construct a permutation bloq from a dense permutation of size `N`. + + Args: + N: the upper limit, i.e. permutation is on range `[0, N)` + cycle_lengths: a tuple of lengths of each non-trivial cycle (i.e. length at least 2). + """ + cycles = tuple( + Shaped((cycle_len,)) + for cycle_len in cycle_lengths + if is_symbolic(cycle_len) or cycle_len >= 2 + ) + return cls(N, cycles) + + def build_composite_bloq(self, bb: 'BloqBuilder', x: 'Soquet') -> dict[str, 'SoquetT']: + if self.is_symbolic(): + raise DecomposeTypeError(f"cannot decompose symbolic {self}") + + assert not isinstance(self.cycles, Shaped) + for cycle in self.cycles: + x = bb.add(PermutationCycle(self.N, cycle), x=x) + + return {'x': x} + + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + if self.is_symbolic(): + # worst case cost: single cycle of length N + cycle = Shaped((self.N,)) + return {(PermutationCycle(self.N, cycle), 1)} + + return super().build_call_graph(ssa) + + +@bloq_example +def _permutation() -> Permutation: + permutation = Permutation.from_dense_permutation([1, 3, 0, 2]) + return permutation + + +@bloq_example +def _permutation_symb() -> Permutation: + import sympy + + from qualtran.symbolics import Shaped + + N, k = sympy.symbols("N k") + permutation_symb = Permutation(N, Shaped((k,))) + return permutation_symb + + +@bloq_example +def _permutation_symb_with_cycles() -> Permutation: + import sympy + + from qualtran.symbolics import Shaped + + N = sympy.symbols("N") + n_cycles = 4 + d = sympy.IndexedBase('d', shape=(n_cycles,)) + permutation_symb_with_cycles = Permutation(N, tuple(Shaped((d[i],)) for i in range(n_cycles))) + return permutation_symb_with_cycles + + +@bloq_example +def _sparse_permutation() -> Permutation: + sparse_permutation = Permutation.from_partial_permutation_map( + 16, {0: 1, 1: 3, 2: 8, 3: 15, 4: 12} + ) + return sparse_permutation + + +_PERMUTATION_DOC = BloqDocSpec( + bloq_cls=Permutation, + import_line='from qualtran.bloqs.arithmetic.permutation import Permutation', + examples=[_permutation, _permutation_symb, _permutation_symb_with_cycles, _sparse_permutation], +) diff --git a/qualtran/bloqs/arithmetic/permutation_test.py b/qualtran/bloqs/arithmetic/permutation_test.py new file mode 100644 index 000000000..2f0520313 --- /dev/null +++ b/qualtran/bloqs/arithmetic/permutation_test.py @@ -0,0 +1,138 @@ +# 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. +# +# 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 sympy + +from qualtran import QBit +from qualtran.bloqs.arithmetic.permutation import ( + _permutation, + _permutation_cycle, + _permutation_cycle_symb, + _permutation_symb, + _permutation_symb_with_cycles, + _sparse_permutation, + Permutation, + PermutationCycle, +) +from qualtran.bloqs.basic_gates import CNOT, TGate, XGate +from qualtran.bloqs.bookkeeping import Allocate, ArbitraryClifford, Free +from qualtran.resource_counting.generalizers import ignore_split_join +from qualtran.symbolics import bit_length, slen + + +def test_examples(bloq_autotester): + bloq_autotester(_permutation_cycle) + bloq_autotester(_permutation) + bloq_autotester(_sparse_permutation) + + +def test_symbolic_examples(bloq_autotester): + bloq_autotester(_permutation_cycle_symb) + bloq_autotester(_permutation_symb) + bloq_autotester(_permutation_symb_with_cycles) + + +def test_permutation_cycle_unitary_and_call_graph(): + bloq = PermutationCycle(4, (0, 1, 2)) + + np.testing.assert_allclose( + bloq.tensor_contract(), np.array([[0, 0, 1, 0], [1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 0, 1]]) + ) + + _, sigma = bloq.call_graph(generalizer=ignore_split_join) + assert sigma == { + CNOT(): 8, + TGate(): 16, + ArbitraryClifford(n=2): 76, + Allocate(QBit()): 1, + Free(QBit()): 1, + } + + +def test_permutation_cycle_symbolic_call_graph(): + bloq = _permutation_cycle_symb() + logN, L = bit_length(bloq.N - 1), slen(bloq.cycle) + + _, sigma = bloq.call_graph() + assert sigma == { + ArbitraryClifford(n=2): (L + 1) * (13 * logN - 13), + TGate(): (L + 1) * (4 * logN - 4), + CNOT(): L * logN + L + 1, + } + + +def test_permutation_unitary_and_call_graph(): + bloq = Permutation.from_dense_permutation([1, 2, 0, 4, 3, 5, 6]) + + np.testing.assert_allclose( + bloq.tensor_contract(), + np.array( + [ + [0, 0, 1, 0, 0, 0, 0, 0], + [1, 0, 0, 0, 0, 0, 0, 0], + [0, 1, 0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 1, 0, 0, 0], + [0, 0, 0, 1, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 1, 0, 0], + [0, 0, 0, 0, 0, 0, 1, 0], + [0, 0, 0, 0, 0, 0, 0, 1], + ] + ), + ) + + _, sigma = bloq.call_graph(generalizer=ignore_split_join) + assert sigma == { + CNOT(): 17, + TGate(): 56, + XGate(): 56, + ArbitraryClifford(n=2): 182, + Allocate(QBit()): 2, + Free(QBit()): 2, + } + + +def test_sparse_permutation_classical_sim(): + N = 50 + perm_map = {0: 1, 1: 20, 2: 30, 3: 40} + bloq = Permutation.from_partial_permutation_map(N, perm_map) + assert bloq.bitsize == 6 + assert bloq.cycles == ((0, 1, 20), (2, 30), (3, 40)) + + for i, x in perm_map.items(): + assert bloq.call_classically(x=i) == (x,) + + +def test_permutation_symbolic_call_graph(): + N = sympy.Symbol("N") + logN = bit_length(N - 1) + bloq = _permutation_symb() + + _, sigma = bloq.call_graph() + assert sigma == { + ArbitraryClifford(n=2): (N + 1) * (13 * logN - 13), + TGate(): (N + 1) * (4 * logN - 4), + CNOT(): N * logN + N + 1, + } diff --git a/qualtran/bloqs/mcmt/and_bloq.py b/qualtran/bloqs/mcmt/and_bloq.py index 1f9cbb572..2836ae48e 100644 --- a/qualtran/bloqs/mcmt/and_bloq.py +++ b/qualtran/bloqs/mcmt/and_bloq.py @@ -260,7 +260,7 @@ class MultiAnd(Bloq): @cvs.validator def _validate_cvs(self, field, val): if not is_symbolic(val) and len(val) < 3: - raise ValueError("MultiAnd must cvshave at least 3 control values `cvs`.") + raise ValueError("MultiAnd must have at least 3 control values `cvs`.") @property def n_ctrls(self) -> SymbolicInt: diff --git a/qualtran/conftest.py b/qualtran/conftest.py index 241fd04dd..7c5fcd56c 100644 --- a/qualtran/conftest.py +++ b/qualtran/conftest.py @@ -103,6 +103,8 @@ def assert_bloq_example_serializes_for_pytest(bloq_ex: BloqExample): 'apply_lth_bloq', 'linear_combination_block_encoding', 'phase_block_encoding', + 'sparse_permutation', + 'permutation_cycle_symb', ]: pytest.xfail("Skipping serialization test for bloq examples that cannot yet be serialized.") diff --git a/qualtran/linalg/permutation.py b/qualtran/linalg/permutation.py new file mode 100644 index 000000000..b09681cc0 --- /dev/null +++ b/qualtran/linalg/permutation.py @@ -0,0 +1,58 @@ +# 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. +from typing import Iterator, Sequence, TypeAlias + +CycleT: TypeAlias = tuple[int, ...] + + +def decompose_permutation_into_cycles(permutation: Sequence[int]) -> Iterator[CycleT]: + """Generate all non-trivial (more than one element) cycles of a permutation of [0, ..., N - 1]""" + import networkx as nx + + G = nx.DiGraph(enumerate(permutation)) + for cycle in nx.simple_cycles(G): + if len(cycle) >= 2: + yield tuple(cycle) + + +def decompose_permutation_map_into_cycles(permutation_map: dict[int, int]) -> Iterator[CycleT]: + r"""Given a (partial) permutation map, return non-trivial cycles requiring minimum swaps. + + We are given a partial permutation on $N$ as a python dictionary. This procedure generates a + sequence of cycles such that the number of swaps required to implement this partial mapping + is minimized. + + >>> list(decompose_permutation_map_into_cycles({0: 1, 1: 5, 5: 0, 2: 6, 3: 3}, N=10)) + [(0, 1, 5), (2, 6)] + + Args: + permutation_map: a (partial) map defining the permutation. + """ + seen = set() + + for i in permutation_map: + if i in seen: + continue + + # compute the cycle starting at `i` + cycle = [] + while True: + seen.add(i) + cycle.append(i) + if i not in permutation_map: + break + i = permutation_map[i] + + if len(cycle) >= 2: + yield tuple(cycle) diff --git a/qualtran/linalg/permutation_test.py b/qualtran/linalg/permutation_test.py new file mode 100644 index 000000000..04c172eeb --- /dev/null +++ b/qualtran/linalg/permutation_test.py @@ -0,0 +1,28 @@ +# 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. +from qualtran.linalg.permutation import ( + decompose_permutation_into_cycles, + decompose_permutation_map_into_cycles, +) + + +def test_decompose_permutation_into_cycles(): + assert list(decompose_permutation_into_cycles([0, 1, 2])) == [] + assert list(decompose_permutation_into_cycles([1, 2, 0])) == [(0, 1, 2)] + assert sorted(decompose_permutation_into_cycles([1, 0, 2, 4, 5, 3])) == [(0, 1), (3, 4, 5)] + + +def test_decompose_sparse_prefix_permutation_into_cycles(): + assert list(decompose_permutation_map_into_cycles({0: 1, 1: 20})) == [(0, 1, 20)] + assert sorted(decompose_permutation_map_into_cycles({0: 30, 1: 50})) == [(0, 30), (1, 50)] diff --git a/qualtran/serialization/resolver_dict.py b/qualtran/serialization/resolver_dict.py index a14d73143..41c65c744 100644 --- a/qualtran/serialization/resolver_dict.py +++ b/qualtran/serialization/resolver_dict.py @@ -15,10 +15,12 @@ from typing import Type import qualtran.bloqs.arithmetic.addition +import qualtran.bloqs.arithmetic.bitwise import qualtran.bloqs.arithmetic.comparison import qualtran.bloqs.arithmetic.conversions import qualtran.bloqs.arithmetic.hamming_weight import qualtran.bloqs.arithmetic.multiplication +import qualtran.bloqs.arithmetic.permutation import qualtran.bloqs.arithmetic.sorting import qualtran.bloqs.basic_gates.cnot import qualtran.bloqs.basic_gates.hadamard @@ -133,6 +135,7 @@ "qualtran.bloqs.arithmetic.addition.Add": qualtran.bloqs.arithmetic.addition.Add, "qualtran.bloqs.arithmetic.addition.OutOfPlaceAdder": qualtran.bloqs.arithmetic.addition.OutOfPlaceAdder, "qualtran.bloqs.arithmetic.addition.AddK": qualtran.bloqs.arithmetic.AddK, + "qualtran.bloqs.arithmetic.bitwise.XorK": qualtran.bloqs.arithmetic.bitwise.XorK, "qualtran.bloqs.arithmetic.comparison.BiQubitsMixer": qualtran.bloqs.arithmetic.comparison.BiQubitsMixer, "qualtran.bloqs.arithmetic.comparison.EqualsAConstant": qualtran.bloqs.arithmetic.comparison.EqualsAConstant, "qualtran.bloqs.arithmetic.comparison.GreaterThan": qualtran.bloqs.arithmetic.comparison.GreaterThan, @@ -151,6 +154,8 @@ "qualtran.bloqs.arithmetic.multiplication.Square": qualtran.bloqs.arithmetic.multiplication.Square, "qualtran.bloqs.arithmetic.multiplication.SquareRealNumber": qualtran.bloqs.arithmetic.multiplication.SquareRealNumber, "qualtran.bloqs.arithmetic.multiplication.SumOfSquares": qualtran.bloqs.arithmetic.multiplication.SumOfSquares, + "qualtran.bloqs.arithmetic.permutation.Permutation": qualtran.bloqs.arithmetic.permutation.Permutation, + "qualtran.bloqs.arithmetic.permutation.PermutationCycle": qualtran.bloqs.arithmetic.permutation.PermutationCycle, "qualtran.bloqs.arithmetic.sorting.BitonicMerge": qualtran.bloqs.arithmetic.sorting.BitonicMerge, "qualtran.bloqs.arithmetic.sorting.BitonicSort": qualtran.bloqs.arithmetic.sorting.BitonicSort, "qualtran.bloqs.arithmetic.sorting.Comparator": qualtran.bloqs.arithmetic.sorting.Comparator, 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 3/3] 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,