add property to PrepareOracle instead

qualtran/bloqs/for_testing/qubitization_walk_test.py

@@ -12,18 +12,19 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 from functools import cached_property
-from typing import Tuple
+from typing import Optional, Tuple
 
 import attrs
 import cirq
 import scipy
-from numpy._typing import NDArray
+from numpy.typing import NDArray
 
 from qualtran import Signature
 from qualtran.bloqs.multiplexers.select_pauli_lcu import SelectPauliLCU
 from qualtran.bloqs.qubitization_walk_operator import QubitizationWalkOperator
 from qualtran.bloqs.select_and_prepare import PrepareOracle
 from qualtran.bloqs.state_preparation import PrepareUniformSuperposition
+from qualtran.resource_counting.symbolic_counting_utils import SymbolicFloat
 
 
 @attrs.frozen
@@ -32,6 +33,7 @@ class PrepareUniformSuperpositionTest(PrepareOracle):
     cvs: Tuple[int, ...] = attrs.field(
         converter=lambda v: (v,) if isinstance(v, int) else tuple(v), default=()
     )
+    qlambda: Optional[float] = None
 
     @cached_property
     def selection_registers(self) -> Signature:
@@ -41,6 +43,10 @@ def selection_registers(self) -> Signature:
     def junk_registers(self) -> Signature:
         return Signature.build()
 
+    @cached_property
+    def l1_norm_of_coeffs(self) -> Optional['SymbolicFloat']:
+        return self.qlambda
+
     def decompose_from_registers(
         self, *, context: cirq.DecompositionContext, **quregs: NDArray[cirq.Qid]  # type: ignore[type-var]
     ) -> cirq.OP_TREE:
@@ -57,10 +63,8 @@ def get_uniform_pauli_qubitized_walk(target_bitsize: int):
     ham_dps = [ps.dense(q) for ps in ham]
 
     assert scipy.linalg.ishermitian(ham.matrix())
-    prepare = PrepareUniformSuperpositionTest(len(ham_coeff))
+    prepare = PrepareUniformSuperpositionTest(len(ham_coeff), qlambda=sum(ham_coeff))
     select = SelectPauliLCU(
         (len(ham_coeff) - 1).bit_length(), select_unitaries=ham_dps, target_bitsize=target_bitsize
     )
-    return ham, QubitizationWalkOperator(
-        select=select, prepare=prepare, sum_of_lcu_coefficients=sum(ham_coeff)
-    )
+    return ham, QubitizationWalkOperator(select=select, prepare=prepare)

qualtran/bloqs/for_testing/random_select_and_prepare.py

@@ -56,6 +56,10 @@ def __attrs_post_init__(self):
     def selection_registers(self) -> tuple[Register, ...]:
         return (Register('selection', BoundedQUInt(bitsize=self.U.bitsize)),)
 
+    @property
+    def l1_norm_of_coeffs(self) -> float:
+        return 1.0
+
     @classmethod
     def create(cls, bitsize: int, *, random_state: np.random.RandomState):
         matrix = RandomGate.create(bitsize, random_state=random_state).matrix
@@ -209,4 +213,4 @@ def random_qubitization_walk_operator(
         ]
     )
 
-    return QubitizationWalkOperator(prepare=prepare, select=select, sum_of_lcu_coefficients=1), ham
+    return QubitizationWalkOperator(prepare=prepare, select=select), ham

qualtran/bloqs/hubbard_model.py

