blob: dee776860085fdc0e1b723a1ca86fd34b5fd91fc [file] [log] [blame]
 use crate::{Distribution, Standard, StandardNormal}; use num_traits::Float; use rand::Rng; /// Error type returned from `InverseGaussian::new` #[derive(Debug, PartialEq)] pub enum Error { /// `mean <= 0` or `nan`. MeanNegativeOrNull, /// `shape <= 0` or `nan`. ShapeNegativeOrNull, } /// The [inverse Gaussian distribution](https://en.wikipedia.org/wiki/Inverse_Gaussian_distribution) #[derive(Debug)] pub struct InverseGaussian where F: Float, StandardNormal: Distribution, Standard: Distribution, { mean: F, shape: F, } impl InverseGaussian where F: Float, StandardNormal: Distribution, Standard: Distribution, { /// Construct a new `InverseGaussian` distribution with the given mean and /// shape. pub fn new(mean: F, shape: F) -> Result, Error> { let zero = F::zero(); if !(mean > zero) { return Err(Error::MeanNegativeOrNull); } if !(shape > zero) { return Err(Error::ShapeNegativeOrNull); } Ok(Self { mean, shape }) } } impl Distribution for InverseGaussian where F: Float, StandardNormal: Distribution, Standard: Distribution, { fn sample(&self, rng: &mut R) -> F where R: Rng + ?Sized { let mu = self.mean; let l = self.shape; let v: F = rng.sample(StandardNormal); let y = mu * v * v; let mu_2l = mu / (F::from(2.).unwrap() * l); let x = mu + mu_2l * (y - (F::from(4.).unwrap() * l * y + y * y).sqrt()); let u: F = rng.gen(); if u <= mu / (mu + x) { return x; } mu * mu / x } } #[cfg(test)] mod tests { use super::*; #[test] fn test_inverse_gaussian() { let inv_gauss = InverseGaussian::new(1.0, 1.0).unwrap(); let mut rng = crate::test::rng(210); for _ in 0..1000 { inv_gauss.sample(&mut rng); } } #[test] fn test_inverse_gaussian_invalid_param() { assert!(InverseGaussian::new(-1.0, 1.0).is_err()); assert!(InverseGaussian::new(-1.0, -1.0).is_err()); assert!(InverseGaussian::new(1.0, -1.0).is_err()); assert!(InverseGaussian::new(1.0, 1.0).is_ok()); } }