wb_cache/test/simulation/scriptwriter/
math.rs1use 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#[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 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}