Move all Pauli measurements to the front before Pauli presimulation (#…

…209)

This commit adds a method `move_pauli_measurements_to_the_front` that applies
exchange rules to move all Pauli measurements to the front of a
pattern.  This allows all the Pauli measurements to be presimulated,
even if their initial domains depend on other measurements.

On random circuits of 10 qubits (depth 10), number of remaining nodes
after Pauli presimulation is reduced by a factor of 4 in average, and
the time spent on state vector simulation is reduced by a factor
75 (including the time for Pauli presimulation).

graphix/pattern.py

@@ -1416,12 +1416,30 @@ def run_pattern(self, backend, **kwargs):
         result = exe.run()
         return result
 
-    def perform_pauli_measurements(self, leave_input=False, use_rustworkx=False):
+    def perform_pauli_measurements(
+        self, leave_input: bool = False, use_rustworkx: bool = False, ignore_pauli_with_deps: bool = False
+    ) -> None:
         """Perform Pauli measurements in the pattern using efficient stabilizer simulator.
 
+        Parameters
+        ----------
+        leave_input : bool
+            Optional (`False` by default).
+            If `True`, measurements on input nodes are preserved as-is in the pattern.
+        use_rustworkx : bool
+            Optional (`False` by default).
+            If `True`, `rustworkx` is used for fast graph processing.
+            If `False`, `networkx` is used.
+        ignore_pauli_with_deps : bool
+            Optional (`False` by default).
+            If `True`, Pauli measurements with domains depending on other measures are preserved as-is in the pattern.
+            If `False`, all Pauli measurements are preprocessed. Formally, measurements are swapped so that all Pauli measurements are applied first, and domains are updated accordingly.
+
         .. seealso:: :func:`measure_pauli`
 
         """
