[open-systems] Phased classical simulation (#1590)

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

~~This branch is based on #1584~~

## `Bloq.basis_state_phase`

This new method can be overridden for phased-classical gates. In this
PR, I show how `CZ` can support this new method. This leaves the basis
states alone but applies a phase when `11` is encountered.

## `PhasedClassicalSimState(ClassicalSimState)`

I refactored the classical simulator into a (mutable) class with
initialization `.from_cbloq(...)`, binst-by-binst stepping `.step()`,
and finalization `.finalize()`. These can be combined with
`.simulate()`, which is the new backend for the familiar entry points
`Bloq.call_classically(...)` and friends.

Doing ordinary classical simulation will raise an error if it encounters
a bloq with a non-trivial phase. The phased classical simulator will
raise an error if it encounters a bloq without classical action;
otherwise it will propagate phase. If a bloq does not specifically
annotate a phase (but does indeed have classical action), the trivial +1
phase is assumed.

## Non-changes (yet)

This PR does not directly support any "open systems" things. The
ultimate goal is to support fuzz-testing measurement-based
uncomputation. Since MBUC involves tracking phases, this PR is a
necessary but not sufficient pre-requisite.

## Related

 - This PR unblocks #1527 
 - #1499 easier to test
- https://arxiv.org/abs/2105.13410 has a ton of cool circuit
constructions that can be tested with this

Author: mpharrigan

qualtran/_infra/bloq.py

@@ -60,7 +60,7 @@
         GeneralizerT,
         SympySymbolAllocator,
     )
-    from qualtran.simulation.classical_sim import ClassicalValT
+    from qualtran.simulation.classical_sim import ClassicalValRetT, ClassicalValT
 
 
 def _decompose_from_build_composite_bloq(bloq: 'Bloq') -> 'CompositeBloq':
@@ -184,13 +184,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 +215,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', ...]:

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/bloqs/basic_gates/z_basis.py

@@ -52,6 +52,7 @@
 
     from qualtran.cirq_interop import CirqQuregT
     from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator
+    from qualtran.simulation.classical_sim import ClassicalValT
 
 _ZERO = np.array([1, 0], dtype=np.complex128)
 _ONE = np.array([0, 1], dtype=np.complex128)
@@ -359,6 +360,15 @@ def get_ctrl_system(self, ctrl_spec: 'CtrlSpec') -> Tuple['Bloq', 'AddControlled
             self, ctrl_spec, current_ctrl_bit=1, bloq_with_ctrl=self, ctrl_reg_name='q1'
         )
 
+    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, q1: int, q2: int) -> Optional[complex]:
+        if q1 == 1 and q2 == 1:
+            return -1
+        return 1
+
 
 @bloq_example
 def _cz() -> CZ:

qualtran/bloqs/basic_gates/z_basis_test.py

@@ -24,6 +24,7 @@
     IntState,
     OneEffect,
     OneState,
+    XGate,
     ZeroEffect,
     ZeroState,
     ZGate,
@@ -245,3 +246,33 @@ def test_cz_manual():
 
     assert ZGate().controlled() == CZ()
     assert t_complexity(cz) == TComplexity(clifford=1)
+
+    with pytest.raises(ValueError, match='.*phase.*'):
+        cz.call_classically(q1=1, q2=1)
+
+
+def test_cz_phased_classical():
+    cz = CZ()
+    from qualtran.simulation.classical_sim import do_phased_classical_simulation
+
+    final_vals, phase = do_phased_classical_simulation(cz, {'q1': 0, 'q2': 1})
+    assert final_vals['q1'] == 0
+    assert final_vals['q2'] == 1
+    assert phase == 1
+
+    final_vals, phase = do_phased_classical_simulation(cz, {'q1': 1, 'q2': 1})
+    assert final_vals['q1'] == 1
+    assert final_vals['q2'] == 1
+    assert phase == -1
+
+    bb = BloqBuilder()
+    q1 = bb.add(ZeroState())
+    q2 = bb.add(ZeroState())
+    q1 = bb.add(XGate(), q=q1)
+    q2 = bb.add(XGate(), q=q2)
+    q1, q2 = bb.add(CZ(), q1=q1, q2=q2)
+    cbloq = bb.finalize(q1=q1, q2=q2)
+    final_vals, phase = do_phased_classical_simulation(cbloq, {})
+    assert final_vals['q1'] == 1
+    assert final_vals['q2'] == 1
+    assert phase == -1