distribution.py

"""Distribution plots: box, violin, and histogram.

Three functions:

* ``box_plot``        — box-and-whisker for one or more groups.
* ``violin_plot``     — violin (kernel-density) shape across groups.
* ``histogram_plot``  — frequency histogram with optional KDE overlay.

Example
-------
>>> import numpy as np
>>> from academic_plot import box_plot, savefig
>>> groups = [np.random.normal(5, 1, 50), np.random.normal(7, 1.2, 50)]
>>> fig = box_plot(groups, ["Control", "Treatment"], ylabel="Response (mV)")
>>> savefig(fig, "box_demo")
"""

from __future__ import annotations

from typing import Sequence

import matplotlib.pyplot as plt
import numpy as np

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


def box_plot(
    data: Sequence[np.ndarray],
    labels: Sequence[str],
    *,
    ylabel: str = "Value",
    title: str | None = None,
    box_widths: float = 0.45,
    box_alpha: float = 0.55,
    median_color: str = "#1A1A1A",
    median_linewidth: float = 1.0,
    whisker_linewidth: float = 0.6,
    cap_linewidth: float = 0.6,
    flier_markersize: float = 3,
    flier_markerfacecolor: str = "#888888",
    show_grid: bool = True,
    grid_alpha: float = GRID_ALPHA,
    grid_linewidth: float = GRID_LINEWIDTH,
    grid_color: str = GRID_COLOR,
    figsize: tuple[float, float] = FIGSIZES["single"],
) -> plt.Figure:
    """Box-and-whisker plot for one or more groups.

    Parameters
    ----------
    data : sequence of np.ndarray
        One 1-D array per group.
    labels : sequence of str
        Group names (x-axis tick labels).
    ylabel : str
        y-axis label.
    title : str or None
        Subplot title.
    box_widths : float
        Width of each box.
    box_alpha : float
        Fill opacity of boxes.
    median_color : str
        Colour of the median line inside each box.
    median_linewidth : float
        Width of the median line.
    whisker_linewidth : float
        Width of the whisker lines.
    cap_linewidth : float
        Width of the whisker caps.
    flier_markersize : float
        Size of outlier markers.
    flier_markerfacecolor : str
        Colour of outlier markers.
    show_grid : bool
        Show horizontal grid lines.
    grid_alpha, grid_linewidth, grid_color
        Grid styling.
    figsize : tuple[float, float]
        Figure size in inches.

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

    # patch_artist=True allows per-box facecolour customisation
    bp = ax.boxplot(
        data, labels=labels, patch_artist=True, widths=box_widths,
        medianprops=dict(color=median_color, linewidth=median_linewidth),
        whiskerprops=dict(linewidth=whisker_linewidth),
        capprops=dict(linewidth=cap_linewidth),
        flierprops=dict(
            marker="o", markersize=flier_markersize,
            markerfacecolor=flier_markerfacecolor, markeredgewidth=0.3,
        ),
    )

    # Style each box with the palette colours
    for patch, c in zip(bp["boxes"], PALETTE[: len(data)]):
        patch.set_facecolor(c)
        patch.set_alpha(box_alpha)
        patch.set_edgecolor("#333")
        patch.set_linewidth(0.6)

    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, axis="y")

    fig.tight_layout(pad=0.3)
    return fig


def violin_plot(
    data: Sequence[np.ndarray],
    labels: Sequence[str],
    *,
    ylabel: str = "Value",
    title: str | None = None,
    violin_alpha: float = 0.45,
    violin_edgecolor: str = "#333",
    violin_edgewidth: float = 0.6,
    median_color: str = "#1A1A1A",
    median_linewidth: float = 0.8,
    show_grid: bool = True,
    grid_alpha: float = GRID_ALPHA,
    grid_linewidth: float = GRID_LINEWIDTH,
    grid_color: str = GRID_COLOR,
    figsize: tuple[float, float] = FIGSIZES["single"],
) -> plt.Figure:
    """Violin plot showing kernel-density shape across groups.

    Parameters
    ----------
    data : sequence of np.ndarray
        One 1-D array per group.
    labels : sequence of str
        Group names.
    ylabel : str
        y-axis label.
    title : str or None
        Subplot title.
    violin_alpha : float
        Fill opacity of each violin.
    violin_edgecolor : str
        Outline colour of each violin.
    violin_edgewidth : float
        Outline thickness.
    median_color : str
        Colour of the median indicator line.
    median_linewidth : float
        Width of the median indicator.
    show_grid : bool
        Show horizontal grid lines.
    grid_alpha, grid_linewidth, grid_color
        Grid styling.
    figsize : tuple[float, float]
        Figure size in inches.

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

    # showmedians=True draws a horizontal line at the median of each violin
    parts = ax.violinplot(data, showmeans=False, showmedians=True,
                          showextrema=True)

    # Style each violin body with palette colours
    for i, pc in enumerate(parts["bodies"]):
        pc.set_facecolor(PALETTE[i % len(PALETTE)])
        pc.set_alpha(violin_alpha)
        pc.set_edgecolor(violin_edgecolor)
        pc.set_linewidth(violin_edgewidth)

    # Style internal indicators
    parts["cmedians"].set_color(median_color)
    parts["cmedians"].set_linewidth(median_linewidth)
    parts["cbars"].set_linewidth(0.5)
    parts["cmins"].set_linewidth(0.5)
    parts["cmaxes"].set_linewidth(0.5)

    # violinplot uses 1-based integer positions
    ax.set_xticks(np.arange(1, len(labels) + 1))
    ax.set_xticklabels(labels)
    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, axis="y")

    fig.tight_layout(pad=0.3)
    return fig


