Deep dive · Visualization

Wigner Functions for Single-Qubit Quantum Gates

Pauli rotations rendered as Wigner functions, animated in two and three dimensions, with the density-operator and Bloch-vector groundwork.

Onri Jay BenallyQuantum Hardware Engineer, University of Minnesota

·6 August 2026·17 min readvisualization quantum

About the author

Onri Jay Benally

Quantum Hardware Engineer · Doctoral Researcher, University of Minnesota

Onri fabricates the hardware he writes about. Eight years in the J.P. Wang group at Minnesota have gone into nanoelectronics fabrication and materials integration: cryogenic MRAM, spin-orbit torque devices, spintronic quantum processor chips. He has trained twenty engineers and physicists on high-yield fabrication in an ISO 4 cleanroom, and has worked as a quantum hardware engineer at IBM's T.J. Watson lab. By his own account what he does best is turn technical jargon into practical layman's terms, which is what these notebooks are for.

GitHub ↗Original notebook ↗

Republished with the author's permission from the original notebook and shared under CC-BY-4.0. The text, code, and figures below are Onri's, as published.

1. Core Objects and Symbols

Pauli matrices (in the computational basis ( ${\lvert 0\rangle,\lvert 1\rangle}$ )):

\sigma_x = \begin{pmatrix} 0 & 1 \\ 1 & 0 \end{pmatrix}, \quad \sigma_y = \begin{pmatrix} 0 & -i \\ i & 0 \end{pmatrix}, \quad \sigma_z = \begin{pmatrix} 1 & 0 \\ 0 & -1 \end{pmatrix}

Density operator and Bloch vector ( $\mathbf r\in\mathbb R^3$ ):

\rho = \tfrac{1}{2}\big(I + \mathbf r\cdot\boldsymbol{\sigma}\big),\qquad r_j = \operatorname{Tr}(\rho,\sigma_j),\ j\in{x,y,z}.

For a pure state ( $\rho = \lvert\psi\rangle\langle\psi\rvert$ ), one has ( $\lVert\mathbf r\rVert = 1$ ).

Unit vector on the sphere (spherical angles ( $\theta\in[0,\pi],\ \phi\in[0,2\pi$ ))):

\mathbf n(\theta,\phi) = \big(\sin\theta\cos\phi,\ \sin\theta\sin\phi,\ \cos\theta\big).

2. Stratonovich–Weyl (SW) Kernel and Qubit Wigner Function

A convenient SW kernel for a qubit is

\Delta(\theta,\phi) = \tfrac{1}{2}\Big(I + \sqrt{3};\mathbf n(\theta,\phi)\cdot\boldsymbol{\sigma}\Big),

which yields the spin Wigner function (on ( $S^2$ ))

W_\rho(\theta,\phi) = \operatorname{Tr}\big[\rho,\Delta(\theta,\phi)\big] = \tfrac{1}{2} + \frac{\sqrt{3}}{2}; \mathbf r\cdot\mathbf n(\theta,\phi). \tag{SW–Qubit}

For a pure state (( $\lVert\mathbf r\rVert=1$ )), the extrema are

W_{\max}=\tfrac{1}{2}+\tfrac{\sqrt{3}}{2},\qquad W_{\min}=\tfrac{1}{2}-\tfrac{\sqrt{3}}{2}.

Covariance (rigid rotation): for any unitary (U),

W_{U\rho U^{\dagger}}(\theta,\phi) = W_{\rho}!\big(R^{-1}\mathbf n(\theta,\phi)\big), \tag{Covariance}

where ( $R\in\mathrm{SO}(3)$ ) is the 3D rotation associated with (U) (see §3–§4).

3. Pauli Gates as ( $\pi$ )-Rotations

The single‑qubit Pauli gates are equivalent (up to a global phase) to ( $\pi$ )-rotations about Cartesian axes:

X = e^{-i\frac{\pi}{2},\sigma_x},\qquad Y = e^{-i\frac{\pi}{2},\sigma_y},\qquad Z = e^{-i\frac{\pi}{2},\sigma_z}.

They act on the Bloch vector by the corresponding rotation matrices ( $R_x(\pi),R_y(\pi),R_z(\pi)$ ), so (Covariance) implies the Wigner pattern simply rotates rigidly on ( $S^2$ ).

4. General SU(2) Rotation and the SO(3) Map

A general single‑qubit rotation about a unit axis ( $\hat{\mathbf u}$ ) by angle ( $\alpha$ ) is

U(\alpha,\hat{\mathbf u}) = \exp!\big(-i,\tfrac{\alpha}{2},\hat{\mathbf u}\cdot\boldsymbol{\sigma}\big).

