taktwerk 0.6.0

Fixed-step model execution engine for Linux with its own OPC UA server
//! Random samples of a campaign: distributions, uniform and Latin hypercube draws.

use serde::{Deserialize, Serialize};
use taktwerk_core::rng::{Rng, derive};

/// How a sampled value is distributed.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
#[serde(rename_all = "snake_case")]
pub enum Dist {
    /// Uniform on `[a, b]`.
    Uniform([f64; 2]),
    /// `exp` of uniform on `[ln a, ln b]`; `0 < a < b`.
    LogUniform([f64; 2]),
    /// Normal with mean and standard deviation.
    Normal([f64; 2]),
    /// One of these values, equally likely.
    Choice(Vec<toml::Value>),
}

impl Dist {
    /// Check the parameters.
    ///
    /// # Errors
    /// Bounds out of order, a non-positive log bound or deviation, or an empty choice.
    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()),
        }
    }

    /// The value at probability `p` in `[0, 1)`.
    #[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()
            }
        }
    }
}

/// How the probabilities are drawn.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum Method {
    /// Independent uniform probabilities.
    Uniform,
    /// Latin hypercube: each variable's `n` probabilities fall one in each of `n` equal strata,
    /// strata shuffled independently per variable.
    #[default]
    Lhs,
}

/// `n` probabilities in `[0, 1)` for the variable `name` of a campaign seeded `seed`. Each
/// variable draws from its own stream, so adding one leaves the others' draws as they were.
#[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();
            // Fisher-Yates.
            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()
        }
    }
}

/// The standard normal quantile at `p`, by Acklam's rational approximation (relative error
/// below 1.2e-9); infinite at 0 and 1.
#[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);
    }
}