diff --git a/qualtran/_infra/controlled.py b/qualtran/_infra/controlled.py index e43a2d2b9..f2f946030 100644 --- a/qualtran/_infra/controlled.py +++ b/qualtran/_infra/controlled.py @@ -23,6 +23,7 @@ Sequence, Tuple, TYPE_CHECKING, + TypeAlias, Union, ) @@ -45,6 +46,9 @@ from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator from qualtran.simulation.classical_sim import ClassicalValT +ControlBit: TypeAlias = int +"""A control bit, either 0 or 1.""" + def _cvs_convert( cvs: Union[ @@ -250,6 +254,18 @@ def from_cirq_cv( bloq_cvs.append(curr_cvs) return CtrlSpec(tuple(qdtypes), tuple(bloq_cvs)) + def get_single_ctrl_bit(self) -> ControlBit: + """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`. diff --git a/qualtran/_infra/controlled_test.py b/qualtran/_infra/controlled_test.py index 34b26def8..77d72432c 100644 --- a/qualtran/_infra/controlled_test.py +++ b/qualtran/_infra/controlled_test.py @@ -100,6 +100,26 @@ def test_ctrl_spec_to_cirq_cv_roundtrip(): assert CtrlSpec.from_cirq_cv(cirq_cv, qdtypes=ctrl_spec.qdtypes, shapes=ctrl_spec.shapes) +@pytest.mark.parametrize( + "ctrl_spec", [CtrlSpec(), CtrlSpec(cvs=[1]), CtrlSpec(cvs=np.atleast_2d([1]))] +) +def test_ctrl_spec_single_bit_one(ctrl_spec: CtrlSpec): + assert ctrl_spec.get_single_ctrl_bit() == 1 + + +@pytest.mark.parametrize( + "ctrl_spec", [CtrlSpec(cvs=0), CtrlSpec(cvs=[0]), CtrlSpec(cvs=np.atleast_2d([0]))] +) +def test_ctrl_spec_single_bit_zero(ctrl_spec: CtrlSpec): + assert ctrl_spec.get_single_ctrl_bit() == 0 + + +@pytest.mark.parametrize("ctrl_spec", [CtrlSpec(cvs=[1, 1]), CtrlSpec(qdtypes=QUInt(2), cvs=0)]) +def test_ctrl_spec_single_bit_raises(ctrl_spec: CtrlSpec): + with pytest.raises(ValueError): + ctrl_spec.get_single_ctrl_bit() + + def _test_cirq_equivalence(bloq: Bloq, gate: 'cirq.Gate'): import cirq diff --git a/qualtran/bloqs/chemistry/hubbard_model/qubitization/select_hubbard.py b/qualtran/bloqs/chemistry/hubbard_model/qubitization/select_hubbard.py index aa889f91f..f99a56971 100644 --- a/qualtran/bloqs/chemistry/hubbard_model/qubitization/select_hubbard.py +++ b/qualtran/bloqs/chemistry/hubbard_model/qubitization/select_hubbard.py @@ -20,9 +20,19 @@ 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 @@ -30,7 +40,7 @@ @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,19 @@ def __str__(self): return f'C{s}' return s + def get_ctrl_system(self, ctrl_spec: 'CtrlSpec') -> Tuple['Bloq', 'AddControlledT']: + from qualtran.bloqs.mcmt.specialized_ctrl import get_ctrl_system_1bit_cv + + return get_ctrl_system_1bit_cv( + self, + ctrl_spec=ctrl_spec, + current_ctrl_bit=self.control_val, + get_ctrl_bloq_and_ctrl_reg_name=lambda cv: ( + attrs.evolve(self, control_val=cv), + 'control', + ), + ) + @bloq_example def _sel_hubb() -> SelectHubbard: diff --git a/qualtran/bloqs/chemistry/sparse/select_bloq.py b/qualtran/bloqs/chemistry/sparse/select_bloq.py index 06d9c42c1..1c7fd5595 100644 --- a/qualtran/bloqs/chemistry/sparse/select_bloq.py +++ b/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,19 @@ 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} + def get_ctrl_system(self, ctrl_spec: 'CtrlSpec') -> Tuple['Bloq', 'AddControlledT']: + from qualtran.bloqs.mcmt.specialized_ctrl import get_ctrl_system_1bit_cv + + return get_ctrl_system_1bit_cv( + self, + ctrl_spec=ctrl_spec, + current_ctrl_bit=self.control_val, + get_ctrl_bloq_and_ctrl_reg_name=lambda cv: ( + attrs.evolve(self, control_val=cv), + 'control', + ), + ) + @bloq_example def _sel_sparse() -> SelectSparse: diff --git a/qualtran/bloqs/chemistry/thc/select_bloq.py b/qualtran/bloqs/chemistry/thc/select_bloq.py index 5ed430599..a3b825950 100644 --- a/qualtran/bloqs/chemistry/thc/select_bloq.py +++ b/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,16 @@ def build_composite_bloq(self, bb: 'BloqBuilder', **soqs: 'SoquetT') -> Dict[str return out_soqs + def get_ctrl_system(self, ctrl_spec: 'CtrlSpec') -> Tuple['Bloq', 'AddControlledT']: + from qualtran.bloqs.mcmt.specialized_ctrl import get_ctrl_system_1bit_cv + + return get_ctrl_system_1bit_cv( + self, + ctrl_spec=ctrl_spec, + current_ctrl_bit=self.control_val, + get_ctrl_bloq_and_ctrl_reg_name=lambda cv: (evolve(self, control_val=cv), 'control'), + ) + @bloq_example def _thc_sel() -> SelectTHC: diff --git a/qualtran/bloqs/for_testing/random_select_and_prepare.py b/qualtran/bloqs/for_testing/random_select_and_prepare.py index c5eed3569..6e1a607f5 100644 --- a/qualtran/bloqs/for_testing/random_select_and_prepare.py +++ b/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,19 @@ def decompose_from_registers( op = op.controlled_by(*quregs['control'], control_values=[self.control_val]) yield op + def get_ctrl_system(self, ctrl_spec: 'CtrlSpec') -> Tuple['Bloq', 'AddControlledT']: + from qualtran.bloqs.mcmt.specialized_ctrl import get_ctrl_system_1bit_cv + + return get_ctrl_system_1bit_cv( + self, + ctrl_spec=ctrl_spec, + current_ctrl_bit=self.control_val, + get_ctrl_bloq_and_ctrl_reg_name=lambda cv: ( + attrs.evolve(self, control_val=cv), + 'control', + ), + ) + def random_qubitization_walk_operator( select_bitsize: int, target_bitsize: int, *, random_state: np.random.RandomState diff --git a/qualtran/bloqs/mcmt/controlled_via_and.py b/qualtran/bloqs/mcmt/controlled_via_and.py index b9f7389a0..a0ae58dee 100644 --- a/qualtran/bloqs/mcmt/controlled_via_and.py +++ b/qualtran/bloqs/mcmt/controlled_via_and.py @@ -50,8 +50,7 @@ def _is_single_bit_control(self) -> bool: @cached_property def _single_control_value(self) -> int: - assert self._is_single_bit_control() - return self.ctrl_spec._cvs_tuple[0] + return self.ctrl_spec.get_single_ctrl_bit() def adjoint(self) -> 'ControlledViaAnd': return ControlledViaAnd(self.subbloq.adjoint(), self.ctrl_spec) diff --git a/qualtran/bloqs/mcmt/specialized_ctrl.py b/qualtran/bloqs/mcmt/specialized_ctrl.py new file mode 100644 index 000000000..8c70df3f6 --- /dev/null +++ b/qualtran/bloqs/mcmt/specialized_ctrl.py @@ -0,0 +1,255 @@ +# 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 functools import cached_property +from typing import Callable, cast, Iterable, Optional, Sequence, TYPE_CHECKING + +import attrs +import numpy as np + +from qualtran import Bloq, QBit, Register, Signature + +if TYPE_CHECKING: + from qualtran import AddControlledT, BloqBuilder, CtrlSpec, SoquetT + from qualtran._infra.controlled import ControlBit + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator + + +@attrs.frozen +class _MultiControlledFromSinglyControlled(Bloq): + """Helper bloq implementing a multi-controlled-U given access to controlled-U. + + This is for internal use only. For reducing multiple controls to a single control, + see :class:`qualtran.bloqs.mcmt.ControlledViaAnd` and + :meth:`qualtran.bloqs.mcmt.specialized_ctrl.get_ctrl_system_1bit_cv`. + + This bloq is used as an intermediate bloq by `get_ctrl_system_1bit_cv` in the + controlled-controlled-bloq case. To cleanly support further controlling this bloq, + the `cvs` attribute accepts a tuple (of at least two controls), and defers to + `ControlledViaAnd` whenever possible, and only extends the `cvs` in the edge cases. + """ + + cvs: tuple[int, ...] + ctrl_bloq: Bloq + ctrl_reg_name: str + + def __attrs_post_init__(self): + assert len(self.cvs) >= 2, f"{self} must have at least 2 controls, got {self.cvs=}" + + @cached_property + def signature(self) -> 'Signature': + return Signature( + [Register(self.ctrl_reg_name, dtype=QBit(), shape=(len(self.cvs),))] + + [reg for reg in self.ctrl_bloq.signature if reg.name != self.ctrl_reg_name] + ) + + @cached_property + def _and_bloq(self) -> Bloq: + from qualtran.bloqs.mcmt import And, MultiAnd + + if len(self.cvs) == 2: + return And(*self.cvs) + else: + return MultiAnd(self.cvs) + + def build_composite_bloq(self, bb: 'BloqBuilder', **soqs: 'SoquetT') -> dict[str, 'SoquetT']: + and_soqs = bb.add_d(self._and_bloq, ctrl=soqs.pop(self.ctrl_reg_name)) + + soqs |= {self.ctrl_reg_name: and_soqs.pop('target')} + soqs = bb.add_d(self.ctrl_bloq, **soqs) + and_soqs |= {'target': soqs.pop(self.ctrl_reg_name)} + + soqs |= {self.ctrl_reg_name: bb.add(self._and_bloq.adjoint(), **and_soqs)} + + return soqs + + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': + return {self._and_bloq: 1, self.ctrl_bloq: 1, self._and_bloq.adjoint(): 1} + + def get_ctrl_system(self, ctrl_spec: 'CtrlSpec') -> tuple['Bloq', 'AddControlledT']: + if ctrl_spec.num_qubits != 1: + return super().get_ctrl_system(ctrl_spec=ctrl_spec) + + ctrl_bloq = attrs.evolve(self, cvs=(ctrl_spec.get_single_ctrl_bit(),) + self.cvs) + + def _adder(bb, ctrl_soqs, in_soqs): + in_soqs[self.ctrl_reg_name] = np.concatenate(ctrl_soqs, in_soqs[self.ctrl_reg_name]) + ctrls, *out_soqs = bb.add_t(ctrl_bloq, **in_soqs) + return ctrls[:1], [*ctrls[1:], *out_soqs] + + return ctrl_bloq, _adder + + def __str__(self): + return f'C[{len(self.cvs)-1}][{self.ctrl_bloq}]' + + +def _get_ctrl_system_1bit_cv( + bloq: 'Bloq', + ctrl_spec: 'CtrlSpec', + *, + current_ctrl_bit: Optional['ControlBit'], + get_ctrl_bloq_and_ctrl_reg_name: Callable[['ControlBit'], Optional[tuple['Bloq', str]]], +) -> tuple['Bloq', 'AddControlledT']: + """Internal method to build the control system for a bloq using single-qubit controlled variants. + + Uses the provided specialized implementation when a singly-controlled variant of the bloq is + requested. When controlled by multiple qubits, the controls are reduced to a single qubit + and the singly-controlled bloq is used. + + The user can provide specializations for the bloq controlled by `1` and (optionally) by `0`. + The specialization for control bit `1` must be provided. + In case a specialization for a control bit `0` is not provided, the default fallback is used + instead, which wraps the bloq using the `Controlled` metabloq. + + Args: + bloq: The current bloq. + ctrl_spec: The control specification + current_ctrl_bit: The control bit of the current bloq, one of `0, 1, None`. + get_ctrl_bloq_and_ctrl_reg_name: A callable that accepts a control bit (`0` or `1`), + and returns the controlled variant of this bloq and the name of the control register. + If the callable returns `None`, then the default fallback is used. + """ + from qualtran import 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() + + ctrl_bit = ctrl_spec.get_single_ctrl_bit() + + if current_ctrl_bit is None: + # the easy case: use the controlled bloq + ctrl_bloq_and_ctrl_reg_name = get_ctrl_bloq_and_ctrl_reg_name(ctrl_bit) + if ctrl_bloq_and_ctrl_reg_name is None: + assert ctrl_bit != 1, "invalid usage: controlled-by-1 variant must be provided" + return _get_default_fallback() + + ctrl_bloq, ctrl_reg_name = ctrl_bloq_and_ctrl_reg_name + + 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_1_bloq_and_reg_name = get_ctrl_bloq_and_ctrl_reg_name(1) + assert ( + ctrl_1_bloq_and_reg_name is not None + ), "invalid usage: controlled-by-1 variant must be provided" + ctrl_1_bloq, ctrl_reg_name = ctrl_1_bloq_and_reg_name + + ctrl_bloq = _MultiControlledFromSinglyControlled( + cvs=(ctrl_bit, current_ctrl_bit), ctrl_bloq=ctrl_1_bloq, ctrl_reg_name=ctrl_reg_name + ) + + 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_reg_name) + + ctrl0 = cast(Soquet, ctrl0) + ctrl1 = cast(Soquet, ctrl1) + + # add the singly controlled bloq + in_soqs |= {ctrl_reg_name: np.array([ctrl0, ctrl1])} + ctrls, *out_soqs = bb.add_t(ctrl_bloq, **in_soqs) + assert isinstance(ctrls, np.ndarray) + ctrl0, ctrl1 = ctrls + + return [ctrl0], [ctrl1, *out_soqs] + + return ctrl_bloq, _adder + + +def get_ctrl_system_1bit_cv( + bloq: 'Bloq', + ctrl_spec: 'CtrlSpec', + *, + current_ctrl_bit: Optional['ControlBit'], + get_ctrl_bloq_and_ctrl_reg_name: Callable[['ControlBit'], tuple['Bloq', str]], +) -> tuple['Bloq', 'AddControlledT']: + """Build the control system for a bloq with specialized single-qubit controlled variants. + + Uses the provided specialized implementation when a singly-controlled variant of the bloq is + requested. When controlled by multiple qubits, the controls are reduced to a single qubit + and the singly-controlled bloq is used. + + The user must provide two specializations for the bloq: controlled by `1` and by `0`. + + When only one specialization (controlled by `1`) is known, use + :meth:`get_ctrl_system_1bit_cv_from_bloqs` instead. + + Args: + bloq: The current bloq. + ctrl_spec: The control specification + current_ctrl_bit: The control bit of the current bloq, one of `0, 1, None`. + get_ctrl_bloq_and_ctrl_reg_name: A callable that accepts a control bit (`0` or `1`), + and returns the controlled variant of this bloq and the name of the control register. + """ + return _get_ctrl_system_1bit_cv( + bloq, + ctrl_spec, + current_ctrl_bit=current_ctrl_bit, + get_ctrl_bloq_and_ctrl_reg_name=get_ctrl_bloq_and_ctrl_reg_name, + ) + + +def get_ctrl_system_1bit_cv_from_bloqs( + bloq: 'Bloq', + ctrl_spec: 'CtrlSpec', + *, + current_ctrl_bit: Optional['ControlBit'], + bloq_with_ctrl: 'Bloq', + ctrl_reg_name: 'str', +) -> tuple['Bloq', 'AddControlledT']: + """Helper to construct the control system given a singly-controlled variant of a bloq. + + Uses the provided specialized implementation when a singly-controlled (by `1`) variant of + the bloq is requested. When controlled by multiple qubits, the controls are reduced to a + single qubit and the singly-controlled bloq is used. + + When specializations for both cases - controlled by `1` and by `0` - are known, use + :meth:`get_ctrl_system_1bit_cv` instead. + + Args: + bloq: The current bloq. + ctrl_spec: The control specification + current_ctrl_bit: The control bit of the current bloq, one of `0, 1, None`. + bloq_with_ctrl: The variant of this bloq controlled by a single qubit in the `1` basis state. + ctrl_reg_name: The name of the control register for the controlled bloq variant(s). + """ + + def get_ctrl_bloq_and_ctrl_reg_name(cv: 'ControlBit') -> Optional[tuple['Bloq', str]]: + if cv == 1: + return bloq_with_ctrl, ctrl_reg_name + else: + return None + + return _get_ctrl_system_1bit_cv( + bloq, + ctrl_spec, + current_ctrl_bit=current_ctrl_bit, + get_ctrl_bloq_and_ctrl_reg_name=get_ctrl_bloq_and_ctrl_reg_name, + ) diff --git a/qualtran/bloqs/mcmt/specialized_ctrl_test.py b/qualtran/bloqs/mcmt/specialized_ctrl_test.py new file mode 100644 index 000000000..01282c1f2 --- /dev/null +++ b/qualtran/bloqs/mcmt/specialized_ctrl_test.py @@ -0,0 +1,203 @@ +# 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. +import itertools +from typing import Optional, Sequence, Tuple +from unittest.mock import ANY + +import attrs +import pytest + +from qualtran import ( + AddControlledT, + Bloq, + BloqBuilder, + CtrlSpec, + QAny, + QBit, + Register, + Signature, + SoquetT, +) +from qualtran.bloqs.mcmt import And +from qualtran.bloqs.mcmt.specialized_ctrl import ( + get_ctrl_system_1bit_cv, + get_ctrl_system_1bit_cv_from_bloqs, +) +from qualtran.resource_counting import CostKey, GateCounts, get_cost_value, QECGatesCost + + +@attrs.frozen +class AtomWithSpecializedControl(Bloq): + cv: Optional[int] = None + ctrl_reg_name: str = 'ctrl' + target_reg_name: str = 'q' + + @property + def signature(self) -> 'Signature': + n_ctrl = 1 if self.cv is not None else 0 + reg_name_map = {self.ctrl_reg_name: n_ctrl, self.target_reg_name: 2} + return Signature.build(**reg_name_map) + + def get_ctrl_system(self, ctrl_spec: 'CtrlSpec') -> Tuple['Bloq', 'AddControlledT']: + return get_ctrl_system_1bit_cv( + self, + ctrl_spec, + current_ctrl_bit=self.cv, + get_ctrl_bloq_and_ctrl_reg_name=lambda cv: ( + attrs.evolve(self, cv=cv), + self.ctrl_reg_name, + ), + ) + + @staticmethod + def cost_expr_for_cv(cv: Optional[int]): + import sympy + + c_unctrl = sympy.Symbol("_c_target_") + c_ctrl = sympy.Symbol("_c_ctrl_") + + if cv is None: + return c_unctrl + return c_unctrl + c_ctrl + + def my_static_costs(self, cost_key: 'CostKey'): + if cost_key == QECGatesCost(): + r = self.cost_expr_for_cv(self.cv) + return GateCounts(rotation=r) + + return NotImplemented + + +def ON(n: int = 1) -> CtrlSpec: + return CtrlSpec(cvs=[1] * n) + + +def OFF(n: int = 1) -> CtrlSpec: + return CtrlSpec(cvs=[0] * n) + + +@pytest.mark.parametrize( + 'ctrl_specs', + [ + [ON()], + [OFF()], + [OFF(), OFF()], + [OFF(4)], + [OFF(2), OFF(2)], + [ON(), OFF(5)], + [ON(), ON(), ON()], + [OFF(4), ON(3), OFF(5)], + ], +) +@pytest.mark.parametrize('ctrl_reg_name', ['ctrl', 'control']) +def test_custom_controlled(ctrl_specs: Sequence[CtrlSpec], ctrl_reg_name: str): + bloq: Bloq = AtomWithSpecializedControl(ctrl_reg_name=ctrl_reg_name) + for ctrl_spec in ctrl_specs: + bloq = bloq.controlled(ctrl_spec) + n_ctrls = sum(ctrl_spec.num_qubits for ctrl_spec in ctrl_specs) + + gc = get_cost_value(bloq, QECGatesCost()) + assert gc == GateCounts( + and_bloq=n_ctrls - 1, + rotation=AtomWithSpecializedControl.cost_expr_for_cv(1), + clifford=ANY, + measurement=ANY, + ) + + +@attrs.frozen +class TestAtom(Bloq): + tag: str + + @property + def signature(self) -> 'Signature': + return Signature.build(q=2) + + def get_ctrl_system(self, ctrl_spec: 'CtrlSpec') -> Tuple['Bloq', 'AddControlledT']: + return get_ctrl_system_1bit_cv_from_bloqs( + self, + ctrl_spec, + current_ctrl_bit=None, + bloq_with_ctrl=CTestAtom(self.tag), + ctrl_reg_name='ctrl', + ) + + +@attrs.frozen +class CTestAtom(Bloq): + tag: str + + @property + def signature(self) -> 'Signature': + return Signature.build(ctrl=1, q=2) + + def get_ctrl_system(self, ctrl_spec: 'CtrlSpec') -> Tuple['Bloq', 'AddControlledT']: + return get_ctrl_system_1bit_cv_from_bloqs( + self, ctrl_spec, current_ctrl_bit=1, bloq_with_ctrl=self, ctrl_reg_name='ctrl' + ) + + +def test_bloq_with_controlled_bloq(): + assert TestAtom('g').controlled() == CTestAtom('g') + + def _keep_and(b): + # TODO remove this after https://github.com/quantumlib/Qualtran/issues/1346 is resolved. + return isinstance(b, And) + + ctrl_bloq = CTestAtom('g').controlled() + _, sigma = ctrl_bloq.call_graph(keep=_keep_and) + assert sigma == {And(): 1, CTestAtom('g'): 1, And().adjoint(): 1} + + ctrl_bloq = CTestAtom('n').controlled(CtrlSpec(cvs=0)) + _, sigma = ctrl_bloq.call_graph(keep=_keep_and) + assert sigma == {And(0, 1): 1, CTestAtom('n'): 1, And(0, 1).adjoint(): 1} + + ctrl_bloq = TestAtom('nn').controlled(CtrlSpec(cvs=[0, 0])) + _, sigma = ctrl_bloq.call_graph(keep=_keep_and) + assert sigma == {And(0, 0): 1, CTestAtom('nn'): 1, And(0, 0).adjoint(): 1} + + +@attrs.frozen +class TestBloqWithDecompose(Bloq): + ctrl_reg_name: str + target_reg_name: str + + @property + def signature(self) -> 'Signature': + return Signature( + [Register(self.ctrl_reg_name, QBit()), Register(self.target_reg_name, QAny(2))] + ) + + def build_composite_bloq(self, bb: 'BloqBuilder', **soqs: 'SoquetT') -> dict[str, 'SoquetT']: + for _ in range(2): + soqs = bb.add_d( + AtomWithSpecializedControl( + cv=1, ctrl_reg_name=self.ctrl_reg_name, target_reg_name=self.target_reg_name + ), + **soqs, + ) + return soqs + + +@pytest.mark.parametrize( + ('ctrl_reg_name', 'target_reg_name'), + [ + (ctrl, targ) + for (ctrl, targ) in itertools.product(['ctrl', 'control', 'a', 'b'], repeat=2) + if ctrl != targ + ], +) +def test_get_ctrl_system(ctrl_reg_name: str, target_reg_name: str): + bloq = TestBloqWithDecompose(ctrl_reg_name, target_reg_name).controlled() + _ = bloq.decompose_bloq().flatten() diff --git a/qualtran/bloqs/multiplexers/apply_lth_bloq.py b/qualtran/bloqs/multiplexers/apply_lth_bloq.py index a0de211a6..5ff1e6c74 100644 --- a/qualtran/bloqs/multiplexers/apply_lth_bloq.py +++ b/qualtran/bloqs/multiplexers/apply_lth_bloq.py @@ -17,12 +17,21 @@ import cirq import numpy as np -from attrs import field, frozen +from attrs import evolve, field, frozen from numpy.typing import NDArray -from qualtran import Bloq, bloq_example, BloqDocSpec, BQUInt, QBit, Register, Side +from qualtran import ( + AddControlledT, + Bloq, + bloq_example, + BloqDocSpec, + BQUInt, + CtrlSpec, + QBit, + Register, + Side, +) from qualtran._infra.gate_with_registers import merge_qubits -from qualtran._infra.single_qubit_controlled import SpecializedSingleQubitControlledExtension from qualtran.bloqs.multiplexers.select_base import SelectOracle from qualtran.bloqs.multiplexers.unary_iteration_bloq import UnaryIterationGate from qualtran.resource_counting import BloqCountT @@ -30,7 +39,7 @@ @frozen -class ApplyLthBloq(UnaryIterationGate, SpecializedSingleQubitControlledExtension, SelectOracle): # type: ignore[misc] +class ApplyLthBloq(UnaryIterationGate, SelectOracle): # type: ignore[misc] r"""A SELECT operation that executes one of a list of bloqs $U_l$ based on a quantum index: $$ @@ -108,6 +117,16 @@ def nth_operation( target_qubits = merge_qubits(bloq.signature, **targets) return bloq.controlled().on(control, *target_qubits) + def get_ctrl_system(self, ctrl_spec: 'CtrlSpec') -> Tuple['Bloq', 'AddControlledT']: + from qualtran.bloqs.mcmt.specialized_ctrl import get_ctrl_system_1bit_cv + + return get_ctrl_system_1bit_cv( + self, + ctrl_spec=ctrl_spec, + current_ctrl_bit=self.control_val, + get_ctrl_bloq_and_ctrl_reg_name=lambda cv: (evolve(self, control_val=cv), 'control'), + ) + @bloq_example def _apply_lth_bloq() -> ApplyLthBloq: diff --git a/qualtran/bloqs/multiplexers/select_pauli_lcu.py b/qualtran/bloqs/multiplexers/select_pauli_lcu.py index 840512a0b..f94b890ac 100644 --- a/qualtran/bloqs/multiplexers/select_pauli_lcu.py +++ b/qualtran/bloqs/multiplexers/select_pauli_lcu.py @@ -22,8 +22,17 @@ import numpy as np from numpy.typing import NDArray -from qualtran import bloq_example, BloqDocSpec, BQUInt, QAny, QBit, Register -from qualtran._infra.single_qubit_controlled import SpecializedSingleQubitControlledExtension +from qualtran import ( + AddControlledT, + Bloq, + bloq_example, + BloqDocSpec, + BQUInt, + CtrlSpec, + QAny, + QBit, + Register, +) from qualtran.bloqs.multiplexers.select_base import SelectOracle from qualtran.bloqs.multiplexers.unary_iteration_bloq import UnaryIterationGate from qualtran.resource_counting.generalizers import ( @@ -39,7 +48,7 @@ def _to_tuple(x: Iterable[cirq.DensePauliString]) -> Sequence[cirq.DensePauliStr @attrs.frozen -class SelectPauliLCU(SelectOracle, UnaryIterationGate, SpecializedSingleQubitControlledExtension): # type: ignore[misc] +class SelectPauliLCU(SelectOracle, UnaryIterationGate): # type: ignore[misc] r"""A SELECT bloq for selecting and applying operators from an array of `PauliString`s. $$ @@ -117,6 +126,19 @@ def nth_operation( # type: ignore[override] ps = self.select_unitaries[selection].on(*target) return ps.with_coefficient(np.sign(complex(ps.coefficient).real)).controlled_by(control) + def get_ctrl_system(self, ctrl_spec: 'CtrlSpec') -> Tuple['Bloq', 'AddControlledT']: + from qualtran.bloqs.mcmt.specialized_ctrl import get_ctrl_system_1bit_cv + + return get_ctrl_system_1bit_cv( + self, + ctrl_spec=ctrl_spec, + current_ctrl_bit=self.control_val, + get_ctrl_bloq_and_ctrl_reg_name=lambda cv: ( + attrs.evolve(self, control_val=cv), + 'control', + ), + ) + @bloq_example(generalizer=[cirq_to_bloqs, ignore_split_join, ignore_cliffords]) def _select_pauli_lcu() -> SelectPauliLCU: diff --git a/qualtran/bloqs/reflections/reflection_using_prepare.py b/qualtran/bloqs/reflections/reflection_using_prepare.py index 885f9317f..a3cd4c2b3 100644 --- a/qualtran/bloqs/reflections/reflection_using_prepare.py +++ b/qualtran/bloqs/reflections/reflection_using_prepare.py @@ -20,9 +20,17 @@ import numpy as np from numpy.typing import NDArray -from qualtran import Bloq, bloq_example, BloqDocSpec, CtrlSpec, QBit, Register, Signature +from qualtran import ( + AddControlledT, + Bloq, + bloq_example, + BloqDocSpec, + CtrlSpec, + QBit, + Register, + Signature, +) from qualtran._infra.gate_with_registers import GateWithRegisters, merge_qubits, total_bits -from qualtran._infra.single_qubit_controlled import SpecializedSingleQubitControlledExtension from qualtran.bloqs.basic_gates.global_phase import GlobalPhase from qualtran.bloqs.basic_gates.x_basis import XGate from qualtran.bloqs.mcmt import MultiControlZ @@ -41,7 +49,7 @@ @attrs.frozen(cache_hash=True) -class ReflectionUsingPrepare(GateWithRegisters, SpecializedSingleQubitControlledExtension): # type: ignore[misc] +class ReflectionUsingPrepare(GateWithRegisters): r"""Applies reflection around a state prepared by `prepare_gate` Applies $R_{s, g=1} = g (I - 2|s\rangle\langle s|)$ using $R_{s} = @@ -186,6 +194,19 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': def adjoint(self) -> 'ReflectionUsingPrepare': return self + def get_ctrl_system(self, ctrl_spec: 'CtrlSpec') -> Tuple['Bloq', 'AddControlledT']: + from qualtran.bloqs.mcmt.specialized_ctrl import get_ctrl_system_1bit_cv + + return get_ctrl_system_1bit_cv( + self, + ctrl_spec=ctrl_spec, + current_ctrl_bit=self.control_val, + get_ctrl_bloq_and_ctrl_reg_name=lambda cv: ( + attrs.evolve(self, control_val=cv), + 'control', + ), + ) + @bloq_example(generalizer=ignore_split_join) def _refl_using_prep() -> ReflectionUsingPrepare: