# SPDX-License-Identifier: MIT
# Copyright (c) 2025 Lucy Eduardo Arias
"""Inverse Gaussian (Wald) distribution family implementation.
The Inverse Gaussian distribution is suitable for modeling positive continuous
data with right skew, such as failure times, reaction times, or durations.
Mathematical Framework
----------------------
The Inverse Gaussian distribution has density:
**Density**:
f(y; μ, λ) = √(λ/(2πy³)) × exp{-λ(y-μ)²/(2μ²y)}
where:
- y > 0 is the response
- μ > 0 is the mean
- λ > 0 is the shape (dispersion) parameter
**Properties**:
- E[Y] = μ
- Var(Y) = μ³/λ
- Also known as the Wald distribution
- Related to first passage time of Brownian motion
**Variance function**:
V(μ) = μ³
**Canonical link**: Inverse square g(μ) = 1/μ²
**Common alternatives**: Log, Inverse, Identity
References
----------
- Tweedie, M. C. K. (1957). "Statistical properties of inverse Gaussian
distributions." The Annals of Mathematical Statistics, 28(2), 362-377.
- Chhikara, R. S., & Folks, J. L. (1989). The Inverse Gaussian Distribution.
Marcel Dekker.
- McCullagh, P., & Nelder, J. A. (1989). Generalized Linear Models (2nd ed.).
Chapman and Hall/CRC.
See Also
--------
aurora.distributions.links.common.InverseSquareLink : Canonical link
aurora.distributions.links.common.LogLink : Common alternative
aurora.distributions.links.common.InverseLink : Alternative link
"""
from __future__ import annotations
import warnings
import numpy as np
from .._utils import as_namespace_array, ensure_positive, namespace
from ..base import Family, LinkFunction
from ..links import InverseSquareLink
try: # pragma: no cover - optional dependency
import torch
except ImportError: # pragma: no cover - optional dependency
torch = None # type: ignore[assignment]
try: # pragma: no cover - optional dependency
import jax.numpy as jnp
except ImportError: # pragma: no cover - optional dependency
jnp = None # type: ignore[assignment]
[docs]
class InverseGaussianFamily(Family):
"""Inverse Gaussian (Wald) distribution for positive continuous data.
Suitable for modeling time-to-event data (failure times), reaction times,
right-skewed positive continuous data, and insurance claim amounts. The
distribution arises naturally as the first passage time of Brownian motion
with positive drift.
Parameters
----------
lambda_ : float or 'estimate', default=1.0
Shape parameter (controls dispersion).
- lambda large: Low variance, concentrated
- lambda small: High variance, dispersed
- 'estimate': Estimate from data via method-of-moments
link : LinkFunction or None, default=None
Link function. If None, uses InverseSquareLink (canonical).
Common alternatives: LogLink, InverseLink, IdentityLink.
Attributes
----------
lambda_ : float or str
Shape parameter or 'estimate' for automatic estimation.
default_link : LinkFunction
The link function used to map the mean to the linear predictor.
See Also
--------
aurora.distributions.families.GammaFamily : Gamma for positive continuous data
aurora.distributions.families.TweedieFamily : Tweedie with power=3 is inverse Gaussian
aurora.distributions.links.InverseSquareLink : Canonical link for inverse Gaussian
Examples
--------
>>> import numpy as np
>>> from aurora.distributions.families import InverseGaussianFamily
>>> from aurora.distributions.links import LogLink
>>> # Basic usage with canonical link
>>> family = InverseGaussianFamily(lambda_=2.0)
>>> mu = np.array([1.0, 2.0, 3.0])
>>> family.variance(mu) # V(mu) = mu^3
array([ 1., 8., 27.])
>>> # With log link (more common in practice)
>>> family = InverseGaussianFamily(lambda_=2.0, link=LogLink())
Notes
-----
The variance function is V(mu) = mu^3, which means variance increases
rapidly with the mean. This makes the inverse Gaussian suitable for
data where larger values are more variable.
The canonical link (inverse square) is rarely used in practice;
the log or inverse link are more common choices.
For failure time data, the inverse Gaussian arises naturally as the
distribution of the first passage time for Brownian motion with drift.
"""
[docs]
def __init__(self, lambda_: float | str = 1.0, link: LinkFunction | None = None) -> None:
"""Initialize Inverse Gaussian family.
Parameters
----------
lambda_ : float or 'estimate', default=1.0
Shape parameter. Use 'estimate' for method-of-moments.
link : LinkFunction or None
Link function (default: InverseSquareLink).
"""
if isinstance(lambda_, str):
if lambda_ != "estimate":
raise ValueError("lambda_ must be a positive float or 'estimate'")
self._lambda: float | str = lambda_
else:
if lambda_ <= 0:
raise ValueError("lambda_ must be positive")
self._lambda = float(lambda_)
self._link = link or InverseSquareLink()
@property
def lambda_(self) -> float | str:
"""Return the shape parameter."""
return self._lambda
def _get_lambda(self, xp, like, y=None, mu=None, **params):
"""Get lambda value, estimating if needed."""
lambda_param = params.get("lambda_", self._lambda)
if lambda_param == "estimate" and y is not None and mu is not None:
# Method-of-moments estimation
lambda_param = self._estimate_lambda_mm(y, mu, xp)
elif lambda_param == "estimate":
warnings.warn(
"lambda_='estimate' requires observed/fitted values to estimate "
"the shape parameter; falling back to lambda_=1.0.",
RuntimeWarning,
stacklevel=3,
)
lambda_param = 1.0
if isinstance(lambda_param, str):
lambda_param = 1.0
return ensure_positive(as_namespace_array(lambda_param, xp, like=like), xp)
def _estimate_lambda_mm(self, y, mu, xp):
"""Estimate lambda via method-of-moments.
For Inverse Gaussian:
Var(Y) = μ³/λ
Using residuals r = y - μ:
λ ≈ E[μ³] / E[(y-μ)² × μ]
More stable: use the deviance-based estimator
λ = n / Σ[(y-μ)² / (μ²y)]
"""
if xp is torch: # type: ignore[comparison-overlap]
y_arr = y
mu_arr = mu
y.shape[0]
elif xp is jnp: # type: ignore[comparison-overlap]
y_arr = y
mu_arr = mu
y.shape[0]
else:
y_arr = np.asarray(y)
mu_arr = np.asarray(mu)
len(y_arr)
# Deviance-based estimator
# From McCullagh & Nelder: φ = (1/n) Σ[(y-μ)² / (μ²y)]
# λ = 1/φ
mu_safe = ensure_positive(mu_arr, xp)
y_safe = ensure_positive(y_arr, xp)
dev_contrib = (y_safe - mu_safe) ** 2 / (mu_safe**2 * y_safe)
if xp is torch: # type: ignore[comparison-overlap]
phi_est = torch.mean(dev_contrib)
lambda_est = 1.0 / phi_est
return max(float(lambda_est), 0.01)
elif xp is jnp: # type: ignore[comparison-overlap]
phi_est = jnp.mean(dev_contrib)
lambda_est = 1.0 / phi_est
return max(float(lambda_est), 0.01)
else:
phi_est = np.mean(dev_contrib)
lambda_est = 1.0 / phi_est
return max(lambda_est, 0.01)
[docs]
def log_likelihood(self, y, mu, **params): # noqa: ANN001 - match Family signature
"""Compute log-likelihood for Inverse Gaussian distribution.
Parameters
----------
y : array
Observed positive values.
mu : array
Mean parameter (positive).
**params : dict
Optional 'lambda_' parameter override.
Returns
-------
scalar
Sum of log-likelihoods.
Notes
-----
Log-likelihood:
ℓ = 0.5 × [log(λ) - log(2π) - 3log(y)] - λ(y-μ)² / (2μ²y)
"""
xp = namespace(y, mu)
y_arr = ensure_positive(as_namespace_array(y, xp, like=mu), xp)
mu_arr = ensure_positive(as_namespace_array(mu, xp, like=y_arr), xp)
lambda_val = self._get_lambda(xp, mu_arr, y=y_arr, mu=mu_arr, **params)
# Log-likelihood
# f(y; μ, λ) = √(λ/(2πy³)) × exp{-λ(y-μ)²/(2μ²y)}
# log f = 0.5[log λ - log(2π) - 3 log y] - λ(y-μ)²/(2μ²y)
# Use namespace-compatible constant for 2π
two_pi = as_namespace_array(2 * np.pi, xp, like=y_arr)
log_lik = 0.5 * (xp.log(lambda_val) - xp.log(two_pi) - 3 * xp.log(y_arr)) - lambda_val * (
y_arr - mu_arr
) ** 2 / (2 * mu_arr**2 * y_arr)
w_arr = params.get("weights")
if w_arr is not None:
log_lik = as_namespace_array(w_arr, xp, like=mu_arr) * log_lik
if xp is torch: # type: ignore[comparison-overlap]
return torch.sum(log_lik)
return np.sum(log_lik)
[docs]
def deviance(self, y, mu, **params): # noqa: ANN001 - match Family signature
"""Compute deviance for Inverse Gaussian distribution.
The unit deviance is:
d(y, μ) = (y - μ)² / (μ² y)
Total deviance: D = Σ d(y, μ)
Parameters
----------
y : array
Observed values.
mu : array
Fitted mean values.
**params : dict
Optional parameters (not used for deviance).
Returns
-------
scalar
Deviance value.
Notes
-----
This is the "scaled deviance" without λ. The full deviance
would be λ × D.
"""
xp = namespace(y, mu)
y_arr = ensure_positive(as_namespace_array(y, xp, like=mu), xp)
mu_arr = ensure_positive(as_namespace_array(mu, xp, like=y_arr), xp)
# Unit deviance: (y - μ)² / (μ² y)
unit_dev = (y_arr - mu_arr) ** 2 / (mu_arr**2 * y_arr)
w_arr = params.get("weights")
if w_arr is not None:
unit_dev = as_namespace_array(w_arr, xp, like=mu_arr) * unit_dev
if xp is torch: # type: ignore[comparison-overlap]
return torch.sum(unit_dev)
return np.sum(unit_dev)
[docs]
def variance(self, mu, **params): # noqa: ANN001 - match Family signature
"""Compute variance function V(μ) = μ³.
Parameters
----------
mu : array
Mean parameter (positive).
**params : dict
Optional parameters (not used).
Returns
-------
array
Variance function values.
Notes
-----
The true variance is Var(Y) = μ³/λ, but the variance function
(without dispersion) is V(μ) = μ³.
"""
xp = namespace(mu)
mu_arr = ensure_positive(as_namespace_array(mu, xp, like=mu), xp)
return mu_arr**3
[docs]
def full_variance(self, mu, **params):
"""Compute full variance: Var(Y) = μ³/λ.
Unlike variance(), this includes the shape parameter.
Parameters
----------
mu : array
Mean parameter.
**params : dict
Optional 'lambda_' parameter.
Returns
-------
array
Full variance values.
"""
xp = namespace(mu)
mu_arr = ensure_positive(as_namespace_array(mu, xp, like=mu), xp)
lambda_val = self._get_lambda(xp, mu_arr, **params)
return mu_arr**3 / lambda_val
[docs]
def initialize(self, y): # noqa: ANN001 - match Family signature
"""Initialize mean parameter from data.
Parameters
----------
y : array
Observed positive values.
Returns
-------
array
Initial mean values (sample mean, broadcast to array shape).
Notes
-----
Uses the sample mean as initial value, which is a natural
choice for positive data.
"""
xp = namespace(y)
y_arr = ensure_positive(as_namespace_array(y, xp, like=y), xp)
if xp is torch: # type: ignore[comparison-overlap]
mean_val = torch.mean(y_arr)
return torch.full_like(y_arr, mean_val)
elif xp is jnp: # type: ignore[comparison-overlap]
mean_val = jnp.mean(y_arr)
return jnp.full_like(y_arr, mean_val)
else:
mean_val = np.mean(y_arr)
return np.full_like(y_arr, mean_val, dtype=y_arr.dtype)
[docs]
def estimate_lambda(self, y, mu):
"""Estimate shape parameter from data and fitted values.
Parameters
----------
y : array
Observed values.
mu : array
Fitted mean values.
Returns
-------
float
Estimated shape parameter λ.
"""
xp = namespace(y, mu)
y_arr = ensure_positive(as_namespace_array(y, xp, like=mu), xp)
mu_arr = ensure_positive(as_namespace_array(mu, xp, like=y_arr), xp)
return self._estimate_lambda_mm(y_arr, mu_arr, xp)
@property
def default_link(self) -> LinkFunction:
"""Return the default link function (InverseSquare)."""
return self._link
# Alias for compatibility
WaldFamily = InverseGaussianFamily
__all__ = ["InverseGaussianFamily", "WaldFamily"]