Axis-aware S³ compression: emit rx/ry/rz for Cartesian-aligned fused rotations

src/rqm_optimize/fusion.py

@@ -11,7 +11,13 @@
 
 import numpy as np
 
-from .geometry import remove_global_phase
+from .geometry import (
+    quaternion_canonicalize,
+    quaternion_multiply,
+    quaternion_to_su2,
+    remove_global_phase,
+    su2_to_quaternion,
+)
 from .qiskit_adapter import extract_matrix, is_fuseable_single_qubit
 
 if TYPE_CHECKING:
@@ -107,24 +113,35 @@ def flush_all() -> None:
 
 
 def _fuse_matrices(run: list["CircuitInstruction"]) -> np.ndarray:
-    """Multiply gate matrices left-to-right (circuit order) and normalise.
+    """Accumulate a run of single-qubit gates as quaternion products on S³.
 
-    In quantum circuit convention the first gate applied is the leftmost
-    matrix.  We accumulate U_n @ ... @ U_1 by applying right-to-left
-    composition so that the overall effect is the same as applying gates
-    sequentially.
+    Each ``CircuitInstruction`` in *run* is inspected for its 2×2 unitary
+    matrix, which is then mapped to a unit quaternion via the exact SU(2)
+    isomorphism.  The quaternions are multiplied in circuit order (gate[0]
+    applied first), and the canonical accumulated result is converted back to
+    a 2×2 SU(2) matrix.
+
+    Using quaternion multiplication instead of raw matrix multiplication keeps
+    intermediate results on the unit 3-sphere S³, avoids accumulated complex
+    phase drift, and produces a canonical output via
+    :func:`~rqm_optimize.geometry.quaternion_canonicalize`.
+
+    Circuit order: gate[0] is applied first.  Quaternion product ``q_n * … *
+    q_1`` is accumulated with the last gate on the left (outer position),
+    matching the matrix convention ``U_n @ … @ U_1``.
 
     Args:
         run: List of ``CircuitInstruction`` objects in circuit order.
 
     Returns:
-        2×2 SU(2)-normalised combined unitary matrix.
+        2×2 SU(2)-normalized combined unitary matrix.
     """
-    # Start with identity and compose in circuit order.
-    # circuit order: gate[0] first, gate[1] second, ...
-    # combined unitary = gate[-1] @ ... @ gate[1] @ gate[0]
-    combined = np.eye(2, dtype=complex)
+    # Identity quaternion: q = (1, 0, 0, 0).
+    q_accum = np.array([1.0, 0.0, 0.0, 0.0])
     for instr in run:
         mat = extract_matrix(instr.operation)
-        combined = mat @ combined
-    return remove_global_phase(combined)
+        q_gate = su2_to_quaternion(remove_global_phase(mat))
+        # Apply q_gate after the current accumulation: q_gate * q_accum.
+        q_accum = quaternion_multiply(q_gate, q_accum)
+    q_accum = quaternion_canonicalize(q_accum)
+    return quaternion_to_su2(q_accum)

src/rqm_optimize/geometry.py

@@ -2,6 +2,28 @@
 
 Keeps math helpers small and rigorous.  Does not duplicate rqm-core; this
 module only contains the minimum needed to support circuit-level optimization.
+
+Quaternion convention
+---------------------
+A unit quaternion is represented as a length-4 NumPy array ``q = [w, x, y, z]``
+where *w* is the scalar part and *x, y, z* are the pure-imaginary components.
+
+The exact isomorphism between unit quaternions and SU(2) matrices used here is::
+
+    q = (w, x, y, z)  <->  U = [[ w - iz,  -(y + ix)],
+                                  [ y - ix,    w + iz ]]
+
+This convention is consistent with the standard quantum rotation gates:
+
+* ``Rx(theta)``  ->  ``(cos(theta/2), sin(theta/2), 0,            0          )``
+* ``Ry(theta)``  ->  ``(cos(theta/2), 0,            sin(theta/2), 0          )``
+* ``Rz(theta)``  ->  ``(cos(theta/2), 0,            0,            sin(theta/2))``
+
+Quaternion multiplication ``quaternion_multiply(q2, q1)`` represents applying
+*q1* first, then *q2*, matching matrix composition ``U2 @ U1``.
+
+For canonicalization the representative with non-negative scalar part (``w >= 0``)
+is preferred, keeping the quaternion on the shortest geodesic arc on S³.
 """
 
 from __future__ import annotations
