[open-systems] Superoperator tensor contraction (#1596)

This pull request is the third in a series for supporting open system
(classical data & measurement), roadmap item
#445.

~~This branch is based on #1590 (for my convenience; it doesn't use any
of that functionality)~~

## `Bloq.tensor_contract(superoperator=True)`

This is the new interface. It will return a 0-, 2-, or 4-dimensional
tensor for constants, density matrices, and superoperators (resp) to
support simulating programs with measurements and non-unitary maps. Read
the docs for the indexing convention for the superoperator; but if
you're really doing something with the superoperators you may want to
operate on the quimb tensor network with friendly indices itself.

## `cbloq_to_superquimb`

The workhorse of the new functionality is this function that converts a
composite bloq to a tensor network that represents the density matrix
and/or superoperator of the program. Some operations like measurement or
discarding qubits cannot be represented in a pure-state statevector /
unitary picture.

`cbloq_to_superquimb` still uses `Bloq.my_tensors` to construct the
network. It adds each tensor twice: once in the "forward" direction and
again in the "backward" (adjoint) direction.

## `DiscardInd`

The `Bloq.my_tensors` method can now return a `DiscardInd` object, which
signifies that that tensor index should be discarded by tracing it out.

It turns out that this simple extension lets you represent any CPTP map
using the "system+environment" modeling approach.


## `MeasZ`, `Discard`

This PR includes two non-unitary maps: measuring in the Z basis and
discarding a qubit or classical bit.


## Manipulating quimb tensor networks

This PR makes some changes to the way we handle outer indices in the
quimb tensor network. This should only be of interest to folks who use
the "hands on" approach of directly manipulating the tensor networks.
The changes make sense and clean up some of the docs.

The outer indices are now re-mapped to tuples of qualtran objects that
actually make sense rather than keeping them as `Connection`s to
`DanglingT` objects. There is a new flag `friendly_indices=True` which
will make the outer indices friendly strings.

**breaking change** I renamed and made-private the
`get_left_and_right_inds` function. The original idea was to use this as
a public way of getting a handle on the outer indices, but it was always
tricky to use; and needed to be changed to support the new variety of
outer indices present in the superoperator networks. No one should need
a replacement for this function since the `qtn.TensorNetwork` object
outer indices now just intrinsically make sense.

## Non-changes

This includes the minimal set of non-unitary CPTP bloqs to support the
docs notebooks.

Author: mpharrigan

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

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)

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