#![allow(clippy::many_single_char_names)]
use self::ChiSquaredRepr::*;
use self::GammaRepr::*;
use crate::normal::StandardNormal;
use num_traits::Float;
use crate::{Distribution, Exp, Exp1, Open01};
use rand::Rng;
use core::fmt;
#[cfg(feature = "serde1")]
use serde::{Serialize, Deserialize};
#[derive(Clone, Copy, Debug, PartialEq)]
#[cfg_attr(feature = "serde1", derive(Serialize, Deserialize))]
pub struct Gamma<F>
where
F: Float,
StandardNormal: Distribution<F>,
Exp1: Distribution<F>,
Open01: Distribution<F>,
{
repr: GammaRepr<F>,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Error {
ShapeTooSmall,
ScaleTooSmall,
ScaleTooLarge,
}
impl fmt::Display for Error {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str(match self {
Error::ShapeTooSmall => "shape is not positive in gamma distribution",
Error::ScaleTooSmall => "scale is not positive in gamma distribution",
Error::ScaleTooLarge => "scale is infinity in gamma distribution",
})
}
}
#[cfg(feature = "std")]
#[cfg_attr(doc_cfg, doc(cfg(feature = "std")))]
impl std::error::Error for Error {}
#[derive(Clone, Copy, Debug, PartialEq)]
#[cfg_attr(feature = "serde1", derive(Serialize, Deserialize))]
enum GammaRepr<F>
where
F: Float,
StandardNormal: Distribution<F>,
Exp1: Distribution<F>,
Open01: Distribution<F>,
{
Large(GammaLargeShape<F>),
One(Exp<F>),
Small(GammaSmallShape<F>),
}
#[derive(Clone, Copy, Debug, PartialEq)]
#[cfg_attr(feature = "serde1", derive(Serialize, Deserialize))]
struct GammaSmallShape<F>
where
F: Float,
StandardNormal: Distribution<F>,
Open01: Distribution<F>,
{
inv_shape: F,
large_shape: GammaLargeShape<F>,
}
#[derive(Clone, Copy, Debug, PartialEq)]
#[cfg_attr(feature = "serde1", derive(Serialize, Deserialize))]
struct GammaLargeShape<F>
where
F: Float,
StandardNormal: Distribution<F>,
Open01: Distribution<F>,
{
scale: F,
c: F,
d: F,
}
impl<F> Gamma<F>
where
F: Float,
StandardNormal: Distribution<F>,
Exp1: Distribution<F>,
Open01: Distribution<F>,
{
#[inline]
pub fn new(shape: F, scale: F) -> Result<Gamma<F>, Error> {
if !(shape > F::zero()) {
return Err(Error::ShapeTooSmall);
}
if !(scale > F::zero()) {
return Err(Error::ScaleTooSmall);
}
let repr = if shape == F::one() {
One(Exp::new(F::one() / scale).map_err(|_| Error::ScaleTooLarge)?)
} else if shape < F::one() {
Small(GammaSmallShape::new_raw(shape, scale))
} else {
Large(GammaLargeShape::new_raw(shape, scale))
};
Ok(Gamma { repr })
}
}
impl<F> GammaSmallShape<F>
where
F: Float,
StandardNormal: Distribution<F>,
Open01: Distribution<F>,
{
fn new_raw(shape: F, scale: F) -> GammaSmallShape<F> {
GammaSmallShape {
inv_shape: F::one() / shape,
large_shape: GammaLargeShape::new_raw(shape + F::one(), scale),
}
}
}
impl<F> GammaLargeShape<F>
where
F: Float,
StandardNormal: Distribution<F>,
Open01: Distribution<F>,
{
fn new_raw(shape: F, scale: F) -> GammaLargeShape<F> {
let d = shape - F::from(1. / 3.).unwrap();
GammaLargeShape {
scale,
c: F::one() / (F::from(9.).unwrap() * d).sqrt(),
d,
}
}
}
impl<F> Distribution<F> for Gamma<F>
where
F: Float,
StandardNormal: Distribution<F>,
Exp1: Distribution<F>,
Open01: Distribution<F>,
{
fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> F {
match self.repr {
Small(ref g) => g.sample(rng),
One(ref g) => g.sample(rng),
Large(ref g) => g.sample(rng),
}
}
}
impl<F> Distribution<F> for GammaSmallShape<F>
where
F: Float,
StandardNormal: Distribution<F>,
Open01: Distribution<F>,
{
fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> F {
let u: F = rng.sample(Open01);
self.large_shape.sample(rng) * u.powf(self.inv_shape)
}
}
impl<F> Distribution<F> for GammaLargeShape<F>
where
F: Float,
StandardNormal: Distribution<F>,
Open01: Distribution<F>,
{
fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> F {
loop {
let x: F = rng.sample(StandardNormal);
let v_cbrt = F::one() + self.c * x;
if v_cbrt <= F::zero() {
continue;
}
let v = v_cbrt * v_cbrt * v_cbrt;
let u: F = rng.sample(Open01);
let x_sqr = x * x;
if u < F::one() - F::from(0.0331).unwrap() * x_sqr * x_sqr
|| u.ln() < F::from(0.5).unwrap() * x_sqr + self.d * (F::one() - v + v.ln())
{
return self.d * v * self.scale;
}
}
}
}
#[derive(Clone, Copy, Debug, PartialEq)]
#[cfg_attr(feature = "serde1", derive(Serialize, Deserialize))]
pub struct ChiSquared<F>
where
F: Float,
StandardNormal: Distribution<F>,
Exp1: Distribution<F>,
Open01: Distribution<F>,
{
repr: ChiSquaredRepr<F>,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[cfg_attr(feature = "serde1", derive(Serialize, Deserialize))]
pub enum ChiSquaredError {
DoFTooSmall,
}
impl fmt::Display for ChiSquaredError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str(match self {
ChiSquaredError::DoFTooSmall => {
"degrees-of-freedom k is not positive in chi-squared distribution"
}
})
}
}
#[cfg(feature = "std")]
#[cfg_attr(doc_cfg, doc(cfg(feature = "std")))]
impl std::error::Error for ChiSquaredError {}
#[derive(Clone, Copy, Debug, PartialEq)]
#[cfg_attr(feature = "serde1", derive(Serialize, Deserialize))]
enum ChiSquaredRepr<F>
where
F: Float,
StandardNormal: Distribution<F>,
Exp1: Distribution<F>,
Open01: Distribution<F>,
{
DoFExactlyOne,
DoFAnythingElse(Gamma<F>),
}
impl<F> ChiSquared<F>
where
F: Float,
StandardNormal: Distribution<F>,
Exp1: Distribution<F>,
Open01: Distribution<F>,
{
pub fn new(k: F) -> Result<ChiSquared<F>, ChiSquaredError> {
let repr = if k == F::one() {
DoFExactlyOne
} else {
if !(F::from(0.5).unwrap() * k > F::zero()) {
return Err(ChiSquaredError::DoFTooSmall);
}
DoFAnythingElse(Gamma::new(F::from(0.5).unwrap() * k, F::from(2.0).unwrap()).unwrap())
};
Ok(ChiSquared { repr })
}
}
impl<F> Distribution<F> for ChiSquared<F>
where
F: Float,
StandardNormal: Distribution<F>,
Exp1: Distribution<F>,
Open01: Distribution<F>,
{
fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> F {
match self.repr {
DoFExactlyOne => {
let norm: F = rng.sample(StandardNormal);
norm * norm
}
DoFAnythingElse(ref g) => g.sample(rng),
}
}
}
#[derive(Clone, Copy, Debug, PartialEq)]
#[cfg_attr(feature = "serde1", derive(Serialize, Deserialize))]
pub struct FisherF<F>
where
F: Float,
StandardNormal: Distribution<F>,
Exp1: Distribution<F>,
Open01: Distribution<F>,
{
numer: ChiSquared<F>,
denom: ChiSquared<F>,
dof_ratio: F,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[cfg_attr(feature = "serde1", derive(Serialize, Deserialize))]
pub enum FisherFError {
MTooSmall,
NTooSmall,
}
impl fmt::Display for FisherFError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str(match self {
FisherFError::MTooSmall => "m is not positive in Fisher F distribution",
FisherFError::NTooSmall => "n is not positive in Fisher F distribution",
})
}
}
#[cfg(feature = "std")]
#[cfg_attr(doc_cfg, doc(cfg(feature = "std")))]
impl std::error::Error for FisherFError {}
impl<F> FisherF<F>
where
F: Float,
StandardNormal: Distribution<F>,
Exp1: Distribution<F>,
Open01: Distribution<F>,
{
pub fn new(m: F, n: F) -> Result<FisherF<F>, FisherFError> {
let zero = F::zero();
if !(m > zero) {
return Err(FisherFError::MTooSmall);
}
if !(n > zero) {
return Err(FisherFError::NTooSmall);
}
Ok(FisherF {
numer: ChiSquared::new(m).unwrap(),
denom: ChiSquared::new(n).unwrap(),
dof_ratio: n / m,
})
}
}
impl<F> Distribution<F> for FisherF<F>
where
F: Float,
StandardNormal: Distribution<F>,
Exp1: Distribution<F>,
Open01: Distribution<F>,
{
fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> F {
self.numer.sample(rng) / self.denom.sample(rng) * self.dof_ratio
}
}
#[derive(Clone, Copy, Debug, PartialEq)]
#[cfg_attr(feature = "serde1", derive(Serialize, Deserialize))]
pub struct StudentT<F>
where
F: Float,
StandardNormal: Distribution<F>,
Exp1: Distribution<F>,
Open01: Distribution<F>,
{
chi: ChiSquared<F>,
dof: F,
}
impl<F> StudentT<F>
where
F: Float,
StandardNormal: Distribution<F>,
Exp1: Distribution<F>,
Open01: Distribution<F>,
{
pub fn new(n: F) -> Result<StudentT<F>, ChiSquaredError> {
Ok(StudentT {
chi: ChiSquared::new(n)?,
dof: n,
})
}
}
impl<F> Distribution<F> for StudentT<F>
where
F: Float,
StandardNormal: Distribution<F>,
Exp1: Distribution<F>,
Open01: Distribution<F>,
{
fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> F {
let norm: F = rng.sample(StandardNormal);
norm * (self.dof / self.chi.sample(rng)).sqrt()
}
}
#[derive(Clone, Copy, Debug, PartialEq)]
#[cfg_attr(feature = "serde1", derive(Serialize, Deserialize))]
enum BetaAlgorithm<N> {
BB(BB<N>),
BC(BC<N>),
}
#[derive(Clone, Copy, Debug, PartialEq)]
#[cfg_attr(feature = "serde1", derive(Serialize, Deserialize))]
struct BB<N> {
alpha: N,
beta: N,
gamma: N,
}
#[derive(Clone, Copy, Debug, PartialEq)]
#[cfg_attr(feature = "serde1", derive(Serialize, Deserialize))]
struct BC<N> {
alpha: N,
beta: N,
kappa1: N,
kappa2: N,
}
#[derive(Clone, Copy, Debug, PartialEq)]
#[cfg_attr(feature = "serde1", derive(Serialize, Deserialize))]
pub struct Beta<F>
where
F: Float,
Open01: Distribution<F>,
{
a: F, b: F, switched_params: bool,
algorithm: BetaAlgorithm<F>,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[cfg_attr(feature = "serde1", derive(Serialize, Deserialize))]
pub enum BetaError {
AlphaTooSmall,
BetaTooSmall,
}
impl fmt::Display for BetaError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str(match self {
BetaError::AlphaTooSmall => "alpha is not positive in beta distribution",
BetaError::BetaTooSmall => "beta is not positive in beta distribution",
})
}
}
#[cfg(feature = "std")]
#[cfg_attr(doc_cfg, doc(cfg(feature = "std")))]
impl std::error::Error for BetaError {}
impl<F> Beta<F>
where
F: Float,
Open01: Distribution<F>,
{
pub fn new(alpha: F, beta: F) -> Result<Beta<F>, BetaError> {
if !(alpha > F::zero()) {
return Err(BetaError::AlphaTooSmall);
}
if !(beta > F::zero()) {
return Err(BetaError::BetaTooSmall);
}
let (a0, b0) = (alpha, beta);
let (a, b, switched_params) = if a0 < b0 {
(a0, b0, false)
} else {
(b0, a0, true)
};
if a > F::one() {
let alpha = a + b;
let beta = ((alpha - F::from(2.).unwrap())
/ (F::from(2.).unwrap()*a*b - alpha)).sqrt();
let gamma = a + F::one() / beta;
Ok(Beta {
a, b, switched_params,
algorithm: BetaAlgorithm::BB(BB {
alpha, beta, gamma,
})
})
} else {
let (a, b, switched_params) = (b, a, !switched_params);
let alpha = a + b;
let beta = F::one() / b;
let delta = F::one() + a - b;
let kappa1 = delta
* (F::from(1. / 18. / 4.).unwrap() + F::from(3. / 18. / 4.).unwrap()*b)
/ (a*beta - F::from(14. / 18.).unwrap());
let kappa2 = F::from(0.25).unwrap()
+ (F::from(0.5).unwrap() + F::from(0.25).unwrap()/delta)*b;
Ok(Beta {
a, b, switched_params,
algorithm: BetaAlgorithm::BC(BC {
alpha, beta, kappa1, kappa2,
})
})
}
}
}
impl<F> Distribution<F> for Beta<F>
where
F: Float,
Open01: Distribution<F>,
{
fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> F {
let mut w;
match self.algorithm {
BetaAlgorithm::BB(algo) => {
loop {
let u1 = rng.sample(Open01);
let u2 = rng.sample(Open01);
let v = algo.beta * (u1 / (F::one() - u1)).ln();
w = self.a * v.exp();
let z = u1*u1 * u2;
let r = algo.gamma * v - F::from(4.).unwrap().ln();
let s = self.a + r - w;
if s + F::one() + F::from(5.).unwrap().ln()
>= F::from(5.).unwrap() * z {
break;
}
let t = z.ln();
if s >= t {
break;
}
if !(r + algo.alpha * (algo.alpha / (self.b + w)).ln() < t) {
break;
}
}
},
BetaAlgorithm::BC(algo) => {
loop {
let z;
let u1 = rng.sample(Open01);
let u2 = rng.sample(Open01);
if u1 < F::from(0.5).unwrap() {
let y = u1 * u2;
z = u1 * y;
if F::from(0.25).unwrap() * u2 + z - y >= algo.kappa1 {
continue;
}
} else {
z = u1 * u1 * u2;
if z <= F::from(0.25).unwrap() {
let v = algo.beta * (u1 / (F::one() - u1)).ln();
w = self.a * v.exp();
break;
}
if z >= algo.kappa2 {
continue;
}
}
let v = algo.beta * (u1 / (F::one() - u1)).ln();
w = self.a * v.exp();
if !(algo.alpha * ((algo.alpha / (self.b + w)).ln() + v)
- F::from(4.).unwrap().ln() < z.ln()) {
break;
};
}
},
};
if !self.switched_params {
if w == F::infinity() {
return F::one();
}
w / (self.b + w)
} else {
self.b / (self.b + w)
}
}
}
#[cfg(test)]
mod test {
use super::*;
#[test]
fn test_chi_squared_one() {
let chi = ChiSquared::new(1.0).unwrap();
let mut rng = crate::test::rng(201);
for _ in 0..1000 {
chi.sample(&mut rng);
}
}
#[test]
fn test_chi_squared_small() {
let chi = ChiSquared::new(0.5).unwrap();
let mut rng = crate::test::rng(202);
for _ in 0..1000 {
chi.sample(&mut rng);
}
}
#[test]
fn test_chi_squared_large() {
let chi = ChiSquared::new(30.0).unwrap();
let mut rng = crate::test::rng(203);
for _ in 0..1000 {
chi.sample(&mut rng);
}
}
#[test]
#[should_panic]
fn test_chi_squared_invalid_dof() {
ChiSquared::new(-1.0).unwrap();
}
#[test]
fn test_f() {
let f = FisherF::new(2.0, 32.0).unwrap();
let mut rng = crate::test::rng(204);
for _ in 0..1000 {
f.sample(&mut rng);
}
}
#[test]
fn test_t() {
let t = StudentT::new(11.0).unwrap();
let mut rng = crate::test::rng(205);
for _ in 0..1000 {
t.sample(&mut rng);
}
}
#[test]
fn test_beta() {
let beta = Beta::new(1.0, 2.0).unwrap();
let mut rng = crate::test::rng(201);
for _ in 0..1000 {
beta.sample(&mut rng);
}
}
#[test]
#[should_panic]
fn test_beta_invalid_dof() {
Beta::new(0., 0.).unwrap();
}
#[test]
fn test_beta_small_param() {
let beta = Beta::<f64>::new(1e-3, 1e-3).unwrap();
let mut rng = crate::test::rng(206);
for i in 0..1000 {
assert!(!beta.sample(&mut rng).is_nan(), "failed at i={}", i);
}
}
#[test]
fn gamma_distributions_can_be_compared() {
assert_eq!(Gamma::new(1.0, 2.0), Gamma::new(1.0, 2.0));
}
#[test]
fn beta_distributions_can_be_compared() {
assert_eq!(Beta::new(1.0, 2.0), Beta::new(1.0, 2.0));
}
#[test]
fn chi_squared_distributions_can_be_compared() {
assert_eq!(ChiSquared::new(1.0), ChiSquared::new(1.0));
}
#[test]
fn fisher_f_distributions_can_be_compared() {
assert_eq!(FisherF::new(1.0, 2.0), FisherF::new(1.0, 2.0));
}
#[test]
fn student_t_distributions_can_be_compared() {
assert_eq!(StudentT::new(1.0), StudentT::new(1.0));
}
}