mgplot.bar_plot

  1"""Create bar plots using Matplotlib.
  2
  3Note: bar plots in Matplotlib are not the same as bar charts in other
  4libraries. Bar plots are used to represent categorical data with
  5rectangular bars. As a result, bar plots and line plots typically
  6cannot be plotted on the same axes.
  7"""
  8
  9from collections.abc import Callable, Sequence
 10from typing import Any, Final, NotRequired, TypedDict, Unpack
 11
 12import matplotlib.patheffects as pe
 13import numpy as np
 14from matplotlib.axes import Axes
 15from pandas import DataFrame, Period, Series
 16
 17from mgplot.axis_utils import map_periodindex, map_stringindex, set_labels
 18from mgplot.keyword_checking import BaseKwargs, report_kwargs, validate_kwargs
 19from mgplot.settings import DataT, get_setting
 20from mgplot.utilities import (
 21    apply_defaults,
 22    constrain_data,
 23    default_rounding,
 24    get_axes,
 25    get_color_list,
 26)
 27
 28# --- constants
 29ME: Final[str] = "bar_plot"
 30MAX_ANNOTATIONS: Final[int] = 30
 31ADJUSTMENT_FACTOR: Final[float] = 0.02
 32MIN_BAR_WIDTH: Final[float] = 0.0
 33MAX_BAR_WIDTH: Final[float] = 1.0
 34DEFAULT_GROUPED_WIDTH: Final[float] = 0.8
 35DEFAULT_BAR_OFFSET: Final[float] = 0.5
 36DEFAULT_MAX_TICKS: Final[int] = 10
 37
 38
 39class BarKwargs(BaseKwargs):
 40    """Keyword arguments for the bar_plot function."""
 41
 42    # --- options for the entire bar plot
 43    ax: NotRequired[Axes | None]
 44    stacked: NotRequired[bool]
 45    horizontal: NotRequired[bool]
 46    max_ticks: NotRequired[int]
 47    tick_relabel: NotRequired[Callable[[str], str]]
 48    plot_from: NotRequired[int | Period]
 49    label_rotation: NotRequired[int | float]
 50    # --- options for each bar ...
 51    color: NotRequired[str | Sequence[str]]
 52    label_series: NotRequired[bool | Sequence[bool]]
 53    width: NotRequired[float | int | Sequence[float | int]]
 54    zorder: NotRequired[int | float | Sequence[int | float]]
 55    # --- options for bar annotations
 56    annotate: NotRequired[bool]
 57    fontsize: NotRequired[int | float | str]
 58    fontname: NotRequired[str]
 59    rounding: NotRequired[int]
 60    rotation: NotRequired[int | float]
 61    annotate_color: NotRequired[str]
 62    above: NotRequired[bool]
 63
 64
 65# --- functions
 66class AnnoKwargs(TypedDict, total=False):
 67    """TypedDict for the kwargs used in annotate_bars."""
 68
 69    annotate: bool
 70    fontsize: int | float | str
 71    fontname: str
 72    color: str
 73    rotation: int | float
 74    foreground: str | Sequence[str]  # used for stroke effect on text (per-bar if a sequence)
 75    above: bool
 76    rounding: bool | int  # if True, uses default rounding; if int, uses that value
 77
 78
 79def annotate_bars(
 80    series: Series,
 81    offset: float,
 82    base: np.ndarray,
 83    axes: Axes,
 84    *,
 85    horizontal: bool = False,
 86    **anno_kwargs: Unpack[AnnoKwargs],
 87) -> None:
 88    """Bar plot annotations.
 89
 90    Annotations are placed along the value axis: the y-axis for vertical
 91    bars, the x-axis when horizontal=True.
 92
 93    Note: "annotate", "fontsize", "fontname", "color", and "rotation" are expected in anno_kwargs.
 94    """
 95    # --- only annotate in limited circumstances
 96    if "annotate" not in anno_kwargs or not anno_kwargs["annotate"]:
 97        return
 98    max_annotations = MAX_ANNOTATIONS
 99    if len(series) > max_annotations:
