check_predict_time_columns¶

def check_predict_time_columns(forecaster, y_test: pl.DataFrame, X_actual_test: pl.DataFrame | None = None) -> None:
    """Check predictions have vintage_time and time columns.

    Parameters
    ----------
    forecaster : BaseForecaster
        Fitted forecaster instance
    y_test : pl.DataFrame
        Test target data
    X_actual_test : pl.DataFrame, optional
        Test features

    Raises
    ------
    AssertionError
        If predictions lack required time columns

    """
    forecasting_horizon = min(3, len(y_test))

    # Check if forecaster is an interval forecaster
    if hasattr(forecaster, "predict_interval"):
        y_pred = forecaster.predict_interval(forecasting_horizon=forecasting_horizon)
    else:
        y_pred = forecaster.predict(forecasting_horizon=forecasting_horizon)

    assert "vintage_time" in y_pred.columns, "Predictions must have 'vintage_time' column"
    assert "time" in y_pred.columns, "Predictions must have 'time' column"

    # Validate shapes. Panel forecasters may stack predictions per group, so the
    # total row count is a multiple of the horizon rather than the horizon
    # itself. Only assert the strict single-group row count for non-panel
    # forecasters (groups_ is None); for panel forecasters, require the per-group
    # row count (rows // n_groups) to equal the horizon.
    groups = getattr(forecaster, "groups_", None)
    if groups is None:
        assert len(y_pred) == forecasting_horizon, (
            f"Predictions should have {forecasting_horizon} rows, got {len(y_pred)}"
        )
    else:
        n_groups = len(groups)
        assert len(y_pred) % forecasting_horizon == 0, (
            f"Panel predictions should have a multiple of {forecasting_horizon} rows, got {len(y_pred)}"
        )
        assert len(y_pred) in (forecasting_horizon, n_groups * forecasting_horizon), (
            f"Panel predictions should have {forecasting_horizon} (wide) or "
            f"{n_groups * forecasting_horizon} (stacked) rows, got {len(y_pred)}"
        )

    # Validate time column types
    assert isinstance(y_pred["vintage_time"].dtype, pl.Datetime | pl.Date), (
        "vintage_time must be Datetime or Date dtype"
    )
    assert isinstance(y_pred["time"].dtype, pl.Datetime | pl.Date), "time must be Datetime or Date dtype"