r2rs-nmath 0.1.1

mod d;
mod p;
mod q;
mod r;

use strafe_type::{LogProbability64, Positive64, Probability64, Real64};

pub(crate) use self::{d::*, p::*, q::*, r::*};
use crate::traits::{Distribution, RNG};

/// # The Poisson Distribution
///
/// ## Description
///
/// Density, distribution function, quantile function and random generation for the Poisson
/// distribution with parameter lambda.
///
/// ## Arguments
///
/// * lambda: (non-negative) means.
///
/// ## Details
///
/// The Poisson distribution has density
///
/// $ p(x) = \lambda^x \frac{exp(-\lambda)}{x!} $
///
/// for x = 0, 1, 2, … . The mean and variance are $ E(X) = Var(X) = \lambda $.
///
/// Note that $ \lambda = 0 $ is really a limit case (setting 0^0 = 1) resulting in a point mass at
/// 0, see also the example.
///
/// If an element of x is not integer, the result of dpois is zero, with a warning. p(x) is
/// computed using Loader's algorithm, see the reference in dbinom.
///
/// The quantile is right continuous: qpois(p, lambda) is the smallest integer x such that
/// $P(X ≤ x) ≥ p$.
///
/// Setting lower.tail = FALSE allows to get much more precise results when the default,
/// lower.tail = TRUE would return 1, see the example below.
///
/// ## Density Plot
///
/// ```rust
/// # use r2rs_base::traits::StatisticalSlice;
/// # use r2rs_nmath::{distribution::PoissonBuilder, traits::Distribution};
/// # use strafe_plot::prelude::{IntoDrawingArea, Line, Plot, PlotOptions, SVGBackend, BLACK};
/// # use strafe_type::FloatConstraint;
/// let pois = PoissonBuilder::new().build();
/// let x = <[f64]>::sequence_by(-1.0, 7.0, 0.001);
/// let y = x
///     .iter()
///     .map(|x| pois.density(x).unwrap())
///     .collect::<Vec<_>>();
///
/// let root = SVGBackend::new("density.svg", (1024, 768)).into_drawing_area();
/// Plot::new()
///     .with_options(PlotOptions {
///         x_axis_label: "x".to_string(),
///         y_axis_label: "density".to_string(),
///         ..Default::default()
///     })
///     .with_plottable(Line {
///         x,
///         y,
///         color: BLACK,
///         ..Default::default()
///     })
///     .plot(&root)
///     .unwrap();
/// # use std::fs::rename;
/// #     drop(root);
/// #     rename(
/// #             format!("density.svg"),
/// #             format!("src/distribution/pois/doctest_out/density.svg"),
/// #     )
/// #     .unwrap();
/// ```
#[cfg_attr(feature = "doc_outputs", cfg_attr(all(), doc = embed_doc_image::embed_image!("density", "src/distribution/pois/doctest_out/density.svg")))]
#[cfg_attr(feature = "doc_outputs", cfg_attr(all(), doc = "![Density][density]"))]
///
/// ## Source
///
/// dpois uses C code contributed by Catherine Loader (see dbinom).
///
/// ppois uses pgamma.
///
/// qpois uses the Cornish–Fisher Expansion to include a skewness correction to a normal
/// approximation, followed by a search.
///
/// rpois uses
///
/// Ahrens, J. H. and Dieter, U. (1982). Computer generation of Poisson deviates from modified
/// normal distributions. ACM Transactions on Mathematical Software, 8, 163–179.
///
/// ## See Also
/// Distributions for other standard distributions, including dbinom for the binomial and dnbinom
/// for the negative binomial distribution.
///
/// poisson.test.
///
/// ## Examples
///
/// // Should be 1
/// ```rust
/// # use r2rs_nmath::{
/// #     distribution::PoissonBuilder, func::gamma, traits::Distribution,
/// # };
/// # use strafe_type::FloatConstraint;
/// let x = (0..=7).collect::<Vec<_>>();
/// let pois = PoissonBuilder::new().with_lambda(1).build();
/// let r = x
///     .iter()
///     .map(|x| -(pois.density(x).unwrap() * gamma(1 + x).unwrap().unwrap()).ln())
///     .collect::<Vec<_>>();
/// println!("{r:?}");
/// # use std::{fs::File, io::Write};
/// # let mut f = File::create("src/distribution/pois/doctest_out/dens.md").unwrap();
/// # writeln!(f, "```output").unwrap();
/// # writeln!(f, "{r:?}").unwrap();
/// # writeln!(f, "```").unwrap();
/// ```
#[cfg_attr(feature = "doc_outputs", cfg_attr(all(), doc = include_str!("doctest_out/dens.md")))]
///
/// ```rust
/// # use std::collections::HashMap;
/// #
/// # use r2rs_nmath::{
/// #     distribution::PoissonBuilder,
/// #     rng::MersenneTwister,
/// #     traits::{Distribution, RNG},
/// # };
/// # use strafe_type::FloatConstraint;
/// let pois = PoissonBuilder::new().with_lambda(4).build();
/// let mut rng = MersenneTwister::new();
/// rng.set_seed(1);
/// let mut r = (0..50)
///     .map(|_| pois.random_sample(&mut rng).unwrap() as usize)
///     .fold(HashMap::new(), |mut acc, r| {
///         *acc.entry(r).or_insert(0) += 1;
///         acc
///     })
///     .into_iter()
///     .collect::<Vec<_>>();
/// r.sort_by(|(i1, _), (i2, _)| i1.cmp(i2));
///
/// for (key, index) in &r {
///     println!("{key:2}: {index}");
/// }
/// # use std::{fs::File, io::Write};
/// # let mut f = File::create("src/distribution/pois/doctest_out/table.md").unwrap();
/// # writeln!(f, "```output").unwrap();
/// # for (key, index) in &r {
/// #     writeln!(f, "{key:2}: {index}").unwrap();
/// # }
/// # writeln!(f, "```").unwrap();
/// ```
#[cfg_attr(feature = "doc_outputs", cfg_attr(all(), doc = include_str!("doctest_out/table.md")))]
///
/// Using lower tail directly fixes the cancellation (values becoming 0)
/// ```rust
/// # use r2rs_nmath::{distribution::PoissonBuilder, traits::Distribution};
/// # use strafe_type::FloatConstraint;
/// let x = (15..=25).collect::<Vec<_>>();
/// let pois = PoissonBuilder::new().with_lambda(100).build();
/// let r1 = x
///     .iter()
///     .map(|x| 1.0 - pois.probability(x * 10, true).unwrap())
///     .collect::<Vec<_>>();
/// println!("{r1:?}");
/// let r2 = x
///     .iter()
///     .map(|x| pois.probability(x * 10, false).unwrap())
///     .collect::<Vec<_>>();
/// println!("{r2:?}");
/// # use std::{fs::File, io::Write};
/// # let mut f = File::create("src/distribution/pois/doctest_out/prob.md").unwrap();
/// # writeln!(f, "```output").unwrap();
/// # writeln!(f, "{r1:?}").unwrap();
/// # writeln!(f, "{r2:?}").unwrap();
/// # writeln!(f, "```").unwrap();
/// ```
#[cfg_attr(feature = "doc_outputs", cfg_attr(all(), doc = include_str!("doctest_out/prob.md")))]
///
/// ```rust
/// # use r2rs_base::traits::StatisticalSlice;
/// # use r2rs_nmath::{
/// #     distribution::{BinomialBuilder, PoissonBuilder},
/// #     traits::Distribution,
/// # };
/// # use strafe_plot::prelude::{IntoDrawingArea, Line, Plot, PlotOptions, SVGBackend, BLACK};
/// # use strafe_type::FloatConstraint;
/// let x = <[f64]>::sequence_by(-0.01, 5.0, 0.01);
///
/// let pois = PoissonBuilder::new().with_lambda(1).build();
/// let y1 = x
///     .iter()
///     .map(|x| pois.probability(x, true).unwrap())
///     .collect::<Vec<_>>();
///
/// let binom = BinomialBuilder::new()
///     .with_size(100)
///     .with_success_probability(0.01)
///     .build();
/// let y2 = x
///     .iter()
///     .map(|x| binom.probability(x, true).unwrap())
///     .collect::<Vec<_>>();
///
/// let root = SVGBackend::new("prob_plots.svg", (1024, 768)).into_drawing_area();
///
/// Plot::new()
///     .with_options(PlotOptions {
///         x_axis_label: "x".to_string(),
///         y_axis_label: "F(x)".to_string(),
///         plot_bottom: 0.5,
///         title: "Poisson(1) CDF".to_string(),
///         ..Default::default()
///     })
///     .with_plottable(Line {
///         x: x.clone(),
///         y: y1,
///         color: BLACK,
///         ..Default::default()
///     })
///     .plot(&root)
///     .unwrap();
///
/// Plot::new()
///     .with_options(PlotOptions {
///         x_axis_label: "x".to_string(),
///         y_axis_label: "F(x)".to_string(),
///         plot_top: 0.5,
///         title: "Binomial(100, 0.01) CDF".to_string(),
///         ..Default::default()
///     })
///     .with_plottable(Line {
///         x,
///         y: y2,
///         color: BLACK,
///         ..Default::default()
///     })
///     .plot(&root)
///     .unwrap();
/// # use std::fs::rename;
/// #     drop(root);
/// #     rename(
/// #             format!("prob_plots.svg"),
/// #             format!("src/distribution/pois/doctest_out/prob_plots.svg"),
/// #     )
/// #     .unwrap();
/// ```
#[cfg_attr(feature = "doc_outputs", cfg_attr(all(), doc = embed_doc_image::embed_image!("prob_plots", "src/distribution/pois/doctest_out/prob_plots.svg")))]
#[cfg_attr(
    feature = "doc_outputs",
    cfg_attr(all(), doc = "![Prob Plots][prob_plots]")
)]
///
/// ```rust
/// # use r2rs_nmath::{distribution::PoissonBuilder, traits::Distribution};
/// # use strafe_type::FloatConstraint;
/// assert_eq!(
///     PoissonBuilder::new()
///         .with_lambda(0)
///         .build()
///         .density(0)
///         .unwrap(),
///     1.0
/// );
/// assert_eq!(
///     PoissonBuilder::new()
///         .with_lambda(0)
///         .build()
///         .probability(0, true)
///         .unwrap(),
///     1.0
/// );
/// assert_eq!(
///     PoissonBuilder::new()
///         .with_lambda(0)
///         .build()
///         .quantile(1, true)
///         .unwrap(),
///     0.0
/// );
/// ```
pub struct Poisson {
    lambda: Positive64,
}

