use super::{normal_quantile_standard, ContinuousDistribution, StatsError};
use crate::special::{erf, erfc};
use crate::FloatScalar;
#[derive(Debug, Clone, Copy)]
pub struct Normal<T> {
mu: T,
sigma: T,
sigma_sqrt_two_pi: T, ln_norm: T, }
impl<T: FloatScalar> Normal<T> {
pub fn new(mu: T, sigma: T) -> Result<Self, StatsError> {
if sigma <= T::zero() {
return Err(StatsError::InvalidParameter);
}
let two = T::one() + T::one();
let pi = T::from(core::f64::consts::PI).unwrap();
Ok(Self {
mu,
sigma,
sigma_sqrt_two_pi: sigma * (two * pi).sqrt(),
ln_norm: -sigma.ln() - (two * pi).ln() / two,
})
}
}
impl<T: FloatScalar> Normal<T> {
pub fn sample(&self, rng: &mut super::Rng) -> T {
self.mu + self.sigma * rng.next_normal::<T>()
}
impl_sample_array!(T, T::zero());
}
impl<T: FloatScalar> ContinuousDistribution<T> for Normal<T> {
fn pdf(&self, x: T) -> T {
let two = T::one() + T::one();
let z = (x - self.mu) / self.sigma;
(-(z * z) / two).exp() / self.sigma_sqrt_two_pi
}
fn ln_pdf(&self, x: T) -> T {
let two = T::one() + T::one();
let z = (x - self.mu) / self.sigma;
self.ln_norm - z * z / two
}
fn cdf(&self, x: T) -> T {
let half = T::from(0.5).unwrap();
let sqrt2 = T::from(core::f64::consts::SQRT_2).unwrap();
let z = (x - self.mu) / (self.sigma * sqrt2);
if z >= T::zero() {
half * (T::one() + erf(z))
} else {
half * erfc(-z)
}
}
fn quantile(&self, p: T) -> T {
self.mu + self.sigma * normal_quantile_standard(p)
}
fn mean(&self) -> T {
self.mu
}
fn variance(&self) -> T {
self.sigma * self.sigma
}
}