# SPDX-License-Identifier: MIT
# Copyright (c) 2025 Lucy Eduardo Arias
"""Beta distribution family implementation.
The Beta distribution is suitable for modeling continuous proportions in (0, 1),
such as success rates, conversion rates, or market shares.
Mathematical Framework
----------------------
The Beta distribution with mean parameterization uses:
**Density**:
f(y; μ, φ) = Γ(φ) / [Γ(μφ)Γ((1-μ)φ)] × y^{μφ-1} × (1-y)^{(1-μ)φ-1}
where:
- y ∈ (0, 1) is the response
- μ ∈ (0, 1) is the mean
- φ > 0 is the precision parameter (larger = less dispersed)
**Properties**:
- E[Y] = μ
- Var(Y) = μ(1-μ) / (φ + 1)
- Shape parameters: α = μφ, β = (1-μ)φ
**Variance function**:
V(μ) = μ(1-μ)
**Canonical link**: Logit (though probit and cloglog also common)
References
----------
- Ferrari, S., & Cribari-Neto, F. (2004). "Beta regression for modelling rates
and proportions." Journal of Applied Statistics, 31(7), 799-815.
- Smithson, M., & Verkuilen, J. (2006). "A better lemon squeezer? Maximum-likelihood
regression with beta-distributed dependent variables." Psychological Methods, 11(1), 54.
See Also
--------
aurora.distributions.links.common.LogitLink : Canonical link function
aurora.distributions.links.common.ProbitLink : Alternative link function
"""
from __future__ import annotations
import numpy as np
from .._utils import (
as_namespace_array,
clip_probability,
ensure_positive,
log_gamma,
namespace,
)
from ..base import Family, LinkFunction
from ..links import LogitLink
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]
def _log_beta(a, b, xp):
"""Compute log(B(a, b)) = log(Γ(a)) + log(Γ(b)) - log(Γ(a+b))."""
return log_gamma(a, xp) + log_gamma(b, xp) - log_gamma(a + b, xp)
[docs]
class BetaFamily(Family):
"""Beta distribution for modeling continuous proportions in (0, 1).
Suitable for modeling success rates, conversion rates, market shares,
and any continuous outcome bounded in (0, 1). The distribution uses a
mean-precision parameterization where the precision parameter phi
controls dispersion independently of the mean.
Parameters
----------
phi : float or 'estimate', default=1.0
Precision parameter (inverse dispersion).
- phi large: Low variance, concentrated around mean
- phi small: High variance, dispersed
- 'estimate': Estimate from data via method-of-moments
link : LinkFunction or None, default=None
Link function. If None, uses LogitLink (canonical).
Also supports ProbitLink, CLogLogLink.
Attributes
----------
phi : float or str
Precision 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.BinomialFamily : Binomial for discrete proportions
aurora.distributions.links.LogitLink : Canonical link for Beta regression
aurora.distributions.links.ProbitLink : Alternative link for Beta regression
Examples
--------
>>> import numpy as np
>>> from aurora.distributions.families import BetaFamily
>>> family = BetaFamily(phi=10.0)
>>> mu = np.array([0.3, 0.5, 0.7])
>>> family.variance(mu) # μ(1-μ)/(φ+1)
array([0.01909091, 0.02272727, 0.01909091])
>>> # Method-of-moments phi estimation
>>> y = np.random.beta(3, 7, size=100) # True mean ≈ 0.3
>>> family = BetaFamily(phi='estimate')
>>> mu_init = family.initialize(y)
Notes
-----
The precision phi controls the dispersion:
- phi = 1: Uniform(0,1) when mu = 0.5
- phi -> inf: Concentration at mu (degenerate)
- phi -> 0: Bimodal distribution at 0 and 1
The variance function is V(mu) = mu(1-mu), same as binomial, but
the actual variance also depends on phi: Var(Y) = V(mu)/(phi+1).
"""
[docs]
def __init__(self, phi: float | str = 1.0, link: LinkFunction | None = None) -> None:
"""Initialize Beta family.
Parameters
----------
phi : float or 'estimate', default=1.0
Precision parameter. Use 'estimate' for method-of-moments.
link : LinkFunction or None
Link function (default: LogitLink).
"""
if isinstance(phi, str):
if phi != "estimate":
raise ValueError("phi must be a positive float or 'estimate'")
self._phi: float | str = phi
else:
if phi <= 0:
raise ValueError("phi must be positive")
self._phi = float(phi)
self._link = link or LogitLink()
@property
def phi(self) -> float | str:
"""Return the precision parameter."""
return self._phi
def _get_phi(self, xp, like, y=None, **params):
"""Get phi value, estimating if needed."""
phi_param = params.get("phi", self._phi)
if phi_param == "estimate" and y is not None:
# Method-of-moments estimation
phi_param = self._estimate_phi_mm(y, xp)
if isinstance(phi_param, str):
# Fallback if estimation not possible
phi_param = 1.0
return ensure_positive(as_namespace_array(phi_param, xp, like=like), xp)
def _estimate_phi_mm(self, y, xp):
"""Estimate phi via method-of-moments.
For Beta(α, β) with μ = α/(α+β) and φ = α+β:
Var(Y) = μ(1-μ)/(φ+1)
Solving for φ:
φ = μ(1-μ)/Var(Y) - 1
"""
if xp is torch: # type: ignore[comparison-overlap]
y_mean = torch.mean(y)
y_var = torch.var(y)
elif xp is jnp: # type: ignore[comparison-overlap]
y_mean = jnp.mean(y)
y_var = jnp.var(y)
else:
y_mean = np.mean(y)
y_var = np.var(y)
# Clamp mean away from boundaries (convert to float for computation)
mu_est = float(np.clip(float(y_mean), 0.01, 0.99))
var_est = float(max(float(y_var), 1e-10))
# phi = mu(1-mu)/var - 1
phi_est = mu_est * (1 - mu_est) / var_est - 1
# Ensure phi is positive and reasonable
return max(phi_est, 0.1)
[docs]
def log_likelihood(self, y, mu, **params): # noqa: ANN001 - match Family signature
"""Compute log-likelihood for Beta distribution.
Parameters
----------
y : array
Observed proportions in (0, 1).
mu : array
Mean parameter in (0, 1).
**params : dict
Optional 'phi' parameter override.
Returns
-------
scalar
Sum of log-likelihoods.
Notes
-----
Log-likelihood:
ℓ = log Γ(φ) - log Γ(μφ) - log Γ((1-μ)φ)
+ (μφ - 1) log(y) + ((1-μ)φ - 1) log(1-y)
"""
xp = namespace(y, mu)
eps = 1e-10
# Ensure y and mu are in valid range
y_arr = clip_probability(as_namespace_array(y, xp, like=mu), xp, eps=eps)
mu_arr = clip_probability(as_namespace_array(mu, xp, like=y_arr), xp, eps=eps)
phi = self._get_phi(xp, mu_arr, y=y_arr, **params)
# Shape parameters
alpha = mu_arr * phi
beta_param = (1.0 - mu_arr) * phi
# Ensure positive
alpha = ensure_positive(alpha, xp)
beta_param = ensure_positive(beta_param, xp)
# Log-likelihood: log f(y; α, β)
# = log Γ(α+β) - log Γ(α) - log Γ(β) + (α-1)log(y) + (β-1)log(1-y)
log_lik = (
log_gamma(alpha + beta_param, xp)
- log_gamma(alpha, xp)
- log_gamma(beta_param, xp)
+ (alpha - 1.0) * xp.log(y_arr)
+ (beta_param - 1.0) * xp.log(1.0 - y_arr)
)
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 Beta distribution.
For Beta regression, the deviance is computed as:
D = 2 × φ × Σ[y log(y/μ) + (1-y) log((1-y)/(1-μ))]
This is the weighted KL-divergence scaled by 2φ.
Parameters
----------
y : array
Observed proportions.
mu : array
Fitted mean values.
**params : dict
Optional 'phi' parameter override.
Returns
-------
scalar
Deviance value.
Notes
-----
For the Beta distribution, deviance is always non-negative.
When y = μ, the deviance is zero.
"""
xp = namespace(y, mu)
eps = 1e-10
y_arr = clip_probability(as_namespace_array(y, xp, like=mu), xp, eps=eps)
mu_arr = clip_probability(as_namespace_array(mu, xp, like=y_arr), xp, eps=eps)
phi = self._get_phi(xp, mu_arr, y=y_arr, **params)
# Unit deviance for Beta: 2 × [y log(y/μ) + (1-y) log((1-y)/(1-μ))]
# This is the KL divergence
term1 = y_arr * xp.log(y_arr / mu_arr)
term2 = (1.0 - y_arr) * xp.log((1.0 - y_arr) / (1.0 - mu_arr))
unit_dev = 2.0 * (term1 + term2)
if xp is torch: # type: ignore[comparison-overlap]
return phi * torch.sum(unit_dev)
return phi * np.sum(unit_dev)
[docs]
def variance(self, mu, **params): # noqa: ANN001 - match Family signature
"""Compute variance function V(μ) = μ(1-μ).
Note: Actual variance is V(μ)/(φ+1), but for GLM/IRLS
we return the variance function without dispersion.
Parameters
----------
mu : array
Mean parameter in (0, 1).
**params : dict
Optional 'phi' parameter override.
Returns
-------
array
Variance function values.
Notes
-----
The true variance is Var(Y) = μ(1-μ)/(φ+1), but in the GLM
framework we separate the variance function V(μ) from the
dispersion parameter. Here we return V(μ) = μ(1-μ).
"""
xp = namespace(mu)
mu_arr = clip_probability(as_namespace_array(mu, xp, like=mu), xp)
return mu_arr * (1.0 - mu_arr)
[docs]
def full_variance(self, mu, **params):
"""Compute full variance: Var(Y) = μ(1-μ)/(φ+1).
Unlike variance(), this includes the precision parameter.
Parameters
----------
mu : array
Mean parameter.
**params : dict
Optional 'phi' parameter.
Returns
-------
array
Full variance values.
"""
xp = namespace(mu)
mu_arr = clip_probability(as_namespace_array(mu, xp, like=mu), xp)
phi = self._get_phi(xp, mu_arr, **params)
return mu_arr * (1.0 - mu_arr) / (phi + 1.0)
[docs]
def initialize(self, y): # noqa: ANN001 - match Family signature
"""Initialize mean parameter from data.
Parameters
----------
y : array
Observed proportions.
Returns
-------
array
Initial mean values, clamped to (0.01, 0.99).
Notes
-----
Uses sample mean clamped away from boundaries to avoid
numerical issues with link function.
"""
xp = namespace(y)
y_arr = as_namespace_array(y, xp, like=y)
# Clamp to valid range
y_arr = clip_probability(y_arr, xp, eps=0.01)
# Return mean of clamped values
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)
[docs]
def estimate_phi(self, y, mu=None):
"""Estimate precision parameter from data.
Uses method-of-moments estimation based on the sample variance:
φ̂ = ȳ(1 − ȳ) / Var(y) − 1
.. note::
**Limitation**: this is the *marginal* MoM estimator — it uses
the sample mean/variance of y and ignores the fitted means
``mu`` (kept for API compatibility only). With heterogeneous
means it is biased; a proper conditional estimator would solve
Σ (yᵢ−μ̂ᵢ)² / [μ̂ᵢ(1−μ̂ᵢ)] = n / (φ+1) style equations.
Parameters
----------
y : array
Observed proportions.
mu : array, optional
Fitted means. Currently unused (see note above).
Returns
-------
float
Estimated precision parameter φ.
"""
xp = namespace(y)
y_arr = clip_probability(as_namespace_array(y, xp, like=y), xp, eps=0.01)
return self._estimate_phi_mm(y_arr, xp)
@property
def default_link(self) -> LinkFunction:
"""Return the default link function (Logit)."""
return self._link
__all__ = ["BetaFamily"]