Its action on ( $\mathbf r$ ) is the adjoint map ( $\mathrm{SU}(2)\to\mathrm{SO}(3)$ ):

\mathbf r;\longmapsto;\mathbf r' = R(\alpha,\hat{\mathbf u}),\mathbf r. \tag{Adjoint}

5. Rodrigues’ Formula (Axis–Angle ( $\to$ ) Rotation Matrix)

Let ( $\hat{\mathbf u}=(u_x,u_y,u_z)$ ) be a unit axis and define the cross‑product matrix

K(\hat{\mathbf u}) = \begin{pmatrix} 0 & -u_z & u_y\ u_z & 0 & -u_x\ -u_y & u_x & 0 \end{pmatrix}.

Then the corresponding rotation matrix is

R(\alpha,\hat{\mathbf u}) = \cos\alpha,I_3 + (1-\cos\alpha),\hat{\mathbf u},\hat{\mathbf u}^{\mathsf T} + \sin\alpha,K(\hat{\mathbf u}). \tag{Rodrigues}

This matrix implements the evolution ( $\mathbf r' = R,\mathbf r$ ), and by (Covariance) the Wigner function obeys

W_\rho(\theta,\phi;\alpha) = \tfrac{1}{2} + \frac{\sqrt{3}}{2}; \big(R(\alpha,\hat{\mathbf u}),\mathbf r\big)\cdot\mathbf n(\theta,\phi).

6. Optional: Discrete ( $2\times2$ ) (Finite‑State) Wigner

On a 4‑point discrete phase space ( $\alpha\in{(q,p)}$ ), define phase‑point operators ( $A_\alpha$ ) with

\sum_{\alpha} A_\alpha = 2I,\quad \operatorname{Tr}(A_\alpha)=1,\quad A_\alpha = A_\alpha^{\dagger}.

The discrete Wigner function is

W_\rho(\alpha) = \tfrac{1}{2},\operatorname{Tr}(\rho,A_\alpha),

with line sums reproducing measurement probabilities in mutually unbiased bases.

7. Color Scale (Pure‑State Bounds)

For ( $\lVert\mathbf r\rVert=1$ ):

W_{\min}=\tfrac{1}{2}-\tfrac{\sqrt{3}}{2}\approx -0.3660254,\qquad W_{\max}=\tfrac{1}{2}+\tfrac{\sqrt{3}}{2}\approx 1.3660254.

These constants are useful for fixing a consistent color scale across frames.

8. Acronym/Notation Glossary

SW — Stratonovich–Weyl (kernel/correspondence).
SU(2) — Special Unitary group of degree two (spin‑(1/2) rotations).
SO(3) — Special Orthogonal group in 3D (real rotation matrices).
WF — Wigner function.
( $\lvert\psi\rangle$ ) — State vector (ket); ( $\langle\psi\rvert$ ) — bra; ( $\rho$ ) — density operator.
Adjoint map — Conjugation action ( $U: \rho\mapsto U\rho U^{\dagger}$ ) inducing ( $\mathbf r\mapsto R\mathbf r$ ).

9. Mind‑Map

Wigner for Qubits

Wigner (spin, S^2)
├─ SW kernel Δ(θ,φ) → W(θ,φ) = Tr[ρΔ]
│  └─ Explicit: ½ + (√3/2) r·n(θ,φ)
├─ State ρ ↔ Bloch vector r
├─ Gates U = exp[-i(α/2) u·σ]
│  ├─ Pauli X/Y/Z: α = π, axes x/y/z
│  └─ Custom axis–angle (Rodrigues)
└─ Covariance: W → W∘R^{-1}

10. Minimal Substitutions (if variations are needed)

If a different SW kernel normalization is chosen, the prefactors in ( $\Delta$ ) and hence ( $W$ ) adjust accordingly, while covariance and linearity remain.
Mixed states use the same formulas with ( $\lVert\mathbf r\rVert<1$ ) (the extrema shrink linearly with ( $\lVert\mathbf r\rVert$ )).

# @title Cell 1 — Imports & global plotting config (DPI=200)
"""
Spin (SU(2)) Wigner-function animations for single-qubit gates.

This notebook renders the qubit Wigner function
    W(θ, φ) = 1/2 + (√3/2) * r · n(θ, φ)
and animates its rigid rotation under Pauli and custom axis–angle gates.

All titles and dynamic text are drawn *inside* the plotted region so they
never clip in the saved GIFs. Use Cells 3 and 4 to generate animations.
"""

from __future__ import annotations

import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt

# Global Matplotlib configuration: DPI = 200 everywhere
mpl.rcParams.update({
    "figure.dpi": 200,       # on-screen DPI
    "savefig.dpi": 200,      # saved file DPI
    "figure.figsize": (6.4, 3.6),  # 1280×720 at DPI=200
    "axes.grid": False,
})

# @title Cell 2 — Math & utilities (PEP 8 / PEP 257)
from dataclasses import dataclass
from pathlib import Path
from typing import Tuple, Iterable
import textwrap

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation, PillowWriter
from IPython.display import Image, display


@dataclass
class Knobs:
    """User controls for Wigner animations.

    Attributes:
        gate_axes: Pauli axes to animate, each in {"x", "y", "z"}.
        initial_state: Starting Bloch eigenstate, or "custom".
        custom_r: Custom Bloch vector if initial_state == "custom".
        frames: Number of frames (≥ 2).
        n_theta: Polar samples (θ ∈ [0, π]).
        n_phi: Azimuth samples (φ ∈ [0, 2π]).
        fps: GIF framerate.
        out_dir: Output directory (e.g., "/content" in Colab).
        figsize: Figure size in inches (DPI set globally).
        info_loc: In-axes anchor ("ul", "ur", "ll", "lr").
        info_fontsize: Font size for in-axes info text.
        wrap_chars: Manual wrap width for info text.
    """
    gate_axes: Tuple[str, ...] = ("x", "y", "z")
    initial_state: str = "+z"
    custom_r: Tuple[float, float, float] = (0.0, 0.0, 1.0)
    frames: int = 24
    n_theta: int = 80
    n_phi: int = 160
    fps: int = 20
    out_dir: str = "/content"
    figsize: Tuple[float, float] = (6.4, 3.6)
    info_loc: str = "ul"
    info_fontsize: int = 9
    wrap_chars: int = 42


# ------------------------------ Bloch helpers ---------------------------------
def bloch_vector(initial_state: str,
                 custom_r: Tuple[float, float, float]) -> np.ndarray:
    """Return a unit Bloch vector for a named eigenstate or custom input."""
    mapping = {
        "+z": (0.0, 0.0, 1.0), "-z": (0.0, 0.0, -1.0),
        "+x": (1.0, 0.0, 0.0), "-x": (-1.0, 0.0, 0.0),
        "+y": (0.0, 1.0, 0.0), "-y": (0.0, -1.0, 0.0),
    }
    r = np.asarray(mapping.get(initial_state.lower(), custom_r), float)
    nrm = np.linalg.norm(r)
    return np.array([0.0, 0.0, 1.0]) if nrm == 0 else r / nrm


# ---------------------------- Axis–angle rotations ----------------------------
def rotation_matrix(axis: str, angle: float) -> np.ndarray:
    """Rodrigues rotation matrix for axis in {'x','y','z'} and angle (radians)."""
    c, s = np.cos(angle), np.sin(angle)
    if axis == "x":
        return np.array([[1, 0, 0], [0, c, -s], [0, s, c]], float)
    if axis == "y":
        return np.array([[c, 0, s], [0, 1, 0], [-s, 0, c]], float)
    if axis == "z":
        return np.array([[c, -s, 0], [s, c, 0], [0, 0, 1]], float)
    raise ValueError("axis must be one of {'x','y','z'}")


def rotation_matrix_axis(axis_vec: Tuple[float, float, float],
                         angle: float) -> np.ndarray:
    """Rodrigues matrix for arbitrary axis vector and angle (radians)."""
    v = np.asarray(axis_vec, float)
    nrm = np.linalg.norm(v)
    if nrm == 0:
        raise ValueError("axis_vec must be nonzero")
    k = v / nrm
    K = np.array([[0.0, -k[2], k[1]], [k[2], 0.0, -k[0]], [-k[1], k[0], 0.0]], float)
    I = np.eye(3)
    c, s = np.cos(angle), np.sin(angle)
    return c * I + (1 - c) * np.outer(k, k) + s * K


# ------------------------------ Wigner on S^2 ---------------------------------
def n_grid(n_theta: int, n_phi: int) -> np.ndarray:
    """Return unit-vector grid n(θ,φ) stacked as (3, n_theta, n_phi)."""
    theta = np.linspace(0.0, np.pi, n_theta)
    phi = np.linspace(0.0, 2.0 * np.pi, n_phi)
    th, ph = np.meshgrid(theta, phi, indexing="ij")
    nx = np.sin(th) * np.cos(ph)
    ny = np.sin(th) * np.sin(ph)
    nz = np.cos(th)
    return np.stack((nx, ny, nz), axis=0)


def spin_wigner_qubit(r_vec: np.ndarray, ngrid: np.ndarray) -> np.ndarray:
    """Spin SW Wigner for a qubit: W = 1/2 + (√3/2) * r·n on the (θ,φ) grid."""
    return 0.5 + (np.sqrt(3.0) / 2.0) * (r_vec.reshape(3, 1, 1) * ngrid).sum(axis=0)


# ----------------------- In-axes title & info placement -----------------------
def _info_xy(loc: str) -> Tuple[float, float, dict]:
    """Map location code to (x, y, kwargs) in axes-fraction coordinates."""
    loc = loc.lower()
    if loc == "ul":
        return 0.015, 0.985, dict(ha="left", va="top")
    if loc == "ur":
        return 0.985, 0.985, dict(ha="right", va="top")
    if loc == "ll":
        return 0.015, 0.015, dict(ha="left", va="bottom")
    if loc == "lr":
        return 0.985, 0.015, dict(ha="right", va="bottom")
    return 0.015, 0.985, dict(ha="left", va="top")


def add_info_box(ax: plt.Axes,
                 text: str,
                 loc: str = "ul",
                 fontsize: int = 9,
                 wrap: int = 42) -> plt.Text:
    """Create a wrapped, padded text box inside the axes (never clips)."""
    x, y, kw = _info_xy(loc)
    return ax.text(
        x, y, textwrap.fill(text, wrap),
        transform=ax.transAxes, fontsize=fontsize, zorder=5,
        bbox=dict(boxstyle="round,pad=0.3", fc="white", alpha=0.75, ec="none"),
        **kw,
    )


def add_title_inside(ax: plt.Axes, text: str) -> plt.Text:
    """Draw a title inside the Axes at the very top center (never clips)."""
    return ax.text(
        0.5, 0.985, text,
        transform=ax.transAxes, ha="center", va="top",
        fontsize=11,
        bbox=dict(boxstyle="round,pad=0.3", fc="white", alpha=0.75, ec="none"),
        zorder=6,
    )

# @title Cell 3 — Pauli X/Y/Z Wigner animations (titles *inside* the plot)
from pathlib import Path

def animate_wigner_pauli(axis: str,
                         r0: np.ndarray,
                         frames: int,
                         n_theta: int,
                         n_phi: int,
                         fps: int,
                         figsize: Tuple[float, float],
                         info_loc: str,
                         info_fontsize: int,
                         wrap_chars: int,
                         out_path: Path) -> Path:
    """Animate Wigner under a Pauli rotation U = exp(-i σ_axis * π * t / 2)."""
    assert frames >= 2, "frames must be ≥ 2"
    ngrid = n_grid(n_theta, n_phi)
    w_min = 0.5 - (np.sqrt(3.0) / 2.0)
    w_max = 0.5 + (np.sqrt(3.0) / 2.0)
    extent = (0.0, 2.0 * np.pi, 0.0, np.pi)

    fig, ax = plt.subplots(figsize=figsize, constrained_layout=True)
    im = ax.imshow(np.zeros((n_theta, n_phi)), origin="lower", extent=extent,
                   vmin=w_min, vmax=w_max, aspect="auto")
    cbar = fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
    cbar.set_label("Wigner value")

    ax.set_xlabel("φ (radians)")
    ax.set_ylabel("θ (radians)")
    add_title_inside(ax, f"Wigner — {axis.upper()} rotation")

    info = add_info_box(ax, "", loc=info_loc, fontsize=info_fontsize,
                        wrap=wrap_chars)

    def update(k: int):
        ang = (k / (frames - 1)) * np.pi  # 0 → π
        r = rotation_matrix(axis.lower(), ang) @ r0
        im.set_data(spin_wigner_qubit(r, ngrid))
        info.set_text(f"θ = {ang:.3f} rad\nframe {k+1}/{frames}")
        return (im, info)

    anim = FuncAnimation(fig, update, frames=frames, interval=60, blit=False)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    anim.save(out_path.as_posix(), writer=PillowWriter(fps=fps),
              savefig_kwargs=dict(facecolor="white"))
    plt.close(fig)
    return out_path


# ---------------------------- Run X, Y, Z animations --------------------------
knobs = Knobs(  # tweak if desired
    gate_axes=("x", "y", "z"),
    initial_state="+z",
    frames=24,                # raise to 48 for smoother motion
    n_theta=80, n_phi=160,    # raise to 120×240 for more detail
    fps=20,
    out_dir="/content",
    figsize=(6.4, 3.6),
    info_loc="ul", info_fontsize=9, wrap_chars=42,
)

r0 = bloch_vector(knobs.initial_state, knobs.custom_r)
outs = []
for axname in knobs.gate_axes:
    path = Path(knobs.out_dir) / f"wigner_pauli_{axname}.gif"
    outs.append(
        animate_wigner_pauli(axname, r0, knobs.frames,
                             knobs.n_theta, knobs.n_phi, knobs.fps,
                             knobs.figsize, knobs.info_loc,
                             knobs.info_fontsize, knobs.wrap_chars,
                             path)
    )

# Display inline
for p in outs:
    display(Image(filename=str(p)))
    print("Saved:", p)

Saved: /content/wigner_pauli_x.gif

Saved: /content/wigner_pauli_y.gif

Saved: /content/wigner_pauli_z.gif

# @title Cell 4 — Custom axis–angle gate (example: 45° about y)
from pathlib import Path

def animate_wigner_custom(axis_vec: Tuple[float, float, float],
                          angle_deg: float,
                          r0: np.ndarray,
                          frames: int,
                          n_theta: int,
                          n_phi: int,
                          fps: int,
                          figsize: Tuple[float, float],
                          info_loc: str,
                          info_fontsize: int,
                          wrap_chars: int,
                          out_path: Path) -> Path:
    """Animate Wigner under an arbitrary axis–angle rotation 0 → angle_deg."""
    assert frames >= 2, "frames must be ≥ 2"
    ngrid = n_grid(n_theta, n_phi)
    w_min = 0.5 - (np.sqrt(3.0) / 2.0)
    w_max = 0.5 + (np.sqrt(3.0) / 2.0)
    extent = (0.0, 2.0 * np.pi, 0.0, np.pi)

    fig, ax = plt.subplots(figsize=figsize, constrained_layout=True)
    im = ax.imshow(np.zeros((n_theta, n_phi)), origin="lower", extent=extent,
                   vmin=w_min, vmax=w_max, aspect="auto")
    cbar = fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
    cbar.set_label("Wigner value")

    ax.set_xlabel("φ (radians)")
    ax.set_ylabel("θ (radians)")
    add_title_inside(ax, "Custom rotation")

    info = add_info_box(ax, "", loc=info_loc, fontsize=info_fontsize,
                        wrap=wrap_chars)

    def update(k: int):
        ang = (k / (frames - 1)) * np.deg2rad(angle_deg)
        R = rotation_matrix_axis(axis_vec, ang)
        r = R @ r0
        im.set_data(spin_wigner_qubit(r, ngrid))
        info.set_text(
            f"axis = ({axis_vec[0]:.2f}, {axis_vec[1]:.2f}, {axis_vec[2]:.2f})\n"
            f"θ = {np.rad2deg(ang):.1f}° / {angle_deg:.1f}°"
        )
        return (im, info)

    anim = FuncAnimation(fig, update, frames=frames, interval=60, blit=False)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    anim.save(out_path.as_posix(), writer=PillowWriter(fps=fps),
              savefig_kwargs=dict(facecolor="white"))
    plt.close(fig)
    return out_path


# Example: 45° about the y-axis, starting from |+z⟩
custom_out = Path(knobs.out_dir) / "wigner_custom_y45.gif"
animate_wigner_custom(
    axis_vec=(0.0, 1.0, 0.0),
    angle_deg=45.0,
    r0=r0,
    frames=24,                 # raise to 48 if you prefer
    n_theta=80, n_phi=160,
    fps=20,
    figsize=knobs.figsize,
    info_loc=knobs.info_loc, info_fontsize=knobs.info_fontsize,
    wrap_chars=knobs.wrap_chars,
    out_path=custom_out,
)

display(Image(filename=str(custom_out)))
print("Saved:", custom_out)

Saved: /content/wigner_custom_y45.gif

The next following scripts are for the 3D Wigner plots.

# @title Cell 5 — 3D Wigner animations for ALL Pauli gates (X, Y, Z)
# PEP 8/257 compliant. Stable: updates facecolors in place; no deleting artists.

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Tuple, Iterable

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation, PillowWriter
from matplotlib.colors import Normalize
from matplotlib.cm import ScalarMappable
from IPython.display import Image, display


@dataclass
class Knobs3D:
    """User controls for 3D Pauli-gate Wigner animations on a colored sphere.

    Attributes:
        initial_state: Starting Bloch eigenstate label or "custom".
        custom_r: Custom Bloch vector if initial_state == "custom".
        frames: Number of frames (≥ 2) for 0 → π rotation.
        n_theta: Polar grid size for the sphere.
        n_phi: Azimuth grid size for the sphere.
        fps: GIF framerate.
        elev: Camera elevation in degrees.
        azim: Camera azimuth in degrees.
        out_dir: Output directory (e.g., "/content" on Colab).
        figsize: Figure size in inches (DPI set globally by rcParams).
        info_fontsize: In-axes info text font size.
    """
    initial_state: str = "+z"
    custom_r: Tuple[float, float, float] = (0.0, 0.0, 1.0)
    frames: int = 24
    n_theta: int = 64
    n_phi: int = 128
    fps: int = 20
    elev: float = 25.0
    azim: float = -60.0
    out_dir: str = "/content"
    figsize: Tuple[float, float] = (6.4, 4.8)
    info_fontsize: int = 9


# ----------------------------- Helpers (self-contained) -----------------------------
def _bloch_vector(label: str, custom=(0.0, 0.0, 1.0)) -> np.ndarray:
    mapping = {
        "+z": (0.0, 0.0, 1.0), "-z": (0.0, 0.0, -1.0),
        "+x": (1.0, 0.0, 0.0), "-x": (-1.0, 0.0, 0.0),
        "+y": (0.0, 1.0, 0.0), "-y": (0.0, -1.0, 0.0),
    }
    r = np.asarray(mapping.get(label.lower(), custom), float)
    nrm = np.linalg.norm(r) or 1.0
    return r / nrm


def _rotation_matrix(axis: str, angle: float) -> np.ndarray:
    c, s = np.cos(angle), np.sin(angle)
    if axis == "x":
        return np.array([[1, 0, 0], [0, c, -s], [0, s, c]], float)
    if axis == "y":
        return np.array([[c, 0, s], [0, 1, 0], [-s, 0, c]], float)
    if axis == "z":
        return np.array([[c, -s, 0], [s, c, 0], [0, 0, 1]], float)
    raise ValueError("axis must be 'x', 'y', or 'z'")


def _n_grid(n_theta: int, n_phi: int):
    """Return (θ, φ, n) with n stacked as (3, n_theta, n_phi)."""
    theta = np.linspace(0.0, np.pi, n_theta)
    phi = np.linspace(0.0, 2.0 * np.pi, n_phi)
    th, ph = np.meshgrid(theta, phi, indexing="ij")
    nx = np.sin(th) * np.cos(ph)
    ny = np.sin(th) * np.sin(ph)
    nz = np.cos(th)
    n = np.stack((nx, ny, nz), axis=0)
    return th, ph, n


def _wigner_qubit(r_vec: np.ndarray, n: np.ndarray) -> np.ndarray:
    """W(θ, φ) = 1/2 + (√3/2) * r·n on the (θ, φ) grid."""
    return 0.5 + (np.sqrt(3.0) / 2.0) * (r_vec.reshape(3, 1, 1) * n).sum(axis=0)


def _animate_one_pauli(axis: str, k: Knobs3D) -> Path:
    """Render one 3D Wigner animation for a chosen Pauli axis; return output path."""
    r0 = _bloch_vector(k.initial_state, k.custom_r)
    th, ph, n = _n_grid(k.n_theta, k.n_phi)

    # Unit sphere geometry
    X = np.sin(th) * np.cos(ph)
    Y = np.sin(th) * np.sin(ph)
    Z = np.cos(th)

    # Fixed normalization (pure-state bounds) for constant color scale
    w_min = 0.5 - (np.sqrt(3.0) / 2.0)
    w_max = 0.5 + (np.sqrt(3.0) / 2.0)
    norm = Normalize(vmin=w_min, vmax=w_max)
    cmap = plt.get_cmap()
    sm = ScalarMappable(norm=norm, cmap=cmap)
    sm.set_array([])  # required for colorbar from a ScalarMappable

    fig = plt.figure(figsize=k.figsize, constrained_layout=True)
    ax = fig.add_subplot(111, projection="3d")
    ax.view_init(elev=k.elev, azim=k.azim)
    if hasattr(ax, "set_box_aspect"):
        ax.set_box_aspect((1, 1, 1))

    # Initial facecolors (use cell-centered (M-1, N-1) colors for quads)
    W0 = _wigner_qubit(_rotation_matrix(axis, 0.0) @ r0, n)
    FC0 = cmap(norm(W0[:-1, :-1]))  # shape: (n_theta-1, n_phi-1, 4)

    surf = ax.plot_surface(
        X, Y, Z,
        facecolors=FC0,
        rstride=1, cstride=1,
        antialiased=False, linewidth=0,
        shade=False,  # use given facecolors directly
    )

    # Colorbar (bound to ScalarMappable using same norm+cmap)
    cbar = fig.colorbar(sm, ax=ax, fraction=0.046, pad=0.04)
    cbar.set_label("Wigner value")

    # Labels and in-axes title and info
    ax.set_xlabel("x")
    ax.set_ylabel("y")
    ax.set_zlabel("z")
    title_txt = ax.text2D(
        0.5, 0.98, f"Wigner — {axis.upper()} rotation (3D sphere)",
        transform=ax.transAxes, ha="center", va="top", fontsize=11,
        bbox=dict(boxstyle="round,pad=0.3", fc="white", alpha=0.75, ec="none"),
        zorder=6,
    )
    info_txt = ax.text2D(
        0.02, 0.96, "", transform=ax.transAxes, ha="left", va="top",
        fontsize=k.info_fontsize,
        bbox=dict(boxstyle="round,pad=0.3", fc="white", alpha=0.75, ec="none"),
        zorder=6,
    )

    def update(frame: int):
        angle = (frame / (k.frames - 1)) * np.pi  # 0 → π
        r_now = _rotation_matrix(axis, angle) @ r0
        W_now = _wigner_qubit(r_now, n)
        FC_now = cmap(norm(W_now[:-1, :-1]))        # (M-1, N-1, 4)
        surf.set_facecolors(FC_now.reshape(-1, 4))  # update in place
        info_txt.set_text(f"θ = {angle:.3f} rad\nframe {frame+1}/{k.frames}")
        return (surf, info_txt, title_txt)

    anim = FuncAnimation(fig, update, frames=k.frames, interval=60, blit=False)
    out = Path(k.out_dir) / f"wigner3d_pauli_{axis}.gif"
    out.parent.mkdir(parents=True, exist_ok=True)
    anim.save(out.as_posix(), writer=PillowWriter(fps=k.fps),
              savefig_kwargs=dict(facecolor="white"))
    plt.close(fig)
    return out


def animate_all_paulis_3d(k: Knobs3D,
                          axes: Tuple[str, ...] = ("x", "y", "z")) -> Iterable[Path]:
    """Generate and display GIFs for all requested Pauli axes; yield file paths."""
    for axname in axes:
        path = _animate_one_pauli(axname.lower(), k)
        display(Image(filename=str(path)))
        print(f"Saved: {path}")
        yield path


# ------------------------------ Run all three gates ------------------------------
K3D = Knobs3D(  # adjust as desired
    initial_state="+z",
    frames=24, n_theta=64, n_phi=128, fps=20,
    elev=25.0, azim=-60.0,
    out_dir="/content",
    figsize=(6.4, 4.8),
    info_fontsize=9,
)

list(animate_all_paulis_3d(K3D, axes=("x", "y", "z")))

Saved: /content/wigner3d_pauli_x.gif

Saved: /content/wigner3d_pauli_y.gif

Saved: /content/wigner3d_pauli_z.gif

# @title Cell 6 — 3D Wigner animation (custom axis–angle; colors-only, no geometry edits)
# PEP 8/257 compliant. Rotates Bloch vector about an arbitrary axis by a target angle.

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Tuple

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation, PillowWriter
from matplotlib.colors import Normalize
from matplotlib.cm import ScalarMappable
from IPython.display import Image, display


@dataclass
class Knobs3DCustom:
    """User controls for 3D custom axis–angle Wigner animation."""
    axis_vec: Tuple[float, float, float] = (0.0, 1.0, 0.0)  # default: y-axis
    angle_deg: float = 45.0
    initial_state: str = "+z"
    custom_r: Tuple[float, float, float] = (0.0, 0.0, 1.0)
    frames: int = 24
    n_theta: int = 64
    n_phi: int = 128
    fps: int = 20
    elev: float = 25.0
    azim: float = -60.0
    out_path: str = "/content/wigner3d_custom_y45.gif"
    figsize: Tuple[float, float] = (6.4, 4.8)
    info_fontsize: int = 9


# ------------- helpers (self-contained; safe for this cell alone) -------------
def _bloch_vector(label: str, custom=(0.0, 0.0, 1.0)) -> np.ndarray:
    mapping = {
        "+z": (0.0, 0.0, 1.0), "-z": (0.0, 0.0, -1.0),
        "+x": (1.0, 0.0, 0.0), "-x": (-1.0, 0.0, 0.0),
        "+y": (0.0, 1.0, 0.0), "-y": (0.0, -1.0, 0.0),
    }
    r = np.asarray(mapping.get(label.lower(), custom), float)
    nrm = np.linalg.norm(r) or 1.0
    return r / nrm


def _rotation_matrix_axis(axis_vec: Tuple[float, float, float], angle: float) -> np.ndarray:
    v = np.asarray(axis_vec, float)
    nrm = np.linalg.norm(v)
    if nrm == 0:
        raise ValueError("axis_vec must be nonzero")
    k = v / nrm
    K = np.array([[0.0, -k[2], k[1]], [k[2], 0.0, -k[0]], [-k[1], k[0], 0.0]], float)
    I = np.eye(3)
    c, s = np.cos(angle), np.sin(angle)
    return c * I + (1 - c) * np.outer(k, k) + s * K


def _n_grid(n_theta: int, n_phi: int):
    theta = np.linspace(0.0, np.pi, n_theta)
    phi = np.linspace(0.0, 2.0 * np.pi, n_phi)
    th, ph = np.meshgrid(theta, phi, indexing="ij")
    nx = np.sin(th) * np.cos(ph)
    ny = np.sin(th) * np.sin(ph)
    nz = np.cos(th)
    n = np.stack((nx, ny, nz), axis=0)
    return th, ph, n


def _wigner_qubit(r_vec: np.ndarray, n: np.ndarray) -> np.ndarray:
    return 0.5 + (np.sqrt(3.0) / 2.0) * (r_vec.reshape(3, 1, 1) * n).sum(axis=0)


def animate_wigner_custom_3d(k: Knobs3DCustom) -> Path:
    """Animate a 3D sphere with colors driven by Wigner under axis–angle rotation."""
    r0 = _bloch_vector(k.initial_state, k.custom_r)
    th, ph, n = _n_grid(k.n_theta, k.n_phi)

    X = np.sin(th) * np.cos(ph)
    Y = np.sin(th) * np.sin(ph)
    Z = np.cos(th)

    w_min = 0.5 - (np.sqrt(3.0) / 2.0)
    w_max = 0.5 + (np.sqrt(3.0) / 2.0)
    norm = Normalize(vmin=w_min, vmax=w_max)
    cmap = plt.get_cmap()
    sm = ScalarMappable(norm=norm, cmap=cmap)
    sm.set_array([])

    fig = plt.figure(figsize=k.figsize, constrained_layout=True)
    ax = fig.add_subplot(111, projection="3d")
    ax.view_init(elev=k.elev, azim=k.azim)
    if hasattr(ax, "set_box_aspect"):
        ax.set_box_aspect((1, 1, 1))

    # Initial facecolors
    angle0 = 0.0
    r_init = _rotation_matrix_axis(k.axis_vec, angle0) @ r0
    W0 = _wigner_qubit(r_init, n)
    FC0 = cmap(norm(W0[:-1, :-1]))  # (M-1,N-1,4)

    surf = ax.plot_surface(
        X, Y, Z,
        facecolors=FC0,
        rstride=1, cstride=1,
        antialiased=False, linewidth=0,
        shade=False,
    )

    cbar = fig.colorbar(sm, ax=ax, fraction=0.046, pad=0.04)
    cbar.set_label("Wigner value")

    ax.set_xlabel("x")
    ax.set_ylabel("y")
    ax.set_zlabel("z")
    title_txt = ax.text2D(
        0.5, 0.98, "Wigner - custom angle",
        transform=ax.transAxes, ha="center", va="top", fontsize=11,
        bbox=dict(boxstyle="round,pad=0.3", fc="white", alpha=0.75, ec="none"),
        zorder=6,
    )
    info_txt = ax.text2D(
        0.02, 0.96, "", transform=ax.transAxes, ha="left", va="top",
        fontsize=k.info_fontsize,
        bbox=dict(boxstyle="round,pad=0.3", fc="white", alpha=0.75, ec="none"),
        zorder=6,
    )

    def update(frame: int):
        ang = (frame / (k.frames - 1)) * np.deg2rad(k.angle_deg)
        r_now = _rotation_matrix_axis(k.axis_vec, ang) @ r0
        W_now = _wigner_qubit(r_now, n)
        FC_now = cmap(norm(W_now[:-1, :-1]))          # (M-1,N-1,4)
        surf.set_facecolors(FC_now.reshape(-1, 4))    # in-place update
        info_txt.set_text(
            f"axis = ({k.axis_vec[0]:.2f}, {k.axis_vec[1]:.2f}, {k.axis_vec[2]:.2f})\n"
            f"θ = {np.rad2deg(ang):.1f}° / {k.angle_deg:.1f}°"
        )
        return (surf, info_txt, title_txt)

    anim = FuncAnimation(fig, update, frames=k.frames, interval=60, blit=False)
    out = Path(k.out_path)
    out.parent.mkdir(parents=True, exist_ok=True)
    anim.save(out.as_posix(), writer=PillowWriter(fps=k.fps),
              savefig_kwargs=dict(facecolor="white"))
    plt.close(fig)
    return out


# ----------------------------- Example run ------------------------------------
K3C = Knobs3DCustom(
    axis_vec=(0.0, 1.0, 0.0), angle_deg=45.0,
    initial_state="+z", out_path="/content/wigner3d_custom_y45.gif",
)
saved_custom = animate_wigner_custom_3d(K3C)
display(Image(filename=saved_custom.as_posix()))
print("Saved:", saved_custom)