ETPlot/plotstyle/figures.py

"""Figure creation, colorbar, and saving helpers."""

from __future__ import annotations

from pathlib import Path
from typing import Mapping, Sequence, Union

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.axes import Axes
from matplotlib.colorbar import Colorbar
from matplotlib.figure import Figure
from mpl_toolkits.axes_grid1 import make_axes_locatable

from .colors import INK

FIGSIZES = {
    "thesis-single": (6.0, 4.0),
    "thesis-wide": (8.0, 4.5),
    "slide-16x9": (10.0, 5.625),
    "square": (5.0, 5.0),
}


def _style_ticks(ax: Axes) -> None:
    """Major tick labels read as primary ("black") ink; minor as muted grey.

    rcParams only expose a single labelcolor per axis (no major/minor split),
    so this distinction has to be applied per-Axes in code.
    """
    ax.tick_params(axis="both", which="major", labelcolor=INK["primary"])
    ax.tick_params(axis="both", which="minor", labelcolor=INK["muted"])


def _format_params(params: Mapping) -> str:
    return " | ".join(f"{key}: {value}" for key, value in params.items())


def _set_figure_title(fig: Figure, title: Union[str, None], params: Union[Mapping, None]) -> None:
    # A single Text artist gets one color under matplotlib's usetex rendering
    # (dvipng rasterizes it as one greyscale mask, tinted uniformly — any
    # in-source \color/\textcolor is ignored), so the subtitle is set apart
    # from the title by size (\small) only, not color.
    lines = []
    if title is not None:
        lines.append(r"\textbf{" + title + "}")
    if params:
        lines.append(r"{\small " + _format_params(params) + "}")
    if lines:
        fig.suptitle("\n".join(lines), x=0.0, ha="left")


def new_figure(
    preset: Union[str, tuple] = "thesis-single",
    *,
    title: str = None,
    params: Mapping = None,
    **subplots_kwargs,
):
    """Create a figure/axes pair sized for a named preset or an explicit (w, h) tuple.

    Presets (inches): thesis-single, thesis-wide, slide-16x9, square.

    `title` sets a left-aligned, bold figure-level title (`fig.suptitle`) —
    this is preferred over an axes title even when there's a single axes, so
    it stays consistent for single- and multi-panel figures alike. `params`
    is an optional dict rendered as a smaller, muted subtitle line below the
    title, formatted as "key1: value1 | key2: value2 | ...".
    """
    if isinstance(preset, str):
        try:
            figsize = FIGSIZES[preset]
        except KeyError as exc:
            raise ValueError(
                f"Unknown figure preset {preset!r}. Choose from {sorted(FIGSIZES)} or pass an (w, h) tuple."
            ) from exc
    else:
        figsize = preset

    subplots_kwargs.setdefault("figsize", figsize)
    fig, axes = plt.subplots(**subplots_kwargs)

    for ax in [axes] if isinstance(axes, Axes) else np.ravel(axes):
        _style_ticks(ax)

    _set_figure_title(fig, title, params)

    return fig, axes


def no_spines(ax: Union[Axes, np.ndarray]) -> None:
    """Hide every spine on `ax` — pass a single Axes or an array of them.

    `use()` keeps only the bottom spine visible, since most plots have a
    meaningful x baseline. Pixel/bin-indexed plots (`imshow`, `pcolormesh`,
    2D histograms) don't — there's no "zero" the bottom spine anchors to —
    so call this on their Axes instead of leaving the themed bottom spine on
    or hand-rolling `ax.spines[...].set_visible(False)`.
    """
    for single_ax in [ax] if isinstance(ax, Axes) else np.ravel(ax):
        for spine in single_ax.spines.values():
            spine.set_visible(False)


def colorbar(mappable, ax: Axes, size: str = "5%", pad: float = 0.05, **kwargs) -> Colorbar:
    """Add a colorbar matched to `ax`'s actual on-screen size.

    `fig.colorbar(mappable, ax=ax)` sizes the colorbar to the axes' nominal
    bounding box, which is taller than the axes once things like
    `ax.set_aspect("equal")` have visually shrunk it (e.g. a non-square
    `imshow`). This appends a same-size axes via `make_axes_locatable`
    instead, so the colorbar always matches what's actually drawn.
    """
    divider = make_axes_locatable(ax)
    cax = divider.append_axes("right", size=size, pad=pad)
    cb = ax.get_figure().colorbar(mappable, cax=cax, **kwargs)
    # Colorbar draws its own border (a dedicated "outline" spine) that isn't
    # covered by axes.spines.{left,right,top} — without this it'd pick up
    # the bold black bottom-spine color/width from the main theme as a box
    # around the whole colorbar, which reads as a stray, heavier-than-intended
    # edge rather than the plain baseline it's styled to be elsewhere.
    cb.outline.set_visible(False)
    return cb


def savefig(fig: Figure, path: Union[str, Path], formats: Sequence[str] = ("pdf",), dpi: int = 300) -> list[Path]:
    """Save `fig` to `path` once per format, creating parent directories as needed.

    `path` should have no extension — it's appended per format, e.g.
    savefig(fig, "plots/my_plot", formats=("pdf", "png")) writes
    plots/my_plot.pdf and plots/my_plot.png.
    """
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)

    written = []
    for fmt in formats:
        out_path = path.with_suffix(f".{fmt}")
        fig.savefig(out_path, format=fmt, dpi=dpi)
        written.append(out_path)
    return written