scatter_error.py

"""Scatter plots with asymmetric error bars.

Two functions:

* ``errorbar_plot``       — single series with x/y error bars.
* ``multi_errorbar_plot`` — multiple series overlaid on shared axes.

Example
-------
>>> import numpy as np
>>> x = np.linspace(0, 10, 12)
>>> y = 2.5 * x + 1
>>> yerr = np.random.uniform(0.5, 1.5, 12)
>>> fig = errorbar_plot(x, y, yerr=yerr, xlabel="V", ylabel="I (mA)")
"""

from __future__ import annotations

import matplotlib.pyplot as plt
import numpy as np

from .style import COLORS, PALETTE, FIGSIZES, GRID_ALPHA, GRID_LINEWIDTH, GRID_COLOR, Z_ORDER, savefig, apply_grid
from .utils import validate_arrays


def errorbar_plot(
    x: np.ndarray,
    y: np.ndarray,
    *,
    xerr: np.ndarray | float | None = None,
    yerr: np.ndarray | float | None = None,
    xlabel: str = "x",
    ylabel: str = "y",
    title: str | None = None,
    label: str | None = None,
    color: str | None = None,
    marker: str = "o",
    marker_size: float = 4,
    marker_edgecolor: str = "white",
    marker_edgewidth: float = 0.3,
    capsize: float = 2.5,
    capthick: float = 0.6,
    error_linewidth: float = 0.7,
    show_grid: bool = True,
    grid_alpha: float = GRID_ALPHA,
    grid_linewidth: float = GRID_LINEWIDTH,
    grid_color: str = GRID_COLOR,
    legend_loc: str = "best",
    figsize: tuple[float, float] = FIGSIZES["single"],
) -> plt.Figure:
    """Scatter plot with optional asymmetric error bars on x and/or y.

    Error arrays can be:
    * ``None``     — no error bars on that axis.
    * ``float``    — symmetric constant error.
    * ``1-D array`` — symmetric per-point error.
    * ``2×N array`` — asymmetric: ``[lower_errors, upper_errors]``.

    Parameters
    ----------
    x, y : np.ndarray
        Data coordinates.
    xerr, yerr : np.ndarray, float, or None
        Error magnitudes (see description above).
    xlabel, ylabel : str
        Axis labels.
    title : str or None
        Subplot title.
    label : str or None
        Legend label.
    color : str or None
        Marker / line colour.  Falls back to ``PALETTE[0]``.
    marker : str
        Matplotlib marker character.
    marker_size : float
        Marker diameter.
    marker_edgecolor : str
        Marker border colour.
    marker_edgewidth : float
        Marker border width.
    capsize : float
        Length of error-bar caps.
    capthick : float
        Thickness of error-bar caps.
    error_linewidth : float
        Thickness of error-bar lines.
    show_grid : bool
        Show background grid.
    grid_alpha, grid_linewidth, grid_color
        Grid styling.
    legend_loc : str
        Legend location.
    figsize : tuple[float, float]
        Figure size in inches.

    Returns
    -------
    plt.Figure
    """
    validate_arrays(x, y, names=["x", "y"])
    fig, ax = plt.subplots(figsize=figsize)
    c = color or PALETTE[0]

    ax.errorbar(
        x, y, xerr=xerr, yerr=yerr,
        fmt=marker, markersize=marker_size, markerfacecolor=c,
        markeredgecolor=marker_edgecolor, markeredgewidth=marker_edgewidth,
        color=c, ecolor=c, elinewidth=error_linewidth, capsize=capsize,
        capthick=capthick, zorder=Z_ORDER["data"], label=label,
    )

    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    if title:
        ax.set_title(title)
    apply_grid(ax, show=show_grid, alpha=grid_alpha, linewidth=grid_linewidth, color=grid_color)
    ax.minorticks_on()
    if label:
        ax.legend(loc=legend_loc, frameon=False)

    fig.tight_layout(pad=0.3)
    return fig


def multi_errorbar_plot(
    datasets: list[dict],
    *,
    xlabel: str = "x",
    ylabel: str = "y",
    title: str | None = None,
    marker_size: float = 3.5,
    capsize: float = 2.2,
    capthick: float = 0.5,
    error_linewidth: float = 0.6,
    show_grid: bool = True,
    grid_alpha: float = GRID_ALPHA,
    grid_linewidth: float = GRID_LINEWIDTH,
    grid_color: str = GRID_COLOR,
    legend_loc: str = "best",
    figsize: tuple[float, float] = FIGSIZES["wide_large"],
) -> plt.Figure:
    """Multiple error-bar series on shared axes.

    Each dict in *datasets* may contain:

    ========== ============ =============================================
    Key        Required     Description
    ========== ============ =============================================
    ``x``      Yes          1-D array.
    ``y``      Yes          1-D array.
    ``xerr``   No           Error on x (see ``errorbar_plot`` for formats).
    ``yerr``   No           Error on y.
    ``label``  No           Legend label.
    ``color``  No           Hex colour (auto-cycled if omitted).
    ========== ============ =============================================

    Parameters
    ----------
    datasets : list of dict
        Series definitions (see table above).
    xlabel, ylabel : str
        Axis labels.
    title : str or None
        Subplot title.
    marker_size : float
        Marker diameter for all series.
    capsize : float
        Error-bar cap length.
    capthick : float
        Error-bar cap thickness.
    error_linewidth : float
        Error-bar line width.
    show_grid : bool
        Show background grid.
    grid_alpha, grid_linewidth, grid_color
        Grid styling.
    legend_loc : str
        Legend location.
    figsize : tuple[float, float]
        Figure size in inches.

    Returns
    -------
    plt.Figure
    """
    fig, ax = plt.subplots(figsize=figsize)

    for i, ds in enumerate(datasets):
        validate_arrays(ds["x"], ds["y"], names=[f"datasets[{i}]['x']", f"datasets[{i}]['y']"])
        c = ds.get("color", PALETTE[i % len(PALETTE)])
        ax.errorbar(
            ds["x"], ds["y"],
            xerr=ds.get("xerr"), yerr=ds.get("yerr"),
            fmt="o", markersize=marker_size, markerfacecolor=c,
            markeredgecolor="white", markeredgewidth=0.3,
            color=c, ecolor=c, elinewidth=error_linewidth, capsize=capsize,
            capthick=capthick, zorder=Z_ORDER["data"],
            label=ds.get("label", f"Series {i+1}"),
        )

    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    if title:
        ax.set_title(title)
    apply_grid(ax, show=show_grid, alpha=grid_alpha, linewidth=grid_linewidth, color=grid_color)
    ax.minorticks_on()
    ax.legend(loc=legend_loc, frameon=False)

    fig.tight_layout(pad=0.3)
    return fig