mgplot.line_plot

  1"""Plot a series or a dataframe with lines."""
  2
  3import math
  4from collections.abc import Callable, Sequence
  5from typing import Any, Final, NotRequired, TypedDict, Unpack
  6
  7from matplotlib.axes import Axes
  8from pandas import DataFrame, Period, PeriodIndex, Series
  9from pandas.api.types import is_numeric_dtype
 10
 11from mgplot.axis_utils import map_periodindex, set_labels
 12from mgplot.keyword_checking import BaseKwargs, report_kwargs, validate_kwargs
 13from mgplot.settings import DataT, get_setting
 14from mgplot.utilities import (
 15    apply_defaults,
 16    check_clean_timeseries,
 17    constrain_data,
 18    default_rounding,
 19    get_axes,
 20    get_color_list,
 21)
 22
 23# --- constants
 24ME: Final[str] = "line_plot"
 25
 26
 27class LineKwargs(BaseKwargs):
 28    """Keyword arguments for the line_plot function."""
 29
 30    # --- options for the entire line plot
 31    ax: NotRequired[Axes | None]
 32    style: NotRequired[str | Sequence[str]]
 33    width: NotRequired[float | int | Sequence[float | int]]
 34    color: NotRequired[str | Sequence[str]]
 35    alpha: NotRequired[float | Sequence[float]]
 36    drawstyle: NotRequired[str | Sequence[str] | None]
 37    marker: NotRequired[str | Sequence[str] | None]
 38    markersize: NotRequired[float | Sequence[float] | int | None]
 39    zorder: NotRequired[int | float | Sequence[int | float]]
 40    dropna: NotRequired[bool | Sequence[bool]]
 41    annotate: NotRequired[bool | Sequence[bool]]
 42    rounding: NotRequired[Sequence[int | bool] | int | bool | None]
 43    fontsize: NotRequired[Sequence[str | int | float] | str | int | float]
 44    fontname: NotRequired[str | Sequence[str]]
 45    rotation: NotRequired[Sequence[int | float] | int | float]
 46    annotate_color: NotRequired[str | Sequence[str] | bool | Sequence[bool] | None]
 47    plot_from: NotRequired[int | Period | None]
 48    label_series: NotRequired[bool | Sequence[bool] | None]
 49    max_ticks: NotRequired[int]
 50    tick_relabel: NotRequired[Callable[[str], str]]
 51
 52
 53class AnnotateKwargs(TypedDict):
 54    """Keyword arguments for the annotate_series function."""
 55
 56    color: str
 57    rounding: int | bool
 58    fontsize: str | int | float
 59    fontname: str
 60    rotation: int | float
 61
 62
 63# --- functions
 64def annotate_series(
 65    series: Series,
 66    axes: Axes,
 67    **kwargs: Unpack[AnnotateKwargs],
 68) -> None:
 69    """Annotate the right-hand end-point of a line-plotted series."""
 70    # --- check the series has a value to annotate
 71    latest: Series = series.dropna()
 72    if latest.empty or not is_numeric_dtype(latest):
 73        return
 74    x: int | float = latest.index[-1]  # type: ignore[assignment]
 75    y: int | float = latest.iloc[-1]
 76    if y is None or math.isnan(y):
 77        return
 78
 79    # --- extract fontsize - could be None, bool, int or str.
 80    fontsize = kwargs.get("fontsize", "small")
 81    if fontsize is None or isinstance(fontsize, bool):
 82        fontsize = "small"
 83    fontname = kwargs.get("fontname", "Helvetica")
 84    rotation = kwargs.get("rotation", 0)
 85
 86    # --- add the annotation
 87    color = kwargs.get("color")
 88    if color is None:
 89        raise ValueError("color is required for annotation")
 90    rounding = default_rounding(value=y, provided=kwargs.get("rounding"))
 91    r_string = f"  {y:.{rounding}f}" if rounding > 0 else f"  {int(y)}"
 92    axes.text(
 93        x=x,
 94        y=y,
 95        s=r_string,
 96        ha="left",
 97        va="center",
 98        fontsize=fontsize,
 99        font=fontname,
100        rotation=rotation,
101        color=color,
102    )
103
104
105def get_style_width_color_etc(
106    item_count: int,
107    num_data_points: int,
108    **kwargs: Unpack[LineKwargs],
109) -> tuple[dict[str, list | tuple], dict[str, Any]]:
110    """Get the plot-line attributes arguemnts.
111
112    Args:
113        item_count: Number of data series to plot (columns in DataFrame)
114        num_data_points: Number of data points in the series (rows in DataFrame)
115        kwargs: LineKwargs - other arguments
116
117    Returns a tuple comprising:
118        - swce: dict[str, list | tuple] - style, width, color, etc. for each line
119        - kwargs_d: dict[str, Any] - the kwargs with defaults applied for the line plot
120
121    """
122    data_point_thresh = 151  # switch from wide to narrow lines
123    force_lines_styles = 4
124
125    line_defaults: dict[str, Any] = {
126        "style": ("solid" if item_count <= force_lines_styles else ["solid", "dashed", "dashdot", "dotted"]),
127        "width": (
128            get_setting("line_normal") if num_data_points > data_point_thresh else get_setting("line_wide")
129        ),
130        "color": get_color_list(item_count),
131        "alpha": 1.0,
132        "drawstyle": None,
133        "marker": None,
134        "markersize": 10,
135        "zorder": None,
136        "dropna": True,
137        "annotate": False,
138        "rounding": True,
139        "fontsize": "small",
140        "fontname": "Helvetica",
141        "rotation": 0,
142        "annotate_color": True,
143        "label_series": True,
144    }
145
146    return apply_defaults(item_count, line_defaults, dict(kwargs))
147
148
149def line_plot(data: DataT, **kwargs: Unpack[LineKwargs]) -> Axes:
150    """Build a single or multi-line plot.
151
152    Args:
153        data: DataFrame | Series - data to plot
154        kwargs: LineKwargs - keyword arguments for the line plot
155
156    Returns:
157    - axes: Axes - the axes object for the plot
158
159    """
160    # --- check the kwargs
161    report_kwargs(caller=ME, **kwargs)
162    validate_kwargs(schema=LineKwargs, caller=ME, **kwargs)
163
164    # --- check the data
165    data = check_clean_timeseries(data, ME)
166    df = DataFrame(data)  # we are only plotting DataFrames
167    df, kwargs_d = constrain_data(df, **kwargs)
168
169    # --- convert PeriodIndex to Integer Index
170    saved_pi = map_periodindex(df)
171    if saved_pi is not None:
172        df = saved_pi[0]
173
174    if isinstance(df.index, PeriodIndex):
175        print("Internal error: data is still a PeriodIndex - come back here and fix it")
176
177    # --- Let's plot
178    axes, kwargs_d = get_axes(**kwargs_d)  # get the axes to plot on
179    if df.empty or df.isna().all().all():
180        # Note: finalise plot should ignore an empty axes object
181        print(f"Warning: No data to plot in {ME}().")
182        return axes
183
184    # --- get the arguments for each line we will plot ...
185    item_count = len(df.columns)
186    num_data_points = len(df)
187    swce, kwargs_d = get_style_width_color_etc(item_count, num_data_points, **kwargs_d)
188
189    for i, column in enumerate(df.columns):
190        series = df[column]
191        series = series.dropna() if "dropna" in swce and swce["dropna"][i] else series
192        if series.empty or series.isna().all():
193            print(f"Warning: No data to plot for {column} in line_plot().")
194            continue
195
196        axes.plot(
197            # using matplotlib, as pandas can set xlabel/ylabel
198            series.index,  # x
199            series,  # y
200            ls=swce["style"][i],
201            lw=swce["width"][i],
202            color=swce["color"][i],
203            alpha=swce["alpha"][i],
204            marker=swce["marker"][i],
205            ms=swce["markersize"][i],
206            drawstyle=swce["drawstyle"][i],
207            zorder=swce["zorder"][i],
208            label=(column if "label_series" in swce and swce["label_series"][i] else f"_{column}_"),
209        )
210
211        if swce["annotate"][i] is None or not swce["annotate"][i]:
212            continue
213
214        color = swce["color"][i] if swce["annotate_color"][i] is True else swce["annotate_color"][i]
215        annotate_series(
216            series,
217            axes,
218            color=color,
219            rounding=swce["rounding"][i],
220            fontsize=swce["fontsize"][i],
221            fontname=swce["fontname"][i],
222            rotation=swce["rotation"][i],
223        )
224
225    # --- set the labels
226    if saved_pi is not None:
227        set_labels(
228            axes,
229            saved_pi[1],
230            kwargs_d.get("max_ticks", get_setting("max_ticks")),
231            tick_relabel=kwargs_d.get("tick_relabel"),
232        )
233
234    return axes