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