Skip to content

Linear solvers¤

GLM fitting in jaxqtl reduces each IRLS iteration to a (weighted) least-squares solve. The solver choice can impact both speed and numerical stability.

jaxqtl.infer.AbstractLinearSolve

jaxqtl.infer.AbstractLinearSolve(equinox.Module) [source] ¤

Base interface for linear solvers used inside IRLS/GLM fitting.

During iteratively reweighted least squares (IRLS), each iteration reduces to a (weighted) least-squares solve. Implementations provide:

  • wgt_lstsq: weighted least squares.
  • lstsq: unweighted least squares.

These are used by jaxqtl.infer.irls and jaxqtl.infer.lstsq.

wgt_lstsq(self, X: jax.Array, r: jax.Array, weights: jax.Array) -> jax.Array ¤

Solve a weighted least-squares problem.

This computes \(\hat\beta\) solving: \(\hat\beta = \arg\min_\beta \sum_{i=1}^n w_i (r_i - x_i^\top \beta)^2\).

Arguments:

  • X: Design matrix \(X\) with shape (n, p).
  • r: Working response vector \(r\) with shape (n,).
  • weights: Non-negative weights \(w\) with shape (n,).

Returns:

Coefficient vector \(\hat\beta\) with shape (p,).

lstsq(self, X: jax.Array, r: jax.Array) -> jax.Array ¤

Solve an unweighted least-squares problem.

This computes \(\hat\beta\) solving: \(\hat\beta = \arg\min_\beta \lVert r - X\beta \rVert_2^2\).

Arguments:

  • X: Design matrix \(X\) with shape (n, p).
  • r: Response vector \(r\) with shape (n,).

Returns:

Coefficient vector \(\hat\beta\) with shape (p,).

jaxqtl.infer.QRSolve(jaxqtl.infer.AbstractLinearSolve) [source] ¤

Solve least-squares problems using a QR decomposition.

QR-based solvers are typically more numerically stable than normal-equation solvers for ill-conditioned designs.

jaxqtl.infer.CholeskySolve(jaxqtl.infer.AbstractLinearSolve) [source] ¤

Solve least-squares problems via normal equations and a Cholesky factorization.

This forms \((X^\top W X)\beta = X^\top W r\) (or \((X^\top X)\beta = X^\top r\)) and solves using a Cholesky factorization. It is often fast for well-conditioned designs.

jaxqtl.infer.CGSolve(jaxqtl.infer.AbstractLinearSolve) [source] ¤

Solve least-squares problems using conjugate gradients (via lineax).

This is useful when forming \((X^\top W X)\) is expensive. The returned solution may be approximate depending on the stopping criteria used by the CG solver.