mgplot.scatter_plot

  1"""Plot one numeric column against another as a scatter plot.
  2
  3Unlike the other mgplot plot functions, the index is not plotted. It is only
  4used to find the latest observation (for highlight_latest) and, when it is a
  5PeriodIndex or an integer index, to apply plot_from.
  6"""
  7
  8from typing import Any, Final, NotRequired, Unpack
  9
 10import numpy as np
 11from matplotlib.axes import Axes
 12from pandas import DataFrame, Period, PeriodIndex, Series
 13from pandas.api.types import is_integer_dtype, is_numeric_dtype
 14
 15from mgplot.keyword_checking import BaseKwargs, report_kwargs, validate_kwargs
 16from mgplot.settings import get_setting
 17from mgplot.utilities import constrain_data, get_axes, get_color_list
 18
 19# --- constants
 20ME: Final[str] = "scatter_plot"
 21REQUIRED_COLUMNS: Final[int] = 2  # x first, y second
 22DEFAULT_SIZE: Final[float] = 12  # marker area (matplotlib s=)
 23DEFAULT_ALPHA: Final[float] = 0.6
 24DEFAULT_MARKER: Final[str] = "o"
 25DEFAULT_POINTS_ZORDER: Final[float] = 1  # matplotlib's default for collections
 26LATEST_SIZE_MULTIPLE: Final[float] = 16  # highlighted point area relative to the others
 27DIAGONAL_LABEL: Final[str] = "45° (equal)"
 28DIAGONAL_COLOR: Final[str] = "darkred"
 29CORR_DECIMALS: Final[int] = 2
 30
 31
 32class ScatterKwargs(BaseKwargs):
 33    """Keyword arguments for the scatter_plot function."""
 34
 35    ax: NotRequired[Axes | None]
 36    color: NotRequired[str]
 37    size: NotRequired[float | int]
 38    alpha: NotRequired[float]
 39    marker: NotRequired[str]
 40    label: NotRequired[str | None]
 41    zorder: NotRequired[int | float]
 42    dropna: NotRequired[bool]
 43    plot_from: NotRequired[int | Period | None]
 44    diagonal: NotRequired[bool | dict[str, Any]]
 45    fit: NotRequired[bool | dict[str, Any]]
 46    highlight_latest: NotRequired[bool | dict[str, Any]]
 47    report_corr: NotRequired[bool]
 48
 49
 50# --- private functions
 51def _check_data(data: object) -> DataFrame:
 52    """Confirm data is a DataFrame with exactly two numeric columns (x, y)."""
 53    if not isinstance(data, DataFrame):
 54        raise TypeError(
 55            f"{ME}() expects a DataFrame with two numeric columns (x, y), not a {type(data).__name__}."
 56        )
 57    if len(data.columns) != REQUIRED_COLUMNS:
 58        raise ValueError(
 59            f"{ME}() expects exactly two columns (x first, y second), but got {len(data.columns)}."
 60        )
 61    non_numeric = [str(c) for c in data.columns if not is_numeric_dtype(data[c])]
 62    if non_numeric:
 63        raise ValueError(f"{ME}() expects numeric columns, but these are not: {', '.join(non_numeric)}.")
 64    return data
 65
 66
 67def _constrain(df: DataFrame, kwargs_d: dict[str, Any]) -> tuple[DataFrame, dict[str, Any]]:
 68    """Apply plot_from, when provided, to the data."""
 69    if kwargs_d.get("plot_from") is None:
 70        kwargs_d.pop("plot_from", None)
 71        return df, kwargs_d
 72    if not isinstance(df.index, PeriodIndex) and not is_integer_dtype(df.index):
 73        print(f"Warning: plot_from ignored in {ME}(); the index is neither a PeriodIndex nor integer.")
 74        kwargs_d.pop("plot_from", None)
 75        return df, kwargs_d
 76    return constrain_data(df, **kwargs_d)
 77
 78
 79def _style(spec: object, house: dict[str, Any]) -> dict[str, Any] | None:
 80    """Resolve an overlay spec: True -> house style, dict -> merged over house style, else None."""
 81    if spec is True:
 82        return dict(house)
 83    if isinstance(spec, dict):
 84        return house | spec
 85    return None
 86
 87
 88def _points_label(label: object, x: Series, y: Series, *, report_corr: bool) -> str | None:
 89    """Return the legend label for the points, with the correlation appended if requested."""
 90    text = label if isinstance(label, str) else None
 91    if not report_corr:
 92        return text
 93    r = x.corr(y)
 94    if not isinstance(r, float) or np.isnan(r):
 95        print(f"Warning: correlation is undefined in {ME}(), so it has not been reported.")
 96        return text
 97    corr = f"r = {r:.{CORR_DECIMALS}f}"
 98    return corr if text is None else f"{text} ({corr})"
 99
