Optimize PhasedXZGate._has_stabilizer_effect_ a bit more (#7852)

- avoid creating a temporary canonical instance of PhasedXZGate
- avoid np.allclose and its argument-checking logic
- return early if any exponent cannot match Clifford values

Partially implements #7797

Author: pavoljuhas

cirq-core/cirq/ops/phased_x_z_gate.py

@@ -18,7 +18,7 @@
 import math
 import numbers
 from collections.abc import Iterator, Sequence, Set
-from typing import Any, TYPE_CHECKING
+from typing import Any, Final, TYPE_CHECKING
 
 import numpy as np
 import sympy
@@ -331,13 +331,34 @@ def _json_dict_(self) -> dict[str, Any]:
     def _has_stabilizer_effect_(self) -> bool:
         if not self._has_unitary_():
             return False
-        c = self._canonical()
-        actual = (c._x_exponent, c._z_exponent, c._axis_phase_exponent)
-        rounded = tuple(round(v, 2) for v in actual)  # for numerical stability.
-        return (
-            np.allclose(actual, rounded)
-            and tuple(v % 2 for v in rounded) in _clifford_as_phasedzx_params()
+        tol: Final = 1e-8
+        result = (
+            abs((x := round(self._x_exponent, 2)) - self._x_exponent) <= tol
+            and abs((z := round(self._z_exponent, 2)) - self._z_exponent) <= tol
+            and abs((a := round(self._axis_phase_exponent, 2)) - self._axis_phase_exponent) <= tol
+            and _canonical_xza_mod_2(x, z, a) in _clifford_as_phasedzx_params()
         )
+        return result
+
+
+def _canonical_xza_mod_2(
+    x_exponent: float, z_exponent: float, axis_phase_exponent: float
+) -> tuple[float, float, float]:
+    """Return canonical values of PhasedXZGate parameters modulo 2.
+
+    Optimized helper for `PhasedXZGate._has_stabilizer_effect_`.
+    """
+    # The result must be consistent with PhasedXZGate._canonical
+    x = x_exponent % 2
+    a = 0.0 if x == 0 else axis_phase_exponent % 2
+    z = z_exponent % 2
+    if x == 1 and z != 0:
+        a = (a + z / 2) % 2
+        z = 0
+    if 0.5 < a <= 1.5:
+        a = (a - 1) % 2
+        x = 2 - x if x else x
+    return (x, z, a)
 
 
 @functools.cache

cirq-core/cirq/ops/phased_x_z_gate_test.py

@@ -366,3 +366,19 @@ def test_has_stabilizer_effect_false_for_non_cliffords(gate: cirq.PhasedXZGate)
 def test_has_stabilizer_effect_returns_false_for_symbolic_unitary() -> None:
     gate = cirq.PhasedXZGate(x_exponent=0, z_exponent=0, axis_phase_exponent=sympy.Symbol('a'))
     assert not cirq.has_stabilizer_effect(gate)
+
+
+@pytest.mark.parametrize(['x', 'z', 'a'], np.random.uniform(-3, 3, (100, 3)), ids=range(100))
+def test_canonical_xza_mod_2_matches_canonical(x: float, z: float, a: float) -> None:
+    gate = cirq.PhasedXZGate(x_exponent=x, z_exponent=z, axis_phase_exponent=a)._canonical()
+    xza_expected = (gate.x_exponent % 2, gate.z_exponent % 2, gate.axis_phase_exponent % 2)
+    xza_actual = cirq.ops.phased_x_z_gate._canonical_xza_mod_2(x, z, a)
+    assert xza_actual == xza_expected
+    # check at boundary values
+    xb = round(4 * x) / 4
+    zb = round(4 * z) / 4
+    ab = round(4 * a) / 4
+    gateb = cirq.PhasedXZGate(x_exponent=xb, z_exponent=zb, axis_phase_exponent=ab)._canonical()
+    xzab_expected = (gateb.x_exponent % 2, gateb.z_exponent % 2, gateb.axis_phase_exponent % 2)
+    xzab_actual = cirq.ops.phased_x_z_gate._canonical_xza_mod_2(xb, zb, ab)
+    assert xzab_actual == xzab_expected