use crate::FloatScalar;
use crate::special::{erf, erfc};
use super::{ContinuousDistribution, StatsError, normal_quantile_standard};
#[derive(Debug, Clone, Copy)]
pub struct Normal<T> {
mu: T,
sigma: T,
}
impl<T: FloatScalar> Normal<T> {
pub fn new(mu: T, sigma: T) -> Result<Self, StatsError> {
if sigma <= T::zero() {
return Err(StatsError::InvalidParameter);
}
Ok(Self { mu, sigma })
}
}
impl<T: FloatScalar> Normal<T> {
pub fn sample(&self, rng: &mut super::Rng) -> T {
self.mu + self.sigma * rng.next_normal::<T>()
}
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 Normal<T> {
fn pdf(&self, x: T) -> T {
let two = T::one() + T::one();
let pi = T::from(core::f64::consts::PI).unwrap();
let z = (x - self.mu) / self.sigma;
(-(z * z) / two).exp() / (self.sigma * (two * pi).sqrt())
}
fn ln_pdf(&self, x: T) -> T {
let two = T::one() + T::one();
let pi = T::from(core::f64::consts::PI).unwrap();
let z = (x - self.mu) / self.sigma;
-self.sigma.ln() - (two * pi).ln() / two - 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
}
}