multivariate_optimization/
distributions.rs

1//! Univariate and multivariate probability distributions.
2
3pub use crate::triangular::Triangular;
4
5use num::traits::{Float, FloatConst, NumAssignOps, NumCast};
6use rand::{
7    distributions::{uniform::SampleUniform, Distribution, OpenClosed01},
8    Rng,
9};
10
11/// Numbers supported by generic items of this module.
12pub trait Num
13where
14    Self: Float + FloatConst + NumCast + NumAssignOps,
15    Self: SampleUniform,
16{
17    /// Generate number between zero (exclusive) and one (inclusive).
18    ///
19    /// This method is needed due to Rust issue [#20671]:
20    /// it avoids having to add `OpenClosed01: Distribution<T>` bounds.
21    ///
22    /// [#20671]: https://github.com/rust-lang/rust/issues/20671
23    fn sample_open_closed_01<R: Rng + ?Sized>(rng: &mut R) -> Self;
24}
25
26impl<T> Num for T
27where
28    T: Float + FloatConst + NumCast + NumAssignOps,
29    T: SampleUniform,
30    OpenClosed01: Distribution<T>,
31{
32    fn sample_open_closed_01<R: Rng + ?Sized>(rng: &mut R) -> Self {
33        OpenClosed01.sample(rng)
34    }
35}
36
37/// Multivariate distribution.
38pub trait MultivarDist<T>: Distribution<Vec<T>> {
39    /// Dimensionality
40    fn dim(&self) -> usize;
41}
42
43/// Univariate standard normal distribtion.
44#[derive(Clone, Copy, Default, Debug)]
45pub struct StdNormDist;
46
47impl<T: Num> Distribution<T> for StdNormDist {
48    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> T {
49        let pi: T = T::PI();
50        let r = (T::from(-2.0).unwrap() * T::ln(T::sample_open_closed_01(rng))).sqrt();
51        let sin = rng.gen_range(-pi..pi).sin();
52        r * sin
53    }
54}
55
56/// Univariate standard normal distribtion with pairwise output.
57#[derive(Clone, Copy, Default, Debug)]
58pub struct StdNormDistPair;
59
60impl<T: Num> Distribution<[T; 2]> for StdNormDistPair {
61    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> [T; 2] {
62        let pi: T = T::PI();
63        let r = (T::from(-2.0).unwrap() * T::ln(T::sample_open_closed_01(rng))).sqrt();
64        let (sin, cos) = rng.gen_range(-pi..pi).sin_cos();
65        [r * sin, r * cos]
66    }
67}
68
69impl<T: Num> Distribution<(T, T)> for StdNormDistPair {
70    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> (T, T) {
71        let [x, y]: [T; 2] = self.sample(rng);
72        (x, y)
73    }
74}
75
76/// Univariate standard normal distribtion with variable length output.
77#[derive(Clone, Copy, Default, Debug)]
78pub struct StdNormDistVec(
79    /// Number of elements in [`Vec`] sample.
80    pub usize,
81);
82
83impl<T: Num> Distribution<Vec<T>> for StdNormDistVec {
84    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Vec<T> {
85        let dim = self.0;
86        let mut vec = Vec::with_capacity(dim);
87        for _ in 0..dim / 2 {
88            let (x, y) = StdNormDistPair.sample(rng);
89            vec.push(x);
90            vec.push(y);
91        }
92        if dim % 2 == 1 {
93            vec.push(StdNormDist.sample(rng));
94        }
95        vec
96    }
97}
98
99/// Univariate (non-standard) normal distribtion with given average and
100/// standard deviation.
101#[derive(Clone, Debug)]
102pub struct NormDist<T> {
103    /// Average of created samples.
104    pub average: T,
105    /// Standard deviation of created samples.
106    pub stddev: T,
107}
108
109impl<T> NormDist<T> {
110    /// Create distribution with given average and standard deviation.
111    pub const fn new(average: T, stddev: T) -> NormDist<T> {
112        NormDist { average, stddev }
113    }
114}
115
116impl<T: Num> Distribution<T> for NormDist<T> {
117    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> T {
118        let x: T = StdNormDist.sample(rng);
119        x * self.stddev + self.average
120    }
121}
122
123/// Multivariate normal distribution for given
124/// average vectors and covariance matrices.
125///
126/// # Example
127///
128/// ```
129/// use multivariate_optimization::distributions::MultivarNormDist;
130/// use multivariate_optimization::triangular::Triangular;
131/// use rand::{thread_rng, Rng};
132///
133/// let averages: Vec<f64> = vec![1.51, 1.57, -0.29];
134/// let covariances = Triangular::<f64>::new(3, |(i, j)| {
135///     // needs to handle only cases for i >= j
136///     // assert!(i >= j);
137///     if i == j { 1.0 } else { 0.5 }
138/// });
139/// let dist = MultivarNormDist::new(averages, covariances);
140///
141/// let mut rng = thread_rng();
142/// let value = rng.sample(dist);
143/// println!("Sample vector: {:?}", value);
144/// ```
145#[derive(Debug)]
146pub struct MultivarNormDist<T> {
147    averages: Vec<T>,
148    factors: Triangular<T>,
149}
150
151impl<T: Num> MultivarNormDist<T> {
152    /// Create multivariate normal distribution for given averages and
153    /// covariances.
154    pub fn new(averages: Vec<T>, covariances: Triangular<T>) -> MultivarNormDist<T> {
155        let dim = covariances.dim();
156        assert_eq!(
157            averages.len(),
158            dim,
159            "dimensions of averages and covariances do not match"
160        );
161        let mut factors = covariances;
162        for i in 0..dim {
163            for j in 0..=i {
164                // SAFETY:
165                //  *  `i` is smaller than `factors.len()`
166                //  *  `j` is equal to or smaller than `i`
167                //  *  `k` is equal to or smaller than `j`
168                unsafe {
169                    let mut sum = *factors.get_unchecked((i, j));
170                    for k in 0..j {
171                        sum -= *factors.get_unchecked((i, k)) * *factors.get_unchecked((j, k));
172                    }
173                    let v = if i == j {
174                        sum.sqrt()
175                    } else {
176                        sum / *factors.get_unchecked((j, j))
177                    };
178                    *factors.get_unchecked_mut((i, j)) = if v.is_finite() { v } else { T::zero() }
179                }
180            }
181        }
182        MultivarNormDist { averages, factors }
183    }
184}
185
186impl<T: Num> MultivarDist<T> for MultivarNormDist<T> {
187    fn dim(&self) -> usize {
188        self.factors.dim()
189    }
190}
191
192impl<T: Num> Distribution<Vec<T>> for MultivarNormDist<T> {
193    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Vec<T> {
194        let dim = self.dim();
195        let mut result: Vec<T> = StdNormDistVec(dim).sample(rng);
196        // SAFETY:
197        //  *  `i` is smaller than `self.dim()` and thus smaller
198        //     than `result.len()` and `self.factors.dim()`
199        //  *  `j` is equal to or smaller than `i`
200        unsafe {
201            for i in (0..dim).rev() {
202                let mut sum = T::zero();
203                for j in 0..=i {
204                    sum += *result.get_unchecked(j) * *self.factors.get_unchecked((i, j));
205                }
206                *result.get_unchecked_mut(i) = sum;
207            }
208            for i in 0..dim {
209                *result.get_unchecked_mut(i) += *self.averages.get_unchecked(i);
210            }
211        }
212        result
213    }
214}
215
216#[cfg(test)]
217mod tests {
218    use super::{MultivarNormDist, StdNormDist};
219    use crate::triangular::Triangular;
220    use rand::{distributions::Distribution, thread_rng};
221    const STATCNT: usize = 10000;
222    #[test]
223    fn test_std_norm_dist() {
224        let values: Vec<f64> = StdNormDist
225            .sample_iter(thread_rng())
226            .take(STATCNT)
227            .collect();
228        let (mut avg, mut var): (f64, f64) = (0.0, 0.0);
229        for value in values.iter() {
230            avg += value;
231            var += value * value;
232        }
233        avg /= STATCNT as f64;
234        var /= STATCNT as f64;
235        assert!(avg.abs() < 0.1);
236        assert!((var - 1.0).abs() < 0.1);
237    }
238    #[test]
239    fn test_multivar_norm_dist() {
240        let avg_goal: Vec<f64> = vec![0.0, 0.0, 0.0];
241        let cov_goal = Triangular::<f64>::new(3, |idx| match idx {
242            (0, 0) => 1.0,
243            (1, 1) => 0.5,
244            (2, 2) => 2.0,
245            (2, 0) => 0.25,
246            (0, 2) => 0.25,
247            _ => 0.0,
248        });
249        let cov_goal = cov_goal;
250        let dist = MultivarNormDist::new(avg_goal.clone(), cov_goal.clone());
251        let values: Vec<Vec<f64>> = dist.sample_iter(thread_rng()).take(STATCNT).collect();
252        let mut avg: Vec<f64> = vec![0.0; 3];
253        for value in values.iter() {
254            for i in 0..3 {
255                avg[i] += value[i];
256            }
257        }
258        for i in 0..3 {
259            avg[i] /= STATCNT as f64;
260        }
261        let mut cov = Triangular::<f64>::new(3, |(i, j)| {
262            let mut sum = 0.0;
263            for value in values.iter() {
264                sum += (value[i] - avg[i]) * (value[j] - avg[j]);
265            }
266            sum
267        });
268        for i in 0..3 {
269            for j in 0..=i {
270                cov[(i, j)] /= STATCNT as f64;
271            }
272        }
273        for i in 0..3 {
274            assert!((avg[i] - avg_goal[i]).abs() < 0.1);
275            for j in 0..=i {
276                assert!((cov[(i, j)] - cov_goal[(i, j)]).abs() < 0.1);
277            }
278        }
279    }
280}
multivariate_optimization/distributions.rs

multivariate_optimization/
distributions.rs
