# SPDX-License-Identifier: MIT
# Copyright (c) 2025 Lucy Eduardo Arias
"""High-level convenience functions for Aurora-GLM.
This module provides unified helper functions that work with any Aurora result:
- summary(result): Formatted model summary
- plot(result): Diagnostic plots
- compare(*results): Model comparison
These are also available at the package level:
>>> import aurora
>>> aurora.summary(result)
>>> aurora.plot(result)
>>> aurora.compare(model1, model2, model3)
Examples
--------
>>> from aurora import fit_glm, summary, plot, compare
>>>
>>> # Fit models
>>> result1 = fit_glm(X, y, family=Gaussian())
>>> result2 = fit_glm(X, y, family=Poisson())
>>>
>>> # Get summaries
>>> print(summary(result1))
>>>
>>> # Plot diagnostics
>>> plot(result1, kind='diagnostics')
>>>
>>> # Compare models
>>> compare(result1, result2, criterion='aic')
"""
from __future__ import annotations
from collections.abc import Sequence
from typing import TYPE_CHECKING, Any, Literal
import numpy as np
if TYPE_CHECKING:
pass
# =============================================================================
# Summary Function
# =============================================================================
[docs]
def summary(
result: Any,
*,
style: Literal["default", "brief", "detailed"] = "default",
print_output: bool = True,
) -> str:
"""Generate a formatted summary of a fitted model result.
This is a unified interface for summarizing any Aurora model result.
It automatically dispatches to the appropriate summary method based
on the result type.
Parameters
----------
result : ModelResult
A fitted model result from fit_glm, fit_gam, fit_gamm, etc.
style : {'default', 'brief', 'detailed'}, default='default'
Summary style:
- 'default': Standard coefficient table with fit statistics
- 'brief': One-line summary with key metrics
- 'detailed': Full output including diagnostics
print_output : bool, default=True
If True, print the summary. If False, only return the string.
Returns
-------
summary_str : str
Formatted summary string.
Examples
--------
>>> from aurora import fit_glm, summary, Gaussian
>>> result = fit_glm(X, y, family=Gaussian())
>>> print(summary(result))
======================================================================
Linear Model Results
======================================================================
No. Observations: 100 R-squared: 0.8542
...
>>> # Brief summary
>>> summary(result, style='brief')
'GLM(n=100, converged=True, R²=0.854)'
>>> # Detailed with diagnostics
>>> summary(result, style='detailed')
Notes
-----
The function works with any result object that has a `.summary()` method.
For custom result classes, implement `.summary()` to integrate with this.
See Also
--------
plot : Diagnostic plots for model results
compare : Compare multiple models
"""
# Check for summary method
if not hasattr(result, "summary"):
raise TypeError(
f"Cannot summarize object of type {type(result).__name__}. "
"Expected a model result with a .summary() method."
)
# Generate summary based on style
if style == "default":
summary_str = result.summary()
elif style == "brief":
summary_str = _brief_summary(result)
elif style == "detailed":
summary_str = _detailed_summary(result)
else:
raise ValueError(f"Unknown style '{style}'. Use 'default', 'brief', or 'detailed'.")
if print_output:
print(summary_str)
return str(summary_str)
def _brief_summary(result: Any) -> str:
"""Generate a one-line brief summary of a model result.
Extracts key attributes (model type, number of observations,
convergence status, R-squared, log-likelihood, AIC) into a
compact ``TypeName(key=val, ...)`` string.
Parameters
----------
result : ModelResult
A fitted model result object.
Returns
-------
summary_str : str
One-line summary string.
"""
parts = []
# Class name
class_name = type(result).__name__
if class_name.endswith("Result"):
class_name = class_name[:-6] # Remove 'Result' suffix
parts.append(class_name)
# Observations
if hasattr(result, "n_obs_"):
parts.append(f"n={result.n_obs_}")
elif hasattr(result, "n_observations"):
parts.append(f"n={result.n_observations}")
# Convergence
if hasattr(result, "converged_"):
parts.append(f"converged={result.converged_}")
elif hasattr(result, "converged"):
parts.append(f"converged={result.converged}")
# R² for linear models
if hasattr(result, "r_squared"):
r2 = result.r_squared
if not np.isnan(r2):
parts.append(f"R²={r2:.3f}")
# Log-likelihood for mixed models
if hasattr(result, "log_likelihood_"):
parts.append(f"LL={result.log_likelihood_:.2f}")
# AIC if available
if hasattr(result, "aic"):
aic = result.aic
if not np.isnan(aic):
parts.append(f"AIC={aic:.1f}")
return f"{parts[0]}({', '.join(parts[1:])})"
def _detailed_summary(result: Any) -> str:
"""Generate detailed summary with diagnostics.
Produces the standard summary output followed by a Diagnostics
section containing residual quantiles, R-squared, adjusted
R-squared, AIC, and BIC (where available).
Parameters
----------
result : ModelResult
A fitted model result object with a ``.summary()`` method.
Returns
-------
summary_str : str
Multi-line detailed summary string.
"""
lines = []
# Base summary
lines.append(result.summary())
lines.append("")
# Diagnostics section
lines.append("=" * 70)
lines.append("Diagnostics")
lines.append("=" * 70)
# Residual statistics
if hasattr(result, "residuals"):
try:
resid = result.residuals
lines.append("Residuals:")
lines.append(f" Min: {np.min(resid):>12.4f}")
lines.append(f" 1Q: {np.percentile(resid, 25):>12.4f}")
lines.append(f" Median: {np.median(resid):>12.4f}")
lines.append(f" 3Q: {np.percentile(resid, 75):>12.4f}")
lines.append(f" Max: {np.max(resid):>12.4f}")
lines.append("")
except Exception:
pass
# Fit statistics
if hasattr(result, "r_squared"):
r2 = result.r_squared
if not np.isnan(r2):
lines.append(f"R-squared: {r2:>12.4f}")
if hasattr(result, "adj_r_squared"):
adj_r2 = result.adj_r_squared
if not np.isnan(adj_r2):
lines.append(f"Adjusted R-squared: {adj_r2:>12.4f}")
if hasattr(result, "aic"):
aic = getattr(result, "aic", np.nan)
if not np.isnan(aic):
lines.append(f"AIC: {aic:>12.2f}")
if hasattr(result, "bic"):
bic = getattr(result, "bic", np.nan)
if not np.isnan(bic):
lines.append(f"BIC: {bic:>12.2f}")
lines.append("=" * 70)
return "\n".join(lines)
# =============================================================================
# Plot Function
# =============================================================================
[docs]
def plot(
result: Any,
*,
kind: Literal["diagnostics", "residuals", "qq", "smooth", "all"] = "diagnostics",
ax: Any = None,
**kwargs,
) -> Any:
"""Create diagnostic plots for a fitted model result.
This is a unified interface for plotting any Aurora model result.
It automatically creates appropriate plots based on the model type.
Parameters
----------
result : ModelResult
A fitted model result from fit_glm, fit_gam, fit_gamm, etc.
kind : {'diagnostics', 'residuals', 'qq', 'smooth', 'all'}, default='diagnostics'
Type of plot:
- 'diagnostics': Standard 4-panel diagnostic plot
- 'residuals': Residuals vs fitted values
- 'qq': Q-Q plot of residuals
- 'smooth': Smooth function plots (GAM/GAMM only)
- 'all': All available plots
ax : matplotlib.axes.Axes, optional
Axes to plot on. If None, creates new figure.
**kwargs
Additional arguments passed to the underlying plot function.
Returns
-------
fig : matplotlib.figure.Figure or list of Figures
The created figure(s).
Examples
--------
>>> from aurora import fit_glm, plot, Gaussian
>>> result = fit_glm(X, y, family=Gaussian())
>>>
>>> # Standard diagnostics
>>> plot(result)
>>>
>>> # Just residuals
>>> plot(result, kind='residuals')
>>>
>>> # For GAM, plot smooth terms
>>> gam_result = fit_gam(X, y)
>>> plot(gam_result, kind='smooth')
See Also
--------
summary : Formatted model summary
compare : Compare multiple models
plot_diagnostics : Low-level diagnostic plots
plot_smooth : Low-level smooth term plots
"""
# Import visualization functions lazily
from .visualization import (
plot_all_smooths,
plot_diagnostics_panel,
)
if kind == "diagnostics":
# GLM-family results carry their own panel method; the GAMM-only
# plot_diagnostics_panel accepts (result, figsize) — no `ax`.
if hasattr(result, "plot_diagnostics"):
return result.plot_diagnostics(**kwargs)
return plot_diagnostics_panel(result, **kwargs)
elif kind == "residuals":
return _plot_residuals(result, ax=ax, **kwargs)
elif kind == "qq":
return _plot_qq(result, ax=ax, **kwargs)
elif kind == "smooth":
# Check if GAM result
if hasattr(result, "smooth_terms") or hasattr(result, "smooth_info"):
return plot_all_smooths(result, **kwargs)
else:
raise ValueError("kind='smooth' requires a GAM or GAMM result with smooth terms.")
elif kind == "all":
figs = []
if hasattr(result, "plot_diagnostics"):
figs.append(result.plot_diagnostics(**kwargs))
else:
figs.append(plot_diagnostics_panel(result, **kwargs))
if hasattr(result, "smooth_terms") or hasattr(result, "smooth_info"):
figs.append(plot_all_smooths(result, **kwargs))
return figs
else:
raise ValueError(
f"Unknown plot kind '{kind}'. Use 'diagnostics', 'residuals', 'qq', 'smooth', or 'all'."
)
def _plot_residuals(result: Any, ax: Any = None, **kwargs) -> Any:
"""Plot residuals vs fitted values.
Parameters
----------
result : ModelResult
A fitted model result with ``fitted_values`` and ``residuals``
attributes.
ax : matplotlib.axes.Axes, optional
Axes to draw on. If ``None``, a new figure is created.
**kwargs
Additional keyword arguments forwarded to ``ax.scatter``.
Returns
-------
fig : matplotlib.figure.Figure
The figure containing the plot.
"""
import matplotlib.pyplot as plt
if ax is None:
fig, ax = plt.subplots(figsize=(8, 6))
else:
fig = ax.get_figure()
fitted = result.fitted_values
residuals = result.residuals
ax.scatter(fitted, residuals, alpha=0.6, **kwargs)
ax.axhline(y=0, color="red", linestyle="--", linewidth=1)
ax.set_xlabel("Fitted values")
ax.set_ylabel("Residuals")
ax.set_title("Residuals vs Fitted")
return fig
def _plot_qq(result: Any, ax: Any = None, **kwargs) -> Any:
"""Create a normal Q-Q plot of standardized residuals.
Parameters
----------
result : ModelResult
A fitted model result with a ``residuals`` attribute.
ax : matplotlib.axes.Axes, optional
Axes to draw on. If ``None``, a new figure is created.
**kwargs
Additional keyword arguments (currently unused, reserved for
future extensions).
Returns
-------
fig : matplotlib.figure.Figure
The figure containing the Q-Q plot.
"""
import matplotlib.pyplot as plt
from scipy import stats
if ax is None:
fig, ax = plt.subplots(figsize=(6, 6))
else:
fig = ax.get_figure()
residuals = result.residuals
# Standardize residuals
std_resid = (residuals - np.mean(residuals)) / np.std(residuals)
stats.probplot(std_resid, dist="norm", plot=ax)
ax.set_title("Normal Q-Q Plot")
return fig
# =============================================================================
# Compare Function
# =============================================================================
[docs]
def compare(
*results: Any,
criterion: Literal["aic", "bic", "loglik", "all"] = "aic",
names: Sequence[str] | None = None,
print_output: bool = True,
) -> dict[str, Any]:
"""Compare multiple fitted models.
This function provides a unified interface for comparing different
model specifications using information criteria or likelihood ratio tests.
Parameters
----------
*results : ModelResult
Two or more fitted model results to compare.
criterion : {'aic', 'bic', 'loglik', 'all'}, default='aic'
Criterion for comparison:
- 'aic': Akaike Information Criterion
- 'bic': Bayesian Information Criterion
- 'loglik': Log-likelihood
- 'all': All available criteria
names : sequence of str, optional
Names for the models. If None, uses 'Model 1', 'Model 2', etc.
print_output : bool, default=True
If True, print the comparison table.
Returns
-------
comparison : dict
Dictionary with comparison results:
- 'summary': Formatted comparison table (str)
- 'criteria': DataFrame or dict with criterion values
- 'best': Index of best model
- 'delta': Differences from best model
Examples
--------
>>> from aurora import fit_glm, compare, Gaussian, Poisson
>>>
>>> # Fit different models
>>> m1 = fit_glm(X, y, family=Gaussian())
>>> m2 = fit_glm(X, y, family=Poisson())
>>> m3 = fit_glm(X[:, :2], y, family=Gaussian()) # Fewer predictors
>>>
>>> # Compare models
>>> compare(m1, m2, m3, criterion='aic')
======================================================================
Model Comparison (AIC)
======================================================================
Model AIC ΔAIC Rank
----------------------------------------------------------------------
Model 1 234.56 0.00 1 ← best
Model 2 245.12 10.56 2
Model 3 267.89 33.33 3
======================================================================
>>> # With custom names
>>> compare(m1, m2, names=['Linear', 'Poisson'], criterion='bic')
Notes
-----
- Lower AIC/BIC indicates better fit (with parsimony)
- For nested models, likelihood ratio tests can also be performed
- Models should be fit on the same data for meaningful comparison
See Also
--------
summary : Individual model summaries
"""
if len(results) < 2:
raise ValueError("compare() requires at least 2 models to compare.")
# Generate names if not provided
if names is None:
names = [f"Model {i + 1}" for i in range(len(results))]
elif len(names) != len(results):
raise ValueError(
f"Number of names ({len(names)}) must match number of models ({len(results)})."
)
# Extract criteria values
metrics = _extract_comparison_metrics(results, criterion)
# Create comparison summary
summary_str = _format_comparison_table(names, metrics, criterion)
# Find best model
primary_criterion = criterion if criterion != "all" else "aic"
if primary_criterion in metrics:
values = metrics[primary_criterion]
if primary_criterion == "loglik":
best_idx = int(np.argmax(values)) # Higher is better
else:
best_idx = int(np.argmin(values)) # Lower is better
delta = np.array(values) - values[best_idx]
else:
best_idx = 0
delta = np.zeros(len(results))
if print_output:
print(summary_str)
return {
"summary": summary_str,
"criteria": metrics,
"names": list(names),
"best": best_idx,
"best_model": names[best_idx],
"delta": delta,
}
def _extract_comparison_metrics(results: tuple[Any, ...], criterion: str) -> dict[str, list[float]]:
"""Extract comparison metrics from a sequence of model results.
Parameters
----------
results : tuple of ModelResult
Fitted model result objects.
criterion : str
One of ``'aic'``, ``'bic'``, ``'loglik'``, or ``'all'``.
Returns
-------
metrics : dict[str, list[float]]
Dictionary mapping metric names (``'n_obs'``, ``'n_params'``,
plus the requested criteria) to lists of values, one per model.
"""
metrics: dict[str, list[float]] = {}
# Always try to get n_obs and n_params
n_obs_list = []
n_params_list = []
for result in results:
# Number of observations
if hasattr(result, "n_obs_"):
n_obs_list.append(result.n_obs_)
elif hasattr(result, "n_observations"):
n_obs_list.append(result.n_observations)
else:
n_obs_list.append(np.nan)
# Number of parameters
if hasattr(result, "df_model"):
n_params_list.append(result.df_model + 1)
elif hasattr(result, "coefficients"):
n_params_list.append(len(result.coefficients))
else:
n_params_list.append(np.nan)
metrics["n_obs"] = n_obs_list
metrics["n_params"] = n_params_list
# Extract requested criteria
criteria_to_extract = ["aic", "bic", "loglik"] if criterion == "all" else [criterion]
for crit in criteria_to_extract:
values = []
for result in results:
if crit == "aic":
values.append(_get_aic(result))
elif crit == "bic":
values.append(_get_bic(result))
elif crit == "loglik":
values.append(_get_loglik(result))
else:
values.append(np.nan)
metrics[crit] = values
return metrics
def _get_aic(result: Any) -> float:
"""Extract or compute the Akaike Information Criterion (AIC).
If ``result.aic`` exists it is returned directly; otherwise AIC is
computed as ``-2 * loglik + 2 * k`` where ``k`` is the number of
parameters.
Parameters
----------
result : ModelResult
A fitted model result.
Returns
-------
aic : float
The AIC value, or ``np.nan`` if it cannot be computed.
"""
if hasattr(result, "aic"):
return float(result.aic)
# Try to compute from log-likelihood
ll = _get_loglik(result)
if np.isnan(ll):
return np.nan
if hasattr(result, "df_model"):
k = result.df_model + 1
elif hasattr(result, "coefficients"):
k = len(result.coefficients)
else:
return np.nan
return float(-2 * ll + 2 * k)
def _get_bic(result: Any) -> float:
"""Extract or compute the Bayesian Information Criterion (BIC).
If ``result.bic`` exists it is returned directly; otherwise BIC is
computed as ``-2 * loglik + k * log(n)`` where ``k`` is the number
of parameters and ``n`` is the number of observations.
Parameters
----------
result : ModelResult
A fitted model result.
Returns
-------
bic : float
The BIC value, or ``np.nan`` if it cannot be computed.
"""
if hasattr(result, "bic"):
return float(result.bic)
# Try to compute from log-likelihood
ll = _get_loglik(result)
if np.isnan(ll):
return np.nan
if hasattr(result, "n_obs_"):
n = result.n_obs_
elif hasattr(result, "n_observations"):
n = result.n_observations
else:
return np.nan
if hasattr(result, "df_model"):
k = result.df_model + 1
elif hasattr(result, "coefficients"):
k = len(result.coefficients)
else:
return np.nan
return float(-2 * ll + k * np.log(n))
def _get_loglik(result: Any) -> float:
"""Extract the log-likelihood from a model result.
Checks for ``log_likelihood_``, ``loglik``, or ``llf`` attributes
in order.
Parameters
----------
result : ModelResult
A fitted model result.
Returns
-------
loglik : float
The log-likelihood value, or ``np.nan`` if unavailable.
"""
if hasattr(result, "log_likelihood_"):
return float(result.log_likelihood_)
if hasattr(result, "loglik"):
return float(result.loglik)
if hasattr(result, "llf"):
return float(result.llf)
return np.nan
def _format_comparison_table(
names: Sequence[str], metrics: dict[str, list[float]], criterion: str
) -> str:
"""Format a model comparison table as a string.
Parameters
----------
names : sequence of str
Display names for each model.
metrics : dict[str, list[float]]
Metric values keyed by metric name, each value being a list
of floats with one entry per model.
criterion : str
The primary criterion used for ranking (``'aic'``,
``'bic'``, ``'loglik'``, or ``'all'``).
Returns
-------
table : str
A formatted multi-line comparison table.
"""
lines = []
sep = "=" * 70
lines.append(sep)
lines.append(f"{'Model Comparison':^70}")
lines.append(sep)
# Determine which criteria to show
criteria_to_show = ["aic", "bic", "loglik"] if criterion == "all" else [criterion]
# Header
header_parts = [f"{'Model':>15}"]
for crit in criteria_to_show:
header_parts.append(f"{crit.upper():>12}")
header_parts.append(f"{'Rank':>8}")
lines.append(" ".join(header_parts))
lines.append("-" * 70)
# Determine ranking (by first criterion)
primary = criteria_to_show[0]
if primary in metrics:
values = np.array(metrics[primary])
if primary == "loglik":
order = np.argsort(-values) # Higher is better
else:
order = np.argsort(values) # Lower is better
ranks = np.empty_like(order)
ranks[order] = np.arange(1, len(order) + 1)
else:
ranks = np.arange(1, len(names) + 1)
# Data rows
for i, name in enumerate(names):
row_parts = [f"{name:>15}"]
for crit in criteria_to_show:
if crit in metrics:
val = metrics[crit][i]
if np.isnan(val):
row_parts.append(f"{'N/A':>12}")
else:
row_parts.append(f"{val:>12.2f}")
else:
row_parts.append(f"{'N/A':>12}")
rank_str = str(int(ranks[i]))
if ranks[i] == 1:
rank_str += " ← best"
row_parts.append(f"{rank_str:>8}")
lines.append(" ".join(row_parts))
lines.append(sep)
return "\n".join(lines)
__all__ = [
"summary",
"plot",
"compare",
]