internal-docs/arxiv-style/fig-scripts/draw_transformer_lower_half_axes.py

#!/usr/bin/env python3
"""
Transformer section lower-half visuals WITH AXES ONLY:
- Axes spines visible
- NO numbers (tick labels hidden)
- NO words (axis labels removed)
- Transparent background
- One SVG output

Run:
  uv run python draw_transformer_lower_half_axes_only.py --out ./assets/transformer_lower_half_axes_only.svg
"""

from __future__ import annotations

import argparse
from dataclasses import dataclass
from pathlib import Path

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle


@dataclass
class Params:
    seed: int = 7
    seconds: float = 10.0
    fs: int = 300

    # Continuous channels
    n_curves: int = 3
    curve_lw: float = 2.4

    # Discrete bars
    n_bins: int = 40
    bar_height: float = 0.55   # fraction of the discrete-axis y-range
    bar_gap: float = 0.08      # fraction of bar width

    # Figure size
    width_in: float = 6.6
    height_in: float = 2.6


def _smooth(x: np.ndarray, win: int) -> np.ndarray:
    win = max(3, int(win) | 1)  # odd
    k = np.ones(win, dtype=float)
    k /= k.sum()
    return np.convolve(x, k, mode="same")


def make_continuous_curves(p: Params) -> tuple[np.ndarray, np.ndarray]:
    rng = np.random.default_rng(p.seed)
    T = int(p.seconds * p.fs)
    t = np.linspace(0, p.seconds, T, endpoint=False)

    Y = []
    base_freqs = [0.12, 0.09, 0.15]
    mid_freqs = [0.65, 0.85, 0.75]

    for i in range(p.n_curves):
        f1 = base_freqs[i % len(base_freqs)]
        f2 = mid_freqs[i % len(mid_freqs)]
        ph = rng.uniform(0, 2 * np.pi)

        y = (
            1.00 * np.sin(2 * np.pi * f1 * t + ph)
            + 0.35 * np.sin(2 * np.pi * f2 * t + 0.7 * ph)
        )

        bumps = np.zeros_like(t)
        for _ in range(2):
            mu = rng.uniform(0.8, p.seconds - 0.8)
            sig = rng.uniform(0.35, 0.75)
            bumps += np.exp(-0.5 * ((t - mu) / sig) ** 2)
        y += 0.55 * bumps

        noise = _smooth(rng.normal(0, 1, size=T), win=int(p.fs * 0.04))
        y += 0.12 * noise

        y = (y - y.mean()) / (y.std() + 1e-9)
        y *= 0.42
        Y.append(y)

    return t, np.vstack(Y)


def make_discrete_bars(p: Params) -> np.ndarray:
    rng = np.random.default_rng(p.seed + 123)
    n = p.n_bins

    ids = np.zeros(n, dtype=int)
    cur = rng.integers(0, 8)
    for i in range(n):
        if i == 0 or rng.random() < 0.25:
            cur = rng.integers(0, 8)
        ids[i] = cur
    return ids


def _axes_only(ax: plt.Axes) -> None:
    """Keep spines (axes lines), remove all ticks/labels/words."""
    # No labels
    ax.set_xlabel("")
    ax.set_ylabel("")
    ax.set_title("")

    # Keep spines as the only axes element
    for side in ("top", "right", "bottom", "left"):
        ax.spines[side].set_visible(True)

    # Remove tick marks and tick labels entirely
    ax.set_xticks([])
    ax.set_yticks([])
    ax.tick_params(
        axis="both",
        which="both",
        bottom=False,
        left=False,
        top=False,
        right=False,
        labelbottom=False,
        labelleft=False,
    )

    # No grid
    ax.grid(False)


def draw_transformer_lower_half_svg(out_path: Path, p: Params) -> None:
    fig = plt.figure(figsize=(p.width_in, p.height_in), dpi=200)
    fig.patch.set_alpha(0.0)

    # Two axes sharing x (top curves, bottom bars)
    ax_curves = fig.add_axes([0.10, 0.38, 0.86, 0.56])
    ax_bars = fig.add_axes([0.10, 0.14, 0.86, 0.18], sharex=ax_curves)

    ax_curves.patch.set_alpha(0.0)
    ax_bars.patch.set_alpha(0.0)

    # Data
    t, Y = make_continuous_curves(p)
    ids = make_discrete_bars(p)

    # Top: continuous curves
    curve_colors = ["#1f77b4", "#ff7f0e", "#2ca02c"]  # blue / orange / green
    for i in range(Y.shape[0]):
        ax_curves.plot(t, Y[i], linewidth=p.curve_lw, color=curve_colors[i % len(curve_colors)])

    ymin, ymax = float(Y.min()), float(Y.max())
    ypad = 0.10 * (ymax - ymin + 1e-9)
    ax_curves.set_xlim(t[0], t[-1])
    ax_curves.set_ylim(ymin - ypad, ymax + ypad)

    # Bottom: discrete bars (colored strip)
    bar_palette = [
        "#e41a1c", "#377eb8", "#4daf4a", "#984ea3",
        "#ff7f00", "#ffff33", "#a65628", "#f781bf",
    ]

    x0, x1 = t[0], t[-1]
    total = x1 - x0
    n = len(ids)
    bin_w = total / n
    gap = p.bar_gap * bin_w

    ax_bars.set_xlim(x0, x1)
    ax_bars.set_ylim(0, 1)

    bar_y = (1 - p.bar_height) / 2
    for i, cat in enumerate(ids):
        left = x0 + i * bin_w + gap / 2
        width = bin_w - gap
        color = bar_palette[int(cat) % len(bar_palette)]
        ax_bars.add_patch(Rectangle((left, bar_y), width, p.bar_height, facecolor=color, edgecolor="none"))

    # Apply "axes only" styling (no numbers/words)
    _axes_only(ax_curves)
    _axes_only(ax_bars)

    out_path.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(out_path, format="svg", transparent=True, bbox_inches="tight", pad_inches=0.0)
    plt.close(fig)


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", type=Path, default=Path("transformer_lower_half_axes_only.svg"))
    ap.add_argument("--seed", type=int, default=7)
    ap.add_argument("--seconds", type=float, default=10.0)
    ap.add_argument("--fs", type=int, default=300)
    ap.add_argument("--bins", type=int, default=40)
    ap.add_argument("--curves", type=int, default=3)
    args = ap.parse_args()

    p = Params(seed=args.seed, seconds=args.seconds, fs=args.fs, n_bins=args.bins, n_curves=args.curves)
    draw_transformer_lower_half_svg(args.out, p)
    print(f"Wrote: {args.out}")


if __name__ == "__main__":
    main()