plot_seasonality¶

def plot_seasonality(
    df: pl.DataFrame,
    *,
    columns: str | list[str] | None = None,
    seasonality: str = "month",
    highlight: int | list[int] | None = None,
    groups: list[str] | None = None,
    facet_by: Literal["group", "member"] | None = "member",
    facet_n_cols: int = 2,
    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,
    connect_gaps: bool = False,
    resampler: bool | Literal["widget"] | None = None,
    line_width: float = 2.0,
    highlight_width: float = 3.0,
    fade_opacity: float = 0.25,
    opacity_power: float = 2.0,
) -> go.Figure:
    """Plot seasonal overlay.

    Shows one line per cycle (e.g. year) overlaid on the same axes, with the
    position within each season (e.g. month 1-12) on the x-axis.  This makes
    it easy to compare how the same season changes across cycles.

    For panel data, a bold mean line across all cycles is overlaid on each
    subplot; the non-panel path does not add this mean line.

    Parameters
    ----------
    df : pl.DataFrame
        Input DataFrame with 'time' column and numeric columns.
    columns : str | list[str] | None, default=None
        Column(s) to plot. If None, uses all numeric columns except 'time'.
    seasonality : str, default="month"
        Seasonal frequency: "month", "quarter", "weekday", "week", "hour".
    highlight : int | list[int] | None, default=None
        Cycle(s) to emphasise (e.g. specific years). Highlighted cycles get a
        thicker line; others are faded. If None, all cycles are equally visible.
    groups : list[str] | None, default=None
        Panel group prefixes to plot.
    facet_by : Literal["group", "member"] | None, default="member"
        Faceting axis for panel data.  ``"group"`` creates one subplot per
        group, ``"member"`` one per member.  ``None`` disables faceting.
        Ignored for non-panel data.
    facet_n_cols : int, default=2
        Number of columns in facet grid.
    color_palette : list[str] | None, default=None
        Custom color palette (one color per cycle).
    show_legend : bool, default=True
        Whether to display the legend.
    title : str | None, default=None
        Plot title.
    x_label : str | None, default=None
        X-axis label.
    y_label : str | None, default=None
        Y-axis label.
    width : int | None, default=None
        Plot width in pixels.
    height : int | None, default=None
        Plot height in pixels.
    connect_gaps : bool, default=False
        If True, connect lines across missing data gaps.
    resampler : bool | Literal["widget"] | None, default=None
        Enable plotly-resampler for large datasets.  ``True`` returns a
        ``FigureResampler``, ``"widget"`` a ``FigureWidgetResampler``,
        ``None`` reads from `get_config`.
    line_width : float, default=2.0
        Line width for cycle traces.
    highlight_width : float, default=3.0
        Line width for highlighted cycle traces.
    fade_opacity : float, default=0.25
        Opacity for non-highlighted cycles when ``highlight`` is set.
    opacity_power : float, default=2.0
        Power curve exponent for time-ordered opacity fade.

    Returns
    -------
    go.Figure
        Plotly figure object.

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

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

    >>> # Plot monthly seasonality
    >>> fig = plot_seasonality(df, columns="y", seasonality="month")
    >>> len(fig.data) > 0
    True

    See Also
    --------
    [`plot_subseasonality`][yohou.plotting.plot_subseasonality] : Plot seasonal subseries.
    [`plot_time_series`][yohou.plotting.plot_time_series] : Plot basic time series.
    """
    # Resolve highlight to a set
    if highlight is None:
        highlight_set: set[int] = set()
    elif isinstance(highlight, int):
        highlight_set = {highlight}
    else:
        highlight_set = set(highlight)

    _oldest_alpha = 0.1
    _newest_alpha = 0.9

    validate_plotting_data(df, min_rows=2)
    validate_plotting_params(width=width, height=height)

    # Auto-detect panel data
    if groups is None and columns is None and _auto_detect_panel(df):
        groups = []

    if groups is not None:
        _color_mgr = PanelColorManager(color_palette)
        _legend_tracker = LegendTracker(show_legend=show_legend)

        def _render_season(ctx: RenderContext) -> None:
            """Render seasonal overlay for a single panel group."""
            base = [c for c in ctx.sub_df.columns if c != "time"][0]
            dfs = _add_season_and_cycle(ctx.sub_df, seasonality)
            cycles = sorted(dfs["cycle"].unique().to_list())
            member_color = _color_mgr.get_color(ctx.display_name)
            slabels = _SEASON_LABELS_MAP.get(seasonality)
            n_cycles = len(cycles)

            is_legend_member = _legend_tracker.should_show(ctx.display_name)
            r, g, b = int(member_color[1:3], 16), int(member_color[3:5], 16), int(member_color[5:7], 16)

            # Invisible trace for an opaque legend swatch
            if is_legend_member:
                ctx.fig.add_trace(
                    go.Scatter(
                        x=[None],
                        y=[None],
                        mode="lines",
                        line={"color": member_color, "width": line_width},
                        name=ctx.display_name,
                        legendgroup=ctx.display_name,
                        showlegend=True,
                    ),
                    row=ctx.row,
                    col=ctx.col,
                )

            for ci, cyc in enumerate(cycles):
                cyc_df = dfs.filter(pl.col("cycle") == cyc).sort("season")
                x_vals = (
                    [slabels[int(s) - 1] for s in cyc_df["season"].to_list()]
                    if slabels and all(s <= len(slabels) for s in cyc_df["season"].to_list())
                    else cyc_df["season"].to_list()
                )
                y_vals = cyc_df[base].to_list()

                is_hl = bool(highlight_set) and cyc in highlight_set
                lw = highlight_width if is_hl else line_width
                # Time-ordered opacity via RGBA so legend swatch stays opaque
                if is_hl:
                    alpha = 1.0
                elif highlight_set:
                    alpha = fade_opacity
                else:
                    t = ci / max(n_cycles - 1, 1)
                    alpha = _oldest_alpha + (_newest_alpha - _oldest_alpha) * t ** (1.0 / opacity_power)

                line_rgba = f"rgba({r},{g},{b},{alpha:.2f})"

                ctx.fig.add_trace(
                    go.Scatter(
                        x=x_vals,
                        y=y_vals,
                        mode="lines",
                        line={"color": line_rgba, "width": lw},
                        connectgaps=connect_gaps,
                        hovertemplate=f"<b>{ctx.display_name} - {cyc}</b><br>%{{x}}: %{{y:.2f}}<extra></extra>",
                        name=ctx.display_name,
                        legendgroup=ctx.display_name,
                        showlegend=False,
                    ),
                    row=ctx.row,
                    col=ctx.col,
                )

            # Add aggregated mean line (bold, on top)
            mean_df = dfs.group_by("season", maintain_order=False).agg(pl.col(base).mean()).sort("season")
            if mean_df.height > 0:
                mean_seasons = mean_df["season"].to_list()
                use_labels = slabels is not None and all(s <= len(slabels) for s in mean_seasons)
                mean_x = [slabels[int(s) - 1] for s in mean_seasons] if use_labels else mean_seasons
                mean_y = [float("nan") if v is None else float(v) for v in mean_df[base].to_list()]
                ctx.fig.add_trace(
                    go.Scatter(
                        x=mean_x,
                        y=mean_y,
                        mode="lines",
                        line={"color": member_color, "width": 3},
                        connectgaps=connect_gaps,
                        hovertemplate=f"<b>{ctx.display_name} Mean</b><br>%{{x}}: %{{y:.2f}}<extra></extra>",
                        name=ctx.display_name,
                        legendgroup=ctx.display_name,
                        showlegend=False,
                    ),
                    row=ctx.row,
                    col=ctx.col,
                )

        effective_facet_by = facet_by or "member"
        fig = facet_figure(
            df,
            _render_season,
            groups=groups,
            columns=columns,
            facet_by=effective_facet_by,
            facet_n_cols=facet_n_cols,
            title=title or "Seasonal Pattern",
            x_label=x_label or seasonality.capitalize(),
            y_label=y_label,
            width=width,
            height=height,
            shared_xaxes=False,
            resampler=resampler,
        )
        fig.update_layout(showlegend=show_legend)
        return fig

    # Non-panel case: column-mode facet_figure
    plot_columns = validate_plotting_data(df, columns=columns, exclude=["time"])
    season_labels = _SEASON_LABELS_MAP.get(seasonality)
    df_aug = _add_season_and_cycle(df, seasonality)
    cycles = sorted(df_aug["cycle"].unique().to_list())
    _colors = resolve_color_palette(color_palette, len(plot_columns))
    _col_colors = dict(zip(plot_columns, _colors, strict=False))
    n_cycles = len(cycles)

    def _render_season(ctx: RenderContext) -> None:
        """Render seasonal pattern for one column into a subplot."""
        base = ctx.display_name
        col_color = _col_colors[base]
        r, g, b = int(col_color[1:3], 16), int(col_color[3:5], 16), int(col_color[5:7], 16)
        for ci, cyc in enumerate(cycles):
            cyc_df = df_aug.filter(pl.col("cycle") == cyc).sort("season")
            x_vals = (
                [season_labels[int(s) - 1] for s in cyc_df["season"].to_list()]
                if season_labels and all(s <= len(season_labels) for s in cyc_df["season"].to_list())
                else cyc_df["season"].to_list()
            )
            is_hl = bool(highlight_set) and cyc in highlight_set
            lw = highlight_width if is_hl else line_width
            if is_hl:
                alpha = 1.0
            elif highlight_set:
                alpha = fade_opacity
            else:
                t = ci / max(n_cycles - 1, 1)
                alpha = _oldest_alpha + (_newest_alpha - _oldest_alpha) * t ** (1.0 / opacity_power)
            line_rgba = f"rgba({r},{g},{b},{alpha:.2f})"
            ctx.fig.add_trace(
                go.Scatter(
                    x=x_vals,
                    y=cyc_df[base].to_list(),
                    mode="lines",
                    line={"color": line_rgba, "width": lw},
                    name=base,
                    legendgroup=base,
                    showlegend=False,
                    connectgaps=connect_gaps,
                    hovertemplate=f"<b>{base} - {cyc}</b><br>%{{x}}: %{{y:.2f}}<extra></extra>",
                ),
                row=ctx.row,
                col=ctx.col,
            )

    fig = facet_figure(
        df,
        _render_season,
        columns=plot_columns,
        facet_n_cols=facet_n_cols,
        title=title or "Seasonal Pattern",
        x_label=x_label or seasonality.capitalize(),
        y_label=y_label or "Value",
        width=width,
        height=height,
        shared_xaxes=False,
        resampler=resampler,
    )
    if season_labels:
        fig.update_xaxes(
            tickmode="array",
            tickvals=list(range(len(season_labels))),
            ticktext=season_labels,
        )
    fig.update_layout(showlegend=show_legend)

    return fig