mgplot.fill_between_plot

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