Skip to content

Scatter Plot

Scatter and bubble plots from a DataFrame.

label_col is not supported when value_col is a list (ValueError), and size_col must contain numeric, non-negative values.

plot(df, value_col, x_label=None, y_label=None, title=None, eyebrow=None, subtitle=None, x_col=None, group_col=None, size_col=None, size_scale=1.0, figsize=None, ax=None, source_text=None, legend_title=None, move_legend_outside=False, label_col=None, label_kwargs=None, **kwargs)

Plot a scatter chart for value_col over x_col (or the index), optionally grouped.

A value_col list and a group_col are mutually exclusive. When size_col is set and s is also in kwargs, a UserWarning is raised and s is ignored. Alpha defaults to 0.7.

Parameters:

Name Type Description Default
df DataFrame or Series

Frame to plot; a Series is converted to a single-column frame.

required
value_col str or list[str]

Column(s) to plot.

required
x_label str

X-axis label; omitted means no label.

None
y_label str

Y-axis label; omitted means no label.

None
title str

Plot title.

None
eyebrow str

Uppercase label rendered above the title.

None
subtitle str

Supporting copy rendered below the title.

None
x_col str

Column used as the x-axis; if None, the index is used.

None
group_col str

Column used to split the data into separate series.

None
size_col str

Column with the values used for point sizes (bubble chart); when set with no group_col, the same sizes apply to every series.

None
size_scale float

Multiplier; actual size = size_col value * size_scale.

1.0
figsize tuple[int, int]

Figure size, used only when ax is None.

None
ax Axes

Axes to plot on.

None
source_text str

Source attribution rendered at the bottom.

None
legend_title str

Legend title.

None
move_legend_outside bool

Move the legend outside the plot.

False
label_col str

Column with a text label per point; not supported when value_col is a list.

None
label_kwargs dict

Extra args passed to textalloc.allocate (draw_lines defaults to False).

None
**kwargs Any

Forwarded to matplotlib scatter.

{}

Raises:

Type Description
ValueError

If value_col is a list and group_col are both provided.

ValueError

If label_col is provided when value_col is a list.

KeyError

If label_col or size_col is not a column of df.

ValueError

If size_col is non-numeric or contains negative values, or size_scale is not positive.

Source code in openretailscience/plots/scatter.py
def plot(  # noqa: PLR0913, PLR0917
    df: pd.DataFrame | pd.Series,
    value_col: str | list[str],
    x_label: str | None = None,
    y_label: str | None = None,
    title: str | None = None,
    eyebrow: str | None = None,
    subtitle: str | None = None,
    x_col: str | None = None,
    group_col: str | None = None,
    size_col: str | None = None,
    size_scale: float = 1.0,
    figsize: tuple[int, int] | None = None,
    ax: Axes | None = None,
    source_text: str | None = None,
    legend_title: str | None = None,
    move_legend_outside: bool = False,
    label_col: str | None = None,
    label_kwargs: dict[str, Any] | None = None,
    **kwargs: Any,  # noqa: ANN401
) -> SubplotBase:
    """Plot a scatter chart for `value_col` over `x_col` (or the index), optionally grouped.

    A `value_col` list and a `group_col` are mutually exclusive. When `size_col` is set and
    `s` is also in kwargs, a UserWarning is raised and `s` is ignored. Alpha defaults to 0.7.

    Args:
        df (pd.DataFrame or pd.Series): Frame to plot; a Series is converted to a single-column frame.
        value_col (str or list[str]): Column(s) to plot.
        x_label (str, optional): X-axis label; omitted means no label.
        y_label (str, optional): Y-axis label; omitted means no label.
        title (str, optional): Plot title.
        eyebrow (str, optional): Uppercase label rendered above the title.
        subtitle (str, optional): Supporting copy rendered below the title.
        x_col (str, optional): Column used as the x-axis; if None, the index is used.
        group_col (str, optional): Column used to split the data into separate series.
        size_col (str, optional): Column with the values used for point sizes (bubble chart);
            when set with no group_col, the same sizes apply to every series.
        size_scale (float, optional): Multiplier; actual size = size_col value * size_scale.
        figsize (tuple[int, int], optional): Figure size, used only when ax is None.
        ax (Axes, optional): Axes to plot on.
        source_text (str, optional): Source attribution rendered at the bottom.
        legend_title (str, optional): Legend title.
        move_legend_outside (bool, optional): Move the legend outside the plot.
        label_col (str, optional): Column with a text label per point; not supported when
            value_col is a list.
        label_kwargs (dict, optional): Extra args passed to textalloc.allocate (draw_lines
            defaults to False).
        **kwargs: Forwarded to matplotlib scatter.

    Raises:
        ValueError: If `value_col` is a list and `group_col` are both provided.
        ValueError: If `label_col` is provided when `value_col` is a list.
        KeyError: If `label_col` or `size_col` is not a column of df.
        ValueError: If `size_col` is non-numeric or contains negative values, or size_scale is not positive.

    """
    if isinstance(df, pd.Series):
        df = df.to_frame()

    if isinstance(value_col, list) and group_col:
        raise ValueError("Cannot use both a list for `value_col` and a `group_col`. Choose one.")

    if label_col is not None:
        if isinstance(value_col, list):
            raise ValueError(
                "label_col is not supported when value_col is a list. "
                "Please use a single value_col or create separate plots.",
            )

        if label_col not in df.columns:
            msg = f"label_col '{label_col}' not found in DataFrame"
            raise KeyError(msg)

    # Validate size parameters and resolve conflicting kwargs
    kwargs = _handle_size_params(df, size_col, size_scale, kwargs)

    if group_col is None:
        pivot_df = df.set_index(x_col if x_col is not None else df.index)[
            [value_col] if isinstance(value_col, str) else value_col
        ]
    else:
        pivot_df = (
            df.pivot(columns=group_col, values=value_col)
            if x_col is None
            else df.pivot(index=x_col, columns=group_col, values=value_col)
        )

    is_multi_scatter = (group_col is not None) or (isinstance(value_col, list) and len(value_col) > 1)

    num_colors = len(pivot_df.columns) if is_multi_scatter else 1
    default_colors = get_plot_colors(num_colors)

    # Handle color parameter - can be single color or list of colors
    color = kwargs.pop("color", default_colors)
    colors = [color] * num_colors if not isinstance(color, list) else color

    # Process size data if size_col is specified
    size_data = _process_size_data(df, size_col, size_scale, x_col, group_col)

    if ax is None:
        _, ax = plt.subplots(figsize=figsize)
    alpha = kwargs.pop("alpha", 0.7)

    _create_scatter_plot(ax, pivot_df, colors, size_data, group_col, is_multi_scatter, alpha, **kwargs)

    # Add labels if requested
    if label_col is not None:
        _add_point_labels(
            ax=ax,
            df=df,
            value_col=value_col,
            label_col=label_col,
            x_col=x_col,
            label_kwargs=label_kwargs,
        )

    return standard_graph_styles(
        ax=ax,
        title=title,
        eyebrow=eyebrow,
        subtitle=subtitle,
        x_label=x_label,
        y_label=y_label,
        legend_title=legend_title,
        move_legend_outside=move_legend_outside,
        source_text=source_text,
    )