plot_lag_scatter

yohou.plotting.diagnostics.plot_lag_scatter(df, *, columns=None, lags=None, seasonality=None, show_diagonal=True, show_regression=False, groups=None, facet_n_cols=3, color_palette=None, show_legend=True, title=None, x_label=None, y_label=None, width=None, height=None, marker_size=4.0, marker_opacity=0.6)

Plot scatter plots of y(t) vs y(t-lag) for analysing temporal dependencies.

Creates scatter plots showing the relationship between current values and lagged values, useful for identifying AR patterns and lag-specific correlations. When lags has more than one entry, each lag gets its own subplot arranged in a grid. Optionally colour points by season using the seasonality parameter.

For panel data, each member gets its own figure with one subplot per lag and all groups overlaid by colour. When there are multiple members, a dict[str, go.Figure] keyed by member name is returned.

For non-panel multivariate data, each column gets its own figure with lag subplots, and a dict[str, go.Figure] keyed by column name is returned.

Parameters

Name	Type	Description	Default
df	DataFrame	Input DataFrame with 'time' column and numeric columns.	required
columns	str | list[str] | None	Column(s) to analyse. If None, uses all numeric columns except 'time'.	None
lags	list[int] | None	Lag values to plot. Each lag gets its own subplot when there are multiple entries. Defaults to [1] when None.	None
seasonality	str | None	When set, colour markers by season. Accepts "month", "quarter", "weekday", "week" or "hour".	None
show_diagonal	bool	Show a diagonal reference line (y = x).	True
show_regression	bool	Show a linear regression line fitted to the data.	False
groups	list[str] | None	Panel group prefixes to plot.	None
facet_n_cols	int	Number of columns in the subplot grid.	3
color_palette	list[str] | None	Custom color palette.	None
show_legend	bool	Whether to show the legend.	True
title	str | None	Plot title.	None
x_label	str | None	X-axis label. Defaults to "y(t-{lag})" per subplot.	None
y_label	str | None	Y-axis label. Defaults to "y(t)" per subplot.	None
width	int | None	Plot width in pixels.	None
height	int | None	Plot height in pixels.	None
marker_size	float	Size of scatter markers.	4.0
marker_opacity	float	Opacity of scatter markers.	0.6

Returns

Type	Description
Figure | dict[str, Figure]	Single figure for univariate data (one column or one panel member). Dictionary of figures keyed by column/member name for multivariate data.

Examples

>>> import polars as pl
>>> from yohou.plotting import plot_lag_scatter

>>> # Create sample time series
>>> df = pl.DataFrame({
...     "time": pl.date_range(pl.date(2020, 1, 1), pl.date(2020, 1, 31), "1d", eager=True),
...     "y": [100 + i % 10 for i in range(31)],
... })

>>> # Plot lag-1 scatter
>>> fig = plot_lag_scatter(df, columns="y", lags=[1])
>>> len(fig.data) > 0
True

See Also

plot_autocorrelation : Plot autocorrelation function.

Source Code

View on GitHub

Show/Hide source

