diff --git a/qualtran/_infra/controlled.py b/qualtran/_infra/controlled.py index 196e87cb6..3c9c83eef 100644 --- a/qualtran/_infra/controlled.py +++ b/qualtran/_infra/controlled.py @@ -29,6 +29,7 @@ import attrs import cirq import numpy as np +import quimb.tensor as qtn from numpy.typing import NDArray from .bloq import Bloq @@ -250,7 +251,8 @@ class AddControlledT(Protocol): def __call__( self, bb: 'BloqBuilder', ctrl_soqs: Sequence['SoquetT'], in_soqs: Dict[str, 'SoquetT'] - ) -> Tuple[Iterable['SoquetT'], Iterable['SoquetT']]: ... + ) -> Tuple[Iterable['SoquetT'], Iterable['SoquetT']]: + ... def _get_nice_ctrl_reg_names(reg_names: List[str], n: int) -> Tuple[str, ...]: @@ -327,13 +329,16 @@ def ctrl_reg_names(self) -> Sequence[str]: return _get_nice_ctrl_reg_names(reg_names, n) @cached_property - def signature(self) -> 'Signature': - # Prepend register(s) corresponding to `ctrl_spec`. - ctrl_regs = tuple( + def ctrl_regs(self) -> Tuple[Register, ...]: + return tuple( Register(name=self.ctrl_reg_names[i], dtype=qdtype, shape=shape, side=Side.THRU) for i, (qdtype, shape) in enumerate(self.ctrl_spec.activation_function_dtypes()) ) - return Signature(ctrl_regs + tuple(self.subbloq.signature)) + + @cached_property + def signature(self) -> 'Signature': + # Prepend register(s) corresponding to `ctrl_spec`. + return Signature(self.ctrl_regs + tuple(self.subbloq.signature)) def decompose_bloq(self) -> 'CompositeBloq': # Use subbloq's decomposition but wire up the additional ctrl_soqs. @@ -376,6 +381,35 @@ def on_classical_vals(self, **vals: 'ClassicalValT') -> Dict[str, 'ClassicalValT return vals + def add_my_tensors( + self, + tn: 'qtn.TensorNetwork', + tag: Any, + *, + incoming: Dict[str, 'SoquetT'], + outgoing: Dict[str, 'SoquetT'], + ): + from qualtran._infra.composite_bloq import _flatten_soquet_collection + from qualtran.simulation.tensor._tensor_data_manipulation import ( + active_space_for_ctrl_spec, + eye_tensor_for_signature, + tensor_shape_from_signature, + ) + + # Create an identity tensor corresponding to the signature of current Bloq + data = eye_tensor_for_signature(self.signature) + # Verify it has the right shape + in_ind = _flatten_soquet_collection(incoming[reg.name] for reg in self.signature.lefts()) + out_ind = _flatten_soquet_collection(outgoing[reg.name] for reg in self.signature.rights()) + assert data.shape == tuple(2**soq.reg.bitsize for ind in [out_ind, in_ind] for soq in ind) + # Figure out the ctrl indexes for which the ctrl is "active" + active_idx = active_space_for_ctrl_spec(self.signature, self.ctrl_spec) + # Put the subbloq tensor at indices where ctrl is active. + subbloq_shape = tensor_shape_from_signature(self.subbloq.signature) + data[active_idx] = self.subbloq.tensor_contract().reshape(subbloq_shape) + # Add the data to the tensor network. + tn.add(qtn.Tensor(data=data, inds=out_ind + in_ind, tags=[self.short_name(), tag])) + def wire_symbol(self, soq: 'Soquet') -> 'WireSymbol': if soq.reg.name not in self.ctrl_reg_names: # Delegate to subbloq diff --git a/qualtran/_infra/controlled_test.py b/qualtran/_infra/controlled_test.py index edb2448b8..701f6fec5 100644 --- a/qualtran/_infra/controlled_test.py +++ b/qualtran/_infra/controlled_test.py @@ -11,17 +11,41 @@ # 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 List +from typing import Dict, List, Tuple +import attrs import cirq import numpy as np import pytest import qualtran.testing as qlt_testing -from qualtran import Bloq, Controlled, CtrlSpec, QBit, QInt, QUInt +from qualtran import ( + Bloq, + BloqBuilder, + Controlled, + CtrlSpec, + QBit, + QInt, + QUInt, + Register, + Side, + Signature, +) from qualtran._infra.gate_with_registers import get_named_qubits, merge_qubits -from qualtran.bloqs.basic_gates import Swap, XGate +from qualtran.bloqs.basic_gates import ( + CSwap, + IntEffect, + IntState, + OneState, + Swap, + XGate, + YGate, + ZeroState, + ZGate, +) from qualtran.bloqs.for_testing import TestAtom, TestParallelCombo, TestSerialCombo +from qualtran.bloqs.mcmt import And +from qualtran.cirq_interop.testing import GateHelper from qualtran.drawing import get_musical_score_data from qualtran.drawing.musical_score import Circle, SoqData, TextBox @@ -294,3 +318,84 @@ def test_classical_sim_int_multi_reg(): @pytest.mark.notebook def test_notebook(): qlt_testing.execute_notebook('../Controlled') + + +def _verify_ctrl_tensor_for_unitary(ctrl_spec: CtrlSpec, bloq: Bloq, gate: cirq.Gate): + cbloq = Controlled(bloq, ctrl_spec) + cgate = gate.controlled(control_values=ctrl_spec.to_cirq_cv()) + np.testing.assert_array_equal(cbloq.tensor_contract(), cirq.unitary(cgate)) + + +interesting_ctrl_specs = [ + CtrlSpec(), + CtrlSpec(cvs=0), + CtrlSpec(qdtypes=QUInt(4), cvs=0b0110), + CtrlSpec(cvs=[0, 1, 1, 0]), + CtrlSpec(qdtypes=[QBit(), QBit()], cvs=[[0, 1], [1, 0]]), +] + + +@pytest.mark.parametrize('ctrl_spec', interesting_ctrl_specs) +def test_controlled_tensor_for_unitary(ctrl_spec: CtrlSpec): + # Test one qubit unitaries + _verify_ctrl_tensor_for_unitary(ctrl_spec, XGate(), cirq.X) + _verify_ctrl_tensor_for_unitary(ctrl_spec, YGate(), cirq.Y) + _verify_ctrl_tensor_for_unitary(ctrl_spec, ZGate(), cirq.Z) + # Test multi-qubit unitaries with non-trivial signature + _verify_ctrl_tensor_for_unitary(ctrl_spec, CSwap(3), CSwap(3)) + + +@attrs.frozen +class TestCtrlStatePrepAnd(Bloq): + """Decomposes into a Controlled-AND gate + int effects & targets where ctrl is active. + + Tensor contraction should give the output state vector corresponding to applying an + `And(and_ctrl)`; assuming all the control bits are active. + """ + + ctrl_spec: CtrlSpec + and_ctrl: Tuple[int, int] + + @property + def signature(self) -> 'Signature': + return Signature([Register('x', QBit(), shape=(3,), side=Side.RIGHT)]) + + def build_composite_bloq(self, bb: 'BloqBuilder') -> Dict[str, 'SoquetT']: + one_or_zero = [ZeroState(), OneState()] + cbloq = Controlled(And(*self.and_ctrl), ctrl_spec=self.ctrl_spec) + + ctrl_soqs = {} + for reg, cvs in zip(cbloq.ctrl_regs, self.ctrl_spec.cvs): + soqs = np.empty(shape=reg.shape, dtype=object) + for idx in reg.all_idxs(): + soqs[idx] = bb.add(IntState(val=cvs[idx], bitsize=reg.dtype.num_qubits)) + ctrl_soqs[reg.name] = soqs + + and_ctrl = [bb.add(one_or_zero[cv]) for cv in self.and_ctrl] + + ctrl_soqs = bb.add_d(cbloq, **ctrl_soqs, ctrl=and_ctrl) + out_soqs = [*ctrl_soqs.pop('ctrl'), ctrl_soqs.pop('target')] + + for reg, cvs in zip(cbloq.ctrl_regs, self.ctrl_spec.cvs): + for idx in reg.all_idxs(): + ctrl_soq = np.asarray(ctrl_soqs[reg.name])[idx] + bb.add(IntEffect(val=cvs[idx], bitsize=reg.dtype.num_qubits), val=ctrl_soq) + return {'x': out_soqs} + + +def _verify_ctrl_tensor_for_and(ctrl_spec: CtrlSpec, and_ctrl: Tuple[int, int]): + cbloq = TestCtrlStatePrepAnd(ctrl_spec, and_ctrl) + bloq_tensor = cbloq.tensor_contract() + cirq_state_vector = GateHelper(And(*and_ctrl)).circuit.final_state_vector( + initial_state=and_ctrl + (0,) + ) + np.testing.assert_allclose(bloq_tensor, cirq_state_vector, atol=1e-8) + + +@pytest.mark.parametrize('ctrl_spec', interesting_ctrl_specs) +def test_controlled_tensor_for_and_bloq(ctrl_spec: CtrlSpec): + # Test AND gate with one-sided signature (aka controlled state preparation). + _verify_ctrl_tensor_for_and(ctrl_spec, (1, 1)) + _verify_ctrl_tensor_for_and(ctrl_spec, (1, 0)) + _verify_ctrl_tensor_for_and(ctrl_spec, (0, 1)) + _verify_ctrl_tensor_for_and(ctrl_spec, (0, 0)) diff --git a/qualtran/cirq_interop/_cirq_to_bloq.py b/qualtran/cirq_interop/_cirq_to_bloq.py index 7b404837f..cce2b383b 100644 --- a/qualtran/cirq_interop/_cirq_to_bloq.py +++ b/qualtran/cirq_interop/_cirq_to_bloq.py @@ -15,7 +15,6 @@ """Cirq gates/circuits to Qualtran Bloqs conversion.""" import abc import itertools -from collections import defaultdict from functools import cached_property from typing import Any, Dict, List, Optional, Sequence, Tuple, TYPE_CHECKING, Union @@ -46,6 +45,9 @@ ) from qualtran.cirq_interop._interop_qubit_manager import InteropQubitManager from qualtran.cirq_interop.t_complexity_protocol import t_complexity, TComplexity +from qualtran.simulation.tensor._tensor_data_manipulation import ( + tensor_data_from_unitary_and_signature, +) if TYPE_CHECKING: from qualtran.drawing import WireSymbol @@ -194,37 +196,7 @@ def _add_my_tensors_from_gate( f"CirqGateAsBloq.add_my_tensors is currently supported only for unitary gates. " f"Found {gate}." ) - unitary_shape = [] - reg_to_idx = defaultdict(list) - for reg in signature: - start = len(unitary_shape) - for i in range(int(np.prod(reg.shape))): - reg_to_idx[reg.name].append(start + i) - unitary_shape.append(2**reg.bitsize) - - unitary_shape = (*unitary_shape, *unitary_shape) - unitary = cirq.unitary(gate).reshape(unitary_shape) - idx: List[Union[int, slice]] = [slice(x) for x in unitary_shape] - n = len(unitary_shape) // 2 - for reg in signature: - if reg.side == Side.LEFT: - for i in reg_to_idx[reg.name]: - # LEFT register ends, extract right subspace that's equivalent to 0. - idx[i] = 0 - if reg.side == Side.RIGHT: - for i in reg_to_idx[reg.name]: - # Right register begins, extract the left subspace that's equivalent to 0. - idx[i + n] = 0 - unitary = unitary[tuple(idx)] - new_shape = tuple( - [ - *itertools.chain.from_iterable( - (2**reg.bitsize,) * int(np.prod(reg.shape)) - for reg in [*signature.rights(), *signature.lefts()] - ) - ] - ) - assert unitary.shape == new_shape + unitary = tensor_data_from_unitary_and_signature(cirq.unitary(gate), signature) incoming_list = [ *itertools.chain.from_iterable( [np.array(incoming[reg.name]).flatten() for reg in signature.lefts()] diff --git a/qualtran/simulation/tensor/__init__.py b/qualtran/simulation/tensor/__init__.py index 8c4c36856..6f5b1f814 100644 --- a/qualtran/simulation/tensor/__init__.py +++ b/qualtran/simulation/tensor/__init__.py @@ -16,3 +16,10 @@ from ._dense import bloq_to_dense, get_right_and_left_inds from ._flattening import bloq_has_custom_tensors, flatten_for_tensor_contraction from ._quimb import cbloq_as_contracted_tensor, cbloq_to_quimb +from ._tensor_data_manipulation import ( + active_space_for_ctrl_spec, + eye_tensor_for_signature, + tensor_data_from_unitary_and_signature, + tensor_out_inp_shape_from_signature, + tensor_shape_from_signature, +) diff --git a/qualtran/simulation/tensor/_tensor_data_manipulation.py b/qualtran/simulation/tensor/_tensor_data_manipulation.py new file mode 100644 index 000000000..97f8852e2 --- /dev/null +++ b/qualtran/simulation/tensor/_tensor_data_manipulation.py @@ -0,0 +1,136 @@ +# 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. + +"""Utility methods to generate and manipulate tensor data for Bloqs.""" +import itertools +from typing import List, Tuple, Union + +import attrs +import numpy as np + +from qualtran import CtrlSpec, Side, Signature + + +def tensor_out_inp_shape_from_signature( + signature: Signature, +) -> Tuple[Tuple[int, ...], Tuple[int, ...]]: + """Returns a tuple for tensor data corresponding to signature. + + Tensor data for a bloq with a given `signature` can be expressed as a ndarray of + shape `out_indices_shape + inp_indices_shape` where + + 1. out_indices_shape - A tuple of values `2 ** soq.reg.bitsize` for every soquet `soq` + corresponding to the RIGHT registers in signature. + 2. inp_indices_shape - A tuple of values `2 ** soq.reg.bitsize` for every soquet `soq` + corresponding to the LEFT registers in signature. + + This method returns a tuple of (out_indices_shape, inp_indices_shape). + """ + inp_indices_shape = [2**reg.bitsize for reg in signature.lefts() for _ in reg.all_idxs()] + out_indices_shape = [2**reg.bitsize for reg in signature.rights() for _ in reg.all_idxs()] + return tuple(out_indices_shape), tuple(inp_indices_shape) + + +def tensor_shape_from_signature(signature: Signature) -> Tuple[int, ...]: + """Returns a tuple for tensor data corresponding to signature. + + Tensor data for a bloq with a given `signature` can be expressed as a ndarray of + shape `out_indices_shape + inp_indices_shape` where + + 1. out_indices_shape - A tuple of values `2 ** soq.reg.bitsize` for every soquet `soq` + corresponding to the RIGHT registers in signature. + 2. inp_indices_shape - A tuple of values `2 ** soq.reg.bitsize` for every soquet `soq` + corresponding to the LEFT registers in signature. + + This method returns a tuple of (*out_indices_shape, *inp_indices_shape). + """ + out_shape, inp_shape = tensor_out_inp_shape_from_signature(signature) + return out_shape + inp_shape + + +def active_space_for_ctrl_spec( + signature: Signature, ctrl_spec: CtrlSpec +) -> Tuple[Union[int, slice], ...]: + """Returns the "active" subspace corresponding to `signature` and `ctrl_spec`. + + Assumes first n-registers for `signature` are control registers corresponding to `ctrl_spec`. + Returns a tuple of indices/slices that can be used to address into the ndarray, representing + tensor data of shape `tensor_shape_from_signature(signature)`, and access the active subspace. + """ + out_ind, inp_ind = tensor_out_inp_shape_from_signature(signature) + data_shape = out_ind + inp_ind + active_idx: List[Union[int, slice]] = [slice(x) for x in data_shape] + ctrl_idx = 0 + for cv in ctrl_spec.cvs: + for idx in itertools.product(*[range(sh) for sh in cv.shape]): + active_idx[ctrl_idx] = int(cv[idx]) + active_idx[ctrl_idx + len(out_ind)] = int(cv[idx]) + ctrl_idx += 1 + return tuple(active_idx) + + +def _n_qubits(signature: Signature) -> int: + return sum(reg.total_bits() for reg in signature) + + +def eye_tensor_for_signature(signature: Signature) -> np.ndarray: + """Returns an identity tensor with shape `tensor_shape_from_signature(signature)`""" + return tensor_data_from_unitary_and_signature( + np.eye(2 ** _n_qubits(signature), dtype=np.complex128), signature + ) + + +def tensor_data_from_unitary_and_signature(unitary: np.ndarray, signature: Signature) -> np.ndarray: + """Returns tensor data respecting `signature` corresponding to `unitary` + + For a given input unitary, we extract the action of the unitary on a subspace where + input qubits corresponding to LEFT registers and output qubits corresponding to RIGHT + registers in `signature` are 0. + + The input unitary is assumed to act on `_n_qubits(signature)`, and thus is of shape + `(2 ** _n_qubits(signature), 2 ** _n_qubits(signature))` where `_n_qubits(signature)` + is `sum(reg.total_bits() for reg in signature)`. + + The shape of the returned tensor matches `tensor_shape_from_signature(signature)`. + """ + + # Reshape the unitary into correct shape assuming all registers are THRU registers. + assert unitary.shape == (2 ** _n_qubits(signature),) * 2 + signature_ignoring_sides = Signature([attrs.evolve(reg, side=Side.THRU) for reg in signature]) + unitary_shape = tensor_shape_from_signature(signature_ignoring_sides) + n = len(unitary_shape) // 2 + unitary = unitary.reshape(unitary_shape) + + # Find the subspace corresponding to registers with sides. + idx: List[Union[int, slice]] = [slice(x) for x in unitary_shape] + curr_idx = 0 + for reg in signature: + if reg.side == Side.LEFT: + for _ in reg.all_idxs(): + # LEFT register ends, extract right subspace that's equivalent to 0. + idx[curr_idx] = 0 + curr_idx += 1 + if reg.side == Side.RIGHT: + for _ in reg.all_idxs(): + # Right register begins, extract the left subspace that's equivalent to 0. + idx[curr_idx + n] = 0 + curr_idx += 1 + if reg.side == Side.THRU: + curr_idx += int(np.prod(reg.shape)) + + # Extract the subspace, assert it has the correct shape corresponding to `signature` and + # return the result. + unitary = unitary[tuple(idx)] + assert unitary.shape == tensor_shape_from_signature(signature) + return unitary diff --git a/qualtran/simulation/tensor/_tensor_data_manipulation_test.py b/qualtran/simulation/tensor/_tensor_data_manipulation_test.py new file mode 100644 index 000000000..5f4142b91 --- /dev/null +++ b/qualtran/simulation/tensor/_tensor_data_manipulation_test.py @@ -0,0 +1,90 @@ +# 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 numpy as np +import pytest + +from qualtran import CtrlSpec, QAny, Register, Side, Signature +from qualtran.simulation.tensor import ( + active_space_for_ctrl_spec, + eye_tensor_for_signature, + tensor_out_inp_shape_from_signature, + tensor_shape_from_signature, +) + + +def test_tensor_shape_from_signature(): + # Test trivial case + assert tensor_shape_from_signature(Signature.build(x=1)) == (2, 2) + + # Test left / right / thru cases + left_register = Register('left', QAny(1), side=Side.LEFT, shape=()) + right_register = Register('right', QAny(2), side=Side.RIGHT, shape=(2,)) + thru_register = Register('thru', QAny(3), side=Side.THRU, shape=(2, 2)) + + assert tensor_shape_from_signature(Signature([left_register])) == (2,) + assert tensor_shape_from_signature(Signature([right_register])) == (4, 4) + assert tensor_shape_from_signature(Signature([thru_register])) == (8, 8, 8, 8) * 2 + + # Test all 3 and asser that ordering is preserved + signature = Signature([left_register, right_register, thru_register]) + inp_shape = (2,) + (8, 8, 8, 8) + out_shape = (4, 4) + (8, 8, 8, 8) + assert tensor_shape_from_signature(signature) == out_shape + inp_shape + assert tensor_out_inp_shape_from_signature(signature) == (out_shape, inp_shape) + + +def test_eye_tensor_for_signature(): + # Test left / right / thru cases individually + left_register = Register('left', QAny(1), side=Side.LEFT) + right_register = Register('right', QAny(1), side=Side.RIGHT) + thru_register = Register('thru', QAny(1), side=Side.THRU) + left_data = right_data = np.array([1, 0]) + np.testing.assert_allclose(eye_tensor_for_signature(Signature([left_register])), left_data) + np.testing.assert_allclose(eye_tensor_for_signature(Signature([right_register])), right_data) + thru_data = np.eye(2) + np.testing.assert_allclose(eye_tensor_for_signature(Signature([thru_register])), thru_data) + + # Test LEFT + RIGHT case + actual = eye_tensor_for_signature(Signature([left_register, right_register])) + expected = np.array([[1, 0], [0, 0]]) # 1 only when LEFT is 0 and RIGHT is 0 + np.testing.assert_allclose(actual, expected) + + # It's helpful to take einsum for more complicated cases, as shown below, because each + # variable corresponds to an output / input index. + + # Test LEFT + THRU case + actual = eye_tensor_for_signature(Signature([left_register, thru_register])) + expected = np.einsum('jk,i->jik', thru_data, left_data) + np.testing.assert_allclose(actual, expected) + + # Test RIGHT + THRU case + actual = eye_tensor_for_signature(Signature([right_register, thru_register])) + expected = np.einsum('jk,i->ijk', thru_data, right_data) + np.testing.assert_allclose(actual, expected) + + # Test LEFT + RIGHT + THRU case + actual = eye_tensor_for_signature(Signature([left_register, right_register, thru_register])) + expected = np.einsum('i,jk,l->ijlk', right_data, thru_data, left_data) + np.testing.assert_allclose(actual, expected) + + +@pytest.mark.parametrize('cv', [0, 1]) +def test_active_space_for_ctrl_spec(cv: int): + ctrl_spec = CtrlSpec(cvs=cv) + signature = Signature([Register('ctrl', QAny(1)), Register('q', QAny(1))]) + assert active_space_for_ctrl_spec(signature, ctrl_spec) == (cv, slice(2), cv, slice(2)) + signature = Signature([Register('ctrl', QAny(1)), Register('q', QAny(1), side=Side.LEFT)]) + assert active_space_for_ctrl_spec(signature, ctrl_spec) == (cv, cv, slice(2)) + signature = Signature([Register('ctrl', QAny(1)), Register('q', QAny(1), side=Side.RIGHT)]) + assert active_space_for_ctrl_spec(signature, ctrl_spec) == (cv, slice(2), cv)