mgplot.line_plot

  1"""Plot a series or a dataframe with lines."""
  2
  3import math
  4from collections.abc import Callable, Sequence
  5from typing import TYPE_CHECKING, Any, Final, NotRequired, TypedDict, Unpack
  6
  7from matplotlib.axes import Axes
  8from matplotlib.text import Text
  9from pandas import DataFrame, Period, PeriodIndex, Series
 10from pandas.api.types import is_numeric_dtype
 11
 12from mgplot.annotation_utils import AnnotationOptions, register_annotations
 13from mgplot.axis_utils import map_periodindex, set_labels
 14from mgplot.keyword_checking import BaseKwargs, report_kwargs, validate_kwargs
 15from mgplot.settings import DataT, get_setting
 16from mgplot.utilities import (
 17    apply_defaults,
 18    check_clean_timeseries,
 19    constrain_data,
 20    default_rounding,
 21    get_axes,
 22    get_color_list,
 23)
 24
 25if TYPE_CHECKING:
 26    from matplotlib.lines import Line2D
 27
 28# --- constants
 29ME: Final[str] = "line_plot"
 30DEFAULT_NEAR_END: Final[float] = 0.1  # fraction-of-width threshold for snapping labels to the edge
 31
 32
 33class LineKwargs(BaseKwargs):
 34    """Keyword arguments for the line_plot function."""
 35
 36    # --- options for the entire line plot
 37    ax: NotRequired[Axes | None]
 38    style: NotRequired[str | Sequence[str]]
 39    width: NotRequired[float | int | Sequence[float | int]]
 40    color: NotRequired[str | Sequence[str]]
 41    alpha: NotRequired[float | Sequence[float]]
 42    drawstyle: NotRequired[str | Sequence[str] | None]
 43    marker: NotRequired[str | Sequence[str] | None]
 44    markersize: NotRequired[float | Sequence[float] | int | None]
 45    zorder: NotRequired[int | float | Sequence[int | float]]
 46    dropna: NotRequired[bool | Sequence[bool]]
 47    annotate: NotRequired[bool | str | Sequence[bool | str]]
 48    rounding: NotRequired[Sequence[int | bool] | int | bool | None]
 49    fontsize: NotRequired[Sequence[str | int | float] | str | int | float]
 50    fontname: NotRequired[str | Sequence[str]]
 51    rotation: NotRequired[Sequence[int | float] | int | float]
 52    annotate_color: NotRequired[str | Sequence[str] | bool | Sequence[bool] | None]
 53    near_end: NotRequired[float]
 54    force_right: NotRequired[bool]
 55    leader_lines: NotRequired[bool]
 56    plot_from: NotRequired[int | Period | None]
 57    label_series: NotRequired[bool | Sequence[bool] | None]
 58    max_ticks: NotRequired[int]
 59    tick_relabel: NotRequired[Callable[[str], str]]
 60
 61
 62class AnnotateKwargs(TypedDict):
 63    """Keyword arguments for the annotate_series function."""
 64
 65    color: str
 66    rounding: int | bool
 67    fontsize: str | int | float
 68    fontname: str
 69    rotation: int | float
 70    text: str | None
 71
 72
 73# --- functions
 74def annotate_series(
 75    series: Series,
 76    axes: Axes,
 77    **kwargs: Unpack[AnnotateKwargs],
 78) -> Text | None:
 79    """Annotate the right-hand end-point of a line-plotted series.
 80
 81    The label text is the `text` keyword argument when that is a string,
 82    otherwise the end-point value rounded to `rounding` places. Either way the
 83    series must have a numeric end-point, which is what the label is anchored
 84    to.
 85
 86    Returns the created Text artist (or None if there was nothing to annotate)
 87    so the caller can register it for end-of-line collision resolution.
 88    """
 89    # --- check the series has a value to annotate
 90    latest: Series = series.dropna()
 91    if latest.empty or not is_numeric_dtype(latest):
 92        return None
 93    x: int | float = latest.index[-1]  # type: ignore[assignment]
 94    y: int | float = latest.iloc[-1]
 95    if y is None or math.isnan(y):
 96        return None
 97
 98    # --- extract fontsize - could be None, bool, int or str.
 99    fontsize = kwargs.get("fontsize", "small")
