mgplot.line_plot

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