statgpu 0.1.0__py3-none-any.whl

statgpu/glm_core/_family.py

@@ -0,0 +1,362 @@
+"""
+Link and Family abstractions for GLM.
+
+Extracted from the duplicated IRLS loops in _logistic.py across CPU/GPU/Torch backends.
+Each Family defines: link function, variance function, and IRLS weights/working response.
+
+All operations are backend-aware: numpy/cupy/torch via _xp dispatch.
+"""
+
+from abc import ABC, abstractmethod
+from typing import Any
+
+import numpy as np
+
+from statgpu.backends._array_ops import _clip, _log, _xp
+from statgpu.inference._distributions_backend import get_distribution
+
+__all__ = [
+    "Link", "IdentityLink", "LogLink", "InversePowerLink", "InverseSquaredLink",
+    "GLMFamily", "Gaussian", "Binomial", "Poisson", "Gamma",
+    "InverseGaussian", "NegativeBinomial", "Tweedie",
+]
+
+
+def _backend_name(arr):
+    """Infer backend name from array type."""
+    mod = type(arr).__module__
+    if mod.startswith("cupy"):
+        return "cupy"
+    if mod.startswith("torch"):
+        return "torch"
+    return "numpy"
+
+
+# _exp with overflow protection (clips input to [-500, 500])
+_ETA_CLIP_EXP = 500.0
+
+
+def _exp(arr):
+    """Exponential with overflow protection."""
+    return _xp(arr).exp(_clip(arr, -_ETA_CLIP_EXP, _ETA_CLIP_EXP))
+
+
+def _sqrt(arr):
+    """Square root with clamp to prevent NaN from negative values."""
+    return _xp(arr).sqrt(_clip(arr, 0, None))
+
+
+def _ones_like(arr):
+    return _xp(arr).ones_like(arr)
+
+
+def _cdf(arr):
+    """Standard normal CDF (Phi)."""
+    backend = _backend_name(arr)
+    return get_distribution("norm", backend=backend).cdf(arr)
+
+
+def _ppf(arr):
+    """Standard normal PPF (inverse CDF, Phi^{-1})."""
+    backend = _backend_name(arr)
+    return get_distribution("norm", backend=backend).ppf(arr)
+
+
+def _pdf(arr):
+    """Standard normal PDF (phi)."""
+    backend = _backend_name(arr)
+    return get_distribution("norm", backend=backend).pdf(arr)
+
+
+# ─── Link Functions ────────────────────────────────────────────────────────
+
+
+class Link(ABC):
+    """Link function abstract base class.
+
+    Maps between mean (mu) and linear predictor (eta):
+        eta = link(mu)
+        mu  = inverse(eta)
+    """
+
+    name: str
+
+    @abstractmethod
+    def link(self, mu):
+        """eta = g(mu)."""
+        pass
+
+    @abstractmethod
+    def inverse(self, eta):
+        """mu = g^{-1}(eta)."""
+        pass
+
+    @abstractmethod
+    def derivative(self, mu):
+        """g'(mu) = d eta / d mu."""
+        pass
+
+
+class LogitLink(Link):
+    name = "logit"
+
+    def link(self, mu):
+        return _log(mu / (1 - mu))
+
+    def inverse(self, eta):
+        return 1.0 / (1.0 + _exp(-_clip(eta, -500, 500)))
+
+    def derivative(self, mu):
+        return 1.0 / (mu * (1 - mu))
+
+
+class ProbitLink(Link):
+    """Probit link: inverse of standard normal CDF (Phi)."""
+
+    name = "probit"
+
+    def link(self, mu):
+        return _ppf(_clip(mu, 1e-10, 1 - 1e-10))
+
+    def inverse(self, eta):
+        return _cdf(eta)
+
+    def derivative(self, mu):
+        return 1.0 / _pdf(
+            _ppf(_clip(mu, 1e-10, 1 - 1e-10))
+        )
+
+
+class LogLink(Link):
+    name = "log"
+
+    def link(self, mu):
+        return _log(_clip(mu, 1e-10, None))
+
+    def inverse(self, eta):
+        return _exp(eta)
+
+    def derivative(self, mu):
+        return 1.0 / mu
+
+
+class IdentityLink(Link):
+    name = "identity"
+
+    def link(self, mu):
+        return mu
+
+    def inverse(self, eta):
+        return eta
+
+    def derivative(self, mu):
+        return _ones_like(mu)
+
+
+class InversePowerLink(Link):
+    """Inverse power link: eta = 1/mu (canonical for Gamma)."""
+
+    name = "inverse_power"
+    _ETA_LO = 1e-4
+    _ETA_HI = 1e3
+
+    def link(self, mu):
+        return 1.0 / _clip(mu, 1e-10, None)
+
+    def inverse(self, eta):
+        return 1.0 / _clip(eta, self._ETA_LO, self._ETA_HI)
+
+    def derivative(self, mu):
+        return -1.0 / (mu * mu)
+
+
+class InverseSquaredLink(Link):
+    """Inverse squared link: eta = 1/mu^2 (canonical for InverseGaussian)."""
+
+    name = "inverse_squared"
+
+    def link(self, mu):
+        return 1.0 / _clip(mu * mu, 1e-10, None)
+
+    def inverse(self, eta):
+        eta_c = _clip(eta, 1e-20, None)
+        return 1.0 / _clip(_sqrt(eta_c), 1e-10, None)
+
+    def derivative(self, mu):
+        return -2.0 / (mu * mu * mu)
+
+
+# ─── Families ──────────────────────────────────────────────────────────────
+
+
+class GLMFamily(ABC):
+    """GLM distribution family.
+
+    Each family defines:
+    - link function: eta <-> mu mapping
+    - variance function: Var(Y) = phi * V(mu)
+    - IRLS weights and working response computation
+    """
+
+    name: str
+    link: Link
+
+    @abstractmethod
+    def variance(self, mu):
+        """Variance function V(mu)."""
+        pass
+
+    def irls_weights(self, mu, y):
+        """IRLS working weights.
+
+        W = 1 / (V(mu) * (g'(mu))^2)
+
+        Default uses the inverse Fisher weights for the WLS step in IRLS.
+        Subclasses can override for more efficient implementations.
+        """
+        denom = self.variance(mu) * self.link.derivative(mu) ** 2
+        return 1.0 / _clip(denom, 1e-10, None)
+
+    def irls_working_response(self, mu, y, eta):
+        """Working response z = eta + (y - mu) * link'(mu)."""
+        return eta + (y - mu) * self.link.derivative(mu)
+
+
+class Gaussian(GLMFamily):
+    """Gaussian family with identity link (standard linear regression)."""
+
+    name = "gaussian"
+    link = IdentityLink()
+
+    def variance(self, mu):
+        return _ones_like(mu)
+
+
+class Binomial(GLMFamily):
+    """Binomial family with configurable link (logistic/probit regression)."""
+
+    name = "binomial"
+
+    def __init__(self, link=None):
+        self.link = link if link is not None else LogitLink()
+
+    def variance(self, mu):
+        return mu * (1 - mu)
+
+    def irls_weights(self, mu, y):
+        mu_c = _clip(mu, 1e-10, 1 - 1e-10)
+        return mu_c * (1 - mu_c)
+
+    def irls_working_response(self, mu, y, eta):
+        mu_c = _clip(mu, 1e-10, 1 - 1e-10)
+        var = mu_c * (1 - mu_c)
+        return eta + (y - mu_c) / var
+
+
+class Poisson(GLMFamily):
+    """Poisson family with log link (Poisson regression)."""
+
+    name = "poisson"
+    link = LogLink()
+
+    def variance(self, mu):
+        return mu
+
+    def irls_weights(self, mu, y):
+        return _clip(mu, 1e-10, None)
+
+    def irls_working_response(self, mu, y, eta):
+        return eta + (y - mu) / _clip(mu, 1e-10, None)
+
+
+class Gamma(GLMFamily):
+    """Gamma family (positive continuous outcomes).
+
+    Default link is log for numerical stability. Canonical link is inverse_power.
+    """
+
+    name = "gamma"
+
+    def __init__(self, link=None):
+        self.link = link if link is not None else LogLink()
+
+    def variance(self, mu):
+        return mu * mu
+
+
+class InverseGaussian(GLMFamily):
+    """Inverse Gaussian family (positive continuous, right-skewed).
+
+    Default link is log for numerical stability.
+    """
+
+    name = "inverse_gaussian"
+    link = LogLink()
+
+    def variance(self, mu):
+        return mu * mu * mu
+
+    def irls_weights(self, mu, y):
+        mu_c = _clip(mu, 1e-10, None)
+        return _ones_like(mu) / mu_c
+
+    def irls_working_response(self, mu, y, eta):
+        mu_c = _clip(mu, 1e-10, None)
+        # z = eta + (y - mu) * g'(mu) = eta + (y - mu) / mu  (log link)
+        return eta + (y - mu_c) / mu_c
+
+
+class NegativeBinomial(GLMFamily):
+    """Negative Binomial family (overdispersed count data).
+
+    Uses log link. Dispersion parameter ``alpha`` controls overdispersion:
+    Var(Y) = mu + alpha * mu^2. When alpha -> 0, approaches Poisson.
+    """
+
+    name = "negative_binomial"
+    link = LogLink()
+
+    def __init__(self, alpha=1.0):
+        if not np.isfinite(alpha) or alpha <= 0.0:
+            raise ValueError("alpha must be a finite positive scalar for negative binomial family")
+        self.alpha = alpha
+
+    def variance(self, mu):
+        return mu + self.alpha * mu * mu
+
+    def irls_weights(self, mu, y):
+        mu_c = _clip(mu, 1e-10, None)
+        return mu_c / (1.0 + self.alpha * mu_c)
+
+    def irls_working_response(self, mu, y, eta):
+        mu_c = _clip(mu, 1e-10, None)
+        return eta + (y - mu_c) / mu_c
+
+
+class Tweedie(GLMFamily):
+    """Tweedie family (power variance function).
+
+    Variance function: V(mu) = mu^power.
+    - power=0: Gaussian
+    - power=1: Poisson
+    - power=2: Gamma
+    - 1 < power < 2: compound Poisson-Gamma (most common usage)
+    """
+
+    name = "tweedie"
+    link = LogLink()
+
+    def __init__(self, power=1.5):
+        self.power = power
+
+    def variance(self, mu):
+        return _clip(mu, 1e-10, None) ** self.power
+
+    def irls_weights(self, mu, y):
+        mu_c = _clip(mu, 1e-10, None)
+        # w = 1 / (V(mu) * g'(mu)^2) = 1 / (mu^p * (1/mu)^2) = mu^(2-p)
+        return mu_c ** (2.0 - self.power)
+
+    def irls_working_response(self, mu, y, eta):
+        mu_c = _clip(mu, 1e-10, None)
+        return eta + (y - mu_c) / mu_c

