Support ActOnStabilizerCHFormArgs in the common gates and testing (#3203)

* Add ActOnStabilizerCHFormArgs and related logic

* Some comment updates

* Fix _update_sum call

* Update parameterized condition

* Undo already split out change to stabilizerstatechfrom

* revert split out change

* Merge and clean

* Improve error handling and documentation

* Add comments with reference to the relevant sections of the paper

* Fix the year in copyright

* Use H instead of HPowGate()

* Address comments

* Hide traceback and correct copyright year

Author: smitsanghavi

cirq/ops/common_gates.py

@@ -103,6 +103,17 @@ def _act_on_(self, args: Any):
                 tableau.xs[:, q] ^= tableau.zs[:, q]
             return True
 
+        if isinstance(args, clifford.ActOnStabilizerCHFormArgs):
+            if protocols.is_parameterized(self) or self.exponent % 0.5 != 0:
+                return NotImplemented
+            assert all(
+                gate._act_on_(args) for gate in  # type: ignore
+                [H, ZPowGate(exponent=self._exponent), H])
+            # Adjust the global phase based on the global_shift parameter.
+            args.state.omega *= np.exp(1j * np.pi * self.global_shift *
+                                       self.exponent)
+            return True
+
         return NotImplemented
 
     def in_su2(self) -> 'XPowGate':
@@ -322,6 +333,30 @@ def _act_on_(self, args: Any):
                                                         tableau.xs[:, q].copy())
             return True
 
+        if isinstance(args, clifford.ActOnStabilizerCHFormArgs):
+            if protocols.is_parameterized(self) or self.exponent % 0.5 != 0:
+                return NotImplemented
+            effective_exponent = self._exponent % 2
+            state = args.state
+            if effective_exponent == 0.5:
+                assert all(
+                    gate._act_on_(args)  # type: ignore
+                    for gate in [ZPowGate(), H])
+                state.omega *= (1 + 1j) / (2**0.5)  # type: ignore
+            elif effective_exponent == 1:
+                assert all(
+                    gate._act_on_(args) for gate in  # type: ignore
+                    [ZPowGate(), H, ZPowGate(), H])
+                state.omega *= 1j  # type: ignore
+            elif effective_exponent == 1.5:
+                assert all(
+                    gate._act_on_(args)  # type: ignore
+                    for gate in [H, ZPowGate()])
+                state.omega *= (1 - 1j) / (2**0.5)  # type: ignore
+            # Adjust the global phase based on the global_shift parameter.
+            args.state.omega *= np.exp(1j * np.pi * self.global_shift *
+                                       self.exponent)
+            return True
         return NotImplemented
 
     def in_su2(self) -> 'YPowGate':
@@ -490,6 +525,22 @@ def _act_on_(self, args: Any):
                 tableau.zs[:, q] ^= tableau.xs[:, q]
             return True
 
+        if isinstance(args, clifford.ActOnStabilizerCHFormArgs):
+            if protocols.is_parameterized(self) or self.exponent % 0.5 != 0:
+                return NotImplemented
+            q = args.axes[0]
+            effective_exponent = self._exponent % 2
+            state = args.state
+            for _ in range(int(effective_exponent * 2)):
+                # Prescription for S left multiplication.
+                # Reference: https://arxiv.org/abs/1808.00128 Proposition 4 end
+                state.M[q, :] ^= state.G[q, :]
+                state.gamma[q] = (state.gamma[q] - 1) % 4
+            # Adjust the global phase based on the global_shift parameter.
+            args.state.omega *= np.exp(1j * np.pi * self.global_shift *
+                                       self.exponent)
+            return True
+
         return NotImplemented
 
     def _decompose_into_clifford_with_qubits_(self, qubits):
@@ -756,18 +807,43 @@ def _act_on_(self, args: Any):
         from cirq.sim import clifford
 
         if isinstance(args, clifford.ActOnCliffordTableauArgs):
-            if protocols.is_parameterized(self) or self.exponent % 0.5 != 0:
+            if protocols.is_parameterized(self) or self.exponent % 1 != 0:
                 return NotImplemented
             tableau = args.tableau
             q = args.axes[0]
-            if self._exponent % 1 != 0:
-                return NotImplemented
             if self._exponent % 2 == 1:
                 (tableau.xs[:, q], tableau.zs[:, q]) = (tableau.zs[:, q].copy(),
                                                         tableau.xs[:, q].copy())
                 tableau.rs[:] ^= (tableau.xs[:, q] & tableau.zs[:, q])
             return True
 
