Improve interface for bloqs with specialized single-qubit-controlled versions

qualtran/_infra/controlled.py

@@ -250,6 +250,18 @@ def from_cirq_cv(
             bloq_cvs.append(curr_cvs)
         return CtrlSpec(tuple(qdtypes), tuple(bloq_cvs))
 
+    def get_single_control_bit(self) -> Union[0, 1]:
+        """If controlled by a single qubit, return the control bit, otherwise raise"""
+        if self.num_qubits != 1:
+            raise ValueError(f"expected a single qubit control, got {self.num_qubits}")
+
+        (qdtype,) = self.qdtypes
+        (cv,) = self.cvs
+        (idx,) = Register('', qdtype, cv.shape).all_idxs()
+        (control_bit,) = qdtype.to_bits(cv[idx])
+
+        return int(control_bit)
+
 
 class AddControlledT(Protocol):
     """The signature for the `add_controlled` callback part of `ctrl_system`.

qualtran/bloqs/chemistry/hubbard_model/qubitization/select_hubbard.py

@@ -20,17 +20,27 @@
 import numpy as np
 from numpy.typing import NDArray
 
-from qualtran import bloq_example, BloqDocSpec, BQUInt, QAny, QBit, Register, Signature
+from qualtran import (
+    AddControlledT,
+    Bloq,
+    bloq_example,
+    BloqDocSpec,
+    BQUInt,
+    CtrlSpec,
+    QAny,
+    QBit,
+    Register,
+    Signature,
+)
 from qualtran._infra.gate_with_registers import total_bits
-from qualtran._infra.single_qubit_controlled import SpecializedSingleQubitControlledExtension
 from qualtran.bloqs.basic_gates import CSwap
 from qualtran.bloqs.multiplexers.apply_gate_to_lth_target import ApplyGateToLthQubit
 from qualtran.bloqs.multiplexers.select_base import SelectOracle
 from qualtran.bloqs.multiplexers.selected_majorana_fermion import SelectedMajoranaFermion
 
 
 @attrs.frozen
-class SelectHubbard(SelectOracle, SpecializedSingleQubitControlledExtension):  # type: ignore[misc]
+class SelectHubbard(SelectOracle):
     r"""The SELECT operation optimized for the 2D Hubbard model.
 
     In contrast to SELECT for an arbitrary chemistry Hamiltonian, we:
@@ -180,6 +190,24 @@ def __str__(self):
             return f'C{s}'
         return s
 
+    @property
+    def cv(self):
+        return self.control_val
+
+    def with_cv(self, *, cv: Optional[int]) -> 'SelectHubbard':
+        return attrs.evolve(self, control_val=cv)
+
+    @property
+    def ctrl_reg_name(self) -> str:
+        return 'control'
+
+    def get_ctrl_system(self, ctrl_spec: 'CtrlSpec') -> Tuple['Bloq', 'AddControlledT']:
+        from qualtran.bloqs.mcmt.bloq_with_specialized_single_qubit_control import (
+            get_ctrl_system_for_bloq_with_specialized_single_qubit_control,
+        )
+
+        return get_ctrl_system_for_bloq_with_specialized_single_qubit_control(self, ctrl_spec)
+
 
 @bloq_example
 def _sel_hubb() -> SelectHubbard:

qualtran/bloqs/chemistry/sparse/select_bloq.py

@@ -16,11 +16,23 @@
 from functools import cached_property
 from typing import Dict, Optional, Tuple, TYPE_CHECKING
 
+import attrs
 import cirq
 from attrs import frozen
 
-from qualtran import bloq_example, BloqBuilder, BloqDocSpec, BQUInt, QAny, QBit, Register, SoquetT
-from qualtran._infra.single_qubit_controlled import SpecializedSingleQubitControlledExtension
+from qualtran import (
+    AddControlledT,
+    Bloq,
+    bloq_example,
+    BloqBuilder,
+    BloqDocSpec,
+    BQUInt,
+    CtrlSpec,
+    QAny,
+    QBit,
+    Register,
+    SoquetT,
+)
 from qualtran.bloqs.basic_gates import SGate
 from qualtran.bloqs.multiplexers.select_base import SelectOracle
 from qualtran.bloqs.multiplexers.selected_majorana_fermion import SelectedMajoranaFermion
@@ -30,7 +42,7 @@
 
 
 @frozen
-class SelectSparse(SpecializedSingleQubitControlledExtension, SelectOracle):  # type: ignore[misc]
+class SelectSparse(SelectOracle):
     r"""SELECT oracle for the sparse Hamiltonian.
 
     Implements the two applications of Fig. 13.
@@ -157,6 +169,24 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT':
         c_maj_y = SelectedMajoranaFermion(sel_pa, target_gate=cirq.Y)
         return {SGate(): 1, maj_x: 1, c_maj_x: 1, maj_y: 1, c_maj_y: 1}
 
+    @property
+    def cv(self):
+        return self.control_val
+
+    def with_cv(self, *, cv: Optional[int]) -> 'SelectSparse':
+        return attrs.evolve(self, control_val=cv)
+
+    @property
+    def ctrl_reg_name(self) -> str:
+        return 'control'
+
+    def get_ctrl_system(self, ctrl_spec: 'CtrlSpec') -> Tuple['Bloq', 'AddControlledT']:
+        from qualtran.bloqs.mcmt.bloq_with_specialized_single_qubit_control import (
+            get_ctrl_system_for_bloq_with_specialized_single_qubit_control,
+        )
+
+        return get_ctrl_system_for_bloq_with_specialized_single_qubit_control(self, ctrl_spec)
+
 
 @bloq_example
 def _sel_sparse() -> SelectSparse:

qualtran/bloqs/chemistry/thc/select_bloq.py

@@ -20,18 +20,19 @@
 from attrs import evolve, frozen
 
 from qualtran import (
+    AddControlledT,
     Bloq,
     bloq_example,
     BloqBuilder,
     BloqDocSpec,
     BQUInt,
+    CtrlSpec,
     QAny,
     QBit,
     Register,
     Signature,
     SoquetT,
 )
-from qualtran._infra.single_qubit_controlled import SpecializedSingleQubitControlledExtension
 from qualtran.bloqs.basic_gates import CSwap, Toffoli, XGate
 from qualtran.bloqs.chemistry.black_boxes import ApplyControlledZs
 from qualtran.bloqs.multiplexers.select_base import SelectOracle
@@ -120,7 +121,7 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT':
 
 
 @frozen
-class SelectTHC(SpecializedSingleQubitControlledExtension, SelectOracle):  # type: ignore[misc]
+class SelectTHC(SelectOracle):
     r"""SELECT for THC Hamiltonian.
 
     Args:
@@ -313,6 +314,24 @@ def build_composite_bloq(self, bb: 'BloqBuilder', **soqs: 'SoquetT') -> Dict[str
 
         return out_soqs
 
+    @property
+    def cv(self):
+        return self.control_val
+
+    def with_cv(self, *, cv: Optional[int]) -> 'SelectTHC':
+        return evolve(self, control_val=cv)
+
+    @property
+    def ctrl_reg_name(self) -> str:
+        return 'control'
+
+    def get_ctrl_system(self, ctrl_spec: 'CtrlSpec') -> Tuple['Bloq', 'AddControlledT']:
+        from qualtran.bloqs.mcmt.bloq_with_specialized_single_qubit_control import (
+            get_ctrl_system_for_bloq_with_specialized_single_qubit_control,
+        )
+
+        return get_ctrl_system_for_bloq_with_specialized_single_qubit_control(self, ctrl_spec)
+
 
 @bloq_example
 def _thc_sel() -> SelectTHC:

qualtran/bloqs/for_testing/random_select_and_prepare.py

@@ -14,13 +14,13 @@
 from functools import cached_property
 from typing import Iterator, Optional, Tuple
 
+import attrs
 import cirq
 import numpy as np
 from attrs import frozen
 from numpy.typing import NDArray
 
-from qualtran import BloqBuilder, BQUInt, QBit, Register, SoquetT
-from qualtran._infra.single_qubit_controlled import SpecializedSingleQubitControlledExtension
+from qualtran import AddControlledT, Bloq, BloqBuilder, BQUInt, CtrlSpec, QBit, Register, SoquetT
 from qualtran.bloqs.block_encoding.lcu_block_encoding import SelectBlockEncoding
 from qualtran.bloqs.for_testing.matrix_gate import MatrixGate
 from qualtran.bloqs.multiplexers.select_base import SelectOracle
@@ -84,7 +84,7 @@ def alphas(self):
 
 
 @frozen
-class TestPauliSelectOracle(SpecializedSingleQubitControlledExtension, SelectOracle):  # type: ignore[misc]
+class TestPauliSelectOracle(SelectOracle):  # type: ignore[misc]
     r"""Paulis acting on $m$ qubits, controlled by an $n$-qubit register.
 
     Given $2^n$ multi-qubit-Paulis (acting on $m$ qubits) $U_j$,
@@ -149,6 +149,24 @@ def decompose_from_registers(
                 op = op.controlled_by(*quregs['control'], control_values=[self.control_val])
             yield op
 
+    @property
+    def cv(self):
+        return self.control_val
+
+    def with_cv(self, *, cv: Optional[int]) -> 'TestPauliSelectOracle':
+        return attrs.evolve(self, control_val=cv)
+
+    @property
+    def ctrl_reg_name(self) -> str:
+        return 'control'
+
+    def get_ctrl_system(self, ctrl_spec: 'CtrlSpec') -> Tuple['Bloq', 'AddControlledT']:
+        from qualtran.bloqs.mcmt.bloq_with_specialized_single_qubit_control import (
+            get_ctrl_system_for_bloq_with_specialized_single_qubit_control,
+        )
+
+        return get_ctrl_system_for_bloq_with_specialized_single_qubit_control(self, ctrl_spec)
+
 
 def random_qubitization_walk_operator(
     select_bitsize: int, target_bitsize: int, *, random_state: np.random.RandomState

qualtran/bloqs/mcmt/bloq_with_specialized_single_qubit_control.py

@@ -0,0 +1,97 @@
+#  Copyright 2024 Google LLC
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+from typing import cast, Iterable, Optional, Sequence, TYPE_CHECKING
+
+import numpy as np
+
+if TYPE_CHECKING:
+    from qualtran import AddControlledT, Bloq, BloqBuilder, CtrlSpec, SoquetT
+
+
+ControlBit = Optional[0 | 1]
+
+
+def get_ctrl_system_for_bloq_with_specialized_single_qubit_control(
+    bloq: 'Bloq',
+    ctrl_spec: 'CtrlSpec',
+    *,
+    current_control_bit: ControlBit,
+    bloq_with_ctrl: 'Bloq',
+    ctrl_reg_name: str,
+    bloq_without_ctrl: 'Bloq',
+    bloq_with_ctrl_0: Optional['Bloq'] = None,
+) -> tuple['Bloq', 'AddControlledT']:
+    """Build the control system for a bloq with a specialized single-qubit controlled variant.
+
+    Args:
+        bloq: The current bloq of which a controlled version is requested.
+        ctrl_spec: The control specification
+        current_control_bit: The control bit of the current bloq, one of `0, 1, None`.
+        bloq_with_ctrl: The variant of this bloq with control bit `1`.
+        ctrl_reg_name: The name of the control qubit register.
+        bloq_without_ctrl: The variant of this bloq without a control.
+        bloq_with_ctrl_0: (optional) the variant of this bloq controlled by a qubit in the 0 state.
+    """
+    from qualtran import Bloq, CtrlSpec, Soquet
+    from qualtran.bloqs.mcmt import ControlledViaAnd
+
+    def _get_default_fallback():
+        return ControlledViaAnd.make_ctrl_system(bloq=bloq, ctrl_spec=ctrl_spec)
+
+    if ctrl_spec.num_qubits != 1:
+        return _get_default_fallback()
+
+    control_bit = ctrl_spec.get_single_control_bit()
+
+    if current_control_bit is None:
+        # the easy case: use the controlled bloq
+        ctrl_bloq = bloq_with_ctrl if control_bit == 1 else bloq_with_ctrl_0
+        if ctrl_bloq is None:
+            return _get_default_fallback()
+
+        def _adder(
+            bb: 'BloqBuilder', ctrl_soqs: Sequence['SoquetT'], in_soqs: dict[str, 'SoquetT']
+        ) -> tuple[Iterable['SoquetT'], Iterable['SoquetT']]:
+            (ctrl,) = ctrl_soqs
+            in_soqs |= {ctrl_reg_name: ctrl}
+
+            out_soqs = bb.add_d(ctrl_bloq, **in_soqs)
+
+            ctrl = out_soqs.pop(ctrl_reg_name)
+            return [ctrl], out_soqs.values()
+
+    else:
+        # the difficult case: must combine the two controls into one
+        ctrl_bloq = ControlledViaAnd(
+            bloq_without_ctrl, CtrlSpec(cvs=[control_bit, current_control_bit])
+        )
+
+        def _adder(
+            bb: 'BloqBuilder', ctrl_soqs: Sequence['SoquetT'], in_soqs: dict[str, 'SoquetT']
+        ) -> tuple[Iterable['SoquetT'], Iterable['SoquetT']]:
+            # extract the two control bits
+            (ctrl0,) = ctrl_soqs
+            ctrl1 = in_soqs.pop('ctrl')
+
+            ctrl0 = cast(Soquet, ctrl0)
+            ctrl1 = cast(Soquet, ctrl1)
+
+            # add the singly controlled bloq
+            ctrls, *out_soqs = bb.add_t(ctrl_bloq, ctrl=[ctrl0, ctrl1], **in_soqs)
+            assert isinstance(ctrls, np.ndarray)
+            ctrl0, ctrl1 = ctrls
+
+            return [ctrl0], [ctrl1, *out_soqs]
+
+    return ctrl_bloq, _adder