discrete_distributions.hpp

Go to the documentation of this file.

 
#pragma once
 
#include "statcpp/special_functions.hpp"
#include "statcpp/random_engine.hpp"
 
#include <algorithm>
#include <cmath>
#include <cstdint>
#include <limits>
#include <random>
#include <stdexcept>
 
namespace statcpp {
 
// ============================================================================
// Helper: Log factorial and binomial coefficient
// ============================================================================
 
inline double log_factorial(std::uint64_t n)
{
    if (n <= 1) return 0.0;
    return lgamma(static_cast<double>(n + 1));
}
 
inline double log_binomial_coef(std::uint64_t n, std::uint64_t k)
{
    if (k > n) return -std::numeric_limits<double>::infinity();
    if (k == 0 || k == n) return 0.0;
    return log_factorial(n) - log_factorial(k) - log_factorial(n - k);
}
 
inline double binomial_coef(std::uint64_t n, std::uint64_t k)
{
    if (k > n) return 0.0;
    return std::exp(log_binomial_coef(n, k));
}
 
// ============================================================================
// Binomial Distribution
// ============================================================================
 
inline double binomial_pmf(std::uint64_t k, std::uint64_t n, double p)
{
    if (p < 0.0 || p > 1.0) {
        throw std::invalid_argument("statcpp::binomial_pmf: p must be in [0, 1]");
    }
    if (k > n) return 0.0;
 
    if (p == 0.0) return (k == 0) ? 1.0 : 0.0;
    if (p == 1.0) return (k == n) ? 1.0 : 0.0;
 
    double log_pmf = log_binomial_coef(n, k) + k * std::log(p) + (n - k) * std::log(1.0 - p);
    return std::exp(log_pmf);
}
 
inline double binomial_cdf(std::uint64_t k, std::uint64_t n, double p)
{
    if (p < 0.0 || p > 1.0) {
        throw std::invalid_argument("statcpp::binomial_cdf: p must be in [0, 1]");
    }
    if (k >= n) return 1.0;
 
    // Use regularized incomplete beta function
    // P(X <= k) = I_{1-p}(n-k, k+1)
    return betainc(static_cast<double>(n - k), static_cast<double>(k + 1), 1.0 - p);
}
 
inline std::uint64_t binomial_quantile(double prob, std::uint64_t n, double p)
{
    if (p < 0.0 || p > 1.0) {
        throw std::invalid_argument("statcpp::binomial_quantile: p must be in [0, 1]");
    }
    if (prob < 0.0 || prob > 1.0) {
        throw std::invalid_argument("statcpp::binomial_quantile: prob must be in [0, 1]");
    }
    if (prob == 0.0) return 0;
    if (prob == 1.0) return n;
 
    // Binary search
    std::uint64_t lo = 0;
    std::uint64_t hi = n;
 
    while (lo < hi) {
        std::uint64_t mid = lo + (hi - lo) / 2;
        if (binomial_cdf(mid, n, p) < prob) {
            lo = mid + 1;
        } else {
            hi = mid;
        }
    }
 
    return lo;
}
 
template <typename Engine = default_random_engine>
std::uint64_t binomial_rand(std::uint64_t n, double p, Engine& engine)
{
    if (p < 0.0 || p > 1.0) {
        throw std::invalid_argument("statcpp::binomial_rand: p must be in [0, 1]");
    }
    std::binomial_distribution<std::uint64_t> dist(n, p);
    return dist(engine);
}
 
inline std::uint64_t binomial_rand(std::uint64_t n, double p)
{
    return binomial_rand(n, p, get_random_engine());
}
 
// ============================================================================
// Poisson Distribution
// ============================================================================
 
inline double poisson_pmf(std::uint64_t k, double lambda)
{
    if (lambda < 0.0) {
        throw std::invalid_argument("statcpp::poisson_pmf: lambda must be non-negative");
    }
    if (lambda == 0.0) return (k == 0) ? 1.0 : 0.0;
 
    double log_pmf = k * std::log(lambda) - lambda - log_factorial(k);
    return std::exp(log_pmf);
}
 
inline double poisson_cdf(std::uint64_t k, double lambda)
{
    if (lambda < 0.0) {
        throw std::invalid_argument("statcpp::poisson_cdf: lambda must be non-negative");
    }
    if (lambda == 0.0) return 1.0;
 
    // P(X <= k) = Q(k+1, lambda) = 1 - P(k+1, lambda) (upper regularized incomplete gamma)
    return gammainc_upper(static_cast<double>(k + 1), lambda);
}
 
inline std::uint64_t poisson_quantile(double p, double lambda)
{
    if (lambda < 0.0) {
        throw std::invalid_argument("statcpp::poisson_quantile: lambda must be non-negative");
    }
    if (p < 0.0 || p > 1.0) {
        throw std::invalid_argument("statcpp::poisson_quantile: p must be in [0, 1]");
    }
    if (p == 0.0) return 0;
    // Infinite support: prob == 1.0 has no finite quantile. Return the maximum
    // representable value instead of casting +inf to uint64 (undefined behavior).
    if (p == 1.0) return std::numeric_limits<std::uint64_t>::max();
    if (lambda == 0.0) return 0;
 
    // Start with Gaussian approximation
    double z = norm_quantile(p);
    double guess = lambda + z * std::sqrt(lambda);
    std::uint64_t k = static_cast<std::uint64_t>(std::max(0.0, guess));
 
    // Adjust up or down
    constexpr std::uint64_t kMaxIter = 1000000;
    std::uint64_t iter = 0;
    while (k > 0 && poisson_cdf(k - 1, lambda) >= p) {
        --k;
        if (++iter >= kMaxIter) return k;
    }
    while (poisson_cdf(k, lambda) < p) {
        ++k;
        if (++iter >= kMaxIter) return k;
    }
 
    return k;
}
 
template <typename Engine = default_random_engine>
std::uint64_t poisson_rand(double lambda, Engine& engine)
{
    if (lambda < 0.0) {
        throw std::invalid_argument("statcpp::poisson_rand: lambda must be non-negative");
    }
    std::poisson_distribution<std::uint64_t> dist(lambda);
    return dist(engine);
}
 
inline std::uint64_t poisson_rand(double lambda)
{
    return poisson_rand(lambda, get_random_engine());
}
 
// ============================================================================
// Geometric Distribution
// ============================================================================
 
inline double geometric_pmf(std::uint64_t k, double p)
{
    if (p <= 0.0 || p > 1.0) {
        throw std::invalid_argument("statcpp::geometric_pmf: p must be in (0, 1]");
    }
    if (p == 1.0) return (k == 0) ? 1.0 : 0.0;
 
    return std::pow(1.0 - p, static_cast<double>(k)) * p;
}
 
inline double geometric_cdf(std::uint64_t k, double p)
{
    if (p <= 0.0 || p > 1.0) {
        throw std::invalid_argument("statcpp::geometric_cdf: p must be in (0, 1]");
    }
    if (p == 1.0) return 1.0;
 
    return 1.0 - std::pow(1.0 - p, static_cast<double>(k + 1));
}
 
inline std::uint64_t geometric_quantile(double prob, double p)
{
    if (p <= 0.0 || p > 1.0) {
        throw std::invalid_argument("statcpp::geometric_quantile: p must be in (0, 1]");
    }
    if (prob < 0.0 || prob > 1.0) {
        throw std::invalid_argument("statcpp::geometric_quantile: prob must be in [0, 1]");
    }
    if (prob == 0.0) return 0;
    // Infinite support: prob == 1.0 has no finite quantile. Return the maximum
    // representable value instead of casting +inf to uint64 (undefined behavior).
    if (prob == 1.0) return std::numeric_limits<std::uint64_t>::max();
    if (p == 1.0) return 0;
 
    // Q(prob) = ceil(log(1 - prob) / log(1 - p)) - 1
    double k_real = std::ceil(std::log(1.0 - prob) / std::log(1.0 - p)) - 1.0;
    return static_cast<std::uint64_t>(std::max(0.0, k_real));
}
 
template <typename Engine = default_random_engine>
std::uint64_t geometric_rand(double p, Engine& engine)
{
    if (p <= 0.0 || p > 1.0) {
        throw std::invalid_argument("statcpp::geometric_rand: p must be in (0, 1]");
    }
    std::geometric_distribution<std::uint64_t> dist(p);
    return dist(engine);
}
 
inline std::uint64_t geometric_rand(double p)
{
    return geometric_rand(p, get_random_engine());
}
 
// ============================================================================
// Hypergeometric Distribution
// ============================================================================
 
inline double hypergeom_pmf(std::uint64_t k, std::uint64_t N, std::uint64_t K, std::uint64_t n)
{
    if (K > N) {
        throw std::invalid_argument("statcpp::hypergeom_pmf: K must be <= N");
    }
    if (n > N) {
        throw std::invalid_argument("statcpp::hypergeom_pmf: n must be <= N");
    }
 
    // k must be in valid range
    std::uint64_t k_min = (n > N - K) ? n - (N - K) : 0;
    std::uint64_t k_max = std::min(n, K);
 
    if (k < k_min || k > k_max) return 0.0;
 
    double log_pmf = log_binomial_coef(K, k) + log_binomial_coef(N - K, n - k) - log_binomial_coef(N, n);
    return std::exp(log_pmf);
}
 
inline double hypergeom_cdf(std::uint64_t k, std::uint64_t N, std::uint64_t K, std::uint64_t n)
{
    if (K > N) {
        throw std::invalid_argument("statcpp::hypergeom_cdf: K must be <= N");
    }
    if (n > N) {
        throw std::invalid_argument("statcpp::hypergeom_cdf: n must be <= N");
    }
 
    std::uint64_t k_min = (n > N - K) ? n - (N - K) : 0;
    std::uint64_t k_max = std::min(n, K);
 
    if (k >= k_max) return 1.0;
 
    double sum = 0.0;
    for (std::uint64_t i = k_min; i <= k; ++i) {
        sum += hypergeom_pmf(i, N, K, n);
    }
    return std::min(1.0, sum);
}
 
inline std::uint64_t hypergeom_quantile(double p, std::uint64_t N, std::uint64_t K, std::uint64_t n)
{
    if (K > N) {
        throw std::invalid_argument("statcpp::hypergeom_quantile: K must be <= N");
    }
    if (n > N) {
        throw std::invalid_argument("statcpp::hypergeom_quantile: n must be <= N");
    }
    if (p < 0.0 || p > 1.0) {
        throw std::invalid_argument("statcpp::hypergeom_quantile: p must be in [0, 1]");
    }
 
    std::uint64_t k_min = (n > N - K) ? n - (N - K) : 0;
    std::uint64_t k_max = std::min(n, K);
 
    if (p == 0.0) return k_min;
    if (p == 1.0) return k_max;
 
    double cum = 0.0;
    for (std::uint64_t k = k_min; k <= k_max; ++k) {
        cum += hypergeom_pmf(k, N, K, n);
        if (cum >= p) return k;
    }
 
    return k_max;
}
 
template <typename Engine = default_random_engine>
std::uint64_t hypergeom_rand(std::uint64_t N, std::uint64_t K, std::uint64_t n, Engine& engine)
{
    if (K > N) {
        throw std::invalid_argument("statcpp::hypergeom_rand: K must be <= N");
    }
    if (n > N) {
        throw std::invalid_argument("statcpp::hypergeom_rand: n must be <= N");
    }
 
    // Direct sampling (for moderate n)
    std::uint64_t successes = 0;
    std::uint64_t population = N;
    std::uint64_t success_states = K;
 
    std::uniform_real_distribution<double> uniform(0.0, 1.0);
 
    for (std::uint64_t i = 0; i < n; ++i) {
        double p = static_cast<double>(success_states) / static_cast<double>(population);
        if (uniform(engine) < p) {
            ++successes;
            --success_states;
        }
        --population;
    }
 
    return successes;
}
 
inline std::uint64_t hypergeom_rand(std::uint64_t N, std::uint64_t K, std::uint64_t n)
{
    return hypergeom_rand(N, K, n, get_random_engine());
}
 
// ============================================================================
// Negative Binomial Distribution
// ============================================================================
 
inline double nbinom_pmf(std::uint64_t k, double r, double p)
{
    if (r <= 0.0) {
        throw std::invalid_argument("statcpp::nbinom_pmf: r must be positive");
    }
    if (p <= 0.0 || p > 1.0) {
        throw std::invalid_argument("statcpp::nbinom_pmf: p must be in (0, 1]");
    }
    if (p == 1.0) return (k == 0) ? 1.0 : 0.0;
 
    double log_pmf = lgamma(k + r) - log_factorial(k) - lgamma(r)
                   + r * std::log(p) + k * std::log(1.0 - p);
    return std::exp(log_pmf);
}
 
inline double nbinom_cdf(std::uint64_t k, double r, double p)
{
    if (r <= 0.0) {
        throw std::invalid_argument("statcpp::nbinom_cdf: r must be positive");
    }
    if (p <= 0.0 || p > 1.0) {
        throw std::invalid_argument("statcpp::nbinom_cdf: p must be in (0, 1]");
    }
    if (p == 1.0) return 1.0;
 
    return betainc(r, static_cast<double>(k + 1), p);
}
 
inline std::uint64_t nbinom_quantile(double prob, double r, double p)
{
    if (r <= 0.0) {
        throw std::invalid_argument("statcpp::nbinom_quantile: r must be positive");
    }
    if (p <= 0.0 || p > 1.0) {
        throw std::invalid_argument("statcpp::nbinom_quantile: p must be in (0, 1]");
    }
    if (prob < 0.0 || prob > 1.0) {
        throw std::invalid_argument("statcpp::nbinom_quantile: prob must be in [0, 1]");
    }
    if (prob == 0.0) return 0;
    // Infinite support: prob == 1.0 has no finite quantile. Return the maximum
    // representable value instead of casting +inf to uint64 (undefined behavior).
    if (prob == 1.0) return std::numeric_limits<std::uint64_t>::max();
    if (p == 1.0) return 0;
 
    // Start with Gaussian approximation
    double mean_val = r * (1.0 - p) / p;
    double var_val = r * (1.0 - p) / (p * p);
    double z = norm_quantile(prob);
    double guess = mean_val + z * std::sqrt(var_val);
    std::uint64_t k = static_cast<std::uint64_t>(std::max(0.0, guess));
 
    // Adjust
    constexpr std::uint64_t kMaxIter = 1000000;
    std::uint64_t iter = 0;
    while (k > 0 && nbinom_cdf(k - 1, r, p) >= prob) {
        --k;
        if (++iter >= kMaxIter) return k;
    }
    while (nbinom_cdf(k, r, p) < prob) {
        ++k;
        if (++iter >= kMaxIter) return k;
    }
 
    return k;
}
 
template <typename Engine = default_random_engine>
std::uint64_t nbinom_rand(double r, double p, Engine& engine)
{
    if (r <= 0.0) {
        throw std::invalid_argument("statcpp::nbinom_rand: r must be positive");
    }
    if (p <= 0.0 || p > 1.0) {
        throw std::invalid_argument("statcpp::nbinom_rand: p must be in (0, 1]");
    }
 
    // Negative binomial as Poisson-Gamma mixture
    // X ~ NB(r, p) can be generated as Poisson(Y) where Y ~ Gamma(r, p/(1-p))
    std::gamma_distribution<double> gamma_dist(r, (1.0 - p) / p);
    double y = gamma_dist(engine);
    std::poisson_distribution<std::uint64_t> poisson_dist(y);
    return poisson_dist(engine);
}
 
inline std::uint64_t nbinom_rand(double r, double p)
{
    return nbinom_rand(r, p, get_random_engine());
}
 
// ============================================================================
// Bernoulli Distribution
// ============================================================================
 
inline double bernoulli_pmf(std::uint64_t k, double p)
{
    if (p < 0.0 || p > 1.0) {
        throw std::invalid_argument("statcpp::bernoulli_pmf: p must be in [0, 1]");
    }
    if (k == 0) {
        return 1.0 - p;
    } else if (k == 1) {
        return p;
    } else {
        return 0.0;
    }
}
 
inline double bernoulli_cdf(std::uint64_t k, double p)
{
    if (p < 0.0 || p > 1.0) {
        throw std::invalid_argument("statcpp::bernoulli_cdf: p must be in [0, 1]");
    }
    if (k == 0) {
        return 1.0 - p;
    } else {
        return 1.0;
    }
}
 
inline std::uint64_t bernoulli_quantile(double prob, double p)
{
    if (prob < 0.0 || prob > 1.0) {
        throw std::invalid_argument("statcpp::bernoulli_quantile: prob must be in [0, 1]");
    }
    if (p < 0.0 || p > 1.0) {
        throw std::invalid_argument("statcpp::bernoulli_quantile: p must be in [0, 1]");
    }
    if (prob <= 1.0 - p) return 0;
    return 1;
}
 
template <typename Engine = default_random_engine>
std::uint64_t bernoulli_rand(double p, Engine& engine)
{
    if (p < 0.0 || p > 1.0) {
        throw std::invalid_argument("statcpp::bernoulli_rand: p must be in [0, 1]");
    }
    std::bernoulli_distribution dist(p);
    return dist(engine) ? 1 : 0;
}
 
inline std::uint64_t bernoulli_rand(double p)
{
    return bernoulli_rand(p, get_random_engine());
}
 
// ============================================================================
// Discrete Uniform Distribution
// ============================================================================
 
inline double discrete_uniform_pmf(std::int64_t k, std::int64_t a, std::int64_t b)
{
    if (a > b) {
        throw std::invalid_argument("statcpp::discrete_uniform_pmf: a must be <= b");
    }
    if (k < a || k > b) {
        return 0.0;
    }
    return 1.0 / static_cast<double>(b - a + 1);
}
 
inline double discrete_uniform_cdf(std::int64_t k, std::int64_t a, std::int64_t b)
{
    if (a > b) {
        throw std::invalid_argument("statcpp::discrete_uniform_cdf: a must be <= b");
    }
    if (k < a) {
        return 0.0;
    }
    if (k >= b) {
        return 1.0;
    }
    return static_cast<double>(k - a + 1) / static_cast<double>(b - a + 1);
}
 
inline std::int64_t discrete_uniform_quantile(double p, std::int64_t a, std::int64_t b)
{
    if (a > b) {
        throw std::invalid_argument("statcpp::discrete_uniform_quantile: a must be <= b");
    }
    if (p < 0.0 || p > 1.0) {
        throw std::invalid_argument("statcpp::discrete_uniform_quantile: p must be in [0, 1]");
    }
    std::int64_t range = b - a + 1;
    std::int64_t k = a + static_cast<std::int64_t>(std::ceil(p * static_cast<double>(range) - 1.0));
    if (k < a) k = a;
    if (k > b) k = b;
    return k;
}
 
template <typename Engine = default_random_engine>
std::int64_t discrete_uniform_rand(std::int64_t a, std::int64_t b, Engine& engine)
{
    if (a > b) {
        throw std::invalid_argument("statcpp::discrete_uniform_rand: a must be <= b");
    }
    std::uniform_int_distribution<std::int64_t> dist(a, b);
    return dist(engine);
}
 
inline std::int64_t discrete_uniform_rand(std::int64_t a, std::int64_t b)
{
    return discrete_uniform_rand(a, b, get_random_engine());
}
 
} // namespace statcpp