plot_splits¶

def plot_splits(
    y: pl.DataFrame,
    splitter: BaseSplitter,
    *,
    X_actual: pl.DataFrame | None = None,
    train_color: str | None = None,
    test_color: str | None = None,
    train_name: str = "Train",
    test_name: str = "Test",
    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,
    resampler: bool | Literal["widget"] | None = None,
    line_width: float = 10.0,
) -> go.Figure:
    """
    Plot cross-validation splits as a timeline visualization.

    Creates a horizontal bar chart showing train/test splits for each fold,
    useful for understanding temporal CV strategies like expanding or sliding windows.

    Parameters
    ----------
    y : pl.DataFrame
        Target time series with mandatory "time" column.
    splitter : BaseSplitter
        A yohou splitter instance (e.g., ExpandingWindowSplitter, SlidingWindowSplitter).
    X_actual : pl.DataFrame or None, default=None
        Actual features passed to ``splitter.split()``. Not used for
        splitting but accepted for API consistency.
    train_color : str | None, default=None
        Color for train segments. If None, uses first color from yohou palette.
    test_color : str | None, default=None
        Color for test segments. If None, uses second color from yohou palette.
    train_name : str, default="Train"
        Legend label for the train segments.
    test_name : str, default="Test"
        Legend label for the test segments.
    show_legend : bool, default=True
        Whether to show the legend.
    title : str | None, default=None
        Plot title. Defaults to "Cross-Validation Splits".
    x_label : str | None, default=None
        X-axis label. Defaults to "Time".
    y_label : str | None, default=None
        Y-axis label. Defaults to "Fold".
    width : int | None, default=None
        Plot width in pixels.
    height : int | None, default=None
        Plot height in pixels. Defaults to 300 + n_splits * 30.
    resampler : bool | Literal["widget"] | None, default=None
        Enable plotly-resampler for large datasets.  ``True`` creates a
        ``FigureResampler`` (Dash-based server callback); ``"widget"`` creates
        a ``FigureWidgetResampler`` (notebook-native widget); ``False`` or
        ``None`` uses a plain ``go.Figure``.
    line_width : float, default=10.0
        Width of the train/test bars.

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

    Raises
    ------
    TypeError
        If y is not a Polars DataFrame or splitter is not a BaseSplitter.
    ValueError
        If DataFrame is empty or missing 'time' column.
        If the splitter produces no splits for the given ``y`` (check the
        ``n_splits`` and ``test_size`` configuration).

    Examples
    --------
    >>> import polars as pl
    >>> from yohou.plotting import plot_splits
    >>> from yohou.model_selection import ExpandingWindowSplitter

    >>> # Create sample data
    >>> y = pl.DataFrame({
    ...     "time": pl.date_range(pl.date(2020, 1, 1), pl.date(2020, 12, 31), "1d", eager=True),
    ...     "value": list(range(366)),
    ... })

    >>> # Create splitter and plot
    >>> splitter = ExpandingWindowSplitter(n_splits=3, test_size=30)
    >>> fig = plot_splits(y, splitter)
    >>> len(fig.data) > 0
    True

    See Also
    --------
    [`plot_cv_results_scatter`][yohou.plotting.plot_cv_results_scatter] : Plot hyperparameter search results.
    [`ExpandingWindowSplitter`][yohou.model_selection.ExpandingWindowSplitter] : Expanding window cross-validation.
    [`SlidingWindowSplitter`][yohou.model_selection.SlidingWindowSplitter] : Sliding window cross-validation.
    """
    # Validate inputs
    validate_plotting_data(y)
    validate_plotting_params(width=width, height=height)

    if not isinstance(splitter, BaseSplitter):
        msg = f"Expected BaseSplitter, got {type(splitter).__name__}"
        raise TypeError(msg)

    # Get colors
    colors = resolve_color_palette(None, 2)
    train_color = train_color or colors[0]
    test_color = test_color or colors[1]

    # Get splits
    splits = list(splitter.split(y, X_actual))
    n_splits = len(splits)
    if n_splits == 0:
        msg = (
            "Splitter produced no splits. Check that 'y' has enough rows for the configured 'n_splits' and 'test_size'."
        )
        raise ValueError(msg)

    # Create figure
    fig = _create_figure(resampler)

    # Get time column
    times = y["time"]

    # Plot each split (one row per fold: train segment + test segment).
    for i, (train_idx, test_idx) in enumerate(splits):
        _add_split_row(
            fig,
            times,
            train_idx,
            test_idx,
            f"Fold {i + 1}",
            train_color=train_color,
            test_color=test_color,
            train_name=train_name,
            test_name=test_name,
            line_width=line_width,
            show_train_legend=(i == 0),
            show_test_legend=(i == 0),
        )

    # Set default labels
    title_default = title or "Cross-Validation Splits"
    x_label_default = x_label or "Time"
    y_label_default = y_label or "Fold"

    # Calculate height based on number of splits
    height_default = height or (300 + n_splits * 30)

    fig = apply_default_layout(
        fig,
        title=title_default,
        x_label=x_label_default,
        y_label=y_label_default,
        width=width,
        height=height_default,
    )
    fig.update_layout(showlegend=show_legend)

    return fig