impl Distribution for Poisson {
    fn density<R: Into<Real64>>(&self, x: R) -> Real64 {
        dpois(x, self.lambda, false)
    }

    fn log_density<R: Into<Real64>>(&self, x: R) -> Real64 {
        dpois(x, self.lambda, true)
    }

    fn probability<R: Into<Real64>>(&self, q: R, lower_tail: bool) -> Probability64 {
        ppois(q, self.lambda, lower_tail)
    }

    fn log_probability<R: Into<Real64>>(&self, q: R, lower_tail: bool) -> LogProbability64 {
        log_ppois(q, self.lambda, lower_tail)
    }

    fn quantile<P: Into<Probability64>>(&self, p: P, lower_tail: bool) -> Real64 {
        qpois(p, self.lambda, lower_tail)
    }

    fn log_quantile<LP: Into<LogProbability64>>(&self, p: LP, lower_tail: bool) -> Real64 {
        log_qpois(p, self.lambda, lower_tail)
    }

    fn random_sample<R: RNG>(&self, rng: &mut R) -> Real64 {
        rpois(self.lambda, rng)
    }
}

pub struct PoissonBuilder {
    lambda: Option<Positive64>,
}

impl PoissonBuilder {
    pub fn new() -> Self {
        Self { lambda: None }
    }

    pub fn with_lambda<P: Into<Positive64>>(&mut self, lambda: P) -> &mut Self {
        self.lambda = Some(lambda.into());
        self
    }

    pub fn build(&self) -> Poisson {
        let lambda = self.lambda.unwrap_or(1.0.into());

        Poisson { lambda }
    }
}

#[cfg(test)]
mod tests;

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

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