Alias sampling and QROM state preparation

docs/source/_static/examples/python/hamiltonian_unitary_builder.py

@@ -148,3 +148,18 @@
 print(unitary.get_summary())
 # end-cell-run-lcu
 ################################################################################
+
+################################################################################
+# start-cell-lcu-prepare
+from qdk_chemistry.data import AlgorithmRef
+
+# Load the PREPARE amplitudes with alias sampling instead of dense state preparation
+psp_mapper = create(
+    "circuit_mapper",
+    "prepare_select_prepare",
+    prepare=AlgorithmRef("state_prep", "alias_sampling", bits_precision=6),
+)
+lcu_circuit = psp_mapper.run(unitary)
+print(lcu_circuit.estimate()["logicalCounts"])
+# end-cell-lcu-prepare
+################################################################################

docs/source/_static/examples/python/state_preparation.py

@@ -90,6 +90,6 @@
 from qdk_chemistry.algorithms import registry
 
 print(registry.available("state_prep"))
-# ['sparse_isometry', 'dense_pure_state', 'qiskit_regular_isometry']
+# ['dense_pure_state', 'sparse_isometry', 'alias_sampling', 'qrom', 'qiskit_regular_isometry']
 # end-cell-list-implementations
 ################################################################################

docs/source/user/comprehensive/algorithms/hamiltonian_unitary_builder.rst

@@ -347,13 +347,23 @@ The walk operator has eigenvalues :math:`e^{\pm i \arccos(E_k/\lambda)}` where :
 
 .. rubric:: Example
 
-::
+.. tab:: Python API
 
-    from qdk_chemistry.algorithms import registry
+   .. literalinclude:: ../../../_static/examples/python/hamiltonian_unitary_builder.py
+      :language: python
+      :start-after: # start-cell-run-lcu
+      :end-before: # end-cell-run-lcu
+
+To use alias sampling for the PREPARE oracle, pass an
+:class:`~qdk_chemistry.data.AlgorithmRef` to the
+:class:`~qdk_chemistry.algorithms.circuit_mapper.psp_mapper.PSPMapper`:
 
-    lcu = registry.create("hamiltonian_unitary_builder", "lcu")
-    lcu.settings().update({"quantum_walk": True})
-    unitary = lcu.run(qubit_hamiltonian)
+.. tab:: Python API
+
+   .. literalinclude:: ../../../_static/examples/python/hamiltonian_unitary_builder.py
+      :language: python
+      :start-after: # start-cell-lcu-prepare
+      :end-before: # end-cell-lcu-prepare
 
 The resulting :class:`~qdk_chemistry.data.UnitaryRepresentation` wraps an ``LCUContainer`` containing the Prepare and Select oracles.
 

docs/source/user/comprehensive/algorithms/state_preparation.rst

