[open-systems] Testing Measurement-Based Uncomputation

qualtran/_infra/bloq.py

@@ -60,7 +60,7 @@
         GeneralizerT,
         SympySymbolAllocator,
     )
-    from qualtran.simulation.classical_sim import ClassicalValRetT, ClassicalValT
+    from qualtran.simulation.classical_sim import ClassicalValRetT, ClassicalValT, MeasurementPhase
     from qualtran.simulation.tensor import DiscardInd
 
 
@@ -220,7 +220,9 @@ 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]:
+    def basis_state_phase(
+        self, **vals: 'ClassicalValT'
+    ) -> Union[complex, 'MeasurementPhase', None]:
         """How this bloq phases classical basis states.
 
         Override this method if your bloq represents classical logic with basis-state
@@ -231,7 +233,8 @@ def basis_state_phase(self, **vals: 'ClassicalValT') -> Union[complex, None]:
         (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.
+        call `qualtran.simulation.classical_sim.do_phased_classical_simulation` or use
+        `qualtran.simulation.classical_sim.PhasedClassicalSimState`.
 
         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

qualtran/_infra/composite_bloq.py

@@ -512,6 +512,42 @@ def _binst_to_cxns(
     return pred_cxns, succ_cxns
 
 
+def _get_soquet(
+    binst: 'BloqInstance',
+    reg_name: str,
+    right: bool = False,
+    idx: Tuple[int, ...] = (),
+    *,
+    binst_graph: nx.DiGraph,
+) -> 'Soquet':
+    """Retrieve a soquet given identifying information.
+
+    We can uniquely address a Soquet by the arguments to this function.
+
+    Args:
+        binst: The bloq instance associated with the desired soquet.
+        reg_name: The name of the register associated with the desired soquet.
+        right: If False, get the input, left soquet. Otherwise: the right, output soquet
+        idx: The index of the soquet within a multidimensional register, or the empty
+            tuple for basic registers.
+    """
+    preds, succs = _binst_to_cxns(binst, binst_graph=binst_graph)
+    if right:
+        for suc in succs:
+            me = suc.left
+            if me.reg.name == reg_name and me.idx == idx:
+                return me
+    else:
+        for pred in preds:
+            me = pred.right
+            if me.reg.name == reg_name and me.idx == idx:
+                return me
+
+    raise ValueError(
+        f"Could not find the requested soquet with {binst=}, {reg_name=}, {right=}, {idx=}"
+    )
+
+
 def _cxns_to_soq_dict(
     regs: Iterable[Register],
     cxns: Iterable[Connection],

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 ClassicalValRetT, ClassicalValT
+    from qualtran.simulation.classical_sim import ClassicalValRetT, ClassicalValT, MeasurementPhase
 
 ControlBit: TypeAlias = int
 """A control bit, either 0 or 1."""
@@ -380,10 +380,7 @@ def ctrl_spec(self) -> 'CtrlSpec':
 
     @cached_property
     def _thru_registers_only(self) -> bool:
-        for reg in self.subbloq.signature:
-            if reg.side != Side.THRU:
-                return False
-        return True
+        return self.signature.thru_registers_only
 
     @staticmethod
     def _make_ctrl_system(cb: '_ControlledBase') -> Tuple['_ControlledBase', 'AddControlledT']:
@@ -453,7 +450,9 @@ def on_classical_vals(self, **vals: 'ClassicalValT') -> Mapping[str, 'ClassicalV
 
         return vals
 
-    def basis_state_phase(self, **vals: 'ClassicalValT') -> Union[complex, None]:
+    def basis_state_phase(
+        self, **vals: 'ClassicalValT'
+    ) -> Union[complex, 'MeasurementPhase', None]:
         """Phasing action of controlled bloqs.
 
         This involves conditionally doing the phasing action of `subbloq`. All implementers
@@ -533,7 +532,15 @@ def wire_symbol(self, reg: Optional[Register], idx: Tuple[int, ...] = tuple()) -
         from qualtran.drawing import Text
 
         if reg is None:
-            return Text(f'C[{self.subbloq}]')
+            sub_title = self.subbloq.wire_symbol(None, idx)
+            if not isinstance(sub_title, Text):
+                raise ValueError(
+                    f"{self.subbloq} should return a `Text` object for reg=None wire symbol."
+                )
+            if sub_title.text == '':
+                return Text('')
+
+            return Text(f'C[{sub_title.text}]')
         if reg.name not in self.ctrl_reg_names:
             # Delegate to subbloq
             return self.subbloq.wire_symbol(reg, idx)
