# SPDX-License-Identifier: MIT
# Copyright (c) 2025 Lucy Eduardo Arias
"""Laplace Approximation for Generalized Linear Mixed Models.
This module implements Laplace approximation for GLMMs, which provides
more accurate inference than PQL but at higher computational cost.
Laplace Approximation
---------------------
The Laplace approximation approximates the marginal likelihood by
integrating out the random effects using a normal approximation:
L(β, θ) ≈ f(y | b̂, β) p(b̂ | θ) (2π)^{q/2} |H|^(-1/2)
Where:
- b̂ = argmax_b f(y | b, β) p(b | θ) is the conditional mode
- H = -∂²/∂b² log[f(y | b, β) p(b | θ)] at b = b̂ is the Hessian
The algorithm alternates between:
1. Finding b̂ given current β, θ (inner optimization)
2. Updating β, θ given b̂ (outer optimization)
References
----------
- Breslow, N.E. & Lin, X. (1995). "Bias correction in generalized linear
mixed models with a single component of dispersion".
- Wood, S.N. (2017). Generalized Additive Models: An Introduction with R (2nd ed.).
Chapter 6.10: Laplace approximation.
"""
from __future__ import annotations
from dataclasses import dataclass
from typing import TYPE_CHECKING
import numpy as np
from scipy import linalg, optimize
if TYPE_CHECKING:
from numpy.typing import NDArray
from aurora.distributions.base import Family
from aurora.distributions.families import BinomialFamily, GaussianFamily, PoissonFamily
[docs]
@dataclass
class LaplaceResult:
"""Result of Laplace approximation.
Attributes
----------
beta : ndarray, shape (p,)
Estimated fixed effect coefficients
b : ndarray, shape (q,)
Estimated random effect coefficients at conditional mode
psi : ndarray, shape (n_effects, n_effects)
Estimated variance-covariance matrix for random effects
sigma2 : float
Estimated residual variance (for Gaussian; 1.0 for others)
fitted_values : ndarray, shape (n,)
Fitted values on response scale (μ)
linear_predictor : ndarray, shape (n,)
Linear predictor (η = Xβ + Zb)
hessian : ndarray, shape (q, q)
Hessian of negative log joint density at b̂
converged : bool
Whether the algorithm converged
n_iter : int
Number of outer iterations
log_likelihood : float
Approximate marginal log-likelihood
"""
beta: NDArray[np.floating]
b: NDArray[np.floating]
psi: NDArray[np.floating]
sigma2: float
fitted_values: NDArray[np.floating]
linear_predictor: NDArray[np.floating]
hessian: NDArray[np.floating]
converged: bool
n_iter: int
log_likelihood: float
[docs]
def fit_laplace(
X: np.ndarray,
Z: np.ndarray,
y: np.ndarray,
family: Family | str,
S: np.ndarray | None = None,
lambda_: float = 0.0,
psi_init: np.ndarray | None = None,
maxiter: int = 50,
tol: float = 1e-6,
) -> LaplaceResult:
"""Fit a GLMM using Laplace approximation.
Parameters
----------
X : ndarray, shape (n, p)
Fixed effects design matrix
Z : ndarray, shape (n, q)
Random effects design matrix
y : ndarray, shape (n,)
Response variable
family : Family or str
Distribution family ('gaussian', 'poisson', 'binomial')
S : ndarray, shape (p, p), optional
Penalty matrix for smooth terms (default: zeros)
lambda_ : float, default=0.0
Smoothing parameter
psi_init : ndarray, shape (n_effects, n_effects), optional
Initial variance-covariance matrix (default: identity).
Required for models with more than one random effect per group
(e.g. random intercept + slope): its dimension determines how the
columns of Z are grouped. Without it, random-intercept structure
(one effect per group) is assumed.
maxiter : int, default=50
Maximum iterations
tol : float, default=1e-6
Convergence tolerance
Returns
-------
LaplaceResult
Fitted model results
Notes
-----
Laplace approximation is more accurate than PQL, especially for:
- Binary data with small group sizes
- Count data with small means
- Models with large random effects
However, it is computationally more expensive due to:
- Inner optimization to find conditional mode b̂
- Hessian computation at each iteration
Examples
--------
>>> # Poisson GLMM with Laplace approximation
>>> import numpy as np
>>> from aurora.models.gamm import fit_laplace
>>>
>>> n_groups, n_per_group = 10, 20
>>> n = n_groups * n_per_group
>>> groups = np.repeat(np.arange(n_groups), n_per_group)
>>> x = np.random.randn(n)
>>> X = np.column_stack([np.ones(n), x])
>>>
>>> # Random effects design
>>> Z = np.zeros((n, n_groups))
>>> for i in range(n):
... Z[i, groups[i]] = 1
>>>
>>> # Generate data
>>> b_true = np.random.randn(n_groups) * 0.5
>>> eta = 1.0 + 0.5 * x + b_true[groups]
>>> y = np.random.poisson(np.exp(eta))
>>>
>>> # Fit model
>>> result = fit_laplace(X, Z, y, family='poisson')
>>> print(f"Log-likelihood: {result.log_likelihood:.2f}")
"""
# Get family object
if isinstance(family, str):
family = _get_family(family)
n, p = X.shape
q = Z.shape[1]
# Initialize penalty matrix
if S is None:
S = np.zeros((p, p))
# Infer number of effects per group from the design / initialization.
# With psi_init=None we cannot distinguish "q groups × 1 effect" from
# "q/2 groups × 2 effects", so random-slope models (n_effects > 1)
# require an explicit psi_init of matching dimension.
if psi_init is None:
n_effects = 1
psi = np.eye(n_effects)
else:
n_effects = psi_init.shape[0]
psi = psi_init.copy()
if q % n_effects != 0:
raise ValueError(
f"Z has q={q} columns, not divisible by n_effects={n_effects} "
f"(inferred from psi_init shape "
f"{psi_init.shape if psi_init is not None else None}). For random "
"intercept + slope models, pass psi_init with shape "
"(n_effects, n_effects)."
)
# Initialize parameters
beta = np.zeros(p)
b = np.zeros(q)
# Get link function
link = family.default_link
# Initial linear predictor
eta = X @ beta + Z @ b
mu = link.inverse(eta)
converged = False
# Main optimization loop
for iter_num in range(maxiter): # noqa: B007
beta_old = beta.copy()
b_old = b.copy()
psi_old = psi.copy()
# Step 1: Find conditional mode b̂ given β, θ
b, hessian = _find_conditional_mode(X, Z, y, beta, psi, n_effects, family)
# Step 2: Update β given b̂
beta = _update_fixed_effects(X, Z, y, b, psi, n_effects, family, S, lambda_)
# Step 3: Update θ (variance components) given β, b̂, H
psi = _update_variance_laplace(b, hessian, n_effects)
# Update predictions
eta = X @ beta + Z @ b
mu = link.inverse(eta)
# Check convergence
delta_beta = np.max(np.abs(beta - beta_old))
delta_b = np.max(np.abs(b - b_old))
delta_psi = np.max(np.abs(psi - psi_old))
if delta_beta < tol and delta_b < tol and delta_psi < tol:
converged = True
break
# Compute approximate marginal log-likelihood
log_lik = _compute_laplace_log_likelihood(y, mu, b, psi, hessian, n_effects, family)
# Residual variance (Gaussian only)
if isinstance(family, GaussianFamily):
residuals = y - mu
sigma2 = np.sum(residuals**2) / (n - p)
else:
sigma2 = 1.0 # Fixed for non-Gaussian
return LaplaceResult(
beta=beta,
b=b,
psi=psi,
sigma2=sigma2,
fitted_values=mu,
linear_predictor=eta,
hessian=hessian,
converged=converged,
n_iter=iter_num + 1,
log_likelihood=log_lik,
)
def _find_conditional_mode(
X: np.ndarray,
Z: np.ndarray,
y: np.ndarray,
beta: np.ndarray,
psi: np.ndarray,
n_effects: int,
family: Family,
) -> tuple[np.ndarray, np.ndarray]:
"""Find conditional mode of random effects given fixed effects.
Maximizes: log p(b | y, β, θ) = log f(y | b, β) + log p(b | θ) - const
Parameters
----------
X : ndarray, shape (n, p)
Fixed effects design
Z : ndarray, shape (n, q)
Random effects design
y : ndarray, shape (n,)
Response
beta : ndarray, shape (p,)
Current fixed effects
psi : ndarray, shape (n_effects, n_effects)
Current variance-covariance matrix
n_effects : int
Number of effects per group
family : Family
Distribution family
Returns
-------
b_mode : ndarray, shape (q,)
Conditional mode
hessian : ndarray, shape (q, q)
Hessian at mode
"""
q = Z.shape[1]
n_groups = q // n_effects
# Compute Ψ⁻¹
try:
psi_inv = linalg.inv(psi)
except linalg.LinAlgError:
psi_inv = linalg.inv(psi + 1e-6 * np.eye(n_effects))
# Expand to block diagonal
psi_inv_block = np.kron(np.eye(n_groups), psi_inv)
# Get link function
link = family.default_link
# Define negative log conditional density
def neg_log_conditional(b):
eta = X @ beta + Z @ b
mu = link.inverse(eta)
# Negative log-likelihood
neg_log_lik = -family.log_likelihood(y, mu).sum()
# Negative log prior
neg_log_prior = 0.5 * (b @ psi_inv_block @ b)
return neg_log_lik + neg_log_prior
# Gradient
def grad_neg_log_conditional(b):
eta = X @ beta + Z @ b
mu = link.inverse(eta)
deta_dmu = link.derivative(mu)
dmu_deta = 1.0 / (deta_dmu + 1e-10)
# Score from likelihood
residuals = (y - mu) / family.variance(mu)
grad_log_lik = Z.T @ (residuals * dmu_deta)
# Gradient of log prior
grad_log_prior = -psi_inv_block @ b
return -(grad_log_lik + grad_log_prior)
# Optimize using L-BFGS-B
result = optimize.minimize(
neg_log_conditional,
x0=np.zeros(q),
jac=grad_neg_log_conditional,
method="L-BFGS-B",
options={"maxiter": 100, "ftol": 1e-8},
)
b_mode = result.x
# Compute Hessian at mode numerically
hessian = _compute_hessian(X, Z, y, beta, b_mode, psi_inv_block, family)
return b_mode, hessian
def _compute_hessian(
X: np.ndarray,
Z: np.ndarray,
y: np.ndarray,
beta: np.ndarray,
b: np.ndarray,
psi_inv_block: np.ndarray,
family: Family,
) -> np.ndarray:
"""Compute Hessian of negative log conditional density.
H = Z'WZ + Ψ⁻¹
Where W = diag(w_i), w_i = (dμ/dη)² / Var(μ) · |d²μ/dη²|
Parameters
----------
X : ndarray, shape (n, p)
Fixed effects design
Z : ndarray, shape (n, q)
Random effects design
y : ndarray, shape (n,)
Response
beta : ndarray, shape (p,)
Fixed effects
b : ndarray, shape (q,)
Random effects
psi_inv_block : ndarray, shape (q, q)
Block-diagonal inverse variance-covariance
family : Family
Distribution family
Returns
-------
hessian : ndarray, shape (q, q)
Hessian matrix
"""
link = family.default_link
eta = X @ beta + Z @ b
mu = link.inverse(eta)
deta_dmu = link.derivative(mu)
dmu_deta = 1.0 / (deta_dmu + 1e-10)
var_mu = family.variance(mu)
# Weights for Hessian
w = (dmu_deta**2) / var_mu
w = np.maximum(w, 1e-8)
# Hessian from likelihood
W = np.diag(w)
hess_lik = Z.T @ W @ Z
# Total Hessian
hessian = hess_lik + psi_inv_block
return np.asarray(hessian)
def _update_fixed_effects(
X: np.ndarray,
Z: np.ndarray,
y: np.ndarray,
b: np.ndarray,
psi: np.ndarray,
n_effects: int,
family: Family,
S: np.ndarray,
lambda_: float,
) -> np.ndarray:
"""Update fixed effects given random effects.
Uses penalized IRLS with b fixed.
Parameters
----------
X : ndarray, shape (n, p)
Fixed effects design
Z : ndarray, shape (n, q)
Random effects design
y : ndarray, shape (n,)
Response
b : ndarray, shape (q,)
Current random effects
psi : ndarray, shape (n_effects, n_effects)
Variance-covariance matrix
n_effects : int
Number of effects per group
family : Family
Distribution family
S : ndarray, shape (p, p)
Penalty matrix
lambda_ : float
Smoothing parameter
Returns
-------
beta : ndarray, shape (p,)
Updated fixed effects
"""
p = X.shape[1]
beta = np.zeros(p)
# Get link function
link = family.default_link
# IRLS for β with offset Zb
offset = Z @ b
for _ in range(10): # Max 10 iterations
eta = X @ beta + offset
mu = link.inverse(eta)
deta_dmu = link.derivative(mu)
dmu_deta = 1.0 / (deta_dmu + 1e-10)
var_mu = family.variance(mu)
# Working response
z = eta + (y - mu) / dmu_deta - offset
# Weights
w = (dmu_deta**2) / var_mu
w = np.maximum(w, 1e-8)
# Weighted penalized least squares
W_sqrt = np.sqrt(w)
X_w = X * W_sqrt[:, np.newaxis]
z_w = z * W_sqrt
A = X_w.T @ X_w + lambda_ * S
rhs = X_w.T @ z_w
try:
beta_new = linalg.solve(A, rhs, assume_a="pos")
except linalg.LinAlgError:
beta_new = linalg.solve(A + 1e-6 * np.eye(p), rhs)
if np.max(np.abs(beta_new - beta)) < 1e-6:
beta = beta_new
break
beta = beta_new
return beta
def _update_variance_laplace(
b: np.ndarray,
hessian: np.ndarray,
n_effects: int,
) -> np.ndarray:
"""Update variance components using Laplace approximation.
Uses the empirical Bayes second-moment estimate corrected for the
posterior uncertainty of the random effects:
Ψ = (1/m) Σᵢ [ b̂ᵢ b̂ᵢᵀ + (H⁻¹)ᵢᵢ ]
where (H⁻¹)ᵢᵢ is the diagonal block of the inverse Hessian for group
i — the Laplace approximation to the posterior covariance of b_i.
H = ZᵀWZ + Ψ_full⁻¹ is block-diagonal across groups (both terms are),
so the diagonal blocks of H⁻¹ coincide with the per-group marginal
blocks.
Parameters
----------
b : ndarray, shape (q,)
Random effects at conditional mode
hessian : ndarray, shape (q, q)
Hessian at conditional mode
n_effects : int
Number of effects per group
Returns
-------
psi : ndarray, shape (n_effects, n_effects)
Updated variance-covariance matrix
"""
q = len(b)
n_groups = q // n_effects
# Reshape b into groups
b_matrix = b.reshape(n_groups, n_effects)
# Invert Hessian (posterior precision) once
try:
H_inv = linalg.inv(hessian)
except linalg.LinAlgError:
H_inv = linalg.pinv(hessian)
# Empirical second moment + posterior covariance correction
psi_new = (b_matrix.T @ b_matrix) / n_groups
for i in range(n_groups):
sl = slice(i * n_effects, (i + 1) * n_effects)
psi_new += H_inv[sl, sl] / n_groups
# Symmetrize and ensure positive definiteness
psi_new = 0.5 * (psi_new + psi_new.T)
eigvals = np.linalg.eigvalsh(psi_new)
if np.min(eigvals) < 1e-6:
psi_new = psi_new + (1e-6 - np.min(eigvals)) * np.eye(n_effects)
return np.asarray(psi_new)
def _compute_laplace_log_likelihood(
y: np.ndarray,
mu: np.ndarray,
b: np.ndarray,
psi: np.ndarray,
hessian: np.ndarray,
n_effects: int,
family: Family,
) -> float:
"""Compute approximate marginal log-likelihood via Laplace.
log L(β, θ) ≈ log f(y | b̂, β) + log p(b̂ | θ) + (q/2)log(2π) − (1/2) log |H|
Parameters
----------
y : ndarray, shape (n,)
Response
mu : ndarray, shape (n,)
Fitted means
b : ndarray, shape (q,)
Random effects at mode
psi : ndarray, shape (n_effects, n_effects)
Variance-covariance matrix
hessian : ndarray, shape (q, q)
Hessian at mode
n_effects : int
Number of effects per group
family : Family
Distribution family
Returns
-------
log_lik : float
Approximate marginal log-likelihood
"""
# Log-likelihood at mode
log_lik_data = np.sum(family.log_likelihood(y, mu))
# Log prior density at mode
n_groups = len(b) // n_effects
b_matrix = b.reshape(n_groups, n_effects)
try:
L = linalg.cholesky(psi, lower=True)
log_det_psi = 2 * np.sum(np.log(np.diag(L)))
psi_inv_b = linalg.cho_solve((L, True), b_matrix.T).T
except linalg.LinAlgError:
log_det_psi = np.linalg.slogdet(psi)[1]
psi_inv_b = linalg.solve(psi, b_matrix.T, assume_a="pos").T
log_prior = -0.5 * (
n_groups * (n_effects * np.log(2 * np.pi) + log_det_psi) + np.sum(b_matrix * psi_inv_b)
)
# Laplace correction: (q/2)·log(2π) − (1/2)·log|H|, where the
# (2π)^{q/2} factor comes from the q-dimensional Gaussian integral.
# The constant is needed for comparability with lme4's reported
# log-likelihoods.
try:
_, log_det_hess = np.linalg.slogdet(hessian)
except np.linalg.LinAlgError:
log_det_hess = 0.0 # Fallback
laplace_correction = 0.5 * len(b) * np.log(2 * np.pi) - 0.5 * log_det_hess
return float(log_lik_data + log_prior + laplace_correction)
def _get_family(family_name: str) -> Family:
"""Get family object from name.
Parameters
----------
family_name : str
Family name ('gaussian', 'poisson', 'binomial')
Returns
-------
family : Family
Family object with default link
Raises
------
ValueError
If family name not recognized
"""
families: dict[str, type[Family]] = {
"gaussian": GaussianFamily,
"poisson": PoissonFamily,
"binomial": BinomialFamily,
}
if family_name.lower() not in families:
raise ValueError(f"Unknown family '{family_name}'. Must be one of {list(families.keys())}")
return families[family_name.lower()]()