[open-systems] Superoperator tensor contraction

qualtran/_infra/bloq.py

@@ -61,13 +61,18 @@
         SympySymbolAllocator,
     )
     from qualtran.simulation.classical_sim import ClassicalValRetT, ClassicalValT
+    from qualtran.simulation.tensor import DiscardInd
 
 
 def _decompose_from_build_composite_bloq(bloq: 'Bloq') -> 'CompositeBloq':
     from qualtran import BloqBuilder
 
     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

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))

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)
 
 

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,
+)

qualtran/bloqs/basic_gates/discard.py

@@ -0,0 +1,68 @@
+#  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 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))]

qualtran/bloqs/basic_gates/discard_test.py

@@ -0,0 +1,82 @@
+#  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 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)

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).