regression-diagnostics 0.2.0

//! Exponential-dispersion **families** for the generalized linear model.
//!
//! A [`Family`] bundles the four things that make one GLM differ from another —
//! the **link**, its **inverse** and **derivative**, and the **variance
//! function** `V(μ)` — plus the per-observation **deviance** and
//! **log-likelihood** the diagnostics are built from. The IRLS solver
//! [`GlmFit`](super::GlmFit) is written once against this trait; each concrete
//! family is a small, self-contained set of formulas.
//!
//! # Link choice
//!
//! Every family here uses the **log link** `η = ln μ`. For Poisson it is the
//! canonical link; for Gamma and the negative binomial it is not, but it is the
//! conventional, numerically robust choice because it keeps the fitted mean
//! strictly positive (which all three families require) for any linear
//! predictor. Fisher scoring is used throughout, so a non-canonical link costs
//! nothing in correctness — only the working weights change, and those come
//! straight off the trait.

use statrs::function::gamma::ln_gamma;

use crate::error::{RegressionError, Result};

/// A generalized-linear-model family: link, variance function, deviance and
/// log-likelihood.
///
/// Implementors are small value types ([`Poisson`], [`Gamma`],
/// [`NegativeBinomial`]) carrying at most a fixed shape parameter. The solver
/// treats the family as data, so adding a family is a matter of implementing
/// this trait — no changes to the IRLS loop or the diagnostics.
///
/// All three provided families use the **log link**, so [`Family::link`] is
/// `ln`, [`Family::inverse_link`] is `exp`, and [`Family::dmu_deta`] equals `μ`.
/// The methods are left general so a future canonical-link family can override
/// them.
pub trait Family: Clone + std::fmt::Debug {
    /// Short human-readable name, e.g. `"Poisson"`, used in reports.
    fn name(&self) -> &'static str;

    /// Reject responses outside the family's support (counts `≥ 0` for Poisson
    /// and the negative binomial, strictly positive reals for Gamma).
    ///
    /// # Errors
    ///
    /// [`RegressionError::InvalidResponse`] with a message describing the first
    /// offending value.
    fn validate(&self, y: &[f64]) -> Result<()>;

    /// A starting fitted mean for observation `y`, nudged into the family's
    /// open support so the first `η = ln μ` is finite.
    fn init_mu(&self, y: f64) -> f64;

    /// Link function `η = g(μ)`.
    fn link(&self, mu: f64) -> f64;

    /// Inverse link `μ = g⁻¹(η)`.
    fn inverse_link(&self, eta: f64) -> f64;

    /// Derivative `dμ/dη`, evaluated at `η`. For the log link this is `μ`.
    fn dmu_deta(&self, eta: f64) -> f64;

    /// Variance function `V(μ)`: the response variance is `φ · V(μ)`.
    fn variance(&self, mu: f64) -> f64;

    /// Per-observation **unit deviance** `d(y, μ) ≥ 0`; the residual deviance is
    /// its sum, and the signed square root is the deviance residual.
    fn unit_deviance(&self, y: f64, mu: f64) -> f64;

    /// Per-observation log-likelihood `ℓ(y; μ, φ)`. Families with a fixed
    /// dispersion (Poisson, negative binomial) ignore `dispersion`.
    fn loglik(&self, y: f64, mu: f64, dispersion: f64) -> f64;

    /// Whether the dispersion `φ` is known and fixed at `1` (Poisson, negative
    /// binomial) rather than estimated from the residuals (Gamma).
    ///
    /// This drives two things: whether the Wald reference distribution is the
    /// normal or Student's *t*, and whether the estimated `φ` counts as an extra
    /// parameter in AIC/BIC.
    fn dispersion_known(&self) -> bool;
}

/// Clamp a fitted mean away from `0` so `ln μ`, `1/μ`, and `V(μ)` stay finite.
const MU_EPS: f64 = 1e-10;

// ---------------------------------------------------------------------------
// Poisson
// ---------------------------------------------------------------------------

