Generalized Linear Models

"""Poisson GLM by iteratively reweighted least squares - the math, transcribed.

Every IRLS step is a weighted least-squares fit to the working response
z = eta + (y - mu) * g'(mu), with weights w = 1 / (V(mu) * g'(mu)^2). For a log
link on Poisson data, g'(mu) = 1/mu and V(mu) = mu, so the weight is simply mu
and the working response is eta + (y - mu) / mu.
"""

import math


def solve(matrix_a, vector_b):
    """Gaussian elimination with partial pivoting - the d x d system per step."""
    d = len(vector_b)
    augmented = [list(row) + [vector_b[i]] for i, row in enumerate(matrix_a)]

    for column in range(d):
        pivot = column
        for row in range(column + 1, d):
            if abs(augmented[row][column]) > abs(augmented[pivot][column]):
                pivot = row
        if abs(augmented[pivot][column]) < 1e-12:
            raise ValueError("singular system: features are collinear")
        augmented[column], augmented[pivot] = augmented[pivot], augmented[column]

        for row in range(column + 1, d):
            factor = augmented[row][column] / augmented[column][column]
            for k in range(column, d + 1):
                augmented[row][k] -= factor * augmented[column][k]

    solution = [0.0] * d
    for row in reversed(range(d)):
        acc = augmented[row][d]
        for k in range(row + 1, d):
            acc -= augmented[row][k] * solution[k]
        solution[row] = acc / augmented[row][row]
    return solution


def fit(X, y, steps=25):
    n = len(X)
    d = len(X[0])
    beta = [0.0] * d

    for _ in range(steps):
        # Accumulate X^T W X and X^T W z one observation at a time.
        hessian = [[0.0] * d for _ in range(d)]
        rhs = [0.0] * d

        for i in range(n):
            eta = 0.0
            for j in range(d):
                eta += X[i][j] * beta[j]

            mu = math.exp(eta)                 # inverse link
            weight = mu                        # 1 / (V(mu) * g'(mu)^2) = mu
            working = eta + (y[i] - mu) / mu   # z_i

            for j in range(d):
                rhs[j] += weight * X[i][j] * working
                for k in range(d):
                    hessian[j][k] += weight * X[i][j] * X[i][k]

        beta = solve(hessian, rhs)

    return beta

"""Generalized linear models by IRLS - typed, family-parameterized, validated."""

from collections.abc import Callable
from dataclasses import dataclass

import numpy as np
from numpy.typing import NDArray

Vector = NDArray[np.float64]
Matrix = NDArray[np.float64]

_MU_FLOOR = 1e-10


@dataclass(frozen=True)
class Family:
    """An exponential-family member, defined by exactly what IRLS needs of it.

    Everything separating Poisson from binomial from Gamma regression lives in
    these three functions. The fitting loop below never changes.
    """

    name: str
    inverse_link: Callable[[Vector], Vector]
    link_derivative: Callable[[Vector], Vector]    # g'(mu)
    variance: Callable[[Vector], Vector]           # V(mu)


POISSON_LOG = Family(
    name="poisson(log)",
    inverse_link=np.exp,
    link_derivative=lambda mu: 1.0 / mu,
    variance=lambda mu: mu,
)

BINOMIAL_LOGIT = Family(
    name="binomial(logit)",
    inverse_link=lambda eta: 1.0 / (1.0 + np.exp(-eta)),
    link_derivative=lambda mu: 1.0 / (mu * (1.0 - mu)),
    variance=lambda mu: mu * (1.0 - mu),
)


@dataclass(frozen=True)
class GlmModel:
    coefficients: Vector
    family: Family
    iterations: int
    converged: bool

    def predict(self, X: Matrix) -> Vector:
        """Returns the conditional mean, not the linear predictor."""
        if X.ndim != 2 or X.shape[1] != self.coefficients.size:
            raise ValueError(
                f"expected (n, {self.coefficients.size}) design matrix, got {X.shape}"
            )
        return self.family.inverse_link(X @ self.coefficients)