+        if isinstance(args, clifford.ActOnStabilizerCHFormArgs):
+            if protocols.is_parameterized(self) or self.exponent % 1 != 0:
+                return NotImplemented
+            q = args.axes[0]
+            state = args.state
+            if self._exponent % 2 == 1:
+                # Prescription for H left multiplication
+                # Reference: https://arxiv.org/abs/1808.00128
+                # Equations 48, 49 and Proposition 4
+                t = state.s ^ (state.G[q, :] & state.v)
+                u = state.s ^ (state.F[q, :] &
+                               (~state.v)) ^ (state.M[q, :] & state.v)
+
+                alpha = sum(state.G[q, :] & (~state.v) & state.s) % 2
+                beta = sum(state.M[q, :] & (~state.v) & state.s)
+                beta += sum(state.F[q, :] & state.v & state.M[q, :])
+                beta += sum(state.F[q, :] & state.v & state.s)
+                beta %= 2
+
+                delta = (state.gamma[q] + 2 * (alpha + beta)) % 4
+
+                state.update_sum(t, u, delta=delta, alpha=alpha)
+            # Adjust the global phase based on the global_shift parameter.
+            args.state.omega *= np.exp(1j * np.pi * self.global_shift *
+                                       self.exponent)
+            return True
+
         return NotImplemented
 
     def _decompose_(self, qubits):
@@ -900,6 +976,22 @@ def _act_on_(self, args: Any):
                 tableau.rs[:] ^= (tableau.xs[:, q2] & tableau.zs[:, q2])
             return True
 
+        if isinstance(args, clifford.ActOnStabilizerCHFormArgs):
+            if protocols.is_parameterized(self) or self.exponent % 1 != 0:
+                return NotImplemented
+            q1 = args.axes[0]
+            q2 = args.axes[1]
+            state = args.state
+            if self._exponent % 2 == 1:
+                # Prescription for CZ left multiplication.
+                # Reference: https://arxiv.org/abs/1808.00128 Proposition 4 end
+                state.M[q1, :] ^= state.G[q2, :]
+                state.M[q2, :] ^= state.G[q1, :]
+            # Adjust the global phase based on the global_shift parameter.
+            args.state.omega *= np.exp(1j * np.pi * self.global_shift *
+                                       self.exponent)
+            return True
+
         return NotImplemented
 
     def _pauli_expansion_(self) -> value.LinearDict[str]:
@@ -1098,6 +1190,26 @@ def _act_on_(self, args: Any):
                 tableau.zs[:, q1] ^= tableau.zs[:, q2]
             return True
 
+        if isinstance(args, clifford.ActOnStabilizerCHFormArgs):
+            if protocols.is_parameterized(self) or self.exponent % 1 != 0:
+                return NotImplemented
+            q1 = args.axes[0]
+            q2 = args.axes[1]
+            state = args.state
+            if self._exponent % 2 == 1:
+                # Prescription for CX left multiplication.
+                # Reference: https://arxiv.org/abs/1808.00128 Proposition 4 end
+                state.gamma[q1] = (
+                    state.gamma[q1] + state.gamma[q2] + 2 *
+                    (sum(state.M[q1, :] & state.F[q2, :]) % 2)) % 4
+                state.G[q2, :] ^= state.G[q1, :]
+                state.F[q1, :] ^= state.F[q2, :]
+                state.M[q1, :] ^= state.M[q2, :]
+            # Adjust the global phase based on the global_shift parameter.
+            args.state.omega *= np.exp(1j * np.pi * self.global_shift *
+                                       self.exponent)
+            return True
+
         return NotImplemented
 
     def _pauli_expansion_(self) -> value.LinearDict[str]:

cirq/ops/common_gates_test.py

@@ -290,7 +290,7 @@ def test_h_str():
     assert str(cirq.H**0.5) == 'H^0.5'
 
 
-def test_x_act_on():
+def test_x_act_on_tableau():
     with pytest.raises(TypeError, match="Failed to act"):
         cirq.act_on(cirq.X, object())
     original_tableau = cirq.CliffordTableau(num_qubits=5, initial_state=31)
@@ -326,14 +326,21 @@ def test_x_act_on():
         cirq.act_on(cirq.X**foo, args)
 
 
-class PhaserGate(cirq.SingleQubitGate):
+class iZGate(cirq.SingleQubitGate):
     """Equivalent to an iZ gate without _act_on_ defined on it."""
 
     def _unitary_(self):
         return np.array([[1j, 0], [0, -1j]])
 
 
