From cd638c749a97d7d48a2a8d7312173477d95a9e7b Mon Sep 17 00:00:00 2001
From: anurudhp <anurudhp97@gmail.com>
Date: Wed, 12 Jun 2024 11:26:16 -0700
Subject: [PATCH 01/19] WIP sparse SP alias sampling

---
 .../state_preparation_alias_sampling.py       | 202 +++++++++++++++++-
 1 file changed, 201 insertions(+), 1 deletion(-)

diff --git a/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py b/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py
index 45d45457b..d9bf8e111 100644
--- a/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py
+++ b/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py
@@ -45,7 +45,7 @@
     ignore_cliffords,
     ignore_split_join,
 )
-from qualtran.symbolics import bit_length, Shaped, SymbolicFloat, SymbolicInt
+from qualtran.symbolics import bit_length, Shaped, SymbolicFloat, SymbolicInt, slen
 
 if TYPE_CHECKING:
     from qualtran.resource_counting import BloqCountT, SympySymbolAllocator
@@ -235,6 +235,206 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']:
         }
 
 
+@cirq.value_equality()
+@attrs.frozen
+class SparseStatePreparationAliasSampling(PrepareOracle):
+    r"""Initialize a state with $L$ unique non-zero coefficients using coherent alias sampling.
+
+    In particular, we take the zero state to:
+
+    $$
+    \sum_{\ell=0}^{L-1} \sqrt{p_\ell} |\ell\rangle |\mathrm{temp}_\ell\rangle
+    $$
+
+    where the probabilities $p_\ell$ are $\mu$-bit binary approximations to the true values and
+    where the temporary register must be treated with care, see the details in Section III.D. of
+    the reference.
+
+    The preparation is equivalent to [classical alias sampling]
+    (https://en.wikipedia.org/wiki/Alias_method): we sample `l` with probability `p[l]` by first
+    selecting `l` uniformly at random and then returning it with probability `keep[l] / 2**mu`;
+    otherwise returning `alt[l]`.
+
+    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
+     - QROM addressed by the selection register into the alt and keep signature.
+     - LessThanEqualGate comparing the keep and sigma signature.
+     - Coherent swap between the selection register and alt register if the comparison
+       returns True.
+
+    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.
+
+    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.
+    """
+    selection_registers: Tuple[Register, ...] = attrs.field(
+        converter=lambda v: (v,) if isinstance(v, Register) else tuple(v)
+    )
+    index: Union[Shaped, NDArray[np.int_]]
+    alt: Union[Shaped, NDArray[np.int_]]
+    keep: Union[Shaped, NDArray[np.int_]]
+    mu: SymbolicInt
+    sum_of_lcu_coeffs: SymbolicFloat
+
+    @classmethod
+    def from_lcu_probs(
+        cls, lcu_probabilities: Sequence[float], *, probability_epsilon: float = 1.0e-5
+    ) -> 'StatePreparationAliasSampling':
+        """Factory to construct the state preparation gate for a given set of LCU coefficients.
+
+        Args:
+            lcu_probabilities: The LCU coefficients.
+            probability_epsilon: 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.
+        """
+        alt, keep, mu = preprocess_lcu_coefficients_for_reversible_sampling(
+            lcu_coefficients=lcu_probabilities, epsilon=probability_epsilon
+        )
+        N = len(lcu_probabilities)
+
+        is_nonzero = ~np.isclose(lcu_probabilities, 0)
+        index = np.arange(N)[is_nonzero]
+        alt = np.array(alt)[is_nonzero]
+        keep = np.array(keep)[is_nonzero]
+
+        return StatePreparationAliasSampling(
+            selection_registers=Register('selection', BoundedQUInt((N - 1).bit_length(), N)),
+            index=index,
+            alt=alt,
+            keep=keep,
+            mu=mu,
+            sum_of_lcu_coeffs=sum(abs(x) for x in lcu_probabilities),
+        )
+
+    @classmethod
+    def from_n_coeff(
+        cls,
+        n_coeff: SymbolicInt,
+        n_nonzero_coeff: SymbolicInt,
+        sum_of_lcu_coeffs: SymbolicFloat,
+        *,
+        probability_epsilon: SymbolicFloat = 1.0e-5,
+    ) -> 'StatePreparationAliasSampling':
+        """Factory to construct the state preparation gate for symbolic number of LCU coefficients.
+
+        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_lcu_coeffs: Sum of absolute values of coefficients of the prepared state.
+            probability_epsilon: 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, probability_epsilon)
+        selection_bitsize = bit_length(n_coeff - 1)
+        return StatePreparationAliasSampling(
+            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_lcu_coeffs=sum_of_lcu_coeffs,
+        )
+
+    @property
+    def n_coeff(self) -> SymbolicInt:
+        return self.selection_registers[0].dtype.iteration_length_or_zero()
+
+    @cached_property
+    def l1_norm_of_coeffs(self) -> 'SymbolicFloat':
+        return self.sum_of_lcu_coeffs
+
+    @cached_property
+    def sigma_mu_bitsize(self) -> SymbolicInt:
+        return self.mu
+
+    @cached_property
+    def sparse_index_bitsize(self) -> SymbolicInt:
+        return bit_length(slen(self.index))
+
+    @cached_property
+    def alternates_bitsize(self) -> SymbolicInt:
+        return total_bits(self.selection_registers)
+
+    @cached_property
+    def keep_bitsize(self) -> SymbolicInt:
+        return self.mu
+
+    @cached_property
+    def selection_bitsize(self) -> SymbolicInt:
+        return total_bits(self.selection_registers)
+
+    @cached_property
+    def junk_registers(self) -> Tuple[Register, ...]:
+        return tuple(
+            Signature.build(
+                sigma_mu=self.sigma_mu_bitsize,
+                sparse_index=self.sparse_index_bitsize,
+                alt=self.alternates_bitsize,
+                keep=self.keep_bitsize,
+                less_than_equal=1,
+            )
+        )
+
+    def _value_equality_values_(self):
+        return (
+            self.selection_registers,
+            self.index if isinstance(self.index, Shaped) else tuple(self.index.ravel()),
+            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(
+            (self.index, self.alt, self.keep),
+            (self.sparse_index_bitsize,),
+            (self.selection_bitsize, self.alternates_bitsize, self.keep_bitsize),
+        )
+
+    def decompose_from_registers(
+        self,
+        *,
+        context: cirq.DecompositionContext,
+        **quregs: NDArray[cirq.Qid],  # type:ignore[type-var]
+    ) -> Iterator[cirq.OP_TREE]:
+        yield PrepareUniformSuperposition(self.n).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)
+
+    def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']:
+        return {
+            (PrepareUniformSuperposition(self.n_coeff), 1),
+            (self.qrom_bloq, 1),
+            (LessThanEqual(self.mu, self.mu), 1),
+            (CSwap(self.selection_bitsize), 1),
+            (Hadamard(), self.mu),
+        }
+
+
 @bloq_example(generalizer=[cirq_to_bloqs, ignore_split_join, ignore_cliffords])
 def _state_prep_alias() -> StatePreparationAliasSampling:
     coeffs = [1.0, 1, 3, 2]

From 4dcd9d6155cb777a47c909f36f8db41314a55c26 Mon Sep 17 00:00:00 2001
From: anurudhp <anurudhp97@gmail.com>
Date: Wed, 12 Jun 2024 12:12:10 -0700
Subject: [PATCH 02/19] remove cirq.value equality

---
 .../state_preparation_alias_sampling.py       | 35 ++++++-------------
 1 file changed, 11 insertions(+), 24 deletions(-)

diff --git a/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py b/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py
index d9bf8e111..124d0226c 100644
--- a/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py
+++ b/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py
@@ -51,7 +51,12 @@
     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$ unique coefficients using coherent alias sampling.
@@ -99,8 +104,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_lcu_coeffs: SymbolicFloat
 
@@ -193,14 +198,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(
@@ -235,7 +232,6 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']:
         }
 
 
-@cirq.value_equality()
 @attrs.frozen
 class SparseStatePreparationAliasSampling(PrepareOracle):
     r"""Initialize a state with $L$ unique non-zero coefficients using coherent alias sampling.
@@ -283,9 +279,9 @@ class SparseStatePreparationAliasSampling(PrepareOracle):
     selection_registers: Tuple[Register, ...] = attrs.field(
         converter=lambda v: (v,) if isinstance(v, Register) else tuple(v)
     )
