r2rs-nmath 0.1.1

//! # The Multinomial Distribution
//!
//! ## Description
//!
//! Generate multinomially distributed random number vectors and compute multinomial probabilities.
//!
//! ## Arguments
//! * x: vector of length K of integers in 0:size.
//! * n: number of random vectors to draw.
//! * size: integer, say N, specifying the total number of objects that are put into K boxes in the typical multinomial experiment. For dmultinom, it defaults to sum(x).
//! * prob: numeric non-negative vector of length K, specifying the probability for the K classes; is internally normalized to sum 1. Infinite and missing values are not allowed.
//! * log: logical; if TRUE, log probabilities are computed.
//!
//! ## Details
//!
//! If x is a K-component vector, dmultinom(x, prob) is the probability
//!
//! $ P(X\[1\]=x\[1\], … , X\[K\]=x\[k\]) = C * prod(j=1 , …, K) p\[j\]^x\[j\] $
//!
//! where C is the ‘multinomial coefficient’
//! $ C = N! / (x\[1\]! * … * x\[K\]!) and N = sum(j=1, …, K) x\[j\] $.
//! By definition, each component $ X\[j\] $ is binomially distributed as
//! $ Bin(size, prob\[j\]) for j = 1, …, K $.
//!
//! The rmultinom() algorithm draws binomials $ X\[j\] $ from $ Bin(n\[j\], P\[j\]) $ sequentially, where
//! $ n\[1\] = N (N := size), P\[1\] = p\[1\] $ (p is prob scaled to sum 1), and for $ j ≥ 2 $,
//! recursively, $ n\[j\] = N - sum(k=1, …, j-1) X\[k\] $ and $ P\[j\] = p\[j\] / (1 - sum(p\[1:(j-1)\])) $.
//!
//! ## Value
//!
//! For rmultinom(), an integer K x n matrix where each column is a random vector generated
//! according to the desired multinomial law, and hence summing to size. Whereas the transposed
//! result would seem more natural at first, the returned matrix is more efficient because of
//! columnwise storage.
//!
//! ## Note
//!
//! dmultinom is currently not vectorized at all and has no C interface (API); this may be amended
//! in the future.
//!
//! ## See Also
//!
//! Distributions for standard distributions, including dbinom which is a special case conceptually.
//!
//! ## Examples
//!
//! ```r
//! rmultinom(10, size = 12, prob = c(0.1,0.2,0.8))
//!
//! pr <- c(1,3,6,10) # normalization not necessary for generation
//! rmultinom(10, 20, prob = pr)
//!
//! ## all possible outcomes of Multinom(N = 3, K = 3)
//! X <- t(as.matrix(expand.grid(0:3, 0:3))); X <- X[, colSums(X) <= 3]
//! X <- rbind(X, 3:3 - colSums(X)); dimnames(X) <- list(letters[1:3], NULL)
//! X
//! round(apply(X, 2, function(x) dmultinom(x, prob = c(1,2,5))), 3)
//! ```

use strafe_type::{pos64, r64, Checked, LogProbability64, Positive64, Probability64, Real64};

use crate::{
    distribution::multinom::{d::dmultinom, r::rmultinom},
    traits::{Distribution, RNG},
};

mod d;
mod r;

pub struct Multinomial {
    probability: Vec<Checked<Positive64>>,
}

impl Distribution for Multinomial {
    type X = Vec<Checked<Real64>>;

    fn density(&self, x: Self::X) -> Checked<Real64> {
        dmultinom(&x, &self.probability, false)
    }

    fn log_density(&self, x: Self::X) -> Checked<Real64> {
        dmultinom(&x, &self.probability, true)
    }

    fn probability(&self, _q: Checked<Real64>, _lower_tail: bool) -> Checked<Probability64> {
        unimplemented!()
    }

    fn log_probability(&self, _q: Checked<Real64>, _lower_tail: bool) -> Checked<LogProbability64> {
        unimplemented!()
    }

    fn quantile(&self, _p: Checked<Probability64>, _lower_tail: bool) -> Checked<Real64> {
        unimplemented!()
    }

    fn log_quantile(&self, _p: Checked<LogProbability64>, _lower_tail: bool) -> Checked<Real64> {
        unimplemented!()
    }

    fn random_sample<R>(&self, rng: &mut R) -> Self::X
    where
        R: RNG,
    {
        let k = self.probability.len();
        let mut ret = vec![r64!(0.0); k];
        rmultinom(1, &self.probability, k, &mut ret, rng);
        ret
    }
}

pub struct MultinomialBuilder {
    size: Option<usize>,
    probability: Option<Vec<Checked<Positive64>>>,
}

impl MultinomialBuilder {
    pub fn new() -> Self {
        Self {
            size: None,
            probability: None,
        }
    }

    pub fn set_size(&mut self, size: usize) {
        self.size = Some(size);
    }

    pub fn set_probability(&mut self, probability: Vec<Checked<Positive64>>) {
        self.probability = Some(probability);
    }

    pub fn build(self) -> Multinomial {
        let probability = if let Some(size) = self.size {
            vec![pos64!(1.0 / size as f64); size]
        } else {
            self.probability.unwrap_or_default()
        };

        Multinomial { probability }
    }
}

#[cfg(test)]
mod tests;

#[cfg(all(test, feature = "enable_proptest"))]
mod proptests;

#[cfg(all(test, feature = "enable_covtest"))]
mod covtests;