Skip to content

Linear and generalized linear models¤

jaxqtl implements linear models and generalized linear models (GLMs) used during variant-level association testing.

jaxqtl.infer.AbstractLinearModel

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

Abstract base class for linear and generalized linear models.

Attributes:

  • family: Response distribution and link function.
  • solver: Linear solver used for least-squares subproblems.
fit(self, X: ArrayLike, y: ArrayLike, offset: ArrayLike = 0.0, std_err: AbstractVarianceEstimator = FisherInfoError()) -> ModelResult ¤

Fit a model and return a summary state.

Arguments:

  • X: Design matrix with shape (n, p).
  • y: Response vector with shape (n,).
  • offset: Offset broadcastable to y (either scalar or (n,)).
  • std_err: Coefficient covariance estimator implementing jaxqtl.infer.AbstractVarianceEstimator.

Returns:

A jaxqtl.infer.ModelResult containing fitted coefficients, standard errors, and auxiliary quantities.

jaxqtl.infer.LinearModel(jaxqtl.infer.AbstractLinearModel) [source] ¤

Gaussian linear regression with a fast least-squares implementation.

This model requires an identity link and avoids the full IRLS loop.

Attributes:

  • family: Gaussian response family with an identity link.
  • solver: Linear solver used for the least-squares fit. Defaults to jaxqtl.infer.CholeskySolve.

Raises:

  • ValueError: If family is not Gaussian with an identity link.
fit(self, X: ArrayLike, y: ArrayLike, offset: ArrayLike = 0.0, std_err: AbstractVarianceEstimator = FisherInfoError(), *, df_resid: int | None = None) -> ModelResult ¤

Fit a Gaussian linear model and return a summary state.

Arguments:

  • X: Design matrix with shape (n, p).
  • y: Response vector with shape (n,).
  • offset: Offset broadcastable to y (either scalar or (n,)).
  • std_err: Coefficient covariance estimator implementing jaxqtl.infer.AbstractVarianceEstimator.
  • df_resid: Residual degrees of freedom used for both the residual-dispersion denominator and the Student's t reference distribution. Defaults to n - p; callers that fit a residualized submodel may provide the degrees of freedom from the corresponding full model.

Returns:

A jaxqtl.infer.ModelResult containing fitted coefficients, standard errors, and auxiliary quantities.

Raises:

  • ValueError: If df_resid is not positive. The default also raises when the design has no residual degrees of freedom.

jaxqtl.infer.GeneralizedLinearModel(jaxqtl.infer.AbstractLinearModel) [source] ¤

Generalized linear model (GLM) fitted via IRLS.

This class wraps a family (distribution + link) and a linear solver, and fits coefficients using iteratively reweighted least squares (IRLS).

Negative Binomial fits start from a Poisson fit and a Poisson-corrected moment estimate of dispersion. Initialization takes at most one accepted dispersion update using step_size with backtracking at fixed Poisson means; subsequent IRLS iterations backtrack coefficient and dispersion updates jointly.

Attributes:

  • family: Response distribution and link function. Defaults to jaxqtl.distribution.Gaussian.
  • solver: Linear solver for each IRLS subproblem. Defaults to jaxqtl.infer.CholeskySolve.
  • max_iter: Maximum number of IRLS iterations. Defaults to 1000.
  • tol: Absolute tolerance on the change in total negative log likelihood, which triggers a gradient check. Defaults to 1e-3.
  • step_size: Initial trial step for IRLS backtracking. Rejected trials use successively halved steps. Defaults to 1.
  • gtol: Per-observation gradient tolerance, with RMS column scaling for coefficients and bound projection for NB2 dispersion. Both the likelihood and gradient checks must pass. Defaults to 1e-3.
fit(self, X: ArrayLike, y: ArrayLike, offset: ArrayLike = 0.0, std_err: AbstractVarianceEstimator = FisherInfoError()) -> ModelResult ¤