@@ -69,3 +91,196 @@ def matrices_close(
     n = a.shape[0]
     inner = np.trace(a.conj().T @ b)
     return bool(abs(abs(inner) - n) < atol * n)
+
+
+# ---------------------------------------------------------------------------
+# Quaternion helpers
+# ---------------------------------------------------------------------------
+
+def su2_to_quaternion(matrix: NDArray[np.complexfloating]) -> NDArray[np.floating]:
+    """Return the unit quaternion ``[w, x, y, z]`` for an SU(2) matrix.
+
+    Uses the exact algebraic isomorphism::
+
+        U = [[ w - iz,  -(y + ix)],
+             [ y - ix,    w + iz ]]  <->  q = (w, x, y, z)
+
+    The result is normalized and has ``w >= 0`` (canonical shortest-path
+    representative on S³).
+
+    Args:
+        matrix: 2×2 complex SU(2) array (determinant +1, unitary).
+
+    Returns:
+        Length-4 float array ``[w, x, y, z]`` representing the unit quaternion.
+    """
+    # Extract components from matrix entries:
+    #   U[0,0] = w - iz  =>  w = Re(U[0,0]),  z = -Im(U[0,0])
+    #   U[0,1] = -(y+ix) =>  y = -Re(U[0,1]), x = -Im(U[0,1])
+    w = float(matrix[0, 0].real)
+    x = float(-matrix[0, 1].imag)
+    y = float(-matrix[0, 1].real)
+    z = float(-matrix[0, 0].imag)
+    q = np.array([w, x, y, z], dtype=float)
+    return quaternion_canonicalize(q)
+
+
+def quaternion_to_su2(q: NDArray[np.floating]) -> NDArray[np.complexfloating]:
+    """Return the SU(2) matrix for a unit quaternion ``[w, x, y, z]``.
+
+    Inverse of :func:`su2_to_quaternion`.  Uses the isomorphism::
+
+        q = (w, x, y, z)  <->  U = [[ w - iz,  -(y + ix)],
+                                      [ y - ix,    w + iz ]]
+
+    Args:
+        q: Length-4 float array ``[w, x, y, z]``.
+
+    Returns:
+        2×2 complex SU(2) array.
+    """
+    w, x, y, z = float(q[0]), float(q[1]), float(q[2]), float(q[3])
+    return np.array(
+        [
+            [complex(w, -z), complex(-y, -x)],
+            [complex(y, -x), complex(w,  z)],
+        ],
+        dtype=complex,
+    )
+
+
+def quaternion_multiply(
+    q2: NDArray[np.floating],
+    q1: NDArray[np.floating],
+) -> NDArray[np.floating]:
+    """Return the quaternion product ``q2 * q1``.
+
+    Represents applying gate *q1* first, then gate *q2*, matching matrix
+    composition ``U2 @ U1``.
+
+    Uses the standard Hamilton product formula::
+
+        (w2, x2, y2, z2) * (w1, x1, y1, z1) = (
+            w2*w1 - x2*x1 - y2*y1 - z2*z1,
+            w2*x1 + x2*w1 + y2*z1 - z2*y1,
+            w2*y1 - x2*z1 + y2*w1 + z2*x1,
+            w2*z1 + x2*y1 - y2*x1 + z2*w1,
+        )
+
+    Args:
+        q2: Unit quaternion for the second (outer) gate.
+        q1: Unit quaternion for the first (inner) gate.
+
+    Returns:
+        Product quaternion (not yet canonicalized or re-normalized).
+    """
+    w2, x2, y2, z2 = float(q2[0]), float(q2[1]), float(q2[2]), float(q2[3])
+    w1, x1, y1, z1 = float(q1[0]), float(q1[1]), float(q1[2]), float(q1[3])
+    return np.array(
+        [
+            w2 * w1 - x2 * x1 - y2 * y1 - z2 * z1,
+            w2 * x1 + x2 * w1 + y2 * z1 - z2 * y1,
+            w2 * y1 - x2 * z1 + y2 * w1 + z2 * x1,
+            w2 * z1 + x2 * y1 - y2 * x1 + z2 * w1,
+        ],
+        dtype=float,
+    )
+
+
+def quaternion_canonicalize(q: NDArray[np.floating]) -> NDArray[np.floating]:
+    """Return the canonical unit quaternion with non-negative scalar part.
+
+    Both *q* and *-q* represent the same physical SU(2) rotation.  This
+    function normalizes the quaternion and selects the representative with
+    ``w >= 0``, keeping it on the shortest geodesic arc (canonical
+    representative) on S³.
+
+    Args:
+        q: Length-4 float array ``[w, x, y, z]``.
+
+    Returns:
+        Normalized quaternion with ``w >= 0``.
+    """
+    norm = float(np.linalg.norm(q))
+    if norm < _ATOL:
+        return np.array([1.0, 0.0, 0.0, 0.0], dtype=float)
+    q = q / norm
+    if q[0] < 0.0:
+        q = -q
+    return q
+
+
+def quaternion_to_axis_angle(
+    q: NDArray[np.floating],
+) -> tuple[NDArray[np.floating], float]:
+    """Extract the rotation axis and angle from a unit quaternion.
+
+    The physical Bloch-sphere rotation angle is ``theta = 2 * arccos(w)``
+    (the factor of 2 arises from the half-angle spinor convention).
+
+    Args:
+        q: Length-4 float array ``[w, x, y, z]``, assumed to be a unit
+           quaternion.  Canonicalization is applied internally.
+
+    Returns:
+        A tuple ``(axis, theta)`` where:
+
+        * ``axis`` is a length-3 float array giving the unit rotation axis
+          ``[nx, ny, nz]``.  When the rotation is near-identity (``theta``
+          close to 0), ``axis`` defaults to ``[0, 0, 1]``.
+        * ``theta`` is the rotation angle in radians (in ``[0, pi]`` after
+          canonicalization).
+    """
+    q = quaternion_canonicalize(q)
+    w = float(np.clip(q[0], -1.0, 1.0))
+    theta = 2.0 * float(np.arccos(w))
+    sin_half = float(np.sqrt(max(0.0, 1.0 - w * w)))
+    if sin_half < _ATOL:
+        axis = np.array([0.0, 0.0, 1.0], dtype=float)
+    else:
+        axis = np.array(q[1:], dtype=float) / sin_half
+    return axis, theta
+
+
+# Reference unit vectors for the three Cartesian rotation axes.
+_CARTESIAN_AXES: tuple[tuple[str, NDArray[np.floating]], ...] = (
+    ("x", np.array([1.0, 0.0, 0.0])),
+    ("y", np.array([0.0, 1.0, 0.0])),
+    ("z", np.array([0.0, 0.0, 1.0])),
+)
+
+# Default tolerance for axis-alignment checks.
+_AXIS_ATOL = 1e-6
+
+
+def axis_aligned_rotation(
+    q: NDArray[np.floating],
+    atol: float = _AXIS_ATOL,
+) -> tuple[str, float] | None:
+    """Return ``(axis_name, theta)`` if *q* is a rotation about x, y, or z.
+
+    Checks whether the rotation axis of *q* lies within *atol* of the x, y, or
+    z unit vector.  This is the geometric test that drives the axis-aware
+    compression pass: when the result is non-``None``, the caller can emit a
+    single named rotation gate (``rx``, ``ry``, or ``rz``) instead of a
+    generic ``U`` gate.
+
+    Near-identity rotations (``|theta| < atol``) return ``None`` because the
+    axis is numerically undefined at that scale.
+
+    Args:
+        q: Length-4 float array ``[w, x, y, z]``, a unit quaternion.
+        atol: Absolute tolerance for axis alignment and near-identity detection.
+
+    Returns:
+        ``("x"/"y"/"z", theta)`` when the axis is Cartesian-aligned, or
+        ``None`` when the rotation is generic or near-identity.
+    """
+    q = quaternion_canonicalize(q)
+    axis, theta = quaternion_to_axis_angle(q)
+    if abs(theta) < atol:
+        return None
+    for name, ref in _CARTESIAN_AXES:
+        if np.allclose(axis, ref, atol=atol):
+            return name, theta
+    return None

src/rqm_optimize/qiskit_adapter.py

@@ -178,6 +178,48 @@ def _try_extract_matrix(op: "Instruction") -> np.ndarray | None:
     return None
 
 
+def emit_axis_aligned_gate(
+    matrix: np.ndarray,
+    circuit: "QuantumCircuit",
+    qubit: "Qubit",
+) -> int:
+    """Emit a single Rx/Ry/Rz gate if *matrix* is a Cartesian-axis rotation.
+
+    Inspects the quaternion representation of *matrix* to determine whether the
+    fused rotation is aligned with the x, y, or z axis on S³.  When it is, a
+    single named rotation gate (``rx``, ``ry``, or ``rz``) is appended to
+    *circuit* and 1 is returned.  When the rotation is generic or near-identity,
+    the circuit is left unchanged and 0 is returned.
+
+    This is the axis-aware compression pass: it operates directly on the S³
+    quaternion representation rather than on matrices, and produces a semantically
+    named gate that is directly executable on hardware (in particular, ``rz`` is
+    virtual/"free" on many backends).
+
+    Args:
+        matrix: 2×2 SU(2)-normalized unitary matrix for the fused run.
+        circuit: Target ``QuantumCircuit`` to append the gate to.
+        qubit: The target qubit.
+
+    Returns:
+        1 if an axis-aligned gate was emitted, 0 otherwise.
+    """
+    from .geometry import axis_aligned_rotation, su2_to_quaternion
+
+    q = su2_to_quaternion(matrix)
+    result = axis_aligned_rotation(q)
+    if result is None:
+        return 0
+    axis_name, theta = result
+    if axis_name == "x":
+        circuit.rx(theta, qubit)
+    elif axis_name == "y":
+        circuit.ry(theta, qubit)
+    else:
+        circuit.rz(theta, qubit)
+    return 1
+
+
 def emit_euler_gate(
     unitary: np.ndarray,
     circuit: "QuantumCircuit",
@@ -190,7 +232,7 @@ def emit_euler_gate(
     produce a compact, exact decomposition.
 
     Args:
-        unitary: 2×2 complex SU(2)-normalised unitary matrix.
+        unitary: 2×2 complex SU(2)-normalized unitary matrix.
         circuit: The ``QuantumCircuit`` to append gates to.
         qubit: The target qubit.
         basis: Euler decomposer basis string (e.g. ``"U"``, ``"ZSX"``,
@@ -246,9 +288,19 @@ def build_optimized_circuit(
             matrix: np.ndarray = seg["matrix"]
             original_count: int = seg["original_count"]
 
-            # Build a temporary single-qubit sub-circuit to count decomposition.
             from qiskit import QuantumCircuit as QC
 
+            # First pass: axis-aware compression.
+            # If the fused rotation is aligned with the x, y, or z axis, emit
+            # a single rx/ry/rz gate directly from the S³ quaternion form.
+            # This always produces 1 gate (≤ original_count which is ≥ 2).
+            tmp = QC(1)
+            if emit_axis_aligned_gate(matrix, tmp, tmp.qubits[0]) > 0:
+                for instr in tmp.data:
+                    out.append(instr.operation, [qubit])
+                continue
+
+            # Second pass: Euler decomposer for generic (non-axis-aligned) rotations.
             tmp = QC(1)
             emitted = emit_euler_gate(matrix, tmp, tmp.qubits[0], basis=basis)