Skip to main content

faer/stats/
mod.rs

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/// A generic random value distribution for complex numbers.
17#[derive(Clone, Copy, Debug)]
18pub struct ComplexDistribution<Re, Im = Re> {
19	re: Re,
20	im: Im,
21}
22impl<Re, Im> ComplexDistribution<Re, Im> {
23	/// Creates a complex distribution from independent
24	/// distributions of the real and imaginary parts.
25	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}