[open-systems] Superoperator tensor contraction

qualtran/_infra/bloq.py

@@ -60,14 +60,19 @@
         GeneralizerT,
         SympySymbolAllocator,
     )
-    from qualtran.simulation.classical_sim import ClassicalValT
+    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)
 
 
@@ -184,13 +189,13 @@ def adjoint(self) -> 'Bloq':
 
     def on_classical_vals(
         self, **vals: Union['sympy.Symbol', 'ClassicalValT']
-    ) -> Mapping[str, 'ClassicalValT']:
+    ) -> Mapping[str, 'ClassicalValRetT']:
         """How this bloq operates on classical data.
 
         Override this method if your bloq represents classical, reversible logic. For example:
         quantum circuits composed of X and C^nNOT gates are classically simulable.
 
-        Bloq definers should override this method. If you already have an instance of a `Bloq`,
+        Bloq authors should override this method. If you already have an instance of a `Bloq`,
         consider calling `call_clasically(**vals)` which will do input validation before
         calling this function.
 
@@ -215,6 +220,25 @@ def on_classical_vals(
         except NotImplementedError as e:
             raise NotImplementedError(f"{self} does not support classical simulation: {e}") from e
 
+    def basis_state_phase(self, **vals: 'ClassicalValT') -> Union[complex, None]:
+        """How this bloq phases classical basis states.
+
+        Override this method if your bloq represents classical logic with basis-state
+        dependent phase factors. This corresponds to bloqs whose matrix representation
+        (in the standard basis) is a generalized permutation matrix: a permutation matrix
+        where each entry can be +1, -1 or any complex number with unit absolute value.
+        Alternatively, this corresponds to bloqs composed from classical operations
+        (X, CNOT, Toffoli, ...) and diagonal operations (T, CZ, CCZ, ...).
+
+        Bloq authors should override this method. If you are using an instantiated bloq object,
+        call TODO and not this method directly.
+
+        If this method is implemented, `on_classical_vals` must also be implemented.
+        If `on_classical_vals` is implemented but this method is not implemented, it is assumed
+        that the bloq does not alter the phase.
+        """
+        return None
+
     def call_classically(
         self, **vals: Union['sympy.Symbol', 'ClassicalValT']
     ) -> Tuple['ClassicalValT', ...]:
@@ -239,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.py

@@ -47,7 +47,7 @@
     from qualtran.cirq_interop import CirqQuregT
     from qualtran.drawing import WireSymbol
     from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator
-    from qualtran.simulation.classical_sim import ClassicalValT
+    from qualtran.simulation.classical_sim import ClassicalValRetT, ClassicalValT
 
 ControlBit: TypeAlias = int
 """A control bit, either 0 or 1."""
@@ -432,7 +432,7 @@ def signature(self) -> 'Signature':
         # Prepend register(s) corresponding to `ctrl_spec`.
         return Signature(self.ctrl_regs + tuple(self.subbloq.signature))
 
-    def on_classical_vals(self, **vals: 'ClassicalValT') -> Mapping[str, 'ClassicalValT']:
+    def on_classical_vals(self, **vals: 'ClassicalValT') -> Mapping[str, 'ClassicalValRetT']:
         """Classical action of controlled bloqs.
 
         This involves conditionally doing the classical action of `subbloq`. All implementers
@@ -453,6 +453,18 @@ def on_classical_vals(self, **vals: 'ClassicalValT') -> Mapping[str, 'ClassicalV
 
         return vals
 
+    def basis_state_phase(self, **vals: 'ClassicalValT') -> Union[complex, None]:
+        """Phasing action of controlled bloqs.
+
+        This involves conditionally doing the phasing action of `subbloq`. All implementers
+        of `_ControlledBase` should provide a decomposition that satisfies this phase funciton.
+        """
+        ctrl_vals = [vals[reg_name] for reg_name in self.ctrl_reg_names]
+        other_vals = {reg.name: vals[reg.name] for reg in self.subbloq.signature}
+        if self.ctrl_spec.is_active(*ctrl_vals):
+            return self.subbloq.basis_state_phase(**other_vals)
+        return None
+
     def _tensor_data(self):
         """Dense tensor encoding a controlled unitary.
 

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)