wb_cache/test/simulation/scriptwriter/
math.rs

1use anyhow::Result;
2use fieldx::fxstruct;
3use rand::Rng;
4use statrs::distribution::ContinuousCDF;
5use statrs::distribution::Normal;
6
7pub fn bisect(v0: f64, v1: f64, expected: f64, tolerance: f64, f: impl Fn(f64) -> f64) -> Result<f64> {
8    let mut v0 = v0;
9    let mut v1 = v1;
10
11    for _ in 0..1000 {
12        let v2 = (v0 + v1) / 2.0;
13        let f2 = f(v2);
14        let diff = f2 - expected;
15
16        if diff.abs() < tolerance {
17            return Ok(v2);
18        }
19        if f2 < expected {
20            v0 = v2;
21        }
22        else {
23            v1 = v2;
24        }
25
26        if v0 == v1 {
27            break;
28        }
29    }
30
31    Err(anyhow::anyhow!("Bisection failed to converge"))
32}
33
34/// Skew normal distribution truncated to [a, b]
35#[fxstruct(no_new, builder(post_build), get(copy))]
36pub struct TruncatedSkewNormal {
37    location: f64,
38    scale:    f64,
39    shape:    f64,
40    lower:    f64,
41    upper:    f64,
42
43    #[fieldx(private, get(off), builder(off))]
44    cdf_a: f64,
45    #[fieldx(private, get(off), builder(off))]
46    cdf_b: f64,
47}
48
49impl TruncatedSkewNormal {
50    fn post_build(mut self) -> Self {
51        self.cdf_a = self.skew_cdf(self.lower);
52        self.cdf_b = self.skew_cdf(self.upper);
53        self
54    }
55
56    fn skew_cdf(&self, x: f64) -> f64 {
57        let z = (x - self.location) / self.scale;
58        let normal = Normal::new(0.0, 1.0).unwrap();
59        let phi = normal.cdf(z);
60        let psi = normal.cdf(self.shape * z);
61        2.0 * phi * psi
62    }
63
64    fn skew_ppf(&self, p: f64) -> f64 {
65        // Invert CDF via binary search
66        let mut lo = self.lower;
67        let mut hi = self.upper;
68        for _ in 0..100 {
69            let mid = (lo + hi) / 2.0;
70            let cdf = self.skew_cdf(mid);
71            if cdf < p {
72                lo = mid;
73            }
74            else {
75                hi = mid;
76            }
77        }
78        (lo + hi) / 2.0
79    }
80
81    pub fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> f64 {
82        let u = rng.random_range(self.cdf_a..self.cdf_b);
83        self.skew_ppf(u)
84    }
85}