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