Fix mutable_pauli_string.inplace_after() and inplace_before()

cirq-core/cirq/ops/pauli_string.py

@@ -53,7 +53,6 @@
     identity,
     op_tree,
     pauli_gates,
-    pauli_interaction_gate,
     raw_types,
 )
 
@@ -968,53 +967,13 @@ def conjugated_by(self, clifford: cirq.OP_TREE) -> PauliString:
             # Decompose P = Pc⊗R, where Pc acts on the same qubits as C, R acts on the remaining.
             # Then the conjugation = (C^{-1}⊗I·Pc⊗R·C⊗I) = (C^{-1}·Pc·C)⊗R.
 
-            # Isolate R
-            remain: cirq.PauliString = PauliString(
+            # Conjugation on the qubits of op
+            conjugated = _calc_conjugation(ps, op)
+            # The pauli string on the remaining qubits
+            remain: PauliString = PauliString(
                 *(pauli(q) for q in all_qubits - set(op.qubits) if (pauli := ps.get(q)) is not None)
             )
-
-            # Initialize the conjugation of Pc.
-            conjugated: cirq.DensePauliString = (
-                dense_pauli_string.DensePauliString(pauli_mask=[identity.I for _ in op.qubits])
-                * ps.coefficient
-            )
-
-            # Calculate the conjugation via CliffordGate's clifford_tableau.
-            # Note the clifford_tableau in CliffordGate represents C·P·C^-1 instead of C^-1·P·C.
-            # So we take the inverse of the tableau to match the definition of the conjugation here.
-            gate_in_clifford: cirq.CliffordGate
-            if isinstance(op.gate, clifford_gate.CliffordGate):
-                gate_in_clifford = op.gate
-            else:
-                # Convert the clifford gate to CliffordGate type.
-                gate_in_clifford = clifford_gate.CliffordGate.from_op_list([op], op.qubits)
-            tableau = gate_in_clifford.clifford_tableau.inverse()
-
-            # Calculate the conjugation by `op` via mutiplying the conjugation of each Pauli:
-            #   C^{-1}·(P_1⊗...⊗P_n)·C
-            # = C^{-1}·(P_1⊗I) ...·(P_n⊗I)·C
-            # = (C^{-1}(P_1⊗I)C)·...·(C^{-1}(P_n⊗I)C)
-            # For the Pauli on the kth qubit P_k. The conjugation is calculated as following.
-            #   Puali X_k's conjugation is from the destabilzer table;
-            #   Puali Z_k's conjugation is from the stabilzer table;
-            #   Puali Y_k's conjugation is calcluated according to Y = iXZ. E.g., for the kth qubit,
-            #     C^{-1}·Y_k⊗I·C = C^{-1}·(iX_k⊗I·Z_k⊗I)·C = i (C^{-1}·X_k⊗I·C)·(C^{-1}·Z_k⊗I·C).
-            for qid, qubit in enumerate(op.qubits):
-                pauli = ps.get(qubit)
-                match pauli:
-                    case None:
-                        continue
-                    case pauli_gates.X:
-                        conjugated *= tableau.destabilizers()[qid]
-                    case pauli_gates.Z:
-                        conjugated *= tableau.stabilizers()[qid]
-                    case pauli_gates.Y:
-                        conjugated *= (
-                            1j
-                            * tableau.destabilizers()[qid]  # conj X first
-                            * tableau.stabilizers()[qid]  # then conj Z
-                        )
-            ps = remain * conjugated.on(*op.qubits)
+            ps = remain * conjugated
         return ps
 
     def after(self, ops: cirq.OP_TREE) -> cirq.PauliString:
@@ -1371,7 +1330,18 @@ def inplace_before(self, ops: cirq.OP_TREE) -> cirq.MutablePauliString:
         Returns:
             The mutable pauli string that was mutated.
         """
-        return self.inplace_after(protocols.inverse(ops))
+        # An inplace impl of PauliString.conjugated_by().
+        flattened_ops = list(op_tree.flatten_to_ops(ops))
+        for op in flattened_ops[::-1]:
+            conjugated = _calc_conjugation(self.frozen(), op)
+            self.coefficient = conjugated.coefficient
+            for q in op.qubits:
+                new_pauli_int = PAULI_GATE_LIKE_TO_INDEX_MAP[conjugated.get(q) or 0]
+                if new_pauli_int == 0:
+                    self.pauli_int_dict.pop(cast(TKey, q), None)
+                else:
+                    self.pauli_int_dict[cast(TKey, q)] = new_pauli_int
+        return self
 
     def inplace_after(self, ops: cirq.OP_TREE) -> cirq.MutablePauliString:
         r"""Propagates the pauli string from before to after a Clifford effect.
@@ -1391,43 +1361,7 @@ def inplace_after(self, ops: cirq.OP_TREE) -> cirq.MutablePauliString:
             NotImplementedError: If any ops decompose into an unsupported
                 Clifford gate.
         """
-        for clifford in op_tree.flatten_to_ops(ops):
-            for op in _decompose_into_cliffords(clifford):
-                ps = [self.pauli_int_dict.pop(cast(TKey, q), 0) for q in op.qubits]
-                if not any(ps):
-                    continue
-                gate = op.gate
-
-                if isinstance(gate, clifford_gate.SingleQubitCliffordGate):
-                    out = gate.pauli_tuple(_INT_TO_PAULI[ps[0] - 1])
-                    if out[1]:
-                        self.coefficient *= -1
-                    self.pauli_int_dict[cast(TKey, op.qubits[0])] = PAULI_GATE_LIKE_TO_INDEX_MAP[
-                        out[0]
-                    ]
-
-                elif isinstance(gate, pauli_interaction_gate.PauliInteractionGate):
-                    q0, q1 = op.qubits
-                    p0 = _INT_TO_PAULI_OR_IDENTITY[ps[0]]
-                    p1 = _INT_TO_PAULI_OR_IDENTITY[ps[1]]
-
-                    # Kick across Paulis that anti-commute with the controls.
-                    kickback_0_to_1 = not protocols.commutes(p0, gate.pauli0)
-                    kickback_1_to_0 = not protocols.commutes(p1, gate.pauli1)
-                    kick0 = gate.pauli1 if kickback_0_to_1 else identity.I
-                    kick1 = gate.pauli0 if kickback_1_to_0 else identity.I
-                    self.__imul__({q0: p0, q1: kick0})
-                    self.__imul__({q0: kick1, q1: p1})
-
-                    # Decompose inverted controls into single-qubit operations.
-                    if gate.invert0:
-                        self.inplace_after(gate.pauli1(q1))
-                    if gate.invert1:
-                        self.inplace_after(gate.pauli0(q0))
-
-                else:  # pragma: no cover
-                    raise NotImplementedError(f"Unrecognized decomposed Clifford: {op!r}")
-        return self
+        return self.inplace_before(protocols.inverse(ops))
 
     def _imul_helper(self, other: cirq.PAULI_STRING_LIKE, sign: int):
         """Left-multiplies or right-multiplies by a PAULI_STRING_LIKE.
@@ -1594,35 +1528,6 @@ def __repr__(self) -> str:
         return f'{self.frozen()!r}.mutable_copy()'
 
 
-def _decompose_into_cliffords(op: cirq.Operation) -> list[cirq.Operation]:
-    # An operation that can be ignored?
-    if isinstance(op.gate, global_phase_op.GlobalPhaseGate):
-        return []
-
-    # Already a known Clifford?
-    if isinstance(
-        op.gate,
-        (clifford_gate.SingleQubitCliffordGate, pauli_interaction_gate.PauliInteractionGate),
-    ):
-        return [op]
-
-    # Specifies a decomposition into Cliffords?
-    v = getattr(op, '_decompose_into_clifford_', None)
-    if v is not None:
-        result = v()
-        if result is not None and result is not NotImplemented:
-            return list(op_tree.flatten_to_ops(result))
-
-    # Specifies a decomposition that happens to contain only Cliffords?
-    decomposed = protocols.decompose_once(op, None)
-    if decomposed is not None:
-        return [out for sub_op in decomposed for out in _decompose_into_cliffords(sub_op)]
-
-    raise TypeError(  # pragma: no cover
-        f'Operation is not a known Clifford and did not decompose into known Cliffords: {op!r}'
-    )
-
-
 # Mypy has extreme difficulty with these constants for some reason.
 _i = cast(identity.IdentityGate, identity.I)  # type: ignore
 _x = cast(pauli_gates.Pauli, pauli_gates.X)  # type: ignore
@@ -1667,3 +1572,52 @@ def _pauli_like_to_pauli_int(key: Any, pauli_gate_like: PAULI_GATE_LIKE):
             f"{set(PAULI_GATE_LIKE_TO_INDEX_MAP.keys())!r}"
         )
     return pauli_int
