mgplot.summary_plot

  1"""Produce a summary plot for the data in a given DataFrame."""
  2
  3# system imports
  4from typing import Any, NotRequired, SupportsFloat, Unpack
  5
  6from matplotlib.axes import Axes
  7
  8# analytic third-party imports
  9from numpy import array, ndarray
 10from pandas import DataFrame, Period
 11
 12from mgplot.finalise_plot import make_legend
 13from mgplot.keyword_checking import (
 14    BaseKwargs,
 15    report_kwargs,
 16    validate_kwargs,
 17)
 18
 19# local imports
 20from mgplot.settings import DataT
 21from mgplot.utilities import check_clean_timeseries, constrain_data, get_axes, label_period
 22
 23# --- constants
 24ME = "summary_plot"
 25ZSCORES = "zscores"
 26ZSCALED = "zscaled"
 27
 28# Plot layout constants
 29SPAN_LIMIT = 1.15
 30SPACE_MARGIN = 0.2
 31DEFAULT_FONT_SIZE = 10
 32SMALL_FONT_SIZE = "x-small"
 33SMALL_MARKER_SIZE = 5
 34REFERENCE_LINE_WIDTH = 0.5
 35DEFAULT_MIDDLE = 0.8
 36DEFAULT_PLOT_FROM = 0
 37HIGH_PRECISION_THRESHOLD = 1
 38
 39
 40class SummaryKwargs(BaseKwargs):
 41    """Keyword arguments for the summary_plot function."""
 42
 43    ax: NotRequired[Axes | None]
 44    verbose: NotRequired[bool]
 45    middle: NotRequired[float]
 46    plot_type: NotRequired[str]
 47    plot_from: NotRequired[int | Period]
 48    legend: NotRequired[bool | dict[str, Any] | None]
 49    xlabel: NotRequired[str | None]
 50
 51
 52# --- functions
 53def calc_quantiles(middle: float) -> ndarray:
 54    """Calculate the quantiles for the middle of the data."""
 55    return array([(1 - middle) / 2.0, 1 - (1 - middle) / 2.0])
 56
 57
 58def calculate_z(
 59    original: DataFrame,
 60    middle: float,
 61    *,
 62    verbose: bool = False,
 63) -> tuple[DataFrame, DataFrame]:
 64    """Calculate z-scores, scaled z-scores and middle quantiles.
 65
 66    Args:
 67        original: DataFrame containing the original data.
 68        middle: float, the proportion of data to highlight in the middle (eg. 0.8 for 80%).
 69        verbose: bool, whether to print the summary data.
 70
 71    Returns:
 72        tuple[DataFrame, DataFrame]: z_scores and z_scaled DataFrames.
 73
 74    Raises:
 75        ValueError: If original DataFrame is empty or has zero variance.
 76
 77    """
 78    if original.empty:
 79        raise ValueError("Cannot calculate z-scores for empty DataFrame")
 80
 81    # Check for zero variance
 82    std_dev = original.std()
 83    if (std_dev == 0).any():
 84        raise ValueError("Cannot calculate z-scores when standard deviation is zero")
 85
 86    # Calculate z-scores
 87    z_scores: DataFrame = (original - original.mean()) / std_dev
 88
 89    # Scale z-scores between -1 and +1
 90    z_min = z_scores.min()
 91    z_max = z_scores.max()
 92    z_range = z_max - z_min
 93
 94    # Avoid division by zero in scaling
 95    if (z_range == 0).any():
 96        z_scaled: DataFrame = z_scores.copy() * 0  # All zeros if no variance
 97    else:
 98        z_scaled = (((z_scores - z_min) / z_range) - 0.5) * 2
 99
