1#![allow(missing_docs)]
2mod meanvar;
3pub use meanvar::{NanHandling, col_mean, col_varm, row_mean, row_varm};
4pub mod prelude {
5 pub use super::ComplexDistribution;
6 #[cfg(feature = "rand")]
7 pub use super::{
8 CwiseColDistribution, CwiseMatDistribution, CwiseRowDistribution,
9 DistributionExt, UnitaryMat,
10 };
11 #[cfg(feature = "rand")]
12 pub use rand::prelude::*;
13 #[cfg(feature = "rand")]
14 pub use rand_distr::{Normal, StandardNormal, StandardUniform};
15}
16#[derive(Clone, Copy, Debug)]
18pub struct ComplexDistribution<Re, Im = Re> {
19 re: Re,
20 im: Im,
21}
22impl<Re, Im> ComplexDistribution<Re, Im> {
23 pub fn new(re: Re, im: Im) -> Self {
26 ComplexDistribution { re, im }
27 }
28}
29#[cfg(feature = "rand")]
30pub use self::rand::*;
31#[cfg(feature = "rand")]
32mod rand {
33 use super::ComplexDistribution;
34 use crate::internal_prelude::*;
35 use rand::Rng;
36 use rand::distr::Distribution;
37 pub trait DistributionExt {
38 fn rand<T>(&self, rng: &mut (impl ?Sized + rand::Rng)) -> T
39 where
40 Self: Distribution<T>,
41 {
42 self.sample(rng)
43 }
44 }
45 impl<T: ?Sized> DistributionExt for T {}
46 #[derive(Copy, Clone, Debug)]
47 pub struct CwiseMatDistribution<Rows: Shape, Cols: Shape, D> {
48 pub nrows: Rows,
49 pub ncols: Cols,
50 pub dist: D,
51 }
52 #[derive(Copy, Clone, Debug)]
53 pub struct CwiseColDistribution<Rows: Shape, D> {
54 pub nrows: Rows,
55 pub dist: D,
56 }
57 #[derive(Copy, Clone, Debug)]
58 pub struct CwiseRowDistribution<Cols: Shape, D> {
59 pub ncols: Cols,
60 pub dist: D,
61 }
62 #[derive(Copy, Clone, Debug)]
63 pub struct UnitaryMat<Dim: Shape, D> {
64 pub dim: Dim,
65 pub standard_normal: D,
66 }
67 impl<T, Rows: Shape, Cols: Shape, D: Distribution<T>>
68 Distribution<Mat<T, Rows, Cols>> for CwiseMatDistribution<Rows, Cols, D>
69 {
70 #[inline]
71 fn sample<R: rand::Rng + ?Sized>(
72 &self,
73 rng: &mut R,
74 ) -> Mat<T, Rows, Cols> {
75 Mat::from_fn(self.nrows, self.ncols, |_, _| self.dist.sample(rng))
76 }
77 }
78 impl<T, Rows: Shape, D: Distribution<T>> Distribution<Col<T, Rows>>
79 for CwiseColDistribution<Rows, D>
80 {
81 #[inline]
82 fn sample<R: rand::Rng + ?Sized>(&self, rng: &mut R) -> Col<T, Rows> {
83 Col::from_fn(self.nrows, |_| self.dist.sample(rng))
84 }
85 }
86 impl<T, Cols: Shape, D: Distribution<T>> Distribution<Row<T, Cols>>
87 for CwiseRowDistribution<Cols, D>
88 {
89 #[inline]
90 fn sample<R: rand::Rng + ?Sized>(&self, rng: &mut R) -> Row<T, Cols> {
91 Row::from_fn(self.ncols, |_| self.dist.sample(rng))
92 }
93 }
94 impl<T: ComplexField, D: Distribution<T>> Distribution<Mat<T>>
95 for UnitaryMat<usize, D>
96 {
97 fn sample<R: rand::prelude::Rng + ?Sized>(
98 &self,
99 rng: &mut R,
100 ) -> Mat<T> {
101 let qr = CwiseMatDistribution {
102 nrows: self.dim,
103 ncols: self.dim,
104 dist: &self.standard_normal,
105 }
106 .sample(rng)
107 .qr();
108 let r_diag = qr.R().diagonal().column_vector();
109 let mut q = qr.compute_Q();
110 for j in 0..self.dim {
111 let ref r = r_diag[j];
112 let ref r = if *r == zero() {
113 one()
114 } else {
115 r.mul_real(r.abs().recip())
116 };
117 z!(q.as_mut().col_mut(j)).for_each(|uz!(q)| {
118 *q *= r;
119 });
120 }
121 q
122 }
123 }
124 impl<T, Re, Im> Distribution<Complex<T>> for ComplexDistribution<Re, Im>
125 where
126 Re: Distribution<T>,
127 Im: Distribution<T>,
128 {
129 fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Complex<T> {
130 Complex::new(self.re.sample(rng), self.im.sample(rng))
131 }
132 }
133}