/// **Poisson** family with the canonical **log link** — the default model for
/// count responses `y ∈ {0, 1, 2, …}`.
///
/// Variance function `V(μ) = μ` (mean equals variance) and fixed dispersion
/// `φ = 1`. When the data are over-dispersed (sample variance exceeding the
/// mean) the Poisson standard errors are too small; reach for
/// [`NegativeBinomial`] instead, and compare the two with AIC.
#[derive(Debug, Clone, Copy, Default, PartialEq)]
pub struct Poisson;

impl Family for Poisson {
    fn name(&self) -> &'static str {
        "Poisson"
    }

    fn validate(&self, y: &[f64]) -> Result<()> {
        for &v in y {
            if !v.is_finite() || v < 0.0 {
                return Err(RegressionError::InvalidResponse {
                    msg: format!("Poisson response must be a non-negative count, found {v}"),
                });
            }
        }
        Ok(())
    }

    fn init_mu(&self, y: f64) -> f64 {
        // Standard GLM start: shrink toward the observation but keep μ > 0.
        (y + 0.1).max(MU_EPS)
    }

    fn link(&self, mu: f64) -> f64 {
        mu.max(MU_EPS).ln()
    }

    fn inverse_link(&self, eta: f64) -> f64 {
        eta.exp()
    }

    fn dmu_deta(&self, eta: f64) -> f64 {
        eta.exp().max(MU_EPS)
    }

    fn variance(&self, mu: f64) -> f64 {
        mu.max(MU_EPS)
    }

    fn unit_deviance(&self, y: f64, mu: f64) -> f64 {
        let mu = mu.max(MU_EPS);
        // 2[ y ln(y/μ) − (y − μ) ], with the y·ln term → 0 as y → 0.
        let term = if y > 0.0 { y * (y / mu).ln() } else { 0.0 };
        (2.0 * (term - (y - mu))).max(0.0)
    }

    fn loglik(&self, y: f64, mu: f64, _dispersion: f64) -> f64 {
        let mu = mu.max(MU_EPS);
        y * mu.ln() - mu - ln_gamma(y + 1.0)
    }

    fn dispersion_known(&self) -> bool {
        true
    }
}

// ---------------------------------------------------------------------------
// Gamma
// ---------------------------------------------------------------------------

/// **Gamma** family with the **log link** — a model for strictly positive,
/// right-skewed continuous responses (costs, durations, concentrations) whose
/// standard deviation scales with the mean.
///
/// Variance function `V(μ) = μ²` (constant coefficient of variation) and an
/// **estimated** dispersion `φ` (the squared CV). Because `φ` is estimated, Wald
/// inference uses Student's *t* with `n − p` degrees of freedom and the
/// information criteria charge one extra parameter for it.
#[derive(Debug, Clone, Copy, Default, PartialEq)]
pub struct Gamma;

impl Family for Gamma {
    fn name(&self) -> &'static str {
        "Gamma"
    }

    fn validate(&self, y: &[f64]) -> Result<()> {
        for &v in y {
            if !v.is_finite() || v <= 0.0 {
                return Err(RegressionError::InvalidResponse {
                    msg: format!("Gamma response must be strictly positive, found {v}"),
                });
            }
        }
        Ok(())
    }

    fn init_mu(&self, y: f64) -> f64 {
        y.max(MU_EPS)
    }

    fn link(&self, mu: f64) -> f64 {
        mu.max(MU_EPS).ln()
    }

    fn inverse_link(&self, eta: f64) -> f64 {
        eta.exp()
    }

    fn dmu_deta(&self, eta: f64) -> f64 {
        eta.exp().max(MU_EPS)
    }

    fn variance(&self, mu: f64) -> f64 {
        let mu = mu.max(MU_EPS);
        mu * mu
    }

    fn unit_deviance(&self, y: f64, mu: f64) -> f64 {
        let mu = mu.max(MU_EPS);
        // 2[ −ln(y/μ) + (y − μ)/μ ].
        (2.0 * (-(y / mu).ln() + (y - mu) / mu)).max(0.0)
    }

    fn loglik(&self, y: f64, mu: f64, dispersion: f64) -> f64 {
        // Shape ν = 1/φ, scale μ/ν. Density in mean/shape form:
        //   ln f = ν ln(ν/μ) + (ν−1) ln y − ν y/μ − ln Γ(ν).
        let mu = mu.max(MU_EPS);
        let nu = 1.0 / dispersion.max(MU_EPS);
        nu * (nu / mu).ln() + (nu - 1.0) * y.ln() - nu * y / mu - ln_gamma(nu)
    }

    fn dispersion_known(&self) -> bool {
        false
    }
}

