atelier_quant 0.0.12

//! # Probabilistic generators
//!
//! This module provides implementations for sampling from various probability
//! distributions, including:
//!
//! - Normal
//! - Poisson
//! - Exponential
//!
//! ## References
//!
//! - [rand_distr](https://docs.rs/rand_distr/latest/rand_distr/)
use rand::distr::Uniform;
use rand::prelude::*;
use rand_distr::Normal;

/// Trait for types that can produce random samples.
///
/// The generic parameter `R` defaults to [`ThreadRng`] but can be
/// substituted with any [`rand::Rng`] implementation for
/// deterministic testing.
pub trait Sampling<R: rand::Rng = ThreadRng> {
    /// Draw `n` independent samples from the distribution using `rng`.
    ///
    /// Returns a `Vec<f64>` of length `n`.
    fn sample(&self, rng: &mut R, n: usize) -> Vec<f64>;
}

/// Trait for distributions whose parameters can be estimated from data.
///
/// Implementations update the distribution's internal parameters
/// in-place to match the maximum-likelihood (or method-of-moments)
/// estimates derived from the observed `data`.
pub trait PDF {
    /// Fit the distribution parameters to the provided sample.
    fn fit(&mut self, data: &[f64]);
}

/// Continuous uniform distribution on `[lower, upper)`.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct UniformDistribution {
    /// Inclusive lower bound of the interval.
    pub lower: f64,
    /// Exclusive upper bound of the interval.
    pub upper: f64,
}

impl<R: Rng> Sampling<R> for UniformDistribution {
    fn sample(&self, rng: &mut R, n: usize) -> Vec<f64> {
        let uni = Uniform::new(self.lower, self.upper).unwrap();
        (0..n).map(|_| rng.sample(uni)).collect()
    }
}

/// Normal (Gaussian) distribution with mean `mu` and standard deviation `sigma`.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct NormalDistribution {
    /// Mean (location parameter).
    pub mu: f64,
    /// Standard deviation (scale parameter, must be > 0).
    pub sigma: f64,
}

impl<R: Rng> Sampling<R> for NormalDistribution {
    fn sample(&self, rng: &mut R, n: usize) -> Vec<f64> {
        let normal = Normal::new(self.mu, self.sigma).unwrap();
        normal.sample_iter(rng).take(n).collect()
    }
}

#[derive(Debug, Clone, Copy, PartialEq)]
/// Poisson distribution with rate parameter `lambda`.
///
/// Sampling uses the inverse-transform method (Knuth's algorithm).
/// The [`PDF`] implementation estimates `lambda` as the sample mean.
pub struct Poisson {
    /// Rate parameter λ (expected number of events per unit interval).
    pub lambda: f64,
}

impl Sampling for Poisson {
    fn sample(&self, rng: &mut ThreadRng, n: usize) -> Vec<f64> {
        let mut samples = Vec::with_capacity(n);
        let exp_neg_lambda = (-self.lambda).exp();

        for _ in 0..n {
            let mut x = 0;
            let mut p = 1.0;
            loop {
                let u = rng.random::<f64>();
                p *= u;
                if p <= exp_neg_lambda {
                    break;
                }
                x += 1;
            }
            samples.push(x as f64);
        }
        samples
    }
}

impl PDF for Poisson {
    fn fit(&mut self, data: &[f64]) {
        let sum: f64 = data.iter().sum();
        let count = data.len() as f64;
        self.lambda = if count > 0.0 { sum / count } else { 0.0 };
    }
}

/// Exponential distribution with rate parameter `lambda`.
///
/// Sampling uses the inverse-CDF method: `x = −ln(U) / λ`.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct Exponential {
    /// Rate parameter λ (inverse of the mean, must be > 0).
    pub lambda: f64,
}

impl Sampling for Exponential {
    fn sample(&self, rng: &mut ThreadRng, n: usize) -> Vec<f64> {
        let mut samples = Vec::with_capacity(n);
        let inv_lambda = -1.0 / self.lambda;
        for _ in 0..n {
            samples.push(inv_lambda * rng.random::<f64>().ln());
        }
        samples
    }
}