+
+
+def _calc_conjugation(ps: cirq.PauliString, clifford_op: cirq.Operation) -> cirq.PauliString:
+    """Computes the conjugation of a Pauli string by a single Clifford operation.
+
+    It computes $C^-1 P C$ where P is the Pauli string `ps` and C is the `clifford_op`.
+    """
+
+    # Initialize the conjugation of the pauli string.
+    conjugated = dense_pauli_string.DensePauliString('I' * len(clifford_op.qubits)) * ps.coefficient
+
+    # Calculate the conjugation via CliffordGate's clifford_tableau.
+    # Note the clifford_tableau in CliffordGate represents C·P·C^-1 instead of C^-1·P·C.
+    # So we take the inverse of the tableau to match the definition of the conjugation here.
+    if isinstance(clifford_op.gate, clifford_gate.CliffordGate):
+        gate_in_clifford = clifford_op.gate
+    else:
+        # Convert the clifford gate to CliffordGate type.
+        gate_in_clifford = clifford_gate.CliffordGate.from_op_list(
+            [clifford_op], clifford_op.qubits
+        )
+    tableau = gate_in_clifford.clifford_tableau.inverse()
+
+    # Calculate the conjugation by `clifford_op` via mutiplying the conjugation of each Pauli:
+    #   C^{-1}·(P_1⊗...⊗P_n)·C
+    # = C^{-1}·(P_1⊗I) ...·(P_n⊗I)·C
+    # = (C^{-1}(P_1⊗I)C)·...·(C^{-1}(P_n⊗I)C)
+    # For the Pauli on the kth qubit P_k. The conjugation is calculated as following.
+    #   Pauli X_k's conjugation is from the destabilizer table;
+    #   Pauli Z_k's conjugation is from the stabilizer table;
+    #   Pauli Y_k's conjugation is calculated according to Y = iXZ. E.g., for the kth qubit,
+    #     C^{-1}·Y_k⊗I·C = C^{-1}·(iX_k⊗I·Z_k⊗I)·C = i (C^{-1}·X_k⊗I·C)·(C^{-1}·Z_k⊗I·C).
+    for qid, qubit in enumerate(clifford_op.qubits):
+        pauli = ps.get(qubit)
+        match pauli:
+            case None:
+                continue
+            case pauli_gates.X:
+                conjugated *= tableau.destabilizers()[qid]
+            case pauli_gates.Z:
+                conjugated *= tableau.stabilizers()[qid]
+            case pauli_gates.Y:
+                conjugated *= (
+                    1j
+                    * tableau.destabilizers()[qid]  # conj X first
+                    * tableau.stabilizers()[qid]  # then conj Z
+                )
+
+    return conjugated.on(*clifford_op.qubits)

cirq-core/cirq/ops/pauli_string_test.py

@@ -1714,6 +1714,9 @@ def with_qubits(self, *new_qubits):
         def _decompose_(self):
             return []
 
+        def __pow__(self, power):
+            return []
+
     # No-ops
     p2 = p.inplace_after(cirq.global_phase_operation(1j))
     assert p2 is p and p == cirq.X(a)
@@ -1825,6 +1828,14 @@ def _decompose_(self):
     assert p2 is p and p == cirq.X(a) * cirq.Y(b)
 
 
+def test_mps_inplace_after_clifford_gate_type():
+    q = cirq.LineQubit(0)
+
+    mps = cirq.MutablePauliString(cirq.X(q))
+    mps2 = mps.inplace_after(cirq.CliffordGate.from_op_list([cirq.H(q)], [q]).on(q))
+    assert mps2 is mps and mps == cirq.Z(q)
+
+
 def test_after_before_vs_conjugate_by():
     a, b, c = cirq.LineQubit.range(3)
     p = cirq.X(a) * cirq.Y(b) * cirq.Z(c)