def fit(
    X: Matrix,
    y: Vector,
    family: Family = POISSON_LOG,
    offset: Vector | None = None,
    max_iter: int = 25,
    tol: float = 1e-8,
) -> GlmModel:
    """Fit by Fisher scoring. Raises ValueError on malformed input."""
    if X.ndim != 2:
        raise ValueError(f"X must be 2-D, got shape {X.shape}")
    if X.shape[0] != y.shape[0]:
        raise ValueError(f"X has {X.shape[0]} rows but y has {y.shape[0]}")
    if offset is not None and offset.shape != y.shape:
        raise ValueError(f"offset must match y, got {offset.shape} vs {y.shape}")

    # None, not np.zeros(...): a mutable default argument would be shared
    # across every call to this function.
    baseline = np.zeros_like(y) if offset is None else offset

    beta = np.zeros(X.shape[1], dtype=np.float64)
    converged = False
    iteration = 0

    for iteration in range(1, max_iter + 1):
        eta = X @ beta + baseline
        mu = np.clip(family.inverse_link(eta), _MU_FLOOR, None)
        g_prime = family.link_derivative(mu)

        weights = 1.0 / (family.variance(mu) * g_prime * g_prime)
        working = eta - baseline + (y - mu) * g_prime

        weighted = X * weights[:, None]
        beta_next = np.linalg.solve(X.T @ weighted, weighted.T @ working)

        if np.max(np.abs(beta_next - beta)) < tol:
            beta = beta_next
            converged = True
            break
        beta = beta_next

    return GlmModel(
        coefficients=beta,
        family=family,
        iterations=iteration,
        converged=converged,
    )

"""GLM by IRLS - QR on the whitened design, reused buffers, no normal equations."""

from collections.abc import Callable

import numpy as np
from numpy.typing import NDArray

Vector = NDArray[np.float64]
Matrix = NDArray[np.float64]

_MU_FLOOR = 1e-10


def fit_fast(
    X: Matrix,
    y: Vector,
    inverse_link: Callable[[Vector], Vector],
    link_derivative: Callable[[Vector], Vector],
    variance: Callable[[Vector], Vector],
    max_iter: int = 25,
    tol: float = 1e-8,
) -> Vector:
    """Same fixed point, different linear algebra.

    Forming X^T W X squares the condition number of the design and throws away
    roughly half the available precision. Solving the equivalent least-squares
    problem min ||sqrt(W) (X beta - z)|| by QR does not. This is not a
    micro-optimization: on a Poisson fit the weights ARE the fitted means, so a
    response spanning several orders of magnitude puts several orders of
    magnitude into W - which is exactly the case where the normal equations
    lose the fit and the factorization keeps it.
    """
    design = np.ascontiguousarray(X, dtype=np.float64)
    target = np.ascontiguousarray(y, dtype=np.float64)
    n, d = design.shape
    beta = np.zeros(d, dtype=np.float64)

    # Buffers allocated once and overwritten in place on every iteration.
    eta = np.empty(n, dtype=np.float64)
    whitened = np.empty_like(design)
    working = np.empty(n, dtype=np.float64)

    for _ in range(max_iter):
        np.matmul(design, beta, out=eta)
        mu = np.clip(inverse_link(eta), _MU_FLOOR, None)
        g_prime = link_derivative(mu)

        # sqrt of the IRLS weight, kept in root form so W is never materialized.
        root_w = 1.0 / np.sqrt(variance(mu) * g_prime * g_prime)

        np.subtract(y, mu, out=working)
        np.multiply(working, g_prime, out=working)
        np.add(working, eta, out=working)              # z = eta + (y - mu) g'(mu)

        np.multiply(design, root_w[:, None], out=whitened)
        beta_next, *_ = np.linalg.lstsq(whitened, root_w * working, rcond=None)

        if np.max(np.abs(beta_next - beta)) < tol:
            return beta_next
        beta = beta_next

    return beta