-    index: Union[Shaped, NDArray[np.int_]]
-    alt: Union[Shaped, NDArray[np.int_]]
-    keep: Union[Shaped, NDArray[np.int_]]
+    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_lcu_coeffs: SymbolicFloat
 
@@ -392,15 +388,6 @@ def junk_registers(self) -> Tuple[Register, ...]:
             )
         )
 
-    def _value_equality_values_(self):
-        return (
-            self.selection_registers,
-            self.index if isinstance(self.index, Shaped) else tuple(self.index.ravel()),
-            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(

From 44398f8aca10970e4ee83800ee3f57f1ad913a1f Mon Sep 17 00:00:00 2001
From: anurudhp <anurudhp97@gmail.com>
Date: Wed, 12 Jun 2024 13:09:10 -0700
Subject: [PATCH 03/19] types

---
 .../state_preparation_alias_sampling.py            | 14 +++++++++-----
 1 file changed, 9 insertions(+), 5 deletions(-)

diff --git a/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py b/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py
index 124d0226c..4cafffa37 100644
--- a/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py
+++ b/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py
@@ -45,7 +45,7 @@
     ignore_cliffords,
     ignore_split_join,
 )
-from qualtran.symbolics import bit_length, Shaped, SymbolicFloat, SymbolicInt, slen
+from qualtran.symbolics import bit_length, Shaped, slen, SymbolicFloat, SymbolicInt
 
 if TYPE_CHECKING:
     from qualtran.resource_counting import BloqCountT, SympySymbolAllocator
@@ -288,7 +288,7 @@ class SparseStatePreparationAliasSampling(PrepareOracle):
     @classmethod
     def from_lcu_probs(
         cls, lcu_probabilities: Sequence[float], *, probability_epsilon: float = 1.0e-5
-    ) -> 'StatePreparationAliasSampling':
+    ) -> 'SparseStatePreparationAliasSampling':
         """Factory to construct the state preparation gate for a given set of LCU coefficients.
 
         Args:
@@ -308,7 +308,7 @@ def from_lcu_probs(
         alt = np.array(alt)[is_nonzero]
         keep = np.array(keep)[is_nonzero]
 
-        return StatePreparationAliasSampling(
+        return cls(
             selection_registers=Register('selection', BoundedQUInt((N - 1).bit_length(), N)),
             index=index,
             alt=alt,
@@ -325,7 +325,7 @@ def from_n_coeff(
         sum_of_lcu_coeffs: SymbolicFloat,
         *,
         probability_epsilon: SymbolicFloat = 1.0e-5,
-    ) -> 'StatePreparationAliasSampling':
+    ) -> 'SparseStatePreparationAliasSampling':
         """Factory to construct the state preparation gate for symbolic number of LCU coefficients.
 
         Args:
