diff --git a/qualtran/_infra/bloq.py b/qualtran/_infra/bloq.py index 430486ec0c..8572718445 100644 --- a/qualtran/_infra/bloq.py +++ b/qualtran/_infra/bloq.py @@ -61,6 +61,7 @@ SympySymbolAllocator, ) from qualtran.simulation.classical_sim import ClassicalValRetT, ClassicalValT + from qualtran.simulation.tensor import DiscardInd def _decompose_from_build_composite_bloq(bloq: 'Bloq') -> 'CompositeBloq': @@ -68,6 +69,10 @@ def _decompose_from_build_composite_bloq(bloq: 'Bloq') -> 'CompositeBloq': bb, initial_soqs = BloqBuilder.from_signature(bloq.signature, add_registers_allowed=False) out_soqs = bloq.build_composite_bloq(bb=bb, **initial_soqs) + if not isinstance(out_soqs, dict): + raise ValueError( + f'{bloq}.build_composite_bloq must return a dictionary mapping right register names to output soquets.' + ) return bb.finalize(**out_soqs) @@ -258,21 +263,45 @@ def call_classically( res = self.as_composite_bloq().on_classical_vals(**vals) return tuple(res[reg.name] for reg in self.signature.rights()) - def tensor_contract(self) -> 'NDArray': - """Return a contracted, dense ndarray representing this bloq. + def tensor_contract(self, superoperator: bool = False) -> 'NDArray': + """Return a contracted, dense ndarray encoding of this bloq. + + This method decomposes and flattens this bloq into a factorized CompositeBloq, + turns that composite bloq into a Quimb tensor network, and contracts it into a dense + ndarray. + + The returned array will be 0-, 1-, 2-, or 4-dimensional with indices arranged according to the + bloq's signature and the type of simulation requested via the `superoperator` flag. + + If `superoperator` is set to False (the default), a pure-state tensor network will be + constructed. + - If `bloq` has all thru-registers, the dense tensor will be 2-dimensional with shape `(n, n)` + where `n` is the number of bits in the signature. We follow the linear algebra convention + and order the indices as (right, left) so the matrix-vector product can be used to evolve + a state vector. + - If `bloq` has all left- or all right-registers, the tensor will be 1-dimensional with + shape `(n,)`. Note that we do not distinguish between 'row' and 'column' vectors in this + function. + - If `bloq` has no external registers, the contracted form is a 0-dimensional complex number. + + If `superoperator` is set to True, an open-system tensor network will be constructed. + - States result in a 2-dimensional density matrix with indices (right_forward, right_backward) + or (left_forward, left_backward) depending on whether they're input or output states. + - Operations result in a 4-dimensional tensor with indices (right_forward, right_backward, + left_forward, left_backward). - This constructs a tensor network and then contracts it according to our registers, - i.e. the dangling indices. The returned array will be 0-, 1- or 2-dimensional. If it is - a 2-dimensional matrix, we follow the quantum computing / matrix multiplication convention - of (right, left) indices. + Args: + superoperator: If toggled to True, do an open-system simulation. This supports + non-unitary operations like measurement, but is more costly and results in + higher-dimension resultant tensors. """ from qualtran.simulation.tensor import bloq_to_dense - return bloq_to_dense(self) + return bloq_to_dense(self, superoperator=superoperator) def my_tensors( self, incoming: Dict[str, 'ConnectionT'], outgoing: Dict[str, 'ConnectionT'] - ) -> List['qtn.Tensor']: + ) -> List[Union['qtn.Tensor', 'DiscardInd']]: """Override this method to support native quimb simulation of this Bloq. This method is responsible for returning tensors corresponding to the unitary, state, or diff --git a/qualtran/_infra/composite_bloq.py b/qualtran/_infra/composite_bloq.py index 3747a57cfe..9832b9ffaa 100644 --- a/qualtran/_infra/composite_bloq.py +++ b/qualtran/_infra/composite_bloq.py @@ -371,7 +371,7 @@ def flatten_once( in_soqs = _map_soqs(in_soqs, soq_map) # update `in_soqs` from old to new. if pred(binst): try: - new_out_soqs = bb.add_from(binst.bloq.decompose_bloq(), **in_soqs) + new_out_soqs = bb.add_from(binst.bloq, **in_soqs) did_work = True except (DecomposeTypeError, DecomposeNotImplementedError): new_out_soqs = tuple(soq for _, soq in bb._add_binst(binst, in_soqs=in_soqs)) diff --git a/qualtran/_infra/controlled_test.py b/qualtran/_infra/controlled_test.py index 3eff66f029..0a909d79ef 100644 --- a/qualtran/_infra/controlled_test.py +++ b/qualtran/_infra/controlled_test.py @@ -33,7 +33,7 @@ from qualtran.bloqs.for_testing import TestAtom, TestParallelCombo, TestSerialCombo from qualtran.drawing import get_musical_score_data from qualtran.drawing.musical_score import Circle, SoqData, TextBox -from qualtran.simulation.tensor import cbloq_to_quimb, get_right_and_left_inds +from qualtran.simulation.tensor import cbloq_to_quimb, quimb_to_dense from qualtran.symbolics import Shaped if TYPE_CHECKING: @@ -432,10 +432,8 @@ def test_controlled_tensor_without_decompose(): cgate = cirq.ControlledGate(cirq.CSWAP, control_values=ctrl_spec.to_cirq_cv()) tn = cbloq_to_quimb(ctrl_bloq.as_composite_bloq()) - # pylint: disable=unbalanced-tuple-unpacking - right, left = get_right_and_left_inds(tn, ctrl_bloq.signature) - # pylint: enable=unbalanced-tuple-unpacking - np.testing.assert_allclose(tn.to_dense(right, left), cirq.unitary(cgate), atol=1e-8) + tn_dense = quimb_to_dense(tn, ctrl_bloq.signature) + np.testing.assert_allclose(tn_dense, cirq.unitary(cgate), atol=1e-8) np.testing.assert_allclose(ctrl_bloq.tensor_contract(), cirq.unitary(cgate), atol=1e-8) diff --git a/qualtran/bloqs/basic_gates/__init__.py b/qualtran/bloqs/basic_gates/__init__.py index 66910b89f1..d89e5080a9 100644 --- a/qualtran/bloqs/basic_gates/__init__.py +++ b/qualtran/bloqs/basic_gates/__init__.py @@ -22,6 +22,7 @@ """ from .cnot import CNOT +from .discard import Discard, DiscardQ from .global_phase import GlobalPhase from .hadamard import CHadamard, Hadamard from .identity import Identity @@ -35,4 +36,14 @@ from .toffoli import Toffoli from .x_basis import MinusEffect, MinusState, PlusEffect, PlusState, XGate from .y_gate import CYGate, YGate -from .z_basis import CZ, IntEffect, IntState, OneEffect, OneState, ZeroEffect, ZeroState, ZGate +from .z_basis import ( + CZ, + IntEffect, + IntState, + MeasZ, + OneEffect, + OneState, + ZeroEffect, + ZeroState, + ZGate, +) diff --git a/qualtran/bloqs/basic_gates/discard.py b/qualtran/bloqs/basic_gates/discard.py new file mode 100644 index 0000000000..90df55c367 --- /dev/null +++ b/qualtran/bloqs/basic_gates/discard.py @@ -0,0 +1,68 @@ +# Copyright 2025 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 Dict, List, TYPE_CHECKING + +from attrs import frozen + +from qualtran import Bloq, CBit, ConnectionT, QBit, Register, Side, Signature +from qualtran.simulation.classical_sim import ClassicalValT + +if TYPE_CHECKING: + from qualtran.simulation.tensor import DiscardInd + + +@frozen +class Discard(Bloq): + """Discard a classical bit. + + This is an allowed operation. + """ + + @cached_property + def signature(self) -> 'Signature': + return Signature([Register('c', CBit(), side=Side.LEFT)]) + + def on_classical_vals(self, c: int) -> Dict[str, 'ClassicalValT']: + return {} + + def my_tensors( + self, incoming: Dict[str, 'ConnectionT'], outgoing: Dict[str, 'ConnectionT'] + ) -> List['DiscardInd']: + + from qualtran.simulation.tensor import DiscardInd + + return [DiscardInd((incoming['c'], 0))] + + +@frozen +class DiscardQ(Bloq): + """Discard a qubit. + + This is a dangerous operation that can ruin your computation. This is equivalent to + measuring the qubit and throwing out the measurement operation, so it removes any coherences + involved with the qubit. Use with care. + """ + + @cached_property + def signature(self) -> 'Signature': + return Signature([Register('q', QBit(), side=Side.LEFT)]) + + def my_tensors( + self, incoming: Dict[str, 'ConnectionT'], outgoing: Dict[str, 'ConnectionT'] + ) -> List['DiscardInd']: + + from qualtran.simulation.tensor import DiscardInd + + return [DiscardInd((incoming['q'], 0))] diff --git a/qualtran/bloqs/basic_gates/discard_test.py b/qualtran/bloqs/basic_gates/discard_test.py new file mode 100644 index 0000000000..da90792be4 --- /dev/null +++ b/qualtran/bloqs/basic_gates/discard_test.py @@ -0,0 +1,82 @@ +# Copyright 2025 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 BloqBuilder +from qualtran.bloqs.basic_gates import ( + CNOT, + Discard, + DiscardQ, + MeasZ, + PlusState, + ZeroEffect, + ZeroState, +) + + +def test_discard(): + bb = BloqBuilder() + q = bb.add(ZeroState()) + c = bb.add(MeasZ(), q=q) + bb.add(Discard(), c=c) + cbloq = bb.finalize() + + # We're allowed to discard classical bits in the classical simulator + ret = cbloq.call_classically() + assert ret == () + + k = cbloq.tensor_contract(superoperator=True) + np.testing.assert_allclose(k, 1.0, atol=1e-8) + + +def test_discard_vs_project(): + # Using the ZeroState effect un-physically projects us, giving trace of 0.5 + bb = BloqBuilder() + q = bb.add(PlusState()) + bb.add(ZeroEffect(), q=q) + cbloq = bb.finalize() + k = cbloq.tensor_contract(superoperator=True) + np.testing.assert_allclose(k, 0.5, atol=1e-8) + + # Measure and discard is trace preserving + bb = BloqBuilder() + q = bb.add(PlusState()) + c = bb.add(MeasZ(), q=q) + bb.add(Discard(), c=c) + cbloq = bb.finalize() + k = cbloq.tensor_contract(superoperator=True) + np.testing.assert_allclose(k, 1.0, atol=1e-8) + + +def test_discardq(): + # Completely dephasing map + # https://learning.quantum.ibm.com/course/general-formulation-of-quantum-information/quantum-channels#the-completely-dephasing-channel + bb = BloqBuilder() + q = bb.add_register('q', 1) + env = bb.add(ZeroState()) + q, env = bb.add(CNOT(), ctrl=q, target=env) + bb.add(DiscardQ(), q=env) + cbloq = bb.finalize(q=q) + ss = cbloq.tensor_contract(superoperator=True) + + should_be = np.zeros((2, 2, 2, 2)) + should_be[0, 0, 0, 0] = 1 + should_be[1, 1, 1, 1] = 1 + + np.testing.assert_allclose(ss, should_be, atol=1e-8) + + # Classical simulator will not let you throw out qubits + with pytest.raises(NotImplementedError, match=r'.*classical simulation.*'): + _ = cbloq.call_classically(q=1) diff --git a/qualtran/bloqs/basic_gates/z_basis.py b/qualtran/bloqs/basic_gates/z_basis.py index 7d10054885..b5e5a73844 100644 --- a/qualtran/bloqs/basic_gates/z_basis.py +++ b/qualtran/bloqs/basic_gates/z_basis.py @@ -13,7 +13,18 @@ # limitations under the License. from functools import cached_property -from typing import cast, Dict, Iterable, List, Optional, Sequence, Tuple, TYPE_CHECKING, Union +from typing import ( + cast, + Dict, + Iterable, + List, + Mapping, + Optional, + Sequence, + Tuple, + TYPE_CHECKING, + Union, +) import attrs import numpy as np @@ -27,6 +38,7 @@ bloq_example, BloqBuilder, BloqDocSpec, + CBit, CompositeBloq, ConnectionT, CtrlSpec, @@ -52,7 +64,7 @@ from qualtran.cirq_interop import CirqQuregT from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator - from qualtran.simulation.classical_sim import ClassicalValT + from qualtran.simulation.classical_sim import ClassicalValRetT, ClassicalValT _ZERO = np.array([1, 0], dtype=np.complex128) _ONE = np.array([0, 1], dtype=np.complex128) @@ -379,6 +391,44 @@ def _cz() -> CZ: _CZ_DOC = BloqDocSpec(bloq_cls=CZ, examples=[_cz], call_graph_example=None) +@frozen +class MeasZ(Bloq): + """Measure a qubit in the Z basis. + + Registers: + q [LEFT]: The qubit to measure. + c [RIGHT]: The classical measurement result. + """ + + @cached_property + def signature(self) -> 'Signature': + return Signature( + [Register('q', QBit(), side=Side.LEFT), Register('c', CBit(), side=Side.RIGHT)] + ) + + def on_classical_vals(self, q: int) -> Mapping[str, 'ClassicalValRetT']: + return {'c': q} + + def my_tensors( + self, incoming: Dict[str, 'ConnectionT'], outgoing: Dict[str, 'ConnectionT'] + ) -> List['qtn.Tensor']: + import quimb.tensor as qtn + + from qualtran.simulation.tensor import DiscardInd + + copy = np.zeros((2, 2, 2), dtype=np.complex128) + copy[0, 0, 0] = 1 + copy[1, 1, 1] = 1 + # Tie together q, c, and meas_result with the copy tensor; throw out one of the legs. + meas_result = qtn.rand_uuid('meas_result') + t = qtn.Tensor( + data=copy, + inds=[(incoming['q'], 0), (outgoing['c'], 0), (meas_result, 0)], + tags=[str(self)], + ) + return [t, DiscardInd((meas_result, 0))] + + @frozen class _IntVector(Bloq): """Represent a classical non-negative integer vector (state or effect). diff --git a/qualtran/bloqs/basic_gates/z_basis_test.py b/qualtran/bloqs/basic_gates/z_basis_test.py index 65528e8631..541a3cbb4d 100644 --- a/qualtran/bloqs/basic_gates/z_basis_test.py +++ b/qualtran/bloqs/basic_gates/z_basis_test.py @@ -17,13 +17,16 @@ import pytest import qualtran.testing as qlt_testing -from qualtran import Bloq, BloqBuilder +from qualtran import Bloq, BloqBuilder, QUInt from qualtran.bloqs.basic_gates import ( CZ, IntEffect, IntState, + MeasZ, + MinusState, OneEffect, OneState, + PlusState, XGate, ZeroEffect, ZeroState, @@ -276,3 +279,55 @@ def test_cz_phased_classical(): assert final_vals['q1'] == 1 assert final_vals['q2'] == 1 assert phase == -1 + + +def test_meas_z_supertensor(): + with pytest.raises(ValueError, match=r'.*superoperator.*'): + MeasZ().tensor_contract() + + # Zero -> Zero + bb = BloqBuilder() + q = bb.add(ZeroState()) + c = bb.add(MeasZ(), q=q) + cbloq = bb.finalize(c=c) + rho = cbloq.tensor_contract(superoperator=True) + should_be = np.outer([1, 0], [1, 0]) + np.testing.assert_allclose(rho, should_be, atol=1e-8) + + # One -> One + bb = BloqBuilder() + q = bb.add(OneState()) + c = bb.add(MeasZ(), q=q) + cbloq = bb.finalize(c=c) + rho = cbloq.tensor_contract(superoperator=True) + should_be = np.outer([0, 1], [0, 1]) + np.testing.assert_allclose(rho, should_be, atol=1e-8) + + # Plus -> mixture + bb = BloqBuilder() + q = bb.add(PlusState()) + c = bb.add(MeasZ(), q=q) + cbloq = bb.finalize(c=c) + rho = cbloq.tensor_contract(superoperator=True) + should_be = np.diag([0.5, 0.5]) + np.testing.assert_allclose(rho, should_be, atol=1e-8) + + # Minus -> mixture + bb = BloqBuilder() + q = bb.add(MinusState()) + c = bb.add(MeasZ(), q=q) + cbloq = bb.finalize(c=c) + rho = cbloq.tensor_contract(superoperator=True) + should_be = np.diag([0.5, 0.5]) + np.testing.assert_allclose(rho, should_be, atol=1e-8) + + +def test_meas_z_classical(): + bb = BloqBuilder() + q = bb.add(IntState(val=52, bitsize=8)) + qs = bb.split(q) + for i in range(8): + qs[i] = bb.add(MeasZ(), q=qs[i]) + cbloq = bb.finalize(outs=qs) + (ret,) = cbloq.call_classically() + assert list(ret) == QUInt(8).to_bits(52) # type: ignore[arg-type] diff --git a/qualtran/cirq_interop/_cirq_to_bloq.py b/qualtran/cirq_interop/_cirq_to_bloq.py index 7e009ac1a5..d2977e4339 100644 --- a/qualtran/cirq_interop/_cirq_to_bloq.py +++ b/qualtran/cirq_interop/_cirq_to_bloq.py @@ -236,7 +236,7 @@ def _my_tensors_from_gate( unitary = tensor_data_from_unitary_and_signature(cirq.unitary(gate), signature) inds = _order_incoming_outgoing_indices(signature, incoming=incoming, outgoing=outgoing) unitary = unitary.reshape((2,) * len(inds)) - return [qtn.Tensor(data=unitary, inds=inds, tags=[str(gate)])] + return [qtn.Tensor(data=unitary, inds=inds, tags=[gate.__class__.__name__])] @frozen(eq=False) diff --git a/qualtran/simulation/supertensor.ipynb b/qualtran/simulation/supertensor.ipynb new file mode 100644 index 0000000000..7ab7e347c8 --- /dev/null +++ b/qualtran/simulation/supertensor.ipynb @@ -0,0 +1,413 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "c90db82e-fe68-40ab-8093-0ee1b76f1cba", + "metadata": {}, + "source": [ + "# Super-Tensor Simulation\n", + "\n", + "Simulation of a quantum program involves encoding the state of the system and encoding operators that update the state according to the desired operation.\n", + "\n", + "In *closed* quantum systems, we can encode our state in a *state vector* and our operations as *unitary matrices*. For a given bloq, we can query these objects with the `Bloq.tensor_contract()` method. Bloqs with only right registers correspond to states; and bloqs with thru registers correspond to unitary matrices. The tensor simulation protocol handles composing the states and unitaries to get e.g. the final state or unitary.\n", + "\n", + "In *open* quantum systems, the state of our system is (potentially) a classically-probabalistic mixture over pure states (which are the states found in closed quantum systems). The operations map this mixture of states to a new mixture of states. The superoperator tensor protocol lets you simulate Qualtran programs that include non-unitary operations like measurement or discarding qubits. Any pure state can also be simulated using this protocol, but it is more expensive than the normal tensor simulation protocol; so we encourage you to use the ordinary tensor contraction protocol wherever feasible.\n" + ] + }, + { + "cell_type": "markdown", + "id": "c038ddeb-185a-45f9-8304-de9e96f1b4dc", + "metadata": {}, + "source": [ + "## Density matrix\n", + "\n", + "In this section, we contract all bloqs to thier `numpy.ndarray` numerical representation. There are more indices and fewer conventions about the arrangement of these indices, so practitioners must either pay close attention to the following documentation, or deal exclusively and directly with the Quimb `qtn.TensorNetwork` objects and their named indices (see the next section).\n", + "\n", + "When dealing with open system simulation, the state of the system is no longer represented by a 1-dimensional vector of probability amplitudes, but rather **a 2-dimensional matrix** of classical probabilities along the diagonal and quantum *coherences* off-diagonal.\n", + "\n", + "\n", + "### The $|+\\rangle$ state\n", + "First, let's inspect the statevector and density matrix representation of the $|+\\rangle$ state." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "12c866ec-abf8-4a4f-941e-33922edf8348", + "metadata": {}, + "outputs": [], + "source": [ + "# Ordinary statevector of |+>\n", + "from qualtran.bloqs.basic_gates import PlusState\n", + "PlusState().tensor_contract()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "794c0d72-ee3a-4679-83bd-6a6b3a504dd1", + "metadata": {}, + "outputs": [], + "source": [ + "# Use the superoperator simulation machinery. States are\n", + "# now density _matrices_. \n", + "PlusState().tensor_contract(superoperator=True)" + ] + }, + { + "cell_type": "markdown", + "id": "d25c6f88-baf2-4a9f-bdfb-0f33b25ebbda", + "metadata": {}, + "source": [ + "### Incoherent states\n", + "\n", + "We know that the $|+\\rangle$ state has a 50% chance of being measured in the 0 vs. 1 state. What is the difference between this and a classical coin flip? We'll compare the density matrices between the coherent `PlusState` vs. measuring the result." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5cab8731-7cd0-4486-93d2-69ab79853786", + "metadata": {}, + "outputs": [], + "source": [ + "from qualtran import BloqBuilder\n", + "from qualtran.drawing import show_bloq\n", + "from qualtran.bloqs.basic_gates import ZeroState, Hadamard, MeasZ\n", + "\n", + "# Initialize a qubit, do a Hadamard\n", + "bb = BloqBuilder()\n", + "q = bb.add(ZeroState())\n", + "q = bb.add(Hadamard(), q=q)\n", + "c = bb.add(MeasZ(), q=q)\n", + "coin_flip = bb.finalize(c=c)\n", + "\n", + "# The circuit\n", + "show_bloq(coin_flip, 'musical_score')\n", + "\n", + "# Its density matrix\n", + "rho_coin_flip = coin_flip.tensor_contract(superoperator=True)\n", + "display(rho_coin_flip)" + ] + }, + { + "cell_type": "markdown", + "id": "3fb7e165-2be8-4c37-b7e5-08159e9c96aa", + "metadata": {}, + "source": [ + "## Superoperator evolution with matrix-vector multiplication\n", + "\n", + "Under the hood, the Quimb tensor network package can find efficient contraction orderings for performing superoperator simulation. However, for clarity we can show how the superoperator tensors can be manipulated to evolve a vector representation of the density operator with a matrix representation of the superoperator." + ] + }, + { + "cell_type": "markdown", + "id": "2dc0cc33-3150-405c-ac13-d317ed72c570", + "metadata": {}, + "source": [ + "### Superoperator tensors\n", + "\n", + "Operations (like `Hadamard` below) are encoded in 4-index tensors. You saw above that states are encoded in 2-index tensors (the density matrix)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "fa757980-f244-4c30-828b-82cae9253248", + "metadata": {}, + "outputs": [], + "source": [ + "super_h = Hadamard().tensor_contract(superoperator=True)\n", + "super_h.shape" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6fc177a0-36b2-4403-95e3-fc644d505b0b", + "metadata": {}, + "outputs": [], + "source": [ + "rho_coin_flip.shape" + ] + }, + { + "cell_type": "markdown", + "id": "89bec971-0894-4e48-a016-3ca6536e33cb", + "metadata": {}, + "source": [ + "There is no standard naming scheme for these four indices, but we describe them here:\n", + "\n", + " - Unitary operations in standard, statevector evolution have two indices (i.e. they are matrices). We name the two indices **left** and **right** indices, corresponding to the input and output (resp.) basis states.\n", + " - The density matrix $\\rho$ for a pure state $|\\psi\\rangle$ is given by $|\\psi\\rangle \\langle \\psi|$.\n", + " - Evolution of a pure state by a unitary $U$ corresponds to applying $U|\\psi\\rangle$ to the first part of $\\rho$ and $\\langle \\psi | U^\\dagger$ to the latter part.\n", + " - With a bit of poetic license, we call the indices of the $|\\psi\\rangle$ part the **forward** indices and the $\\langle \\psi|$ part the **backward** indices.\n", + "\n", + "The Qualtran ordering of the superoperator tensor is:\n", + "\n", + " (right_forward, right_backward, left_forward, left_backward)\n", + "\n", + "The ordering of the density matrix indices is the familiar, textbook ordering which—following the terminoligy we set up—is either\n", + "\n", + " (right_forward, right_backward)\n", + " *or*\n", + " (left_foward, left_backward)\n", + "\n", + "depending on whether it is an initializing state or de-allocating state, respectively." + ] + }, + { + "cell_type": "markdown", + "id": "51c1dddd-99cb-4205-83ce-326be9c3d5bf", + "metadata": {}, + "source": [ + "### Reshaping\n", + "\n", + "The index ordering allows reshaping of superoperators into matrices and density matrices into vectors so evolution can be computed by the traditional matrix-vector product.\n", + "\n", + "We'll see that applying a reshaped tensor of `Hadamard` to our reshaped `rho_coin_flip` gives us a random result, but applying it to a coherent state results in a deterministic output." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c8c33878-f71a-4ea3-b116-a1682c3a45e8", + "metadata": {}, + "outputs": [], + "source": [ + "# Applying H to our incoherent, coin-flip state\n", + "(super_h.reshape(4,4) @ rho_coin_flip.reshape(4)).reshape(2,2)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0ca8abb6-7156-42e4-905b-c17f2d71bf64", + "metadata": {}, + "outputs": [], + "source": [ + "# Applying H to our coherent state\n", + "rho_coherent = PlusState().tensor_contract(superoperator=True)\n", + "(super_h.reshape(4,4) @ rho_coherent.reshape(4)).reshape(2,2)" + ] + }, + { + "cell_type": "markdown", + "id": "fd627fc9-d734-4ed0-901c-d09e22480e37", + "metadata": {}, + "source": [ + "## Quimb Tensor Network\n", + "\n", + "The function `cbloq_to_superquimb` returns a `qtn.TensorNetwork` representing the composite bloq. The structure is apparent: there are effectively two pure-state evolutions occuring ('forward' and 'backward'). Non-unitary operations introduce an index coupling the forward and backward evolutions." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "84d7be2f-44ac-4c69-aaa0-cc6f15364530", + "metadata": {}, + "outputs": [], + "source": [ + "from qualtran.simulation.tensor import cbloq_to_superquimb\n", + "tn = cbloq_to_superquimb(coin_flip, friendly_indices=True)\n", + "tn.draw(color=['|0>', 'H', 'MeasZ', 'dag'])" + ] + }, + { + "cell_type": "markdown", + "id": "6375ba0c-1b8c-4de9-9153-bc848f626730", + "metadata": {}, + "source": [ + "## System+Environment modeling\n", + "\n", + "All CPTP maps can be implemented by unitary evolution in a larger \"system + environment\" space. In this section, we show how to build a measurement operation with only standard, unitary bloqs and the ability to discard information.\n", + "\n", + "Any Hermitian operator can be \"measured\" into a fresh ancilla using a simple prescription, see e.g. Nielsen and Chuang Exercise 4.44. We build that construction below." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8cde1066-7e36-49b6-a665-58905250521c", + "metadata": {}, + "outputs": [], + "source": [ + "from qualtran import BloqBuilder, Register, Side, CtrlSpec, QBit, CBit\n", + "from qualtran.bloqs.basic_gates import ZeroState, Hadamard, MeasZ, ZGate\n", + "from qualtran.bloqs.bookkeeping import Cast\n", + "\n", + "bb = BloqBuilder()\n", + "# Take a single qubit as input\n", + "q = bb.add_register(Register(\"q\", QBit()))\n", + "# Set up our output register: we'll return one classical bit\n", + "bb.add_register(Register(\"c\", CBit(), side=Side.RIGHT))\n", + "\n", + "# This construction works for any Hermitian operator. We'll\n", + "# use Z as a familiar first example.\n", + "op = ZGate()\n", + "\n", + "# Allocate space to record the result of our measurement operation\n", + "meas_space = bb.add(ZeroState())\n", + "meas_space = bb.add(Hadamard(), q=meas_space)\n", + "\n", + "# Do Controlled(op)\n", + "_, add_ctrled = op.get_ctrl_system(CtrlSpec())\n", + "(meas_space,), (q,) = add_ctrled(bb, ctrl_soqs=[meas_space], in_soqs={'q': q})\n", + "\n", + "# Final Hadamard, and cast our measurement register\n", + "# into a classical bit.\n", + "meas_space = bb.add(Hadamard(), q=meas_space)\n", + "meas_result = bb.add(Cast(QBit(), CBit(), allow_quantum_to_classical=True), reg=meas_space)\n", + "meas_cbloq = bb.finalize(c = meas_result, q=q)\n", + "show_bloq(meas_cbloq)" + ] + }, + { + "cell_type": "markdown", + "id": "455ab7ed-c93d-4fc3-8f81-a694d4f90a24", + "metadata": {}, + "source": [ + "Note that we've entangled our system \"q\" with a fresh register. We've used a `Cast` operation to denote that the new bit is a classical bit. Practically this means we can no longer perform quantum operations like `Hadamard` to it, and any classical processing can happen on ordinary CPUs. But at a quantum-information level, there is nothing about the tensor structure to show that \"c\" is a 'classical' index. Below, we draw the tensor network encoding of the State/Unitary composite bloq using the ordinary tensor protocol" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "57f300ec-28a5-4b6f-aabd-4cc5c8cf31ed", + "metadata": {}, + "outputs": [], + "source": [ + "from qualtran.simulation.tensor import cbloq_to_quimb\n", + "tn = cbloq_to_quimb(meas_cbloq, friendly_indices=True)\n", + "tn.draw(color=['CZ', 'H'])\n", + "display(tn.contract())" + ] + }, + { + "cell_type": "markdown", + "id": "6ec443b0-d7e9-4d69-89eb-a7c068ed2737", + "metadata": {}, + "source": [ + "### Making a measurement operation\n", + "\n", + "In an open system—like the world we live in—we don't have coherent access to each (qu)bit worth of information. Our measurement apparatus might have $10^{23}$ particles, each recording the result of a measurement. We can simulate the standard measurement channel where information is lost to the environment by using the previous circuit and discarding the coherent qubit wire. Now, the signature of our composite bloq takes in one `QBit()` and returns one `CBit()`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "690955e1-eab3-4f90-863b-53090d13a75c", + "metadata": {}, + "outputs": [], + "source": [ + "from qualtran.bloqs.basic_gates import DiscardQ\n", + "\n", + "bb = BloqBuilder()\n", + "q = bb.add_register(Register(\"q\", QBit(), side=Side.LEFT))\n", + "q, c = bb.add_from(meas_cbloq, q=q)\n", + "bb.add(DiscardQ(), q=q)\n", + "meas2_cbloq = bb.finalize(c=c)\n", + "show_bloq(meas2_cbloq)" + ] + }, + { + "cell_type": "markdown", + "id": "2842fec6-d709-488a-9d31-d8817f47dbb0", + "metadata": {}, + "source": [ + "The ordinary tensor simulation protocol is insufficient to handle discarding a qubit. " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b48228f0-3ef1-4e07-b89a-65c46fd0a865", + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " tn = cbloq_to_quimb(meas2_cbloq, friendly_indices=True)\n", + "except ValueError as e:\n", + " print(\"ValueError:\", e)" + ] + }, + { + "cell_type": "markdown", + "id": "cbec03e4-1a88-4a2b-ae4c-32f1efb6ddc7", + "metadata": {}, + "source": [ + "To remove a qubit, we need to sum over its possible states as we would when computing a marginal probability distribution. But our probability *amplitudes* only sum to 1 when we consider their absolute value squared\n", + "$$\n", + "\\sum a^* a = 1,\n", + "$$\n", + "so our integration proceedure requires indexing into both our state $|\\psi\\rangle$ and its adjoint $\\langle \\psi|$ to remove the offending bit.\n", + "\n", + "If you're using a densitry matrix, this correspond to performing a parital trace. In Qualtran, the superoperator tensor simulation protocol sets up two simultaneous tensor networks for simulating unitary action on the circuit *and* its adjoint. That is, we simulate both $|\\psi\\rangle$ and $\\langle \\psi|$. Discarding a qubit is performed by contracting the qubit's index in the forward network with its corresponding index in the backwards, adjoint network." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ca011792-af4d-463e-bd26-df7c0cdba14a", + "metadata": {}, + "outputs": [], + "source": [ + "tn = cbloq_to_superquimb(meas2_cbloq, friendly_indices=True)\n", + "tn.draw(color=['CZ', 'H', 'dag'], initial_layout='spectral')\n", + "display(tn.shape)" + ] + }, + { + "cell_type": "markdown", + "id": "5dbadc5c-7422-4ea6-8a5d-379a58bc81b8", + "metadata": {}, + "source": [ + "The discard operation couples the two pure evolutions. The superoperator is now a rank-4 tensor.\n", + "\n", + "By the Stinespring dilation theorem, we can actually represent any superoperator (aka quantum channel, aka CPTP map) with only pure state evolution and the ability to discard qubits. This gives rise to the \"System-Environment\" representation of superoperators, which is the native representation in Qualtran. It is quite natural for the open-system operations we're most concerned with (like measurement); but practitioners may have to do some careful translation to encode superoperators traditionally expressed in another representation like the operator-sum (Kraus operator) representation." + ] + }, + { + "cell_type": "markdown", + "id": "91784324-c495-4878-afed-5597ea6782d6", + "metadata": {}, + "source": [ + "## Implementation\n", + "\n", + "The super-tensor protocol is a superset of the ordinary tensor simulation protocol. Bloqs override the same method to define their tensors for open-system simulation: the bloq author overrides `Bloq.my_tensors` and is responsible for returning a list of tensors whose indices follow the same prescribed format as documented in the [tensor protocol](.). \n", + "\n", + "All additional functionality is unlocked by the ability to return a `qualtran.simulation.tensor.DiscardInd` object amongst the ordinary `qtn.Tensor` objects. Per above, any superoperator can be expressed in this system-environment representation. This simple sentinel object flags the named index as subject to \"tracing out\" during construction of the complete network.\n", + "\n", + "When calling `cbloq_to_quimb` the indices are faithfully kept as `(cxn: Connection, j: int)` tuples. During the conversion to the superoperator tensor network, each tensor returned by `Bloq.my_tensors` is added twice:\n", + " - The 'forward' tensor is added. Its indices `(cxn, j)` are transformed to `(cxn, j, True)`\n", + " - The 'backward' tensor is added. Its indices `(cxn, j)` are transformed to `(cxn, j, False)` and we take the element-wise complex conjugate of its data ndarray.\n", + " - A `DiscardInd` will remove the booleans from the named index in tensors which have already been added. This causes the forward and backward indices to be contracted together.\n", + "\n", + "Note that index permutation operations like taking the transpose of the backwards tensor is handled by the structure of the tensor network rather than mutating the data ndarray. \n", + "\n", + "The `my_tensors` override must order its return values such that a `DiscardInd` is encountered *after* the tensor that defines the index to discard. " + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.8" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/qualtran/simulation/tensor.ipynb b/qualtran/simulation/tensor.ipynb index 014193182d..d91373e7d8 100644 --- a/qualtran/simulation/tensor.ipynb +++ b/qualtran/simulation/tensor.ipynb @@ -5,7 +5,7 @@ "id": "8cad0714", "metadata": {}, "source": [ - "# Tensor\n", + "# Tensor Simulation\n", "\n", "The tensor protocol lets you query the tensor (vector, matrix, etc.) representation of a bloq. For example, we can easily inspect the familiar unitary matrix representing the controlled-not operation:" ] @@ -28,7 +28,7 @@ "id": "5a2f84d3", "metadata": {}, "source": [ - "Bloqs can represent states, effects, and non-unitary operations. Below, we see the vector representation of the plus state and zero effect." + "Bloqs can represent states, effects, unitary operations, and compositions of these operations. Below, we see the vector representation of the plus state and zero effect." ] }, { @@ -64,6 +64,14 @@ "And().tensor_contract().shape" ] }, + { + "cell_type": "markdown", + "id": "6f61b73c-50df-4ed3-aeab-f24e9ec2b3e5", + "metadata": {}, + "source": [ + "For a bloq with exclusively thru-registers, the returned tensor will be a matrix with shape `(n, n)` where `n` is the number of bits in the signature. For a bloq with exlusively right- or left-registers, the returned tensor will be a vector with shape `(n,)`. In general, the tensor will be an ndarray of shape `(n_right_bits, n_left_bits)`; but empty dimensions are *dropped*." + ] + }, { "cell_type": "markdown", "id": "912fb914", @@ -140,8 +148,8 @@ "source": [ "from qualtran.simulation.tensor import cbloq_to_quimb\n", "\n", - "tn = cbloq_to_quimb(cbloq)\n", - "tn.draw(show_inds=False)" + "tn = cbloq_to_quimb(cbloq, friendly_indices=True)\n", + "tn.draw()" ] }, { @@ -169,36 +177,29 @@ "source": [ "### Quimb index format\n", "\n", - "In `CompositeBloq`, we form the compute graph by storing a list of nodes and edges. Quimb uses a different strategy for representing the tensor network graphs. To form a tensor network in Quimb, we provide a list of `qtn.Tensor` which contain not only the tensor data but also a list of \"indices\" that can form connections to other tensors. Similar to the Einstein summation convention, if two tensors each have an index with the same name: an edge is created in the tensor network graph and this shared index is summed over. These indices are traditionally short strings like `\"k\"`, but can be any hashable Python object. In `CompositeBloq`, the unique object that identifies a connection between bloqs is `qualtran.Connection`, so we use these connection objects as our indices.\n", + "Above, we used the `friendly_indices=True` argument to give string names to the outer indices of the `qtn.TensorNetwork`. This can be useful for interactive manipulation of the tensor network by a human. Qualtran uses a highly-structured (but less human-readible) format for internal tensor indices, which we describe here.\n", + "\n", + "In `CompositeBloq`, we form the compute graph by storing a list of nodes and edges. Quimb uses a different strategy for representing the tensor network graphs. To form a tensor network in Quimb, we provide a list of `qtn.Tensor` which contain not only the tensor data but also a list of \"indices\" that can form connections to other tensors. Similar to the Einstein summation convention, if two tensors each have an index with the same name: an edge is created in the tensor network graph and this shared index is summed over. In the Quimb documentation (for example), these indices are traditionally short strings like `\"k\"`, but can be any hashable Python object. In `CompositeBloq`, the unique object that identifies a connection between bloqs is `qualtran.Connection`, so we use these connection objects as the first part of our indices.\n", "\n", - "Qualtran and Quimb both support \"wires\" with arbitrary bitsize. Qualtran uses bookkeeping bloqs like `Join` and `Split` to fuse and un-fuse indices. In theory, these operations should be free in the tensor contraction, as they are essentially an identity tensor. In our understanding, Quimb does not have special support for supporting these re-shaping operations within the tensor network graph. In versions of Qualtran prior to v0.5, split and join operations were tensored up to `n`-sized identity tensors. This would form a bottleneck in any contraction ordering. Therefore, we keep all the indices un-fused in the tensor network representation and use tuples of `(cxn, i)` for our tensor indices, where the second integer `i` indexes the individual bits in a register with `reg.dtype.num_qubits` > 1.\n", + "Qualtran and Quimb both support \"wires\" with arbitrary bitsize. Qualtran uses bookkeeping bloqs like `Join` and `Split` to fuse and un-fuse indices. In theory, these operations should be free in the tensor contraction, as they are essentially an identity tensor. In our understanding, Quimb does not have special support for handling these re-shaping operations within the tensor network graph. In versions of Qualtran prior to v0.5, split and join operations were tensored up to `n`-sized identity tensors. This would form a bottleneck in any contraction ordering. Therefore, we keep all the indices un-fused in the tensor network representation and use tuples of `(cxn, j)` for our tensor indices, where the second integer `j` indexes the individual bits in a register with `reg.dtype.num_bits` > 1.\n", "\n", "**In summary:**\n", - " - Each tensor index is a tuple `(cxn, i)`\n", + " - Each tensor index is a tuple `(cxn, j)`\n", " - The `cxn: qualtran.Connection` entry identifies the connection between soquets in a Qualtran compute graph.\n", - " - The second integer `i` is the bit index within high-bitsize registers, which is necessary due to technical restrictions." + " - The second integer `j` is the bit index within high-bitsize registers, which is necessary due to technical restrictions." ] }, { "cell_type": "code", "execution_count": null, - "id": "d9c0ed47-275e-417b-bbf5-b8dff2f84099", + "id": "e0831ffe-2845-4878-a80e-bbe4a533daf4", "metadata": {}, "outputs": [], "source": [ - "# Use `get_right_and_left_inds` to get the quimb indices ordered according to\n", - "# the bloq's signature.\n", - "\n", - "from qualtran.simulation.tensor import get_right_and_left_inds\n", - "left, right = get_right_and_left_inds(tn, bloq.signature)\n", - "\n", - "print(\"Left outer inds:\")\n", - "for cxn, i in left:\n", - " print(' ', cxn.left)\n", - "\n", - "print(\"Right outer inds\")\n", - "for cxn, i in right:\n", - " print(' ', cxn.right)" + "example_inner_index = tn.inner_inds()[0]\n", + "cxn, j = example_inner_index\n", + "print(\"cxn:\", cxn)\n", + "print(\"j: \", j)" ] }, { @@ -212,7 +213,20 @@ "\n", "In Qualtran, we usually avoid flattening bloqs and strongly to prefer to work with one level of decomposition at a time. This is to avoid performance issues with large, abstract algorithms. But typically if the full circuit is large enough to cause performance issues with flattening it is also too large to simulate numerically; so an exception to the general advice is made here.\n", "\n", - "All bloqs in the flattened circuit must provide their explicit tensors. If your bloq's tensors ought to be derived from its decomposition: this is achieved by the previously mentioned flattening operation. If a bloq provides tensors through overriding `Bloq.my_tensors` _and also_ defines a decomposition, the explicit tensors will not be used (by default). This is because any bloq with a decomposition will be flattened. If you would like to control the flattening operation, use the free functions to control the tensor network construction and contraction." + "All bloqs in the flattened circuit must provide their explicit tensors. If your bloq's tensors ought to be derived from its decomposition: this is achieved by the previously mentioned flattening operation. If a bloq provides tensors through overriding `Bloq.my_tensors` _and also_ defines a decomposition, the explicit tensors will not be used (by default). If you'd like to always use annotated tensors, set `bloq_to_dense(..., full_flatten=False)`. If you would like full control over flattening, use the free functions to control the tensor network construction and contraction." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "def8e03b-6608-4cd3-8035-6755cf93d9cc", + "metadata": {}, + "outputs": [], + "source": [ + "# Set `full_flatten=False` to use annotated tensors even\n", + "# if there's a decomposition. (No change in this example since\n", + "# most bloqs don't define both a decomposition and tensors).\n", + "_ = bloq_to_dense(bloq, full_flatten=False)" ] }, { @@ -222,8 +236,7 @@ "metadata": {}, "outputs": [], "source": [ - "# Stop flattening if e.g. a bloq supports explicit tensors _and_ a decomposition.\n", - "# Use the flatten predicate to control.\n", + "# Manually flatten and contract for complete control\n", "custom_flat = bloq.as_composite_bloq().flatten(lambda binst: binst.i != 2)\n", "tn = cbloq_to_quimb(custom_flat)\n", "len(tn.tensors)" @@ -285,7 +298,6 @@ " # We'll reshape our matrix into the more natural n-dimensional\n", " # tensor form.\n", " tensor = matrix.reshape((2,2,2,2))\n", - " \n", "\n", " # This is a simple case: we only need one tensor and\n", " # each register is one bit.\n", @@ -354,7 +366,7 @@ " q1 = bb.add(ZeroState())\n", "\n", " q0, q1 = bb.add(CNOT(), ctrl=q0, target=q1)\n", - " return {'q0': q0, 'q1': q1}\n" + " return {'q0': q0, 'q1': q1}" ] }, { @@ -380,7 +392,7 @@ "id": "ebbe9a2a", "metadata": {}, "source": [ - "If you don't flatten first, an error will be raised" + "If you try to directly use `cbloq_to_quimb` and bypass the flattening operation, an error will be raised." ] }, { @@ -422,18 +434,10 @@ "outputs": [], "source": [ "from qualtran.bloqs.basic_gates import CNOT\n", - "from qualtran.simulation.tensor import (\n", - " cbloq_to_quimb, get_right_and_left_inds\n", - ")\n", + "from qualtran.simulation.tensor import cbloq_to_quimb\n", "\n", "cbloq = CNOT().as_composite_bloq()\n", - "tn = cbloq_to_quimb(cbloq)\n", - "\n", - "# Rename the indices to something less verbose\n", - "lefts, rights = get_right_and_left_inds(tn, cbloq.signature)\n", - "rename = {left: f'{left[0].left.reg.name}_in' for left in lefts}\n", - "rename |= {right: f'{right[0].right.reg.name}_out' for right in rights}\n", - "tn = tn.reindex(rename)\n", + "tn = cbloq_to_quimb(cbloq, friendly_indices=True)\n", "\n", "tn.draw(color=['COPY', 'XOR'], show_tags=True, initial_layout='spectral')\n", "for tensor in tn:\n", @@ -441,6 +445,63 @@ " print(tensor.data)\n", " print()" ] + }, + { + "cell_type": "markdown", + "id": "d4ef3cdf-6edc-4f1a-900c-f0d73cd2cb26", + "metadata": {}, + "source": [ + "### Final state vector from a circuit\n", + "\n", + "In Qualtran, all initial states must be explicitly specified with allocation-like bloqs. For example, if we define the circuit below, the `.tensor_contract` simulation method will only ever return a unitary matrix. If you'd like the state vector that results from applying that circuit to qubits initialized in a particular way (e.g. the all-zeros computational basis state), then you must specify that" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "28effa80-e699-4478-9779-d7d5fc93eb0e", + "metadata": {}, + "outputs": [], + "source": [ + "from qualtran import BloqBuilder\n", + "from qualtran.bloqs.basic_gates import Hadamard\n", + "\n", + "bb = BloqBuilder()\n", + "q1 = bb.add_register('q1', 1)\n", + "q2 = bb.add_register('q2', 1)\n", + "\n", + "q1 = bb.add(Hadamard(), q=q1)\n", + "q1, q2 = bb.add(CNOT(), ctrl=q1, target=q2)\n", + "bell_circuit = bb.finalize(q1=q1, q2=q2)\n", + "\n", + "# This circuit always corresponds to a unitary *matrix*\n", + "show_bloq(bell_circuit, 'musical_score')\n", + "print(bell_circuit.tensor_contract().round(2))" + ] + }, + { + "cell_type": "markdown", + "id": "549cb96d-162e-4112-9e62-88a4aac73a5d", + "metadata": {}, + "source": [ + "We can use the `initialize_from_zero` helper function to get the state vector corresponding to an all-zeros initial state." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8cf3fbba-eccf-4f15-b9ff-9b3bfbb7e65c", + "metadata": {}, + "outputs": [], + "source": [ + "from qualtran.simulation.tensor import initialize_from_zero\n", + "bell_state_cbloq = initialize_from_zero(bell_circuit)\n", + "\n", + "# The new composite bloq consists of a round of intialization and\n", + "# then the circuit from above. Its tensor contraction is the state vector\n", + "show_bloq(bell_state_cbloq)\n", + "print(bell_state_cbloq.tensor_contract().round(2))" + ] } ], "metadata": { diff --git a/qualtran/simulation/tensor/__init__.py b/qualtran/simulation/tensor/__init__.py index de6cf893e0..e656cbaaea 100644 --- a/qualtran/simulation/tensor/__init__.py +++ b/qualtran/simulation/tensor/__init__.py @@ -14,9 +14,9 @@ """Functionality for the `Bloq.tensor_contract()` protocol.""" -from ._dense import bloq_to_dense, get_right_and_left_inds, quimb_to_dense +from ._dense import bloq_to_dense, quimb_to_dense from ._flattening import bloq_has_custom_tensors, flatten_for_tensor_contraction -from ._quimb import cbloq_to_quimb, initialize_from_zero +from ._quimb import cbloq_to_quimb, cbloq_to_superquimb, DiscardInd, initialize_from_zero from ._tensor_data_manipulation import ( active_space_for_ctrl_spec, eye_tensor_for_signature, diff --git a/qualtran/simulation/tensor/_dense.py b/qualtran/simulation/tensor/_dense.py index ea4a35bbed..f4e630ddb0 100644 --- a/qualtran/simulation/tensor/_dense.py +++ b/qualtran/simulation/tensor/_dense.py @@ -13,14 +13,15 @@ # limitations under the License. import logging -from typing import Any, Dict, List, Tuple, TYPE_CHECKING +from collections import defaultdict +from typing import Dict, List, Tuple, TYPE_CHECKING from numpy.typing import NDArray -from qualtran import Bloq, Connection, ConnectionT, LeftDangle, RightDangle, Signature +from qualtran import Bloq, Connection, ConnectionT, Signature from ._flattening import flatten_for_tensor_contraction -from ._quimb import cbloq_to_quimb +from ._quimb import _IndT, cbloq_to_quimb, cbloq_to_superquimb if TYPE_CHECKING: import quimb.tensor as qtn @@ -64,65 +65,59 @@ def _order_incoming_outgoing_indices( return inds -def get_right_and_left_inds(tn: 'qtn.TensorNetwork', signature: Signature) -> List[List[Any]]: - """Return right and left tensor indices. +def _group_outer_inds( + tn: 'qtn.TensorNetwork', signature: Signature, superoperator: bool = False +) -> List[List[_IndT]]: + """Group outer indices of a tensor network. - In general, this will be returned as a list of length-2 corresponding - to the right and left indices, respectively. If there *are no* right - or left indices, that entry will be omitted from the returned list. - - Right indices come first to match the quantum computing / matrix multiplication - convention where U_tot = U_n ... U_2 U_1. + This is used by 'bloq_to_dense` and `quimb_to_dense` to return a 1-, 2-, or 4-dimensional + array depending on the quantity and type of outer indices in the tensor network. See + the docstring for `Bloq.tensor_contract()` for more informaiton. Args: tn: The tensor network to fetch the outer indices, which won't necessarily be ordered. signature: The signature of the bloq used to order the indices. + superoperator: Whether `tn` is a pure-state or open-system tensor network. """ - left_inds = {} - right_inds = {} - - # Each index is a (cxn: Connection, j: int) tuple. - cxn: Connection + reg_name: str + idx: Tuple[int, ...] j: int + group: str + _KeyT = Tuple[str, Tuple[int, ...], int] + ind_groups_d: Dict[str, Dict[_KeyT, _IndT]] = defaultdict(dict) for ind in tn.outer_inds(): - cxn, j = ind - if cxn.left.binst is LeftDangle: - soq = cxn.left - left_inds[soq.reg, soq.idx, j] = ind - elif cxn.right.binst is RightDangle: - soq = cxn.right - right_inds[soq.reg, soq.idx, j] = ind - else: - raise ValueError( - "Outer indices of a tensor network should be " - "connections to LeftDangle or RightDangle" - ) - - left_ordered_inds = [] - for reg in signature.lefts(): - for idx in reg.all_idxs(): - for j in range(reg.dtype.num_bits): - left_ordered_inds.append(left_inds[reg, idx, j]) - - right_ordered_inds = [] - for reg in signature.rights(): - for idx in reg.all_idxs(): - for j in range(reg.dtype.num_bits): - right_ordered_inds.append(right_inds[reg, idx, j]) - - inds = [] - if right_ordered_inds: - inds.append(right_ordered_inds) - if left_ordered_inds: - inds.append(left_ordered_inds) + reg_name, idx, j, group = ind + ind_groups_d[group][reg_name, idx, j] = ind + + ind_groups_l: Dict[str, List[_IndT]] = defaultdict(list) + + def _sort_group(regs, group_name): + for reg in regs: + for idx in reg.all_idxs(): + for j in range(reg.dtype.num_bits): + ind_groups_l[group_name].append(ind_groups_d[group_name][reg.name, idx, j]) + + if superoperator: + _sort_group(signature.lefts(), 'lf') + _sort_group(signature.lefts(), 'lb') + _sort_group(signature.rights(), 'rf') + _sort_group(signature.rights(), 'rb') + group_names = ['rf', 'rb', 'lf', 'lb'] + else: + _sort_group(signature.lefts(), 'l') + _sort_group(signature.rights(), 'r') + group_names = ['r', 'l'] + inds = [ind_groups_l[groupname] for groupname in group_names if ind_groups_l[groupname]] return inds -def quimb_to_dense(tn: 'qtn.TensorNetwork', signature: Signature) -> NDArray: +def quimb_to_dense( + tn: 'qtn.TensorNetwork', signature: Signature, superoperator: bool = False +) -> NDArray: """Contract a quimb tensor network `tn` to a dense matrix consistent with `signature`.""" - inds = get_right_and_left_inds(tn, signature) + inds = _group_outer_inds(tn, signature, superoperator=superoperator) if tn.contraction_width() > 8: tn.full_simplify(inplace=True) @@ -134,18 +129,34 @@ def quimb_to_dense(tn: 'qtn.TensorNetwork', signature: Signature) -> NDArray: return data -def bloq_to_dense(bloq: Bloq, full_flatten: bool = True) -> NDArray: +def bloq_to_dense(bloq: Bloq, full_flatten: bool = True, superoperator: bool = False) -> NDArray: """Return a contracted, dense ndarray representing the composite bloq. This function is also available as the `Bloq.tensor_contract()` method. This function decomposes and flattens a given bloq into a factorized CompositeBloq, turns that composite bloq into a Quimb tensor network, and contracts it into a dense - matrix. - - The returned array will be 0-, 1- or 2- dimensional with indices arranged according to the - bloq's signature. In the case of a 2-dimensional matrix, we follow the - quantum computing / matrix multiplication convention of (right, left) order of dimensions. + ndarray. + + The returned array will be 0-, 1-, 2-, or 4-dimensional with indices arranged according to the + bloq's signature and the type of simulation requested via the `superoperator` flag. + + If `superoperator` is set to False (the default), a pure-state tensor network will be + constructed. + - If `bloq` has all thru-registers, the dense tensor will be 2-dimensional with shape `(n, n)` + where `n` is the number of bits in the signature. We follow the linear algebra convention + and order the indices as (right, left) so the matrix-vector product can be used to evolve + a state vector. + - If `bloq` has all left- or all right-registers, the tensor will be 1-dimensional with + shape `(n,)`. Note that we do not distinguish between 'row' and 'column' vectors in this + function. + - If `bloq` has no external registers, the contracted form is a 0-dimensional complex number. + + If `superoperator` is set to True, an open-system tensor network will be constructed. + - States result in a 2-dimensional density matrix with indices (right_forward, right_backward) + or (left_forward, left_backward) depending on whether they're input or output states. + - Operations result in a 4-dimensional tensor with indices (right_forward, right_backward, + left_forward, left_backward). For fine-grained control over the tensor contraction, use `cbloq_to_quimb` and `TensorNetwork.to_dense` directly. @@ -154,8 +165,14 @@ def bloq_to_dense(bloq: Bloq, full_flatten: bool = True) -> NDArray: bloq: The bloq full_flatten: Whether to completely flatten the bloq into the smallest possible bloqs. Otherwise, stop flattening if custom tensors are encountered. + superoperator: If toggled to True, do an open-system simulation. This supports + non-unitary operations like measurement, but is more costly and results in + higher-dimension resultant tensors. """ logging.info("bloq_to_dense() on %s", bloq) flat_cbloq = flatten_for_tensor_contraction(bloq, full_flatten=full_flatten) - tn = cbloq_to_quimb(flat_cbloq) - return quimb_to_dense(tn, bloq.signature) + if superoperator: + tn = cbloq_to_superquimb(flat_cbloq) + else: + tn = cbloq_to_quimb(flat_cbloq) + return quimb_to_dense(tn, bloq.signature, superoperator=superoperator) diff --git a/qualtran/simulation/tensor/_quimb.py b/qualtran/simulation/tensor/_quimb.py index d4957e4778..7a4583fdc4 100644 --- a/qualtran/simulation/tensor/_quimb.py +++ b/qualtran/simulation/tensor/_quimb.py @@ -12,8 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. import logging -from typing import cast, Dict, Iterable +from typing import Any, cast, Dict, Iterable, Tuple, TypeAlias, Union +import attrs import numpy as np import quimb.tensor as qtn @@ -21,6 +22,7 @@ Bloq, CompositeBloq, Connection, + ConnectionT, LeftDangle, QBit, Register, @@ -32,8 +34,10 @@ logger = logging.getLogger(__name__) +_IndT: TypeAlias = Any -def cbloq_to_quimb(cbloq: CompositeBloq) -> qtn.TensorNetwork: + +def cbloq_to_quimb(cbloq: CompositeBloq, friendly_indices: bool = False) -> qtn.TensorNetwork: """Convert a composite bloq into a tensor network. This function will call `Bloq.my_tensors` on each subbloq in the composite bloq to add @@ -42,6 +46,18 @@ def cbloq_to_quimb(cbloq: CompositeBloq) -> qtn.TensorNetwork: smallest form first. The small bloqs that result from a flattening 1) likely already have their `my_tensors` method implemented; and 2) can enable a more efficient tensor contraction path. + + Args: + cbloq: The composite bloq. + friendly_indices: If set to True, the outer indices of the tensor network will be renamed + from their Qualtran-computer-readable form to human-friendly strings. This may be + useful if you plan on manually manipulating the resulting tensor network but will + preclude any further processing by Qualtran functions. The indices are named + {soq.reg.name}{soq.idx}_{j}{side}, where j is the individual bit index and side is 'l' + or 'r' for left or right, respectively. + + Returns: + The tensor network """ tn = qtn.TensorNetwork([]) @@ -57,34 +73,248 @@ def cbloq_to_quimb(cbloq: CompositeBloq) -> qtn.TensorNetwork: out_d = _cxns_to_cxn_dict(bloq.signature.rights(), succ_cxns, get_me=lambda cxn: cxn.left) for tensor in bloq.my_tensors(inc_d, out_d): + if isinstance(tensor, DiscardInd): + raise ValueError( + f"During tensor simulation, {bloq} tried to discard information. This requires using `tensor_contract(superoperator=True)` or `cbloq_to_superquimb`." + ) tn.add(tensor) - # Special case: Add variables corresponding to all registers that don't connect to any Bloq. - # This is needed because `CompositeBloq.iter_bloqnections` ignores `LeftDangle/RightDangle` - # bloqs, and therefore we never see connections that exist only b/w LeftDangle and - # RightDangle bloqs. + # Special case: Add indices corresponding to unused wires for cxn in cbloq.connections: if cxn.left.binst is LeftDangle and cxn.right.binst is RightDangle: - # This register has no Bloq acting on it, and thus it would not have a variable in - # the tensor network. Add an identity tensor acting on this register to make sure the - # tensor network has variables corresponding to all input / output registers. - - n = cxn.left.reg.bitsize - for j in range(cxn.left.reg.bitsize): - - placeholder = Soquet(None, Register('simulation_placeholder', QBit())) # type: ignore - Connection(cxn.left, placeholder) - tn.add( - qtn.Tensor( - data=np.eye(2), - inds=[ - (Connection(cxn.left, placeholder), j), - (Connection(placeholder, cxn.right), j), - ], - ) - ) + # Connections that directly tie LeftDangle to RightDangle + for id_tensor in _get_placeholder_tensors(cxn): + tn.add(id_tensor) + + return tn.reindex(_get_outer_indices(tn, friendly_indices=friendly_indices)) + + +def _get_placeholder_tensors(cxn): + """Get identity placeholder tensors to directly connect LeftDangle to RightDangle. + + This function is used in `cbloq_to_quimb` and `cbloq_to_superquimb` for the following + contingency: + + >>> for cxn in cbloq.connections: + >>> if cxn.left.binst is LeftDangle and cxn.right.binst is RightDangle: + + This is needed because `CompositeBloq.iter_bloqnections` ignores `LeftDangle/RightDangle` + bloqs, and therefore we never see connections that exist only between LeftDangle and + RightDangle sentinel values. + """ + for j in range(cxn.left.reg.bitsize): + placeholder = Soquet(None, Register('simulation_placeholder', QBit())) # type: ignore + Connection(cxn.left, placeholder) + yield qtn.Tensor( + data=np.eye(2), + inds=[(Connection(cxn.left, placeholder), j), (Connection(placeholder, cxn.right), j)], + ) + + +_OuterIndT = Tuple[str, Tuple[int, ...], int, str] + + +def _get_outer_indices( + tn: 'qtn.TensorNetwork', friendly_indices: bool = False +) -> Dict[_IndT, Union[str, _OuterIndT]]: + """Provide a mapping for a tensor network's outer indices. + + Internal indices effectively use `qualtran.Connection` objects as their indices. The + outer indices correspond to connections to `DanglingT` items, and you end up having to + do logic to disambiguate left dangling indices from right dangling indices. This function + facilitates re-indexing the tensor network's outer indices to better match a bloq's signature. + + In particular, we map each outer index to a tuple (reg.name, soq.idx, j, group) where + group is 'l' or 'r' for left or right indices. + + If `friendly_indices` is set to True, the tuple of items is converted to a string. + + This function is called at the end of `cbloq_to_quimb` as part of a `tn.reindex(...) operation. + """ + ind_name_map: Dict[_IndT, Union[str, _OuterIndT]] = {} + + # Each index is a (cxn: Connection, j: int) tuple. + cxn: Connection + j: int + + for ind in tn.outer_inds(): + cxn, j = ind + if cxn.left.binst is LeftDangle: + soq = cxn.left + group = 'l' + elif cxn.right.binst is RightDangle: + soq = cxn.right + group = 'r' + else: + raise ValueError( + f"Outer indices of a tensor network should be " + f"connections to LeftDangle or RightDangle, not {cxn}" + ) + + if friendly_indices: + # Turn everything to strings + idx_str = f'{soq.idx}' if soq.idx else '' + ind_name_map[ind] = f'{soq.reg.name}{idx_str}_{j}{group}' + else: + # Keep as tuple + ind_name_map[ind] = (soq.reg.name, soq.idx, j, group) + + return ind_name_map + + +@attrs.frozen +class DiscardInd: + """Return `DiscardInd` in `Bloq.my_tensors()` to indicate an index should be discarded. + + We cannot discard an index from a state-vector pure-state simulation, so any bloq that + returns `DiscardInd` in its `my_tensors` method will cause an error in the ordinary + tensor contraction simulator. + + We can discard indices in open-system simulations by tracing out the index. When using + `Bloq.tensor_contract(superoperator=True)`, the index contained in a `DiscardInd` will be + traced out of the superoperator tensor network. + + Args: + ind_tuple: The index to trace out, of the form (cxn, j) where `j` addresses + individual bits. + """ + + ind_tuple: Tuple['ConnectionT', int] + + +def make_forward_tensor(t: qtn.Tensor): + new_inds = [(*ind, True) for ind in t.inds] + + t2 = t.copy() + t2.modify(inds=new_inds) + return t2 + + +def make_backward_tensor(t: qtn.Tensor): + new_inds = [] + for ind in t.inds: + new_inds.append((*ind, False)) + + t2 = t.H + t2.modify(inds=new_inds, tags=t.tags | {'dag'}) + return t2 + + +def cbloq_to_superquimb(cbloq: CompositeBloq, friendly_indices: bool = False) -> qtn.TensorNetwork: + """Convert a composite bloq into a superoperator tensor network. + + This simulation strategy can handle non-unitary dynamics, but is more costly. + + This function will call `Bloq.my_tensors` on each subbloq in the composite bloq to add + tensors to a quimb tensor network. This uses ths system+environment strategy for modeling + open system dynamics. In contrast to `cbloq_to_quimb`, each bloq will have + its tensors added twice: once to the part of the network representing the "forward" + wavefunction, and its conjugate added to the part of the network representing the "backward" + part of the wavefunction. If the bloq returns a sentinel value of the `DiscardInd` class, + that particular index is *traced out*: the forward and backward copies of the index are joined. + This corresponds to removing the qubit from the computation and integrating over its possible + values. Arbitrary non-unitary dynamics can be modeled by unitary interaction of the 'system' + with an 'environment' that is traced out. + + If a bloq returns a value of type `DiscardInd` in its tensors, this function must be + used. The ordinary `cbloq_to_quimb` will raise an error. + + Args: + cbloq: The composite bloq. + friendly_indices: If set to True, the outer indices of the tensor network will be renamed + from their Qualtran-computer-readable form to human-friendly strings. This may be + useful if you plan on manually manipulating the resulting tensor network but will + preclude any further processing by Qualtran functions. The indices are named + {soq.reg.name}{soq.idx}_{j}{side}{direction}, where j is the individual bit index, + side is 'l' or 'r' for left or right (resp.), and direction is 'f' or 'b' for the + forward or backward (adjoint) wavefunctions. + """ + tn = qtn.TensorNetwork([]) + + logging.info( + "Constructing a super tensor network for composite bloq of size %d", + len(cbloq.bloq_instances), + ) + + for binst, pred_cxns, succ_cxns in cbloq.iter_bloqnections(): + bloq = binst.bloq + assert isinstance(bloq, Bloq) + + inc_d = _cxns_to_cxn_dict(bloq.signature.lefts(), pred_cxns, get_me=lambda cxn: cxn.right) + out_d = _cxns_to_cxn_dict(bloq.signature.rights(), succ_cxns, get_me=lambda cxn: cxn.left) + + for tensor in bloq.my_tensors(inc_d, out_d): + if isinstance(tensor, DiscardInd): + dind = tensor.ind_tuple + tn.reindex({(*dind, True): dind, (*dind, False): dind}, inplace=True) + else: + forward_tensor = make_forward_tensor(tensor) + backward_tensor = make_backward_tensor(tensor) + tn.add(forward_tensor) + tn.add(backward_tensor) + + # Special case: Add indices corresponding to unused wires + for cxn in cbloq.connections: + if cxn.left.binst is LeftDangle and cxn.right.binst is RightDangle: + # Connections that directly tie LeftDangle to RightDangle + for id_tensor in _get_placeholder_tensors(cxn): + forward_tensor = make_forward_tensor(id_tensor) + backward_tensor = make_backward_tensor(id_tensor) + tn.add(forward_tensor) + tn.add(backward_tensor) + + return tn.reindex(_get_outer_superindices(tn, friendly_indices=friendly_indices)) + + +_SuperOuterIndT = Tuple[str, Tuple[int, ...], int, str] + + +def _get_outer_superindices( + tn: 'qtn.TensorNetwork', friendly_indices: bool = False +) -> Dict[_IndT, Union[str, _SuperOuterIndT]]: + """Provide a mapping for a super-tensor network's outer indices. + + Internal indices effectively use `qualtran.Connection` objects as their indices. The + outer indices correspond to connections to `DanglingT` items, and you end up having to + do logic to disambiguate left dangling indices from right dangling indices. This function + facilitates re-indexing the tensor network's outer indices to better match a bloq's signature. + + In particular, we map each outer index to a tuple (reg.name, soq.idx, j, group) where + group is 'lf', 'lb', 'rf', or 'rb' corresponding to (left or right) x (forward or backward) + indices. + + If `friendly_indices` is set to True, the tuple of items is converted to a string. + + This function is called at the end of `cbloq_to_superquimb` as part of a `tn.reindex(...) + operation. + """ + # Each index is a (cxn: Connection, j: int, forward: bool) tuple. + cxn: Connection + j: int + forward: bool + + ind_name_map: Dict[_IndT, Union[str, _SuperOuterIndT]] = {} + for ind in tn.outer_inds(): + cxn, j, forward = ind + if cxn.left.binst is LeftDangle: + soq = cxn.left + group = 'lf' if forward else 'lb' + elif cxn.right.binst is RightDangle: + soq = cxn.right + group = 'rf' if forward else 'rb' + else: + raise ValueError( + f"Outer indices of a tensor network should be " + f"connections to LeftDangle or RightDangle, not {cxn}" + ) + + if friendly_indices: + idx_str = f'{soq.idx}' if soq.idx else '' + ind_name_map[ind] = f'{soq.reg.name}{idx_str}_{j}{group}' + else: + ind_name_map[ind] = (soq.reg.name, soq.idx, j, group) - return tn + return ind_name_map def _add_classical_kets(bb: BloqBuilder, registers: Iterable[Register]) -> Dict[str, 'SoquetT']: diff --git a/qualtran/simulation/tensor/_quimb_test.py b/qualtran/simulation/tensor/_quimb_test.py index cde0e06e0e..ceb9a347a0 100644 --- a/qualtran/simulation/tensor/_quimb_test.py +++ b/qualtran/simulation/tensor/_quimb_test.py @@ -19,7 +19,7 @@ import quimb.tensor as qtn from attrs import frozen -from qualtran import Bloq, BloqBuilder, Connection, ConnectionT, DanglingT, QAny, Signature +from qualtran import Bloq, BloqBuilder, ConnectionT, QAny, Signature from qualtran.bloqs.bookkeeping import Join, Split from qualtran.simulation.tensor import cbloq_to_quimb @@ -49,10 +49,11 @@ def test_cbloq_to_quimb(): tn = cbloq_to_quimb(cbloq) assert len(tn.tensors) == 4 - for outer_ind in tn.outer_inds(): - cxn, j = outer_ind - assert isinstance(cxn, Connection) - assert isinstance(cxn.left.binst, DanglingT) or isinstance(cxn.right.binst, DanglingT) + assert sorted(tn.outer_inds()) == [ + # reg_name, idx, j, side_str + ('x', (), 0, 'l'), + ('x', (), 0, 'r'), + ] def test_cbloq_to_quimb_with_no_ops_on_register():