Source code for knowledgespaces.viz.diagnostics

"""Optional matplotlib diagnostics for complete BLIM response tables."""

from __future__ import annotations

from typing import TYPE_CHECKING, Literal, cast

import numpy as np

from knowledgespaces.viz.hasse import _MATPLOTLIB_IMPORT_ERROR

if TYPE_CHECKING:
    from matplotlib.axes import Axes
    from matplotlib.figure import Figure

    from knowledgespaces.estimation.diagnostics import BLIMResiduals




def plot_blim_residuals(
    residuals: BLIMResiduals,
    *,
    kind: Literal["deviance", "pearson"] = "deviance",
    ax: Axes | None = None,
    figsize: tuple[float, float] = (8, 5),
) -> Figure:
    """Plot signed residuals against predicted response probabilities.

    Uses all cells, including zero observed counts. The horizontal zero
    line is a visual reference, not a significance threshold. A table with
    infinite residuals is rejected explicitly; inspect impossible observed
    cells in the numerical result before plotting. Requires the viz extra.
    """
    if kind not in ("deviance", "pearson"):
        raise ValueError("kind must be 'deviance' or 'pearson'.")
    values = getattr(residuals, kind)
    if not np.isfinite(values).all():
        raise ValueError("Residuals are nonfinite; inspect model-impossible observed cells.")
    try:
        import matplotlib.pyplot as plt
    except ImportError as error:
        raise ImportError(_MATPLOTLIB_IMPORT_ERROR) from error
    if ax is None:
        fig, ax = plt.subplots(figsize=figsize)
    else:
        fig = cast("Figure", ax.get_figure())
    ax.scatter(np.exp(residuals.log_probabilities), values, s=18, alpha=0.7)
    ax.axhline(0, color="0.4", linestyle="--", linewidth=1)
    ax.set(xlabel="Predicted response probability", ylabel=f"{kind.capitalize()} residual")
    return fig