alnn_rewrite/analysis/plotting.py

"""Shared figure styling and output helpers.

Centralising this keeps every analysis visually consistent (same palette, same
grid weight, same file formats) and means an analysis body contains only its own
logic, not boilerplate.
"""

import json
import math
import pathlib as pl
from typing import Any, Dict, List, Literal, Sequence, Tuple

import matplotlib

matplotlib.use("Agg")  # headless: analysis runs in a background worker thread

import matplotlib.pyplot as plt  # noqa: E402
from matplotlib.axes import Axes  # noqa: E402
from matplotlib.figure import Figure  # noqa: E402

#: Raster resolution for the .png copies.
DPI = 150

def style_axes(
    ax: Axes, *, grid_axis: Literal["both", "x", "y"] = "both"
) -> None:
    """Apply the shared axis style: recessive grid, no top/right spines."""
    ax.grid(True, axis=grid_axis, color="#000000", alpha=0.08, linewidth=0.7)
    for spine in ("top", "right"):
        ax.spines[spine].set_visible(False)


def panel_grid(
    n_panels: int,
    *,
    n_cols: int = 2,
    panel_size: Tuple[float, float] = (5.0, 3.4),
) -> Tuple[Figure, List[Axes]]:
    """Create a grid of panels sized to the panel count.

    Returns the figure and a flat list of exactly ``n_panels`` axes; any surplus
    axes in the final row are removed so no empty frames are shown.
    """
    n_panels = max(1, n_panels)
    n_cols = max(1, min(n_cols, n_panels))
    n_rows = math.ceil(n_panels / n_cols)
    # constrained layout keeps axis labels from colliding with the titles of the
    # row beneath them once several panels are stacked.
    fig, axes = plt.subplots(
        n_rows,
        n_cols,
        figsize=(panel_size[0] * n_cols, panel_size[1] * n_rows),
        squeeze=False,
        layout="constrained",
    )
    flat = [ax for row in axes for ax in row]
    for extra in flat[n_panels:]:
        extra.remove()
    return fig, flat[:n_panels]


def add_legend(fig: Figure, axes: Sequence[Axes], *, ncol: int = 4) -> None:
    """Attach one shared, de-duplicated legend below the figure.

    A legend is always present when more than one series is drawn, so identity is
    never carried by colour alone.
    """
    handles: List[Any] = []
    labels: List[str] = []
    for ax in axes:
        for handle, label in zip(*ax.get_legend_handles_labels()):
            if label not in labels:
                handles.append(handle)
                labels.append(label)
    if len(labels) < 2:
        return
    fig.legend(
        handles,
        labels,
        loc="outside lower center" if fig.get_layout_engine() else "lower center",
        ncol=min(ncol, len(labels)),
        frameon=False,
    )


#: Figures and the numbers behind them live in separate subdirectories of the
#: analysis output dir, so the plots can be browsed without wading through JSON.
PLOTS_SUBDIR = "plots"
DATA_SUBDIR = "data"


def plots_dir(out_dir: pl.Path) -> pl.Path:
    """``<analysis>/plots`` -- every figure lands here."""
    path = out_dir / PLOTS_SUBDIR
    path.mkdir(parents=True, exist_ok=True)
    return path


def data_dir(out_dir: pl.Path) -> pl.Path:
    """``<analysis>/data`` -- the numbers behind the figures."""
    path = out_dir / DATA_SUBDIR
    path.mkdir(parents=True, exist_ok=True)
    return path


def save_figure(fig: Figure, out_dir: pl.Path, stem: str) -> pl.Path:
    """Write ``plots/stem.png`` (raster) and ``plots/stem.pdf``; close the figure."""
    target = plots_dir(out_dir)
    png_path = target / f"{stem}.png"
    fig.savefig(png_path, dpi=DPI, bbox_inches="tight")
    fig.savefig(target / f"{stem}.pdf", bbox_inches="tight")
    plt.close(fig)
    return png_path


def save_json(payload: Dict[str, Any], out_dir: pl.Path, stem: str) -> pl.Path:
    """Persist the numbers behind a figure so plots can be rebuilt or reused."""
    path = data_dir(out_dir) / f"{stem}.json"
    path.write_text(json.dumps(payload, indent=2, default=_json_default))
    return path


def _json_default(value: Any) -> Any:
    """Fallback encoder for numpy scalars/arrays."""
    if hasattr(value, "tolist"):
        return value.tolist()
    if hasattr(value, "item"):
        return value.item()
    return str(value)