Skip to content

Commit

Permalink
Support symbolic decomposition for state prep via alias sampling (#1084)
Browse files Browse the repository at this point in the history 
* Support symbolic decomposition for state prep via alias sampling

* Fix pylint

* Use assert_valid_bloq_decomposition(bloq)

* Fix formatting

* Update docstrings with args and regenerate notebook

* Address nits and add an assertion for 1D state prep

---------

Co-authored-by: Anurudh Peduri <[email protected]>
  • Loading branch information
@tanujkhattar @anurudhp
tanujkhattar and anurudhp authored Aug 13, 2024
1 parent ec85731 commit d409550
Show file tree
Hide file tree
Showing 4 changed files with 78 additions and 57 deletions.
31 changes: 14 additions & 17 deletions qualtran/bloqs/state_preparation/state_preparation_alias_sampling.ipynb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -62,23 +62,6 @@
"selecting `l` uniformly at random and then returning it with probability `keep[l] / 2**mu`;\n",
"otherwise returning `alt[l]`.\n",
"\n",
"#### Parameters\n",
" - `selection_registers`: The input/output registers to prepare the state on (see Signature).\n",
" - `keep`: The discretized `keep` probabilities for alias sampling.\n",
" - `alt`: The alternate/alias values to swap.\n",
" - `mu`: The number of bits to approximate the `keep` probabilities.\n",
" - `sum_of_unnormalized_probabilities`: The total of the input unnormalized probabilities, i.e., $\\lambda$. This is used as the `PrepareOracle.l1_norm_of_coeffs` property. \n",
"\n",
"Signature:\n",
" selection: The input/output register $|\\ell\\rangle$ of size lg(L) where the desired\n",
" coefficient state is prepared.\n",
" temp: Work space comprised of sub signature:\n",
" - sigma: A mu-sized register containing uniform probabilities for comparison against\n",
" `keep`.\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",
" - one bit for the result of the comparison.\n",
"\n",
"This gate corresponds to the following operations:\n",
" - UNIFORM_L on the selection register\n",
" - H^mu on the sigma register\n",
Expand All @@ -90,6 +73,20 @@
"Total space will be (2 * log(L) + 2 mu + 1) work qubits + log(L) ancillas for QROM.\n",
"The 1 ancilla in work qubits is for the `LessThanEqualGate` followed by coherent swap.\n",
"\n",
"#### Registers\n",
" - `selection`: The input/output register $|\\mathrm{ind}_l\\rangle$ of size lg(L) where the desired coefficient state is prepared. Default name is 'selection' if the builder methods on the class are used. Or else, users can specify custom named registers\n",
" - `sigma_mu`: A mu-sized register containing uniform probabilities for comparison against `keep`.\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",
"#### Parameters\n",
" - `selection_registers`: The input/output registers to prepare the state on (see Registers section).\n",
" - `keep`: The discretized `keep` probabilities for alias sampling.\n",
" - `alt`: The alternate/alias values to swap.\n",
" - `mu`: The number of bits to approximate the `keep` probabilities.\n",
" - `sum_of_unnormalized_probabilities`: The total of the input unnormalized probabilities, i.e., $\\lambda$. This is used as the `PrepareOracle.l1_norm_of_coeffs` property. \n",
"\n",
"#### References\n",
" - [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"
]
Expand Down
90 changes: 55 additions & 35 deletions qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,10 +20,9 @@
largest absolute error that one can tolerate in the prepared amplitudes.
"""
from functools import cached_property
from typing import Iterator, Sequence, Set, Tuple, TYPE_CHECKING, Union
from typing import Sequence, Set, Tuple, TYPE_CHECKING, Union

import attrs
import cirq
import numpy as np
from numpy.typing import NDArray

Expand All @@ -48,7 +47,7 @@
from qualtran.symbolics import bit_length, is_symbolic, Shaped, slen, SymbolicFloat, SymbolicInt

if TYPE_CHECKING:
from qualtran import BloqBuilder, SoquetT
from qualtran import BloqBuilder, Soquet, SoquetT
from qualtran.resource_counting import BloqCountT, SympySymbolAllocator


Expand Down Expand Up @@ -81,24 +80,6 @@ class StatePreparationAliasSampling(PrepareOracle):
selecting `l` uniformly at random and then returning it with probability `keep[l] / 2**mu`;
otherwise returning `alt[l]`.
Args:
selection_registers: The input/output registers to prepare the state on (see Signature).
keep: The discretized `keep` probabilities for alias sampling.
alt: The alternate/alias values to swap.
mu: The number of bits to approximate the `keep` probabilities.
sum_of_unnormalized_probabilities: The total of the input unnormalized probabilities,
i.e., $\lambda$. This is used as the `PrepareOracle.l1_norm_of_coeffs` property.
Signature:
selection: The input/output register $|\ell\rangle$ of size lg(L) where the desired
coefficient state is prepared.
temp: Work space comprised of sub signature:
- sigma: A mu-sized register containing uniform probabilities for comparison against
`keep`.
- alt: A lg(L)-sized register of alternate indices
- keep: a mu-sized register of probabilities of keeping the initially sampled index.
- one bit for the result of the comparison.
This gate corresponds to the following operations:
- UNIFORM_L on the selection register
- H^mu on the sigma register
Expand All @@ -110,6 +91,23 @@ class StatePreparationAliasSampling(PrepareOracle):
Total space will be (2 * log(L) + 2 mu + 1) work qubits + log(L) ancillas for QROM.
The 1 ancilla in work qubits is for the `LessThanEqualGate` followed by coherent swap.
Registers:
selection: The input/output register $|\mathrm{ind}_l\rangle$ of size lg(L) where the desired
coefficient state is prepared. Default name is 'selection' if the builder methods on
the class are used. Or else, users can specify custom named registers
sigma_mu: A mu-sized register containing uniform probabilities for comparison against `keep`.
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.
Args:
selection_registers: The input/output registers to prepare the state on (see Registers section).
keep: The discretized `keep` probabilities for alias sampling.
alt: The alternate/alias values to swap.
mu: The number of bits to approximate the `keep` probabilities.
sum_of_unnormalized_probabilities: The total of the input unnormalized probabilities,
i.e., $\lambda$. This is used as the `PrepareOracle.l1_norm_of_coeffs` property.
References:
[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.
Expand All @@ -123,6 +121,14 @@ class StatePreparationAliasSampling(PrepareOracle):
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 greater than 0.")
if len(self.selection_registers) != 1:
raise ValueError(
f"{type(self)} only supports 1D state preparation. Found multiple {self.selection_registers=}."
)

@classmethod
def from_probabilities(
cls, unnormalized_probabilities: Sequence[float], *, precision: float = 1.0e-5
Expand Down Expand Up @@ -235,21 +241,35 @@ def qrom_bloq(self) -> QROM:
(self.alternates_bitsize, self.keep_bitsize),
)

def decompose_from_registers(
def build_composite_bloq(
self,
*,
context: cirq.DecompositionContext,
**quregs: NDArray[cirq.Qid], # type:ignore[type-var]
) -> Iterator[cirq.OP_TREE]:
yield PrepareUniformSuperposition(self.n_coeff).on(*quregs['selection'])
if self.mu == 0:
return
selection, less_than_equal = quregs['selection'], quregs['less_than_equal']
sigma_mu, alt, keep = quregs['sigma_mu'], quregs['alt'], quregs['keep']
yield cirq.H.on_each(*sigma_mu)
yield self.qrom_bloq.on_registers(selection=selection, target0_=alt, target1_=keep)
yield LessThanEqual(self.mu, self.mu).on(*keep, *sigma_mu, *less_than_equal)
yield CSwap.make_on(ctrl=less_than_equal, x=alt, y=selection)
bb: 'BloqBuilder',
sigma_mu: 'SoquetT',
alt: 'SoquetT',
keep: 'SoquetT',
less_than_equal: 'Soquet',
**soqs: 'SoquetT',
):
selection = soqs.pop(self.selection_registers[0].name)
assert not soqs
selection = bb.add(PrepareUniformSuperposition(self.n_coeff), target=selection)
sigma_mu = bb.add(OnEach(self.mu, Hadamard()), q=sigma_mu)
selection, alt, keep = bb.add(
self.qrom_bloq, selection=selection, target0_=alt, target1_=keep
)
keep, sigma_mu, less_than_equal = bb.add(
LessThanEqual(self.mu, self.mu), x=keep, y=sigma_mu, target=less_than_equal
)
less_than_equal, alt, selection = bb.add(
CSwap(self.selection_bitsize), ctrl=less_than_equal, x=alt, y=selection
)
return {
self.selection_registers[0].name: selection,
'less_than_equal': less_than_equal,
'sigma_mu': sigma_mu,
'alt': alt,
'keep': keep,
}

def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']:
return {
Expand Down
13 changes: 8 additions & 5 deletions qualtran/bloqs/state_preparation/state_preparation_alias_sampling_test.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@

def test_state_prep_alias_sampling_autotest(bloq_autotester):
bloq_autotester(_state_prep_alias)
bloq_autotester(_state_prep_alias_symb)


def test_sparse_state_prep_alias_sampling_autotest(bloq_autotester):
Expand Down Expand Up @@ -89,6 +90,8 @@ def test_state_prep_alias_sampling_symb():
)
)
np.testing.assert_allclose(concrete_t_counts, symb_t_counts, rtol=5e-4)
# Ensure the symbolic bloq can decomposes into a composite bloq
assert_valid_bloq_decomposition(bloq)


def assert_state_preparation_valid_for_coefficient(
Expand Down Expand Up @@ -166,11 +169,11 @@ def test_state_preparation_via_coherent_alias_sampling_diagram():
│ │ │
selection1: ────────target───────In───────────────────×(y)───
│ │
sigma_mu0: ─────────H────────────┼────────In(y)───────┼──────
│ │ │
sigma_mu1: ─────────H────────────┼────────In(y)───────┼──────
│ │ │
sigma_mu2: ─────────H────────────┼────────In(y)───────┼──────
sigma_mu0: ─────────H⨂3──────────┼────────In(y)───────┼──────
│ │ │ │
sigma_mu1: ─────────H⨂3──────────┼────────In(y)───────┼──────
│ │ │ │
sigma_mu2: ─────────H⨂3──────────┼────────In(y)───────┼──────
│ │ │
alt0: ───────────────────────────QROM_a───┼───────────×(x)───
│ │ │
Expand Down
1 change: 1 addition & 0 deletions qualtran/conftest.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -103,6 +103,7 @@ def assert_bloq_example_serializes_for_pytest(bloq_ex: BloqExample):
'apply_lth_bloq',
'linear_combination_block_encoding',
'phase_block_encoding',
'state_prep_alias_symb', # cannot serialize Shaped
'sparse_matrix_block_encoding',
'sparse_matrix_symb_block_encoding',
'sparse_state_prep_alias_symb', # cannot serialize Shaped
Expand Down

0 comments on commit d409550

Please sign in to comment.