Fit a GLM with IRLS and return a summary state.

Arguments:

  • X: Design matrix with shape (n, p).
  • y: Response vector with shape (n,).
  • offset: Offset broadcastable to y (either scalar or (n,)).
  • std_err: Coefficient covariance estimator implementing jaxqtl.infer.AbstractVarianceEstimator.

Returns:

A jaxqtl.infer.ModelResult containing fitted coefficients, standard errors, and auxiliary quantities.

Fitted result¤

jaxqtl.infer.ModelResult(builtins.tuple) [source] ¤

Container for fitted model outputs.

This stores coefficient estimates and derived quantities returned by jaxqtl.infer.LinearModel.fit and jaxqtl.infer.GeneralizedLinearModel.fit. For a design matrix with shape (n, p), coefficient-level fields have shape (p,), observation-level fields have shape (n,), and resid_covar has shape (p, p).

Attributes:

  • beta: Fitted coefficient estimates.
  • se: Coefficient standard errors.
  • z: Wald statistics, computed as beta / se.
  • p: Two-sided coefficient p-values.
  • eta: Complete linear predictor, including the supplied offset.
  • mu: Fitted response mean obtained from the inverse link.
  • glm_wt: IRLS working weights. The Gaussian fast path may return a scalar weight.
  • link_prime: Link derivative evaluated at mu. The Gaussian fast path may return a scalar value.
  • num_iters: Number of solver iterations.
  • converged: Boolean convergence indicator.
  • resid_covar: Coefficient covariance matrix used to compute se.
  • resid: Working residual used by score tests. For the Gaussian identity-link model, this is y - mu.
  • disp: Fitted family dispersion or scale parameter.

Fitting process¤

Generalized linear model fitting consists of:

  1. Selecting a response family and link.
  2. Solving weighted least-squares subproblems with the configured solver.
  3. Estimating dispersion when the family requires it.
  4. Computing a coefficient covariance matrix and Wald statistics.

NB2 initialization fits Poisson means, subtracts the Poisson variance contribution from the moment estimate of dispersion, and takes at most one accepted, backtracked dispersion step at those means. Subsequent IRLS iterations update coefficients and dispersion jointly. Each iteration computes one weighted least-squares direction; backtracking reuses that direction at successively halved step sizes.

NB2 dispersion uses a Riemannian modified-Newton direction on the variance manifold, holding the trial means fixed. If the Riemannian curvature is nonpositive or nonfinite, the update uses the variance metric for a natural-gradient direction. If that metric or the score is invalid, dispersion is unchanged. Dispersion candidates are bounded to [1e-9, 1e9]; the outer IRLS line search accepts only finite, non-increasing joint negative log likelihood. See NegativeBinomial.update_dispersion.

Convergence¤

tol bounds the absolute change in total negative log likelihood. Once this change is small, IRLS checks the gradients at the accepted coefficients and dispersion. A small likelihood change alone is not sufficient. Under vmap, JAX can evaluate both conditional branches, so batched fits may perform these checks more often.

For n observations and design-column RMS magnitudes s_j, the coefficient criterion is max_j(abs(dNLL/dbeta_j) / (n * s_j)) <= gtol. This makes the check insensitive to multiplying a covariate by a constant. It does not whiten correlated covariates or require another linear solve.

NB2 additionally requires abs(projected(dNLL/dalpha)) / n <= gtol. At the lower dispersion bound (1e-9), a nonnegative gradient satisfies the constraint; at the upper bound (1e9), a nonpositive gradient does. Interior gradients must be small in either direction. The log-alpha gradient is converted to the alpha-space gradient so that small alpha alone cannot produce apparent convergence.

Both tolerances default to 1e-3 and are exposed as CLI options --tol and --gtol. NB2's initial Poisson fit uses the same tolerances. converged=False indicates exhausted backtracking, an iteration limit, or an unchanged nonstationary state; reaching a dispersion bound is not by itself convergence.