|
| 1 | +# -*- coding: utf-8 -*- |
| 2 | +# emacs: -*- mode: python; py-indent-offset: 4; indent-tabs-mode: nil -*- |
| 3 | +# vi: set ft=python sts=4 ts=4 sw=4 et: |
| 4 | +r""" |
| 5 | +Robust M-estimator regression and the MAD scale estimator. |
| 6 | +
|
| 7 | +Outlier-resistant linear regression by iteratively reweighted least squares |
| 8 | +(IRLS): each observation is down-weighted by an *influence function* of its |
| 9 | +current standardised residual, so a handful of gross outliers (motion-corrupted |
| 10 | +fMRI frames, a mislabelled subject) cannot dominate the fit the way they do |
| 11 | +ordinary least squares. |
| 12 | +
|
| 13 | +- :func:`mad` -- the median absolute deviation, the robust scale estimator that |
| 14 | + pairs with the M-estimators (normal-consistent by default). |
| 15 | +- :func:`huber_regress` -- Huber's monotone M-estimator: quadratic near zero, |
| 16 | + linear in the tails. Bounded influence but never zero -- large outliers still |
| 17 | + contribute a constant pull. |
| 18 | +- :func:`tukey_bisquare_regress` -- the Tukey bisquare *redescender*: influence |
| 19 | + rises, then falls back to zero beyond the tuning radius, so extreme outliers |
| 20 | + are rejected outright (at the cost of a non-convex objective needing a decent |
| 21 | + start; the OLS warm start suffices in practice). |
| 22 | +
|
| 23 | +Both estimators are pure composition: the IRLS loop is a fixed-iteration |
| 24 | +:func:`jax.lax.fori_loop` over the cuSOLVER-free ``(p, p)`` weighted-least- |
| 25 | +squares solve (:func:`nitrix.linalg._smalllinalg.small_inv_logdet`, the same |
| 26 | +inner solve the GLM/GAM suite uses), so the fit is differentiable and carries no |
| 27 | +factorisation that could invoke a fragile cuSOLVER path. Fit a single response |
| 28 | +``y`` of shape ``(N,)``; ``jax.vmap`` over a stack of responses for the |
| 29 | +mass-univariate case. |
| 30 | +
|
| 31 | +References |
| 32 | +---------- |
| 33 | +Huber PJ (1964). Robust estimation of a location parameter. *Annals of |
| 34 | +Mathematical Statistics*, 35(1), 73-101. |
| 35 | +https://doi.org/10.1214/aoms/1177703732 |
| 36 | +
|
| 37 | +Beaton AE, Tukey JW (1974). The fitting of power series, meaning polynomials, |
| 38 | +illustrated on band-spectroscopic data. *Technometrics*, 16(2), 147-185. |
| 39 | +https://doi.org/10.1080/00401706.1974.10489171 |
| 40 | +""" |
| 41 | + |
| 42 | +from __future__ import annotations |
| 43 | + |
| 44 | +from typing import Callable, NamedTuple, Optional, Union |
| 45 | + |
| 46 | +import jax.lax as lax |
| 47 | +import jax.numpy as jnp |
| 48 | +from jaxtyping import Array, Float |
| 49 | + |
| 50 | +from ..linalg._smalllinalg import small_inv_logdet |
| 51 | + |
| 52 | +__all__ = [ |
| 53 | + 'RobustFit', |
| 54 | + 'huber_regress', |
| 55 | + 'mad', |
| 56 | + 'tukey_bisquare_regress', |
| 57 | +] |
| 58 | + |
| 59 | +# 1 / Phi^{-1}(3/4) = 1 / 0.674489...: scales the raw MAD to a consistent |
| 60 | +# estimator of the standard deviation under Gaussian data. |
| 61 | +_MAD_TO_SIGMA = 1.4826022185056018 |
| 62 | + |
| 63 | + |
| 64 | +class RobustFit(NamedTuple): |
| 65 | + """A fitted robust regression (a NamedTuple of arrays, not a module). |
| 66 | +
|
| 67 | + Attributes |
| 68 | + ---------- |
| 69 | + coef : Float[Array, 'p'] |
| 70 | + The M-estimated coefficient vector. |
| 71 | + scale : Float[Array, ''] |
| 72 | + The final robust residual scale (normal-consistent :func:`mad` of the |
| 73 | + residuals), in the same units as ``y``. |
| 74 | + weights : Float[Array, 'N'] |
| 75 | + The final per-observation IRLS weights in :math:`[0, 1]`; near-zero |
| 76 | + weights flag observations the estimator treated as outliers. |
| 77 | + residuals : Float[Array, 'N'] |
| 78 | + The final residuals ``y - X @ coef``. |
| 79 | + """ |
| 80 | + |
| 81 | + coef: Float[Array, 'p'] |
| 82 | + scale: Float[Array, ''] |
| 83 | + weights: Float[Array, 'N'] |
| 84 | + residuals: Float[Array, 'N'] |
| 85 | + |
| 86 | + |
| 87 | +def mad( |
| 88 | + x: Float[Array, '...'], |
| 89 | + *, |
| 90 | + axis: int = -1, |
| 91 | + center: Optional[Union[float, Float[Array, '...']]] = None, |
| 92 | + normalize: bool = True, |
| 93 | +) -> Float[Array, '...']: |
| 94 | + r"""Median absolute deviation of ``x`` along an axis. |
| 95 | +
|
| 96 | + :math:`\operatorname{MAD}(x) = \operatorname{median}_i\, |
| 97 | + |x_i - \operatorname{median}(x)|`, a breakdown-point-½ robust dispersion |
| 98 | + estimator. With ``normalize=True`` (default) the result is multiplied by |
| 99 | + :math:`1/\Phi^{-1}(3/4) \approx 1.4826`, making it a consistent estimator of |
| 100 | + the standard deviation :math:`\sigma` for Gaussian data (so it is |
| 101 | + interchangeable with the sample standard deviation but resistant to |
| 102 | + outliers). |
| 103 | +
|
| 104 | + Parameters |
| 105 | + ---------- |
| 106 | + x : Float[Array, '...'] |
| 107 | + Input values. |
| 108 | + axis : int, optional |
| 109 | + Axis to reduce. Default ``-1``. |
| 110 | + center : float or Float[Array, '...'], optional |
| 111 | + Centre to deviate about. Default ``None`` uses the median of ``x`` along |
| 112 | + ``axis`` (broadcast back over the reduced axis). |
| 113 | + normalize : bool, optional |
| 114 | + If ``True`` (default), scale by :math:`1.4826` for consistency with the |
| 115 | + Gaussian standard deviation; if ``False``, return the raw MAD. |
| 116 | +
|
| 117 | + Returns |
| 118 | + ------- |
| 119 | + Float[Array, '...'] |
| 120 | + The (optionally normalised) MAD, with ``axis`` reduced. |
| 121 | + """ |
| 122 | + med = jnp.median(x, axis=axis, keepdims=True) if center is None else center |
| 123 | + dev = jnp.median(jnp.abs(x - med), axis=axis) |
| 124 | + factor = _MAD_TO_SIGMA if normalize else 1.0 |
| 125 | + return factor * dev |
| 126 | + |
| 127 | + |
| 128 | +def _huber_weight(u: Float[Array, 'N'], delta: float) -> Float[Array, 'N']: |
| 129 | + """Huber IRLS weight ``w(u) = min(1, delta / |u|)`` (double-``where`` safe).""" |
| 130 | + au = jnp.abs(u) |
| 131 | + inner = au <= delta |
| 132 | + safe = jnp.where(inner, 1.0, au) # avoid 0-div in the unused branch |
| 133 | + return jnp.where(inner, 1.0, delta / safe) |
| 134 | + |
| 135 | + |
| 136 | +def _tukey_weight(u: Float[Array, 'N'], c: float) -> Float[Array, 'N']: |
| 137 | + """Tukey bisquare IRLS weight ``(1 - (u/c)^2)^2`` inside ``|u| <= c``, else 0.""" |
| 138 | + t = u / c |
| 139 | + return jnp.where(jnp.abs(u) <= c, (1.0 - t * t) ** 2, 0.0) |
| 140 | + |
| 141 | + |
| 142 | +def _wls_solve( |
| 143 | + X: Float[Array, 'N p'], |
| 144 | + y: Float[Array, 'N'], |
| 145 | + w: Float[Array, 'N'], |
| 146 | + ridge: float, |
| 147 | + p: int, |
| 148 | +) -> Float[Array, 'p']: |
| 149 | + """One cuSOLVER-free weighted-least-squares solve of the normal equations.""" |
| 150 | + Xw = X * w[:, None] |
| 151 | + xtwx = Xw.T @ X + ridge * jnp.eye(p, dtype=X.dtype) |
| 152 | + inv, _ = small_inv_logdet(xtwx, p) |
| 153 | + return inv @ (Xw.T @ y) |
| 154 | + |
| 155 | + |
| 156 | +def _robust_regress( |
| 157 | + X: Float[Array, 'N p'], |
| 158 | + y: Float[Array, 'N'], |
| 159 | + weight_fn: Callable[[Float[Array, 'N']], Float[Array, 'N']], |
| 160 | + *, |
| 161 | + n_iter: int, |
| 162 | + ridge: float, |
| 163 | +) -> RobustFit: |
| 164 | + """IRLS shared by the Huber and Tukey estimators. |
| 165 | +
|
| 166 | + Warm-starts at OLS, then reweights by ``weight_fn`` of the residuals |
| 167 | + standardised by their (normal-consistent) MAD, solving the weighted normal |
| 168 | + equations each iteration. When the residual scale collapses (a near-exact |
| 169 | + fit), it falls back to unit scale so the reweighting reduces to OLS. |
| 170 | + """ |
| 171 | + n, p = X.shape |
| 172 | + eps = jnp.finfo(X.dtype).eps |
| 173 | + ones = jnp.ones((n,), X.dtype) |
| 174 | + beta0 = _wls_solve(X, y, ones, ridge, p) |
| 175 | + |
| 176 | + def scaled_weights(beta: Float[Array, 'p']) -> Float[Array, 'N']: |
| 177 | + resid = y - X @ beta |
| 178 | + s = mad(resid, normalize=True) |
| 179 | + s = jnp.where(s > eps, s, 1.0) |
| 180 | + return weight_fn(resid / s) |
| 181 | + |
| 182 | + def step(_: int, beta: Float[Array, 'p']) -> Float[Array, 'p']: |
| 183 | + return _wls_solve(X, y, scaled_weights(beta), ridge, p) |
| 184 | + |
| 185 | + beta = lax.fori_loop(0, n_iter, step, beta0) |
| 186 | + resid = y - X @ beta |
| 187 | + scale = jnp.where(mad(resid, normalize=True) > eps, mad(resid), 1.0) |
| 188 | + return RobustFit( |
| 189 | + coef=beta, |
| 190 | + scale=scale, |
| 191 | + weights=weight_fn(resid / scale), |
| 192 | + residuals=resid, |
| 193 | + ) |
| 194 | + |
| 195 | + |
| 196 | +def huber_regress( |
| 197 | + X: Float[Array, 'N p'], |
| 198 | + y: Float[Array, 'N'], |
| 199 | + *, |
| 200 | + delta: float = 1.345, |
| 201 | + n_iter: int = 20, |
| 202 | + ridge: float = 0.0, |
| 203 | +) -> RobustFit: |
| 204 | + r"""Huber M-estimator linear regression. |
| 205 | +
|
| 206 | + Minimises :math:`\sum_i \rho_\delta(r_i / s)` where :math:`\rho_\delta` is |
| 207 | + quadratic for :math:`|r| \le \delta s` and linear beyond, and :math:`s` is |
| 208 | + the robust residual scale. The influence is bounded but non-redescending, so |
| 209 | + the objective is convex and the fit is stable, but very extreme outliers |
| 210 | + still exert a constant (not vanishing) pull -- use |
| 211 | + :func:`tukey_bisquare_regress` to reject them outright. |
| 212 | +
|
| 213 | + Parameters |
| 214 | + ---------- |
| 215 | + X : Float[Array, 'N p'] |
| 216 | + Design matrix (``N`` observations, ``p`` predictors). Include an |
| 217 | + intercept column explicitly if wanted. |
| 218 | + y : Float[Array, 'N'] |
| 219 | + Response vector. ``vmap`` over a stack for many responses. |
| 220 | + delta : float, optional |
| 221 | + Tuning constant in units of the robust scale. Default ``1.345`` gives |
| 222 | + 95% efficiency relative to OLS at the Gaussian model. |
| 223 | + n_iter : int, optional |
| 224 | + IRLS iterations (fixed, for differentiability). Default ``20``. |
| 225 | + ridge : float, optional |
| 226 | + Optional L2 stabiliser on the normal-equation diagonal (for rank- |
| 227 | + deficient or collinear ``X``). Default ``0.0``. |
| 228 | +
|
| 229 | + Returns |
| 230 | + ------- |
| 231 | + RobustFit |
| 232 | + ``(coef, scale, weights, residuals)``. |
| 233 | + """ |
| 234 | + return _robust_regress( |
| 235 | + X, y, lambda u: _huber_weight(u, delta), n_iter=n_iter, ridge=ridge |
| 236 | + ) |
| 237 | + |
| 238 | + |
| 239 | +def tukey_bisquare_regress( |
| 240 | + X: Float[Array, 'N p'], |
| 241 | + y: Float[Array, 'N'], |
| 242 | + *, |
| 243 | + c: float = 4.685, |
| 244 | + n_iter: int = 20, |
| 245 | + ridge: float = 0.0, |
| 246 | +) -> RobustFit: |
| 247 | + r"""Tukey bisquare (biweight) redescending M-estimator regression. |
| 248 | +
|
| 249 | + Minimises :math:`\sum_i \rho_c(r_i / s)` with the bisquare |
| 250 | + :math:`\rho_c`, whose influence rises then *redescends* to zero for |
| 251 | + :math:`|r| > c\,s`: observations beyond the tuning radius are rejected |
| 252 | + outright (zero weight), giving maximal outlier resistance. The objective is |
| 253 | + non-convex, so the estimate is the redescender basin reachable from the OLS |
| 254 | + warm start -- adequate in practice for the moderate-contamination regime this |
| 255 | + targets. |
| 256 | +
|
| 257 | + Parameters |
| 258 | + ---------- |
| 259 | + X : Float[Array, 'N p'] |
| 260 | + Design matrix (``N`` observations, ``p`` predictors). |
| 261 | + y : Float[Array, 'N'] |
| 262 | + Response vector. ``vmap`` over a stack for many responses. |
| 263 | + c : float, optional |
| 264 | + Tuning radius in units of the robust scale. Default ``4.685`` gives 95% |
| 265 | + efficiency relative to OLS at the Gaussian model. |
| 266 | + n_iter : int, optional |
| 267 | + IRLS iterations (fixed, for differentiability). Default ``20``. |
| 268 | + ridge : float, optional |
| 269 | + Optional L2 stabiliser on the normal-equation diagonal. Default ``0.0``. |
| 270 | +
|
| 271 | + Returns |
| 272 | + ------- |
| 273 | + RobustFit |
| 274 | + ``(coef, scale, weights, residuals)``. |
| 275 | + """ |
| 276 | + return _robust_regress( |
| 277 | + X, y, lambda u: _tukey_weight(u, c), n_iter=n_iter, ridge=ridge |
| 278 | + ) |
0 commit comments