# SPDX-License-Identifier: MIT
# Copyright (c) 2025 Lucy Eduardo Arias
"""Tweedie distribution family for zero-inflated continuous data.
The Tweedie distribution is a member of the exponential dispersion family
that naturally handles data with exact zeros and continuous positive values.
It is particularly useful for insurance claims, rainfall, and healthcare costs.
References
----------
.. [1] Jørgensen, B. (1987). "Exponential dispersion models."
Journal of the Royal Statistical Society: Series B, 49(2), 127-162.
.. [2] Smyth, G. K., & Jørgensen, B. (2002).
"Fitting Tweedie's compound Poisson model to insurance claims data."
ASTIN Bulletin, 32(1), 143-157.
.. [3] Dunn, P. K., & Smyth, G. K. (2005).
"Series evaluation of Tweedie exponential dispersion model densities."
Statistics and Computing, 15(4), 267-280.
.. [4] Dunn, P. K., & Smyth, G. K. (2008).
"Evaluation of Tweedie exponential dispersion model densities by Fourier inversion."
Statistics and Computing, 18(1), 73-86.
"""
from __future__ import annotations
from typing import TYPE_CHECKING, cast
import numpy as np
from aurora.distributions._utils import (
as_namespace_array,
namespace,
)
from aurora.distributions.base import Family, LinkFunction
from aurora.distributions.links import IdentityLink, LogLink, PowerLink
if TYPE_CHECKING:
from numpy.typing import NDArray
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 TweedieFamily(Family):
"""Tweedie distribution for zero-inflated positive continuous data.
A compound Poisson-Gamma distribution that naturally handles data with
a point mass at zero (no events) and continuous positive values (sum of
event magnitudes). The variance function is Var(Y) = phi * mu^p, where
the power parameter p controls the shape.
Special cases of the Tweedie family (for reference):
- p = 0: Gaussian
- p = 1: Poisson
- 1 < p < 2: Compound Poisson-Gamma (Tweedie proper)
- p = 2: Gamma
- p = 3: Inverse Gaussian
This class implements only the compound Poisson-Gamma case 1 < p < 2;
for the boundary cases use the dedicated families (GaussianFamily,
PoissonFamily, GammaFamily, InverseGaussianFamily).
Parameters
----------
power : float, default=1.5
Variance power parameter p. For compound Poisson-Gamma, must be
in (1, 2). Common choices:
- p = 1.5: Balanced, works well for many applications
- p = 1.6-1.7: More right-skewed, common for insurance claims
- p closer to 1: More mass at zero
- p closer to 2: Less mass at zero, more Gamma-like
link : str, default='log'
Link function: 'log' (default), 'identity', or 'power'.
phi : float, default=1.0
Dispersion parameter (scale).
Attributes
----------
power : float
Variance power parameter in (1, 2).
phi : float
Dispersion (scale) parameter.
default_link : LinkFunction
The link function used to map the mean to the linear predictor.
See Also
--------
aurora.distributions.families.CompoundPoissonGammaFamily : Explicit compound
Poisson-Gamma alias for TweedieFamily
aurora.distributions.families.GammaFamily : Tweedie with power=2
aurora.distributions.families.InverseGaussianFamily : Tweedie with power=3
aurora.models.glm.fit_glm : Fit a GLM with this family
Notes
-----
The Tweedie distribution can be viewed as:
- N ~ Poisson(lambda) events occur
- Each event has magnitude X_i ~ Gamma(alpha, beta)
- Total Y = X_1 + ... + X_N
This naturally creates exact zeros (N=0) and continuous positives (N>0).
The probability of zero is: P(Y=0) = exp(-lambda)
Examples
--------
>>> from aurora.models.glm import fit_glm
>>> from aurora.distributions.families import TweedieFamily
>>> # Insurance claims (many zeros, heavy right tail)
>>> result = fit_glm(X, claims, family=TweedieFamily(power=1.7))
>>> # Rainfall amounts (zeros for dry days)
>>> result = fit_glm(X, rainfall, family=TweedieFamily(power=1.5))
References
----------
.. [1] Jorgensen (1987). Exponential dispersion models.
.. [2] Smyth & Jorgensen (2002). Fitting Tweedie's compound Poisson model.
"""
name = "tweedie"
# Valid range for responses (non-negative)
valid_y_range = (0, np.inf)
[docs]
def __init__(self, power: float = 1.5, link: str = "log", phi: float = 1.0):
"""Initialize Tweedie family.
Parameters
----------
power : float
Variance power parameter. Must be in (1, 2) for proper Tweedie.
link : str
Link function name.
phi : float
Dispersion parameter (scale).
"""
if not (1 < power < 2):
raise ValueError(
f"For compound Poisson-Gamma Tweedie, power must be in (1, 2). "
f"Got power={power}. Use Gamma (power=2) or Poisson (power=1) instead."
)
self.power = power
self.phi = phi
if link == "log":
self._link: LinkFunction = LogLink()
elif link == "identity":
self._link = IdentityLink()
elif link.startswith("power"):
# Parse power link, e.g., 'power0.5'
try:
link_power = float(link.replace("power", ""))
except ValueError:
link_power = 1 - power / 2 # Canonical power
self._link = PowerLink(link_power)
else:
raise ValueError(f"Unsupported link: {link}. Use 'log', 'identity', or 'power'")
@property
def default_link(self):
"""Default link function (log)."""
return self._link
[docs]
def variance(self, mu: NDArray, **params) -> NDArray:
"""Variance function: V(μ) = μ^p.
Parameters
----------
mu : ndarray
Mean values
**params : dict
May contain 'phi' (dispersion)
Returns
-------
ndarray
Variance at each observation (without phi factor)
"""
xp = namespace(mu)
mu_arr = as_namespace_array(mu, xp, like=mu)
# Clip mu to avoid log(0)
if xp is torch: # type: ignore[comparison-overlap]
mu_arr = torch.clamp(mu_arr, min=1e-10)
elif xp is jnp: # type: ignore[comparison-overlap]
mu_arr = jnp.clip(mu_arr, 1e-10, None)
else:
mu_arr = np.clip(mu_arr, 1e-10, None)
return cast("NDArray", mu_arr**self.power)
[docs]
def initialize(self, y: NDArray) -> NDArray:
"""Initialize mean using positive values.
Parameters
----------
y : ndarray
Response values
Returns
-------
ndarray
Initial mean estimates
"""
xp = namespace(y)
y_arr = as_namespace_array(y, xp, like=y)
# Use mean of positive values, or small value if all zeros
# Convert to numpy for indexing, then determine mu_init
if xp is torch: # type: ignore[comparison-overlap]
y_np = y_arr.cpu().numpy() if y_arr.is_cuda else y_arr.numpy()
elif xp is jnp: # type: ignore[comparison-overlap]
y_np = np.array(y_arr)
else:
y_np = np.asarray(y_arr, dtype=float)
positive_y = y_np[y_np > 0]
if len(positive_y) > 0:
mu_init = float(np.mean(positive_y))
else:
mu_init = 0.1
# Return with correct backend, preserving the input dtype
if xp is torch: # type: ignore[comparison-overlap]
return cast("NDArray", torch.full_like(y_arr, mu_init))
elif xp is jnp: # type: ignore[comparison-overlap]
return cast("NDArray", jnp.full_like(y_arr, mu_init))
else:
return np.full_like(y_arr, mu_init, dtype=float)
[docs]
def deviance(self, y: NDArray, mu: NDArray, **params) -> float:
"""Deviance for Tweedie distribution.
The unit deviance for Tweedie with power p is:
For p ∈ (1, 2):
d(y, μ) = 2 × [y^(2-p)/((1-p)(2-p)) - y×μ^(1-p)/(1-p) + μ^(2-p)/(2-p)]
For y = 0:
d(0, μ) = 2 × μ^(2-p) / (2-p)
Parameters
----------
y : ndarray
Observed values (including zeros)
mu : ndarray
Fitted mean values
**params : dict
Additional parameters
Returns
-------
float
Total deviance
"""
xp = namespace(y, mu)
y_arr = as_namespace_array(y, xp, like=mu)
mu_arr = as_namespace_array(mu, xp, like=y_arr)
# Clip mu to avoid issues
if xp is torch: # type: ignore[comparison-overlap]
mu_arr = torch.clamp(mu_arr, min=1e-10)
elif xp is jnp: # type: ignore[comparison-overlap]
mu_arr = jnp.clip(mu_arr, 1e-10, None)
else:
mu_arr = np.clip(mu_arr, 1e-10, None)
p = self.power
# Unit deviance
# For y > 0
if xp is torch: # type: ignore[comparison-overlap]
d_pos = 2 * (
y_arr ** (2 - p) / ((1 - p) * (2 - p))
- y_arr * mu_arr ** (1 - p) / (1 - p)
+ mu_arr ** (2 - p) / (2 - p)
)
elif xp is jnp: # type: ignore[comparison-overlap]
d_pos = 2 * (
y_arr ** (2 - p) / ((1 - p) * (2 - p))
- y_arr * mu_arr ** (1 - p) / (1 - p)
+ mu_arr ** (2 - p) / (2 - p)
)
else:
with np.errstate(divide="ignore", invalid="ignore"):
d_pos = 2 * (
y_arr ** (2 - p) / ((1 - p) * (2 - p))
- y_arr * mu_arr ** (1 - p) / (1 - p)
+ mu_arr ** (2 - p) / (2 - p)
)
# For y = 0
d_zero = 2 * mu_arr ** (2 - p) / (2 - p)
# Select based on y
if xp is torch: # type: ignore[comparison-overlap]
d = torch.where(y_arr == 0, d_zero, d_pos)
elif xp is jnp: # type: ignore[comparison-overlap]
d = jnp.where(y_arr == 0, d_zero, d_pos)
else:
d = np.where(y_arr == 0, d_zero, d_pos)
return float(xp.sum(d))
[docs]
def log_likelihood(self, y: NDArray, mu: NDArray, **params) -> float:
"""Approximate (quasi-)log-likelihood for Tweedie.
.. warning::
The Tweedie density has no closed form. This method returns the
**quasi-likelihood** ``-D/(2φ)``, which drops the normalization
constant ``c(y, φ, p)`` of the full density (the series
evaluation of Dunn & Smyth 2005 is not implemented). As a
consequence, the returned value — and any AIC/BIC derived from
it — is **not comparable across families, nor across different
values of the power parameter p or the dispersion φ**. It is
only meaningful for comparing models that share the same family,
power and dispersion (e.g. nested mean models).
Parameters
----------
y : ndarray
Observed values
mu : ndarray
Mean values
**params : dict
May contain 'phi'
Returns
-------
float
Approximate (quasi-)log-likelihood
"""
phi = params.get("phi", self.phi)
# Quasi-likelihood approximation
deviance = self.deviance(y, mu, **params)
# Approximate log-likelihood (ignoring constants)
return float(-0.5 * deviance / phi)
[docs]
def weights(self, y: NDArray, mu: NDArray, **params) -> NDArray:
"""IRLS weights for Tweedie.
Weights are 1/V(μ) = 1/μ^p.
Parameters
----------
y : ndarray
Observed values (not used directly)
mu : ndarray
Current mean estimates
Returns
-------
ndarray
IRLS weights
"""
xp = namespace(mu)
mu_arr = as_namespace_array(mu, xp, like=mu)
if xp is torch: # type: ignore[comparison-overlap]
mu_arr = torch.clamp(mu_arr, min=1e-10)
elif xp is jnp: # type: ignore[comparison-overlap]
mu_arr = jnp.clip(mu_arr, 1e-10, None)
else:
mu_arr = np.maximum(mu_arr, 1e-10)
return 1.0 / self.variance(mu_arr)
[docs]
def estimate_power(
self,
y: NDArray,
mu: NDArray,
power_range: tuple[float, float] = (1.1, 1.9),
n_grid: int = 20,
) -> float:
"""Estimate the power parameter via a deviance grid search.
.. note::
This is a **heuristic**: it minimizes the unit deviance over a
grid of powers. It is not a profile likelihood in p — the full
Tweedie density (Dunn & Smyth 2005 series evaluation) is not
implemented, and the quasi-likelihood ignores the normalization
constant that depends on p, so the criterion is not a valid
profile likelihood for the power parameter. Treat the result as
a rough starting value, not an MLE.
The family's ``power`` attribute is always restored to its original
value before returning.
Parameters
----------
y : ndarray
Observed values
mu : ndarray
Fitted means (from current model)
power_range : tuple
Range of powers to search
n_grid : int
Number of grid points
Returns
-------
float
Estimated power parameter
"""
y = np.asarray(y, dtype=float)
mu = np.maximum(mu, 1e-10)
powers = np.linspace(power_range[0], power_range[1], n_grid)
deviances = []
original_power = self.power
try:
for p in powers:
self.power = p
dev = self.deviance(y, mu)
deviances.append(dev)
finally:
self.power = original_power
# Find minimum
best_idx = np.argmin(deviances)
best_power = powers[best_idx]
return float(best_power)
[docs]
def estimate_phi(self, y: NDArray, mu: NDArray, ddof: int = 1) -> float:
"""Estimate dispersion parameter.
Uses Pearson estimator: φ = (1/n) Σ (y-μ)² / V(μ)
Parameters
----------
y : ndarray
Observed values
mu : ndarray
Fitted means
ddof : int
Degrees of freedom adjustment
Returns
-------
float
Estimated dispersion
"""
y = np.asarray(y, dtype=float)
mu = np.maximum(mu, 1e-10)
n = len(y)
# Pearson residuals
resid = (y - mu) / np.sqrt(self.variance(mu))
# Estimate phi
phi = np.sum(resid**2) / (n - ddof)
return float(max(phi, 1e-8))
[docs]
def d_log_likelihood(self, y: NDArray, mu: NDArray, **params) -> NDArray:
"""First derivative of quasi-log-likelihood w.r.t. μ.
Parameters
----------
y : ndarray
Observed values
mu : ndarray
Mean values
Returns
-------
ndarray
Gradient
"""
phi = params.get("phi", self.phi)
xp = namespace(y, mu)
y_arr = as_namespace_array(y, xp, like=mu)
mu_arr = as_namespace_array(mu, xp, like=y_arr)
# Clip mu to avoid division by zero
if xp is torch: # type: ignore[comparison-overlap]
mu_arr = torch.clamp(mu_arr, min=1e-10)
elif xp is jnp: # type: ignore[comparison-overlap]
mu_arr = jnp.clip(mu_arr, 1e-10, None)
else:
mu_arr = np.clip(mu_arr, 1e-10, None)
p = self.power
# d/dμ (-D/2φ) = (y - μ) / (φ μ^p)
grad = (y_arr - mu_arr) / (phi * mu_arr**p)
return cast("NDArray", grad)
[docs]
def d2_log_likelihood(self, y: NDArray, mu: NDArray, **params) -> NDArray:
"""Second derivative of quasi-log-likelihood w.r.t. μ.
Parameters
----------
y : ndarray
Observed values
mu : ndarray
Mean values
Returns
-------
ndarray
Negative Hessian diagonal (Fisher information)
"""
phi = params.get("phi", self.phi)
xp = namespace(mu)
mu_arr = as_namespace_array(mu, xp, like=mu)
# Clip mu to avoid division by zero
if xp is torch: # type: ignore[comparison-overlap]
mu_arr = torch.clamp(mu_arr, min=1e-10)
elif xp is jnp: # type: ignore[comparison-overlap]
mu_arr = jnp.clip(mu_arr, 1e-10, None)
else:
mu_arr = np.clip(mu_arr, 1e-10, None)
p = self.power
# Expected Fisher information: -E[d²/dμ²] = 1/(φ μ^p)
# (using expected information for stability)
info = -1.0 / (phi * mu_arr**p)
return cast("NDArray", info)
[docs]
def probability_zero(self, mu: NDArray, **params) -> NDArray:
"""Probability of observing exactly zero.
For Tweedie with 1 < p < 2:
P(Y=0) = exp(-λ)
where λ = μ^(2-p) / [φ(2-p)]
Parameters
----------
mu : ndarray
Mean values
Returns
-------
ndarray
Probability of zero at each observation
"""
phi = params.get("phi", self.phi)
xp = namespace(mu)
mu_arr = as_namespace_array(mu, xp, like=mu)
# Clip mu to avoid issues
if xp is torch: # type: ignore[comparison-overlap]
mu_arr = torch.clamp(mu_arr, min=1e-10)
elif xp is jnp: # type: ignore[comparison-overlap]
mu_arr = jnp.clip(mu_arr, 1e-10, None)
else:
mu_arr = np.clip(mu_arr, 1e-10, None)
p = self.power
# Poisson rate parameter
lambda_ = mu_arr ** (2 - p) / (phi * (2 - p))
return cast("NDArray", xp.exp(-lambda_))
[docs]
def __repr__(self) -> str:
"""String representation."""
# Get link name from class name (e.g., LogLink -> log)
link_name = self._link.__class__.__name__.replace("Link", "").lower()
return f"TweedieFamily(power={self.power:.3g}, phi={self.phi:.3g}, link='{link_name}')"
[docs]
class CompoundPoissonGammaFamily(TweedieFamily):
"""Compound Poisson-Gamma distribution, an alias for TweedieFamily.
Identical to TweedieFamily but with more descriptive naming for users
familiar with the compound Poisson-Gamma formulation.
The model assumes:
- N ~ Poisson(lambda) events occur
- Each event has magnitude G_i ~ Gamma(alpha, beta)
- Total Y = G_1 + ... + G_N (or 0 if N=0)
Parameters are related to Tweedie by:
- lambda = mu^(2-p) / [phi(2-p)] (Poisson rate)
- alpha = (2-p)/(p-1) (Gamma shape)
- beta = phi(p-1)mu^(p-1) (Gamma rate)
Parameters
----------
power : float, default=1.5
Variance power parameter p. Must be in (1, 2).
link : str, default='log'
Link function name.
phi : float, default=1.0
Dispersion parameter (scale).
Attributes
----------
power : float
Variance power parameter in (1, 2).
phi : float
Dispersion (scale) parameter.
default_link : LinkFunction
The link function used to map the mean to the linear predictor.
See Also
--------
aurora.distributions.families.TweedieFamily : Parent class with identical behavior
Examples
--------
>>> from aurora.distributions.families import CompoundPoissonGammaFamily
>>> family = CompoundPoissonGammaFamily(power=1.5, phi=1.0)
>>> mu = np.array([1.0, 2.0, 3.0])
>>> family.get_poisson_rate(mu) # lambda
array([...])
>>> family.get_gamma_shape() # alpha
1.0
"""
name = "compound_poisson_gamma"
[docs]
def get_poisson_rate(self, mu: NDArray, **params) -> NDArray:
"""Get the underlying Poisson rate λ.
Parameters
----------
mu : ndarray
Mean values
Returns
-------
ndarray
Poisson rate for event count
"""
phi = params.get("phi", self.phi)
xp = namespace(mu)
mu_arr = as_namespace_array(mu, xp, like=mu)
# Clip mu to avoid issues
if xp is torch: # type: ignore[comparison-overlap]
mu_arr = torch.clamp(mu_arr, min=1e-10)
elif xp is jnp: # type: ignore[comparison-overlap]
mu_arr = jnp.clip(mu_arr, 1e-10, None)
else:
mu_arr = np.clip(mu_arr, 1e-10, None)
p = self.power
return cast("NDArray", mu_arr ** (2 - p) / (phi * (2 - p)))
[docs]
def get_gamma_shape(self) -> float:
"""Get the Gamma shape parameter α."""
return (2 - self.power) / (self.power - 1)
[docs]
def get_gamma_rate(self, mu: NDArray, **params) -> NDArray:
"""Get the Gamma rate parameter β.
Parameters
----------
mu : ndarray
Mean values
Returns
-------
ndarray
Gamma rate for event magnitude
"""
phi = params.get("phi", self.phi)
xp = namespace(mu)
mu_arr = as_namespace_array(mu, xp, like=mu)
# Clip mu to avoid issues
if xp is torch: # type: ignore[comparison-overlap]
mu_arr = torch.clamp(mu_arr, min=1e-10)
elif xp is jnp: # type: ignore[comparison-overlap]
mu_arr = jnp.clip(mu_arr, 1e-10, None)
else:
mu_arr = np.clip(mu_arr, 1e-10, None)
p = self.power
return cast("NDArray", phi * (p - 1) * mu_arr ** (p - 1))
__all__ = ["TweedieFamily", "CompoundPoissonGammaFamily"]