-def test_y_act_on():
+class MinusOnePhaseGate(cirq.SingleQubitGate):
+    """Equivalent to a -1 global phase without _act_on_ defined on it."""
+
+    def _unitary_(self):
+        return np.array([[-1, 0], [0, -1]])
+
+
+def test_y_act_on_tableau():
     with pytest.raises(TypeError, match="Failed to act"):
         cirq.act_on(cirq.Y, object())
     original_tableau = cirq.CliffordTableau(num_qubits=5, initial_state=31)
@@ -348,18 +355,18 @@ def test_y_act_on():
 
     cirq.act_on(cirq.Y**0.5, args, allow_decompose=False)
     cirq.act_on(cirq.Y**0.5, args, allow_decompose=False)
-    cirq.act_on(PhaserGate(), args)
+    cirq.act_on(iZGate(), args)
     assert args.log_of_measurement_results == {}
     assert args.tableau == flipped_tableau
 
     cirq.act_on(cirq.Y, args, allow_decompose=False)
-    cirq.act_on(PhaserGate(), args, allow_decompose=True)
+    cirq.act_on(iZGate(), args, allow_decompose=True)
     assert args.log_of_measurement_results == {}
     assert args.tableau == original_tableau
 
     cirq.act_on(cirq.Y**3.5, args, allow_decompose=False)
     cirq.act_on(cirq.Y**3.5, args, allow_decompose=False)
-    cirq.act_on(PhaserGate(), args)
+    cirq.act_on(iZGate(), args)
     assert args.log_of_measurement_results == {}
     assert args.tableau == flipped_tableau
 
@@ -372,9 +379,11 @@ def test_y_act_on():
         cirq.act_on(cirq.Y**foo, args)
 
 
-def test_z_h_act_on():
+def test_z_h_act_on_tableau():
     with pytest.raises(TypeError, match="Failed to act"):
-        cirq.act_on(cirq.Y, object())
+        cirq.act_on(cirq.Z, object())
+    with pytest.raises(TypeError, match="Failed to act"):
+        cirq.act_on(cirq.H, object())
     original_tableau = cirq.CliffordTableau(num_qubits=5, initial_state=31)
     flipped_tableau = cirq.CliffordTableau(num_qubits=5, initial_state=23)
 
@@ -417,18 +426,16 @@ def test_z_h_act_on():
     with pytest.raises(TypeError, match="Failed to act action on state"):
         cirq.act_on(cirq.Z**foo, args)
 
-    foo = sympy.Symbol('foo')
     with pytest.raises(TypeError, match="Failed to act action on state"):
         cirq.act_on(cirq.H**foo, args)
 
-    foo = sympy.Symbol('foo')
     with pytest.raises(TypeError, match="Failed to act action on state"):
         cirq.act_on(cirq.H**1.5, args)
 
 
-def test_cx_act_on():
+def test_cx_act_on_tableau():
     with pytest.raises(TypeError, match="Failed to act"):
