Skip to content

Response families¤

jaxqtl implements generalized linear model (GLM) families as exponential families with an associated link function.

jaxqtl.distribution.ExponentialFamily

jaxqtl.distribution.ExponentialFamily(equinox.Module) [source] ¤

Base interface for one-parameter exponential distribution and their GLM link. A natural exponential family has density \(f(y \mid \theta, \phi) = \exp\left((y \theta - b(\theta))/\phi + c(y, \phi)\right)\), where \(\theta\) is the natural parameter, \(\phi\) is the dispersion/scale, \(b(\theta)\) is the cumulant (log-partition) function, and \(c(y, \phi)\) is the log base measure.

The exponential dispersion model uses mean \(\mu = b'(\theta)\) and variance function \(V(\mu)\) with link mapping \(g: \mu \mapsto \eta\). Subclasses specify \(b(\cdot)\), \(V(\cdot)\), and the link. For most cases, \(V(\mu) := \phi b''(\theta)\); however, Negative Binomial models use \(V(\mu) := \mu + \alpha \mu^2\) with overdispersion \(\alpha\).

Info

Not all links are valid depending on the concrete class; this is checked automatically with a ValueError if invalid.

__init__(self) -> None ¤

Initialize self. See help(type(self)) for accurate signature.

scale(self, X: ArrayLike, y: ArrayLike, mu: ArrayLike) -> jax.Array ¤

Compute a dispersion/scale parameter (i.e. \(\phi\)) given predictors and the mean.

Arguments:

  • X: Design matrix.
  • y: Observed response.
  • mu: Mean parameter for each observation.

Returns:

Dispersion estimate.

negloglikelihood(self, X: ArrayLike, y: ArrayLike, eta: ArrayLike, disp: ScalarLike) -> jax.Array ¤

Compute the negative log-likelihood at a given linear predictor eta and dispersion disp.

Arguments:

  • X: Design matrix.
  • y: Observed response.
  • eta: Linear predictor.
  • disp: Dispersion/scale parameter.

Returns:

Negative log-likelihood.

variance(self, mu: ArrayLike, disp: ScalarLike = 1.0) -> jax.Array ¤

Return the variance as a function of the mean and dispersion.

Arguments:

  • mu: Mean parameter.
  • disp: Dispersion/scale parameter.

Returns:

Variance for each observation.

sample(self, key, eta: ArrayLike, disp: ScalarLike = 1.0) -> jax.Array ¤

Draw a sample given a linear predictor and dispersion.

Arguments:

  • key: JAX PRNG key.
  • eta: Linear predictor.
  • disp: Dispersion/scale parameter.

Returns:

Simulated observations.

calc_weight(self, eta: ArrayLike, disp: ScalarLike = 0.0) -> tuple[jax.Array, jax.Array, jax.Array] ¤

Compute mean, link derivative, and IRLS weights for observations.

Arguments:

  • eta: Linear predictor.
  • disp: Dispersion/scale parameter.

Returns:

Tuple of (mu, link derivative, weights).

init_eta(self, y: ArrayLike) -> jax.Array ¤

Provide a heuristic initializer for the linear predictor.

Arguments:

  • y: Observed response.

Returns:

Initial linear predictor.

update_dispersion(self, X: ArrayLike, y: ArrayLike, eta: ArrayLike, disp: ScalarLike = 1.0, step_size: ScalarLike = 1.0) -> jax.Array ¤

Perform one dispersion update step. If not implemented, defaults to disp argument.

Arguments:

  • X: Design matrix.
  • y: Observed response.
  • eta: Linear predictor.
  • disp: Current dispersion estimate.
  • step_size: Update step size.

Returns:

Updated dispersion estimate.

estimate_dispersion(self, X: ArrayLike, y: ArrayLike, eta: ArrayLike, disp: ScalarLike = 1.0, step_size: ScalarLike = 1.0, tol: ScalarLike = 0.001, max_iter: int = 1000) -> jax.Array ¤

Iteratively estimate dispersion. If not implemented, defaults to disp argument

Arguments:

  • X: Design matrix.
  • y: Observed response.
  • eta: Linear predictor.
  • disp: Initial dispersion estimate.
  • step_size: Update step size.
  • tol: Convergence tolerance.
  • max_iter: Maximum iterations.

Returns:

Estimated dispersion.

Continuous families¤

jaxqtl.distribution.Gaussian(jaxqtl.distribution.ExponentialFamily) [source] ¤

Normal exponential dispersion model with density \(f(y \mid \mu, \phi) = (2\pi \phi)^{-1/2}\exp(-(y-\mu)^2/(2\phi))\). Dispersion \(\phi > 0\) equals the variance, and the mean \(\mu\) lies in \(\mathbb{R}\).

Info

Valid links: jaxqtl.distribution.IdentityLink, jaxqtl.distribution.LogLink, jaxqtl.distribution.PowerLink.

__init__(self, glink: AbstractLink = IdentityLink()) ¤

Arguments:

jaxqtl.distribution.Gamma(jaxqtl.distribution.ExponentialFamily) [source] ¤

Gamma exponential dispersion model with density \(f(y \mid \mu, \phi) = y^{1/\phi-1}\exp(-y/(\mu\phi))/(\Gamma(1/\phi)(\mu\phi)^{1/\phi})\). Dispersion \(\phi > 0\) scales the variance \(\phi \mu^2\), and the mean \(\mu\) lies in \(\mathbb{R}_{+}\).

Info

Valid links: jaxqtl.distribution.IdentityLink, jaxqtl.distribution.InverseLink, jaxqtl.distribution.LogLink.

__init__(self, glink: AbstractLink = InverseLink()) ¤

Arguments:

Discrete families¤

jaxqtl.distribution.Poisson(jaxqtl.distribution.ExponentialFamily) [source] ¤

Poisson exponential family with density \(f(y \mid \mu) = \exp(-\mu) \mu^{y}/y!\) and unit dispersion. The mean \(\mu\) lies in \(\mathbb{R}_{+}\); dispersion is fixed at 1.

Info

Valid links: jaxqtl.distribution.IdentityLink, jaxqtl.distribution.LogLink.

__init__(self, glink: AbstractLink = LogLink()) ¤

Arguments:

jaxqtl.distribution.NegativeBinomial(jaxqtl.distribution.ExponentialFamily) [source] ¤

NB2 parameterization with dispersion \(\alpha\) (variance \(\mu + \alpha \mu^2\)) and density \(f(y \mid \mu, \alpha) = \frac{\Gamma(y+r)}{\Gamma(r)\,y!}\left(\frac{r}{r+\mu}\right)^r \left(\frac{\mu}{r+\mu}\right)^y\) where \(r = 1/\alpha\) for \(\alpha > 0\). The dispersion satisfies \(\alpha \geq 0\), the mean \(\mu\) lies in \(\mathbb{R}_{+}\), and \(\alpha = 0\) evaluates the Poisson limit.

Info

Valid links: jaxqtl.distribution.IdentityLink, jaxqtl.distribution.LogLink, jaxqtl.distribution.NBLink, jaxqtl.distribution.PowerLink.

__init__(self, glink: AbstractLink = LogLink()) ¤

Arguments:

update_dispersion(self, X: ArrayLike, y: ArrayLike, eta: ArrayLike, disp: ScalarLike = 0.1, step_size: ScalarLike = 0.1) -> jax.Array ¤

Take one variance-manifold modified-Newton step for NB2 dispersion.

Arguments:

  • X: Design matrix, retained for the common family API but unused by this update.
  • y: Response vector.
  • eta: Linear predictor used to hold the fitted means fixed during the dispersion step.
  • disp: Current NB2 dispersion $alpha$.
  • step_size: Scalar multiplier for the tangent direction.

Returns:

The updated dispersion, clipped to $[10^{-9}, 10^9]$.

Failure Modes:

A nonpositive or nonfinite Riemannian Hessian uses the variance-manifold metric as a natural-gradient fallback. If that metric or the score is invalid, the clipped input dispersion is returned unchanged.

jaxqtl.distribution.Binomial(jaxqtl.distribution.ExponentialFamily) [source] ¤

Bernoulli/binomial (\(n=1\)) model with density \(f(y \mid \mu) = \mu^{y}(1-\mu)^{1-y}\) and fixed dispersion 1. The mean \(\mu\) lies in \([0, 1]\); there is no additional scale parameter beyond \(\mu\).

Info

Valid links: jaxqtl.distribution.LogitLink, jaxqtl.distribution.LogLink, jaxqtl.distribution.IdentityLink.

__init__(self, glink: AbstractLink = LogitLink()) ¤

Arguments:

Distribution utilities¤

jaxqtl.distribution.t_cdf(value: ArrayLike, df: ArrayLike, loc: ArrayLike = 0.0, scale: ArrayLike = 1.0) -> jax.Array ¤

Evaluate the cumulative distribution function of Student's t distribution.

The arguments follow JAX broadcasting rules. This implementation is written in terms of jax.scipy.special.betainc so it can be transformed by JAX.

Arguments:

  • value: Points at which to evaluate the CDF.
  • df: Positive degrees of freedom.
  • loc: Distribution location. Defaults to 0.
  • scale: Positive distribution scale. Defaults to 1.

Returns:

CDF values with the broadcasted shape of value, df, loc, and scale.

jaxqtl.distribution.ncx2_sf(x: ArrayLike, df: ScalarLike, nc: ScalarLike, terms: int = 100) ¤

Approximate the survival function of a noncentral chi-squared distribution.

The approximation truncates the Poisson-mixture representation after terms components. Larger noncentrality parameters may require more terms to capture the upper tail accurately.

Arguments:

  • x: Points at which to evaluate the survival function.
  • df: Positive degrees of freedom.
  • nc: Non-negative noncentrality parameter.
  • terms: Number of Poisson-mixture components. Defaults to 100.

Returns:

Approximate survival probabilities with the same shape as x.