// ---------------------------------------------------------------------------
// Negative binomial
// ---------------------------------------------------------------------------

/// **Negative binomial** family (NB2) with the **log link** — the standard
/// remedy for **over-dispersed counts**, where the variance grows faster than
/// the mean than Poisson allows.
///
/// The dispersion (shape) parameter `θ > 0` is treated as **known and fixed**
/// (supplied at construction): the variance function is `V(μ) = μ + μ²/θ`, which
/// approaches Poisson's `μ` as `θ → ∞` and admits progressively more
/// over-dispersion as `θ` shrinks. Estimating `θ` jointly with the coefficients
/// is a separate problem left to the caller (e.g. a profile-likelihood search
/// over `θ`, refitting at each value).
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct NegativeBinomial {
    theta: f64,
}

impl NegativeBinomial {
    /// Create a negative-binomial family with fixed dispersion parameter `θ`.
    ///
    /// # Errors
    ///
    /// [`RegressionError::InvalidParameter`] if `θ` is not finite and strictly
    /// positive.
    pub fn new(theta: f64) -> Result<Self> {
        if !theta.is_finite() || theta <= 0.0 {
            return Err(RegressionError::InvalidParameter {
                msg: format!("negative-binomial θ must be finite and > 0, got {theta}"),
            });
        }
        Ok(Self { theta })
    }

    /// The fixed dispersion (shape) parameter `θ`.
    pub fn theta(&self) -> f64 {
        self.theta
    }
}

impl Family for NegativeBinomial {
    fn name(&self) -> &'static str {
        "Negative Binomial"
    }

    fn validate(&self, y: &[f64]) -> Result<()> {
        for &v in y {
            if !v.is_finite() || v < 0.0 {
                return Err(RegressionError::InvalidResponse {
                    msg: format!(
                        "negative-binomial response must be a non-negative count, found {v}"
                    ),
                });
            }
        }
        Ok(())
    }

    fn init_mu(&self, y: f64) -> f64 {
        (y + 0.1).max(MU_EPS)
    }

    fn link(&self, mu: f64) -> f64 {
        mu.max(MU_EPS).ln()
    }

    fn inverse_link(&self, eta: f64) -> f64 {
        eta.exp()
    }

    fn dmu_deta(&self, eta: f64) -> f64 {
        eta.exp().max(MU_EPS)
    }

    fn variance(&self, mu: f64) -> f64 {
        let mu = mu.max(MU_EPS);
        mu + mu * mu / self.theta
    }

    fn unit_deviance(&self, y: f64, mu: f64) -> f64 {
        let mu = mu.max(MU_EPS);
        let theta = self.theta;
        // 2[ y ln(y/μ) − (y + θ) ln((y + θ)/(μ + θ)) ], y·ln term → 0 as y → 0.
        let term1 = if y > 0.0 { y * (y / mu).ln() } else { 0.0 };
        let term2 = (y + theta) * ((y + theta) / (mu + theta)).ln();
        (2.0 * (term1 - term2)).max(0.0)
    }

    fn loglik(&self, y: f64, mu: f64, _dispersion: f64) -> f64 {
        let mu = mu.max(MU_EPS);
        let theta = self.theta;
        // ln Γ(y+θ) − ln Γ(θ) − ln Γ(y+1)
        //   + θ ln(θ/(θ+μ)) + y ln(μ/(μ+θ)), with the y·ln term → 0 as y → 0.
        let y_term = if y > 0.0 { y * (mu / (mu + theta)).ln() } else { 0.0 };
        ln_gamma(y + theta) - ln_gamma(theta) - ln_gamma(y + 1.0)
            + theta * (theta / (theta + mu)).ln()
            + y_term
    }

    fn dispersion_known(&self) -> bool {
        true
    }
}