-        cirq.act_on(cirq.Y, object())
+        cirq.act_on(cirq.CX, object())
     original_tableau = cirq.CliffordTableau(num_qubits=5, initial_state=31)
 
     args = cirq.ActOnCliffordTableauArgs(
@@ -471,7 +478,7 @@ def test_cx_act_on():
         cirq.act_on(cirq.CX**1.5, args)
 
 
-def test_cz_act_on():
+def test_cz_act_on_tableau():
     with pytest.raises(TypeError, match="Failed to act"):
         cirq.act_on(cirq.Y, object())
     original_tableau = cirq.CliffordTableau(num_qubits=5, initial_state=31)
@@ -516,38 +523,102 @@ def test_cz_act_on():
         cirq.act_on(cirq.CZ**1.5, args)
 
 
+foo = sympy.Symbol('foo')
+
+
+@pytest.mark.parametrize('input_gate_sequence, outcome', [
+    ([cirq.X**foo], 'Error'),
+    ([cirq.X**0.25], 'Error'),
+    ([cirq.X**4], 'Original'),
+    ([cirq.X**0.5, cirq.X**0.5], 'Flipped'),
+    ([cirq.X], 'Flipped'),
+    ([cirq.X**3.5, cirq.X**3.5], 'Flipped'),
+    ([cirq.Y**foo], 'Error'),
+    ([cirq.Y**0.25], 'Error'),
+    ([cirq.Y**4], 'Original'),
+    ([cirq.Y**0.5, cirq.Y**0.5, iZGate()], 'Flipped'),
+    ([cirq.Y, iZGate()], 'Flipped'),
+    ([cirq.Y**3.5, cirq.Y**3.5, iZGate()], 'Flipped'),
+    ([cirq.Z**foo], 'Error'),
+    ([cirq.H**foo], 'Error'),
+    ([cirq.H**1.5], 'Error'),
+    ([cirq.Z**4], 'Original'),
+    ([cirq.H**4], 'Original'),
+    ([cirq.H, cirq.S, cirq.S, cirq.H], 'Flipped'),
+    ([cirq.H, cirq.Z, cirq.H], 'Flipped'),
+    ([cirq.H, cirq.Z**3.5, cirq.Z**3.5, cirq.H], 'Flipped'),
+    ([cirq.CX**foo], 'Error'),
+    ([cirq.CX**1.5], 'Error'),
+    ([cirq.CX**4], 'Original'),
+    ([cirq.CX], 'Flipped'),
+    ([cirq.CZ**foo], 'Error'),
+    ([cirq.CZ**1.5], 'Error'),
+    ([cirq.CZ**4], 'Original'),
+    ([cirq.CZ, MinusOnePhaseGate()], 'Original'),
+])
+def test_act_on_ch_form(input_gate_sequence, outcome):
+    original_state = cirq.StabilizerStateChForm(num_qubits=5, initial_state=31)
+    num_qubits = cirq.num_qubits(input_gate_sequence[0])
+    if num_qubits == 1:
+        axes = [1]
+    else:
+        assert num_qubits == 2
+        axes = [0, 1]
+    args = cirq.ActOnStabilizerCHFormArgs(state=original_state.copy(),
+                                          axes=axes)
+
+    flipped_state = cirq.StabilizerStateChForm(num_qubits=5, initial_state=23)
+
+    if outcome == 'Error':
+        with pytest.raises(TypeError, match="Failed to act action on state"):
+            for input_gate in input_gate_sequence:
+                cirq.act_on(input_gate, args)
+        return
+
+    for input_gate in input_gate_sequence:
+        cirq.act_on(input_gate, args)
+
+    if outcome == 'Original':
+        np.testing.assert_allclose(args.state.state_vector(),
+                                   original_state.state_vector())
+
+    if outcome == 'Flipped':
+        np.testing.assert_allclose(args.state.state_vector(),
+                                   flipped_state.state_vector())
+
+
 @pytest.mark.parametrize(
-    'input_gate',
+    'input_gate, assert_implemented',
     [
-        cirq.X,
-        cirq.Y,
-        cirq.Z,
-        cirq.X**0.5,
-        cirq.Y**0.5,
-        cirq.Z**0.5,
-        cirq.X**3.5,
-        cirq.Y**3.5,
-        cirq.Z**3.5,
-        cirq.X**4,
-        cirq.Y**4,
-        cirq.Z**4,
-        cirq.H,
-        cirq.CX,
-        cirq.CZ,
-        cirq.H**4,
-        cirq.CX**4,
-        cirq.CZ**4,
-        # Gates not supported by CliffordTableau should not fail too.
-        cirq.X**0.25,
-        cirq.Y**0.25,
-        cirq.Z**0.25,
-        cirq.H**0.5,
-        cirq.CX**0.5,
-        cirq.CZ**0.5
+        (cirq.X, True),
+        (cirq.Y, True),
+        (cirq.Z, True),
+        (cirq.X**0.5, True),
+        (cirq.Y**0.5, True),
+        (cirq.Z**0.5, True),
+        (cirq.X**3.5, True),
+        (cirq.Y**3.5, True),
+        (cirq.Z**3.5, True),
+        (cirq.X**4, True),
+        (cirq.Y**4, True),
+        (cirq.Z**4, True),
+        (cirq.H, True),
+        (cirq.CX, True),
+        (cirq.CZ, True),
+        (cirq.H**4, True),
+        (cirq.CX**4, True),
+        (cirq.CZ**4, True),
+        # Unsupported gates should not fail too.
+        (cirq.X**0.25, False),
+        (cirq.Y**0.25, False),
+        (cirq.Z**0.25, False),
+        (cirq.H**0.5, False),
+        (cirq.CX**0.5, False),
+        (cirq.CZ**0.5, False),
     ])
-def test_act_on_clifford_tableau(input_gate):
-    cirq.testing.assert_act_on_clifford_tableau_effect_matches_unitary(
-        input_gate)
+def test_act_on_consistency(input_gate, assert_implemented):
+    cirq.testing.assert_all_implemented_act_on_effects_match_unitary(
+        input_gate, assert_implemented, assert_implemented)
 
 
 def test_runtime_types_of_rot_gates():