Fix recursive phased classical simulation (#1731)

The classical simulators should recurse into bloq decompositions. The
phased classical simulator would use the ordinary classical simulator
when recursing since that was "baked into" the `Bloq` method -- and it
would throw an error even if it should have been able to simulate the
(sub-)circuit.

This changes the responsibility of dealing with the possible recursion
to the simulator classes themselves.

Author: mpharrigan

qualtran/_infra/bloq.py

@@ -272,19 +272,10 @@ def on_classical_vals(
                 registers with `shape` will be an ndarray of values of the expected type.
 
         Returns:
-            A dictionary mapping right (or thru) register name to output classical values.
+            A dictionary mapping right (or thru) register name to output classical values or
+                NotImplemented if this is not classical-reversible logic.
         """
-        try:
-            return self.decompose_bloq().on_classical_vals(**vals)
-        except DecomposeTypeError as e:
-            raise NotImplementedError(f"{self} is not classically simulable.") from e
-        except DecomposeNotImplementedError as e:
-            raise NotImplementedError(
-                f"{self} has no decomposition and does not "
-                f"support classical simulation directly"
-            ) from e
-        except NotImplementedError as e:
-            raise NotImplementedError(f"{self} does not support classical simulation: {e}") from e
+        return NotImplemented
 
     def basis_state_phase(
         self, **vals: 'ClassicalValT'

qualtran/_infra/controlled.py

@@ -440,12 +440,12 @@ def on_classical_vals(self, **vals: 'ClassicalValT') -> Mapping[str, 'ClassicalV
         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):
-            rets = {
-                **self.subbloq.on_classical_vals(**other_vals),
-                **{
-                    reg_name: ctrl_val for reg_name, ctrl_val in zip(self.ctrl_reg_names, ctrl_vals)
-                },
-            }
+            rets = self.subbloq.on_classical_vals(**other_vals)
+            if rets is NotImplemented:
+                return NotImplemented
+            rets |= {
+                reg_name: ctrl_val for reg_name, ctrl_val in zip(self.ctrl_reg_names, ctrl_vals)
+            }  # type: ignore[operator]
             return rets
 
         return vals

qualtran/bloqs/basic_gates/z_basis.py

@@ -268,6 +268,15 @@ def my_tensors(
             )
         ]
 
+    def on_classical_vals(self, **vals: 'ClassicalValT') -> Dict[str, 'ClassicalValT']:
+        # Diagonal, but causes phases: see `basis_state_phase`
+        return vals
+
+    def basis_state_phase(self, q: int) -> Optional[complex]:
+        if q == 1:
+            return -1
+        return 1
+
     def get_ctrl_system(self, ctrl_spec: 'CtrlSpec') -> Tuple['Bloq', 'AddControlledT']:
         if ctrl_spec != CtrlSpec():
             # Delegate to the general superclass behavior

qualtran/simulation/classical_sim.py

@@ -34,11 +34,14 @@
 import numpy as np
 import sympy
 from numpy.typing import NDArray
+from typing_extensions import Self
 
 from qualtran import (
     Bloq,
     BloqInstance,
     DanglingT,
+    DecomposeNotImplementedError,
+    DecomposeTypeError,
     LeftDangle,
     Register,
     RightDangle,
@@ -301,31 +304,49 @@ def _update_assign_from_vals(
                 soq = Soquet(binst, reg)
                 self.soq_assign[soq] = val
 
-    def _binst_on_classical_vals(self, binst, in_vals) -> None:
-        """Call `on_classical_vals` on a given bloq instance."""
-        bloq = binst.bloq
+    def _recurse_impl(self, cbloq: 'CompositeBloq', in_vals):
+        """Overridable function to recursively simulate a composite bloq."""
+        out_vals, _ = call_cbloq_classically(cbloq.signature, in_vals, cbloq._binst_graph)
+        phase = None
 
-        out_vals = bloq.on_classical_vals(**in_vals)
-        if not isinstance(out_vals, dict):
-            raise TypeError(
-                f"{bloq.__class__.__name__}.on_classical_vals should return a dictionary."
-            )
-        self._update_assign_from_vals(bloq.signature.rights(), binst, out_vals)
+        return out_vals, phase
 
-    def _binst_basis_state_phase(self, binst, in_vals) -> None:
-        """Call `basis_state_phase` on a given bloq instance.
+    def _recurse(self, binst: 'BloqInstance', in_vals) -> Self:
+        """Recursively simulate a composite bloq.
 
-        This base simulation class will raise an error if the bloq reports any phasing.
-        This method is overwritten in `PhasedClassicalSimState` to support phasing.
+        This handles decomposing the bloq and using the results of the sub-simulation
+        to update the simulator state. Developers should override `_recurse_impl` to
+        customize the sub-simulation.
         """
         bloq = binst.bloq
-        bloq_phase = bloq.basis_state_phase(**in_vals)
+        try:
+            cbloq = bloq.decompose_bloq()
+            out_vals, bloq_phase = self._recurse_impl(cbloq, in_vals)
+
+        except DecomposeTypeError as e:
+            raise NotImplementedError(f"{bloq} is not classically simulable.") from e
+        except DecomposeNotImplementedError as e:
+            raise NotImplementedError(
+                f"{bloq} has no decomposition and does not "
+                f"support classical simulation directly"
+            ) from e
+        except NotImplementedError as e:
+            raise NotImplementedError(f"{bloq} does not support classical simulation: {e}") from e
+
+        self._update(binst, out_vals, bloq_phase)
+        return self
+
+    def _update(self, binst: 'BloqInstance', out_vals, bloq_phase: Union[complex, None]) -> None:
+        """Overridable method to update the current simulator state."""
+        self._update_assign_from_vals(binst.bloq.signature.rights(), binst, out_vals)
+
         if bloq_phase is not None:
             raise ValueError(
-                f"{bloq} imparts a phase, and can't be simulated purely classically. Consider using `do_phased_classical_simulation`."
+                f"{binst.bloq} imparts a phase of {bloq_phase}, and can't be simulated purely classically. "
+                f"Consider using `do_phased_classical_simulation`."
             )
 
-    def step(self) -> 'ClassicalSimState':
+    def step(self) -> Self:
         """Advance the simulation by one bloq instance.
 
         After calling this method, `self.last_binst` will contain the bloq instance that
@@ -349,11 +370,32 @@ def _in_vals(reg: Register):
             return _get_in_vals(binst, reg, soq_assign=self.soq_assign)
 
         bloq = binst.bloq
+        bcls_name = bloq.__class__.__name__
         in_vals = {reg.name: _in_vals(reg) for reg in bloq.signature.lefts()}
+        out_vals = bloq.on_classical_vals(**in_vals)
+        bloq_phase = bloq.basis_state_phase(**in_vals)
+
+        #                     +--+ basis_state_phase
+        #   +- on_classical_vals
+        #   |                 |
+        # dict              None           Use classical values.
+        # dict              number         Use classical values and phase only if doing phased sim
+        # NotImplemented    None           decompose and use the correct simulator type
+        # NotImplemented    number         error
+
+        if out_vals is NotImplemented:
+            if bloq_phase is not None:
+                raise ValueError(
+                    f"`basis_state_phase` defined on {bcls_name}, but not `on_classical_vals`"
+                )
+
+            return self._recurse(binst, in_vals)
+
+        if not isinstance(out_vals, dict):
+            raise TypeError(f"`{bcls_name}.on_classical_vals` should return a dictionary.")
+
+        self._update(binst, out_vals, bloq_phase)
 
-        # Apply methods
-        self._binst_on_classical_vals(binst, in_vals)
-        self._binst_basis_state_phase(binst, in_vals)
         return self
 
     def finalize(self) -> Dict[str, 'ClassicalValT']:
@@ -466,14 +508,19 @@ def from_cbloq(
             random_handler=rnd_handler,
         )
 
-    def _binst_basis_state_phase(self, binst, in_vals):
-        """Call `basis_state_phase` on a given bloq instance.
+    def _recurse_impl(self, cbloq, in_vals):
+        """Use phased classical simulation when recursing."""
+        sim = PhasedClassicalSimState(
+            cbloq.signature, cbloq._binst_graph, in_vals, random_handler=self._random_handler
+        )
+        final_vals = sim.simulate()
+        phase = sim.phase
+        return final_vals, phase
+
+    def _update(self, binst: 'BloqInstance', out_vals, bloq_phase: Union[complex, None]) -> None:
+        """Update the current simulator state, including phase tracking."""
+        self._update_assign_from_vals(binst.bloq.signature.rights(), binst, out_vals)
 
-        If this method returns a value, the current phase will be updated. Otherwise, we
-        leave the phase as-is.
-        """
-        bloq = binst.bloq
-        bloq_phase = bloq.basis_state_phase(**in_vals)
         if isinstance(bloq_phase, MeasurementPhase):
             # In this special case, there is a coupling between the classical result and the
             # phase result (because the classical result is stochastic). We look up the measurement

qualtran/simulation/classical_sim_test.py

@@ -15,6 +15,7 @@
 import itertools
 from typing import Dict, Union
 
+import attrs
 import networkx as nx
 import numpy as np
 import pytest
@@ -38,8 +39,8 @@
     Register,
     Side,
     Signature,
+    SoquetT,
 )
-from qualtran.bloqs.basic_gates import CNOT
 from qualtran.simulation.classical_sim import (
     _BannedClassicalValHandler,
     _FixedClassicalValHandler,
@@ -153,6 +154,8 @@ def test_normal_classical_on_phased():
 
 
 def test_cnot_assign_dict():
+    from qualtran.bloqs.basic_gates import CNOT
+
     cbloq = CNOT().as_composite_bloq()
     binst_graph = cbloq._binst_graph  # pylint: disable=protected-access
     vals = dict(ctrl=1, target=0)
@@ -338,3 +341,38 @@ def test_phased_classical_distribution():
     )
     assert final_values['c'] == 1
     assert phase == 1
+
+
+@attrs.frozen
+class ComposedPhasing(Bloq):
+    n: int = 0
+
+    @property
+    def signature(self) -> 'Signature':
+        return Signature([Register('x', QBit(), side=Side.RIGHT)])
+
+    def build_composite_bloq(self, bb: 'BloqBuilder') -> Dict[str, 'SoquetT']:
+        from qualtran.bloqs.basic_gates import OneState, ZGate
+
+        x = bb.add(OneState())
+        for _ in range(self.n):
+            x = bb.add(ZGate(), q=x)
+        return {'x': x}
+
+
+def test_derive_phasing_from_composed_bloq():
+    vals, phase = do_phased_classical_simulation(ComposedPhasing(0), {})
+    assert vals == {'x': 1}
+    assert phase == +1.0
+
+    vals, phase = do_phased_classical_simulation(ComposedPhasing(1), {})
+    assert vals == {'x': 1}
+    assert phase == -1.0
+
+    vals, phase = do_phased_classical_simulation(ComposedPhasing(2), {})
+    assert vals == {'x': 1}
+    assert phase == +1.0
+
+    vals, phase = do_phased_classical_simulation(ComposedPhasing(3), {})
+    assert vals == {'x': 1}
+    assert phase == -1.0