100    if verbose:
101        if original.index.empty:
102            raise ValueError("Cannot display statistics for empty DataFrame")
103
104        q_middle = calc_quantiles(middle)
105        frame = DataFrame(
106            {
107                "count": original.count(),
108                "mean": original.mean(),
109                "median": original.median(),
110                "min shaded": original.quantile(q=q_middle[0]),
111                "max shaded": original.quantile(q=q_middle[1]),
112                "z-scores": z_scores.iloc[-1],
113                "scaled": z_scaled.iloc[-1],
114            },
115        )
116        print(frame)
117
118    return z_scores, z_scaled
119
120
121def plot_middle_bars(
122    adjusted: DataFrame,
123    middle: float,
124    kwargs: dict[str, Any],
125) -> Axes:
126    """Plot the middle (typically 80%) of the data as a bar."""
127    if adjusted.empty:
128        raise ValueError("Cannot plot bars for empty DataFrame")
129
130    q = calc_quantiles(middle)
131    lo_hi: DataFrame = adjusted.quantile(q=q).T  # get the middle section of data
132
133    low = min(adjusted.iloc[-1].min(), lo_hi.min().min(), -SPAN_LIMIT) - SPACE_MARGIN
134    high = max(adjusted.iloc[-1].max(), lo_hi.max().max(), SPAN_LIMIT) + SPACE_MARGIN
135    kwargs["xlim"] = (low, high)  # update the kwargs with the xlim
136    ax, _ = get_axes(**kwargs)
137    ax.barh(
138        y=lo_hi.index,
139        width=lo_hi[q[1]] - lo_hi[q[0]],
140        left=lo_hi[q[0]],
141        color="#bbbbbb",
142        label=f"Middle {middle * 100:0.0f}% of prints",
143    )
144    return ax
145
146
147def plot_latest_datapoint(
148    ax: Axes,
149    original: DataFrame,
150    adjusted: DataFrame,
151    font_size: int | str,
152) -> None:
153    """Add the latest datapoints to the summary plot."""
154    if adjusted.empty or original.empty:
155        raise ValueError("Cannot plot datapoints for empty DataFrame")
156
157    ax.scatter(adjusted.iloc[-1], adjusted.columns, color="darkorange", label="Latest")
158    row = adjusted.index[-1]
159    for col_num, col_name in enumerate(original.columns):
160        raw_adj = adjusted.at[row, col_name]
161        raw_orig = original.at[row, col_name]
162        if not isinstance(raw_adj, SupportsFloat) or not isinstance(raw_orig, SupportsFloat):
163            raise TypeError(f"Expected numeric data for {col_name}, got {type(raw_orig).__name__}")
164        x_adj = float(raw_adj)
165        x_orig = float(raw_orig)
166        precision = 2 if abs(x_orig) < HIGH_PRECISION_THRESHOLD else 1
167        ax.text(
168            x=x_adj,
169            y=col_num,
170            s=f"{x_orig:.{precision}f}",
171            ha="center",
172            va="center",
173            size=font_size,
174        )
175
176
177def label_extremes(
178    ax: Axes,
179    data: tuple[DataFrame, DataFrame],
180    plot_type: str,
181    font_size: int | str,
182    kwargs: dict[str, Any],  # must be a dictionary, not a splat
183) -> None:
184    """Label the extremes in the scaled plots."""
185    original, adjusted = data
186    low, high = kwargs["xlim"]
187    ax.set_xlim(low, high)  # set the x-axis limits
188    if plot_type == ZSCALED:
189        ax.scatter(
190            adjusted.median(),
191            adjusted.columns,
192            color="darkorchid",
193            marker="x",
194            s=SMALL_MARKER_SIZE,
195            label="Median",
196        )
197        for col_num, col_name in enumerate(original.columns):
198            minima, maxima = original[col_name].min(), original[col_name].max()
199            min_precision = 2 if abs(minima) < HIGH_PRECISION_THRESHOLD else 1
200            max_precision = 2 if abs(maxima) < HIGH_PRECISION_THRESHOLD else 1
201            ax.text(
202                low,
203                col_num,
204                f" {minima:.{min_precision}f}",
205                ha="left",
206                va="center",
207                size=font_size,
208            )
209            ax.text(
210                high,
211                col_num,
212                f"{maxima:.{max_precision}f} ",
213                ha="right",
214                va="center",
215                size=font_size,
216            )
217
218
219def horizontal_bar_plot(
220    original: DataFrame,
221    adjusted: DataFrame,
222    middle: float,
223    plot_type: str,
224    kwargs: dict[str, Any],  # must be a dictionary, not a splat
225) -> Axes:
226    """Plot horizontal bars for the middle of the data."""
227    ax = plot_middle_bars(adjusted, middle, kwargs)
228    font_size = SMALL_FONT_SIZE
229    plot_latest_datapoint(ax, original, adjusted, font_size)
230    label_extremes(ax, data=(original, adjusted), plot_type=plot_type, font_size=font_size, kwargs=kwargs)
231
232    return ax
233
234
235def label_x_axis(plot_from: int | Period, label: str | None, plot_type: str, ax: Axes, df: DataFrame) -> None:
236    """Label the x-axis for the plot."""
237    start: Period = plot_from if isinstance(plot_from, Period) else df.index[plot_from]
238    if label is not None:
239        if not label:
240            if plot_type == ZSCORES:
241                label = f"Z-scores for prints since {label_period(start)}"
242            else:
243                label = f"-1 to 1 scaled z-scores since {label_period(start)}"
244        ax.set_xlabel(label)
245
246
247def mark_reference_lines(plot_type: str, ax: Axes) -> None:
248    """Mark the reference lines for the plot."""
249    line_color = "#555555"
250    line_style = "--"
251
252    if plot_type == ZSCALED:
253        ax.axvline(-1, color=line_color, linewidth=REFERENCE_LINE_WIDTH, linestyle=line_style, label="-1")
254        ax.axvline(1, color=line_color, linewidth=REFERENCE_LINE_WIDTH, linestyle=line_style, label="+1")
255    elif plot_type == ZSCORES:
256        ax.axvline(0, color=line_color, linewidth=REFERENCE_LINE_WIDTH, linestyle=line_style, label="0")
257
258
259def plot_the_data(df: DataFrame, **kwargs: Unpack[SummaryKwargs]) -> tuple[Axes, str]:
260    """Plot the data as a summary plot.
261
262    Args:
263        df: DataFrame - the data to plot.
264        kwargs: SummaryKwargs, additional keyword arguments for the plot.
265
266    Returns:
267        tuple[Axes, str]: A tuple comprising the Axes object and plot type ('zscores' or 'zscaled').
268
269    Raises:
270        ValueError: If middle value is not between 0 and 1, or if plot_type is invalid.
271
272    """
273    verbose = kwargs.pop("verbose", False)
274    middle = float(kwargs.pop("middle", DEFAULT_MIDDLE))
275    plot_type = kwargs.pop("plot_type", ZSCORES)
276
277    # Validate inputs
278    if not 0 < middle < 1:
279        raise ValueError(f"Middle value must be between 0 and 1, got {middle}")
280    if plot_type not in (ZSCORES, ZSCALED):
281        raise ValueError(f"plot_type must be '{ZSCORES}' or '{ZSCALED}', got '{plot_type}'")
282
283    subset, kwargsd = constrain_data(df, **kwargs)
284    z_scores, z_scaled = calculate_z(subset, middle, verbose=verbose)
285
286    # plot as required by the plot_types argument
287    adjusted = z_scores if plot_type == ZSCORES else z_scaled
288    ax = horizontal_bar_plot(subset, adjusted, middle, plot_type, kwargsd)
289    ax.tick_params(axis="y", labelsize="small")
290    make_legend(ax, legend=kwargsd["legend"])
291    ax.set_xlim(kwargsd.get("xlim"))  # provide space for the labels
292
293    return ax, plot_type
294
295
296# --- public
297def summary_plot(data: DataT, **kwargs: Unpack[SummaryKwargs]) -> Axes:
298    """Plot a summary of historical data for a given DataFrame.
299
300    Args:
301        data: DataFrame containing the summary data. The column names are
302              used as labels for the plot.
303        kwargs: Additional arguments for the plot, including middle (float),
304               plot_type (str), verbose (bool), and standard plotting options.
305
306    Returns:
307        Axes: A matplotlib Axes object containing the summary plot.
308
309    Raises:
310        TypeError: If data is not a DataFrame.
311
312    """
313    # --- check the kwargs
314    report_kwargs(caller=ME, **kwargs)
315    validate_kwargs(schema=SummaryKwargs, caller=ME, **kwargs)
316
317    # --- check the data
318    data = check_clean_timeseries(data, ME)
319    if not isinstance(data, DataFrame):
320        raise TypeError("data must be a pandas DataFrame for summary_plot()")
321
322    # --- legend
323    kwargs["legend"] = kwargs.get(
324        "legend",
325        {
326            # put the legend below the x-axis label
327            "loc": "upper center",
328            "fontsize": "xx-small",
329            "bbox_to_anchor": (0.5, -0.125),
330            "ncol": 4,
331        },
332    )
333
334    # --- and plot it ...
335    ax, plot_type = plot_the_data(data, **kwargs)
336    label_x_axis(
337        kwargs.get("plot_from", DEFAULT_PLOT_FROM),
338        label=kwargs.get("xlabel", ""),
339        plot_type=plot_type,
340        ax=ax,
341        df=data,
342    )
343    mark_reference_lines(plot_type, ax)
344
345    return ax