use serde::{Deserialize, Serialize};
use taktwerk_core::rng::{Rng, derive};
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
#[serde(rename_all = "snake_case")]
pub enum Dist {
Uniform([f64; 2]),
LogUniform([f64; 2]),
Normal([f64; 2]),
Choice(Vec<toml::Value>),
}
impl Dist {
pub fn check(&self) -> Result<(), String> {
let finite = |v: &[f64]| v.iter().all(|x| x.is_finite());
match self {
Self::Uniform([a, b]) if finite(&[*a, *b]) && a <= b => Ok(()),
Self::LogUniform([a, b]) if finite(&[*a, *b]) && 0.0 < *a && a <= b => Ok(()),
Self::Normal([m, s]) if finite(&[*m, *s]) && *s >= 0.0 => Ok(()),
Self::Choice(v) if !v.is_empty() => Ok(()),
Self::Uniform(_) => Err("uniform = [a, b] needs finite a <= b".into()),
Self::LogUniform(_) => Err("log_uniform = [a, b] needs 0 < a <= b".into()),
Self::Normal(_) => Err("normal = [mean, sd] needs a finite mean and sd >= 0".into()),
Self::Choice(_) => Err("choice needs at least one value".into()),
}
}
#[must_use]
#[allow(
clippy::cast_possible_truncation,
clippy::cast_sign_loss,
clippy::cast_precision_loss,
reason = "an index from a probability"
)]
pub fn at(&self, p: f64) -> toml::Value {
match self {
Self::Uniform([a, b]) => (a + (b - a) * p).into(),
Self::LogUniform([a, b]) => (a.ln() + (b.ln() - a.ln()) * p).exp().into(),
Self::Normal([m, s]) => (m + s * normal_quantile(p)).into(),
Self::Choice(v) => {
let i = ((p * v.len() as f64) as usize).min(v.len() - 1);
v[i].clone()
}
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum Method {
Uniform,
#[default]
Lhs,
}
#[must_use]
#[allow(
clippy::cast_possible_truncation,
clippy::cast_precision_loss,
reason = "indices and counts"
)]
pub fn probabilities(method: Method, n: usize, seed: u64, name: &str) -> Vec<f64> {
let mut rng = Rng::new(derive(seed, &format!("sample {name}")));
match method {
Method::Uniform => (0..n).map(|_| rng.uniform()).collect(),
Method::Lhs => {
let mut strata: Vec<usize> = (0..n).collect();
for i in (1..n).rev() {
let j = (rng.next_u64() % (i as u64 + 1)) as usize;
strata.swap(i, j);
}
strata
.into_iter()
.map(|k| (k as f64 + rng.uniform()) / n as f64)
.collect()
}
}
}
#[must_use]
pub fn normal_quantile(p: f64) -> f64 {
const A: [f64; 6] = [
-3.969_683_028_665_376e1,
2.209_460_984_245_205e2,
-2.759_285_104_469_687e2,
1.383_577_518_672_69e2,
-3.066_479_806_614_716e1,
2.506_628_277_459_239,
];
const B: [f64; 5] = [
-5.447_609_879_822_406e1,
1.615_858_368_580_409e2,
-1.556_989_798_598_866e2,
6.680_131_188_771_972e1,
-1.328_068_155_288_572e1,
];
const C: [f64; 6] = [
-7.784_894_002_430_293e-3,
-3.223_964_580_411_365e-1,
-2.400_758_277_161_838,
-2.549_732_539_343_734,
4.374_664_141_464_968,
2.938_163_982_698_783,
];
const D: [f64; 4] = [
7.784_695_709_041_462e-3,
3.224_671_290_700_398e-1,
2.445_134_137_142_996,
3.754_408_661_907_416,
];
const LOW: f64 = 0.024_25;
if p <= 0.0 {
return f64::NEG_INFINITY;
}
if p >= 1.0 {
return f64::INFINITY;
}
let tail = |q: f64| {
(((((C[0] * q + C[1]) * q + C[2]) * q + C[3]) * q + C[4]) * q + C[5])
/ ((((D[0] * q + D[1]) * q + D[2]) * q + D[3]) * q + 1.0)
};
if p < LOW {
tail((-2.0 * p.ln()).sqrt())
} else if p > 1.0 - LOW {
-tail((-2.0 * (1.0 - p).ln()).sqrt())
} else {
let q = p - 0.5;
let r = q * q;
(((((A[0] * r + A[1]) * r + A[2]) * r + A[3]) * r + A[4]) * r + A[5]) * q
/ (((((B[0] * r + B[1]) * r + B[2]) * r + B[3]) * r + B[4]) * r + 1.0)
}
}
#[cfg(test)]
#[allow(clippy::unwrap_used, reason = "tests")]
mod tests {
use super::*;
#[test]
fn a_latin_hypercube_puts_one_draw_in_each_stratum() {
let p = probabilities(Method::Lhs, 50, 1, "a");
let mut strata: Vec<usize> = p.iter().map(|x| (x * 50.0) as usize).collect();
strata.sort_unstable();
assert_eq!(strata, (0..50).collect::<Vec<_>>());
assert_eq!(p, probabilities(Method::Lhs, 50, 1, "a"), "seeded");
assert_ne!(
p,
probabilities(Method::Lhs, 50, 1, "b"),
"a stream per variable"
);
assert_ne!(p, probabilities(Method::Lhs, 50, 2, "a"));
assert!(
probabilities(Method::Uniform, 9, 1, "a")
.iter()
.all(|x| (0.0..1.0).contains(x))
);
}
#[test]
fn distributions_map_probabilities() {
let f = |d: Dist, p| d.at(p).as_float().unwrap();
assert_eq!(f(Dist::Uniform([2.0, 4.0]), 0.25), 2.5);
assert!((f(Dist::LogUniform([0.5, 2.0]), 0.5) - 1.0).abs() < 1e-12);
assert!((f(Dist::Normal([10.0, 2.0]), 0.975) - (10.0 + 2.0 * 1.959_963_985)).abs() < 1e-6);
assert_eq!(f(Dist::Normal([0.0, 1.0]), 0.5), 0.0);
let c = Dist::Choice(vec![1.into(), 2.into(), 3.into()]);
assert_eq!(c.at(0.0), 1.into());
assert_eq!(c.at(0.999), 3.into());
assert!(Dist::LogUniform([0.0, 1.0]).check().is_err());
assert!(Dist::Uniform([2.0, 1.0]).check().is_err());
assert!(Dist::Choice(vec![]).check().is_err());
assert!((normal_quantile(0.001) + 3.090_232_306).abs() < 1e-6);
assert!((normal_quantile(0.999) - 3.090_232_306).abs() < 1e-6);
}
}