lars
36ae18e880
Introduces a new plotstyle package providing consistent, presentation/thesis-ready matplotlib figures: a KIT corporate-design color palette (categorical, sequential, diverging, status), a theme applied via use() (KIT-black bottom spine only, left/bottom ticks, horizontal gridlines, left-aligned titles, LaTeX text in Latin Modern Sans), new_figure() with size presets and figure-level title/params subtitles, a same-size colorbar() helper, style_legend() and panel_label() building blocks, and savefig(). Ships as an optional "plotting" extra so core gallery installs stay lightweight, with a full test suite and a runnable example notebook. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
128 lines
4.6 KiB
Python
128 lines
4.6 KiB
Python
"""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 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
|