# SPDX-License-Identifier: MIT
# Copyright (c) 2025 Lucy Eduardo Arias
"""Visualization tools for GAMM models.
This module provides plotting functions for diagnosing and exploring
Generalized Additive Mixed Models, with focus on random effects visualization.
Functions include:
- Caterpillar plots: Visualize random effects with confidence intervals
- Q-Q plots: Check normality of random effects
- Density plots: Show distributions of random effects
- Diagnostic plots: Residual analysis and model checking
References
----------
.. [1] Pinheiro & Bates (2000). Mixed-Effects Models in S and S-PLUS.
.. [2] Gelman & Hill (2007). Data Analysis Using Regression and Multilevel/
Hierarchical Models. Chapter 12: Multilevel linear models.
.. [3] Wood (2017). Generalized Additive Models: An Introduction with R, 2nd ed.
"""
from __future__ import annotations
from typing import TYPE_CHECKING, Literal
import matplotlib.pyplot as plt
import numpy as np
import scipy.stats as stats
from matplotlib.figure import Figure
if TYPE_CHECKING:
from matplotlib.axes import Axes
from aurora.models.gamm.fitting import GAMMResult
[docs]
def plot_caterpillar(
result: GAMMResult,
grouping: str | int | None = None,
effect_index: int = 0,
confidence: float = 0.95,
sort: bool = True,
figsize: tuple[float, float] | None = None,
ax: Axes | None = None,
) -> tuple[Figure, Axes]:
"""Create caterpillar plot for random effects.
Caterpillar plots show the estimated random effects (BLUPs) for each
group with confidence intervals, sorted by magnitude. They are useful
for identifying groups with unusually large or small effects.
Parameters
----------
result : GAMMResult
Fitted GAMM model result.
grouping : str or int, optional
Which grouping variable to plot. If None, uses the first random
effect term. Required if model has multiple grouping variables.
effect_index : int, default=0
Which random effect to plot (0=intercept, 1=first slope, etc.).
confidence : float, default=0.95
Confidence level for intervals (e.g., 0.95 for 95% CI).
sort : bool, default=True
Whether to sort groups by effect magnitude.
figsize : tuple, optional
Figure size (width, height) in inches.
ax : Axes, optional
Matplotlib axes to plot on. If None, creates new figure.
Returns
-------
fig : Figure
Matplotlib figure.
ax : Axes
Matplotlib axes with plot.
Examples
--------
>>> import numpy as np
>>> from aurora.models.gamm import fit_gamm
>>> from aurora.models.gamm.plotting import plot_caterpillar
>>>
>>> # Fit model
>>> data = pd.DataFrame({
... 'y': np.random.randn(100),
... 'x': np.random.randn(100),
... 'subject': np.repeat(np.arange(10), 10)
... })
>>> result = fit_gamm(formula="y ~ x + (1 | subject)", data=data)
>>>
>>> # Plot random intercepts
>>> fig, ax = plot_caterpillar(result)
>>> plt.show()
Notes
-----
The confidence intervals are approximate and computed using the
posterior variance of the random effects:
Var(b | y) ≈ (Z'WZ + Ψ⁻¹)⁻¹
Groups whose confidence intervals do not include zero have effects
that are significantly different from the population mean.
"""
# Get random effects
if result._Z_info is None or len(result._Z_info) == 0:
raise ValueError("Model does not contain random effects")
z_info_all: list[dict] = result._Z_info
# Determine which grouping to plot
if grouping is None:
if len(z_info_all) > 1:
raise ValueError(
f"Model has {len(z_info_all)} grouping variables. "
"Please specify which one to plot using the 'grouping' parameter."
)
grouping_to_plot = z_info_all[0]["grouping"]
info = z_info_all[0]
else:
# Find the Z_info entry for this grouping
info = None
for z_info in z_info_all:
if z_info["grouping"] == grouping:
info = z_info
grouping_to_plot = grouping
break
if info is None:
available = [z["grouping"] for z in z_info_all]
raise ValueError(f"Grouping '{grouping}' not found in model. Available: {available}")
assert info is not None
# Check effect_index is valid
if effect_index >= info["n_effects"]:
raise ValueError(
f"effect_index={effect_index} but grouping '{grouping_to_plot}' "
f"only has {info['n_effects']} effects (0-indexed)"
)
# Extract random effects for this grouping
# GAMMResult already has random_effects in the correct format
random_effects = result.random_effects[grouping_to_plot]
# Get groups and effects
groups = info["groups"]
n_groups = len(groups)
effects = np.array([random_effects[g][effect_index] for g in groups])
# Compute approximate standard errors
# Use posterior variance: Var(b|y) ≈ (Z'WZ + Ψ⁻¹)⁻¹
# For simplicity, use empirical SE scaled by estimated variance
var_component = result.variance_components[z_info_all.index(info)]
if var_component.ndim == 2:
# Covariance matrix
effect_var = var_component[effect_index, effect_index]
else:
# Scalar variance
effect_var = var_component
# Approximate SE (shrinkage-adjusted)
# In practice, would use full posterior variance
se = np.sqrt(effect_var / np.sqrt(n_groups))
se_array = np.full(n_groups, se)
# Compute confidence intervals
z_crit = stats.norm.ppf((1 + confidence) / 2)
ci_lower = effects - z_crit * se_array
ci_upper = effects + z_crit * se_array
# Sort if requested
if sort:
sort_idx = np.argsort(effects)
effects = effects[sort_idx]
ci_lower = ci_lower[sort_idx]
ci_upper = ci_upper[sort_idx]
groups = groups[sort_idx]
# Create plot
if ax is None:
if figsize is None:
figsize = (8, max(4, n_groups * 0.2))
fig, ax = plt.subplots(figsize=figsize)
else:
fig = ax.figure # type: ignore[assignment] # Figure | SubFigure
# Plot horizontal lines for CIs
y_positions = np.arange(n_groups)
ax.hlines(y_positions, ci_lower, ci_upper, color="gray", linewidth=1.5)
# Plot points for estimates
ax.plot(effects, y_positions, "o", color="steelblue", markersize=6)
# Add reference line at zero
ax.axvline(0, color="red", linestyle="--", linewidth=1, alpha=0.7)
# Labels
effect_names = ["Intercept"] + [f"Slope {i}" for i in range(1, info["n_effects"])]
ax.set_xlabel(f"Random Effect: {effect_names[effect_index]}")
ax.set_ylabel(f"Group ({grouping_to_plot})")
ax.set_yticks(y_positions)
ax.set_yticklabels([str(g) for g in groups])
ax.set_title(
f"Caterpillar Plot: {effect_names[effect_index]} by {grouping_to_plot}\n"
f"{confidence * 100:.0f}% Confidence Intervals"
)
ax.grid(axis="x", alpha=0.3)
fig.tight_layout()
return fig, ax
[docs]
def plot_random_effects_qq(
result: GAMMResult,
grouping: str | int | None = None,
effect_index: int = 0,
figsize: tuple[float, float] = (6, 6),
ax: Axes | None = None,
) -> tuple[Figure, Axes]:
"""Create Q-Q plot for random effects normality check.
Q-Q (quantile-quantile) plots compare the distribution of random effects
to a theoretical normal distribution. Points falling along the diagonal
line indicate normality.
Parameters
----------
result : GAMMResult
Fitted GAMM model result.
grouping : str or int, optional
Which grouping variable to plot. If None, uses the first random
effect term.
effect_index : int, default=0
Which random effect to plot (0=intercept, 1=first slope, etc.).
figsize : tuple, default=(6, 6)
Figure size (width, height) in inches.
ax : Axes, optional
Matplotlib axes to plot on. If None, creates new figure.
Returns
-------
fig : Figure
Matplotlib figure.
ax : Axes
Matplotlib axes with plot.
Examples
--------
>>> fig, ax = plot_random_effects_qq(result)
>>> plt.show()
Notes
-----
The random effects are assumed to follow a normal distribution:
b ~ N(0, Ψ)
Departures from normality (heavy tails, skewness) can be identified
visually from Q-Q plots. Consider:
- S-shaped pattern: Heavy tails (outliers)
- Systematic curve: Skewness
- Points far from line: Individual outliers
"""
# Get random effects
if result._Z_info is None or len(result._Z_info) == 0:
raise ValueError("Model does not contain random effects")
z_info_all: list[dict] = result._Z_info
# Determine which grouping to plot
if grouping is None:
if len(z_info_all) > 1:
raise ValueError(
"Model has multiple grouping variables. "
"Please specify which one to plot using the 'grouping' parameter."
)
grouping_to_plot = z_info_all[0]["grouping"]
info = z_info_all[0]
else:
info = None
for z_info in z_info_all:
if z_info["grouping"] == grouping:
info = z_info
grouping_to_plot = grouping
break
if info is None:
available = [z["grouping"] for z in z_info_all]
raise ValueError(f"Grouping '{grouping}' not found. Available: {available}")
assert info is not None
# Check effect_index
if effect_index >= info["n_effects"]:
raise ValueError(
f"effect_index={effect_index} but grouping '{grouping_to_plot}' "
f"only has {info['n_effects']} effects"
)
# Extract random effects
# GAMMResult already has random_effects in the correct format
random_effects = result.random_effects[grouping_to_plot]
# Get effects for this index
groups = info["groups"]
effects = np.array([random_effects[g][effect_index] for g in groups])
# Standardize effects
effects_std = (effects - effects.mean()) / effects.std()
# Create Q-Q plot
if ax is None:
fig, ax = plt.subplots(figsize=figsize)
else:
fig = ax.figure # type: ignore[assignment] # Figure | SubFigure
# Compute theoretical and sample quantiles
stats.probplot(effects_std, dist="norm", plot=ax)
# Styling
effect_names = ["Intercept"] + [f"Slope {i}" for i in range(1, info["n_effects"])]
ax.set_title(
f"Q-Q Plot: {effect_names[effect_index]} by {grouping_to_plot}\n"
f"(Checking Normality Assumption)"
)
ax.set_xlabel("Theoretical Quantiles")
ax.set_ylabel("Standardized Random Effects")
ax.grid(alpha=0.3)
fig.tight_layout()
return fig, ax
[docs]
def plot_random_effects_density(
result: GAMMResult,
grouping: str | int | None = None,
effect_index: int = 0,
show_normal: bool = True,
figsize: tuple[float, float] = (8, 5),
ax: Axes | None = None,
) -> tuple[Figure, Axes]:
"""Plot density of random effects with optional normal overlay.
Parameters
----------
result : GAMMResult
Fitted GAMM model result.
grouping : str or int, optional
Which grouping variable to plot.
effect_index : int, default=0
Which random effect to plot.
show_normal : bool, default=True
Whether to overlay theoretical normal distribution.
figsize : tuple, default=(8, 5)
Figure size.
ax : Axes, optional
Matplotlib axes to plot on.
Returns
-------
fig : Figure
Matplotlib figure.
ax : Axes
Matplotlib axes with plot.
Examples
--------
>>> fig, ax = plot_random_effects_density(result)
>>> plt.show()
Notes
-----
The histogram shows the empirical distribution of random effects,
while the overlaid curve shows the theoretical normal distribution
N(0, Ψ) implied by the model.
"""
# Get random effects
if result._Z_info is None or len(result._Z_info) == 0:
raise ValueError("Model does not contain random effects")
z_info_all: list[dict] = result._Z_info
# Determine grouping
if grouping is None:
if len(z_info_all) > 1:
raise ValueError("Model has multiple grouping variables. Please specify which one.")
grouping_to_plot = z_info_all[0]["grouping"]
info = z_info_all[0]
else:
info = None
for z_info in z_info_all:
if z_info["grouping"] == grouping:
info = z_info
grouping_to_plot = grouping
break
if info is None:
available = [z["grouping"] for z in z_info_all]
raise ValueError(f"Grouping '{grouping}' not found. Available: {available}")
assert info is not None
# Check effect_index
if effect_index >= info["n_effects"]:
raise ValueError(f"effect_index={effect_index} invalid for {info['n_effects']} effects")
# Extract random effects
# GAMMResult already has random_effects in the correct format
random_effects = result.random_effects[grouping_to_plot]
groups = info["groups"]
effects = np.array([random_effects[g][effect_index] for g in groups])
# Create plot
if ax is None:
fig, ax = plt.subplots(figsize=figsize)
else:
fig = ax.figure # type: ignore[assignment] # Figure | SubFigure
# Histogram
ax.hist(
effects,
bins=min(30, len(effects) // 2),
density=True,
alpha=0.6,
color="steelblue",
edgecolor="black",
label="Observed",
)
# Overlay normal if requested
if show_normal:
var_component = result.variance_components[z_info_all.index(info)]
if var_component.ndim == 2:
effect_var = var_component[effect_index, effect_index]
else:
effect_var = var_component
# Theoretical normal
x_range = np.linspace(
effects.min() - 0.5 * effects.std(),
effects.max() + 0.5 * effects.std(),
200,
)
y_normal = stats.norm.pdf(x_range, loc=0, scale=np.sqrt(effect_var))
ax.plot(x_range, y_normal, "r-", linewidth=2, label=f"N(0, {effect_var:.3f})")
# Labels
effect_names = ["Intercept"] + [f"Slope {i}" for i in range(1, info["n_effects"])]
ax.set_xlabel(f"Random Effect: {effect_names[effect_index]}")
ax.set_ylabel("Density")
ax.set_title(
f"Distribution of {effect_names[effect_index]} by {grouping_to_plot}\n"
f"({len(effects)} groups)"
)
ax.legend()
ax.grid(alpha=0.3)
fig.tight_layout()
return fig, ax
[docs]
def plot_diagnostics(
result: GAMMResult,
plot_type: Literal["residuals", "fitted", "qq", "scale-location"] = "residuals",
figsize: tuple[float, float] = (8, 6),
ax: Axes | None = None,
) -> tuple[Figure, Axes]:
"""Create diagnostic plots for GAMM models.
Parameters
----------
result : GAMMResult
Fitted GAMM model result.
plot_type : {'residuals', 'fitted', 'qq', 'scale-location'}
Type of diagnostic plot:
- 'residuals': Residuals vs fitted values
- 'fitted': Fitted vs observed values
- 'qq': Q-Q plot of residuals
- 'scale-location': Scale-location plot (sqrt(|residuals|) vs fitted)
figsize : tuple, default=(8, 6)
Figure size.
ax : Axes, optional
Matplotlib axes to plot on.
Returns
-------
fig : Figure
Matplotlib figure.
ax : Axes
Matplotlib axes with plot.
Examples
--------
>>> # Residuals vs fitted
>>> fig, ax = plot_diagnostics(result, plot_type='residuals')
>>> plt.show()
>>>
>>> # Q-Q plot of residuals
>>> fig, ax = plot_diagnostics(result, plot_type='qq')
>>> plt.show()
Notes
-----
**Residuals vs Fitted:**
- Should show no pattern (randomness around zero)
- Fan shape indicates heteroscedasticity
- Curvature indicates nonlinearity
**Q-Q Plot:**
- Points should fall on diagonal line
- Departures indicate non-normality of residuals
**Scale-Location:**
- Check homoscedasticity (constant variance)
- Should show horizontal line with random scatter
**Fitted vs Observed:**
- Points should fall near y=x line
- Shows overall model fit quality
"""
if ax is None:
fig, ax = plt.subplots(figsize=figsize)
else:
fig = ax.figure # type: ignore[assignment] # Figure | SubFigure
# Compute residuals
residuals = result.residuals
fitted = result.fitted_values
if plot_type == "residuals":
# Residuals vs fitted
ax.scatter(fitted, residuals, alpha=0.5, edgecolor="black", linewidth=0.5)
ax.axhline(0, color="red", linestyle="--", linewidth=1)
# Add smooth line
if len(fitted) > 10:
from scipy.ndimage import gaussian_filter1d
sort_idx = np.argsort(fitted)
fitted_sorted = fitted[sort_idx]
residuals_sorted = residuals[sort_idx]
smooth = gaussian_filter1d(residuals_sorted, sigma=max(1, len(fitted) // 20))
ax.plot(fitted_sorted, smooth, "b-", linewidth=2, alpha=0.7)
ax.set_xlabel("Fitted Values")
ax.set_ylabel("Residuals")
ax.set_title("Residuals vs Fitted\n(Should show random scatter around zero)")
ax.grid(alpha=0.3)
elif plot_type == "fitted":
# Fitted vs observed
observed = fitted + residuals
ax.scatter(fitted, observed, alpha=0.5, edgecolor="black", linewidth=0.5)
# Add y=x line
lims = [
min(fitted.min(), observed.min()),
max(fitted.max(), observed.max()),
]
ax.plot(lims, lims, "r--", linewidth=1, label="y=x")
ax.set_xlabel("Fitted Values")
ax.set_ylabel("Observed Values")
ax.set_title("Fitted vs Observed\n(Points should cluster near y=x line)")
ax.legend()
ax.grid(alpha=0.3)
elif plot_type == "qq":
# Q-Q plot of residuals
stats.probplot(residuals, dist="norm", plot=ax)
ax.set_title("Q-Q Plot of Residuals\n(Checking normality assumption)")
ax.grid(alpha=0.3)
elif plot_type == "scale-location":
# Scale-location plot
sqrt_std_resid = np.sqrt(np.abs(residuals / residuals.std()))
ax.scatter(fitted, sqrt_std_resid, alpha=0.5, edgecolor="black", linewidth=0.5)
# Add smooth line
if len(fitted) > 10:
from scipy.ndimage import gaussian_filter1d
sort_idx = np.argsort(fitted)
fitted_sorted = fitted[sort_idx]
sqrt_std_resid_sorted = sqrt_std_resid[sort_idx]
smooth = gaussian_filter1d(sqrt_std_resid_sorted, sigma=max(1, len(fitted) // 20))
ax.plot(fitted_sorted, smooth, "r-", linewidth=2, alpha=0.7)
ax.set_xlabel("Fitted Values")
ax.set_ylabel("√|Standardized Residuals|")
ax.set_title("Scale-Location Plot\n(Check homoscedasticity)")
ax.grid(alpha=0.3)
else:
raise ValueError(
f"Unknown plot_type '{plot_type}'. "
"Must be one of: 'residuals', 'fitted', 'qq', 'scale-location'"
)
fig.tight_layout()
return fig, ax
[docs]
def plot_random_effects_summary(
result: GAMMResult,
grouping: str | int | None = None,
effect_index: int = 0,
figsize: tuple[float, float] = (14, 10),
) -> Figure:
"""Create comprehensive summary of random effects diagnostics.
Produces a 2x2 grid with:
1. Caterpillar plot
2. Q-Q plot
3. Density plot
4. Residuals vs fitted
Parameters
----------
result : GAMMResult
Fitted GAMM model result.
grouping : str or int, optional
Which grouping variable to plot.
effect_index : int, default=0
Which random effect to plot.
figsize : tuple, default=(14, 10)
Figure size.
Returns
-------
fig : Figure
Matplotlib figure with 2x2 subplot grid.
Examples
--------
>>> fig = plot_random_effects_summary(result)
>>> plt.show()
Notes
-----
This is a convenience function that creates a comprehensive
diagnostic summary in a single figure.
"""
fig, axes = plt.subplots(2, 2, figsize=figsize)
axes = axes.flatten()
# 1. Caterpillar plot
plot_caterpillar(result, grouping=grouping, effect_index=effect_index, ax=axes[0])
# 2. Q-Q plot
plot_random_effects_qq(result, grouping=grouping, effect_index=effect_index, ax=axes[1])
# 3. Density plot
plot_random_effects_density(result, grouping=grouping, effect_index=effect_index, ax=axes[2])
# 4. Residuals vs fitted
plot_diagnostics(result, plot_type="residuals", ax=axes[3])
fig.suptitle("Random Effects Diagnostic Summary", fontsize=14, y=0.995)
fig.tight_layout()
return fig
[docs]
def plot_diagnostics_panel(
result: GAMMResult,
figsize: tuple[float, float] = (12, 10),
) -> Figure:
"""Create 2x2 diagnostic panel for GAMM model.
Standard diagnostic panel inspired by R's plot.lm():
1. Residuals vs Fitted (top-left)
2. Q-Q Plot of residuals (top-right)
3. Scale-Location plot (bottom-left)
4. Residuals vs Leverage with Cook's distance (bottom-right)
Parameters
----------
result : GAMMResult
Fitted GAMM model result.
figsize : tuple, default=(12, 10)
Figure size (width, height) in inches.
Returns
-------
fig : Figure
Matplotlib figure with 2x2 diagnostic panel.
Examples
--------
>>> from aurora.models.gamm import fit_gamm
>>> result = fit_gamm("y ~ x + (1|group)", data=df)
>>> fig = plot_diagnostics_panel(result)
>>> plt.show()
Notes
-----
This produces a publication-quality diagnostic panel following
the classic R layout:
**Residuals vs Fitted (top-left):**
- Should show random scatter around zero
- Patterns indicate model misspecification
**Normal Q-Q (top-right):**
- Points should fall on diagonal
- Departures indicate non-normal residuals
**Scale-Location (bottom-left):**
- Should show horizontal trend
- Increasing spread indicates heteroscedasticity
**Residuals vs Leverage (bottom-right):**
- Identifies influential observations
- Cook's distance contours show influence
References
----------
.. [1] Belsley, D. A., Kuh, E., & Welsch, R. E. (1980).
Regression Diagnostics. Wiley.
.. [2] Cook, R. D., & Weisberg, S. (1982).
Residuals and Influence in Regression. Chapman and Hall.
"""
fig, axes = plt.subplots(2, 2, figsize=figsize)
# Extract data
residuals = result.residuals
fitted = result.fitted_values
n = len(residuals)
# Standardized residuals
resid_std = residuals / np.std(residuals)
# 1. Residuals vs Fitted (top-left)
ax = axes[0, 0]
ax.scatter(fitted, residuals, alpha=0.6, edgecolor="black", linewidth=0.3)
ax.axhline(0, color="red", linestyle="--", linewidth=1)
# Add LOWESS smooth
if n > 10:
try:
from scipy.ndimage import gaussian_filter1d
sort_idx = np.argsort(fitted)
smooth = gaussian_filter1d(residuals[sort_idx], sigma=max(1, n // 15))
ax.plot(fitted[sort_idx], smooth, "b-", linewidth=2, alpha=0.7)
except ImportError:
pass
ax.set_xlabel("Fitted values")
ax.set_ylabel("Residuals")
ax.set_title("Residuals vs Fitted")
ax.grid(alpha=0.3)
# 2. Q-Q Plot (top-right)
ax = axes[0, 1]
stats.probplot(resid_std, dist="norm", plot=ax)
ax.get_lines()[0].set_markerfacecolor("steelblue")
ax.get_lines()[0].set_markeredgecolor("black")
ax.get_lines()[0].set_markersize(5)
ax.get_lines()[1].set_color("red")
ax.set_title("Normal Q-Q")
ax.set_xlabel("Theoretical Quantiles")
ax.set_ylabel("Standardized Residuals")
ax.grid(alpha=0.3)
# 3. Scale-Location (bottom-left)
ax = axes[1, 0]
sqrt_abs_resid = np.sqrt(np.abs(resid_std))
ax.scatter(fitted, sqrt_abs_resid, alpha=0.6, edgecolor="black", linewidth=0.3)
if n > 10:
try:
from scipy.ndimage import gaussian_filter1d
sort_idx = np.argsort(fitted)
smooth = gaussian_filter1d(sqrt_abs_resid[sort_idx], sigma=max(1, n // 15))
ax.plot(fitted[sort_idx], smooth, "r-", linewidth=2, alpha=0.7)
except ImportError:
pass
ax.set_xlabel("Fitted values")
ax.set_ylabel("√|Standardized residuals|")
ax.set_title("Scale-Location")
ax.grid(alpha=0.3)
# 4. Residuals vs Leverage with Cook's distance (bottom-right)
ax = axes[1, 1]
# Compute leverage (hat values) if available
# For now, use simplified version based on fitted values
# Approximate leverage: h_ii ≈ 1/n + (x_i - x̄)²/SSx
fitted_centered = fitted - np.mean(fitted)
leverage = 1 / n + fitted_centered**2 / np.sum(fitted_centered**2)
leverage = np.clip(leverage, 0.01, 0.99) # Ensure valid range
ax.scatter(leverage, resid_std, alpha=0.6, edgecolor="black", linewidth=0.3)
ax.axhline(0, color="gray", linestyle="--", linewidth=0.5)
# Add Cook's distance contours
# Cook's D ≈ r²_i × h_i / (p × (1-h_i)²)
# For contours at D = 0.5, 1.0
h_range = np.linspace(0.01, 0.5, 100)
p = 2 # Approximate number of parameters
for cook_d in [0.5, 1.0]:
# r² = cook_d × p × (1-h)² / h
r_pos = np.sqrt(cook_d * p * (1 - h_range) ** 2 / h_range)
r_neg = -r_pos
ax.plot(h_range, r_pos, "r--", alpha=0.5, linewidth=0.8)
ax.plot(h_range, r_neg, "r--", alpha=0.5, linewidth=0.8)
# Label
ax.text(h_range[-1], r_pos[-1], f"D={cook_d}", fontsize=8, color="red", alpha=0.7)
ax.set_xlabel("Leverage")
ax.set_ylabel("Standardized residuals")
ax.set_title("Residuals vs Leverage")
ax.grid(alpha=0.3)
# Add overall title
fig.suptitle("GAMM Diagnostic Plots", fontsize=14, y=0.995)
fig.tight_layout()
return fig
[docs]
def plot_smooth_effect(
result: GAMMResult,
term_name: str,
data: dict | None = None,
n_points: int = 100,
level: float = 0.95,
show_data: bool = True,
show_residuals: bool = False,
figsize: tuple[float, float] = (8, 6),
ax: Axes | None = None,
) -> tuple[Figure, Axes]:
"""Plot estimated smooth function with confidence band.
Creates a plot of the estimated smooth term f(x) with:
- Smooth curve (estimated effect)
- Pointwise confidence interval
- Rug plot showing data density
- Optional partial residuals
Parameters
----------
result : GAMMResult
Fitted GAMM/GAM model result.
term_name : str
Name of the smooth term to plot (e.g., 's(age)' or 'age').
data : dict or DataFrame, optional
Original data used for fitting. If None, attempts to extract from result.
n_points : int, default=100
Number of points for smooth curve evaluation.
level : float, default=0.95
Confidence level for intervals.
show_data : bool, default=True
Whether to show rug plot of data density.
show_residuals : bool, default=False
Whether to show partial residuals.
figsize : tuple, default=(8, 6)
Figure size (width, height) in inches.
ax : Axes, optional
Matplotlib axes to plot on. If None, creates new figure.
Returns
-------
fig : Figure
Matplotlib figure.
ax : Axes
Matplotlib axes with plot.
Examples
--------
>>> from aurora.models.gamm import fit_gamm
>>> result = fit_gamm("y ~ s(age) + (1|group)", data=df)
>>> fig, ax = plot_smooth_effect(result, term_name='age', data=df)
>>> plt.show()
Notes
-----
The confidence band is a pointwise interval, not a simultaneous band.
For simultaneous confidence bands, additional adjustments are needed.
Partial residuals are computed as:
e_i = f̂(x_i) + (y_i - ŷ_i)
This allows visualizing the data on the scale of the smooth effect.
References
----------
.. [1] Wood, S. N. (2017). GAMs: An Introduction with R, 2nd ed. Chapter 4.
.. [2] Hastie, T., & Tibshirani, R. (1990). GAMs. Chapman and Hall.
"""
# Try to extract smooth term info
if not hasattr(result, "smooth_terms") or result.smooth_terms is None:
raise ValueError("Model does not contain smooth terms")
# Find the smooth term
smooth_info = None
for term in result.smooth_terms:
if term.get("name") == term_name or term.get("variable") == term_name:
smooth_info = term
break
# Also check for 's(name)' format
if f"s({term_name})" == term.get("name"):
smooth_info = term
break
if smooth_info is None:
available = [t.get("name", t.get("variable", "?")) for t in result.smooth_terms]
raise ValueError(f"Smooth term '{term_name}' not found. Available: {available}")
# Get x values for the smooth
var_name = smooth_info.get("variable", term_name)
if data is not None:
if hasattr(data, "values"): # DataFrame
x_data = data[var_name].values
else: # dict
x_data = np.asarray(data[var_name])
elif hasattr(result, "_data") and result._data is not None:
x_data = result._data[var_name]
else:
# Use range from smooth info if available
x_data = np.linspace(smooth_info.get("x_min", 0), smooth_info.get("x_max", 1), n_points)
# Create evaluation grid
x_grid = np.linspace(x_data.min(), x_data.max(), n_points)
# Get smooth coefficients and basis
coef_start = smooth_info.get("coef_start", 0)
coef_end = smooth_info.get("coef_end", coef_start + smooth_info.get("n_basis", 10))
smooth_coefs = result.coefficients[coef_start:coef_end]
# Build basis matrix for grid points
smooth_info.get("basis", "cr") # Default to cubic regression splines
n_basis = smooth_info.get("n_basis", 10)
smooth_info.get("knots", None)
# Simple B-spline basis construction
try:
from aurora.smoothing.splines.bspline import BSplineBasis
degree = 3
# BSplineBasis: n_basis_ = len(knots) - degree - 1, so the knot
# vector needs n_basis + degree + 1 entries for n_basis columns.
knots = np.linspace(x_grid.min(), x_grid.max(), n_basis + degree + 1)
basis = BSplineBasis(knots=knots, degree=degree)
X_grid = basis.basis_matrix(x_grid)
except Exception:
# Fallback: polynomial basis
X_grid = np.column_stack([x_grid**i for i in range(n_basis)])
# Compute smooth values
smooth_values = X_grid @ smooth_coefs
# Compute confidence intervals
# Approximate variance: Var(f(x)) = B(x) @ Cov(β) @ B(x).T
if hasattr(result, "covariance") and result.covariance is not None:
cov_smooth = result.covariance[coef_start:coef_end, coef_start:coef_end]
var_smooth = np.diag(X_grid @ cov_smooth @ X_grid.T)
se_smooth = np.sqrt(np.maximum(var_smooth, 0))
else:
# Approximate SE from residual variance
residual_var = np.var(result.residuals)
se_smooth = np.sqrt(residual_var / n_points) * np.ones(n_points)
z_crit = stats.norm.ppf((1 + level) / 2)
ci_lower = smooth_values - z_crit * se_smooth
ci_upper = smooth_values + z_crit * se_smooth
# Create plot
if ax is None:
fig, ax = plt.subplots(figsize=figsize)
else:
fig = ax.figure # type: ignore[assignment] # Figure | SubFigure
# Confidence band
ax.fill_between(
x_grid,
ci_lower,
ci_upper,
alpha=0.3,
color="steelblue",
label=f"{level * 100:.0f}% CI",
)
# Smooth curve
ax.plot(x_grid, smooth_values, "b-", linewidth=2, label="Smooth effect")
# Partial residuals
if show_residuals and data is not None:
# Compute partial residuals at data points
try:
basis_data = BSplineBasis(knots=knots, degree=3)
X_data = basis_data.basis_matrix(x_data)
except Exception:
X_data = np.column_stack([x_data**i for i in range(n_basis)])
smooth_at_data = X_data @ smooth_coefs
partial_resid = smooth_at_data + result.residuals
ax.scatter(x_data, partial_resid, alpha=0.3, s=20, c="gray", label="Partial residuals")
# Rug plot
if show_data:
ax.plot(
x_data,
np.full_like(x_data, ax.get_ylim()[0]),
"|",
color="black",
alpha=0.3,
markersize=10,
)
# Labels
ax.set_xlabel(var_name)
ax.set_ylabel(f"s({var_name})")
ax.set_title(f"Smooth Effect: s({var_name})")
ax.legend(loc="best")
ax.grid(alpha=0.3)
fig.tight_layout()
return fig, ax
[docs]
def plot_all_smooth_effects(
result: GAMMResult,
data: dict | None = None,
n_cols: int = 2,
figsize_per_plot: tuple[float, float] = (5, 4),
**kwargs,
) -> Figure:
"""Plot all smooth effects in a grid layout.
Parameters
----------
result : GAMMResult
Fitted GAMM/GAM model result.
data : dict or DataFrame, optional
Original data used for fitting.
n_cols : int, default=2
Number of columns in the grid.
figsize_per_plot : tuple, default=(5, 4)
Size of each subplot.
**kwargs : dict
Additional arguments passed to plot_smooth_effect.
Returns
-------
fig : Figure
Matplotlib figure with grid of smooth effect plots.
Examples
--------
>>> result = fit_gamm("y ~ s(x1) + s(x2) + s(x3) + (1|group)", data=df)
>>> fig = plot_all_smooth_effects(result, data=df)
>>> plt.show()
"""
if not hasattr(result, "smooth_terms") or result.smooth_terms is None:
raise ValueError("Model does not contain smooth terms")
n_smooths = len(result.smooth_terms)
if n_smooths == 0:
raise ValueError("Model has no smooth terms to plot")
n_rows = int(np.ceil(n_smooths / n_cols))
figsize = (figsize_per_plot[0] * n_cols, figsize_per_plot[1] * n_rows)
fig, axes = plt.subplots(n_rows, n_cols, figsize=figsize)
if n_smooths == 1:
axes = np.array([axes])
axes = axes.flatten()
for i, term in enumerate(result.smooth_terms):
term_name = term.get("variable", term.get("name", f"term_{i}"))
try:
plot_smooth_effect(result, term_name, data=data, ax=axes[i], **kwargs)
except Exception as e:
axes[i].text(
0.5,
0.5,
f"Error: {str(e)[:30]}...",
ha="center",
va="center",
transform=axes[i].transAxes,
)
axes[i].set_title(f"s({term_name}) - Error")
# Hide unused subplots
for j in range(n_smooths, len(axes)):
axes[j].set_visible(False)
fig.suptitle("Smooth Effect Plots", fontsize=14, y=1.02)
fig.tight_layout()
return fig
__all__ = [
"plot_caterpillar",
"plot_random_effects_qq",
"plot_random_effects_density",
"plot_diagnostics",
"plot_diagnostics_panel",
"plot_random_effects_summary",
"plot_smooth_effect",
"plot_all_smooth_effects",
]