check_class_proba_prediction_structure¶

def check_class_proba_prediction_structure(forecaster, y_test: pl.DataFrame) -> None:
    """Check class-probability predictions have correct column structure.

    Validates that ``predict_class_proba`` output contains ``"vintage_time"``
    and ``"time"`` columns, plus one probability column per target and class.
    For non-panel data the columns follow ``{target}_proba_{class}``; for panel
    data the format is ``{group_prefix}__{target}_proba_{class}`` where
    ``group_prefix`` is the entity identifier (e.g. ``store_1``).

    Parameters
    ----------
    forecaster : BaseClassProbaForecaster
        Fitted class-probability forecaster instance.
    y_test : pl.DataFrame
        Test target data.

    Raises
    ------
    AssertionError
        If prediction structure is incorrect.

    """
    forecasting_horizon = min(3, len(y_test))
    y_pred = forecaster.predict_class_proba(forecasting_horizon=forecasting_horizon)

    assert "vintage_time" in y_pred.columns, "Class-proba predictions must have 'vintage_time'"
    assert "time" in y_pred.columns, "Class-proba predictions must have 'time'"

    # Should NOT have interval columns
    interval_cols = [col for col in y_pred.columns if "_lower_" in col or "_upper_" in col]
    assert len(interval_cols) == 0, f"Class-proba predictions should not have interval columns, found: {interval_cols}"

    # Should have _proba_ columns for each target and class
    _, y_panel_groups = inspect_panel(y_test)
    if len(y_panel_groups) > 0:
        # Panel data: classes_ keys are unprefixed target names
        for target_col, class_labels in forecaster.classes_.items():
            for group_prefix in y_panel_groups:
                for label in class_labels:
                    col_name = f"{group_prefix}__{target_col}_proba_{label}"
                    assert col_name in y_pred.columns, f"Missing probability column: {col_name}"
    else:
        for target_col, class_labels in forecaster.classes_.items():
            for label in class_labels:
                col_name = f"{target_col}_proba_{label}"
                assert col_name in y_pred.columns, f"Missing probability column: {col_name}"