use crate::FloatScalar;
use crate::special::{lbeta, betainc};
use super::{ContinuousDistribution, StatsError, quantile_newton};
#[derive(Debug, Clone, Copy)]
pub struct Beta<T> {
alpha: T,
beta: T,
}
impl<T: FloatScalar> Beta<T> {
pub fn new(alpha: T, beta: T) -> Result<Self, StatsError> {
if alpha <= T::zero() || beta <= T::zero() {
return Err(StatsError::InvalidParameter);
}
Ok(Self { alpha, beta })
}
}
impl<T: FloatScalar> Beta<T> {
pub fn sample(&self, rng: &mut super::Rng) -> T {
let x = rng.next_gamma(self.alpha);
let y = rng.next_gamma(self.beta);
x / (x + y)
}
pub fn sample_array<const K: usize>(&self, rng: &mut super::Rng) -> [T; K] {
let mut out = [T::zero(); K];
for v in out.iter_mut() {
*v = self.sample(rng);
}
out
}
}
impl<T: FloatScalar> ContinuousDistribution<T> for Beta<T> {
fn pdf(&self, x: T) -> T {
if x < T::zero() || x > T::one() {
return T::zero();
}
self.ln_pdf(x).exp()
}
fn ln_pdf(&self, x: T) -> T {
if x < T::zero() || x > T::one() {
return T::neg_infinity();
}
let one = T::one();
(self.alpha - one) * x.ln() + (self.beta - one) * (one - x).ln()
- lbeta(self.alpha, self.beta)
}
fn cdf(&self, x: T) -> T {
if x <= T::zero() {
return T::zero();
}
if x >= T::one() {
return T::one();
}
betainc(self.alpha, self.beta, x).unwrap_or(T::nan())
}
fn quantile(&self, p: T) -> T {
let x0 = self.mean();
let eps = T::epsilon();
quantile_newton(
|x| self.cdf(x),
|x| self.pdf(x),
p,
x0,
eps,
T::one() - eps,
)
}
fn mean(&self) -> T {
self.alpha / (self.alpha + self.beta)
}
fn variance(&self) -> T {
let ab = self.alpha + self.beta;
self.alpha * self.beta / (ab * ab * (ab + T::one()))
}
}