Skip to content

Fitting strategies¤

glmax.fit defaults to IRLSFitter. Pass fitter= to swap strategies without changing anything else in the workflow.

fitted_irls   = glmax.fit(family, X, y)
fitted_newton = glmax.fit(family, X, y, fitter=glmax.NewtonFitter())

Both strategies return the same FittedGLM noun and work with JIT compilation and JAX transforms.

Fitter strategies¤

glmax.AbstractFitter

glmax.AbstractFitter ¤

Abstract base for fit strategies used by fit(family, X, y, fitter=...).

Subclasses must declare a concrete solver: lx.AbstractLinearSolver field and implement fit. The default concrete strategy is IRLSFitter.

fit(self, family: glmax.ExponentialDispersionFamily, X: jax.Array, y: jax.Array, offset: jax.Array, weights: jax.Array | None, init: glmax.Params | None = None) -> glmax.FitResult ¤

Fit family against data and return a FitResult.

Concrete fitters return the full fit contract, not just the fitted coefficient vector.

Arguments:

  • family: glmax.ExponentialDispersionFamily instance.
  • X: covariate matrix, shape (n, p).
  • y: response vector, shape (n,).
  • offset: offset vector, shape (n,).
  • weights: optional per-sample weight vector, shape (n,).
  • init: optional glmax.Params for warm-starting; None uses the family default.

Returns:

glmax.FitResult carrying all fit artifacts.

glmax.IRLSFitter(glmax.AbstractFitter) ¤

Iteratively Reweighted Least Squares (IRLS) fit strategy.

The default glmax.AbstractFitter used by glmax.fit. At each iteration IRLS forms the adjusted response

\[z_i = \eta_i + s \cdot g'(\mu_i)(y_i - \mu_i)\]

and solves the weighted normal equations

\[(X^\top W X)\,\hat\beta = X^\top W z\]

where \(W = \text{diag}(w)\) are the GLM working weights from glmax.ExponentialDispersionFamily.calc_weight and \(s\) is step_size. The linear system is solved with a lineax solver and the linear predictor is updated as \(\eta \leftarrow X\hat\beta + \text{offset}\).

IRLS is mathematically equivalent to Fisher scoring Newton but expressed as a sequence of weighted least-squares problems. It converges in one iteration for Gaussian/identity, where the working weights are constant and the objective is exactly quadratic. For all other families the weights depend on the current mean estimate, so multiple iterations are required. The fixed step_size controls the update magnitude; use glmax.NewtonFitter if you want automatic backtracking line search.

__init__(self, solver: lineax.AbstractLinearSolver = lineax.Cholesky(), step_size: float = 1.0, tol: float = 0.001, max_iter: int = 1000) ¤

Construct an IRLS fitter.

Arguments:

  • solver: lineax solver used for each IRLS weighted least-squares step. Defaults to lx.Cholesky(). Any lx.AbstractLinearSolver that handles symmetric positive-semidefinite systems works here.
  • step_size: IRLS update step-size multiplier. Defaults to 1.0.
  • tol: convergence tolerance on the objective change. Defaults to 1e-3.
  • max_iter: maximum number of IRLS iterations. Defaults to 1000.

glmax.NewtonFitter(glmax.AbstractFitter) ¤

Fisher scoring Newton fit strategy with backtracking Armijo line search.

An glmax.AbstractFitter that solves the GLM log-likelihood directly using Newton's method with the Fisher information matrix as the Hessian. Each Newton step solves

\[\Delta\beta = (X^\top W X)^{-1} X^\top [w \odot g'(\mu) \odot (\mu - y)]\]

where \(W = \text{diag}(w)\) are the GLM working weights from glmax.ExponentialDispersionFamily.calc_weight. Step size is chosen by backtracking Armijo: starting from step_size and shrinking by armijo_factor until the sufficient-decrease condition holds.

Compared to glmax.IRLSFitter, NewtonFitter takes fewer outer iterations on problems where the IRLS fixed-step update overshoots — in particular with non-canonical links or near-boundary means.

__init__(self, solver: lineax.AbstractLinearSolver = lineax.Cholesky(), step_size: float = 1.0, tol: float = 1e-06, max_iter: int = 200, armijo_c: float = 0.1, armijo_factor: float = 0.5) ¤

Construct a Newton fitter.

Arguments:

  • solver: lineax solver for each Newton weighted normal-equations step. Defaults to lx.Cholesky(). Any lx.AbstractLinearSolver that handles symmetric positive-semidefinite systems works here.
  • step_size: initial trial step size for the Armijo line search. Defaults to 1.0 (pure Newton step). Shrunk geometrically by armijo_factor until sufficient decrease is satisfied.
  • tol: convergence tolerance on the objective change between consecutive iterations. Defaults to 1e-6.
  • max_iter: maximum number of Newton iterations. Defaults to 200.
  • armijo_c: sufficient-decrease constant \(c\) in the Armijo condition \(f(\beta - s\Delta\beta) \leq f(\beta) - cs \cdot \nabla f^\top \Delta\beta\). Defaults to 0.1.
  • armijo_factor: geometric backtracking multiplier. Each rejected trial step is shrunk by this factor. Defaults to 0.5.

Linear solvers¤

Each fitter uses a lineax solver for its internal linear system. Pass any lineax.AbstractLinearSolver as solver= when constructing a fitter. lineax.Cholesky() is the default and the fastest option for well-conditioned systems. Use lineax.QR() when the design matrix is rank-deficient.