100        return
101
102    # --- internal logic check
103    if len(base) != len(series):
104        print(f"Warning: base array length {len(base)} does not match series length {len(series)}.")
105        return
106
107    # --- assemble the annotation parameters
108    above: Final[bool | None] = anno_kwargs.get("above", False)  # None is also False-ish
109    annotate_style: dict[str, Any] = {
110        "fontsize": anno_kwargs.get("fontsize"),
111        "fontname": anno_kwargs.get("fontname"),
112        "color": anno_kwargs.get("color"),
113        "rotation": anno_kwargs.get("rotation"),
114    }
115    rounding = default_rounding(series=series, provided=anno_kwargs.get("rounding"))
116    adjustment = (series.max() - series.min()) * ADJUSTMENT_FACTOR
117    zero_correction = series.index.min()
118
119    # --- annotate each bar
120    for index, value in zip(series.index.astype(int), series, strict=True):
121        position = base[index - zero_correction] + (adjustment if value >= 0 else -adjustment)
122        if above:
123            position += value
124        if horizontal:
125            placement: dict[str, Any] = {
126                "x": position,
127                "y": index + offset,
128                "ha": "left" if value >= 0 else "right",
129                "va": "center",
130            }
131        else:
132            placement = {
133                "x": index + offset,
134                "y": position,
135                "ha": "center",
136                "va": "bottom" if value >= 0 else "top",
137            }
138        text = axes.text(
139            s=f"{value:.{rounding}f}",
140            **placement,
141            **annotate_style,
142        )
143        if not above and "foreground" in anno_kwargs:
144            # apply a stroke-effect to within bar annotations
145            # to make them more readable with very small bars.
146            foreground = anno_kwargs.get("foreground")
147            if isinstance(foreground, Sequence) and not isinstance(foreground, str):
148                foreground = foreground[index - zero_correction]  # per-bar colours
149            text.set_path_effects([pe.withStroke(linewidth=2, foreground=foreground)])
150
151
152class GroupedKwargs(TypedDict):
153    """TypedDict for the kwargs used in grouped."""
154
155    color: Sequence[str]
156    width: Sequence[float | int]
157    label_series: Sequence[bool]
158    zorder: Sequence[int | float | None]
159
160
161def grouped(
162    axes: Axes,
163    df: DataFrame,
164    anno_args: AnnoKwargs,
165    *,
166    horizontal: bool = False,
167    **kwargs: Unpack[GroupedKwargs],
168) -> None:
169    """Plot a grouped bar plot."""
170    series_count = len(df.columns)
171
172    for i, col in enumerate(df.columns):
173        series = df[col]
174        if series.isna().all():
175            continue
176        width = kwargs["width"][i]
177        if width < MIN_BAR_WIDTH or width > MAX_BAR_WIDTH:
178            width = DEFAULT_GROUPED_WIDTH
179        adjusted_width = width / series_count
180        # far-left + margin + halfway through one grouped column
181        left = -DEFAULT_BAR_OFFSET + ((1 - width) / 2.0) + (adjusted_width / 2.0)
182        offset = left + (i * adjusted_width)
183        foreground = kwargs["color"][i]
184        common: dict[str, Any] = {
185            "color": foreground,
186            "zorder": kwargs["zorder"][i],
187            "label": col if kwargs["label_series"][i] else f"_{col}_",
188        }
189        if horizontal:
190            axes.barh(y=series.index + offset, width=series, height=adjusted_width, **common)
191        else:
192            axes.bar(x=series.index + offset, height=series, width=adjusted_width, **common)
193        anno_args["foreground"] = foreground
194        annotate_bars(
195            series=series,
196            offset=offset,
197            base=np.zeros(len(series)),
198            axes=axes,
199            horizontal=horizontal,
200            **anno_args,
201        )
202
203
204class StackedKwargs(TypedDict):
205    """TypedDict for the kwargs used in stacked."""
206
207    color: Sequence[str]
208    width: Sequence[float | int]
209    label_series: Sequence[bool]
210    zorder: Sequence[int | float | None]
211
212
213def stacked(
214    axes: Axes,
215    df: DataFrame,
216    anno_args: AnnoKwargs,
217    *,
218    horizontal: bool = False,
219    **kwargs: Unpack[StackedKwargs],
220) -> None:
221    """Plot a stacked bar plot."""
222    row_count = len(df)
223    base_plus: np.ndarray = np.zeros(shape=row_count, dtype=np.float64)
224    base_minus: np.ndarray = np.zeros(shape=row_count, dtype=np.float64)
225    for i, col in enumerate(df.columns):
226        series = df[col]
227        base = np.where(series >= 0, base_plus, base_minus)
228        foreground = kwargs["color"][i]
229        common: dict[str, Any] = {
230            "color": foreground,
231            "zorder": kwargs["zorder"][i],
232            "label": col if kwargs["label_series"][i] else f"_{col}_",
233        }
234        if horizontal:
235            axes.barh(y=series.index, width=series, left=base, height=kwargs["width"][i], **common)
236        else:
237            axes.bar(x=series.index, height=series, bottom=base, width=kwargs["width"][i], **common)
238        anno_args["foreground"] = foreground
239        annotate_bars(
240            series=series,
241            offset=0,
242            base=base,
243            axes=axes,
244            horizontal=horizontal,
245            **anno_args,
246        )
247        base_plus += np.where(series >= 0, series, 0)
248        base_minus += np.where(series < 0, series, 0)
249
250
251def bar_plot(data: DataT, **kwargs: Unpack[BarKwargs]) -> Axes:
252    """Create a bar plot from the given data.
253
254    Each column in the DataFrame will be stacked on top of each other,
255    with positive values above zero and negative values below zero.
256
257    Args:
258        data: Series | DataFrame - The data to plot. Can be a DataFrame or a Series.
259        **kwargs: BarKwargs - Additional keyword arguments for customization.
260        (see BarKwargs for details)
261
262    Note: This function does not assume all data is timeseries with a PeriodIndex.
263
264    Returns:
265        axes: Axes - The axes for the plot.
266
267    """
268    # --- check the kwargs
269    report_kwargs(caller=ME, **kwargs)
270    validate_kwargs(schema=BarKwargs, caller=ME, **kwargs)
271
272    # --- get the data
273    # no call to check_clean_timeseries here, as bar plots are not
274    # necessarily timeseries data. If the data is a Series, it will be
275    # converted to a DataFrame with a single column.
276    df = DataFrame(data)  # really we are only plotting DataFrames
277    df, kwargs_d = constrain_data(df, **kwargs)
278    item_count = len(df.columns)
279
280    # --- deal with string indices
281    saved_strings = map_stringindex(df)
282    if saved_strings is not None:
283        df = saved_strings[0]
284
285    # --- deal with complete PeriodIndex indices
286    saved_pi = map_periodindex(df)
287    if saved_pi is not None:
288        df = saved_pi[0]  # extract the reindexed DataFrame from the PeriodIndex
289
290    # --- set up the default arguments
291    chart_defaults: dict[str, bool | int] = {
292        "stacked": False,
293        "horizontal": False,
294        "max_ticks": DEFAULT_MAX_TICKS,
295        "label_series": item_count > 1,
296        "label_rotation": 0,
297    }
298    chart_args = {k: kwargs_d.get(k, v) for k, v in chart_defaults.items()}
299
300    # --- horizontal bars are for categorical data, not PeriodIndex timeseries
301    horizontal = bool(chart_args["horizontal"])
302    if horizontal and saved_pi is not None:
303        print(f"Warning: horizontal=True is not supported with a PeriodIndex in {ME}(); plotting vertical.")
304        horizontal = False
305
306    # --- single series + one colour per bar => per-bar colours
307    user_color = kwargs_d.get("color")
308    if (
309        item_count == 1
310        and isinstance(user_color, Sequence)
311        and not isinstance(user_color, str)
312        and len(user_color) == len(df)
313    ):
314        kwargs_d["color"] = [list(user_color)]  # one series whose colour is a per-bar list
315
316    bar_defaults = {
317        "color": get_color_list(item_count),
318        "width": get_setting("bar_width"),
319        "label_series": item_count > 1,
320        "zorder": None,
321    }
322    above = kwargs_d.get("above", False)
323    anno_args: AnnoKwargs = {
324        "annotate": kwargs_d.get("annotate", False),
325        "fontsize": kwargs_d.get("fontsize", "small"),
326        "fontname": kwargs_d.get("fontname", "Helvetica"),
327        "rotation": kwargs_d.get("rotation", 0),
328        "rounding": kwargs_d.get("rounding", True),
329        "color": kwargs_d.get("annotate_color", "black" if above else "white"),
330        "above": above,
331    }
332    bar_args, remaining_kwargs = apply_defaults(item_count, bar_defaults, kwargs_d)
333
334    # --- plot the data
335    axes, remaining_kwargs = get_axes(**dict(remaining_kwargs))
336    if chart_args["stacked"]:
337        stacked(axes, df, anno_args, horizontal=horizontal, **bar_args)
338    else:
339        grouped(axes, df, anno_args, horizontal=horizontal, **bar_args)
340
341    # --- handle index labels and rotation
342    if saved_strings is not None:
343        if horizontal:
344            axes.set_yticks(range(len(saved_strings[1])))
345            axes.set_yticklabels(saved_strings[1], rotation=chart_args["label_rotation"])
346        else:
347            axes.set_xticks(range(len(saved_strings[1])))
348            axes.set_xticklabels(saved_strings[1], rotation=chart_args["label_rotation"])
349    elif saved_pi is not None:
350        set_labels(
351            axes,
352            saved_pi[1],
353            chart_args["max_ticks"],
354            rotation=chart_args["label_rotation"],
355            tick_relabel=kwargs_d.get("tick_relabel"),
356        )
357
358    return axes