100    if fontsize is None or isinstance(fontsize, bool):
101        fontsize = "small"
102    fontname = kwargs.get("fontname", "Helvetica")
103    rotation = kwargs.get("rotation", 0)
104
105    # --- add the annotation
106    color = kwargs.get("color")
107    if color is None:
108        raise ValueError("color is required for annotation")
109    text = kwargs.get("text")
110    if isinstance(text, str):
111        r_string = f"  {text}"
112    else:
113        rounding = default_rounding(value=y, provided=kwargs.get("rounding"))
114        r_string = f"  {y:.{rounding}f}" if rounding > 0 else f"  {int(y)}"
115    return axes.text(
116        x=x,
117        y=y,
118        s=r_string,
119        ha="left",
120        va="center",
121        fontsize=fontsize,
122        font=fontname,
123        rotation=rotation,
124        color=color,
125    )
126
127
128def get_style_width_color_etc(
129    item_count: int,
130    num_data_points: int,
131    **kwargs: Unpack[LineKwargs],
132) -> tuple[dict[str, list | tuple], dict[str, Any]]:
133    """Get the plot-line attributes arguemnts.
134
135    Args:
136        item_count: Number of data series to plot (columns in DataFrame)
137        num_data_points: Number of data points in the series (rows in DataFrame)
138        kwargs: LineKwargs - other arguments
139
140    Returns a tuple comprising:
141        - swce: dict[str, list | tuple] - style, width, color, etc. for each line
142        - kwargs_d: dict[str, Any] - the kwargs with defaults applied for the line plot
143
144    """
145    data_point_thresh = 151  # switch from wide to narrow lines
146    force_lines_styles = 4
147
148    line_defaults: dict[str, Any] = {
149        "style": ("solid" if item_count <= force_lines_styles else ["solid", "dashed", "dashdot", "dotted"]),
150        "width": (
151            get_setting("line_normal") if num_data_points > data_point_thresh else get_setting("line_wide")
152        ),
153        "color": get_color_list(item_count),
154        "alpha": 1.0,
155        "drawstyle": None,
156        "marker": None,
157        "markersize": 10,
158        "zorder": None,
159        "dropna": True,
160        "annotate": False,
161        "rounding": True,
162        "fontsize": "small",
163        "fontname": "Helvetica",
164        "rotation": 0,
165        "annotate_color": True,
166        "label_series": True,
167    }
168
169    return apply_defaults(item_count, line_defaults, dict(kwargs))
170
171
172def line_plot(data: DataT, **kwargs: Unpack[LineKwargs]) -> Axes:
173    """Build a single or multi-line plot.
174
175    Args:
176        data: DataFrame | Series - data to plot
177        kwargs: LineKwargs - keyword arguments for the line plot
178
179    The annotate key labels the right-hand end of each line. It takes either a
180    flag or a string, as a scalar broadcast to every series or as a per-series
181    sequence:
182        annotate=True             - label with the end-point value (default
183                                    behaviour)
184        annotate="26Q2"           - label every series with that text
185        annotate=["26Q2", "26Q3"] - label each series with its own text
186        annotate=False (or "")    - no label
187    A string label is printed as given, so rounding has no effect on it;
188    annotate_color still applies and remains the way to colour the label.
189
190    Returns:
191    - axes: Axes - the axes object for the plot
192
193    """
194    # --- check the kwargs
195    report_kwargs(caller=ME, **kwargs)
196    validate_kwargs(schema=LineKwargs, caller=ME, **kwargs)
197
198    # --- check the data
199    data = check_clean_timeseries(data, ME)
200    df = DataFrame(data)  # we are only plotting DataFrames
201    df, kwargs_d = constrain_data(df, **kwargs)
202
203    # --- convert PeriodIndex to Integer Index
204    saved_pi = map_periodindex(df)
205    if saved_pi is not None:
206        df = saved_pi[0]
207
208    if isinstance(df.index, PeriodIndex):
209        print("Internal error: data is still a PeriodIndex - come back here and fix it")
210
211    # --- Let's plot
212    axes, kwargs_d = get_axes(**kwargs_d)  # get the axes to plot on
213    if df.empty or df.isna().all().all():
214        # Note: finalise plot should ignore an empty axes object
215        print(f"Warning: No data to plot in {ME}().")
216        return axes
217
218    # --- get the arguments for each line we will plot ...
219    item_count = len(df.columns)
220    num_data_points = len(df)
221    swce, kwargs_d = get_style_width_color_etc(item_count, num_data_points, **kwargs_d)
222
223    drawn_lines: list[Line2D] = []  # every data line - obstacles for label de-collision
224    annotations: list[tuple[Text, Line2D | None]] = []  # (label, the line it annotates)
225    for i, column in enumerate(df.columns):
226        series = df[column]
227        series = series.dropna() if "dropna" in swce and swce["dropna"][i] else series
228        if series.empty or series.isna().all():
229            print(f"Warning: No data to plot for {column} in line_plot().")
230            continue
231
232        lines = axes.plot(
233            # using matplotlib, as pandas can set xlabel/ylabel
234            series.index,  # x
235            series,  # y
236            ls=swce["style"][i],
237            lw=swce["width"][i],
238            color=swce["color"][i],
239            alpha=swce["alpha"][i],
240            marker=swce["marker"][i],
241            ms=swce["markersize"][i],
242            drawstyle=swce["drawstyle"][i],
243            zorder=swce["zorder"][i],
244            label=(column if "label_series" in swce and swce["label_series"][i] else f"_{column}_"),
245        )
246        drawn_lines.extend(lines)
247
248        if swce["annotate"][i] is None or not swce["annotate"][i]:
249            continue
250
251        color = swce["color"][i] if swce["annotate_color"][i] is True else swce["annotate_color"][i]
252        text = annotate_series(
253            series,
254            axes,
255            color=color,
256            rounding=swce["rounding"][i],
257            fontsize=swce["fontsize"][i],
258            fontname=swce["fontname"][i],
259            rotation=swce["rotation"][i],
260            text=swce["annotate"][i],
261        )
262        if text is not None:
263            annotations.append((text, lines[0] if lines else None))
264
265    # --- register annotations so finalise_plot() can de-collide them
266    if annotations:
267        register_annotations(
268            axes,
269            annotations,
270            drawn_lines,
271            AnnotationOptions(
272                near_end=kwargs_d.get("near_end", DEFAULT_NEAR_END),
273                force_right=bool(kwargs_d.get("force_right", False)),
274                leader_lines=bool(kwargs_d.get("leader_lines", False)),
275            ),
276        )
277
278    # --- set the labels
279    if saved_pi is not None:
280        set_labels(
281            axes,
282            saved_pi[1],
283            kwargs_d.get("max_ticks", get_setting("max_ticks")),
284            tick_relabel=kwargs_d.get("tick_relabel"),
285        )
286
287    return axes