100
101def _draw_diagonal(axes: Axes, spec: object) -> None:
102    """Draw the y = x line across the range covering the data on both axes."""
103    style = _style(
104        spec,
105        {"color": DIAGONAL_COLOR, "ls": "--", "lw": get_setting("line_narrow"), "label": DIAGONAL_LABEL},
106    )
107    if style is None:
108        return
109    limits = axes.dataLim  # covers everything already plotted on these axes
110    low = min(limits.x0, limits.y0)
111    high = max(limits.x1, limits.y1)
112    axes.plot([low, high], [low, high], **style)
113
114
115def _draw_fit(axes: Axes, x: Series, y: Series, spec: object, color: str) -> None:
116    """Draw the OLS line of best fit across the x-range of the points."""
117    style = _style(spec, {"color": color, "lw": get_setting("line_normal")})
118    if style is None:
119        return
120    if x.nunique() < REQUIRED_COLUMNS:
121        print(f"Warning: a line of best fit needs at least two distinct x values in {ME}().")
122        return
123    slope, intercept = np.polyfit(x.to_numpy(dtype=float), y.to_numpy(dtype=float), 1)
124    ends = np.array([x.min(), x.max()], dtype=float)
125    axes.plot(ends, slope * ends + intercept, **style)
126
127
128def _draw_latest(axes: Axes, df: DataFrame, spec: object, points: dict[str, Any]) -> None:
129    """Redraw the point with the latest index value, larger and as a star."""
130    position = int(df.index.argmax())
131    latest = df.index[position]
132    style = _style(
133        spec,
134        {
135            "color": points["color"],
136            "s": points["s"] * LATEST_SIZE_MULTIPLE,
137            "marker": "*",
138            "edgecolors": "black",
139            "linewidths": get_setting("line_narrow"),
140            "zorder": points["zorder"] + 1,
141            "label": f"Latest ({latest})",
142        },
143    )
144    if style is None:
145        return
146    row = df.iloc[position]
147    axes.scatter([row.iloc[0]], [row.iloc[1]], **style)
148
149
150# --- public functions
151def scatter_plot(data: DataFrame, **kwargs: Unpack[ScatterKwargs]) -> Axes:
152    """Plot the second column of a DataFrame (y) against the first (x).
153
154    Args:
155        data: DataFrame - exactly two numeric columns, x first and y second.
156            The index is not plotted; it identifies the latest point.
157        kwargs: ScatterKwargs - keyword arguments for the scatter plot
158
159    The overlays - diagonal, fit and highlight_latest - each take True for the
160    house style, or a dict of matplotlib arguments merged over the house style.
161        diagonal         - the y = x line (dashed, labelled "45° (equal)")
162        fit              - an OLS line of best fit, in the colour of the points
163        highlight_latest - the point with the latest index value, as a star
164    report_corr=True appends the correlation to the label, e.g. "GDP (r = 0.83)".
165
166    To layer several groups (each with its own colour and fit line), make
167    repeated calls with the same ax=, then call finalise_plot().
168
169    Returns:
170    - axes: Axes - the axes object for the plot
171
172    """
173    # --- check the kwargs
174    report_kwargs(caller=ME, **kwargs)
175    validate_kwargs(schema=ScatterKwargs, caller=ME, **kwargs)
176
177    # --- check and prepare the data
178    df, kwargs_d = _constrain(_check_data(data), dict(kwargs))
179    if kwargs_d.pop("dropna", True):
180        df = df.dropna()
181
182    # --- Let's plot
183    axes, kwargs_d = get_axes(**kwargs_d)
184    if df.empty or df.isna().all().all():
185        # Note: finalise plot should ignore an empty axes object
186        print(f"Warning: No data to plot in {ME}().")
187        return axes
188
189    x, y = df.iloc[:, 0], df.iloc[:, 1]
190    points: dict[str, Any] = {
191        "color": kwargs_d.get("color", get_color_list(1)[0]),
192        "s": kwargs_d.get("size", DEFAULT_SIZE),
193        "alpha": kwargs_d.get("alpha", DEFAULT_ALPHA),
194        "marker": kwargs_d.get("marker", DEFAULT_MARKER),
195        "zorder": kwargs_d.get("zorder", DEFAULT_POINTS_ZORDER),
196    }
197    label = _points_label(kwargs_d.get("label"), x, y, report_corr=bool(kwargs_d.get("report_corr")))
198    axes.scatter(x, y, label=label, **points)
199
200    # --- overlays
201    _draw_fit(axes, x, y, kwargs_d.get("fit"), points["color"])
202    _draw_latest(axes, df, kwargs_d.get("highlight_latest"), points)
203    _draw_diagonal(axes, kwargs_d.get("diagonal"))  # last, so its range covers the points
204
205    return axes