# SPDX-License-Identifier: MIT
# Copyright (c) 2025 Lucy Eduardo Arias
"""Common link function implementations."""
from __future__ import annotations
import math
import numpy as np
from .._utils import (
as_namespace_array,
clip_probability,
ensure_positive,
namespace,
ones_like,
)
from ..base import LinkFunction
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 IdentityLink(LinkFunction):
"""Identity link function: g(mu) = mu.
The simplest link where the linear predictor equals the response directly.
This is the canonical link for the Gaussian family.
Examples
--------
>>> import numpy as np
>>> from aurora.distributions.links import IdentityLink
>>> link = IdentityLink()
>>> mu = np.array([1.0, 2.0, 3.0])
>>> eta = link.link(mu)
>>> np.allclose(eta, mu)
True
>>> mu_back = link.inverse(eta)
>>> np.allclose(mu_back, mu)
True
See Also
--------
LogLink : Log link for positive data.
GaussianFamily : Default family for identity link.
"""
[docs]
def link(self, mu): # noqa: ANN001 - signature from base class
xp = namespace(mu)
return as_namespace_array(mu, xp, like=mu)
[docs]
def inverse(self, eta): # noqa: ANN001 - signature from base class
xp = namespace(eta)
return as_namespace_array(eta, xp, like=eta)
[docs]
def derivative(self, mu): # noqa: ANN001 - signature from base class
mu_arr = as_namespace_array(mu, namespace(mu), like=mu)
return ones_like(mu_arr)
[docs]
class LogLink(LinkFunction):
"""Log link function: g(mu) = log(mu).
Ensures mu > 0. This is the canonical link for the Poisson family
and is commonly used for any strictly positive response.
Notes
-----
The inverse is g^{-1}(eta) = exp(eta). The linear predictor eta is
clamped to [-700, 700] to prevent numerical overflow.
Examples
--------
>>> import numpy as np
>>> from aurora.distributions.links import LogLink
>>> link = LogLink()
>>> mu = np.array([1.0, 2.0, 3.0])
>>> eta = link.link(mu)
>>> np.allclose(eta, np.log(mu))
True
>>> mu_back = link.inverse(eta)
>>> np.allclose(mu_back, mu)
True
See Also
--------
IdentityLink : Identity link.
PoissonFamily : Default family for log link.
GammaFamily : Often used with log link.
"""
[docs]
def link(self, mu): # noqa: ANN001 - signature from base class
xp = namespace(mu)
mu_arr = ensure_positive(as_namespace_array(mu, xp, like=mu), xp)
return xp.log(mu_arr)
[docs]
def inverse(self, eta): # noqa: ANN001 - signature from base class
xp = namespace(eta)
eta_arr = as_namespace_array(eta, xp, like=eta)
# Clamp eta to prevent overflow: log(max_float64) ≈ 709
if xp is np:
eta_clamped = np.clip(eta_arr, -700, 700)
elif hasattr(xp, "clamp"): # PyTorch
eta_clamped = xp.clamp(eta_arr, -700, 700)
else: # JAX or other
eta_clamped = xp.clip(eta_arr, -700, 700)
return xp.exp(eta_clamped)
[docs]
def derivative(self, mu): # noqa: ANN001 - signature from base class
xp = namespace(mu)
mu_arr = ensure_positive(as_namespace_array(mu, xp, like=mu), xp)
return 1.0 / mu_arr
[docs]
class LogitLink(LinkFunction):
"""Logit link function: g(mu) = log(mu / (1 - mu)).
Maps probabilities in (0, 1) to the real line. This is the canonical link
for the binomial family. The inverse is the logistic sigmoid:
g^{-1}(eta) = 1 / (1 + exp(-eta)).
Examples
--------
>>> import numpy as np
>>> from aurora.distributions.links import LogitLink
>>> link = LogitLink()
>>> mu = np.array([0.2, 0.5, 0.8])
>>> eta = link.link(mu)
>>> mu_back = link.inverse(eta)
>>> np.allclose(mu, mu_back)
True
See Also
--------
ProbitLink : Normal CDF inverse link.
CLogLogLink : Complementary log-log link.
BinomialFamily : Default family for logit link.
BetaFamily : Often used with logit link.
"""
[docs]
def link(self, mu): # noqa: ANN001 - signature from base class
xp = namespace(mu)
mu_arr = clip_probability(as_namespace_array(mu, xp, like=mu), xp)
return xp.log(mu_arr / (1.0 - mu_arr))
[docs]
def inverse(self, eta): # noqa: ANN001 - signature from base class
xp = namespace(eta)
eta_arr = as_namespace_array(eta, xp, like=eta)
# Clamp eta to prevent overflow in exp(-eta): exp(±700) stays finite
# in float64. Beyond this range the sigmoid is numerically 0 or 1.
if xp is np:
eta_arr = np.clip(eta_arr, -700, 700)
elif hasattr(xp, "clamp"): # PyTorch
eta_arr = xp.clamp(eta_arr, -700, 700)
else: # JAX or other
eta_arr = xp.clip(eta_arr, -700, 700)
return 1.0 / (1.0 + xp.exp(-eta_arr))
[docs]
def derivative(self, mu): # noqa: ANN001 - signature from base class
xp = namespace(mu)
mu_arr = clip_probability(as_namespace_array(mu, xp, like=mu), xp)
return 1.0 / (mu_arr * (1.0 - mu_arr))
[docs]
class InverseLink(LinkFunction):
"""Inverse link function: g(mu) = 1 / mu.
Ensures mu > 0. This is the canonical link for the Gamma family.
The inverse is g^{-1}(eta) = 1 / eta.
Examples
--------
>>> import numpy as np
>>> from aurora.distributions.links import InverseLink
>>> link = InverseLink()
>>> mu = np.array([0.5, 1.0, 2.0])
>>> eta = link.link(mu)
>>> np.allclose(eta, 1.0 / mu)
True
>>> mu_back = link.inverse(eta)
>>> np.allclose(mu_back, mu)
True
See Also
--------
InverseSquareLink : Inverse square link for inverse Gaussian.
LogLink : Log link (alternative for Gamma).
GammaFamily : Default family for inverse link.
"""
[docs]
def link(self, mu): # noqa: ANN001 - signature from base class
xp = namespace(mu)
mu_arr = ensure_positive(as_namespace_array(mu, xp, like=mu), xp)
return 1.0 / mu_arr
[docs]
def inverse(self, eta): # noqa: ANN001 - signature from base class
xp = namespace(eta)
eta_arr = ensure_positive(as_namespace_array(eta, xp, like=eta), xp)
return 1.0 / eta_arr
[docs]
def derivative(self, mu): # noqa: ANN001 - signature from base class
xp = namespace(mu)
mu_arr = ensure_positive(as_namespace_array(mu, xp, like=mu), xp)
return -1.0 / (mu_arr**2)
[docs]
class CLogLogLink(LinkFunction):
"""Complementary log-log link function: g(mu) = log(-log(1 - mu)).
Maps probabilities in (0, 1) to the real line. Asymmetric link
that approaches 0 more slowly than 1. Commonly used for
binomial regression when the probability of success increases
slowly at low doses but rapidly at high doses.
The inverse is g^{-1}(eta) = 1 - exp(-exp(eta)).
Examples
--------
>>> import numpy as np
>>> from aurora.distributions.links import CLogLogLink
>>> link = CLogLogLink()
>>> mu = np.array([0.2, 0.5, 0.8])
>>> eta = link.link(mu)
>>> mu_back = link.inverse(eta)
>>> np.allclose(mu, mu_back)
True
See Also
--------
LogitLink : Symmetric link for binomial.
ProbitLink : Normal CDF inverse link.
BinomialFamily : Distribution family using these links.
"""
[docs]
def link(self, mu): # noqa: ANN001 - signature from base class
xp = namespace(mu)
mu_arr = clip_probability(as_namespace_array(mu, xp, like=mu), xp)
# -log(1 - mu) via log1p to keep full precision for mu close to 1
log_term = -xp.log1p(-mu_arr)
log_term = ensure_positive(log_term, xp)
return xp.log(log_term)
[docs]
def inverse(self, eta): # noqa: ANN001 - signature from base class
xp = namespace(eta)
eta_arr = as_namespace_array(eta, xp, like=eta)
# 1 - exp(-exp(eta)) = -expm1(-exp(eta)): expm1 keeps full precision
# in the left tail (eta << 0), where 1 - exp(-exp(eta)) underflows
# to exactly 0 (e.g. eta <= -37 in float64).
return -xp.expm1(-xp.exp(eta_arr))
[docs]
def derivative(self, mu): # noqa: ANN001 - signature from base class
xp = namespace(mu)
mu_arr = clip_probability(as_namespace_array(mu, xp, like=mu), xp)
one_minus = 1.0 - mu_arr
one_minus = ensure_positive(one_minus, xp)
log_term = -xp.log(one_minus)
return 1.0 / (log_term * one_minus)
[docs]
class SqrtLink(LinkFunction):
"""Square root link function: g(mu) = sqrt(mu).
Ensures mu > 0. Useful for count data where variance is proportional
to the mean. Common alternative to the log link for Poisson-like data.
The inverse is g^{-1}(eta) = eta^2.
Examples
--------
>>> import numpy as np
>>> from aurora.distributions.links import SqrtLink
>>> link = SqrtLink()
>>> mu = np.array([1.0, 4.0, 9.0])
>>> eta = link.link(mu)
>>> np.allclose(eta, np.sqrt(mu))
True
>>> mu_back = link.inverse(eta)
>>> np.allclose(mu_back, mu)
True
See Also
--------
LogLink : Log link for count data.
IdentityLink : Identity link.
PoissonFamily : Common family for count data.
"""
[docs]
def link(self, mu): # noqa: ANN001 - signature from base class
xp = namespace(mu)
mu_arr = ensure_positive(as_namespace_array(mu, xp, like=mu), xp)
return xp.sqrt(mu_arr)
[docs]
def inverse(self, eta): # noqa: ANN001 - signature from base class
xp = namespace(eta)
eta_arr = as_namespace_array(eta, xp, like=eta)
return eta_arr**2
[docs]
def derivative(self, mu): # noqa: ANN001 - signature from base class
xp = namespace(mu)
mu_arr = ensure_positive(as_namespace_array(mu, xp, like=mu), xp)
return 0.5 / xp.sqrt(mu_arr)
[docs]
class PowerLink(LinkFunction):
"""Power link function: g(mu) = mu^power.
General power transformation that generalizes several other link
functions as special cases:
- power = 1: Identity link
- power = 0: Log link (limit as power approaches 0)
- power = -1: Inverse link
- power = 0.5: Square root link
- power = -2: Inverse square link
Parameters
----------
power : float, default=1.0
Power parameter for the transformation.
Notes
-----
The Box-Cox transformation is a special case when properly normalized.
For power = 0, this class uses the log link as the limit.
Examples
--------
>>> import numpy as np
>>> from aurora.distributions.links import PowerLink
>>> # Square root link via power
>>> link = PowerLink(power=0.5)
>>> mu = np.array([1.0, 4.0, 9.0])
>>> eta = link.link(mu)
>>> np.allclose(eta, np.sqrt(mu))
True
>>> # Inverse link via power
>>> link = PowerLink(power=-1)
>>> eta = link.link(mu)
>>> np.allclose(eta, 1.0 / mu)
True
See Also
--------
IdentityLink : Power link with power=1.
LogLink : Power link with power=0 (limit).
InverseLink : Power link with power=-1.
InverseSquareLink : Power link with power=-2.
TweedieFamily : Uses canonical power link.
"""
[docs]
def __init__(self, power: float = 1.0):
"""Initialize power link.
Parameters
----------
power : float
Power parameter.
"""
self.power = power
self.name = f"power{power}"
[docs]
def link(self, mu): # noqa: ANN001 - signature from base class
xp = namespace(mu)
mu_arr = ensure_positive(as_namespace_array(mu, xp, like=mu), xp)
if abs(self.power) < 1e-10:
# Use log for power ≈ 0
return xp.log(mu_arr)
return mu_arr**self.power
[docs]
def inverse(self, eta): # noqa: ANN001 - signature from base class
xp = namespace(eta)
eta_arr = as_namespace_array(eta, xp, like=eta)
if abs(self.power) < 1e-10:
# Use exp for power ≈ 0
return xp.exp(eta_arr)
# eta = mu^power with mu > 0 is always positive, for any nonzero
# power. Guard against negative/NaN-producing input for both signs
# of power (previously power < 0 with eta < 0 returned NaN or a
# negative — invalid — mean).
eta_arr = ensure_positive(eta_arr, xp)
return eta_arr ** (1.0 / self.power)
[docs]
def derivative(self, mu): # noqa: ANN001 - signature from base class
xp = namespace(mu)
mu_arr = ensure_positive(as_namespace_array(mu, xp, like=mu), xp)
if abs(self.power) < 1e-10:
# Log link derivative: 1/μ
return 1.0 / mu_arr
return self.power * mu_arr ** (self.power - 1)
[docs]
class InverseSquareLink(LinkFunction):
"""Inverse square link function: g(mu) = 1 / mu^2.
Ensures mu > 0. This is the canonical link for the inverse Gaussian
distribution. The inverse is g^{-1}(eta) = 1 / sqrt(eta).
Examples
--------
>>> import numpy as np
>>> from aurora.distributions.links import InverseSquareLink
>>> link = InverseSquareLink()
>>> mu = np.array([1.0, 2.0, 3.0])
>>> eta = link.link(mu)
>>> np.allclose(eta, 1.0 / mu**2)
True
>>> mu_back = link.inverse(eta)
>>> np.allclose(mu_back, mu)
True
See Also
--------
InverseLink : Inverse link g(mu) = 1/mu.
InverseGaussianFamily : Default family for inverse square link.
PowerLink : General power link (power=-2 is equivalent).
"""
[docs]
def link(self, mu): # noqa: ANN001 - signature from base class
xp = namespace(mu)
mu_arr = ensure_positive(as_namespace_array(mu, xp, like=mu), xp)
return 1.0 / (mu_arr**2)
[docs]
def inverse(self, eta): # noqa: ANN001 - signature from base class
xp = namespace(eta)
eta_arr = ensure_positive(as_namespace_array(eta, xp, like=eta), xp)
return 1.0 / xp.sqrt(eta_arr)
[docs]
def derivative(self, mu): # noqa: ANN001 - signature from base class
xp = namespace(mu)
mu_arr = ensure_positive(as_namespace_array(mu, xp, like=mu), xp)
return -2.0 / (mu_arr**3)
[docs]
class ProbitLink(LinkFunction):
"""Probit link ``g(mu) = Φ^{-1}(mu)``.
The probit link uses the inverse cumulative distribution function (CDF)
of the standard normal distribution. Common alternative to logit for
binomial and beta regression models.
Properties:
- Lighter tails than logit
- Assumes underlying normally distributed latent variable
- Results similar to logit in practice for μ ∈ [0.2, 0.8]
Mathematical details:
- Link: g(μ) = Φ^{-1}(μ) where Φ is the standard normal CDF
- Inverse: μ = Φ(η)
- Derivative: dg/dμ = 1/φ(Φ^{-1}(μ)) where φ is the normal PDF
Examples
--------
>>> from aurora.distributions.links import ProbitLink
>>> link = ProbitLink()
>>> import numpy as np
>>> mu = np.array([0.1, 0.5, 0.9])
>>> eta = link.link(mu) # Transform to linear predictor
>>> mu_back = link.inverse(eta) # Should equal mu
>>> np.allclose(mu, mu_back)
True
Notes
-----
The probit link is preferred when:
- There is a theoretical latent normal process
- Lighter tails than logit are desired
- Compatibility with other software using probit (e.g., econometrics)
Comparison with logit:
- Both are symmetric around 0.5
- Logit has heavier tails (more robust to outliers)
- Probit: π × logit(μ) / √3 is a good approximation
References
----------
- Bliss, C. I. (1934). "The method of probits." Science, 79, 38-39.
- McCullagh, P., & Nelder, J. A. (1989). Generalized Linear Models.
"""
[docs]
def link(self, mu): # noqa: ANN001 - signature from base class
"""Transform probability to linear predictor: η = Φ^{-1}(μ)."""
xp = namespace(mu)
mu_arr = clip_probability(as_namespace_array(mu, xp, like=mu), xp)
# Φ^{-1}(μ) - inverse normal CDF, native to each backend
if xp is torch: # type: ignore[comparison-overlap]
return torch.special.ndtri(mu_arr)
elif xp is jnp: # type: ignore[comparison-overlap]
from jax.scipy.special import ndtri
return ndtri(mu_arr)
from scipy.stats import norm
return norm.ppf(mu_arr)
[docs]
def inverse(self, eta): # noqa: ANN001 - signature from base class
"""Transform linear predictor to probability: μ = Φ(η).
No clamping is applied: Φ(η) is evaluable over the whole float64
range (it saturates to 0 or 1 only around |η| ≈ 38). Families are
responsible for clipping probabilities away from the boundaries.
"""
xp = namespace(eta)
eta_arr = as_namespace_array(eta, xp, like=eta)
if xp is torch: # type: ignore[comparison-overlap]
return torch.special.ndtr(eta_arr)
elif xp is jnp: # type: ignore[comparison-overlap]
from jax.scipy.special import ndtr
return ndtr(eta_arr)
from scipy.stats import norm
return norm.cdf(eta_arr)
[docs]
def derivative(self, mu): # noqa: ANN001 - signature from base class
"""Compute derivative: dg/dμ = 1/φ(Φ^{-1}(μ)).
The derivative is the reciprocal of the normal PDF evaluated
at the quantile corresponding to μ.
"""
xp = namespace(mu)
mu_arr = clip_probability(as_namespace_array(mu, xp, like=mu), xp)
# z = Φ^{-1}(μ); φ(z) = exp(-z²/2) / √(2π), all backend-native
z = self.link(mu_arr)
pdf_z = xp.exp(-0.5 * z**2) / math.sqrt(2.0 * math.pi)
pdf_z = ensure_positive(pdf_z, xp, eps=1e-10)
return 1.0 / pdf_z
__all__ = [
"IdentityLink",
"LogLink",
"LogitLink",
"InverseLink",
"CLogLogLink",
"SqrtLink",
"PowerLink",
"InverseSquareLink",
"ProbitLink",
]