+        if not ignore_pauli_with_deps:
+            self.move_pauli_measurements_to_the_front()
         measure_pauli(self, leave_input, copy=False, use_rustworkx=use_rustworkx)
 
     def draw_graph(
@@ -1521,14 +1539,87 @@ def to_qasm3(self, filename):
     def copy(self) -> Pattern:
         """Return a copy of the pattern."""
         result = self.__new__(self.__class__)
-        result.__seq = [cmd.model_copy() for cmd in self.__seq]
+        result.__seq = [cmd.model_copy(deep=True) for cmd in self.__seq]
         result.__input_nodes = self.__input_nodes.copy()
         result.__output_nodes = self.__output_nodes.copy()
         result.__n_node = self.__n_node
         result._pauli_preprocessed = self._pauli_preprocessed
         result.results = self.results.copy()
         return result
 
+    def move_pauli_measurements_to_the_front(self, leave_nodes: set[int] | None = None) -> None:
+        """Move all the Pauli measurements to the front of the sequence (except nodes in `leave_nodes`)."""
+        if leave_nodes is None:
+            leave_nodes = set()
+        self.standardize()
+        pauli_nodes = {}
+        shift_domains = {}
+
+        def expand_domain(domain: set[int]) -> None:
+            for node in domain & shift_domains.keys():
+                domain ^= shift_domains[node]
+
+        for cmd in self:
+            if cmd.kind == CommandKind.X or cmd.kind == CommandKind.Z:
+                expand_domain(cmd.domain)
+            if cmd.kind == CommandKind.M:
+                expand_domain(cmd.s_domain)
+                expand_domain(cmd.t_domain)
+                pm = PauliMeasurement.try_from(
+                    cmd.plane, cmd.angle
+                )  # None returned if the measurement is not in Pauli basis
+                if pm is not None and cmd.node not in leave_nodes:
+                    if pm.axis == Axis.X:
+                        # M^X X^s Z^t = M^{XY,0} X^s Z^t
+                        #             = M^{XY,(-1)^s·0+tπ}
+                        #             = S^t M^X
+                        # M^{-X} X^s Z^t = M^{XY,π} X^s Z^t
+                        #                = M^{XY,(-1)^s·π+tπ}
+                        #                = S^t M^{-X}
+                        shift_domains[cmd.node] = cmd.t_domain
+                    elif pm.axis == Axis.Y:
+                        # M^Y X^s Z^t = M^{XY,π/2} X^s Z^t
+                        #             = M^{XY,(-1)^s·π/2+tπ}
+                        #             = M^{XY,π/2+(s+t)π}      (since -π/2 = π/2 - π ≡ π/2 + π (mod 2π))
+                        #             = S^{s+t} M^Y
+                        # M^{-Y} X^s Z^t = M^{XY,-π/2} X^s Z^t
+                        #                = M^{XY,(-1)^s·(-π/2)+tπ}
+                        #                = M^{XY,-π/2+(s+t)π}  (since π/2 = -π/2 + π)
+                        #                = S^{s+t} M^{-Y}
+                        shift_domains[cmd.node] = cmd.s_domain ^ cmd.t_domain
+                    elif pm.axis == Axis.Z:
+                        # M^Z X^s Z^t = M^{XZ,0} X^s Z^t
+                        #             = M^{XZ,(-1)^t((-1)^s·0+sπ)}
+                        #             = M^{XZ,(-1)^t·sπ}
+                        #             = M^{XZ,sπ}              (since (-1)^t·π ≡ π (mod 2π))
+                        #             = S^s M^Z
+                        # M^{-Z} X^s Z^t = M^{XZ,π} X^s Z^t
+                        #                = M^{XZ,(-1)^t((-1)^s·π+sπ)}
+                        #                = M^{XZ,(s+1)π}
+                        #                = S^s M^{-Z}
+                        shift_domains[cmd.node] = cmd.s_domain
+                    else:
+                        typing_extensions.assert_never(pm.axis)
+                    cmd.s_domain = set()
+                    cmd.t_domain = set()
+                    pauli_nodes[cmd.node] = cmd
+
+        # Create a new sequence with all Pauli nodes to the front
+        new_seq = []
+        pauli_nodes_inserted = False
+        for cmd in self:
+            if cmd.kind == CommandKind.M:
+                if cmd.node not in pauli_nodes:
+                    if not pauli_nodes_inserted:
+                        new_seq.extend(pauli_nodes.values())
+                        pauli_nodes_inserted = True
+                    new_seq.append(cmd)
+            else:
+                new_seq.append(cmd)
+        if not pauli_nodes_inserted:
+            new_seq.extend(pauli_nodes.values())
+        self.__seq = new_seq
+
 
 class CommandNode:
     """A node decorated with a distributed command sequence.

tests/test_pattern.py

@@ -11,9 +11,9 @@
 from numpy.random import PCG64, Generator
 
 from graphix import clifford
-from graphix.command import C, E, M, N, X, Z
+from graphix.command import C, CommandKind, E, M, N, X, Z
 from graphix.pattern import CommandNode, Pattern, shift_outcomes
-from graphix.pauli import Plane
+from graphix.pauli import PauliMeasurement, Plane
 from graphix.random_objects import rand_circuit, rand_gate
 from graphix.sim.density_matrix import DensityMatrix
 from graphix.simulator import PatternSimulator
@@ -172,6 +172,23 @@ def test_pauli_measurement_random_circuit(
         state_mbqc = pattern.simulate_pattern(backend, rng=rng)
         assert compare_backend_result_with_statevec(backend, state_mbqc, state) == pytest.approx(1)
 
+    @pytest.mark.parametrize("jumps", range(1, 11))
+    @pytest.mark.parametrize("ignore_pauli_with_deps", (False, True))
+    def test_pauli_measurement_random_circuit_all_paulis(
+        self, fx_bg: PCG64, jumps: int, ignore_pauli_with_deps: bool, use_rustworkx: bool = True
+    ) -> None:
+        rng = Generator(fx_bg.jumped(jumps))
+        nqubits = 3
+        depth = 3
+        circuit = rand_circuit(nqubits, depth, rng)
+        pattern = circuit.transpile().pattern
+        pattern.standardize(method="global")
+        pattern.shift_signals(method="global")
+        pattern.perform_pauli_measurements(use_rustworkx=use_rustworkx, ignore_pauli_with_deps=ignore_pauli_with_deps)
+        assert ignore_pauli_with_deps or not any(
+            PauliMeasurement.try_from(cmd.plane, cmd.angle) for cmd in pattern if cmd.kind == CommandKind.M
+        )
+
     @pytest.mark.parametrize("plane", Plane)
     @pytest.mark.parametrize("angle", [0.0, 0.5, 1.0, 1.5])
     def test_pauli_measurement_single(self, plane: Plane, angle: float, use_rustworkx: bool = True) -> None: