# SPDX-License-Identifier: MIT
# Copyright (c) 2025 Lucy Eduardo Arias
"""Iteratively Reweighted Least Squares (IRLS) for Generalized Linear Models.
Mathematical Framework
----------------------
IRLS is the standard algorithm for fitting Generalized Linear Models (GLMs).
It is equivalent to Fisher scoring when using the canonical link, and closely
related to Newton-Raphson for non-canonical links.
GLM Structure
-------------
A GLM relates the response Y to predictors X through:
g(μ) = η = X^T β
where:
- μ = E[Y|X] is the mean
- g(·) is the link function
- η is the linear predictor
- β are regression coefficients
**Distribution**: Y ~ ExpFamily(μ, φ) with variance Var(Y) = φ V(μ)
IRLS Algorithm
--------------
Starting from initial estimates β⁽⁰⁾, iterate until convergence:
**Step 1: Compute working response and weights**
For iteration t, given current estimates β⁽ᵗ⁾:
η⁽ᵗ⁾ = X β⁽ᵗ⁾
μ⁽ᵗ⁾ = g⁻¹(η⁽ᵗ⁾)
Working response (adjusted dependent variable):
z⁽ᵗ⁾ = η⁽ᵗ⁾ + (y - μ⁽ᵗ⁾) g'(μ⁽ᵗ⁾)
Working weights (iterative weights):
w⁽ᵗ⁾ = 1 / [V(μ⁽ᵗ⁾) (g'(μ⁽ᵗ⁾))²]
**Step 2: Solve weighted least squares**
Update coefficients by solving:
β⁽ᵗ⁺¹⁾ = argmin_β ||W^(1/2)(z - Xβ)||²
where W = diag(w⁽ᵗ⁾).
**Normal equations**:
(X^T W X) β⁽ᵗ⁺¹⁾ = X^T W z⁽ᵗ⁾
**Step 3: Check convergence**
Stop when the relative parameter change is below tolerance:
||β⁽ᵗ⁺¹⁾ − β⁽ᵗ⁾|| / (||β⁽ᵗ⁾|| + ε) < tol (both dense and sparse routes).
Equivalence to Fisher Scoring
------------------------------
IRLS is exactly Fisher scoring for GLMs. The Fisher information matrix is:
I(β) = E[-∂²ℓ/∂β∂β^T] = X^T W X
where the expected information equals the observed information for
exponential families.
**Fisher scoring update**:
β⁽ᵗ⁺¹⁾ = β⁽ᵗ⁾ + [I(β⁽ᵗ⁾)]⁻¹ ∂ℓ/∂β|_{β⁽ᵗ⁾}
= β⁽ᵗ⁾ + (X^T W X)⁻¹ X^T W (y - μ⁽ᵗ⁾) / g'(μ⁽ᵗ⁾)
This is equivalent to solving:
(X^T W X) β⁽ᵗ⁺¹⁾ = X^T W [η⁽ᵗ⁾ + (y - μ⁽ᵗ⁾)g'(μ⁽ᵗ⁾)]
= X^T W z⁽ᵗ⁾
Convergence Properties
----------------------
**Theorem** (Green, 1984): For canonical link functions and regular
exponential families, IRLS converges to the maximum likelihood estimate
if it exists.
**Rate**: Locally quadratic convergence near the solution (Fisher scoring
inherits Newton-Raphson's convergence rate for canonical links).
**Conditions for convergence**:
1. Design matrix X has full column rank
2. Starting values yield finite μ⁽⁰⁾ in the support
3. V(μ) > 0 and g'(μ) ≠ 0 throughout iteration path
**Failure modes**:
- Separation (perfect prediction) in logistic regression → β → ∞
- Zero responses in Poisson/Gamma models → boundary issues
- Near-collinearity → ill-conditioned X^T W X
Numerical Stability
-------------------
**Techniques used**:
1. **Safe division**: Clip denominators V(μ) (g'(μ))² away from zero
- Minimum threshold: 10⁻¹²
- Prevents overflow in weights
2. **Direct solve via Cholesky**: For positive-definite X^T W X
- Computational cost: O(p³) for p predictors
- Numerically stable when well-conditioned
3. **Backend-agnostic**: Supports NumPy, PyTorch, JAX
- Automatic differentiation not required
- Pure linear algebra operations
4. **Sparse matrix support**: Automatic detection and handling
- Uses scipy.sparse operations when X is sparse
- The (dense) p × p normal equations are solved via Cholesky,
with a pivoted least-squares fallback
- Significant speedup for high-dimensional categorical data
Sparse Matrix Support
---------------------
When the design matrix X is sparse (scipy.sparse), IRLS automatically uses
sparse-aware operations for computational efficiency:
**Sparse operations**:
1. **Matrix-vector products**: X @ β uses sparse BLAS level 2
2. **Weighted normal equations**: X^T W X computed via sparse matmul, then
densified (it is only p × p)
3. **Linear solve**: Dense Cholesky on X^T W X, falling back to
``numpy.linalg.lstsq`` when X^T W X is not positive definite
**Complexity with sparse X** (nnz = number of non-zeros):
- Matrix-vector product: O(nnz) instead of O(np)
- X^T W X formation: O(nnz × p) when result is dense
- Total per iteration: O(nnz × p + p³) vs O(np² + p³) dense
**When to use sparse**:
- High-dimensional categorical features (one-hot encoded)
- B-spline basis matrices (compact support → banded/sparse)
- Interaction terms with many zero entries
- Design matrices with > 90% zeros typically benefit
**Detection**: Sparse matrices are auto-detected via `scipy.sparse.issparse()`.
Computational Complexity
------------------------
Per iteration, for n observations and p predictors:
**Operations**:
- Link computations (η → μ, g'(μ)): O(n)
- Variance function V(μ): O(n)
- Weight computation: O(n)
- Form X^T W X: O(np² + p²)
- Solve (X^T W X)β = X^T W z: O(p³)
**Total per iteration**: O(np² + p³)
**Typical iterations**: 3-10 for well-conditioned problems
**Total cost**: O(k(np² + p³)) where k is iteration count
Comparison with Other Methods
------------------------------
**vs Newton-Raphson**:
- IRLS uses expected information (Fisher), NR uses observed
- For canonical links: IRLS = NR
- IRLS more stable for GLMs (guaranteed positive-definite)
**vs gradient descent**:
- Faster convergence (quadratic vs linear)
- No learning rate tuning required
- More expensive per iteration: O(p³) vs O(np)
**vs L-BFGS**:
- IRLS exploits GLM structure (closed-form weights)
- L-BFGS more general, uses quasi-Newton approximation
- For large p: L-BFGS may be faster (avoids full Hessian)
Implementation Notes
--------------------
**Multi-backend support**:
All operations work transparently with NumPy, PyTorch, and JAX arrays
through the backend abstraction.
**Offset terms**:
Supports known offset η₀:
η = X^T β + η₀
Useful for log-exposure models in Poisson regression.
**Initial values**:
Caller must provide β⁽⁰⁾. Common strategies:
- GLM with identity link: β⁽⁰⁾ = (X^T X)⁻¹ X^T g(y)
- Zeros: often works for well-scaled data
- Previous fit: warm-starting
References
----------
**Core IRLS theory**:
- Green, P. J. (1984). \"Iteratively reweighted least squares for maximum
likelihood estimation, and some robust and resistant alternatives.\"
*Journal of the Royal Statistical Society: Series B*, 46(2), 149-192.
https://doi.org/10.1111/j.2517-6161.1984.tb01288.x
- McCullagh, P., & Nelder, J. A. (1989). *Generalized Linear Models* (2nd ed.).
Chapman and Hall/CRC. Chapter 2.5: Fitting GLMs.
https://doi.org/10.1007/978-1-4899-3242-6
**Convergence analysis**:
- Wedderburn, R. W. M. (1976). \"On the existence and uniqueness of the
maximum likelihood estimates for certain generalized linear models.\"
*Biometrika*, 63(1), 27-32.
https://doi.org/10.1093/biomet/63.1.27
**Numerical methods**:
- Dempster, A. P., Laird, N. M., & Rubin, D. B. (1977). \"Maximum likelihood
from incomplete data via the EM algorithm.\" *Journal of the Royal
Statistical Society: Series B*, 39(1), 1-22.
(IRLS as special case of EM)
**Sparse linear algebra**:
- Davis, T. A. (2006). *Direct Methods for Sparse Linear Systems*.
SIAM. (SuperLU, UMFPACK algorithms)
- Saad, Y. (2003). *Iterative Methods for Sparse Linear Systems* (2nd ed.).
SIAM. (CG, GMRES for large sparse systems)
**Separation and boundary issues**:
- Albert, A., & Anderson, J. A. (1984). \"On the existence of maximum
likelihood estimates in logistic regression models.\" *Biometrika*,
71(1), 1-10.
https://doi.org/10.1093/biomet/71.1.1
See Also
--------
aurora.models.glm.fitting : GLM fitting using IRLS
aurora.distributions.base : Variance functions and link functions
aurora.core.optimization.newton : Newton-Raphson method
Notes
-----
For detailed mathematical derivations, see REFERENCES.md in the repository root.
IRLS is the workhorse algorithm for GLM fitting, combining computational
efficiency with statistical optimality (Fisher scoring achieves the
Cramér-Rao lower bound asymptotically).
"""
from __future__ import annotations
from collections.abc import Callable
from typing import Any
import numpy as np
from ..types import Array, OptimizationCallback
from .result import OptimizationResult
# Condition number monitoring threshold
_ILL_CONDITIONED_THRESHOLD = 1e8
# =============================================================================
# Sparse matrix utilities
# =============================================================================
def _is_sparse(X) -> bool:
"""Check if X is a scipy sparse matrix."""
try:
from scipy import sparse
return bool(sparse.issparse(X))
except ImportError:
return False
def _sparse_diag_matmul(X, w):
"""Compute diag(w) @ X for sparse X.
For sparse X, this is equivalent to row-scaling: X_i,: *= w_i
Parameters
----------
X : sparse matrix (n, p)
Design matrix
w : ndarray (n,)
Diagonal weights
Returns
-------
sparse matrix (n, p)
Row-scaled matrix
"""
from scipy import sparse
# Create diagonal matrix and multiply
W = sparse.diags(w)
return W @ X
def _sparse_weighted_lstsq(X, w, z):
"""Solve weighted least squares with sparse X.
Solves: β = argmin_β ||W^(1/2)(z - Xβ)||²
This is equivalent to solving the normal equations:
(X^T W X) β = X^T W z
For sparse X with dense result (p small), we form X^T W X densely
and solve with standard dense solver. For very large p with sparse
X^T W X, an iterative solver would be preferred.
Parameters
----------
X : sparse matrix (n, p)
Design matrix
w : ndarray (n,)
Weights (positive)
z : ndarray (n,)
Working response
Returns
-------
beta : ndarray (p,)
Least squares solution
Notes
-----
The algorithm:
1. Form W^(1/2) X via row scaling (sparse)
2. Form X^T W X = (W^(1/2) X)^T (W^(1/2) X) (dense if p small)
3. Form X^T W z (dense vector)
4. Solve via Cholesky or direct solver
Complexity: O(nnz × p + p³) where nnz = number of non-zeros in X
"""
from scipy.linalg import cho_factor, cho_solve
n, p = X.shape
# Compute sqrt(W) @ X by row scaling
sqrt_w = np.sqrt(w)
WhalfX = _sparse_diag_matmul(X, sqrt_w)
# Form X^T W X (dense) - this is p × p
# WhalfX^T @ WhalfX
XtWX = (WhalfX.T @ WhalfX).toarray()
# Form X^T W z (dense vector)
XtWz = X.T @ (w * z)
# Solve via Cholesky (X^T W X is positive definite when X full rank)
try:
L, lower = cho_factor(XtWX, lower=True)
beta = cho_solve((L, lower), XtWz)
except np.linalg.LinAlgError:
# Fallback to least squares if Cholesky fails
beta, _, _, _ = np.linalg.lstsq(XtWX, XtWz, rcond=None)
return beta
def _irls_sparse(
loss_fn: Callable,
init_params: np.ndarray,
X, # scipy.sparse matrix
y: np.ndarray,
link,
variance_fn: Callable,
offset: np.ndarray | None,
max_iter: int,
tol: float,
callback,
args: tuple,
kwargs: dict,
) -> OptimizationResult:
"""IRLS implementation for sparse design matrices.
Uses scipy.sparse operations for matrix-vector products and
forms the normal equations X^T W X as dense (suitable for p << n).
Parameters
----------
loss_fn : callable
Loss function for convergence monitoring
init_params : ndarray (p,)
Initial parameter estimates
X : sparse matrix (n, p)
Design matrix
y : ndarray (n,)
Response vector
link : LinkFunction
Link function with inverse() and derivative() methods
variance_fn : callable
Variance function V(μ)
offset : ndarray (n,) or None
Known offset term
max_iter : int
Maximum iterations
tol : float
Convergence tolerance
callback : callable or None
Progress callback
args, kwargs : tuple, dict
Additional arguments for loss_fn
Returns
-------
OptimizationResult
Optimization result with solution and diagnostics
"""
beta = np.asarray(init_params, dtype=np.float64)
offset_arr = np.zeros_like(y) if offset is None else np.asarray(offset)
nfev = 0
total_backtrack = 0
max_cond_number = 0.0
for iteration in range(max_iter):
# Linear predictor without offset: η* = Xβ
eta_star = np.asarray(X @ beta).ravel()
# Full linear predictor with offset: η = η* + offset
eta = eta_star + offset_arr
# Mean: μ = g⁻¹(η)
mu = link.inverse(eta)
# Link derivative: g'(μ)
g_prime = link.derivative(mu)
# Variance function: V(μ)
var = variance_fn(mu)
# Working weights: w = 1 / [V(μ) (g'(μ))²]
denom = var * (g_prime**2)
denom = np.clip(denom, 1e-12, None) # Numerical stability
weights = 1.0 / denom
# Working response: z* = η* + (y - μ) g'(μ)
# Note: We use η* (without offset) so WLS solves for β correctly
# This is equivalent to: z* = y - offset (for identity link)
z_star = eta_star + (y - mu) * g_prime
# Solve weighted least squares: β = (X^T W X)^{-1} X^T W z*
beta_new = _sparse_weighted_lstsq(X, weights, z_star)
# Check convergence
delta = beta_new - beta
step_norm = np.sqrt(np.sum(delta**2))
# Monitor condition number of X^T W X
from ..linalg import weighted_condition_number
cond_number = weighted_condition_number(np.asarray(X.todense()), weights)
if cond_number > max_cond_number:
max_cond_number = cond_number
if cond_number > _ILL_CONDITIONED_THRESHOLD:
import warnings
warnings.warn(
f"Ill-conditioned system detected: κ = {cond_number:.2e} > 1e8. "
"Results may be numerically unstable.",
RuntimeWarning,
stacklevel=3,
)
# Evaluate current loss before step-halving
nfev += 1
current_loss = float(loss_fn(beta, *args, **kwargs))
# Step-halving (backtracking line search)
step_size = 1.0
alpha = 0.5
max_backtrack = 10
min_step_size = 1.0 / 64.0
for backtrack in range(max_backtrack + 1):
trial_beta = beta + step_size * delta
nfev += 1
trial_loss = float(loss_fn(trial_beta, *args, **kwargs))
if trial_loss <= current_loss or step_size < min_step_size:
beta = trial_beta
loss_value = trial_loss
if backtrack > 0:
total_backtrack += 1
break
step_size *= alpha
else:
beta = trial_beta
loss_value = trial_loss
total_backtrack += 1
if callback is not None:
callback(iteration, beta.copy(), loss_value)
# Relative convergence test, aligned with the dense route
# (scipy.optimize / R glm convention).
beta_norm = float(np.sqrt(np.sum(beta**2)))
if step_norm / (beta_norm + 1e-12) < tol:
return OptimizationResult(
x=beta,
fun=loss_value,
grad=None,
success=True,
message="Converged: parameter change below tolerance",
nit=iteration + 1,
nfev=nfev,
njev=0,
nhev=0,
backtrack_iterations=total_backtrack,
condition_number=max_cond_number,
)
# Max iterations reached
loss_value = float(loss_fn(beta, *args, **kwargs))
nfev += 1
return OptimizationResult(
x=beta,
fun=loss_value,
grad=None,
success=False,
message="Maximum iterations reached",
nit=max_iter,
nfev=nfev,
njev=0,
nhev=0,
backtrack_iterations=total_backtrack,
condition_number=max_cond_number,
)
# =============================================================================
# Main IRLS interface
# =============================================================================
[docs]
def irls(
loss_fn: Callable,
init_params: Array,
*,
backend=None,
args: tuple = (),
kwargs: dict | None = None,
max_iter: int = 100,
tol: float = 1e-6,
callback: OptimizationCallback | None = None,
design_matrix: Array | None = None,
response: Array | None = None,
link: Any | None = None,
variance_fn: Callable[[Array], Array] | None = None,
offset: Array | None = None,
) -> OptimizationResult:
"""Run the IRLS procedure for generalized linear models.
This function automatically detects sparse design matrices and
uses optimized sparse operations when beneficial.
Parameters
----------
loss_fn : callable
Loss function taking ``(params, *args, **kwargs)``
init_params : array-like (p,)
Initial parameter estimates
backend : Backend, optional
Computational backend (numpy, pytorch, jax)
args : tuple
Positional arguments for loss_fn
kwargs : dict
Keyword arguments for loss_fn
max_iter : int, default=100
Maximum IRLS iterations
tol : float, default=1e-6
Convergence tolerance on parameter change
callback : callable, optional
Callback function(iteration, params, loss)
design_matrix : array-like (n, p)
Design matrix X. May be dense or scipy.sparse.
response : array-like (n,)
Response vector y
link : LinkFunction
Link function with inverse() and derivative()
variance_fn : callable
Variance function V(μ)
offset : array-like (n,), optional
Known offset term
Returns
-------
OptimizationResult
Optimization result containing:
- x: final parameters
- fun: final loss value
- success: convergence status
- nit: number of iterations
- message: status message
Notes
-----
For sparse design matrices (scipy.sparse), the algorithm automatically
uses sparse matrix operations:
1. Matrix-vector products X @ β use sparse BLAS
2. Normal equations X^T W X are formed efficiently
3. The p × p system is solved with Cholesky factorization
This is efficient when:
- X is sparse (> 90% zeros)
- Number of predictors p is moderate (< 10,000)
- X^T W X fits in memory as dense
For very large sparse X^T W X, iterative solvers (CG) would be needed.
"""
if kwargs is None:
kwargs = {}
if design_matrix is None or response is None or link is None or variance_fn is None:
raise ValueError(
"IRLS requires design_matrix, response, link, and variance_fn keyword arguments."
)
if not hasattr(link, "inverse") or not hasattr(link, "derivative"):
raise TypeError("link must expose inverse() and derivative() methods")
# Check for sparse matrix - use specialized path
if _is_sparse(design_matrix):
return _irls_sparse(
loss_fn=loss_fn,
init_params=np.asarray(init_params),
X=design_matrix,
y=np.asarray(response),
link=link,
variance_fn=variance_fn,
offset=np.asarray(offset) if offset is not None else None,
max_iter=max_iter,
tol=tol,
callback=callback,
args=args,
kwargs=kwargs,
)
# Dense path - use backend abstraction
if backend is None:
from ..backends import get_backend
backend = get_backend("jax")
X = backend.array(design_matrix)
y = backend.array(response)
beta = backend.array(init_params)
offset_arr = backend.array(offset) if offset is not None else y * 0
converted_args = tuple(_convert_to_backend(backend, value) for value in args)
converted_kwargs = {key: _convert_to_backend(backend, value) for key, value in kwargs.items()}
nfev = 0
total_backtrack = 0
max_cond_number = 0.0
def _to_backend(data):
return backend.array(data, dtype=getattr(beta, "dtype", None))
for iteration in range(max_iter):
# Linear predictor without offset: η* = Xβ
eta_star = X @ beta
# Full linear predictor with offset: η = η* + offset
eta = eta_star + offset_arr
mu = link.inverse(eta)
g_prime = link.derivative(mu)
var = variance_fn(mu)
weights = _safe_divide(backend, 1.0, var * (g_prime**2))
# Working response on the η* scale (offset removed) so the WLS step
# solves against X alone and does not absorb the offset into β.
z = eta_star + (y - mu) * g_prime
sqrt_w = _sqrt(backend, weights)
WX = X * sqrt_w.unsqueeze(-1) if hasattr(sqrt_w, "unsqueeze") else X * sqrt_w[:, None]
Wz = z * sqrt_w
beta_new = None
try: # Prefer backend-native linear algebra when available
import torch # type: ignore
if isinstance(WX, torch.Tensor):
lhs = WX.T @ WX
rhs = WX.T @ Wz
if rhs.ndim == 1:
rhs = rhs.unsqueeze(-1)
beta_new = torch.linalg.solve(lhs, rhs).squeeze(-1)
beta_new = beta_new.to(beta.device).type_as(beta)
except ImportError: # pragma: no cover - optional dependency missing
beta_new = None
if beta_new is None:
WX_np = backend.as_numpy(WX)
Wz_np = backend.as_numpy(Wz)
lhs = WX_np.T @ WX_np
rhs = WX_np.T @ Wz_np
solution = np.linalg.solve(lhs, rhs)
beta_new = _to_backend(solution)
delta = beta_new - beta
step_norm = backend.as_numpy((delta * delta).sum() ** 0.5)
# Monitor condition number of weighted normal equations
from ..linalg import weighted_condition_number
X_np = backend.as_numpy(X)
weights_np = backend.as_numpy(weights)
cond_number = weighted_condition_number(X_np, weights_np)
if cond_number > max_cond_number:
max_cond_number = cond_number
if cond_number > _ILL_CONDITIONED_THRESHOLD:
import warnings
warnings.warn(
f"Ill-conditioned system detected: κ = {cond_number:.2e} > 1e8. "
"Results may be numerically unstable.",
RuntimeWarning,
stacklevel=3,
)
# Evaluate current loss before step-halving
nfev += 1
current_loss = float(backend.as_numpy(loss_fn(beta, *converted_args, **converted_kwargs)))
# Step-halving (backtracking line search)
step_size = 1.0
alpha = 0.5
max_backtrack = 10
min_step_size = 1.0 / 64.0
for backtrack in range(max_backtrack + 1):
trial_beta = beta + step_size * delta
nfev += 1
trial_loss = float(
backend.as_numpy(loss_fn(trial_beta, *converted_args, **converted_kwargs))
)
if trial_loss <= current_loss or step_size < min_step_size:
beta = trial_beta
loss_value = trial_loss
if backtrack > 0:
total_backtrack += 1
break
step_size *= alpha
else:
beta = trial_beta
loss_value = trial_loss
total_backtrack += 1
if callback is not None:
callback(iteration, backend.as_numpy(beta), float(backend.as_numpy(loss_value)))
# Relative convergence test (matches scipy.optimize / R glm convention):
# ||delta|| / (||beta|| + eps) < tol. The +eps guards against
# division by zero when beta collapses to 0 (e.g., intercept-only).
beta_norm = float(backend.as_numpy((beta * beta).sum() ** 0.5))
if step_norm / (beta_norm + 1e-12) < tol:
return OptimizationResult(
x=backend.as_numpy(beta),
fun=float(backend.as_numpy(loss_value)),
grad=None,
success=True,
message="Converged: relative parameter change below tolerance",
nit=iteration + 1,
nfev=nfev,
njev=0,
nhev=0,
backtrack_iterations=total_backtrack,
condition_number=max_cond_number,
)
loss_value = loss_fn(beta, *converted_args, **converted_kwargs)
nfev += 1
return OptimizationResult(
x=backend.as_numpy(beta),
fun=float(backend.as_numpy(loss_value)),
grad=None,
success=False,
message="Maximum iterations reached",
nit=max_iter,
nfev=nfev,
njev=0,
nhev=0,
backtrack_iterations=total_backtrack,
condition_number=max_cond_number,
)
def _safe_divide(backend, numerator, denominator):
eps = 1e-12
if hasattr(denominator, "clamp_min"):
denominator = denominator.clamp_min(eps)
return numerator / denominator
denom_np = backend.as_numpy(denominator)
denom_np = np.clip(denom_np, eps, None)
return numerator / backend.array(denom_np)
def _sqrt(backend, value):
if hasattr(value, "sqrt"):
return value.sqrt()
value_np = backend.as_numpy(value)
return backend.array(np.sqrt(value_np))
def _convert_to_backend(backend, value):
if isinstance(value, (tuple, list)):
converted = [_convert_to_backend(backend, item) for item in value]
return type(value)(converted)
if isinstance(value, dict):
return {key: _convert_to_backend(backend, item) for key, item in value.items()}
try:
return backend.array(value)
except Exception: # pragma: no cover - fallback when conversion is not applicable
return value
__all__ = ["irls"]