use std::str::FromStr;
use strafe_testing::{
r::{RString, RTester},
r_assert_relative_equal, r_assert_relative_equal_result,
};
use strafe_type::{tof64, Checked, Positive64, Real64};
use crate::{
distribution::multinom::MultinomialBuilder,
reset_statics,
rng::MarsagliaMulticarry,
traits::{Distribution, RNG},
};
pub fn density_inner(
x: Vec<Checked<Real64>>,
prob: Option<Vec<Checked<Positive64>>>,
size: Option<usize>,
) {
let mut builder = MultinomialBuilder::new();
if let Some(prob) = prob {
builder.set_probability(prob);
} else if let Some(size) = size {
builder.set_size(size);
}
let multinomial = builder.build();
let xs = x.iter().map(|x| tof64!(x)).collect::<Vec<_>>();
let ps = multinomial
.probability
.iter()
.map(|p| tof64!(p))
.collect::<Vec<_>>();
let rust_ans = multinomial.density(x);
let r_ans = RTester::new()
.set_display(&RString::from_string(format!(
"dmultinom({}, prob={}, log={})",
RString::from_f64_slice(&xs),
RString::from_f64_slice(&ps),
RString::from_bool(false),
)))
.run()
.expect("R running error")
.as_f64();
r_assert_relative_equal_result!(r_ans, rust_ans);
}
pub fn log_density_inner(
x: Vec<Checked<Real64>>,
prob: Option<Vec<Checked<Positive64>>>,
size: Option<usize>,
) {
let mut builder = MultinomialBuilder::new();
if let Some(prob) = prob {
builder.set_probability(prob);
} else if let Some(size) = size {
builder.set_size(size);
}
let multinomial = builder.build();
let xs = x.iter().map(|x| tof64!(x)).collect::<Vec<_>>();
let ps = multinomial
.probability
.iter()
.map(|p| tof64!(p))
.collect::<Vec<_>>();
let rust_ans = multinomial.log_density(x);
let r_ans = RTester::new()
.set_display(&RString::from_string(format!(
"dmultinom({}, prob={}, log={})",
RString::from_f64_slice(&xs),
RString::from_f64_slice(&ps),
RString::from_bool(true),
)))
.run()
.expect("R running error")
.as_f64();
r_assert_relative_equal_result!(r_ans, rust_ans);
}
pub fn random_inner(seed: u16, prob: Option<Vec<Checked<Positive64>>>, size: Option<usize>) {
let num = 50;
let mut builder = MultinomialBuilder::new();
if let Some(prob) = prob {
builder.set_probability(prob);
} else if let Some(size) = size {
builder.set_size(size);
}
let multinomial = builder.build();
let ps = multinomial
.probability
.iter()
.map(|p| tof64!(p))
.collect::<Vec<_>>();
let r_state = RTester::new()
.set_seed(seed)
.set_display(&RString::from_str(".Random.seed").unwrap())
.run()
.unwrap()
.as_f64_vec()
.unwrap()
.into_iter()
.skip(1)
.map(|i| i as i32)
.collect::<Vec<_>>();
let mut rust_rng = MarsagliaMulticarry::new();
rust_rng.set_seed(seed as u32);
let rust_state = rust_rng.get_state();
for (r, rust) in r_state.iter().zip(rust_state.iter()) {
assert_eq!(r, rust, "States not equal!");
}
let r_ret = RTester::new()
.set_seed(seed)
.set_display(&RString::from_string(format!(
"rmultinom({}, prob={}, size={})",
RString::from_f64(num as f64),
RString::from_f64_slice(&ps),
RString::from_f64(multinomial.probability.len()),
)))
.run()
.unwrap()
.as_f64_matrix()
.unwrap()
.column_iter()
.map(|r_i| r_i.data.into_slice().to_vec())
.collect::<Vec<_>>();
reset_statics();
let rust_ret = (0..num)
.map(|_| multinomial.random_sample(&mut rust_rng))
.collect::<Vec<_>>();
for (r, rust) in r_ret.iter().zip(rust_ret.iter()) {
for (r_i, rust_i) in r.iter().zip(rust.iter()) {
r_assert_relative_equal!(*r_i, tof64!(rust_i));
}
}
}