statgpu/glm_core/_fused.py

@@ -0,0 +1,149 @@
+"""GLM-specific fused loss+gradient functions.
+
+These avoid redundant X @ coef computation by computing value and gradient
+in a single pass. They are called by GLMLoss subclasses' fused_value_and_gradient()
+methods for performance.
+
+NOT part of the generic solver interface — these are GLM internal optimizations.
+"""
+
+from statgpu.backends._utils import _to_float_scalar, _get_xp
+from statgpu.backends import _to_numpy
+
+
+def _fused_logistic(eta, X, y, n, loss):
+    from statgpu.backends._array_ops import _sigmoid, _softplus, _sum
+    p = _sigmoid(eta)
+    log1pexp = _softplus(eta)
+    val = _sum(-y * eta + log1pexp) / n
+    grad = X.T @ (p - y) / n
+    return val, grad
+
+
+def _fused_poisson(eta, X, y, n, loss):
+    from statgpu.backends._array_ops import _exp, _log, _clip, _sum
+    mu = _exp(_clip(eta, -30, 30))
+    mu_c = _clip(mu, 1e-10, None)
+    val = _sum(mu - y * _log(mu_c)) / n
+    grad = X.T @ (mu - y) / n
+    return val, grad
+
+
+def _fused_gamma(eta, X, y, n, loss):
+    from statgpu.backends._array_ops import _exp, _log, _clip, _sum
+    gamma_link = getattr(loss, "link_name", getattr(loss, "link", "log"))
+    if gamma_link == "inverse_power":
+        eta_lo = float(getattr(loss, "_ETA_LO", 1e-2))
+        eta_hi = float(getattr(loss, "_ETA_HI", 1e3))
+        eta_c = _clip(eta, eta_lo, eta_hi)
+        mu = 1.0 / eta_c
+        val = _sum(y * eta_c - _log(eta_c)) / n
+        grad = X.T @ (y - mu) / n
+        return val, grad
+    mu = _exp(_clip(eta, -30, 30))
+    mu_c = _clip(mu, 1e-10, None)
+    val = _sum(y / mu_c + _log(mu_c)) / n
+    grad = X.T @ ((mu_c - y) / mu_c) / n
+    return val, grad
+
+
+def _fused_negative_binomial(eta, X, y, n, loss):
+    from statgpu.backends._array_ops import _exp, _log, _clip, _sum
+    a = float(getattr(loss, "alpha", 1.0))
+    mu = _exp(_clip(eta, -30, 30))
+    mu_c = _clip(mu, 1e-300, None)
+    one_plus_a_mu = 1.0 + a * mu_c
+    val = (
+        _sum(-y * _log(mu_c / one_plus_a_mu) + (1.0 / a) * _log(one_plus_a_mu)) / n
+    )
+    grad = X.T @ ((mu_c - y) / one_plus_a_mu) / n
+    return val, grad
+
+
+def _fused_tweedie(eta, X, y, n, loss):
+    from statgpu.backends._array_ops import _exp, _clip, _sum, _log
+    pw = float(getattr(loss, "power", 1.5))
+    mu = _exp(_clip(eta, -50, 50))
+    mu_c = _clip(mu, 1e-10, 1e6)
+    log_mu = _log(mu_c)
+    d1 = 1.0 - pw
+    d2 = 2.0 - pw
+    if abs(d1) < 0.01:
+        term1 = -y * log_mu
+    else:
+        term1 = -y * mu_c**d1 / d1
+    if abs(d2) < 0.01:
+        term2 = log_mu
+    else:
+        term2 = mu_c**d2 / d2
+    val = _sum(term1 + term2) / n
+    grad = X.T @ (mu_c**d1 * (mu_c - y)) / n
+    return val, grad
+
+
+def _fused_inverse_gaussian(eta, X, y, n, loss):
+    from statgpu.backends._array_ops import _exp, _clip, _sum
+    mu = _exp(_clip(eta, -30, 30))
+    mu_c = _clip(mu, 5e-2, 1e3)
+    val = _sum(y / (2.0 * mu_c * mu_c) - 1.0 / mu_c) / n
+    grad = X.T @ ((mu_c - y) / (mu_c * mu_c)) / n
+    return val, grad
+
+
+# Module-level dispatch table (avoids rebuilding dict on every call)
+_FUSED_DISPATCH = {
+    "logistic": _fused_logistic,
+    "poisson": _fused_poisson,
+    "gamma": _fused_gamma,
+    "negative_binomial": _fused_negative_binomial,
+    "tweedie": _fused_tweedie,
+    "inverse_gaussian": _fused_inverse_gaussian,
+}
+
+
+def _fused_glm_value_and_gradient(loss, X, y, coef):
+    """Dispatch to fused kernel based on loss name (GLM-specific)."""
+    n = X.shape[0]
+    eta = X @ coef
+    loss_name = getattr(loss, "name", "")
+    if loss_name in _FUSED_DISPATCH:
+        return _FUSED_DISPATCH[loss_name](eta, X, y, n, loss)
+    return loss.value(X, y, coef), loss.gradient(X, y, coef)
+
+
+def _weighted_loss_and_grad(loss, X, y, coef, sample_weight):
+    """Weighted loss+gradient (GLM-specific fast paths)."""
+    from statgpu.backends import _resolve_backend
+    n = X.shape[0]
+    _backend = _resolve_backend("auto", X)
+    xp = _get_xp(_backend)
+    _sw_np = _to_numpy(sample_weight)
+    if hasattr(X, "device"):
+        _sw = xp.asarray(_sw_np, dtype=X.dtype, device=X.device)
+    else:
+        _sw = xp.asarray(_sw_np, dtype=X.dtype)
+    sw_sum = _to_float_scalar(xp.sum(_sw))
+
+    loss_name = getattr(loss, "name", "")
+    if loss_name == "squared_error":
+        resid = X @ coef - y
+        grad = X.T @ (_sw * resid) / sw_sum
+        val = 0.5 * _to_float_scalar(xp.sum(_sw * resid * resid)) / sw_sum
+        return val, grad
+
+    if hasattr(loss, "fused_value_and_gradient"):
+        try:
+            return loss.fused_value_and_gradient(
+                X, y, coef, sample_weight=sample_weight
+            )
+        except TypeError:
+            pass
+
+    try:
+        val = loss.value(X, y, coef, sample_weight=sample_weight)
+        grad = loss.gradient(X, y, coef, sample_weight=sample_weight)
+        return val, grad
+    except TypeError:
+        val = loss.value(X, y, coef)
+        grad = loss.gradient(X, y, coef)
+        return val, grad

