[open-systems] Phased classical simulation

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