plot_calibration¶

def plot_calibration(
    y_pred: pl.DataFrame,
    y_truth: pl.DataFrame,
    coverage_rates: list[StrictFloat] | None = None,
    *,
    type: Literal["interval", "quantile"] = "quantile",
    columns: str | list[str] | None = None,
    target: str | None = None,
    n_bins: int = 10,
    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,
    line_width: float = 2.0,
    line_opacity: float = 1.0,
    reference_color: str = "#1e293b",
    reference_width: float = 3.0,
    reference_dash: str = "dash",
) -> go.Figure:
    """Plot calibration for interval or class-probability forecasts.

    Automatically detects the prediction type from column names:

    - **Interval / quantile predictions**: columns named
      ``"{target}_upper_{rate}"`` / ``"{target}_lower_{rate}"``.  With
      ``type="interval"`` the two-sided band coverage
      ``P(lower <= y <= upper)`` is compared against the nominal
      *coverage_rates*.  With ``type="quantile"`` each bound is treated as
      a single quantile (the band at coverage ``c`` yields quantile levels
      ``(1 - c) / 2`` and ``(1 + c) / 2``, and the median/point column adds
      the ``0.5`` level), and the empirical CDF coverage ``P(y <= q)`` is
      compared against the nominal quantile level.
    - **Class-probability predictions**: columns containing
      ``"_proba_"`` are binned and compared against empirical
      frequencies (reliability diagram).  *type* is ignored.

    A well-calibrated model has points close to the diagonal.

    Parameters
    ----------
    y_pred : pl.DataFrame
        Predicted values.  Either prediction intervals with
        ``"{target}_upper_{rate}"`` / ``"{target}_lower_{rate}"``
        columns, or class-probability predictions with
        ``"{target}_proba_{class}"`` columns.
    y_truth : pl.DataFrame
        Ground truth values.
    coverage_rates : list of float or None, default=None
        Nominal coverage rates for interval calibration (e.g.,
        ``[0.9, 0.95]``).  When ``None``, every coverage rate present as a
        matched ``_lower_{rate}`` / ``_upper_{rate}`` pair in *y_pred* is
        detected and plotted.  Ignored for class-probability predictions.
    type : {"interval", "quantile"}, default="quantile"
        Calibration mode for bound-based predictions.  ``"quantile"``
        reinterprets each bound (and the median column) as a single quantile
        and plots the empirical CDF coverage ``P(y <= q)`` against the nominal
        quantile level; ``"interval"`` plots two-sided band coverage against
        the nominal coverage rate.  Ignored for class-probability predictions.
    columns : str | list[str] | None, default=None
        Target column name(s) for interval calibration.  When
        *groups* is set this acts as a member postfix
        filter.  Ignored for class-probability predictions (use
        *target* instead).
    target : str or None, default=None
        Target column name for class-probability calibration.  Required
        when multiple targets are present in the probability columns.
        Ignored for interval calibration.
    n_bins : int, default=10
        Number of bins for class-probability calibration.  Ignored for
        interval calibration.
    groups : list[str] | None, default=None
        Panel group prefixes for faceted subplots.  When provided, each
        resolved panel column gets its own subplot.
    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 the facet grid when *groups* is
        used.
    color_palette : list[str] | None, default=None
        Custom color palette as hex codes. If None, uses yohou palette.
    show_legend : bool, default=True
        Whether to show the legend.
    title : str | None, default=None
        Plot title. Defaults to ``"Calibration plot"`` for intervals or
        ``"Reliability Diagram"`` for class probabilities.
    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.
    line_width : float, default=2.0
        Width of the calibration line.
    line_opacity : float, default=1.0
        Opacity of the calibration line.  Ignored for class-probability
        predictions.
    reference_color : str, default="#1e293b"
        Colour of the perfect-calibration reference line.
    reference_width : float, default=3.0
        Width of the reference line.  Ignored for class-probability
        predictions (the reference width is fixed at 2.0 there).
    reference_dash : str, default="dash"
        Dash style of the reference line.

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

    Raises
    ------
    ValueError
        If interval columns are missing, *coverage_rates* is not
        provided for interval predictions, or class-probability columns
        are ambiguous.

    Examples
    --------
    Interval calibration:

    >>> import polars as pl
    >>> import numpy as np
    >>> from datetime import datetime, timedelta
    >>> from yohou.plotting import plot_calibration

    >>> # Create sample data
    >>> n = 100
    >>> rng = np.random.default_rng(0)
    >>> times = [datetime(2020, 1, 1) + timedelta(days=i) for i in range(n)]
    >>> y = rng.standard_normal(n)
    >>> y_truth = pl.DataFrame({"time": times, "y": y})
    >>> # Deterministic symmetric prediction intervals centred on y.
    >>> y_pred_int = pl.DataFrame({
    ...     "time": times,
    ...     "y_upper_0.9": y + 1.65,
    ...     "y_lower_0.9": y - 1.65,
    ...     "y_upper_0.95": y + 1.96,
    ...     "y_lower_0.95": y - 1.96,
    ... })

    >>> # Plot calibration
    >>> fig = plot_calibration(y_pred_int, y_truth, coverage_rates=[0.9, 0.95])
    >>> len(fig.data)
    2

    Class-probability calibration (reliability diagram):

    >>> from datetime import datetime
    >>> y_pred_proba = pl.DataFrame({
    ...     "time": [datetime(2020, 1, i) for i in range(1, 11)],
    ...     "w_proba_sunny": [0.8, 0.7, 0.6, 0.5, 0.4, 0.3, 0.2, 0.1, 0.9, 0.85],
    ...     "w_proba_rainy": [0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 0.1, 0.15],
    ... })
    >>> y_truth_cat = pl.DataFrame({
    ...     "time": [datetime(2020, 1, i) for i in range(1, 11)],
    ...     "w": ["sunny", "sunny", "sunny", "rainy", "rainy", "rainy", "rainy", "rainy", "sunny", "sunny"],
    ... })
    >>> fig = plot_calibration(y_pred_proba, y_truth_cat)
    >>> isinstance(fig, go.Figure)
    True

    See Also
    --------
    [`plot_forecast`][yohou.plotting.plot_forecast] : Plot forecast with optional prediction intervals.
    [`plot_residuals`][yohou.plotting.plot_residuals] : Residual diagnostics with panel facets.
    """
    # Validate inputs
    if not isinstance(y_truth, pl.DataFrame):
        msg = f"Expected pl.DataFrame for y_truth, got {y_truth.__class__.__name__}"
        raise TypeError(msg)
    if not isinstance(y_pred, pl.DataFrame):
        msg = f"Expected pl.DataFrame for y_pred, got {y_pred.__class__.__name__}"
        raise TypeError(msg)
    validate_plotting_params(width=width, height=height)

    # Detect class-probability columns
    proba_cols = [c for c in y_pred.columns if "_proba_" in c]
    if proba_cols:
        return _plot_calibration_class_proba(
            y_pred=y_pred,
            y_truth=y_truth,
            target=target,
            n_bins=n_bins,
            color_palette=color_palette,
            title=title,
            x_label=x_label,
            y_label=y_label,
            width=width,
            height=height,
            show_legend=show_legend,
            line_width=line_width,
            reference_color=reference_color,
            reference_dash=reference_dash,
        )

    # Interval / quantile calibration path. When the caller does not pin the
    # rates, plot every coverage rate present as a matched lower/upper pair.
    if coverage_rates is None:
        coverage_rates = _detect_coverage_rates(y_pred)
        if not coverage_rates:
            msg = (
                "No interval columns found in y_pred for calibration. Expected "
                "'{base}_lower_{rate}' / '{base}_upper_{rate}' columns, or pass "
                "coverage_rates explicitly (e.g., [0.9, 0.95])."
            )
            raise ValueError(msg)

    # Mode-dependent labels: quantile mode plots one-sided CDF coverage against
    # the nominal quantile level, interval mode two-sided coverage vs the rate.
    _is_quantile = type == "quantile"
    _default_title = "Quantile calibration" if _is_quantile else "Calibration plot"
    _default_x_label = "Nominal quantile level" if _is_quantile else "Nominal coverage"
    _hover_x = "Quantile" if _is_quantile else "Nominal"

    def _calibration_xy(panel_col: str) -> tuple[list[float], list[float]]:
        """Return ``(x, empirical)`` for one column under the active mode."""
        truth_vals = y_truth[panel_col].to_numpy().flatten()
        if _is_quantile:
            return _compute_quantile_coverages(truth_vals, y_pred, panel_col, coverage_rates)
        return list(coverage_rates), _compute_empirical_coverages(truth_vals, y_pred, panel_col, coverage_rates)

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

    if groups is not None:
        panel_cols = resolve_panel_columns(y_truth, groups, columns)

        # Group columns by panel prefix, collect unique members
        grouped, all_members = _group_panel_columns(panel_cols)

        n_groups = len(grouped)
        n_cols_grid = min(n_groups, facet_n_cols)
        n_rows = (n_groups + n_cols_grid - 1) // n_cols_grid

        fig = make_subplots(
            rows=n_rows,
            cols=n_cols_grid,
            subplot_titles=list(grouped.keys()),
            shared_xaxes=True,
            shared_yaxes=True,
            vertical_spacing=_subplot_spacing(n_rows),
            horizontal_spacing=0.08,
        )

        palette = resolve_color_palette(color_palette, len(all_members))
        legend_tracker = LegendTracker(show_legend=show_legend)
        for group_idx, (_, group_cols) in enumerate(grouped.items()):
            row = group_idx // n_cols_grid + 1
            col_idx = group_idx % n_cols_grid + 1

            for panel_col in group_cols:
                member_name = _member_name(panel_col)
                member_idx = all_members.index(member_name)

                x_vals, emp_cov = _calibration_xy(panel_col)

                fig.add_trace(
                    go.Scatter(
                        x=x_vals,
                        y=emp_cov,
                        mode="lines+markers",
                        name=member_name,
                        line={"color": palette[member_idx], "width": line_width},
                        opacity=line_opacity,
                        showlegend=legend_tracker.should_show(member_name),
                        legendgroup=member_name,
                        hovertemplate=(
                            f"<b>%{{fullData.name}}</b><br>{_hover_x}: %{{x:.2f}}<br>Coverage: %{{y:.3f}}<extra></extra>"
                        ),
                    ),
                    row=row,
                    col=col_idx,
                )

            fig.add_trace(
                go.Scatter(
                    x=[0, 1],
                    y=[0, 1],
                    mode="lines",
                    name="Perfect",
                    line={"color": reference_color, "width": reference_width, "dash": reference_dash},
                    showlegend=legend_tracker.should_show("Perfect"),
                    legendgroup="perfect",
                ),
                row=row,
                col=col_idx,
            )

        row_height = 300
        default_height = max(row_height * n_rows, 400)
        fig = apply_default_layout(
            fig,
            title=title or _default_title,
            x_label=x_label or _default_x_label,
            y_label=y_label or "Empirical coverage",
            width=width,
            height=height or default_height,
            # Reliability curves and the diagonal must stay linear; never step
            # them via the global line_shape config.
            apply_line_shape=False,
        )
        fig.update_layout(showlegend=show_legend)
        # Calibration axes are coverage probabilities: pin both to the full
        # [0, 1] range and force a square aspect (1:1 scale, box constrained to
        # the domain) so every subplot reads on the same scale and the diagonal
        # sits at a true 45 degrees.  ``make_subplots`` shares axes within each
        # row/column via ``matches``; a matched axis cannot also honour its own
        # ``scaleanchor``, which leaves only the first column square.  Clearing
        # ``matches`` lets each cell apply its own square constraint.
        fig.update_xaxes(range=[0, 1], constrain="domain")
        fig.update_yaxes(range=[0, 1], constrain="domain")
        for _xaxis in [_k for _k in fig.layout if _k.startswith("xaxis")]:
            fig.layout[_xaxis].matches = None
        for _yaxis in [_k for _k in fig.layout if _k.startswith("yaxis")]:
            fig.layout[_yaxis].matches = None
            fig.layout[_yaxis].scaleanchor = "x" + _yaxis.removeprefix("yaxis")
            fig.layout[_yaxis].scaleratio = 1
        return fig

    if columns is not None:
        target_columns = [columns] if isinstance(columns, str) else list(columns)
        for col in target_columns:
            if col not in y_truth.columns:
                msg = f"Target column '{col}' not found in y_truth"
                raise ValueError(msg)
    else:
        target_columns = [c for c in y_truth.columns if c not in ("time", "vintage_time")]
        if not target_columns:
            msg = "y_truth has no non-time columns"
            raise ValueError(msg)

    palette = resolve_color_palette(color_palette, len(target_columns))

    fig = go.Figure()

    for col_idx, target_column in enumerate(target_columns):
        x_vals, emp_cov = _calibration_xy(target_column)

        trace_name = target_column if len(target_columns) > 1 else "Empirical coverage"
        fig.add_trace(
            go.Scatter(
                x=x_vals,
                y=emp_cov,
                mode="lines+markers",
                name=trace_name,
                line={"color": palette[col_idx % len(palette)], "width": line_width},
                opacity=line_opacity,
                hovertemplate=f"<b>{trace_name}</b><br>{_hover_x}: %{{x:.2f}}<br>Coverage: %{{y:.3f}}<extra></extra>",
            )
        )

    # Reference diagonal (always exactly one), spanning the full [0, 1] range.
    fig.add_trace(
        go.Scatter(
            x=[0, 1],
            y=[0, 1],
            mode="lines",
            name="Perfect calibration",
            line={"color": reference_color, "width": reference_width, "dash": reference_dash},
            hovertemplate="<b>Perfect</b><br>Coverage: %{x:.2f}<extra></extra>",
        )
    )

    fig = apply_default_layout(
        fig,
        title=title or _default_title,
        x_label=x_label or _default_x_label,
        y_label=y_label or "Empirical coverage",
        width=width,
        height=height,
        # Reliability curves and the diagonal must stay linear; never step them
        # via the global line_shape config.
        apply_line_shape=False,
    )
    fig.update_layout(showlegend=show_legend)
    # Calibration axes are coverage probabilities: pin both to the full [0, 1]
    # range and force a square aspect (1:1 scale, box constrained to the domain)
    # so the diagonal sits at a true 45 degrees.
    fig.update_xaxes(range=[0, 1], constrain="domain")
    fig.update_yaxes(range=[0, 1], scaleanchor="x", scaleratio=1, constrain="domain")

    return fig