def plot_lag_scatter(
    df: pl.DataFrame,
    *,
    columns: str | list[str] | None = None,
    lags: int | list[int] | None = None,
    seasonality: str | None = None,
    show_diagonal: bool = True,
    show_regression: bool = False,
    groups: list[str] | None = None,
    facet_n_cols: int = 3,
    color_palette: list[str] | None = None,
    show_legend: bool = True,
    title: str | None = None,
    x_label: str | None = None,
    y_label: str | None = None,
    width: int | None = None,
    height: int | None = None,
    marker_size: float = 4.0,
    marker_opacity: float = 0.6,
) -> go.Figure | dict[str, go.Figure]:
    """Plot scatter plots of y(t) vs y(t-lag) for analysing temporal dependencies.

    Creates scatter plots showing the relationship between current values and
    lagged values, useful for identifying AR patterns and lag-specific
    correlations.  When ``lags`` has more than one entry, each lag
    gets its own subplot arranged in a grid.  Optionally colour points by
    season using the ``seasonality`` parameter.

    For panel data, each member gets its own figure with one subplot per lag
    and all groups overlaid by colour.  When there are multiple members, a
    ``dict[str, go.Figure]`` keyed by member name is returned.

    For non-panel multivariate data, each column gets its own figure with lag
    subplots, and a ``dict[str, go.Figure]`` keyed by column name is returned.

    Parameters
    ----------
    df : pl.DataFrame
        Input DataFrame with 'time' column and numeric columns.
    columns : str | list[str] | None, default=None
        Column(s) to analyse. If None, uses all numeric columns except 'time'.
    lags : list[int] | None, default=None
        Lag values to plot. Each lag gets its own subplot when there are
        multiple entries.  Defaults to ``[1]`` when ``None``.
    seasonality : str | None, default=None
        When set, colour markers by season.  Accepts ``"month"``,
        ``"quarter"``, ``"weekday"``, ``"week"`` or ``"hour"``.
    show_diagonal : bool, default=True
        Show a diagonal reference line (y = x).
    show_regression : bool, default=False
        Show a linear regression line fitted to the data.
    groups : list[str] | None, default=None
        Panel group prefixes to plot.
    facet_n_cols : int, default=3
        Number of columns in the subplot grid.
    color_palette : list[str] | None, default=None
        Custom color palette.
    show_legend : bool, default=True
        Whether to show the legend.
    title : str | None, default=None
        Plot title.
    x_label : str | None, default=None
        X-axis label. Defaults to ``"y(t-{lag})"`` per subplot.
    y_label : str | None, default=None
        Y-axis label. Defaults to ``"y(t)"`` per subplot.
    width : int | None, default=None
        Plot width in pixels.
    height : int | None, default=None
        Plot height in pixels.
    marker_size : float, default=4.0
        Size of scatter markers.
    marker_opacity : float, default=0.6
        Opacity of scatter markers.

    Returns
    -------
    go.Figure | dict[str, go.Figure]
        Single figure for univariate data (one column or one panel member).
        Dictionary of figures keyed by column/member name for multivariate
        data.

    Examples
    --------
    >>> import polars as pl
    >>> from yohou.plotting import plot_lag_scatter

    >>> # Create sample time series
    >>> df = pl.DataFrame({
    ...     "time": pl.date_range(pl.date(2020, 1, 1), pl.date(2020, 1, 31), "1d", eager=True),
    ...     "y": [100 + i % 10 for i in range(31)],
    ... })

    >>> # Plot lag-1 scatter
    >>> fig = plot_lag_scatter(df, columns="y", lags=[1])
    >>> len(fig.data) > 0
    True

    See Also
    --------
    [`plot_autocorrelation`][yohou.plotting.plot_autocorrelation] : Plot autocorrelation function.
    """
    # Validate inputs
    validate_plotting_data(df)
    validate_plotting_params(width=width, height=height)

    if lags is None:
        lags = [1]
    elif isinstance(lags, int):
        lags = list(range(1, lags + 1))

    # Auto-detect panel data
    _, _panel_groups = inspect_panel(df)
    if groups is None and columns is None and _panel_groups:
        groups = []

    # Panel path
    if groups is not None:
        panel_cols = resolve_panel_columns(df, groups, columns)
        grouped, all_members = _group_panel_columns(panel_cols)

        n_lags = len(lags)
        figures: dict[str, go.Figure] = {}

        for member in all_members:
            color_mgr = PanelColorManager(color_palette)
            legend_tracker = LegendTracker(show_legend=show_legend)

            ncols_grid = min(facet_n_cols, n_lags)
            nrows_grid = math.ceil(n_lags / ncols_grid)
            subtitles = [f"lag {k}" for k in lags]

            fig = make_subplots(
                rows=nrows_grid,
                cols=ncols_grid,
                subplot_titles=subtitles,
                horizontal_spacing=0.08,
                vertical_spacing=max(0.04, 0.3 / nrows_grid),
            )

            for li, lag in enumerate(lags):
                r = li // ncols_grid + 1
                c = li % ncols_grid + 1
                all_vals: list[float] = []

                for gname, group_cols in grouped.items():
                    col_name = next(
                        (col for col in group_cols if _member_name(col) == member),
                        None,
                    )
                    if col_name is None or col_name not in df.columns:
                        continue

                    dl = (
                        df
                        .select(col_name)
                        .with_columns(
                            pl.col(col_name).shift(lag).alias("lagged"),
                        )
                        .drop_nulls()
                    )
                    if len(dl) == 0:
                        continue

                    group_color = color_mgr.get_color(gname)
                    fig.add_trace(
                        go.Scattergl(
                            x=dl["lagged"],
                            y=dl[col_name],
                            mode="markers",
                            marker={"size": marker_size, "color": group_color, "opacity": marker_opacity},
                            name=gname,
                            legendgroup=gname,
                            showlegend=legend_tracker.should_show(gname),
                            hovertemplate=(
                                f"<b>{gname}: {member}</b><br>y(t-{lag}): %{{x:.2f}}<br>y(t): %{{y:.2f}}<extra></extra>"
                            ),
                        ),
                        row=r,
                        col=c,
                    )

                    all_vals.extend(dl[col_name].to_list())
                    all_vals.extend(dl["lagged"].to_list())

                if show_diagonal and all_vals:
                    vmin = min(all_vals)
                    vmax = max(all_vals)
                    fig.add_trace(
                        go.Scatter(
                            x=[vmin, vmax],
                            y=[vmin, vmax],
                            mode="lines",
                            line={"dash": "dash", "color": "#94a3b8", "width": 1},
                            showlegend=False,
                            hoverinfo="skip",
                        ),
                        row=r,
                        col=c,
                    )

                fig.update_xaxes(title_text=x_label or f"y(t-{lag})", row=r, col=c)
                fig.update_yaxes(title_text=y_label or "y(t)", row=r, col=c)

            cell_size = 260
            default_width = max(DEFAULT_WIDTH, ncols_grid * cell_size + 100)
            default_height = max(400, nrows_grid * cell_size + 100)

            fig = apply_default_layout(
                fig,
                title=title or f"Lag Scatter - {member}",
                x_label=None,
                y_label=None,
                width=width or default_width,
                height=height or default_height,
            )
            fig.update_layout(showlegend=show_legend)
            figures[member] = fig

        if len(figures) == 1:
            return next(iter(figures.values()))
        return figures

    # Non-panel path
    plot_columns = validate_plotting_data(df, columns=columns, exclude=["time"])

    season_labels: list[str] | None = None
    season_colors: list[str] | None = None
    df_aug: pl.DataFrame | None = None

    if seasonality is not None:
        df_aug = _add_season_and_cycle(df, seasonality)
        label_map = _SEASON_LABELS_MAP.get(seasonality)
        seasons_raw = sorted(df_aug["season"].unique().to_list())
        n_seasons = len(seasons_raw)
        season_labels = (
            [label_map[s - 1] for s in seasons_raw] if label_map is not None else [str(s) for s in seasons_raw]
        )
        season_colors = resolve_color_palette(color_palette, n_seasons)

    figures = {}

    for col in plot_columns:
        use_subplots = len(lags) > 1

        if use_subplots:
            n_lags = len(lags)
            ncols = min(facet_n_cols, n_lags)
            nrows = math.ceil(n_lags / ncols)

            subtitles = [f"lag {k}" for k in lags]
            fig = make_subplots(
                rows=nrows,
                cols=ncols,
                subplot_titles=subtitles,
                horizontal_spacing=0.08,
                vertical_spacing=max(0.04, 0.3 / nrows),
            )

            for lag_idx, lag in enumerate(lags):
                r = lag_idx // ncols + 1
                c = lag_idx % ncols + 1
                is_first_cell = lag_idx == 0

                if df_aug is not None and season_labels is not None and season_colors is not None:
                    seasons_raw = sorted(df_aug["season"].unique().to_list())
                    for si, (sval, slabel) in enumerate(zip(seasons_raw, season_labels, strict=False)):
                        sub = df_aug.filter(pl.col("season") == sval)
                        dl = sub.with_columns(pl.col(col).shift(lag).alias("lagged")).drop_nulls()
                        if len(dl) == 0:
                            continue
                        fig.add_trace(
                            go.Scattergl(
                                x=dl["lagged"],
                                y=dl[col],
                                mode="markers",
                                marker={
                                    "size": marker_size,
                                    "color": season_colors[si % len(season_colors)],
                                    "opacity": marker_opacity,
                                },
                                name=slabel,
                                legendgroup=slabel,
                                legendgrouptitle=(
                                    {"text": seasonality.title()}  # ty: ignore[unresolved-attribute]
                                    if is_first_cell and si == 0
                                    else None
                                ),
                                showlegend=is_first_cell,
                                hovertemplate=(
                                    f"<b>{col}</b> ({slabel})<br>"
                                    f"y(t-{lag}): %{{x:.2f}}<br>"
                                    f"y(t): %{{y:.2f}}<extra></extra>"
                                ),
                            ),
                            row=r,
                            col=c,
                        )
                else:
                    colors = resolve_color_palette(color_palette, 1)
                    dl = df.with_columns(pl.col(col).shift(lag).alias("lagged")).drop_nulls()
                    fig.add_trace(
                        go.Scattergl(
                            x=dl["lagged"],
                            y=dl[col],
                            mode="markers",
                            marker={
                                "size": marker_size,
                                "color": colors[0],
                                "opacity": marker_opacity,
                            },
                            name=col,
                            legendgroup=col,
                            showlegend=is_first_cell,
                            hovertemplate=(
                                f"<b>{col}</b><br>y(t-{lag}): %{{x:.2f}}<br>y(t): %{{y:.2f}}<extra></extra>"
                            ),
                        ),
                        row=r,
                        col=c,
                    )

                if show_diagonal:
                    source = df_aug if df_aug is not None else df
                    dl_all = source.with_columns(pl.col(col).shift(lag).alias("lagged")).drop_nulls()
                    vmin = min(
                        float(dl_all[col].min()),  # ty: ignore[invalid-argument-type]
                        float(dl_all["lagged"].min()),  # ty: ignore[invalid-argument-type]
                    )
                    vmax = max(
                        float(dl_all[col].max()),  # ty: ignore[invalid-argument-type]
                        float(dl_all["lagged"].max()),  # ty: ignore[invalid-argument-type]
                    )
                    fig.add_trace(
                        go.Scatter(
                            x=[vmin, vmax],
                            y=[vmin, vmax],
                            mode="lines",
                            line={"dash": "dash", "color": "#94a3b8", "width": 1},
                            showlegend=False,
                            hoverinfo="skip",
                        ),
                        row=r,
                        col=c,
                    )

                fig.update_xaxes(title_text=x_label or f"y(t-{lag})", row=r, col=c)
                fig.update_yaxes(title_text=y_label or "y(t)", row=r, col=c)

            cell_size = 260
            default_width = max(DEFAULT_WIDTH, ncols * cell_size + 100)
            default_height = max(400, nrows * cell_size + 100)

            fig = apply_default_layout(
                fig,
                title=title or "Lag Scatter",
                x_label=None,
                y_label=None,
                width=width or default_width,
                height=height or default_height,
            )
        else:
            fig = go.Figure()
            lag = lags[0]

            if df_aug is not None and season_labels is not None and season_colors is not None:
                seasons_raw = sorted(df_aug["season"].unique().to_list())
                for si, (sval, slabel) in enumerate(zip(seasons_raw, season_labels, strict=False)):
                    sub = df_aug.filter(pl.col("season") == sval)
                    dl = sub.with_columns(pl.col(col).shift(lag).alias("lagged")).drop_nulls()
                    if len(dl) == 0:
                        continue
                    fig.add_trace(
                        go.Scattergl(
                            x=dl["lagged"],
                            y=dl[col],
                            mode="markers",
                            marker={
                                "size": marker_size,
                                "color": season_colors[si % len(season_colors)],
                                "opacity": marker_opacity,
                            },
                            name=slabel,
                            legendgroup=slabel,
                            legendgrouptitle=(
                                {"text": seasonality.title()} if si == 0 else None  # ty: ignore[unresolved-attribute]
                            ),
                            showlegend=True,
                            hovertemplate=(
                                f"<b>{col}</b> ({slabel})<br>y(t-{lag}): %{{x:.2f}}<br>y(t): %{{y:.2f}}<extra></extra>"
                            ),
                        )
                    )
            else:
                colors = resolve_color_palette(color_palette, 1)
                dl = df.with_columns(pl.col(col).shift(lag).alias("lagged")).drop_nulls()
                fig.add_trace(
                    go.Scattergl(
                        x=dl["lagged"],
                        y=dl[col],
                        mode="markers",
                        marker={
                            "size": marker_size,
                            "color": colors[0],
                            "opacity": marker_opacity,
                        },
                        name=f"{col} (lag={lag})",
                        hovertemplate=(f"<b>{col}</b><br>y(t-{lag}): %{{x:.2f}}<br>y(t): %{{y:.2f}}<extra></extra>"),
                    )
                )

            if show_diagonal:
                source = df_aug if df_aug is not None else df
                dl_all = source.with_columns(pl.col(col).shift(lag).alias("lagged")).drop_nulls()
                y_lagged_min = dl_all["lagged"].min()
                y_lagged_max = dl_all["lagged"].max()
                y_current_min = dl_all[col].min()
                y_current_max = dl_all[col].max()
                if y_lagged_min is not None and y_current_min is not None:
                    min_val = min(cast(float, y_lagged_min), cast(float, y_current_min))
                    max_val = max(cast(float, y_lagged_max), cast(float, y_current_max))
                    fig.add_trace(
                        go.Scatter(
                            x=[min_val, max_val],
                            y=[min_val, max_val],
                            mode="lines",
                            line={"dash": "dash", "color": "#94a3b8", "width": 1},
                            showlegend=False,
                            hoverinfo="skip",
                        )
                    )

            if show_regression:
                source = df_aug if df_aug is not None else df
                dl_all = source.with_columns(pl.col(col).shift(lag).alias("lagged")).drop_nulls()
                y_lagged = dl_all["lagged"]
                y_current = dl_all[col]
                x_mean = y_lagged.mean()
                y_mean = y_current.mean()

                numerator = float(((y_lagged - x_mean) * (y_current - y_mean)).sum())
                denominator = float(((y_lagged - x_mean) ** 2).sum())

                if denominator != 0 and x_mean is not None and y_mean is not None:
                    slope = numerator / denominator
                    intercept = cast(float, y_mean) - slope * cast(float, x_mean)

                    y_lagged_min = y_lagged.min()
                    y_lagged_max = y_lagged.max()
                    if y_lagged_min is not None and y_lagged_max is not None:
                        x_line = [cast(float, y_lagged_min), cast(float, y_lagged_max)]
                        y_line = [slope * x + intercept for x in x_line]
                        reg_colors = resolve_color_palette(color_palette, 1)
                        fig.add_trace(
                            go.Scatter(
                                x=x_line,
                                y=y_line,
                                mode="lines",
                                line={"color": reg_colors[0], "width": 2},
                                showlegend=False,
                                hoverinfo="skip",
                            )
                        )

            fig = apply_default_layout(
                fig,
                title=title or f"Lag {lag} Scatter Plot",
                x_label=x_label or f"y(t-{lag})",
                y_label=y_label or "y(t)",
                width=width,
                height=height,
            )

        fig.update_layout(showlegend=show_legend)
        figures[col] = fig

    if len(figures) == 1:
        return next(iter(figures.values()))
    return figures

Tutorials

The following example notebooks use this component:

• How to Visualize Correlations

  ---

  Visualization

  Pairwise correlation heatmaps, scatter matrices, cross-correlation at multiple lags, and lag scatter plots for multivariate time series diagnostics.

  View · Open in marimo