diff --git a/Cargo.toml b/Cargo.toml index faca884c..113364b7 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -43,9 +43,15 @@ optional = true default-features = false features = ["std"] +[[example]] +name = "random_clap" +required-features = ["rand", "nalgebra"] + [dev-dependencies] criterion = "0.5" anyhow = "1.0" +clap = { version = "4.0.32", features = ["derive"] } +clap_derive = "4.0.21" [dev-dependencies.nalgebra] version = "0.33" diff --git a/examples/random_clap.rs b/examples/random_clap.rs new file mode 100644 index 00000000..f0a08957 --- /dev/null +++ b/examples/random_clap.rs @@ -0,0 +1,233 @@ +extern crate statrs; + +use nalgebra as na; +use rand::{thread_rng, Rng}; +use statrs::distribution::{Binomial, Continuous, Discrete, Multinomial, Normal}; +use statrs::statistics::Mode; + +use std::fmt::Display; +use std::io::{self, BufWriter, Write}; +use std::str::{FromStr, Split}; + +use anyhow::Result; +use clap::{ArgAction, Parser, Subcommand}; + +#[derive(Parser)] +#[command(name="random_clapping", author, version="0.0.1", about, long_about = None)] +struct Args { + #[command(subcommand)] + cmd: Commands, +} + +#[derive(Subcommand, Debug)] +enum Commands { + /// for sampling + Sample { + #[arg(value_name = "SAMPLE COUNT")] + count: Option, + #[command(subcommand)] + dist: DistributionAsCommand, + }, + /// for evaluating distribution function density + Density { + /// sample to evaluate density at, default=distribution's mode + #[arg(short, long = "arg", action = ArgAction::Append, value_name = "SAMPLE", help = "sample(s) evaluate at, (space-delimited string for multivariate)")] + args: Vec, + #[command(subcommand)] + dist: DistributionAsCommand, + }, +} + +#[derive(Subcommand, Debug)] +enum DistributionAsCommand { + /// the multinomial distribution + Multinomial { + #[arg(value_name = "trial counts")] + n: u64, + #[arg(value_name = "success probabilities")] + p: Vec, + }, + /// the binomial distribution + Binomial { + #[arg(value_name = "trial counts")] + n: u64, + #[arg(value_name = "success probability", default_value = "0.5")] + p: f64, + }, + /// the normal distribution + Normal { + #[arg(value_name = "mean", default_value = "0.0")] + mu: f64, + #[arg(value_name = "standard deviation", default_value = "1.0")] + sigma: f64, + }, +} + +fn main() -> Result<()> { + let args = Args::parse(); + match args.cmd { + Commands::Sample { count, dist } => run_command_sample(count, dist), + Commands::Density { args, dist } => run_command_density(&args, dist), + }?; + println!(); + Ok(()) +} + +fn run_command_density(args_str: &[String], dist: DistributionAsCommand) -> Result<()> { + let densities = match dist { + DistributionAsCommand::Multinomial { n, p } => { + let dist = Multinomial::new(p, n)?; + if !args_str.is_empty() { + let mut densities = Vec::with_capacity(args_str.len()); + + for arg_str in args_str { + let arg = parse_str_split_to_vec(arg_str.split(' ')); + if arg.len() == dist.p().len() { + densities.push(dist.pmf(&arg.into())); + } else { + anyhow::bail!("dimension mismatch after parsing `--arg {arg_str}`"); + } + } + + densities + } else { + vec![dist.pmf(&dist.mode())] + } + } + DistributionAsCommand::Binomial { n, p } => { + let dist = Binomial::new(p, n)?; + if !args_str.is_empty() { + args_str + .iter() + .map_while(|s| match s.parse() { + Ok(x) => Some(x), + Err(e) => { + eprintln!("could not parse argment, got {e}"); + None + } + }) + .map(|x| dist.pmf(x)) + .collect() + } else { + vec![dist.pmf(dist.mode().unwrap())] + } + } + DistributionAsCommand::Normal { mu, sigma } => { + let dist = Normal::new(mu, sigma)?; + if !args_str.is_empty() { + args_str + .iter() + .map_while(|s| match s.parse() { + Ok(x) => Some(x), + Err(e) => { + eprintln!("could not parse argment, got {e}"); + None + } + }) + .map(|x| dist.pdf(x)) + .collect() + } else { + vec![dist.pdf(dist.mode().unwrap())] + } + } + }; + + util::write_interspersed(&mut BufWriter::new(io::stdout()), densities, ", ")?; + + Ok(()) +} + +fn parse_str_split_to_vec(sp: Split) -> Vec +where + T: FromStr, + E: Display + std::error::Error, +{ + sp.map_while(|si| match si.parse::() { + Ok(x) => Some(x), + Err(e) => { + eprintln!("could not parse argment, got {e}"); + None + } + }) + .collect() +} + +fn run_command_sample(count: Option, dist: DistributionAsCommand) -> Result<()> { + let count = count.unwrap_or(10); + + match dist { + // multinomial should print `count` of Vec + DistributionAsCommand::Multinomial { n, p } => { + let sample_iter = thread_rng().sample_iter(Multinomial::new(p, n)?); + print_multivariate_samples( + count, + sample_iter.map(|v: na::DVector| { + let vec: Vec<_> = v.into_iter().cloned().collect(); + vec + }), + )?; + } + // binomial should print `count` of uint + DistributionAsCommand::Binomial { n, p } => { + let sample_iter = thread_rng().sample_iter::(Binomial::new(p, n)?); + print_samples(count, sample_iter)?; + } + // normal should print `count` of float + DistributionAsCommand::Normal { mu, sigma } => { + let sample_iter = thread_rng().sample_iter(Normal::new(mu, sigma)?); + print_samples(count, sample_iter)? + } + } + + Ok(()) +} + +mod util { + use std::fmt::Display; + use std::io::{self, BufWriter, Write}; + pub(super) fn write_interspersed( + handle: &mut BufWriter, + it: I, + sep: &str, + ) -> io::Result<()> + where + I: IntoIterator, + T: Display, + W: Write, + { + let mut it = it.into_iter(); + if let Some(i) = it.next() { + write!(handle, "{i}")?; + for i in it { + write!(handle, "{sep}{i}")?; + } + } + Ok(()) + } +} + +fn print_multivariate_samples( + count: usize, + samples: impl IntoIterator, +) -> io::Result<()> +where + T: IntoIterator, + S: Display, +{ + let mut handle = io::BufWriter::new(io::stdout()); + + for s in samples.into_iter().take(count) { + util::write_interspersed(&mut handle, s.into_iter(), ", ")?; + writeln!(&mut handle)?; + } + Ok(()) +} + +fn print_samples(count: usize, samples: impl IntoIterator) -> io::Result<()> +where + T: Display, +{ + let mut handle = io::BufWriter::new(io::stdout()); + util::write_interspersed(&mut handle, samples.into_iter().take(count), "\n")?; + Ok(()) +} diff --git a/src/distribution/multinomial.rs b/src/distribution/multinomial.rs index f5b88945..90ff64a9 100644 --- a/src/distribution/multinomial.rs +++ b/src/distribution/multinomial.rs @@ -203,6 +203,23 @@ where res } +impl Mode> for Multinomial +where + D: Dim, + nalgebra::DefaultAllocator: nalgebra::allocator::Allocator, + nalgebra::DefaultAllocator: nalgebra::allocator::Allocator, +{ + fn mode(&self) -> OVector { + let n = self.n() as f64; + let (nr, nc) = self.p.shape_generic(); + OVector::::from_iterator_generic( + nr, + nc, + self.p.iter().map(|&p| ((n + 1.0) * p).floor() as u64), + ) + } +} + impl MeanN> for Multinomial where D: Dim,