@@ -339,7 +339,7 @@ def from_n_coeff(
         """
         mu = sub_bit_prec_from_epsilon(n_coeff, probability_epsilon)
         selection_bitsize = bit_length(n_coeff - 1)
-        return StatePreparationAliasSampling(
+        return cls(
             selection_registers=Register('selection', BoundedQUInt(selection_bitsize, n_coeff)),
             index=Shaped((n_nonzero_coeff,)),
             alt=Shaped((n_nonzero_coeff,)),
@@ -352,6 +352,10 @@ def from_n_coeff(
     def n_coeff(self) -> SymbolicInt:
         return self.selection_registers[0].dtype.iteration_length_or_zero()
 
+    @property
+    def n_nonzero_coeff(self) -> SymbolicInt:
+        return self.selection_registers[0].dtype.iteration_length_or_zero()
+
     @cached_property
     def l1_norm_of_coeffs(self) -> 'SymbolicFloat':
         return self.sum_of_lcu_coeffs

From eaf1d17e0eb4c250d87c40e638e56d886de6154a Mon Sep 17 00:00:00 2001
From: anurudhp <anurudhp97@gmail.com>
Date: Wed, 12 Jun 2024 15:10:17 -0700
Subject: [PATCH 04/19] example

---
 .../state_preparation_alias_sampling.py       | 46 +++++++++++++------
 .../state_preparation_alias_sampling_test.py  |  2 +
 2 files changed, 35 insertions(+), 13 deletions(-)

diff --git a/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py b/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py
index 4cafffa37..0d32a31ea 100644
--- a/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py
+++ b/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py
@@ -53,7 +53,9 @@
 
 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
+        tuple(data_or_shape.flatten())
+        if isinstance(data_or_shape, np.ndarray)
+        else (data_or_shape,)
     )
 
 
@@ -239,20 +241,23 @@ class SparseStatePreparationAliasSampling(PrepareOracle):
     In particular, we take the zero state to:
 
     $$
-    \sum_{\ell=0}^{L-1} \sqrt{p_\ell} |\ell\rangle |\mathrm{temp}_\ell\rangle
+    \sum_{l=0}^{L-1} \sqrt{p_l} |\mathrm{ind}_l\rangle |\mathrm{temp}_\ell\rangle
     $$
 
-    where the probabilities $p_\ell$ are $\mu$-bit binary approximations to the true values and
+    where $\mathrm{ind}_l$ is the index of the $l$-th non-zero coefficient,
+    and the probabilities $p_\ell$ are $\mu$-bit binary approximations to the true values and
     where the temporary register must be treated with care, see the details in Section III.D. of
-    the reference.
+    the reference [2].
 
     The preparation is equivalent to [classical alias sampling]
     (https://en.wikipedia.org/wiki/Alias_method): we sample `l` with probability `p[l]` by first
-    selecting `l` uniformly at random and then returning it with probability `keep[l] / 2**mu`;
+    selecting `l` uniformly at random and then returning `ind[l]` with probability `keep[l] / 2**mu`;
     otherwise returning `alt[l]`.
+    This bloq is nearly identical to :class:`StatePreparationByAliasSampling`, except that this loads the
+    non-zero coefficient indices as well from the QROM.
 
     Signature:
-        selection: The input/output register $|\ell\rangle$ of size lg(L) where the desired
+        selection: The input/output register $|\mathrm{ind}_l\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
@@ -262,7 +267,7 @@ class SparseStatePreparationAliasSampling(PrepareOracle):
             - one bit for the result of the comparison.
 
     This gate corresponds to the following operations:
-     - UNIFORM_L on the selection register
+     - UNIFORM_L on the sparse_index register
      - H^mu on the sigma register
      - QROM addressed by the selection register into the alt and keep signature.
      - LessThanEqualGate comparing the keep and sigma signature.
@@ -273,8 +278,10 @@ class SparseStatePreparationAliasSampling(PrepareOracle):
     The 1 ancilla in work qubits is for the `LessThanEqualGate` followed by coherent swap.
 
     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.
+        [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)
@@ -354,7 +361,7 @@ def n_coeff(self) -> SymbolicInt:
 
     @property
     def n_nonzero_coeff(self) -> SymbolicInt:
-        return self.selection_registers[0].dtype.iteration_length_or_zero()
+        return slen(self.index)
 
     @cached_property
     def l1_norm_of_coeffs(self) -> 'SymbolicFloat':
@@ -366,7 +373,7 @@ def sigma_mu_bitsize(self) -> SymbolicInt:
 
     @cached_property
     def sparse_index_bitsize(self) -> SymbolicInt:
-        return bit_length(slen(self.index))
+        return bit_length(self.n_nonzero_coeff)
 
     @cached_property
     def alternates_bitsize(self) -> SymbolicInt:
@@ -406,13 +413,16 @@ def decompose_from_registers(
         context: cirq.DecompositionContext,
         **quregs: NDArray[cirq.Qid],  # type:ignore[type-var]
     ) -> Iterator[cirq.OP_TREE]:
-        yield PrepareUniformSuperposition(self.n).on(*quregs['selection'])
+        sparse_index = quregs['sparse_index']
+        yield PrepareUniformSuperposition(self.n_nonzero_coeff).on(*sparse_index)
         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 self.qrom_bloq.on_registers(
+            selection=sparse_index, target0_=selection, target1_=alt, target2_=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)
 
@@ -447,6 +457,16 @@ def _state_prep_alias_symb() -> StatePreparationAliasSampling:
     return state_prep_alias_symb
 
 
+@bloq_example(generalizer=[cirq_to_bloqs, ignore_split_join, ignore_cliffords])
+def _sparse_state_prep_alias() -> SparseStatePreparationAliasSampling:
+    coeffs = [1.0, 0, 0, 1, 0, 3, 0, 2, 0]
+    mu = 3
+    state_prep_alias = SparseStatePreparationAliasSampling.from_lcu_probs(
+        coeffs, probability_epsilon=2**-mu / len(coeffs)
+    )
+    return state_prep_alias
+
+
 _STATE_PREP_ALIAS_DOC = BloqDocSpec(
     bloq_cls=StatePreparationAliasSampling,
     import_line='from qualtran.bloqs.state_preparation import StatePreparationAliasSampling',
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 a777e6161..933526f0e 100644
--- a/qualtran/bloqs/state_preparation/state_preparation_alias_sampling_test.py
+++ b/qualtran/bloqs/state_preparation/state_preparation_alias_sampling_test.py
@@ -22,6 +22,7 @@
     _state_prep_alias,
     _state_prep_alias_symb,
     StatePreparationAliasSampling,
+    _sparse_state_prep_alias,
 )
 from qualtran.cirq_interop.testing import GateHelper
 from qualtran.testing import assert_valid_bloq_decomposition, execute_notebook
@@ -29,6 +30,7 @@
 
 def test_state_prep_alias_sampling_autotest(bloq_autotester):
     bloq_autotester(_state_prep_alias)
+    bloq_autotester(_sparse_state_prep_alias)
 
 
 def test_state_prep_alias_sampling_symb():

From 08fe4eee19fa1a9a3171bd5c03781cbd15bef613 Mon Sep 17 00:00:00 2001
From: anurudhp <anurudhp97@gmail.com>
Date: Thu, 13 Jun 2024 15:04:46 -0700
Subject: [PATCH 05/19] add tests

---
 .../state_preparation_alias_sampling.py       | 67 ++++++++++++-------
 .../state_preparation_alias_sampling_test.py  | 30 +++++++--
 qualtran/serialization/resolver_dict.py       |  1 +
 3 files changed, 68 insertions(+), 30 deletions(-)

diff --git a/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py b/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py
index 0d32a31ea..48ecb6b36 100644
--- a/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py
+++ b/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py
@@ -20,7 +20,7 @@
 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 Dict, Iterator, Sequence, Set, Tuple, TYPE_CHECKING, Union
 
 import attrs
 import cirq
@@ -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 (
@@ -48,6 +48,7 @@
 from qualtran.symbolics import bit_length, Shaped, slen, SymbolicFloat, SymbolicInt
 
 if TYPE_CHECKING:
+    from qualtran import BloqBuilder, SoquetT
     from qualtran.resource_counting import BloqCountT, SympySymbolAllocator
 
 
@@ -373,7 +374,7 @@ def sigma_mu_bitsize(self) -> SymbolicInt:
 
     @cached_property
     def sparse_index_bitsize(self) -> SymbolicInt:
-        return bit_length(self.n_nonzero_coeff)
+        return bit_length(self.n_nonzero_coeff - 1)
 
     @cached_property
     def alternates_bitsize(self) -> SymbolicInt:
@@ -407,32 +408,48 @@ def qrom_bloq(self) -> QROM:
             (self.selection_bitsize, 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]:
-        sparse_index = quregs['sparse_index']
-        yield PrepareUniformSuperposition(self.n_nonzero_coeff).on(*sparse_index)
-        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=sparse_index, target0_=selection, target1_=alt, target2_=keep
+        bb: 'BloqBuilder',
+        selection: 'SoquetT',
+        sparse_index: 'SoquetT',
+        alt: 'SoquetT',
+        keep: 'SoquetT',
+        sigma_mu: 'SoquetT',
+        less_than_equal: 'SoquetT',
+    ) -> Dict[str, 'SoquetT']:
+        sparse_index = bb.add(
+            PrepareUniformSuperposition(self.n_nonzero_coeff), target=sparse_index
         )
-        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)
+        if self.mu == 0:
+            sparse_index, selection = bb.add(
+                QROM((self.index,), (self.sparse_index_bitsize,), (self.selection_bitsize,)),
+                selection=sparse_index,
+                target0_=selection,
+            )
+        else:
+            sigma_mu = bb.add(OnEach(self.sigma_mu_bitsize, Hadamard()), q=sigma_mu)
+            sparse_index, selection, alt, keep = bb.add_t(
+                self.qrom_bloq,
+                selection=sparse_index,
+                target0_=selection,
+                target1_=alt,
+                target2_=keep,
+            )
+            keep, sigma_mu, less_than_equal = bb.add_t(
+                LessThanEqual(self.mu, self.mu), x=keep, y=sigma_mu, target=less_than_equal
+            )
+            less_than_equal, alt, selection = bb.add_t(
+                CSwap(self.selection_bitsize), ctrl=less_than_equal, x=alt, y=selection
+            )
 
-    def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']:
         return {
-            (PrepareUniformSuperposition(self.n_coeff), 1),
-            (self.qrom_bloq, 1),
-            (LessThanEqual(self.mu, self.mu), 1),
-            (CSwap(self.selection_bitsize), 1),
-            (Hadamard(), self.mu),
+            'selection': selection,
+            'sparse_index': sparse_index,
+            'alt': alt,
+            'keep': keep,
+            'sigma_mu': sigma_mu,
+            'less_than_equal': less_than_equal,
         }
 
 
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 933526f0e..cd7c09da6 100644
--- a/qualtran/bloqs/state_preparation/state_preparation_alias_sampling_test.py
+++ b/qualtran/bloqs/state_preparation/state_preparation_alias_sampling_test.py
@@ -19,10 +19,11 @@
 
 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,
     _state_prep_alias,
     _state_prep_alias_symb,
+    SparseStatePreparationAliasSampling,
     StatePreparationAliasSampling,
-    _sparse_state_prep_alias,
 )
 from qualtran.cirq_interop.testing import GateHelper
 from qualtran.testing import assert_valid_bloq_decomposition, execute_notebook
@@ -30,6 +31,9 @@
 
 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)
 
 
@@ -57,10 +61,17 @@ def test_state_prep_alias_sampling_symb():
     np.testing.assert_allclose(concrete_t_counts, symb_t_counts, rtol=1e-4)
 
 
-def assert_state_preparation_valid_for_coefficient(lcu_coefficients: np.ndarray, epsilon: float):
-    gate = StatePreparationAliasSampling.from_lcu_probs(
-        lcu_probabilities=lcu_coefficients.tolist(), probability_epsilon=epsilon
-    )
+def assert_state_preparation_valid_for_coefficient(
+    lcu_coefficients: np.ndarray, epsilon: float, *, sparse: bool = False
+):
+    if sparse:
+        gate = SparseStatePreparationAliasSampling.from_lcu_probs(
+            lcu_probabilities=lcu_coefficients.tolist(), probability_epsilon=epsilon
+        )
+    else:
+        gate = StatePreparationAliasSampling.from_lcu_probs(
+            lcu_probabilities=lcu_coefficients.tolist(), probability_epsilon=epsilon
+        )
 
     assert_valid_bloq_decomposition(gate)
     _ = gate.call_graph()
@@ -143,6 +154,15 @@ def test_state_preparation_via_coherent_alias_sampling_diagram():
     )
 
 
+def test_sparse_state_preparation_via_coherent_alias_for_0_mu():
+    one_shot_coeffs = np.zeros(16)
+    one_shot_coeffs[0] = 1
+    assert_state_preparation_valid_for_coefficient(one_shot_coeffs, 2e-1, sparse=True)
+
+    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)
+
+
 @pytest.mark.notebook
 def test_notebook():
     execute_notebook('state_preparation_alias_sampling')
diff --git a/qualtran/serialization/resolver_dict.py b/qualtran/serialization/resolver_dict.py
index 8db239534..48ec91cf0 100644
--- a/qualtran/serialization/resolver_dict.py
+++ b/qualtran/serialization/resolver_dict.py
@@ -323,6 +323,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,

From 189aff9c6586ac49bd82c37ad7fe7bc4ae67faf0 Mon Sep 17 00:00:00 2001
From: anurudhp <anurudhp97@gmail.com>
Date: Thu, 13 Jun 2024 15:07:19 -0700
Subject: [PATCH 06/19] mypy

---
 .../state_preparation/state_preparation_alias_sampling_test.py  | 2 ++
 1 file changed, 2 insertions(+)

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 cd7c09da6..808f16b0a 100644
--- a/qualtran/bloqs/state_preparation/state_preparation_alias_sampling_test.py
+++ b/qualtran/bloqs/state_preparation/state_preparation_alias_sampling_test.py
@@ -17,6 +17,7 @@
 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,
@@ -64,6 +65,7 @@ def test_state_prep_alias_sampling_symb():
 def assert_state_preparation_valid_for_coefficient(
     lcu_coefficients: np.ndarray, epsilon: float, *, sparse: bool = False
 ):
+    gate: Bloq
     if sparse:
         gate = SparseStatePreparationAliasSampling.from_lcu_probs(
             lcu_probabilities=lcu_coefficients.tolist(), probability_epsilon=epsilon

From f356072cb939825365f0746fd74b8ef840d37764 Mon Sep 17 00:00:00 2001
From: anurudhp <anurudhp97@gmail.com>
Date: Thu, 13 Jun 2024 15:26:19 -0700
Subject: [PATCH 07/19] fix decomp

---
 .../state_preparation_alias_sampling.py       | 56 ++++++++-----------
 .../state_preparation_alias_sampling_test.py  |  4 --
 2 files changed, 23 insertions(+), 37 deletions(-)

diff --git a/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py b/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py
index 48ecb6b36..25367ad63 100644
--- a/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py
+++ b/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py
@@ -408,49 +408,39 @@ def qrom_bloq(self) -> QROM:
             (self.selection_bitsize, self.alternates_bitsize, self.keep_bitsize),
         )
 
-    def build_composite_bloq(
-        self,
-        bb: 'BloqBuilder',
-        selection: 'SoquetT',
-        sparse_index: 'SoquetT',
-        alt: 'SoquetT',
-        keep: 'SoquetT',
-        sigma_mu: 'SoquetT',
-        less_than_equal: 'SoquetT',
-    ) -> Dict[str, 'SoquetT']:
-        sparse_index = bb.add(
-            PrepareUniformSuperposition(self.n_nonzero_coeff), target=sparse_index
+    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']
         )
         if self.mu == 0:
-            sparse_index, selection = bb.add(
+            soqs['sparse_index'], soqs['selection'] = bb.add(
                 QROM((self.index,), (self.sparse_index_bitsize,), (self.selection_bitsize,)),
-                selection=sparse_index,
-                target0_=selection,
+                selection=soqs['sparse_index'],
+                target0_=soqs['selection'],
             )
         else:
-            sigma_mu = bb.add(OnEach(self.sigma_mu_bitsize, Hadamard()), q=sigma_mu)
-            sparse_index, selection, alt, keep = bb.add_t(
+            soqs['sigma_mu'] = bb.add(OnEach(self.sigma_mu_bitsize, Hadamard()), q=soqs['sigma_mu'])
+            soqs['sparse_index'], soqs['selection'], soqs['alt'], soqs['keep'] = bb.add_t(
                 self.qrom_bloq,
-                selection=sparse_index,
-                target0_=selection,
-                target1_=alt,
-                target2_=keep,
+                selection=soqs['sparse_index'],
+                target0_=soqs['selection'],
+                target1_=soqs['alt'],
+                target2_=soqs['keep'],
             )
-            keep, sigma_mu, less_than_equal = bb.add_t(
-                LessThanEqual(self.mu, self.mu), x=keep, y=sigma_mu, target=less_than_equal
+            soqs['keep'], soqs['sigma_mu'], less_than_equal = bb.add_t(
+                LessThanEqual(self.mu, self.mu),
+                x=soqs['keep'],
+                y=soqs['sigma_mu'],
+                target=soqs['less_than_equal'],
             )
-            less_than_equal, alt, selection = bb.add_t(
-                CSwap(self.selection_bitsize), ctrl=less_than_equal, x=alt, y=selection
+            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 {
-            'selection': selection,
-            'sparse_index': sparse_index,
-            'alt': alt,
-            'keep': keep,
-            'sigma_mu': sigma_mu,
-            'less_than_equal': less_than_equal,
-        }
+        return soqs
 
 
 @bloq_example(generalizer=[cirq_to_bloqs, ignore_split_join, ignore_cliffords])
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 808f16b0a..3d8e0576c 100644
--- a/qualtran/bloqs/state_preparation/state_preparation_alias_sampling_test.py
+++ b/qualtran/bloqs/state_preparation/state_preparation_alias_sampling_test.py
@@ -157,10 +157,6 @@ def test_state_preparation_via_coherent_alias_sampling_diagram():
 
 
 def test_sparse_state_preparation_via_coherent_alias_for_0_mu():
-    one_shot_coeffs = np.zeros(16)
-    one_shot_coeffs[0] = 1
-    assert_state_preparation_valid_for_coefficient(one_shot_coeffs, 2e-1, sparse=True)
-
     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)
 

From 76d52e724595b7c12dec5df4de6d45f92da3af05 Mon Sep 17 00:00:00 2001
From: anurudhp <anurudhp97@gmail.com>
Date: Thu, 13 Jun 2024 15:48:25 -0700
Subject: [PATCH 08/19] fix validity test

---
 .../state_preparation_alias_sampling.py            |  2 +-
 .../state_preparation_alias_sampling_test.py       | 14 +++++++-------
 2 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py b/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py
index 25367ad63..0ae775e5c 100644
--- a/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py
+++ b/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py
@@ -427,7 +427,7 @@ def build_composite_bloq(self, bb: 'BloqBuilder', **soqs: 'SoquetT') -> Dict[str
                 target1_=soqs['alt'],
                 target2_=soqs['keep'],
             )
-            soqs['keep'], soqs['sigma_mu'], less_than_equal = bb.add_t(
+            soqs['keep'], soqs['sigma_mu'], soqs['less_than_equal'] = bb.add_t(
                 LessThanEqual(self.mu, self.mu),
                 x=soqs['keep'],
                 y=soqs['sigma_mu'],
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 3d8e0576c..246026125 100644
--- a/qualtran/bloqs/state_preparation/state_preparation_alias_sampling_test.py
+++ b/qualtran/bloqs/state_preparation/state_preparation_alias_sampling_test.py
@@ -63,7 +63,7 @@ def test_state_prep_alias_sampling_symb():
 
 
 def assert_state_preparation_valid_for_coefficient(
-    lcu_coefficients: np.ndarray, epsilon: float, *, sparse: bool = False
+    lcu_coefficients: np.ndarray, epsilon: float, *, sparse: bool = False, atol: float = 1e-6
 ):
     gate: Bloq
     if sparse:
@@ -90,14 +90,14 @@ def assert_state_preparation_valid_for_coefficient(
     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():

From b2a0fc2ed3570bc1ff7761eb4f792ad7d389f1ed Mon Sep 17 00:00:00 2001
From: anurudhp <anurudhp97@gmail.com>
Date: Thu, 13 Jun 2024 15:52:09 -0700
Subject: [PATCH 09/19] docstring

---
 .../state_preparation/state_preparation_alias_sampling.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py b/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py
index 0ae775e5c..25af8421a 100644
--- a/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py
+++ b/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py
@@ -247,15 +247,15 @@ class SparseStatePreparationAliasSampling(PrepareOracle):
 
     where $\mathrm{ind}_l$ is the index of the $l$-th non-zero coefficient,
     and the probabilities $p_\ell$ are $\mu$-bit binary approximations to the true values and
-    where the temporary register must be treated with care, see the details in Section III.D. of
-    the reference [2].
+    where the temporary register must be treated with care, see the details in Section 5 of
+    reference [1] and Section III.D. of the reference [2].
 
     The preparation is equivalent to [classical alias sampling]
     (https://en.wikipedia.org/wiki/Alias_method): we sample `l` with probability `p[l]` by first
     selecting `l` uniformly at random and then returning `ind[l]` with probability `keep[l] / 2**mu`;
     otherwise returning `alt[l]`.
-    This bloq is nearly identical to :class:`StatePreparationByAliasSampling`, except that this loads the
-    non-zero coefficient indices as well from the QROM.
+    This bloq is nearly identical to :class:`StatePreparationByAliasSampling`, except that this loads
+    the non-zero coefficient indices as well from the QROM.
 
     Signature:
         selection: The input/output register $|\mathrm{ind}_l\rangle$ of size lg(L) where the desired

From 18193eaf4b108ad16fa61fb444915c39c6077a85 Mon Sep 17 00:00:00 2001
From: anurudhp <anurudhp97@gmail.com>
Date: Thu, 13 Jun 2024 15:55:33 -0700
Subject: [PATCH 10/19] one more test

---
 .../state_preparation/state_preparation_alias_sampling_test.py | 3 +++
 1 file changed, 3 insertions(+)

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 246026125..b4e1a30fd 100644
--- a/qualtran/bloqs/state_preparation/state_preparation_alias_sampling_test.py
+++ b/qualtran/bloqs/state_preparation/state_preparation_alias_sampling_test.py
@@ -160,6 +160,9 @@ def test_sparse_state_preparation_via_coherent_alias_for_0_mu():
     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)
+
 
 @pytest.mark.notebook
 def test_notebook():

From 5424590326d49d3d53bd03652535547f58a197b6 Mon Sep 17 00:00:00 2001
From: anurudhp <anurudhp97@gmail.com>
Date: Thu, 13 Jun 2024 16:23:46 -0700
Subject: [PATCH 11/19] add tolerance parameter for nonzero check

---
 .../state_preparation_alias_sampling.py                  | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py b/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py
index 25af8421a..e47a4d922 100644
--- a/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py
+++ b/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py
@@ -295,7 +295,11 @@ class SparseStatePreparationAliasSampling(PrepareOracle):
 
     @classmethod
     def from_lcu_probs(
-        cls, lcu_probabilities: Sequence[float], *, probability_epsilon: float = 1.0e-5
+        cls,
+        lcu_probabilities: Sequence[float],
+        *,
+        probability_epsilon: float = 1.0e-5,
+        nonzero_epsilon: float = 1e-6,
     ) -> 'SparseStatePreparationAliasSampling':
         """Factory to construct the state preparation gate for a given set of LCU coefficients.
 
@@ -305,13 +309,14 @@ def from_lcu_probs(
                 (which sets mu size and keep/alt integers).
                 See `qualtran.linalg.lcu_util.preprocess_lcu_coefficients_for_reversible_sampling`
                 for more information.
+            nonzero_epsilon: minimum value for a probability entry to be considered non-zero.
         """
         alt, keep, mu = preprocess_lcu_coefficients_for_reversible_sampling(
             lcu_coefficients=lcu_probabilities, epsilon=probability_epsilon
         )
         N = len(lcu_probabilities)
 
-        is_nonzero = ~np.isclose(lcu_probabilities, 0)
+        is_nonzero = ~np.isclose(lcu_probabilities, 0, atol=nonzero_epsilon)
         index = np.arange(N)[is_nonzero]
         alt = np.array(alt)[is_nonzero]
         keep = np.array(keep)[is_nonzero]

From 90434f8b273f6f1b2f790ec776c0ca80a5e8e7bf Mon Sep 17 00:00:00 2001
From: Anurudh Peduri <anurudhp97@gmail.com>
Date: Mon, 8 Jul 2024 15:56:40 -0700
Subject: [PATCH 12/19] update docstring and parameter names

---
 .../state_preparation_alias_sampling.py       | 113 +++++++++---------
 .../state_preparation_alias_sampling_test.py  |   4 +-
 2 files changed, 58 insertions(+), 59 deletions(-)

diff --git a/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py b/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py
index e47a4d922..84f9e26fc 100644
--- a/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py
+++ b/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py
@@ -20,7 +20,7 @@
 largest absolute error that one can tolerate in the prepared amplitudes.
 """
 from functools import cached_property
-from typing import Dict, Iterator, Sequence, Set, Tuple, TYPE_CHECKING, Union
+from typing import Iterator, Sequence, Set, Tuple, TYPE_CHECKING, Union
 
 import attrs
 import cirq
@@ -237,25 +237,25 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']:
 
 @attrs.frozen
 class SparseStatePreparationAliasSampling(PrepareOracle):
-    r"""Initialize a state with $L$ unique non-zero coefficients using coherent alias sampling.
+    r"""Initialize a $d$-sparse state over $L$ indices using coherent alias sampling.
 
     In particular, we take the zero state to:
 
     $$
-    \sum_{l=0}^{L-1} \sqrt{p_l} |\mathrm{ind}_l\rangle |\mathrm{temp}_\ell\rangle
+        \sum_{j=0}^{d-1} \sqrt{p_{\mathrm{ind}_j}} |\mathrm{ind}_j\rangle |\mathrm{temp}_j\rangle
     $$
 
-    where $\mathrm{ind}_l$ is the index of the $l$-th non-zero coefficient,
-    and the probabilities $p_\ell$ are $\mu$-bit binary approximations to the true values and
-    where the temporary register must be treated with care, see the details in Section 5 of
-    reference [1] and Section III.D. of the reference [2].
+    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.
 
-    The preparation is equivalent to [classical alias sampling]
-    (https://en.wikipedia.org/wiki/Alias_method): we sample `l` with probability `p[l]` by first
-    selecting `l` uniformly at random and then returning `ind[l]` with probability `keep[l] / 2**mu`;
-    otherwise returning `alt[l]`.
-    This bloq is nearly identical to :class:`StatePreparationByAliasSampling`, except that this loads
-    the non-zero coefficient indices as well from the QROM.
 
     Signature:
         selection: The input/output register $|\mathrm{ind}_l\rangle$ of size lg(L) where the desired
@@ -263,20 +263,19 @@ class SparseStatePreparationAliasSampling(PrepareOracle):
         temp: Work space comprised of sub signature:
             - sigma: 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.
             - one bit for the result of the comparison.
 
     This gate corresponds to the following operations:
-     - UNIFORM_L on the sparse_index register
-     - H^mu on the sigma register
-     - QROM addressed by the selection register into the alt and keep signature.
-     - LessThanEqualGate comparing the keep and sigma signature.
-     - Coherent swap between the selection register and alt register if the comparison
-       returns True.
+     - 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) + 2 mu + 1) work qubits + log(L) ancillas for QROM.
-    The 1 ancilla in work qubits is for the `LessThanEqualGate` followed by coherent swap.
+    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)
@@ -291,32 +290,44 @@ class SparseStatePreparationAliasSampling(PrepareOracle):
     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_lcu_coeffs: SymbolicFloat
