Also handle Z as the target gate.

Author: maxglick

qualtran/bloqs/multiplexers/selected_majorana_fermion.py

@@ -139,7 +139,7 @@ def nth_operation(  # type: ignore[override]
         yield cirq.CZ(*accumulator, target[target_idx])
 
     def on_classical_vals(self, **vals) -> Dict[str, 'ClassicalValT']:
-        if self.target_gate != cirq.X:
+        if self.target_gate != cirq.X and self.target_gate != cirq.Z:
             return NotImplemented
         if len(self.control_registers) > 1 or len(self.selection_registers) > 1:
             return NotImplemented
@@ -148,13 +148,17 @@ def on_classical_vals(self, **vals) -> Dict[str, 'ClassicalValT']:
         selection_name = self.selection_registers[0].name
         selection = vals[selection_name]
         target = vals['target']
-        if control:
+
+        # When target_gate == cirq.X, the action is (modulo phase) a single bitflip.
+        if control and self.target_gate == cirq.X:
             max_selection = self.selection_registers[0].dtype.iteration_length_or_zero() - 1
             target = (2 ** (max_selection - selection)) ^ target
+        # When target_gate == cirq.Z, the action is only in the phase.
+
         return {control_name: control, selection_name: selection, 'target': target}
 
     def basis_state_phase(self, **vals) -> Union[complex, None]:
-        if self.target_gate != cirq.X:
+        if self.target_gate != cirq.X and self.target_gate != cirq.Z:
             return None
         if len(self.control_registers) > 1 or len(self.selection_registers) > 1:
             return None
@@ -165,7 +169,10 @@ def basis_state_phase(self, **vals) -> Union[complex, None]:
         target = vals['target']
         if control:
             max_selection = self.selection_registers[0].dtype.iteration_length_or_zero() - 1
-            num_phases = (target >> (max_selection - selection + 1)).bit_count()
+            if self.target_gate == cirq.X:
+                num_phases = (target >> (max_selection - selection + 1)).bit_count()
+            else:
+                num_phases = (target >> (max_selection - selection)).bit_count()
             return 1 if (num_phases % 2) == 0 else -1
         return 1
 

qualtran/bloqs/multiplexers/selected_majorana_fermion_test.py

@@ -154,9 +154,10 @@ def test_selected_majorana_fermion_gate_make_on():
 
 
 @pytest.mark.parametrize("selection_bitsize, target_bitsize", [(2, 4), (3, 5)])
-def test_selected_majorana_fermion_classical_action(selection_bitsize, target_bitsize):
+@pytest.mark.parametrize("target_gate", [cirq.X, cirq.Z])
+def test_selected_majorana_fermion_classical_action(selection_bitsize, target_bitsize, target_gate):
     gate = SelectedMajoranaFermion(
-        Register('selection', BQUInt(selection_bitsize, target_bitsize)), target_gate=cirq.X
+        Register('selection', BQUInt(selection_bitsize, target_bitsize)), target_gate=target_gate
     )
     assert_consistent_phased_classical_action(
         gate, selection=range(target_bitsize), target=range(2**target_bitsize), control=range(2)