def histogram_plot(
    data: np.ndarray,
    *,
    bins: int = 20,
    xlabel: str = "Value",
    ylabel: str = "Count",
    title: str | None = None,
    color: str | None = None,
    hist_alpha: float = 0.65,
    edgecolor: str = "white",
    edgewidth: float = 0.4,
    show_kde: bool = False,
    kde_color: str = "#CB181D",
    kde_linewidth: float = 1.2,
    kde_ylabel: str = "Density",
    show_grid: bool = True,
    grid_alpha: float = GRID_ALPHA,
    grid_linewidth: float = GRID_LINEWIDTH,
    grid_color: str = GRID_COLOR,
    figsize: tuple[float, float] = FIGSIZES["single"],
) -> plt.Figure:
    """Histogram with optional kernel density estimate (KDE) overlay.

    Parameters
    ----------
    data : np.ndarray
        1-D array of observations.
    bins : int
        Number of histogram bins.
    xlabel : str
        x-axis label.
    ylabel : str
        y-axis label (for histogram counts).
    title : str or None
        Subplot title.
    color : str or None
        Histogram bar colour.  Falls back to ``PALETTE[0]``.
    hist_alpha : float
        Histogram bar opacity.
    edgecolor : str
        Colour of bar borders.
    edgewidth : float
        Width of bar borders.
    show_kde : bool
        If ``True``, overlay a kernel density estimate curve on a secondary
        y-axis.
    kde_color : str
        Colour of the KDE line.
    kde_linewidth : float
        Thickness of the KDE line.
    kde_ylabel : str
        Label for the secondary y-axis (density).
    show_grid : bool
        Show horizontal grid lines.
    grid_alpha, grid_linewidth, grid_color
        Grid styling.
    figsize : tuple[float, float]
        Figure size in inches.

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

    # Draw the histogram
    ax.hist(
        data, bins=bins, color=c, edgecolor=edgecolor,
        linewidth=edgewidth, alpha=hist_alpha, zorder=Z_ORDER["data"],
    )

    # Optional KDE on a secondary y-axis so the density scale is correct
    if show_kde:
        from scipy.stats import gaussian_kde
        x_kde = np.linspace(data.min(), data.max(), 300)
        kde = gaussian_kde(data)
        ax2 = ax.twinx()
        ax2.plot(x_kde, kde(x_kde), color=kde_color, linewidth=kde_linewidth,
                 label="KDE")
        ax2.set_ylabel(kde_ylabel, fontsize=8)
        ax2.tick_params(labelsize=7)
        ax2.legend(loc="upper right", frameon=False)

    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, axis="y")

    fig.tight_layout(pad=0.3)
    return fig