+    sum_of_unnormalized_probabilities: SymbolicFloat
+
+    @cached_property
+    def junk_registers(self) -> Tuple[Register, ...]:
+        return tuple(
+            Signature.build(
+                sigma_mu=self.sigma_mu_bitsize,
+                sparse_index=self.sparse_index_bitsize,
+                alt=self.alternates_bitsize,
+                keep=self.keep_bitsize,
+                less_than_equal=1,
+            )
+        )
 
     @classmethod
-    def from_lcu_probs(
+    def from_dense_probabilities(
         cls,
-        lcu_probabilities: Sequence[float],
+        unnormalized_probabilities: Sequence[float],
         *,
-        probability_epsilon: float = 1.0e-5,
-        nonzero_epsilon: float = 1e-6,
+        precision: float = 1.0e-5,
+        nonzero_threshold: float = 1e-6,
     ) -> 'SparseStatePreparationAliasSampling':
-        """Factory to construct the state preparation gate for a given set of LCU coefficients.
+        """Factory to construct the state preparation gate for a given set of probability coefficients.
 
         Args:
-            lcu_probabilities: The LCU coefficients.
-            probability_epsilon: The desired accuracy to represent each probability
+            unnormalized_probabilities: A dense list of all probabilities (i.e. including 0s)
+            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.
-            nonzero_epsilon: minimum value for a probability entry to be considered non-zero.
+            nonzero_threshold: minimum value for a probability entry to be considered non-zero.
         """
         alt, keep, mu = preprocess_lcu_coefficients_for_reversible_sampling(
-            lcu_coefficients=lcu_probabilities, epsilon=probability_epsilon
+            lcu_coefficients=unnormalized_probabilities, epsilon=precision
         )
-        N = len(lcu_probabilities)
+        N = len(unnormalized_probabilities)
 
-        is_nonzero = ~np.isclose(lcu_probabilities, 0, atol=nonzero_epsilon)
+        is_nonzero = ~np.isclose(unnormalized_probabilities, 0, atol=nonzero_threshold)
         index = np.arange(N)[is_nonzero]
         alt = np.array(alt)[is_nonzero]
         keep = np.array(keep)[is_nonzero]
@@ -327,7 +338,7 @@ def from_lcu_probs(
             alt=alt,
             keep=keep,
             mu=mu,
-            sum_of_lcu_coeffs=sum(abs(x) for x in lcu_probabilities),
+            sum_of_unnormalized_probabilities=sum(abs(x) for x in unnormalized_probabilities),
         )
 
     @classmethod
@@ -335,22 +346,22 @@ def from_n_coeff(
         cls,
         n_coeff: SymbolicInt,
         n_nonzero_coeff: SymbolicInt,
-        sum_of_lcu_coeffs: SymbolicFloat,
+        sum_of_terms: SymbolicFloat,
         *,
-        probability_epsilon: SymbolicFloat = 1.0e-5,
+        precision: SymbolicFloat = 1.0e-5,
     ) -> 'SparseStatePreparationAliasSampling':
-        """Factory to construct the state preparation gate for symbolic number of LCU coefficients.
+        """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_lcu_coeffs: Sum of absolute values of coefficients of the prepared state.
-            probability_epsilon: The desired accuracy to represent each probability
+            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, probability_epsilon)
+        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)),
@@ -358,7 +369,7 @@ def from_n_coeff(
             alt=Shaped((n_nonzero_coeff,)),
             keep=Shaped((n_nonzero_coeff,)),
             mu=mu,
-            sum_of_lcu_coeffs=sum_of_lcu_coeffs,
+            sum_of_unnormalized_probabilities=sum_of_terms,
         )
 
     @property
@@ -371,7 +382,7 @@ def n_nonzero_coeff(self) -> SymbolicInt:
 
     @cached_property
     def l1_norm_of_coeffs(self) -> 'SymbolicFloat':
-        return self.sum_of_lcu_coeffs
+        return self.sum_of_unnormalized_probabilities
 
     @cached_property
     def sigma_mu_bitsize(self) -> SymbolicInt:
@@ -393,18 +404,6 @@ def keep_bitsize(self) -> SymbolicInt:
     def selection_bitsize(self) -> SymbolicInt:
         return total_bits(self.selection_registers)
 
-    @cached_property
-    def junk_registers(self) -> Tuple[Register, ...]:
-        return tuple(
-            Signature.build(
-                sigma_mu=self.sigma_mu_bitsize,
-                sparse_index=self.sparse_index_bitsize,
-                alt=self.alternates_bitsize,
-                keep=self.keep_bitsize,
-                less_than_equal=1,
-            )
-        )
-
     @cached_property
     def qrom_bloq(self) -> QROM:
         return QROM(
@@ -413,7 +412,7 @@ def qrom_bloq(self) -> QROM:
             (self.selection_bitsize, self.alternates_bitsize, self.keep_bitsize),
         )
 
-    def build_composite_bloq(self, bb: 'BloqBuilder', **soqs: 'SoquetT') -> Dict[str, 'SoquetT']:
+    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']
         )
@@ -473,8 +472,8 @@ def _state_prep_alias_symb() -> StatePreparationAliasSampling:
 def _sparse_state_prep_alias() -> SparseStatePreparationAliasSampling:
     coeffs = [1.0, 0, 0, 1, 0, 3, 0, 2, 0]
     mu = 3
-    state_prep_alias = SparseStatePreparationAliasSampling.from_lcu_probs(
-        coeffs, probability_epsilon=2**-mu / len(coeffs)
+    state_prep_alias = SparseStatePreparationAliasSampling.from_dense_probabilities(
+        coeffs, precision=2**-mu / len(coeffs)
     )
     return state_prep_alias
 
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 b4e1a30fd..2ebcf96ca 100644
--- a/qualtran/bloqs/state_preparation/state_preparation_alias_sampling_test.py
+++ b/qualtran/bloqs/state_preparation/state_preparation_alias_sampling_test.py
@@ -67,8 +67,8 @@ def assert_state_preparation_valid_for_coefficient(
 ):
     gate: Bloq
     if sparse:
-        gate = SparseStatePreparationAliasSampling.from_lcu_probs(
-            lcu_probabilities=lcu_coefficients.tolist(), probability_epsilon=epsilon
+        gate = SparseStatePreparationAliasSampling.from_dense_probabilities(
+            unnormalized_probabilities=lcu_coefficients.tolist(), precision=epsilon
         )
     else:
         gate = StatePreparationAliasSampling.from_lcu_probs(

From ae76c04abd3e9f3830fa138f9626df675c436b6a Mon Sep 17 00:00:00 2001
From: Anurudh Peduri <anurudhp97@gmail.com>
Date: Mon, 15 Jul 2024 23:33:13 -0700
Subject: [PATCH 13/19] call preprocess only on nonzero values

---
 .../state_preparation_alias_sampling.py       | 74 ++++++++++---------
 .../state_preparation_alias_sampling_test.py  |  6 +-
 2 files changed, 42 insertions(+), 38 deletions(-)

diff --git a/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py b/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py
index e6e4174d1..115eae006 100644
--- a/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py
+++ b/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py
@@ -45,7 +45,7 @@
     ignore_cliffords,
     ignore_split_join,
 )
-from qualtran.symbolics import bit_length, Shaped, slen, SymbolicFloat, SymbolicInt
+from qualtran.symbolics import bit_length, Shaped, slen, SymbolicFloat, SymbolicInt, is_symbolic
 
 if TYPE_CHECKING:
     from qualtran import BloqBuilder, SoquetT
@@ -318,6 +318,10 @@ class SparseStatePreparationAliasSampling(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 at least 1")
+
     @cached_property
     def junk_registers(self) -> Tuple[Register, ...]:
         return tuple(
@@ -348,23 +352,28 @@ def from_dense_probabilities(
                 for more information.
             nonzero_threshold: minimum value for a probability entry to be considered non-zero.
         """
-        alt, keep, mu = preprocess_probabilities_for_reversible_sampling(
-            unnormalized_probabilities=unnormalized_probabilities, epsilon=precision
-        )
-        N = len(unnormalized_probabilities)
+        if not all(x >= 0 for x in unnormalized_probabilities):
+            raise ValueError(f"{cls} expects only non-negative probabilities")
 
+        N = len(unnormalized_probabilities)
         is_nonzero = ~np.isclose(unnormalized_probabilities, 0, atol=nonzero_threshold)
+        nonzero_values = np.array(unnormalized_probabilities)[is_nonzero].tolist()
+        d = len(nonzero_values)  # sparsity
+
+        alt_compressed, keep, mu = preprocess_probabilities_for_reversible_sampling(
+            unnormalized_probabilities=nonzero_values, epsilon=precision
+        )
+
         index = np.arange(N)[is_nonzero]
-        alt = np.array(alt)[is_nonzero]
-        keep = np.array(keep)[is_nonzero]
+        alt = np.array([index[idx] for idx in alt_compressed])
 
         return cls(
             selection_registers=Register('selection', BoundedQUInt((N - 1).bit_length(), N)),
             index=index,
             alt=alt,
-            keep=keep,
+            keep=np.array(keep),
             mu=mu,
-            sum_of_unnormalized_probabilities=sum(abs(x) for x in unnormalized_probabilities),
+            sum_of_unnormalized_probabilities=np.sum(unnormalized_probabilities),
         )
 
     @classmethod
@@ -442,33 +451,26 @@ def build_composite_bloq(self, bb: 'BloqBuilder', **soqs: 'SoquetT') -> dict[str
         soqs['sparse_index'] = bb.add(
             PrepareUniformSuperposition(self.n_nonzero_coeff), target=soqs['sparse_index']
         )
-        if self.mu == 0:
-            soqs['sparse_index'], soqs['selection'] = bb.add(
-                QROM((self.index,), (self.sparse_index_bitsize,), (self.selection_bitsize,)),
-                selection=soqs['sparse_index'],
-                target0_=soqs['selection'],
-            )
-        else:
-            soqs['sigma_mu'] = bb.add(OnEach(self.sigma_mu_bitsize, 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'],
-            )
+        soqs['sigma_mu'] = bb.add(OnEach(self.sigma_mu_bitsize, 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
 
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 f29649ab0..48c7dd0ef 100644
--- a/qualtran/bloqs/state_preparation/state_preparation_alias_sampling_test.py
+++ b/qualtran/bloqs/state_preparation/state_preparation_alias_sampling_test.py
@@ -93,7 +93,9 @@ def assert_state_preparation_valid_for_coefficient(
     coeff_precision = epsilon * np.sum(np.abs(lcu_coefficients))
     if sparse:
         gate = SparseStatePreparationAliasSampling.from_dense_probabilities(
-            unnormalized_probabilities=lcu_coefficients.tolist(), precision=coeff_precision
+            unnormalized_probabilities=lcu_coefficients.tolist(),
+            precision=coeff_precision,
+            nonzero_threshold=atol,
         )
     else:
         gate = StatePreparationAliasSampling.from_probabilities(
@@ -181,7 +183,7 @@ def test_state_preparation_via_coherent_alias_sampling_diagram():
     )
 
 
-def test_sparse_state_preparation_via_coherent_alias_for_0_mu():
+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)
 

From f10d7426fc7c0ff9c64660051e07005b916f4a3b Mon Sep 17 00:00:00 2001
From: Anurudh Peduri <anurudhp97@gmail.com>
Date: Mon, 15 Jul 2024 23:38:00 -0700
Subject: [PATCH 14/19] cleanup

---
 .../state_preparation_alias_sampling.py       | 26 +++++--------------
 1 file changed, 7 insertions(+), 19 deletions(-)

diff --git a/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py b/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py
index 115eae006..8a740489e 100644
--- a/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py
+++ b/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py
@@ -326,10 +326,10 @@ def __attrs_post_init__(self):
     def junk_registers(self) -> Tuple[Register, ...]:
         return tuple(
             Signature.build(
-                sigma_mu=self.sigma_mu_bitsize,
+                sigma_mu=self.mu,
                 sparse_index=self.sparse_index_bitsize,
-                alt=self.alternates_bitsize,
-                keep=self.keep_bitsize,
+                alt=self.selection_bitsize,
+                keep=self.mu,
                 less_than_equal=1,
             )
         )
@@ -420,38 +420,26 @@ def l1_norm_of_coeffs(self) -> 'SymbolicFloat':
         return self.sum_of_unnormalized_probabilities
 
     @cached_property
-    def sigma_mu_bitsize(self) -> SymbolicInt:
-        return self.mu
+    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 alternates_bitsize(self) -> SymbolicInt:
-        return total_bits(self.selection_registers)
-
-    @cached_property
-    def keep_bitsize(self) -> SymbolicInt:
-        return self.mu
-
-    @cached_property
-    def selection_bitsize(self) -> SymbolicInt:
-        return total_bits(self.selection_registers)
-
     @cached_property
     def qrom_bloq(self) -> QROM:
         return QROM(
             (self.index, self.alt, self.keep),
             (self.sparse_index_bitsize,),
-            (self.selection_bitsize, self.alternates_bitsize, self.keep_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.sigma_mu_bitsize, Hadamard()), q=soqs['sigma_mu'])
+        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'],

From 83fb615e3bab8b47d71b3653d137745ad44566d3 Mon Sep 17 00:00:00 2001
From: Anurudh Peduri <anurudhp97@gmail.com>
Date: Mon, 15 Jul 2024 23:56:50 -0700
Subject: [PATCH 15/19] support sparse state input as map

---
 .../state_preparation_alias_sampling.py       | 83 +++++++++++++------
 .../state_preparation_alias_sampling_test.py  |  2 +
 2 files changed, 59 insertions(+), 26 deletions(-)

diff --git a/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py b/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py
index 8a740489e..645d004c7 100644
--- a/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py
+++ b/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py
@@ -45,7 +45,7 @@
     ignore_cliffords,
     ignore_split_join,
 )
-from qualtran.symbolics import bit_length, Shaped, slen, SymbolicFloat, SymbolicInt, is_symbolic
+from qualtran.symbolics import bit_length, is_symbolic, Shaped, slen, SymbolicFloat, SymbolicInt
 
 if TYPE_CHECKING:
     from qualtran import BloqBuilder, SoquetT
@@ -335,45 +335,65 @@ def junk_registers(self) -> Tuple[Register, ...]:
         )
 
     @classmethod
-    def from_dense_probabilities(
-        cls,
-        unnormalized_probabilities: Sequence[float],
-        *,
-        precision: float = 1.0e-5,
-        nonzero_threshold: float = 1e-6,
+    def from_sparse_dict(
+        cls, unnormalized_probabilities: dict[int, float], N: int, *, precision: float = 1.0e-5
     ) -> 'SparseStatePreparationAliasSampling':
-        """Factory to construct the state preparation gate for a given set of probability coefficients.
+        """Construct the state preparation gate for a given dictionary of non-zero probabilities.
 
         Args:
-            unnormalized_probabilities: A dense list of all probabilities (i.e. including 0s)
+            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_lcu_coefficients_for_reversible_sampling`
+                See `qualtran.linalg.lcu_util.preprocess_probabilities_for_reversible_sampling`
                 for more information.
-            nonzero_threshold: minimum value for a probability entry to be considered non-zero.
         """
-        if not all(x >= 0 for x in unnormalized_probabilities):
+        if not all(x >= 0 for x in unnormalized_probabilities.values()):
             raise ValueError(f"{cls} expects only non-negative probabilities")
-
-        N = len(unnormalized_probabilities)
-        is_nonzero = ~np.isclose(unnormalized_probabilities, 0, atol=nonzero_threshold)
-        nonzero_values = np.array(unnormalized_probabilities)[is_nonzero].tolist()
-        d = len(nonzero_values)  # sparsity
+        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=nonzero_values, epsilon=precision
+            unnormalized_probabilities=list(unnormalized_probabilities.values()), epsilon=precision
         )
 
-        index = np.arange(N)[is_nonzero]
-        alt = np.array([index[idx] for idx in alt_compressed])
+        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=index,
-            alt=alt,
+            index=np.array(index),
+            alt=np.array(alt),
             keep=np.array(keep),
             mu=mu,
-            sum_of_unnormalized_probabilities=np.sum(unnormalized_probabilities),
+            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
@@ -486,14 +506,25 @@ def _state_prep_alias_symb() -> StatePreparationAliasSampling:
 
 @bloq_example(generalizer=[cirq_to_bloqs, ignore_split_join, ignore_cliffords])
 def _sparse_state_prep_alias() -> SparseStatePreparationAliasSampling:
-    coeffs = [1.0, 0, 0, 1, 0, 3, 0, 2, 0]
+    coeff_map = {0: 1.0, 3: 1.0, 5: 3.0, 7: 2.0}
+    N = 9
     mu = 3
-    state_prep_alias = SparseStatePreparationAliasSampling.from_dense_probabilities(
-        coeffs, precision=2**-mu / len(coeffs)
+    state_prep_alias = SparseStatePreparationAliasSampling.from_sparse_dict(
+        coeff_map, N, precision=2**-mu / len(coeff_map)
     )
     return 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
+
+
 _STATE_PREP_ALIAS_DOC = BloqDocSpec(
     bloq_cls=StatePreparationAliasSampling,
     import_line='from qualtran.bloqs.state_preparation import StatePreparationAliasSampling',
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 48c7dd0ef..0a2583331 100644
--- a/qualtran/bloqs/state_preparation/state_preparation_alias_sampling_test.py
+++ b/qualtran/bloqs/state_preparation/state_preparation_alias_sampling_test.py
@@ -21,6 +21,7 @@
 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,
     _state_prep_alias,
     _state_prep_alias_symb,
     SparseStatePreparationAliasSampling,
@@ -36,6 +37,7 @@ def test_state_prep_alias_sampling_autotest(bloq_autotester):
 
 def test_sparse_state_prep_alias_sampling_autotest(bloq_autotester):
     bloq_autotester(_sparse_state_prep_alias)
+    bloq_autotester(_sparse_state_prep_alias_from_list)
 
 
 def test_mu_from_precision():

From 58df7f2ee8335ea8adfaa1fcaa1b144615b48906 Mon Sep 17 00:00:00 2001
From: Anurudh Peduri <anurudhp97@gmail.com>
Date: Mon, 15 Jul 2024 23:59:57 -0700
Subject: [PATCH 16/19] move examples

---
 qualtran/bloqs/state_preparation/__init__.py  |  1 +
 .../state_preparation_alias_sampling.py       | 56 ++++++++++---------
 2 files changed, 32 insertions(+), 25 deletions(-)

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.py b/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py
index 645d004c7..7faf03dca 100644
--- a/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py
+++ b/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py
@@ -261,6 +261,34 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']:
         }
 
 
+@bloq_example(generalizer=[cirq_to_bloqs, ignore_split_join, ignore_cliffords])
+def _state_prep_alias() -> StatePreparationAliasSampling:
+    coeffs = [1.0, 1, 3, 2]
+    mu = 3
+    state_prep_alias = StatePreparationAliasSampling.from_probabilities(
+        coeffs, precision=2**-mu / len(coeffs) * sum(coeffs)
+    )
+    return state_prep_alias
+
+
+@bloq_example(generalizer=[cirq_to_bloqs, ignore_split_join, ignore_cliffords])
+def _state_prep_alias_symb() -> StatePreparationAliasSampling:
+    import sympy
+
+    n_coeffs, sum_coeff, eps = sympy.symbols(r"L \lambda \epsilon")
+    state_prep_alias_symb = StatePreparationAliasSampling.from_n_coeff(
+        n_coeffs, sum_coeff, precision=eps
+    )
+    return state_prep_alias_symb
+
+
+_STATE_PREP_ALIAS_DOC = BloqDocSpec(
+    bloq_cls=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.
@@ -483,27 +511,6 @@ def build_composite_bloq(self, bb: 'BloqBuilder', **soqs: 'SoquetT') -> dict[str
         return soqs
 
 
-@bloq_example(generalizer=[cirq_to_bloqs, ignore_split_join, ignore_cliffords])
-def _state_prep_alias() -> StatePreparationAliasSampling:
-    coeffs = [1.0, 1, 3, 2]
-    mu = 3
-    state_prep_alias = StatePreparationAliasSampling.from_probabilities(
-        coeffs, precision=2**-mu / len(coeffs) * sum(coeffs)
-    )
-    return state_prep_alias
-
-
-@bloq_example(generalizer=[cirq_to_bloqs, ignore_split_join, ignore_cliffords])
-def _state_prep_alias_symb() -> StatePreparationAliasSampling:
-    import sympy
-
-    n_coeffs, sum_coeff, eps = sympy.symbols(r"L \lambda \epsilon")
-    state_prep_alias_symb = StatePreparationAliasSampling.from_n_coeff(
-        n_coeffs, sum_coeff, precision=eps
-    )
-    return state_prep_alias_symb
-
-
 @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}
@@ -525,8 +532,7 @@ def _sparse_state_prep_alias_from_list() -> SparseStatePreparationAliasSampling:
     return sparse_state_prep_alias_from_list
 
 
-_STATE_PREP_ALIAS_DOC = BloqDocSpec(
-    bloq_cls=StatePreparationAliasSampling,
-    import_line='from qualtran.bloqs.state_preparation import StatePreparationAliasSampling',
-    examples=(_state_prep_alias, _state_prep_alias_symb),
+_SPARSE_STATE_PREP_ALIAS_DOC = BloqDocSpec(
+    bloq_cls=SparseStatePreparationAliasSampling,
+    examples=(_sparse_state_prep_alias, _sparse_state_prep_alias_from_list),
 )

From 8e4959ffb6de814023e124f615d8103523f2a449 Mon Sep 17 00:00:00 2001
From: Anurudh Peduri <anurudhp97@gmail.com>
Date: Tue, 16 Jul 2024 00:02:23 -0700
Subject: [PATCH 17/19] symbolic example

---
 .../state_preparation_alias_sampling.py         | 17 ++++++++++++++++-
 .../state_preparation_alias_sampling_test.py    |  2 ++
 qualtran/conftest.py                            |  1 +
 3 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py b/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py
index 7faf03dca..694c2b672 100644
--- a/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py
+++ b/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py
@@ -532,7 +532,22 @@ def _sparse_state_prep_alias_from_list() -> SparseStatePreparationAliasSampling:
     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),
+    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 0a2583331..3a280e203 100644
--- a/qualtran/bloqs/state_preparation/state_preparation_alias_sampling_test.py
+++ b/qualtran/bloqs/state_preparation/state_preparation_alias_sampling_test.py
@@ -24,6 +24,7 @@
     _sparse_state_prep_alias_from_list,
     _state_prep_alias,
     _state_prep_alias_symb,
+    _sparse_state_prep_alias_symb,
     SparseStatePreparationAliasSampling,
     StatePreparationAliasSampling,
 )
@@ -38,6 +39,7 @@ def test_state_prep_alias_sampling_autotest(bloq_autotester):
 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():
diff --git a/qualtran/conftest.py b/qualtran/conftest.py
index ff39888a7..66a1933ed 100644
--- a/qualtran/conftest.py
+++ b/qualtran/conftest.py
@@ -102,6 +102,7 @@ def assert_bloq_example_serializes_for_pytest(bloq_ex: BloqExample):
         'product_block_encoding_symb',
         'apply_lth_bloq',
         'phase_block_encoding',
+        'sparse_state_prep_alias_symb',  # cannot serialize Shaped
     ]:
         pytest.xfail("Skipping serialization test for bloq examples that cannot yet be serialized.")
 

From e80f3284df78684e13c16f40ea7fdb2e6b2d7ecc Mon Sep 17 00:00:00 2001
From: Anurudh Peduri <anurudhp97@gmail.com>
Date: Tue, 16 Jul 2024 00:08:58 -0700
Subject: [PATCH 18/19] notebook

---
 dev_tools/autogenerate-bloqs-notebooks-v2.py  |   3 +-
 .../state_preparation_alias_sampling.ipynb    | 170 ++++++++++++++++++
 .../state_preparation_alias_sampling.py       |  18 +-
 3 files changed, 180 insertions(+), 11 deletions(-)

diff --git a/dev_tools/autogenerate-bloqs-notebooks-v2.py b/dev_tools/autogenerate-bloqs-notebooks-v2.py
index 24fe5050f..8a3408825 100644
--- a/dev_tools/autogenerate-bloqs-notebooks-v2.py
+++ b/dev_tools/autogenerate-bloqs-notebooks-v2.py
@@ -623,7 +623,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/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 694c2b672..200155568 100644
--- a/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py
+++ b/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py
@@ -311,16 +311,14 @@ class SparseStatePreparationAliasSampling(PrepareOracle):
     See :class:`StatePreparationAliasSampling` for an exposition on alias sampling.
 
 
-    Signature:
+    Registers:
         selection: The input/output register $|\mathrm{ind}_l\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`.
-            - 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.
-            - one bit for the result of the comparison.
+        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.
@@ -516,10 +514,10 @@ def _sparse_state_prep_alias() -> SparseStatePreparationAliasSampling:
     coeff_map = {0: 1.0, 3: 1.0, 5: 3.0, 7: 2.0}
     N = 9
     mu = 3
-    state_prep_alias = SparseStatePreparationAliasSampling.from_sparse_dict(
+    sparse_state_prep_alias = SparseStatePreparationAliasSampling.from_sparse_dict(
         coeff_map, N, precision=2**-mu / len(coeff_map)
     )
-    return state_prep_alias
+    return sparse_state_prep_alias
 
 
 @bloq_example(generalizer=[cirq_to_bloqs, ignore_split_join, ignore_cliffords])

From fb893f49282abbcf3de864aa712f87dd479ce9df Mon Sep 17 00:00:00 2001
From: Anurudh Peduri <anurudhp97@gmail.com>
Date: Tue, 16 Jul 2024 00:37:37 -0700
Subject: [PATCH 19/19] test symbolic t complexity

---
 .../state_preparation_alias_sampling_test.py   | 18 +++++++++++++++++-
 1 file changed, 17 insertions(+), 1 deletion(-)

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 3a280e203..e1cbdbb6e 100644
--- a/qualtran/bloqs/state_preparation/state_preparation_alias_sampling_test.py
+++ b/qualtran/bloqs/state_preparation/state_preparation_alias_sampling_test.py
@@ -22,13 +22,14 @@
 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,
-    _sparse_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
 
 
@@ -195,6 +196,21 @@ def test_sparse_state_preparation_via_coherent_alias():
     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')