Cohort Plot

Cohort heatmap: wraps heatmap.plot with x tick labels on top and percentage formatting.

Input is a pre-aggregated matrix (rows = cohorts, columns = periods since start), such as CohortAnalysis.df; no aggregation is performed.

plot(df, cbar_label, x_label=None, y_label=None, title=None, eyebrow=None, subtitle=None, ax=None, source_text=None, percentage=True, figsize=None, colormap_style='discrete', y_col='min_period_shopped')

Plot a cohort heatmap, with period labels along the top.

Parameters:

Name	Type	Description	Default
df	DataFrame	Pre-aggregated cohort matrix: the y_col cohort labels, then one column per period.	required
cbar_label	str	Colorbar label.	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
ax	Axes	Axes to plot on.	None
source_text	str	Source attribution rendered at the bottom.	None
percentage	bool	If True, cells must be 0-1 fractions (as from CohortAnalysis(percentage=True)) and display as percentages; pass False for raw counts.	True
figsize	tuple[int, int]	Figure size, used only when ax is None.	None
colormap_style	(discrete, continuous)	Use "continuous" when retention differences between cohorts are small enough that the discrete bins lump them together.	'discrete'
y_col	str	Column holding the cohort labels; None uses the index. Defaults to "min_period_shopped", as in CohortAnalysis.df.	'min_period_shopped'

Returns:

Name	Type	Description
SubplotBase	SubplotBase	The matplotlib axes object.

Raises:

Type	Description
ValueError	If y_col is not a column of df, or df is empty or contains no finite values.
TypeError	If a period column is not numeric.

Examples:

>>> from openretailscience.analysis import CohortAnalysis
>>> from openretailscience.plots import cohort
>>> cohorts = CohortAnalysis(
...     df=transactions, aggregation_column="customer_id", period="quarter", percentage=True
... )
>>> ax = cohort.plot(
...     df=cohorts.df, cbar_label="Customers retained",
...     x_label="Quarters since first purchase", y_label="First quarter",
...     eyebrow="Retention", title="Quarterly customer retention", subtitle="Share of each cohort still buying",
... )

Source code in openretailscience/plots/cohort.py

def plot(
    df: pd.DataFrame,
    cbar_label: str,
    x_label: str | None = None,
    y_label: str | None = None,
    title: str | None = None,
    eyebrow: str | None = None,
    subtitle: str | None = None,
    ax: Axes | None = None,
    source_text: str | None = None,
    percentage: bool = True,
    figsize: tuple[int, int] | None = None,
    colormap_style: Literal["discrete", "continuous"] = "discrete",
    y_col: str | None = "min_period_shopped",
) -> SubplotBase:
    """Plot a cohort heatmap, with period labels along the top.

    Args:
        df (pd.DataFrame): Pre-aggregated cohort matrix: the ``y_col`` cohort labels, then one
            column per period.
        cbar_label (str): Colorbar label.
        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.
        ax (Axes, optional): Axes to plot on.
        source_text (str, optional): Source attribution rendered at the bottom.
        percentage (bool, optional): If True, cells must be 0-1 fractions (as from
            ``CohortAnalysis(percentage=True)``) and display as percentages; pass False for raw counts.
        figsize (tuple[int, int], optional): Figure size, used only when ax is None.
        colormap_style ("discrete", "continuous", optional): Use "continuous" when retention
            differences between cohorts are small enough that the discrete bins lump them
            together.
        y_col (str, optional): Column holding the cohort labels; None uses the index. Defaults to
            "min_period_shopped", as in ``CohortAnalysis.df``.

    Returns:
        SubplotBase: The matplotlib axes object.

    Raises:
        ValueError: If ``y_col`` is not a column of df, or df is empty or contains no finite values.
        TypeError: If a period column is not numeric.

    Examples:
        >>> from openretailscience.analysis import CohortAnalysis
        >>> from openretailscience.plots import cohort
        >>> cohorts = CohortAnalysis(
        ...     df=transactions, aggregation_column="customer_id", period="quarter", percentage=True
        ... )
        >>> ax = cohort.plot(
        ...     df=cohorts.df, cbar_label="Customers retained",
        ...     x_label="Quarters since first purchase", y_label="First quarter",
        ...     eyebrow="Retention", title="Quarterly customer retention", subtitle="Share of each cohort still buying",
        ... )
    """
    if y_col is not None and y_col not in df.columns:
        msg = f"{y_col!r} is not a column of df; pass y_col=None to label the rows by the index."
        raise ValueError(msg)
    return heatmap.plot(
        df=df,
        cbar_label=cbar_label,
        x_label=x_label,
        y_label=y_label,
        title=title,
        eyebrow=eyebrow,
        subtitle=subtitle,
        ax=ax,
        source_text=source_text,
        figsize=figsize,
        x_labels_position="top",
        cbar_format="{x:.0%}" if percentage else "{x:g}",
        colormap_style=colormap_style,
        y_col=y_col,
    )