mgplot.fill_between_plot

  1"""Plot a filled region between two bounds."""
  2
  3from typing import Final, NotRequired, Unpack
  4
  5from matplotlib.axes import Axes
  6from pandas import DataFrame
  7
  8from mgplot.axis_utils import map_periodindex, set_labels
  9from mgplot.keyword_checking import BaseKwargs, report_kwargs, validate_kwargs
 10from mgplot.settings import get_setting
 11from mgplot.utilities import check_clean_timeseries, constrain_data, get_axes
 12
 13# --- constants
 14ME: Final[str] = "fill_between_plot"
 15REQUIRED_COLUMNS: Final[int] = 2
 16DEFAULT_COLOR: Final[str] = "steelblue"
 17DEFAULT_ALPHA: Final[float] = 0.3
 18
 19
 20class FillBetweenKwargs(BaseKwargs):
 21    """Keyword arguments for the fill_between_plot function."""
 22
 23    ax: NotRequired[Axes | None]
 24    color: NotRequired[str]
 25    alpha: NotRequired[float]
 26    label: NotRequired[str | None]
 27    linewidth: NotRequired[float]
 28    edgecolor: NotRequired[str | None]
 29    zorder: NotRequired[int | float]
 30    plot_from: NotRequired[int | None]
 31    max_ticks: NotRequired[int]
 32
 33
 34def fill_between_plot(data: DataFrame, **kwargs: Unpack[FillBetweenKwargs]) -> Axes:
 35    """Plot a filled region between lower and upper bounds.
 36
 37    Args:
 38        data: DataFrame - A two-column DataFrame with PeriodIndex.
 39              The first column is the lower bound, the second is the upper bound.
 40        kwargs: FillBetweenKwargs - keyword arguments for the plot.
 41
 42    Returns:
 43        Axes - matplotlib Axes object.
 44
 45    Raises:
 46        TypeError: If data is not a DataFrame.
 47        ValueError: If data does not have exactly two columns.
 48
 49    """
 50    # --- validate inputs
 51    report_kwargs(caller=ME, **kwargs)
 52    validate_kwargs(schema=FillBetweenKwargs, caller=ME, **kwargs)
 53
 54    if not isinstance(data, DataFrame):
 55        raise TypeError(f"data must be a DataFrame for {ME}()")
 56
 57    if len(data.columns) != REQUIRED_COLUMNS:
 58        raise ValueError(f"data must have exactly two columns for {ME}(), got {len(data.columns)}")
 59
 60    # --- check and constrain data
 61    data = check_clean_timeseries(data, ME)
 62    data, kwargs_d = constrain_data(data, **kwargs)
 63
 64    # --- handle PeriodIndex conversion
 65    saved_pi = map_periodindex(data)
 66    if saved_pi is not None:
 67        data = saved_pi[0]
 68
 69    # --- get axes
 70    axes, kwargs_d = get_axes(**kwargs_d)
 71
 72    if data.empty or data.isna().all().all():
 73        print(f"Warning: No data to plot in {ME}().")
 74        return axes
 75
 76    # --- extract bounds
 77    lower = data.iloc[:, 0]
 78    upper = data.iloc[:, 1]
 79
 80    # --- extract plot arguments
 81    color = kwargs_d.get("color", DEFAULT_COLOR)
 82    alpha = kwargs_d.get("alpha", DEFAULT_ALPHA)
 83    label = kwargs_d.get("label", None)
 84    linewidth = kwargs_d.get("linewidth", 0)
 85    edgecolor = kwargs_d.get("edgecolor", None)
 86    zorder = kwargs_d.get("zorder", None)
 87
 88    # --- plot
 89    axes.fill_between(
 90        data.index,
 91        lower,
 92        upper,
 93        color=color,
 94        alpha=alpha,
 95        label=label,
 96        linewidth=linewidth,
 97        edgecolor=edgecolor,
 98        zorder=zorder,
 99    )
100
101    # --- set axis labels
102    if saved_pi is not None:
103        set_labels(axes, saved_pi[1], kwargs_d.get("max_ticks", get_setting("max_ticks")))
104
105    return axes