statgpu/glm_core/_gamma.py

@@ -0,0 +1,111 @@
+"""
+Gamma loss: negative Gamma log-likelihood.
+
+For positive continuous outcomes:
+    loss = (1/n) * sum(y/mu + log(mu))
+where mu is determined by the configured link:
+    - log: mu = exp(X @ coef)
+    - inverse_power: mu = 1 / (X @ coef)
+
+Supports numpy / cupy / torch backends via _array_ops helpers.
+"""
+from statgpu.backends._array_ops import _clip, _exp, _log, _sum, _max_eigval_power, _xp
+from statgpu.glm_core._base import GLMLoss, register_glm_loss
+
+
+@register_glm_loss('gamma')
+class GammaLoss(GLMLoss):
+    name = "gamma"
+    y_type = "positive"
+    smooth_gradient = True
+    has_hessian = True
+    _lipschitz_uses_y = True
+    _lipschitz_safety = 3.0  # Gamma Hessian varies with mu
+    _conservative_momentum_with_nonsmooth = True
+    _gamma_like = True
+
+    _MU_LO = 1e-3
+    _MU_HI = 1e4
+    _ETA_LO = 1e-4
+    _ETA_HI = 1e3
+
+    def __init__(self, link="log"):
+        if link not in ("log", "inverse_power"):
+            raise ValueError(
+                "GammaLoss link must be 'log' or 'inverse_power', "
+                f"got {link!r}."
+            )
+        self.link = link
+        self.link_name = link
+        self._lipschitz_at_init = link == "inverse_power"
+        self._has_constant_hessian = (link == "log")
+
+    def _eta_mu(self, X, coef):
+        eta = X @ coef
+        if self.link == "inverse_power":
+            eta_c = _clip(eta, self._ETA_LO, self._ETA_HI)
+            return eta_c, 1.0 / eta_c
+        z = _clip(eta, -30, 30)
+        return z, _clip(_exp(z), self._MU_LO, self._MU_HI)
+
+    def _mu_from_eta(self, eta):
+        if self.link == "inverse_power":
+            eta_c = _clip(eta, self._ETA_LO, self._ETA_HI)
+            return 1.0 / eta_c
+        return _clip(_exp(_clip(eta, -30, 30)), self._MU_LO, self._MU_HI)
+
+    # ── Per-sample formulas (single source of truth) ──────────────────
+
+    def per_sample_value(self, eta, y):
+        if self.link == "inverse_power":
+            eta_c = _clip(eta, self._ETA_LO, self._ETA_HI)
+            return y * eta_c - _log(eta_c)
+        mu = self._mu_from_eta(eta)
+        return y / mu + _log(mu)
+
+    def per_sample_gradient(self, eta, y):
+        if self.link == "inverse_power":
+            mu = self._mu_from_eta(eta)
+            return y - mu
+        mu = self._mu_from_eta(eta)
+        return 1.0 - y / mu
+
+    def hessian(self, X, y, coef, sample_weight=None):
+        n_eff = float(sample_weight.sum()) if sample_weight is not None else X.shape[0]
+        if self.link == "inverse_power":
+            eta, _ = self._eta_mu(X, coef)
+            W = 1.0 / (eta * eta)
+        else:
+            # Expected Fisher: W(mu) = 1 for Gamma with log link
+            xp = _xp(X)
+            if xp.__name__ == "torch":
+                W = xp.ones(X.shape[0], dtype=X.dtype, device=X.device)
+            else:
+                W = xp.ones(X.shape[0], dtype=X.dtype)
+        if sample_weight is not None:
+            W = W * sample_weight
+        return X.T @ (X * W[:, None]) / n_eff
+
+    def lipschitz(self, X, coef, y=None, sample_weight=None):
+        n_eff = float(sample_weight.sum()) if sample_weight is not None else X.shape[0]
+        if self.link == "inverse_power":
+            eta, _ = self._eta_mu(X, coef)
+            W = 1.0 / (eta * eta)
+        elif y is not None:
+            z = _clip(X @ coef, -30, 30)
+            mu = _clip(_exp(z), self._MU_LO, self._MU_HI)
+            W = y / mu
+        else:
+            XtX = X.T @ X
+            return max(_max_eigval_power(XtX) / n_eff, 1e-8)
+        if sample_weight is not None:
+            W = W * sample_weight
+        XtWX = X.T @ (X * W[:, None])
+        L = _max_eigval_power(XtWX) / n_eff
+        return max(L, 1e-8)
+
+    def predict(self, X, coef):
+        if self.link == "inverse_power":
+            eta = _clip(X @ coef, self._ETA_LO, self._ETA_HI)
+            return 1.0 / eta
+        return _exp(X @ coef)