@@ -688,6 +695,7 @@ def make_ctrl_system_with_correct_metabloq(
        `ControlledViaAnd`, which computes the activation function once and re-uses it
        for each subbloq in the decomposition of `bloq`.
     """
+    from qualtran.bloqs.mcmt.classically_controlled import ClassicallyControlled
     from qualtran.bloqs.mcmt.controlled_via_and import ControlledViaAnd
 
     if ctrl_spec == CtrlSpec():
@@ -710,6 +718,6 @@ def make_ctrl_system_with_correct_metabloq(
     if qdtypes:
         return ControlledViaAnd.make_ctrl_system(bloq, ctrl_spec=ctrl_spec)
     if cdtypes:
-        raise NotImplementedError("Stay tuned...")
+        return ClassicallyControlled.make_ctrl_system(bloq, ctrl_spec=ctrl_spec)
 
     raise ValueError(f"Invalid control spec: {ctrl_spec}")

qualtran/_infra/registers.py

@@ -16,6 +16,7 @@
 import enum
 import itertools
 from collections import defaultdict
+from functools import cached_property
 from typing import cast, Dict, Iterable, Iterator, List, overload, Tuple, Union
 
 import attrs
@@ -179,6 +180,13 @@ def build_from_dtypes(cls, **registers: QCDType) -> 'Signature':
         """
         return cls(Register(name=k, dtype=v) for k, v in registers.items() if v.num_qubits)
 
+    @cached_property
+    def thru_registers_only(self) -> bool:
+        for reg in self:
+            if reg.side != Side.THRU:
+                return False
+        return True
+
     def lefts(self) -> Iterable[Register]:
         """Iterable over all registers that appear on the LEFT as input."""
         yield from self._lefts.values()

qualtran/bloqs/basic_gates/__init__.py

@@ -34,7 +34,7 @@
 from .swap import CSwap, Swap, TwoBitCSwap, TwoBitSwap
 from .t_gate import TGate
 from .toffoli import Toffoli
-from .x_basis import MinusEffect, MinusState, PlusEffect, PlusState, XGate
+from .x_basis import MeasX, MinusEffect, MinusState, PlusEffect, PlusState, XGate
 from .y_gate import CYGate, YGate
 from .z_basis import (
     CZ,

qualtran/bloqs/basic_gates/x_basis.py

@@ -24,6 +24,7 @@
     bloq_example,
     BloqBuilder,
     BloqDocSpec,
+    CBit,
     ConnectionT,
     CtrlSpec,
     QBit,
@@ -33,6 +34,12 @@
     SoquetT,
 )
 from qualtran.drawing import directional_text_box, Text, WireSymbol
+from qualtran.simulation.classical_sim import (
+    ClassicalValDistribution,
+    ClassicalValRetT,
+    ClassicalValT,
+    MeasurementPhase,
+)
 
 if TYPE_CHECKING:
     import cirq
@@ -41,7 +48,6 @@
     from pennylane.wires import Wires
 
     from qualtran.cirq_interop import CirqQuregT
-    from qualtran.simulation.classical_sim import ClassicalValT
 
 _PLUS = np.ones(2, dtype=np.complex128) / np.sqrt(2)
 _MINUS = np.array([1, -1], dtype=np.complex128) / np.sqrt(2)
@@ -85,7 +91,10 @@ def my_tensors(
         ]
 
     def as_cirq_op(
-        self, qubit_manager: 'cirq.QubitManager', **cirq_quregs: 'CirqQuregT'  # type: ignore[type-var]
+        self,
+        qubit_manager: 'cirq.QubitManager',
+        **cirq_quregs: 'CirqQuregT',
+        # type: ignore[type-var]
     ) -> Tuple[Union['cirq.Operation', None], Dict[str, 'CirqQuregT']]:  # type: ignore[type-var]
         if not self.state:
             raise ValueError(f"There is no Cirq equivalent for {self}")
@@ -270,3 +279,47 @@ def wire_symbol(self, reg: Register, idx: Tuple[int, ...] = tuple()) -> 'WireSym
             return Text('X')
 
         return ModPlus()
+
+    def __str__(self):
+        return 'X'
+
+
+@frozen
+class MeasX(Bloq):
+    @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) -> Dict[str, 'ClassicalValRetT']:
+        if q not in [0, 1]:
+            raise ValueError(f"Invalid classical value encountered in {self}: {q}")
+        return {'c': ClassicalValDistribution(2)}
+
+    def basis_state_phase(self, q: int) -> Union[complex, MeasurementPhase]:
+        if q == 0:
+            return 1
+        if q == 1:
+            return MeasurementPhase(reg_name='c')
+        raise ValueError(f"Invalid classical value encountered in {self}: {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
+
+        data = np.array(
+            [
+                [[0.5 + 0.0j, 0.5 + 0.0j], [0.5 + 0.0j, -0.5 + 0.0j]],
+                [[0.5 + 0.0j, -0.5 + 0.0j], [0.5 + 0.0j, 0.5 + 0.0j]],
+            ]
+        )
+
+        q_trace = qtn.rand_uuid('q_trace')
+        t = qtn.Tensor(
+            data=data, inds=[(incoming['q'], 0), (q_trace, 0), (outgoing['c'], 0)], tags=[str(self)]
+        )
+        return [t, DiscardInd((q_trace, 0))]

qualtran/bloqs/basic_gates/x_basis_test.py

@@ -15,7 +15,7 @@
 import numpy as np
 
 from qualtran import BloqBuilder
-from qualtran.bloqs.basic_gates import MinusState, PlusEffect, PlusState, XGate
+from qualtran.bloqs.basic_gates import MeasX, MinusState, PlusEffect, PlusState, XGate
 from qualtran.resource_counting import GateCounts, get_cost_value, QECGatesCost
 from qualtran.simulation.classical_sim import (
     format_classical_truth_table,
@@ -119,3 +119,8 @@ def _keep_and(b):
     bloq = XGate().controlled(CtrlSpec(qdtypes=QUInt(n), cvs=1))
     _, sigma = bloq.call_graph(keep=_keep_and)
     assert sigma == {And(): n - 1, CNOT(): 1, And().adjoint(): n - 1, XGate(): 4 * (n - 1)}
+
+
+def test_meas_x_classical_sim():
+    m = MeasX()
+    m.call_classically(q=0)

qualtran/bloqs/mcmt/and_bloq.py

@@ -103,7 +103,10 @@ def on_classical_vals(
             return {'ctrl': ctrl, 'target': out}
 
         # Uncompute
-        assert target == out
+        if target != out:
+            raise ValueError(
+                f"Inconsistent `target` found for uncomputing `And`: {ctrl=}, {target=}. Expected target={out}"
+            )
         return {'ctrl': ctrl}
 
     def my_tensors(

qualtran/bloqs/mcmt/classically_controlled.py

@@ -0,0 +1,40 @@
+#  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 typing import Tuple
+
+import attrs
+
+from qualtran import AddControlledT, Bloq, CDType, CtrlSpec, QCDType
+from qualtran._infra.controlled import _ControlledBase
+
+
+@attrs.frozen
+class ClassicallyControlled(_ControlledBase):
+
+    subbloq: 'Bloq'
+    ctrl_spec: 'CtrlSpec'
+
+    def __attrs_post_init__(self):
+        for qcdtype in self.ctrl_spec.qdtypes:
+            if not isinstance(qcdtype, QCDType):
+                raise ValueError(f"Invalid type found in `ctrl_spec`: {qcdtype}")
+            if not isinstance(qcdtype, CDType):
+                raise ValueError(f"Invalid type found in `ctrl_spec`: {qcdtype}")
+
+    @classmethod
+    def make_ctrl_system(
+        cls, bloq: 'Bloq', ctrl_spec: 'CtrlSpec'
+    ) -> Tuple['_ControlledBase', 'AddControlledT']:
+        cb = cls(subbloq=bloq, ctrl_spec=ctrl_spec)
+        return cls._make_ctrl_system(cb)

qualtran/resource_counting/_bloq_counts.py

@@ -297,7 +297,15 @@ class QECGatesCost(CostKey[GateCounts]):
     legacy_shims: bool = False
 
     def compute(self, bloq: 'Bloq', get_callee_cost: Callable[['Bloq'], GateCounts]) -> GateCounts:
-        from qualtran.bloqs.basic_gates import GlobalPhase, Identity, Toffoli, TwoBitCSwap
+        from qualtran.bloqs.basic_gates import (
+            Discard,
+            GlobalPhase,
+            Identity,
+            MeasX,
+            MeasZ,
+            Toffoli,
+            TwoBitCSwap,
+        )
         from qualtran.bloqs.basic_gates._shims import Measure
         from qualtran.bloqs.bookkeeping._bookkeeping_bloq import _BookkeepingBloq
         from qualtran.bloqs.mcmt import And, MultiTargetCNOT
@@ -326,7 +334,7 @@ def compute(self, bloq: 'Bloq', get_callee_cost: Callable[['Bloq'], GateCounts])
             return GateCounts(toffoli=1)
 
         # Measurement
-        if isinstance(bloq, Measure):
+        if isinstance(bloq, (Measure, MeasZ, MeasX)):
             return GateCounts(measurement=1)
 
         # 'And' bloqs
@@ -370,9 +378,10 @@ def compute(self, bloq: 'Bloq', get_callee_cost: Callable[['Bloq'], GateCounts])
             return GateCounts()
 
         # Bookkeeping, empty bloqs
-        if isinstance(bloq, _BookkeepingBloq) or isinstance(bloq, (GlobalPhase, Identity)):
+        if isinstance(bloq, _BookkeepingBloq) or isinstance(bloq, (GlobalPhase, Identity, Discard)):
             return GateCounts()
 
+        # Rotations
         if bloq_is_rotation(bloq):
             return GateCounts(rotation=1)