Optimize calculation of PhasedXZGate unitary (#7811)

Use analytical formula for the PhasedXZGate unitary resulting in
about 6-fold decrease in calculation time.  For Eliott's test
colab this improves the speed of

- `CZGaugeTransformer.as_sweep` by 25% (30 to 23s)
- `merge_single_qubit_gates_to_phxz_symbolized` by 15% (77 to 66s)
- `merge_single_qubit_moments_to_phxz` by 40% (29 to 17s)

Also add pytest benchmarks for `PhasedXZGate` unitary and
for the transformation chain in Eliott's colab.

Partially implements #7797

Author: pavoljuhas

benchmarks/dynamical_decoupling_perf.py

@@ -0,0 +1,75 @@
+import itertools
+
+import pytest
+
+import cirq
+
+QUBITS_PER_LAYER = [
+    [cirq.GridQubit(5, 9), cirq.GridQubit(6, 9), cirq.GridQubit(5, 10), cirq.GridQubit(6, 8)],
+    [
+        cirq.GridQubit(8, 8),
+        cirq.GridQubit(6, 9),
+        cirq.GridQubit(7, 9),
+        cirq.GridQubit(5, 10),
+        cirq.GridQubit(6, 10),
+        cirq.GridQubit(7, 8),
+    ],
+    [cirq.GridQubit(6, 9), cirq.GridQubit(7, 9), cirq.GridQubit(6, 10), cirq.GridQubit(7, 8)],
+    [cirq.GridQubit(5, 9), cirq.GridQubit(6, 9), cirq.GridQubit(6, 8), cirq.GridQubit(7, 8)],
+]
+
+LAYERS_GRID_QUBIT_PAIRS = [
+    [(cirq.GridQubit(5, 9), cirq.GridQubit(5, 10)), (cirq.GridQubit(6, 8), cirq.GridQubit(6, 9))],
+    [
+        (cirq.GridQubit(5, 10), cirq.GridQubit(6, 10)),
+        (cirq.GridQubit(6, 9), cirq.GridQubit(7, 9)),
+        (cirq.GridQubit(7, 8), cirq.GridQubit(8, 8)),
+    ],
+    [(cirq.GridQubit(6, 9), cirq.GridQubit(6, 10)), (cirq.GridQubit(7, 8), cirq.GridQubit(7, 9))],
+    [(cirq.GridQubit(5, 9), cirq.GridQubit(6, 9)), (cirq.GridQubit(6, 8), cirq.GridQubit(7, 8))],
+]
+
+
+def make_fake_trotter_circuit(num_cycles: int):
+    all_qubits = sorted(set(itertools.chain.from_iterable(QUBITS_PER_LAYER)))
+    moments = []
+    for layer, qubits in zip(LAYERS_GRID_QUBIT_PAIRS, QUBITS_PER_LAYER):
+        moments.append(cirq.Moment((cirq.Y**0.3)(qubit) for qubit in qubits))
+        moments.append(cirq.Moment(cirq.CZ(*pair) for pair in layer))
+    return cirq.Circuit.from_moments(*moments) * num_cycles + cirq.Moment(cirq.M(*all_qubits))
+
+
+@pytest.mark.parametrize(
+    ["num_cycles", "num_circuits"], [(5, 10), pytest.param(50, 100, marks=pytest.mark.slow)]
+)
+@pytest.mark.benchmark(group="dynamical_decoupling", max_time=10)
+def test_dynamical_decoupling_sweep(benchmark, num_cycles: int, num_circuits: int) -> None:
+    def _f(num_cycles: int, num_circuits: int) -> cirq.Circuit:
+        circuit = make_fake_trotter_circuit(num_cycles)
+        circuit1, sweep1 = cirq.transformers.gauge_compiling.CZGaugeTransformer.as_sweep(
+            circuit, N=num_circuits
+        )
+        circuit2, sweep2 = cirq.merge_single_qubit_gates_to_phxz_symbolized(circuit1, sweep=sweep1)
+        assert len(sweep2)
+        circuit3 = cirq.add_dynamical_decoupling(circuit2)
+        return circuit3
+
+    _ = benchmark(_f, num_cycles=num_cycles, num_circuits=num_circuits)
+
+
+@pytest.mark.parametrize(
+    ["num_cycles", "num_circuits"], [(5, 10), pytest.param(50, 100, marks=pytest.mark.slow)]
+)
+@pytest.mark.benchmark(group="dynamical_decoupling", max_time=10)
+def test_dynamical_decoupling_batch(benchmark, num_cycles: int, num_circuits: int) -> None:
+    def _f(num_cycles: int, num_circuits: int) -> list[cirq.Circuit]:
+        circuit = make_fake_trotter_circuit(num_cycles)
+        circuits_batch1 = [
+            cirq.transformers.gauge_compiling.cz_gauge.CZGaugeTransformer(circuit)
+            for _ in range(num_circuits)
+        ]
+        circuits_batch2 = [cirq.merge_single_qubit_moments_to_phxz(c) for c in circuits_batch1]
+        circuits_batch3 = [cirq.add_dynamical_decoupling(c) for c in circuits_batch2]
+        return circuits_batch3
+
+    _ = benchmark(_f, num_cycles=num_cycles, num_circuits=num_circuits)

benchmarks/phased_x_z_gate_perf.py

@@ -39,3 +39,25 @@ def test_has_stabilizer_effect_for_random_gates(benchmark) -> None:
     f = lambda: any(cirq.has_stabilizer_effect(g) for g in random_phased_xz_gates)
     result = benchmark(f)
     assert not result
+
+
+@pytest.mark.benchmark(group="phased_x_z_gate")
+def test_unitary_for_cliffords(benchmark) -> None:
+    phased_xz_cliffords = [
+        g.to_phased_xz_gate() for g in cirq.SingleQubitCliffordGate.all_single_qubit_cliffords
+    ]
+    f = lambda: all(cirq.unitary(g) is not None for g in phased_xz_cliffords)
+    result = benchmark(f)
+    assert result
+
+
+@pytest.mark.benchmark(group="phased_x_z_gate")
+def test_unitary_for_random_gates(benchmark) -> None:
+    num_calls = len(cirq.SingleQubitCliffordGate.all_single_qubit_cliffords)
+    random_phased_xz_gates = [
+        cirq.PhasedXZGate(x_exponent=x, z_exponent=z, axis_phase_exponent=a)
+        for x, z, a in np.random.uniform(-1, 1, (num_calls, 3))
+    ]
+    f = lambda: all(cirq.unitary(g) is not None for g in random_phased_xz_gates)
+    result = benchmark(f)
+    assert result

cirq-core/cirq/ops/phased_x_z_gate.py

@@ -15,6 +15,7 @@
 from __future__ import annotations
 
 import functools
+import math
 import numbers
 from collections.abc import Iterator, Sequence, Set
 from typing import Any, TYPE_CHECKING
@@ -192,10 +193,24 @@ def _unitary_(self) -> np.ndarray | None:
         """See `cirq.SupportsUnitary`."""
         if self._is_parameterized_():
             return None
-        z_pre = protocols.unitary(ops.Z**-self._axis_phase_exponent)
-        x = protocols.unitary(ops.X**self._x_exponent)
-        z_post = protocols.unitary(ops.Z ** (self._axis_phase_exponent + self._z_exponent))
-        return z_post @ x @ z_pre
+
+        a = self._axis_phase_exponent
+        x = self._x_exponent
+        z = self._z_exponent
+        # evaluate unitary terms ucx = exp(1j * pi * x / 2) * cos(pi * x / 2) and
+        # usx = -1j * exp(1j * pi * x / 2) * sin(pi * x / 2) without round-off
+        # errors at half-integer x
+        if x % 1 == 0.5:
+            ucx = 0.5 + 0.5j * (-1) ** math.floor(x)
+            usx = ucx.conjugate()
+        else:
+            cpxh = 1j**x
+            ucx = cpxh * cpxh.real
+            usx = -1j * cpxh * cpxh.imag
+        u = np.array(
+            [[ucx, usx * 1j ** (-2 * a)], [usx * 1j ** (2 * (z + a)), ucx * 1j ** (2 * z)]]
+        )
+        return u
 
     def _decompose_(self, qubits: Sequence[cirq.Qid]) -> Iterator[cirq.OP_TREE]:
         q = qubits[0]

cirq-core/cirq/ops/phased_x_z_gate_test.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import itertools
 import random
 from typing import cast
 
@@ -227,6 +228,20 @@ def test_from_matrix_close_unitary(unitary: np.ndarray) -> None:
     )
 
 
+@pytest.mark.parametrize(
+    ['x_exponent', 'z_exponent', 'axis_phase_exponent'],
+    itertools.product([-0.5, 0.0, 0.5, 1.0], repeat=3),
+)
+def test_exact_unitary_at_half_integers(
+    x_exponent: float, z_exponent: float, axis_phase_exponent: float
+) -> None:
+    gate = cirq.PhasedXZGate(
+        x_exponent=x_exponent, z_exponent=z_exponent, axis_phase_exponent=axis_phase_exponent
+    )
+    u = cirq.unitary(gate)
+    np.testing.assert_equal(u, np.round(u, decimals=1))
+
+
 @pytest.mark.parametrize(
     'unitary',
     [