@@ -47,7 +47,7 @@
 See the documentation for `PrepareHubbard` and `SelectHubbard` for details.
 """
 from functools import cached_property
-from typing import Collection, Optional, Sequence, Tuple, Union
+from typing import Collection, Optional, Sequence, Tuple, TYPE_CHECKING, Union
 
 import attrs
 import cirq
@@ -67,6 +67,9 @@
     PrepareUniformSuperposition,
 )
 
+if TYPE_CHECKING:
+    from qualtran.resource_counting.symbolic_counting_utils import SymbolicFloat
+
 
 @attrs.frozen
 class SelectHubbard(SelectOracle):
@@ -309,6 +312,13 @@ def selection_registers(self) -> Tuple[Register, ...]:
     def junk_registers(self) -> Tuple[Register, ...]:
         return (Register('temp', QAny(2)),)
 
+    @cached_property
+    def l1_norm_of_coeffs(self) -> 'SymbolicFloat':
+        # https://arxiv.org/abs/1805.03662v2 equation 60
+        N = self.x_dim * self.y_dim * 2
+        qlambda = 2 * N * self.t + (N * self.mu) // 2
+        return qlambda
+
     @cached_property
     def signature(self) -> Signature:
         return Signature([*self.selection_registers, *self.junk_registers])
@@ -321,8 +331,7 @@ def decompose_from_registers(
         temp = quregs['temp']
 
         N = self.x_dim * self.y_dim * 2
-        qlambda = 2 * N * self.t + (N * self.mu) // 2
-        yield cirq.Ry(rads=2 * np.arccos(np.sqrt(self.t * N / qlambda))).on(*V)
+        yield cirq.Ry(rads=2 * np.arccos(np.sqrt(self.t * N / self.l1_norm_of_coeffs))).on(*V)
         yield cirq.Ry(rads=2 * np.arccos(np.sqrt(1 / 5))).on(*U).controlled_by(*V)
         yield PrepareUniformSuperposition(self.x_dim).on_registers(controls=[], target=p_x)
         yield PrepareUniformSuperposition(self.y_dim).on_registers(controls=[], target=p_y)
@@ -353,10 +362,4 @@ def get_walk_operator_for_hubbard_model(
     select = SelectHubbard(x_dim, y_dim)
     prepare = PrepareHubbard(x_dim, y_dim, t, mu)
 
-    # https://arxiv.org/abs/1805.03662v2 equation 60
-    N = x_dim * y_dim
-    sum_of_ham_coeffs = 2 * N * t + N * mu / 2
-
-    return QubitizationWalkOperator(
-        select=select, prepare=prepare, sum_of_lcu_coefficients=sum_of_ham_coeffs
-    )
+    return QubitizationWalkOperator(select=select, prepare=prepare)

qualtran/bloqs/qubitization_walk_operator.py

@@ -67,7 +67,6 @@ class QubitizationWalkOperator(GateWithRegisters):
 
     select: SelectOracle
     prepare: PrepareOracle
-    sum_of_lcu_coefficients: Optional[float] = None
     control_val: Optional[int] = None
     power: int = 1
 
@@ -96,6 +95,10 @@ def signature(self) -> Signature:
     def reflect(self) -> ReflectionUsingPrepare:
         return ReflectionUsingPrepare(self.prepare, control_val=self.control_val, global_phase=-1)
 
+    @cached_property
+    def sum_of_lcu_coefficients(self) -> Optional[float]:
+        return self.prepare.l1_norm_of_coeffs
+
     def decompose_from_registers(
         self,
         context: cirq.DecompositionContext,

qualtran/bloqs/qubitization_walk_operator_test.py

@@ -39,9 +39,7 @@ def walk_operator_for_pauli_hamiltonian(ham: cirq.PauliSum, eps: float) -> Qubit
     select = SelectPauliLCU(
         total_bits(prepare.selection_registers), select_unitaries=ham_dps, target_bitsize=len(q)
     )
-    return QubitizationWalkOperator(
-        select=select, prepare=prepare, sum_of_lcu_coefficients=sum(ham_coeff)
-    )
+    return QubitizationWalkOperator(select=select, prepare=prepare)
 
 
 def get_walk_operator_for_1d_ising_model(num_sites: int, eps: float) -> QubitizationWalkOperator:
@@ -53,15 +51,16 @@ def get_walk_operator_for_1d_ising_model(num_sites: int, eps: float) -> Qubitiza
 def test_qubitization_walk_operator(num_sites: int, eps: float):
     ham = get_1d_ising_hamiltonian(cirq.LineQubit.range(num_sites))
     ham_coeff = [abs(ps.coefficient.real) for ps in ham]
-    qubitization_lambda = np.sum(ham_coeff)
 
     walk = walk_operator_for_pauli_hamiltonian(ham, eps)
     assert_valid_bloq_decomposition(walk)
 
+    qubitization_lambda = walk.sum_of_lcu_coefficients
+
     g, qubit_order, walk_circuit = construct_gate_helper_and_qubit_order(walk)
 
     L_state = np.zeros(2 ** len(g.quregs['selection']))
-    L_state[: len(ham_coeff)] = np.sqrt(ham_coeff / qubitization_lambda)
+    L_state[: len(ham_coeff)] = np.sqrt(np.array(ham_coeff) / qubitization_lambda)
 
     assert len(walk_circuit.all_qubits()) < 23
 

qualtran/bloqs/select_and_prepare.py

@@ -14,10 +14,13 @@
 
 import abc
 from functools import cached_property
-from typing import Tuple
+from typing import Optional, Tuple, TYPE_CHECKING
 
 from qualtran import GateWithRegisters, Register, Signature
 
+if TYPE_CHECKING:
+    from qualtran.resource_counting.symbolic_counting_utils import SymbolicFloat
+
 
 class SelectOracle(GateWithRegisters):
     r"""Abstract base class that defines the API for a SELECT Oracle.
@@ -63,14 +66,15 @@ class PrepareOracle(GateWithRegisters):
     r"""Abstract base class that defines the API for a PREPARE Oracle.
 
     Given a set of coefficients $\{c_0, c_1, ..., c_{N - 1}\}, the PREPARE oracle is used to encode
-    the coefficients as amplitudes of a state $|\Psi\rangle = \sum_{i=0}^{N-1}c_{i}|i\rangle$ using
-    a selection register $|i\rangle$. In order to prepare such a state, the PREPARE circuit is also
-    allowed to use a junk register that is entangled with selection register.
+    the coefficients as amplitudes of a state $|\Psi\rangle = \sum_{i=0}^{N-1} \sqrt{\frac{c_{i}}{\lambda}} |i\rangle$
+    where $\lambda = \sum_i c_i$, using a selection register $|i\rangle$. In order to prepare such
+    a state, the PREPARE circuit is also allowed to use a junk register that is entangled with
+    selection register.
 
     Thus, the action of a PREPARE circuit on an input state $|0\rangle$ can be defined as:
 
     $$
-        PREPARE |0\rangle = \sum_{i=0}^{N-1}c_{i}|i\rangle |junk_{i}\rangle
+        PREPARE |0\rangle = \sum_{i=0}^{N-1} \sqrt{ \frac{c_{i}}{\lambda} } |i\rangle |junk_{i}\rangle
     $$
     """
 
@@ -86,3 +90,11 @@ def junk_registers(self) -> Tuple[Register, ...]:
     @cached_property
     def signature(self) -> Signature:
         return Signature([*self.selection_registers, *self.junk_registers])
+
+    @property
+    def l1_norm_of_coeffs(self) -> Optional['SymbolicFloat']:
+        r"""Sum of the coefficients $c_i$.
+
+        For LCU Hamiltonians, this is usually referred to as $\lambda$ in texts.
+        """
+        return None

qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py

@@ -21,7 +21,7 @@
 """
 
 from functools import cached_property
-from typing import List, Tuple
+from typing import List, Tuple, TYPE_CHECKING
 
 import attrs
 import cirq
@@ -44,6 +44,9 @@
     ignore_split_join,
 )
 
+if TYPE_CHECKING:
+    from qualtran.resource_counting.symbolic_counting_utils import SymbolicFloat
+
 
 @cirq.value_equality()
 @attrs.frozen
@@ -96,6 +99,7 @@ class StatePreparationAliasSampling(PrepareOracle):
     alt: NDArray[np.int_]
     keep: NDArray[np.int_]
     mu: int
+    sum_of_lcu_coeffs: float
 
     @classmethod
     def from_lcu_probs(
@@ -119,8 +123,13 @@ def from_lcu_probs(
             alt=np.array(alt),
             keep=np.array(keep),
             mu=mu,
+            sum_of_lcu_coeffs=sum(lcu_probabilities),
         )
 
+    @cached_property
+    def l1_norm_of_coeffs(self) -> 'SymbolicFloat':
+        return self.sum_of_lcu_coeffs
+
     @cached_property
     def sigma_mu_bitsize(self) -> int:
         return self.mu