use rand::distributions::Distribution;
use serde::{Deserialize, Serialize};
use rand_distr::{Beta, Exp, Gamma, LogNormal, Normal, Triangular, Uniform, Weibull};
use rand_distr::{Bernoulli, Geometric, Poisson, WeightedIndex};
use super::dynamic_rng::DynRng;
use crate::utils::errors::SimulationError;
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub enum Continuous {
Beta { alpha: f64, beta: f64 },
Exp { lambda: f64 },
Gamma { shape: f64, scale: f64 },
LogNormal { mu: f64, sigma: f64 },
Normal { mean: f64, std_dev: f64 },
Triangular { min: f64, max: f64, mode: f64 },
Uniform { min: f64, max: f64 },
Weibull { shape: f64, scale: f64 },
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub enum Boolean {
Bernoulli { p: f64 },
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub enum Discrete {
Geometric {
p: f64,
},
Poisson {
lambda: f64,
},
Uniform {
min: u64,
max: u64,
},
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub enum Index {
Uniform {
min: usize,
max: usize,
},
WeightedIndex {
weights: Vec<u64>,
},
}
impl Continuous {
pub fn random_variate(&mut self, uniform_rng: DynRng) -> Result<f64, SimulationError> {
let mut rng = (*uniform_rng).borrow_mut();
match self {
Continuous::Beta { alpha, beta } => Ok(Beta::new(*alpha, *beta)?.sample(&mut *rng)),
Continuous::Exp { lambda } => Ok(Exp::new(*lambda)?.sample(&mut *rng)),
Continuous::Gamma { shape, scale } => Ok(Gamma::new(*shape, *scale)?.sample(&mut *rng)),
Continuous::LogNormal { mu, sigma } => {
Ok(LogNormal::new(*mu, *sigma)?.sample(&mut *rng))
}
Continuous::Normal { mean, std_dev } => {
Ok(Normal::new(*mean, *std_dev)?.sample(&mut *rng))
}
Continuous::Triangular { min, max, mode } => {
Ok(Triangular::new(*min, *max, *mode)?.sample(&mut *rng))
}
Continuous::Uniform { min, max } => Ok(Uniform::new(*min, *max).sample(&mut *rng)),
Continuous::Weibull { shape, scale } => {
Ok(Weibull::new(*shape, *scale)?.sample(&mut *rng))
}
}
}
}
impl Boolean {
pub fn random_variate(&mut self, uniform_rng: DynRng) -> Result<bool, SimulationError> {
let mut rng = (*uniform_rng).borrow_mut();
match self {
Boolean::Bernoulli { p } => Ok(Bernoulli::new(*p)?.sample(&mut *rng)),
}
}
}
impl Discrete {
pub fn random_variate(&mut self, uniform_rng: DynRng) -> Result<u64, SimulationError> {
let mut rng = (*uniform_rng).borrow_mut();
match self {
Discrete::Geometric { p } => Ok(Geometric::new(*p)?.sample(&mut *rng)),
Discrete::Poisson { lambda } => Ok(Poisson::new(*lambda)?.sample(&mut *rng) as u64),
Discrete::Uniform { min, max } => Ok(Uniform::new(*min, *max).sample(&mut *rng)),
}
}
}
impl Index {
pub fn random_variate(&mut self, uniform_rng: DynRng) -> Result<usize, SimulationError> {
let mut rng = (*uniform_rng).borrow_mut();
match self {
Index::Uniform { min, max } => Ok(Uniform::new(*min, *max).sample(&mut *rng)),
Index::WeightedIndex { weights } => {
Ok(WeightedIndex::new(weights.clone())?.sample(&mut *rng))
}
}
}
}
#[cfg(test)]
mod tests {
use crate::input_modeling::dynamic_rng::default_rng;
use super::*;
enum RandomVariable {
Continuous(Continuous),
Discrete(Discrete),
}
enum ChiSquareTest {
Continuous {
variable: Continuous,
bin_mapping_fn: fn(f64) -> usize,
},
Boolean {
variable: Boolean,
bin_mapping_fn: fn(bool) -> usize,
},
Discrete {
variable: Discrete,
bin_mapping_fn: fn(u64) -> usize,
},
Index {
variable: Index,
bin_mapping_fn: fn(usize) -> usize,
},
}
fn empirical_mean(random_variable: &mut RandomVariable, sample_size: usize) -> f64 {
let uniform_rng = default_rng();
(0..sample_size)
.map(|_| match random_variable {
RandomVariable::Continuous(variable) => {
variable.random_variate(uniform_rng.clone()).unwrap()
}
RandomVariable::Discrete(variable) => {
variable.random_variate(uniform_rng.clone()).unwrap() as f64
}
})
.sum::<f64>()
/ (sample_size as f64)
}
fn chi_square(test: &mut ChiSquareTest, expected_counts: &[usize]) -> f64 {
let mut class_counts = vec![0; expected_counts.len()];
let uniform_rng = default_rng();
let sample_size = expected_counts.iter().sum();
(0..sample_size).for_each(|_| {
let index = match test {
ChiSquareTest::Continuous {
variable,
bin_mapping_fn,
} => bin_mapping_fn(variable.random_variate(uniform_rng.clone()).unwrap()),
ChiSquareTest::Boolean {
variable,
bin_mapping_fn,
} => bin_mapping_fn(variable.random_variate(uniform_rng.clone()).unwrap()),
ChiSquareTest::Discrete {
variable,
bin_mapping_fn,
} => bin_mapping_fn(variable.random_variate(uniform_rng.clone()).unwrap()),
ChiSquareTest::Index {
variable,
bin_mapping_fn,
} => bin_mapping_fn(variable.random_variate(uniform_rng.clone()).unwrap()),
};
class_counts[index] += 1
});
class_counts.iter().zip(expected_counts.iter()).fold(
0.0,
|acc, (class_count, expected_count)| {
let f_class_count = *class_count as f64;
let f_expected_count = *expected_count as f64;
acc + (f_class_count - f_expected_count).powi(2) / f_expected_count
},
)
}
#[test]
fn beta_samples_match_expectation() {
let variable = Continuous::Beta {
alpha: 7.0,
beta: 11.0,
};
let mean = empirical_mean(&mut RandomVariable::Continuous(variable), 10000);
let expected = 7.0 / (7.0 + 11.0);
assert!((mean - expected).abs() / expected < 0.025);
}
#[test]
fn exponential_samples_match_expectation() {
let variable = Continuous::Exp { lambda: 7.0 };
let mean = empirical_mean(&mut RandomVariable::Continuous(variable), 10000);
let expected = 1.0 / 7.0;
assert!((mean - expected).abs() / expected < 0.025);
}
#[test]
fn gamma_samples_match_expectation() {
let variable = Continuous::Gamma {
shape: 7.0,
scale: 11.0,
};
let mean = empirical_mean(&mut RandomVariable::Continuous(variable), 10000);
let expected = 77.0;
assert!((mean - expected).abs() / expected < 0.025);
}
#[test]
fn lognormal_samples_match_expectation() {
let variable = Continuous::LogNormal {
mu: 11.0,
sigma: 1.0,
};
let mean = empirical_mean(&mut RandomVariable::Continuous(variable), 10000);
let expected = (11.0f64 + 1.0f64.powi(2) / 2.0f64).exp();
assert!((mean - expected).abs() / expected < 0.025);
}
#[test]
fn normal_samples_chi_square() {
fn bins_mapping(variate: f64) -> usize {
let mean = 11.0;
let std_dev = 3.0;
if variate < mean - 3.0 * std_dev {
0
} else if variate < mean - 2.0 * std_dev {
1
} else if variate < mean - std_dev {
2
} else if variate < mean {
3
} else if variate < mean + std_dev {
4
} else if variate < mean + 2.0 * std_dev {
5
} else if variate < mean + 3.0 * std_dev {
6
} else {
7
}
}
let variable = Continuous::Normal {
mean: 11.0,
std_dev: 3.0,
};
let expected_counts: [usize; 8] = [20, 210, 1360, 3410, 3410, 1360, 210, 20];
let chi_square_actual = chi_square(
&mut ChiSquareTest::Continuous {
variable: variable,
bin_mapping_fn: bins_mapping,
},
&expected_counts,
);
let chi_square_critical = 18.475;
assert![chi_square_actual < chi_square_critical];
}
#[test]
fn triangular_samples_chi_square() {
fn bins_mapping(variate: f64) -> usize {
((variate - 5.0) / 5.0) as usize
}
let variable = Continuous::Triangular {
min: 5.0,
max: 25.0,
mode: 15.0,
};
let expected_counts: [usize; 4] = [125, 375, 375, 125];
let chi_square_actual = chi_square(
&mut ChiSquareTest::Continuous {
variable: variable,
bin_mapping_fn: bins_mapping,
},
&expected_counts,
);
let chi_square_critical = 11.345;
assert![chi_square_actual < chi_square_critical];
}
#[test]
fn continuous_uniform_samples_chi_square() {
fn bins_mapping(variate: f64) -> usize {
let min = 7.0;
let max = 11.0;
((variate - min) * (max - 1.0)) as usize
}
let variable = Continuous::Uniform {
min: 7.0,
max: 11.0,
};
let expected_counts: [usize; 40] = [250; 40];
let chi_square_actual = chi_square(
&mut ChiSquareTest::Continuous {
variable: variable,
bin_mapping_fn: bins_mapping,
},
&expected_counts,
);
let chi_square_critical = 62.428;
assert![chi_square_actual < chi_square_critical];
}
#[test]
fn weibull_samples_match_expectation() {
let variable = Continuous::Weibull {
shape: 7.0,
scale: 0.5,
};
let mean = empirical_mean(&mut RandomVariable::Continuous(variable), 10000);
let expected = 14.0;
assert!((mean - expected).abs() / expected < 0.025);
}
#[test]
fn bernoulli_samples_chi_square() {
fn bins_mapping(variate: bool) -> usize {
variate as usize
}
let variable = Boolean::Bernoulli { p: 0.3 };
let expected_counts: [usize; 2] = [7000, 3000];
let chi_square_actual = chi_square(
&mut ChiSquareTest::Boolean {
variable: variable,
bin_mapping_fn: bins_mapping,
},
&expected_counts,
);
let chi_square_critical = 6.635;
assert![chi_square_actual < chi_square_critical];
}
#[test]
fn geometric_samples_match_expectation() {
let variable = Discrete::Geometric { p: 0.2 };
let mean = empirical_mean(&mut RandomVariable::Discrete(variable), 10000);
let expected = (1.0 - 0.2) / 0.2;
assert!((mean - expected).abs() / expected < 0.025);
}
#[test]
fn poisson_samples_match_expectation() {
let variable = Discrete::Poisson { lambda: 7.0 };
let mean = empirical_mean(&mut RandomVariable::Discrete(variable), 10000);
let expected = 7.0;
assert!((mean - expected).abs() / expected < 0.025);
}
#[test]
fn discrete_uniform_samples_chi_square() {
fn bins_mapping(variate: u64) -> usize {
let min = 7;
(variate - min) as usize
}
let variable = Discrete::Uniform { min: 7, max: 11 };
let expected_counts: [usize; 4] = [2500; 4];
let chi_square_actual = chi_square(
&mut ChiSquareTest::Discrete {
variable: variable,
bin_mapping_fn: bins_mapping,
},
&expected_counts,
);
let chi_square_critical = 13.277;
assert![chi_square_actual < chi_square_critical];
}
#[test]
fn weighted_index_samples_chi_square() {
fn bins_mapping(variate: usize) -> usize {
variate
}
let variable = Index::WeightedIndex {
weights: vec![1, 2, 3, 4],
};
let expected_counts: [usize; 4] = [1000, 2000, 3000, 4000];
let chi_square_actual = chi_square(
&mut ChiSquareTest::Index {
variable: variable,
bin_mapping_fn: bins_mapping,
},
&expected_counts,
);
let chi_square_critical = 11.345;
assert![chi_square_actual < chi_square_critical];
}
#[test]
fn index_uniform_samples_chi_square() {
fn bins_mapping(variate: usize) -> usize {
let min = 7;
variate - min
}
let variable = Index::Uniform { min: 7, max: 11 };
let expected_counts: [usize; 4] = [2500; 4];
let chi_square_actual = chi_square(
&mut ChiSquareTest::Index {
variable: variable,
bin_mapping_fn: bins_mapping,
},
&expected_counts,
);
let chi_square_critical = 13.277;
assert![chi_square_actual < chi_square_critical];
}
}