@@ -208,6 +208,36 @@ This method uses regular isometry synthesis via `Qiskit <https://quantum.cloud.i
 
 For more details on how QDK/Chemistry interfaces with external packages, see the :ref:`plugin system <plugin-system>` documentation.
 
+Alias Sampling
+~~~~~~~~~~~~~~
+
+.. rubric:: Factory name: ``"alias_sampling"``
+
+This method implements the coherent alias sampling oracle of Babbush et al. :cite:`Babbush2018` (section III.D). Given :math:`L` non-negative coefficients :math:`c_\ell`, it prepares
+
+.. math::
+
+   \sum_{\ell} \sqrt{\tilde{p}_\ell} \left| \ell \right\rangle \left| \mathrm{garbage}_\ell \right\rangle ,
+   \qquad \tilde{p}_\ell \approx \frac{c_\ell^2}{\sum_k c_k^2} ,
+
+where :math:`\tilde{p}` is the target distribution discretized to :math:`\mu` bits.
+
+.. warning::
+   This is a **block-encoding subroutine, not a general state preparation for algorithms like QPE**. The index register is left entangled with ancilla, so the output is only meaningful inside an :term:`LCU` or qubitization circuit where :math:`\mathrm{PREPARE}^\dagger` later uncomputes the garbage. Negative coefficients are not supported.
+
+.. rubric:: Settings
+
+.. list-table::
+   :header-rows: 1
+   :widths: 25 25 50
+
+   * - Setting
+     - Type
+     - Description
+   * - ``bits_precision``
+     - int
+     - Number of bits :math:`\mu` of precision for the alias table's keep probabilities. Each prepared probability is within :math:`1/(L 2^{\mu})` of the target for :math:`L` coefficients. The upper bound of 30 is a sanity limit as :math:`2^{-30}` is far below chemical accuracy. Default is 10.
+
 Related classes
 ---------------
 

python/src/qdk_chemistry/algorithms/circuit_mapper/psp_mapper.py

@@ -13,7 +13,7 @@
 from qdk_chemistry.data.circuit import Circuit, QsharpFactoryData
 from qdk_chemistry.data.unitary_representation.base import UnitaryRepresentation
 from qdk_chemistry.data.unitary_representation.containers.base import UnitaryContainer
-from qdk_chemistry.data.unitary_representation.containers.block_encoding import LCUContainer, Select
+from qdk_chemistry.data.unitary_representation.containers.block_encoding import LCUContainer
 from qdk_chemistry.data.unitary_representation.containers.quantum_walk import LCUWalkContainer
 from qdk_chemistry.utils.qsharp import QSHARP_UTILS
 
@@ -90,36 +90,6 @@ def type_name(self) -> str:
         """
         return "circuit_mapper"
 
-    @staticmethod
-    def _build_pauli_select_op(select: Select):
-        """Build the Pauli SELECT Q# operation from a Select data object.
-
-        Converts each controlled operation's Pauli string into Q# ``Pauli`` enums
-        and packages them with sign phases into a ``PauliSelectParams`` struct.
-
-        Args:
-            select: The SELECT oracle data object containing controlled operations,
-                phases, and qubit layout.
-
-        Returns:
-            A Q# callable implementing the Pauli SELECT oracle.
-
-        """
-        pauli_terms: list[list[qsharp.Pauli]] = []
-        control_states: list[int] = []
-        for op in select.controlled_operations:
-            base_paulis = [qsharp.Pauli.I] * select.num_target_qubits
-            for i, pauli_char in enumerate(reversed(op.operation)):
-                if pauli_char != "I":
-                    base_paulis[i] = getattr(qsharp.Pauli, pauli_char)
-            pauli_terms.append(base_paulis)
-            control_states.append(op.ctrl_state)
-        phases = [int(s) for s in select.phases]
-        select_params = QSHARP_UTILS.Select.PauliSelectParams(
-            pauliTerms=pauli_terms, signs=phases, controlStates=control_states
-        )
-        return QSHARP_UTILS.Select.MakeSelectOp(select_params)
-
     def resolve_lcu(self, container: UnitaryContainer) -> tuple[LCUContainer, bool]:
         """Unwrap a container into its LCU data and whether it is a quantum walk.
 
@@ -142,22 +112,85 @@ def resolve_lcu(self, container: UnitaryContainer) -> tuple[LCUContainer, bool]:
             "PSPMapper requires LCUContainer or LCUWalkContainer."
         )
 
-    def build_prepare_select_ops(self, container: UnitaryContainer) -> tuple[Any, Any, int]:
-        """Return the PREPARE and SELECT Q# oracles and the system register size.
+    def build_select_ops(self, container: UnitaryContainer) -> tuple[Any, int]:
+        """Return the SELECT oracle and the width of the system register it targets.
+
+        Each controlled operation's Pauli string becomes a list of Q# ``Pauli`` enums,
+        packaged with the sign phases into a ``PauliSelectParams`` struct.
 
         Args:
             container: The container held by the unitary representation.
 
         Returns:
-            The PREPARE Q# callable, the SELECT Q# callable, and the system register size.
+            The Q# SELECT callable and the number of system qubits.
 
         """
         lcu, _ = self.resolve_lcu(container)
-        if lcu.prepare is not None:
-            prepare_op = self._create_nested("prepare").run(lcu.prepare)._qsharp_op  # noqa: SLF001
-        else:
-            prepare_op = QSHARP_UTILS.PrepSelPrep.NoOpPrepare
-        return prepare_op, self._build_pauli_select_op(lcu.select), lcu.select.num_target_qubits
+        select = lcu.select
+        pauli_terms: list[list[qsharp.Pauli]] = []
+        control_states: list[int] = []
+        for op in select.controlled_operations:
+            base_paulis = [qsharp.Pauli.I] * select.num_target_qubits
+            for i, pauli_char in enumerate(reversed(op.operation)):
+                if pauli_char != "I":
+                    base_paulis[i] = getattr(qsharp.Pauli, pauli_char)
+            pauli_terms.append(base_paulis)
+            control_states.append(op.ctrl_state)
+        select_params = QSHARP_UTILS.Select.PauliSelectParams(
+            pauliTerms=pauli_terms, signs=[int(s) for s in select.phases], controlStates=control_states
+        )
+        return QSHARP_UTILS.Select.MakeSelectOp(select_params), select.num_target_qubits
+
+    def build_prep_ops(self, container: UnitaryContainer) -> tuple[Any, int, int]:
+        """Return the PREPARE oracle and the widths of the ancilla it acts on.
+
+        Args:
+            container: The container held by the unitary representation.
+
+        Returns:
+            The Q# PREPARE callable, the index width SELECT controls on, and the block
+            ancilla width.
+
+        Raises:
+            ValueError: If the PREPARE circuit carries no Q# operation, does not declare the
+                width it acts on, declares one too narrow for the index SELECT controls on, or
+                expects shared ancilla this mapper does not supply.
+
+        """
+        lcu, _ = self.resolve_lcu(container)
+        if lcu.prepare is None:
+            return QSHARP_UTILS.PrepSelPrep.NoOpPrepare, 0, 0
+
+        prepare_algorithm = self._create_nested("prepare")
+        prepare_circuit = prepare_algorithm.run(lcu.prepare)
+        prepare_op = prepare_circuit._qsharp_op  # noqa: SLF001
+        if prepare_op is None:
+            raise ValueError("The PREPARE circuit has no Q# operation to embed in the block encoding.")
+
+        num_block_ancillas = prepare_circuit.num_qubits
+        num_select_qubits = lcu.num_prepare_ancillas
+        if num_block_ancillas is None:
+            raise ValueError(
+                f"State preparation '{prepare_algorithm.name()}' does not declare num_qubits, so the "
+                "block ancilla register cannot be sized."
+            )
+        if num_block_ancillas <= 0:
+            raise ValueError(
+                f"State preparation '{prepare_algorithm.name()}' declares num_qubits={num_block_ancillas}, "
+                "but a PREPARE oracle must act on at least one qubit."
+            )
+        if prepare_circuit.metadata.num_phase_gradient_ancillas:
+            raise ValueError(
+                f"PREPARE oracle '{prepare_algorithm.name()}' requests "
+                f"{prepare_circuit.metadata.num_phase_gradient_ancillas} phase gradient ancilla, which "
+                "PSPMapper does not supply. Choose a state preparation that owns all the qubits it acts on."
+            )
+        if num_block_ancillas < num_select_qubits:
+            raise ValueError(
+                f"The PREPARE circuit acts on {num_block_ancillas} qubits, but the LCU decomposition indexes "
+                f"{num_select_qubits} of them. SELECT would control on qubits PREPARE does not own."
+            )
+        return prepare_op, num_select_qubits, num_block_ancillas
 
     def _run_impl(self, unitary: UnitaryRepresentation) -> Circuit:
         r"""Construct the block-encoding circuit on the flat ``[system | ancilla]`` register.
@@ -172,12 +205,15 @@ def _run_impl(self, unitary: UnitaryRepresentation) -> Circuit:
 
         """
         container = unitary.get_container()
-        lcu, use_quantum_walk = self.resolve_lcu(container)
-        prepare_op, select_op, num_system = self.build_prepare_select_ops(container)
+        _, use_quantum_walk = self.resolve_lcu(container)
+        select_op, num_system_qubits = self.build_select_ops(container)
+        prepare_op, num_select_qubits, num_block_ancillas = self.build_prep_ops(container)
 
-        qsharp_op = QSHARP_UTILS.PrepSelPrep.MakePrepSelPrepOp(prepare_op, select_op, num_system)
+        qsharp_op = QSHARP_UTILS.PrepSelPrep.MakePrepSelPrepOp(
+            prepare_op, select_op, num_system_qubits, num_select_qubits
+        )
         if use_quantum_walk:
-            reflection_op = QSHARP_UTILS.PrepSelPrep.MakeAncillaReflectionOp(num_system)
+            reflection_op = QSHARP_UTILS.PrepSelPrep.MakeAncillaReflectionOp(num_system_qubits, num_block_ancillas)
             qsharp_op = QSHARP_UTILS.PrepSelPrep.MakeWalkOp(qsharp_op, reflection_op)
 
         if container.power != 1:
@@ -192,8 +228,9 @@ def _run_impl(self, unitary: UnitaryRepresentation) -> Circuit:
             parameter={
                 "prepareOp": prepare_op,
                 "selectOp": select_op,
-                "numSystemQubits": num_system,
-                "numAncillaQubits": lcu.num_prepare_ancillas,
+                "numSystemQubits": num_system_qubits,
+                "numSelectQubits": num_select_qubits,
+                "numBlockAncillaQubits": num_block_ancillas,
                 "power": container.power,
                 "useWalk": use_quantum_walk,
             },
@@ -202,5 +239,5 @@ def _run_impl(self, unitary: UnitaryRepresentation) -> Circuit:
         return Circuit(
             qsharp_factory=qsharp_factory,
             qsharp_op=qsharp_op,
-            num_qubits=num_system + lcu.num_prepare_ancillas,
+            num_qubits=num_system_qubits + num_block_ancillas,
         )

python/src/qdk_chemistry/algorithms/controlled_circuit_mapper/controlled_psp_mapper.py

@@ -110,12 +110,15 @@ def _run_impl(self, unitary: UnitaryRepresentation) -> Circuit:
 
         block_mapper = self._block_mapper()
         container = unitary.get_container()
-        lcu, use_quantum_walk = block_mapper.resolve_lcu(container)
-        prepare_op, select_op, num_system = block_mapper.build_prepare_select_ops(container)
+        _, use_quantum_walk = block_mapper.resolve_lcu(container)
+        select_op, num_system_qubits = block_mapper.build_select_ops(container)
+        prepare_op, num_select_qubits, num_block_ancillas = block_mapper.build_prep_ops(container)
 
-        step_op = QSHARP_UTILS.PrepSelPrep.MakePrepSelPrepOp(prepare_op, select_op, num_system)
+        step_op = QSHARP_UTILS.PrepSelPrep.MakePrepSelPrepOp(
+            prepare_op, select_op, num_system_qubits, num_select_qubits
+        )
         if use_quantum_walk:
-            reflection_op = QSHARP_UTILS.PrepSelPrep.MakeAncillaReflectionOp(num_system)
+            reflection_op = QSHARP_UTILS.PrepSelPrep.MakeAncillaReflectionOp(num_system_qubits, num_block_ancillas)
             step_op = QSHARP_UTILS.PrepSelPrep.MakeWalkOp(step_op, reflection_op)
 
         controlled_op = QSHARP_UTILS.CircuitComposition.MakeControlledOp(step_op)
@@ -130,8 +133,9 @@ def _run_impl(self, unitary: UnitaryRepresentation) -> Circuit:
             parameter={
                 "prepareOp": prepare_op,
                 "selectOp": select_op,
-                "numSystemQubits": num_system,
-                "numAncillaQubits": lcu.num_prepare_ancillas,
+                "numSystemQubits": num_system_qubits,
+                "numSelectQubits": num_select_qubits,
+                "numBlockAncillaQubits": num_block_ancillas,
                 "power": container.power,
                 "useWalk": use_quantum_walk,
             },

python/src/qdk_chemistry/algorithms/phase_estimation/circuit_builder/base.py

@@ -123,6 +123,25 @@ def _create_controlled_circuit(
         circuit = circuit_mapper.run(unitary_rep)
         return circuit, num_ancilla_qubits
 
+    @staticmethod
+    def _validate_state_prep_width(state_preparation: Circuit, num_qubits_passed: int) -> None:
+        """Check that the state preparation fits the register phase estimation hands it.
+
+        Args:
+            state_preparation: The state preparation circuit.
+            num_qubits_passed: Width of the register phase estimation applies it to.
+
+        Raises:
+            ValueError: If the state preparation acts on more qubits than it is given.
+
+        """
+        width = state_preparation.num_qubits
+        if width is not None and width > num_qubits_passed:
+            raise ValueError(
+                f"State preparation acts on {width} qubits but phase estimation applies it to "
+                f"{num_qubits_passed}. Choose a state preparation that fits the system register."
+            )
+
 
 class QpeCircuitBuilderFactory(AlgorithmFactory):
     """Factory class for creating QpeCircuitBuilder instances."""

python/src/qdk_chemistry/algorithms/phase_estimation/circuit_builder/iterative_builder.py

@@ -184,6 +184,7 @@ def _create_circuit_from_qsharp_op(
         """
         state_prep_op = state_preparation._qsharp_op  # noqa: SLF001
         ctrl_unitary_op = controlled_unitary_circuit._qsharp_op  # noqa: SLF001
+        self._validate_state_prep_width(state_preparation, num_system_qubits)
         iterative_parameters = {
             "statePrep": state_prep_op,
             "repControlledUnitary": ctrl_unitary_op,

python/src/qdk_chemistry/algorithms/phase_estimation/circuit_builder/standard_builder.py

@@ -147,6 +147,7 @@ def _create_circuit_from_qsharp_op(
         """
         state_prep_op = state_preparation._qsharp_op  # noqa: SLF001
         ctrl_unitary_ops = [c._qsharp_op for c in controlled_unitary_circuits]  # noqa: SLF001
+        self._validate_state_prep_width(state_preparation, num_system_qubits)
         phase_qubit_prep_op = QSHARP_UTILS.StatePreparation.MakePrepareHadamardAllOp()
         ancillas = list(range(num_bits))
         systems = [i + num_bits for i in range(num_system_qubits)]