statgpu/glm_core/_inverse_gaussian.py

@@ -0,0 +1,62 @@
+"""
+Inverse Gaussian loss: negative log-likelihood with log link.
+
+For positive right-skewed outcomes:
+    loss = (1/n) * sum(y/(2*mu^2) - 1/mu)
+where mu = exp(X @ coef).
+
+Supports numpy / cupy / torch backends via _array_ops helpers.
+"""
+from statgpu.backends._array_ops import _clip, _exp, _sum, _max_eigval_power
+from statgpu.glm_core._base import GLMLoss, register_glm_loss
+
+
+@register_glm_loss('inverse_gaussian')
+class InverseGaussianLoss(GLMLoss):
+    name = "inverse_gaussian"
+    y_type = "positive"
+    smooth_gradient = True
+    has_hessian = True
+    _lipschitz_uses_y = True
+    _lipschitz_safety = 3.0  # 1/mu^3 gradient scaling requires safety factor
+    _skip_momentum = True
+    _inverse_gaussian = True  # 1/mu^3 scaling causes Nesterov oscillation
+
+    _MU_LO = 5e-2
+    _MU_HI = 1e3
+
+    def _mu_from_eta(self, eta):
+        return _clip(_exp(_clip(eta, -30, 30)), self._MU_LO, self._MU_HI)
+
+    # ── Per-sample formulas (single source of truth) ──────────────────
+
+    def per_sample_value(self, eta, y):
+        mu = self._mu_from_eta(eta)
+        return y / (2.0 * mu * mu) - 1.0 / mu
+
+    def per_sample_gradient(self, eta, y):
+        mu = self._mu_from_eta(eta)
+        return (mu - y) / (mu * mu)
+
+    def hessian(self, X, y, coef, sample_weight=None):
+        z = _clip(X @ coef, -30, 30)
+        mu = _clip(_exp(z), self._MU_LO, self._MU_HI)
+        W = 1.0 / mu
+        if sample_weight is not None:
+            W = W * sample_weight
+        n_eff = float(sample_weight.sum()) if sample_weight is not None else X.shape[0]
+        return X.T @ (X * W[:, None]) / n_eff
+
+    def lipschitz(self, X, coef, y=None, sample_weight=None):
+        z = _clip(X @ coef, -30, 30)
+        mu = _clip(_exp(z), self._MU_LO, self._MU_HI)
+        W = 1.0 / mu
+        if sample_weight is not None:
+            W = W * sample_weight
+        n_eff = float(sample_weight.sum()) if sample_weight is not None else X.shape[0]
+        XtWX = X.T @ (X * W[:, None])
+        L = _max_eigval_power(XtWX) / n_eff
+        return max(L, 1e-8)
+
+    def predict(self, X, coef):
+        return _exp(X @ coef)