faer/stats/
meanvar.rs

1use crate::assert;
2use crate::internal_prelude::*;
3use faer_traits::RealReg;
4/// Specifies how missing values should be handled in mean and variance
5/// computations.
6#[derive(Copy, Clone, Debug, PartialEq, Eq)]
7pub enum NanHandling {
8	/// NaNs are passed as-is to arithmetic operators.
9	Propagate,
10	/// NaNs are skipped, and they're not included in the total count of
11	/// entries.
12	Ignore,
13}
14#[inline(always)]
15fn from_usize<T: RealField>(n: usize) -> T {
16	from_f64::<T>(n as u32 as f64)
17		+ from_f64::<T>((n as u64 - (n as u32 as u64)) as f64)
18}
19#[inline(always)]
20fn reduce<T: ComplexField, S: pulp::Simd>(
21	non_nan_count: T::SimdIndex<S>,
22) -> usize {
23	let slice: &[T::Index] =
24		bytemuck::cast_slice(core::slice::from_ref(&non_nan_count));
25	let mut acc = 0usize;
26	for &count in slice {
27		acc += count.zx();
28	}
29	acc
30}
31fn col_mean_row_major_ignore_nan<T: ComplexField>(
32	out: ColMut<'_, T>,
33	mat: MatRef<'_, T, usize, usize, isize, ContiguousFwd>,
34) {
35	struct Impl<'a, T: ComplexField> {
36		out: ColMut<'a, T>,
37		mat: MatRef<'a, T, usize, usize, isize, ContiguousFwd>,
38	}
39	impl<T: ComplexField> pulp::WithSimd for Impl<'_, T> {
40		type Output = ();
41
42		#[inline(always)]
43		fn with_simd<S: pulp::Simd>(self, simd: S) -> Self::Output {
44			let Self { out, mat } = self;
45			with_dim!(M, mat.nrows());
46			with_dim!(N, mat.ncols());
47			let simd = SimdCtx::<T, S>::new(T::simd_ctx(simd), N);
48			let mut out = out.as_row_shape_mut(M);
49			let mat = mat.as_shape(M, N);
50			let indices = simd.batch_indices::<4>();
51			let nan = simd.splat(&nan::<T>());
52			for i in M.indices() {
53				let row = mat.row(i).transpose();
54				let (head, mut body4, body1, tail) = indices.clone();
55				let mut non_nan_count_total = 0usize;
56				#[inline(always)]
57				fn process<'M, T: ComplexField, S: pulp::Simd>(
58					simd: SimdCtx<'M, T, S>,
59					acc: T::SimdVec<S>,
60					non_nan_count: T::SimdIndex<S>,
61					val: T::SimdVec<S>,
62				) -> (T::SimdVec<S>, T::SimdIndex<S>) {
63					let is_not_nan = (*simd).eq(val, val);
64					(
65						simd.select(is_not_nan, simd.add(acc, val), acc),
66						simd.iselect(
67							is_not_nan,
68							simd.iadd(
69								non_nan_count,
70								simd.isplat(T::Index::truncate(1)),
71							),
72							non_nan_count,
73						),
74					)
75				}
76				let mut sum0 = simd.splat(&zero::<T>());
77				let mut sum1 = simd.splat(&zero::<T>());
78				let mut sum2 = simd.splat(&zero::<T>());
79				let mut sum3 = simd.splat(&zero::<T>());
80				let mut non_nan_count0 = simd.isplat(T::Index::truncate(0));
81				let mut non_nan_count1 = simd.isplat(T::Index::truncate(0));
82				let mut non_nan_count2 = simd.isplat(T::Index::truncate(0));
83				let mut non_nan_count3 = simd.isplat(T::Index::truncate(0));
84				if let Some(i) = head {
85					(sum0, non_nan_count0) = process(
86						simd,
87						sum0,
88						non_nan_count0,
89						simd.select(simd.head_mask(), simd.read(row, i), nan),
90					);
91					non_nan_count_total += reduce::<T, S>(non_nan_count0);
92					non_nan_count0 = simd.isplat(T::Index::truncate(0));
93				}
94				loop {
95					if body4.len() == 0 {
96						break;
97					}
98					for [i0, i1, i2, i3] in (&mut body4).take(256) {
99						(sum0, non_nan_count0) = process(
100							simd,
101							sum0,
102							non_nan_count0,
103							simd.read(row, i0),
104						);
105						(sum1, non_nan_count1) = process(
106							simd,
107							sum1,
108							non_nan_count1,
109							simd.read(row, i1),
110						);
111						(sum2, non_nan_count2) = process(
112							simd,
113							sum2,
114							non_nan_count2,
115							simd.read(row, i2),
116						);
117						(sum3, non_nan_count3) = process(
118							simd,
119							sum3,
120							non_nan_count3,
121							simd.read(row, i3),
122						);
123					}
124					non_nan_count0 = simd.iadd(non_nan_count0, non_nan_count1);
125					non_nan_count2 = simd.iadd(non_nan_count2, non_nan_count3);
126					non_nan_count0 = simd.iadd(non_nan_count0, non_nan_count2);
127					non_nan_count_total += reduce::<T, S>(non_nan_count0);
128					non_nan_count0 = simd.isplat(T::Index::truncate(0));
129					non_nan_count1 = simd.isplat(T::Index::truncate(0));
130					non_nan_count2 = simd.isplat(T::Index::truncate(0));
131					non_nan_count3 = simd.isplat(T::Index::truncate(0));
132				}
133				for i in body1 {
134					(sum0, non_nan_count0) =
135						process(simd, sum0, non_nan_count0, simd.read(row, i));
136				}
137				if let Some(i) = tail {
138					(sum0, non_nan_count0) = process(
139						simd,
140						sum0,
141						non_nan_count0,
142						simd.select(simd.tail_mask(), simd.read(row, i), nan),
143					);
144				}
145				non_nan_count_total += reduce::<T, S>(non_nan_count0);
146				sum0 = simd.add(sum0, sum1);
147				sum2 = simd.add(sum2, sum3);
148				sum0 = simd.add(sum0, sum2);
149				let sum = simd.reduce_sum(sum0);
150				out[i] = sum.mul_real(
151					from_usize::<T::Real>(non_nan_count_total).recip(),
152				);
153			}
154		}
155	}
156	T::Arch::default().dispatch(Impl { out, mat });
157}
158fn col_varm_row_major_ignore_nan<T: ComplexField>(
159	out: ColMut<'_, T::Real>,
160	mat: MatRef<'_, T, usize, usize, isize, ContiguousFwd>,
161	col_mean: ColRef<'_, T>,
162) {
163	struct Impl<'a, T: ComplexField> {
164		out: ColMut<'a, T::Real>,
165		mat: MatRef<'a, T, usize, usize, isize, ContiguousFwd>,
166		col_mean: ColRef<'a, T>,
167	}
168	impl<T: ComplexField> pulp::WithSimd for Impl<'_, T> {
169		type Output = ();
170
171		#[inline(always)]
172		fn with_simd<S: pulp::Simd>(self, simd: S) -> Self::Output {
173			let Self { out, mat, col_mean } = self;
174			with_dim!(M, mat.nrows());
175			with_dim!(N, mat.ncols());
176			let simd = SimdCtx::<T, S>::new(T::simd_ctx(simd), N);
177			let mut out = out.as_row_shape_mut(M);
178			let mat = mat.as_shape(M, N);
179			let col_mean = col_mean.as_row_shape(M);
180			let indices = simd.batch_indices::<4>();
181			let nan_v = simd.splat(&nan::<T>());
182			for i in M.indices() {
183				let row = mat.row(i).transpose();
184				let mean = simd.splat(&col_mean[i]);
185				let (head, mut body4, body1, tail) = indices.clone();
186				let mut non_nan_count = 0usize;
187				#[inline(always)]
188				fn process<'M, T: ComplexField, S: pulp::Simd>(
189					simd: SimdCtx<'M, T, S>,
190					mean: T::SimdVec<S>,
191					acc: RealReg<T::SimdVec<S>>,
192					non_nan_count: T::SimdIndex<S>,
193					val: T::SimdVec<S>,
194				) -> (RealReg<T::SimdVec<S>>, T::SimdIndex<S>) {
195					let is_not_nan = (*simd).eq(val, val);
196					let diff = simd.sub(val, mean);
197					(
198						RealReg(simd.select(
199							is_not_nan,
200							simd.abs2_add(diff, acc).0,
201							acc.0,
202						)),
203						simd.iselect(
204							is_not_nan,
205							simd.iadd(
206								non_nan_count,
207								simd.isplat(T::Index::truncate(1)),
208							),
209							non_nan_count,
210						),
211					)
212				}
213				let mut sum0 = RealReg(simd.splat(&zero::<T>()));
214				let mut sum1 = RealReg(simd.splat(&zero::<T>()));
215				let mut sum2 = RealReg(simd.splat(&zero::<T>()));
216				let mut sum3 = RealReg(simd.splat(&zero::<T>()));
217				let mut non_nan_count0 = simd.isplat(T::Index::truncate(0));
218				let mut non_nan_count1 = simd.isplat(T::Index::truncate(0));
219				let mut non_nan_count2 = simd.isplat(T::Index::truncate(0));
220				let mut non_nan_count3 = simd.isplat(T::Index::truncate(0));
221				if let Some(i) = head {
222					(sum0, non_nan_count0) = process(
223						simd,
224						mean,
225						sum0,
226						non_nan_count0,
227						simd.select(simd.head_mask(), simd.read(row, i), nan_v),
228					);
229					non_nan_count += reduce::<T, S>(non_nan_count0);
230					non_nan_count0 = simd.isplat(T::Index::truncate(0));
231				}
232				loop {
233					if body4.len() == 0 {
234						break;
235					}
236					for [i0, i1, i2, i3] in (&mut body4).take(256) {
237						(sum0, non_nan_count0) = process(
238							simd,
239							mean,
240							sum0,
241							non_nan_count0,
242							simd.read(row, i0),
243						);
244						(sum1, non_nan_count1) = process(
245							simd,
246							mean,
247							sum1,
248							non_nan_count1,
249							simd.read(row, i1),
250						);
251						(sum2, non_nan_count2) = process(
252							simd,
253							mean,
254							sum2,
255							non_nan_count2,
256							simd.read(row, i2),
257						);
258						(sum3, non_nan_count3) = process(
259							simd,
260							mean,
261							sum3,
262							non_nan_count3,
263							simd.read(row, i3),
264						);
265					}
266					non_nan_count0 = simd.iadd(non_nan_count0, non_nan_count1);
267					non_nan_count2 = simd.iadd(non_nan_count2, non_nan_count3);
268					non_nan_count0 = simd.iadd(non_nan_count0, non_nan_count2);
269					non_nan_count += reduce::<T, S>(non_nan_count0);
270					non_nan_count0 = simd.isplat(T::Index::truncate(0));
271					non_nan_count1 = simd.isplat(T::Index::truncate(0));
272					non_nan_count2 = simd.isplat(T::Index::truncate(0));
273					non_nan_count3 = simd.isplat(T::Index::truncate(0));
274				}
275				for i in body1 {
276					(sum0, non_nan_count0) = process(
277						simd,
278						mean,
279						sum0,
280						non_nan_count0,
281						simd.read(row, i),
282					);
283				}
284				if let Some(i) = tail {
285					(sum0, non_nan_count0) = process(
286						simd,
287						mean,
288						sum0,
289						non_nan_count0,
290						simd.select(simd.tail_mask(), simd.read(row, i), nan_v),
291					);
292				}
293				non_nan_count += reduce::<T, S>(non_nan_count0);
294				sum0 = RealReg(simd.add(sum0.0, sum1.0));
295				sum2 = RealReg(simd.add(sum2.0, sum3.0));
296				sum0 = RealReg(simd.add(sum0.0, sum2.0));
297				let sum = simd.reduce_sum(sum0.0).real();
298				if non_nan_count == 0 {
299					out[i] = nan();
300				} else if non_nan_count == 1 {
301					out[i] = zero();
302				} else {
303					out[i] = sum.mul_real(
304						from_usize::<T::Real>(non_nan_count - 1).recip(),
305					);
306				}
307			}
308		}
309	}
310	T::Arch::default().dispatch(Impl { out, mat, col_mean });
311}
312fn col_mean_row_major_propagate_nan<T: ComplexField>(
313	out: ColMut<'_, T>,
314	mat: MatRef<'_, T, usize, usize, isize, ContiguousFwd>,
315) {
316	struct Impl<'a, T: ComplexField> {
317		out: ColMut<'a, T>,
318		mat: MatRef<'a, T, usize, usize, isize, ContiguousFwd>,
319	}
320	impl<T: ComplexField> pulp::WithSimd for Impl<'_, T> {
321		type Output = ();
322
323		#[inline(always)]
324		fn with_simd<S: pulp::Simd>(self, simd: S) -> Self::Output {
325			let Self { out, mat } = self;
326			with_dim!(M, mat.nrows());
327			with_dim!(N, mat.ncols());
328			let simd = SimdCtx::<T, S>::new(T::simd_ctx(simd), N);
329			let mut out = out.as_row_shape_mut(M);
330			let mat = mat.as_shape(M, N);
331			let indices = simd.batch_indices::<4>();
332			let ref n = from_usize::<T::Real>(*N).recip();
333			for i in M.indices() {
334				let row = mat.row(i).transpose();
335				let (head, body4, body1, tail) = indices.clone();
336				let mut sum0 = simd.splat(zero::<T>());
337				let mut sum1 = simd.splat(zero::<T>());
338				let mut sum2 = simd.splat(zero::<T>());
339				let mut sum3 = simd.splat(zero::<T>());
340				if let Some(i) = head {
341					sum0 = simd.add(sum0, simd.read(row, i));
342				}
343				for [i0, i1, i2, i3] in body4 {
344					sum0 = simd.add(sum0, simd.read(row, i0));
345					sum1 = simd.add(sum1, simd.read(row, i1));
346					sum2 = simd.add(sum2, simd.read(row, i2));
347					sum3 = simd.add(sum3, simd.read(row, i3));
348				}
349				for i in body1 {
350					sum0 = simd.add(sum0, simd.read(row, i));
351				}
352				if let Some(i) = tail {
353					sum0 = simd.add(sum0, simd.read(row, i));
354				}
355				sum0 = simd.add(sum0, sum1);
356				sum2 = simd.add(sum2, sum3);
357				sum0 = simd.add(sum0, sum2);
358				let sum = simd.reduce_sum(sum0);
359				out[i] = sum.mul_real(n);
360			}
361		}
362	}
363	T::Arch::default().dispatch(Impl { out, mat });
364}
365fn col_varm_row_major_propagate_nan<T: ComplexField>(
366	out: ColMut<'_, T::Real>,
367	mat: MatRef<'_, T, usize, usize, isize, ContiguousFwd>,
368	col_mean: ColRef<'_, T>,
369) {
370	struct Impl<'a, T: ComplexField> {
371		out: ColMut<'a, T::Real>,
372		mat: MatRef<'a, T, usize, usize, isize, ContiguousFwd>,
373		col_mean: ColRef<'a, T>,
374	}
375	impl<T: ComplexField> pulp::WithSimd for Impl<'_, T> {
376		type Output = ();
377
378		#[inline(always)]
379		fn with_simd<S: pulp::Simd>(self, simd: S) -> Self::Output {
380			let Self { out, mat, col_mean } = self;
381			with_dim!(M, mat.nrows());
382			with_dim!(N, mat.ncols());
383			let simd = SimdCtx::<T, S>::new(T::simd_ctx(simd), N);
384			let mut out = out.as_row_shape_mut(M);
385			let mat = mat.as_shape(M, N);
386			let col_mean = col_mean.as_row_shape(M);
387			let indices = simd.batch_indices::<4>();
388			let n = *N;
389			if n == 0 {
390				out.fill(nan());
391			} else if n == 1 {
392				out.fill(zero());
393			} else {
394				let ref n = from_usize::<T::Real>(n - 1).recip();
395				for i in M.indices() {
396					let row = mat.row(i).transpose();
397					let mean = simd.splat(&col_mean[i]);
398					let (head, body4, body1, tail) = indices.clone();
399					let mut sum0 = simd.splat(zero::<T>());
400					let mut sum1 = simd.splat(zero::<T>());
401					let mut sum2 = simd.splat(zero::<T>());
402					let mut sum3 = simd.splat(zero::<T>());
403					if let Some(i0) = head {
404						sum0 = simd.select(
405							simd.head_mask(),
406							simd.abs2_add(
407								simd.sub(simd.read(row, i0), mean),
408								RealReg(sum0),
409							)
410							.0,
411							sum0,
412						);
413					}
414					for [i0, i1, i2, i3] in body4 {
415						sum0 = simd
416							.abs2_add(
417								simd.sub(simd.read(row, i0), mean),
418								RealReg(sum0),
419							)
420							.0;
421						sum1 = simd
422							.abs2_add(
423								simd.sub(simd.read(row, i1), mean),
424								RealReg(sum1),
425							)
426							.0;
427						sum2 = simd
428							.abs2_add(
429								simd.sub(simd.read(row, i2), mean),
430								RealReg(sum2),
431							)
432							.0;
433						sum3 = simd
434							.abs2_add(
435								simd.sub(simd.read(row, i3), mean),
436								RealReg(sum3),
437							)
438							.0;
439					}
440					for i0 in body1 {
441						sum0 = simd
442							.abs2_add(
443								simd.sub(simd.read(row, i0), mean),
444								RealReg(sum0),
445							)
446							.0;
447					}
448					if let Some(i0) = tail {
449						sum0 = simd.select(
450							simd.tail_mask(),
451							simd.abs2_add(
452								simd.sub(simd.read(row, i0), mean),
453								RealReg(sum0),
454							)
455							.0,
456							sum0,
457						);
458					}
459					sum0 = simd.add(sum0, sum1);
460					sum2 = simd.add(sum2, sum3);
461					sum0 = simd.add(sum0, sum2);
462					let sum = simd.reduce_sum(sum0).real();
463					out[i] = sum * n;
464				}
465			}
466		}
467	}
468	T::Arch::default().dispatch(Impl { out, mat, col_mean });
469}
470fn col_mean_ignore_nan_fallback<T: ComplexField>(
471	out: ColMut<'_, T>,
472	mat: MatRef<'_, T>,
473) {
474	with_dim!(M, mat.nrows());
475	with_dim!(N, mat.ncols());
476	let non_nan_count = &mut *alloc::vec![0usize; * M];
477	let non_nan_count = Array::from_mut(non_nan_count, M);
478	let mut out = out.as_row_shape_mut(M);
479	let mat = mat.as_shape(M, N);
480	out.fill(zero());
481	for j in N.indices() {
482		for i in M.indices() {
483			let val = mat[(i, j)].copy();
484			let nan = val.is_nan();
485			let val = if nan { zero::<T>() } else { val };
486			non_nan_count[i] += (!nan) as usize;
487			out[i] += val;
488		}
489	}
490	for i in M.indices() {
491		out[i] =
492			out[i].mul_real(from_usize::<T::Real>(non_nan_count[i]).recip());
493	}
494}
495fn col_varm_ignore_nan_fallback<T: ComplexField>(
496	out: ColMut<'_, T::Real>,
497	mat: MatRef<'_, T>,
498	col_mean: ColRef<'_, T>,
499) {
500	with_dim!(M, mat.nrows());
501	with_dim!(N, mat.ncols());
502	let non_nan_count = &mut *alloc::vec![0usize; * M];
503	let non_nan_count = Array::from_mut(non_nan_count, M);
504	let mut out = out.as_row_shape_mut(M);
505	let mat = mat.as_shape(M, N);
506	let col_mean = col_mean.as_row_shape(M);
507	out.fill(zero());
508	for j in N.indices() {
509		for i in M.indices() {
510			let val = mat[(i, j)].copy();
511			let col_mean = col_mean[i].copy();
512			let nan = val.is_nan();
513			let val = if nan {
514				zero::<T::Real>()
515			} else {
516				(val - col_mean).abs2()
517			};
518			non_nan_count[i] += (!nan) as usize;
519			out[i] += val;
520		}
521	}
522	for i in M.indices() {
523		let non_nan_count = non_nan_count[i];
524		if non_nan_count == 0 {
525			out[i] = nan();
526		} else if non_nan_count == 1 {
527			out[i] = zero();
528		} else {
529			out[i] = out[i]
530				.mul_real(from_usize::<T::Real>(non_nan_count - 1).recip());
531		}
532	}
533}
534fn col_mean_propagate_nan_fallback<T: ComplexField>(
535	out: ColMut<'_, T>,
536	mat: MatRef<'_, T>,
537) {
538	with_dim!(M, mat.nrows());
539	with_dim!(N, mat.ncols());
540	let mut out = out.as_row_shape_mut(M);
541	let mat = mat.as_shape(M, N);
542	out.fill(zero());
543	for j in N.indices() {
544		for i in M.indices() {
545			out[i] += &mat[(i, j)];
546		}
547	}
548	let ref n = from_usize::<T::Real>(*N).recip();
549	for i in M.indices() {
550		out[i] = out[i].mul_real(n);
551	}
552}
553fn col_varm_propagate_nan_fallback<T: ComplexField>(
554	out: ColMut<'_, T::Real>,
555	mat: MatRef<'_, T>,
556	col_mean: ColRef<'_, T>,
557) {
558	with_dim!(M, mat.nrows());
559	with_dim!(N, mat.ncols());
560	let mut out = out.as_row_shape_mut(M);
561	let mat = mat.as_shape(M, N);
562	let col_mean = col_mean.as_row_shape(M);
563	out.fill(zero());
564	for j in N.indices() {
565		for i in M.indices() {
566			let val = (&mat[(i, j)] - &col_mean[i]).abs2();
567			out[i] += val;
568		}
569	}
570	let n = *N;
571	if n == 0 {
572		out.fill(nan());
573	} else if n == 1 {
574		out.fill(zero());
575	} else {
576		let ref n = from_usize::<T::Real>(*N - 1).recip();
577		for i in M.indices() {
578			out[i] = out[i].mul_real(n);
579		}
580	}
581}
582fn col_mean_ignore<T: ComplexField>(out: ColMut<'_, T>, mat: MatRef<'_, T>) {
583	let mat = if mat.col_stride() >= 0 {
584		mat
585	} else {
586		mat.reverse_cols()
587	};
588	let mat = if mat.row_stride() >= 0 {
589		mat
590	} else {
591		mat.reverse_rows()
592	};
593	if const { T::SIMD_CAPABILITIES.is_simd() } {
594		if mat.ncols() > 1 && mat.col_stride() == 1 {
595			col_mean_row_major_ignore_nan(out, mat.try_as_row_major().unwrap());
596		} else {
597			col_mean_ignore_nan_fallback(out, mat);
598		}
599	} else {
600		col_mean_ignore_nan_fallback(out, mat);
601	}
602}
603fn col_varm_ignore<T: ComplexField>(
604	out: ColMut<'_, T::Real>,
605	mat: MatRef<'_, T>,
606	col_mean: ColRef<'_, T>,
607) {
608	let mat = if mat.col_stride() >= 0 {
609		mat
610	} else {
611		mat.reverse_cols()
612	};
613	let mat = if mat.row_stride() >= 0 {
614		mat
615	} else {
616		mat.reverse_rows()
617	};
618	if const { T::SIMD_CAPABILITIES.is_simd() } {
619		if mat.ncols() > 1 && mat.col_stride() == 1 {
620			col_varm_row_major_ignore_nan(
621				out,
622				mat.try_as_row_major().unwrap(),
623				col_mean,
624			);
625		} else {
626			col_varm_ignore_nan_fallback(out, mat, col_mean);
627		}
628	} else {
629		col_varm_ignore_nan_fallback(out, mat, col_mean);
630	}
631}
632fn col_mean_propagate<T: ComplexField>(out: ColMut<'_, T>, mat: MatRef<'_, T>) {
633	let mat = if mat.col_stride() >= 0 {
634		mat
635	} else {
636		mat.reverse_cols()
637	};
638	let mat = if mat.row_stride() >= 0 {
639		mat
640	} else {
641		mat.reverse_rows()
642	};
643	if const { T::SIMD_CAPABILITIES.is_simd() } {
644		if mat.ncols() > 1 && mat.col_stride() == 1 {
645			col_mean_row_major_propagate_nan(
646				out,
647				mat.try_as_row_major().unwrap(),
648			);
649		} else {
650			col_mean_propagate_nan_fallback(out, mat);
651		}
652	} else {
653		col_mean_propagate_nan_fallback(out, mat);
654	}
655}
656fn col_varm_propagate<T: ComplexField>(
657	out: ColMut<'_, T::Real>,
658	mat: MatRef<'_, T>,
659	col_mean: ColRef<'_, T>,
660) {
661	let mat = if mat.col_stride() >= 0 {
662		mat
663	} else {
664		mat.reverse_cols()
665	};
666	let mat = if mat.row_stride() >= 0 {
667		mat
668	} else {
669		mat.reverse_rows()
670	};
671	if const { T::SIMD_CAPABILITIES.is_simd() } {
672		if mat.ncols() > 1 && mat.col_stride() == 1 {
673			col_varm_row_major_propagate_nan(
674				out,
675				mat.try_as_row_major().unwrap(),
676				col_mean,
677			);
678		} else {
679			col_varm_propagate_nan_fallback(out, mat, col_mean);
680		}
681	} else {
682		col_varm_propagate_nan_fallback(out, mat, col_mean);
683	}
684}
685/// computes the mean of the columns of `mat` and stores the result in `out`
686#[track_caller]
687pub fn col_mean<T: ComplexField>(
688	out: ColMut<'_, T>,
689	mat: MatRef<'_, T>,
690	nan: NanHandling,
691) {
692	assert!(all(out.nrows() == mat.nrows()));
693	match nan {
694		NanHandling::Propagate => col_mean_propagate(out, mat),
695		NanHandling::Ignore => col_mean_ignore(out, mat),
696	}
697}
698/// computes the mean of the rows of `mat` and stores the result in `out`
699#[track_caller]
700pub fn row_mean<T: ComplexField>(
701	out: RowMut<'_, T>,
702	mat: MatRef<'_, T>,
703	nan: NanHandling,
704) {
705	assert!(all(out.ncols() == mat.ncols()));
706	col_mean(out.transpose_mut(), mat.transpose(), nan);
707}
708/// computes the variance of the columns of `mat` and stores the result in `out`
709#[track_caller]
710pub fn col_varm<T: ComplexField>(
711	out: ColMut<'_, T::Real>,
712	mat: MatRef<'_, T>,
713	col_mean: ColRef<'_, T>,
714	nan: NanHandling,
715) {
716	assert!(all(
717		out.nrows() == mat.nrows(),
718		col_mean.nrows() == mat.nrows()
719	));
720	match nan {
721		NanHandling::Propagate => col_varm_propagate(out, mat, col_mean),
722		NanHandling::Ignore => col_varm_ignore(out, mat, col_mean),
723	}
724}
725/// computes the variance of the rows of `mat` and stores the result in `out`
726#[track_caller]
727pub fn row_varm<T: ComplexField>(
728	out: RowMut<'_, T::Real>,
729	mat: MatRef<'_, T>,
730	row_mean: RowRef<'_, T>,
731	nan: NanHandling,
732) {
733	assert!(all(
734		out.ncols() == mat.ncols(),
735		row_mean.ncols() == mat.ncols()
736	));
737	col_varm(
738		out.transpose_mut(),
739		mat.transpose(),
740		row_mean.transpose(),
741		nan,
742	);
743}
744#[cfg(test)]
745mod tests {
746	use super::*;
747	use equator::assert;
748	#[test]
749	fn test_meanvar_propagate() {
750		let c32 = c32::new;
751		let A = mat![
752			[c32(1.2, 2.3), c32(3.4, 1.2)],
753			[c32(1.7, -1.0), c32(-3.8, 1.95)],
754		];
755		let mut row_mean = Row::zeros(A.ncols());
756		let mut row_var = Row::zeros(A.ncols());
757		super::row_mean(row_mean.as_mut(), A.as_ref(), NanHandling::Propagate);
758		super::row_varm(
759			row_var.as_mut(),
760			A.as_ref(),
761			row_mean.as_ref(),
762			NanHandling::Propagate,
763		);
764		let mut col_mean = Col::zeros(A.nrows());
765		let mut col_var = Col::zeros(A.nrows());
766		super::col_mean(col_mean.as_mut(), A.as_ref(), NanHandling::Propagate);
767		super::col_varm(
768			col_var.as_mut(),
769			A.as_ref(),
770			col_mean.as_ref(),
771			NanHandling::Propagate,
772		);
773		assert!(
774			row_mean
775				== row![
776					(A[(0, 0)] + A[(1, 0)]) / 2.0,
777					(A[(0, 1)] + A[(1, 1)]) / 2.0,
778				]
779		);
780		assert!(
781			row_var
782				== row![
783					(A[(0, 0)] - row_mean[0]).abs2()
784						+ (A[(1, 0)] - row_mean[0]).abs2(),
785					(A[(0, 1)] - row_mean[1]).abs2()
786						+ (A[(1, 1)] - row_mean[1]).abs2(),
787				]
788		);
789		assert!(
790			col_mean
791				== col![
792					(A[(0, 0)] + A[(0, 1)]) / 2.0,
793					(A[(1, 0)] + A[(1, 1)]) / 2.0,
794				]
795		);
796		assert!(
797			col_var
798				== col![
799					(A[(0, 0)] - col_mean[0]).abs2()
800						+ (A[(0, 1)] - col_mean[0]).abs2(),
801					(A[(1, 0)] - col_mean[1]).abs2()
802						+ (A[(1, 1)] - col_mean[1]).abs2(),
803				]
804		);
805	}
806	#[test]
807	fn test_meanvar_ignore_nan_nonan_c32() {
808		let c32 = c32::new;
809		let A = mat![
810			[c32(1.2, 2.3), c32(3.4, 1.2)],
811			[c32(1.7, -1.0), c32(-3.8, 1.95)],
812		];
813		let mut row_mean = Row::zeros(A.ncols());
814		let mut row_var = Row::zeros(A.ncols());
815		super::row_mean(row_mean.as_mut(), A.as_ref(), NanHandling::Ignore);
816		super::row_varm(
817			row_var.as_mut(),
818			A.as_ref(),
819			row_mean.as_ref(),
820			NanHandling::Ignore,
821		);
822		let mut col_mean = Col::zeros(A.nrows());
823		let mut col_var = Col::zeros(A.nrows());
824		super::col_mean(col_mean.as_mut(), A.as_ref(), NanHandling::Ignore);
825		super::col_varm(
826			col_var.as_mut(),
827			A.as_ref(),
828			col_mean.as_ref(),
829			NanHandling::Ignore,
830		);
831		assert!(
832			row_mean
833				== row![
834					(A[(0, 0)] + A[(1, 0)]) / 2.0,
835					(A[(0, 1)] + A[(1, 1)]) / 2.0,
836				]
837		);
838		assert!(
839			row_var
840				== row![
841					(A[(0, 0)] - row_mean[0]).abs2()
842						+ (A[(1, 0)] - row_mean[0]).abs2(),
843					(A[(0, 1)] - row_mean[1]).abs2()
844						+ (A[(1, 1)] - row_mean[1]).abs2(),
845				]
846		);
847		assert!(
848			col_mean
849				== col![
850					(A[(0, 0)] + A[(0, 1)]) / 2.0,
851					(A[(1, 0)] + A[(1, 1)]) / 2.0,
852				]
853		);
854		assert!(
855			col_var
856				== col![
857					(A[(0, 0)] - col_mean[0]).abs2()
858						+ (A[(0, 1)] - col_mean[0]).abs2(),
859					(A[(1, 0)] - col_mean[1]).abs2()
860						+ (A[(1, 1)] - col_mean[1]).abs2(),
861				]
862		);
863	}
864	#[test]
865	fn test_meanvar_ignore_nan_yesnan_c32() {
866		let c32 = c32::new;
867		let nan = f32::NAN;
868		let A = mat![
869			[c32(1.2, nan), c32(3.4, 1.2)],
870			[c32(1.7, -1.0), c32(-3.8, 1.95)],
871		];
872		let mut row_mean = Row::zeros(A.ncols());
873		let mut row_var = Row::zeros(A.ncols());
874		super::row_mean(row_mean.as_mut(), A.as_ref(), NanHandling::Ignore);
875		super::row_varm(
876			row_var.as_mut(),
877			A.as_ref(),
878			row_mean.as_ref(),
879			NanHandling::Ignore,
880		);
881		let mut col_mean = Col::zeros(A.nrows());
882		let mut col_var = Col::zeros(A.nrows());
883		super::col_mean(col_mean.as_mut(), A.as_ref(), NanHandling::Ignore);
884		super::col_varm(
885			col_var.as_mut(),
886			A.as_ref(),
887			col_mean.as_ref(),
888			NanHandling::Ignore,
889		);
890		assert!(
891			row_mean == row![A[(1, 0)] / 1.0, (A[(0, 1)] + A[(1, 1)]) / 2.0,]
892		);
893		assert!(
894			row_var
895				== row![
896					(A[(1, 0)] - row_mean[0]).abs2(),
897					(A[(0, 1)] - row_mean[1]).abs2()
898						+ (A[(1, 1)] - row_mean[1]).abs2(),
899				]
900		);
901		assert!(
902			col_mean == col![A[(0, 1)] / 1.0, (A[(1, 0)] + A[(1, 1)]) / 2.0,]
903		);
904		assert!(
905			col_var
906				== col![
907					(A[(0, 1)] - col_mean[0]).abs2(),
908					(A[(1, 0)] - col_mean[1]).abs2()
909						+ (A[(1, 1)] - col_mean[1]).abs2(),
910				]
911		);
912	}
913	#[test]
914	fn test_meanvar_ignore_nan_nonan_c64() {
915		let c64 = c64::new;
916		let A = mat![
917			[c64(1.2, 2.3), c64(3.4, 1.2)],
918			[c64(1.7, -1.0), c64(-3.8, 1.95)],
919		];
920		let mut row_mean = Row::zeros(A.ncols());
921		let mut row_var = Row::zeros(A.ncols());
922		super::row_mean(row_mean.as_mut(), A.as_ref(), NanHandling::Ignore);
923		super::row_varm(
924			row_var.as_mut(),
925			A.as_ref(),
926			row_mean.as_ref(),
927			NanHandling::Ignore,
928		);
929		let mut col_mean = Col::zeros(A.nrows());
930		let mut col_var = Col::zeros(A.nrows());
931		super::col_mean(col_mean.as_mut(), A.as_ref(), NanHandling::Ignore);
932		super::col_varm(
933			col_var.as_mut(),
934			A.as_ref(),
935			col_mean.as_ref(),
936			NanHandling::Ignore,
937		);
938		assert!(
939			row_mean
940				== row![
941					(A[(0, 0)] + A[(1, 0)]) / 2.0,
942					(A[(0, 1)] + A[(1, 1)]) / 2.0,
943				]
944		);
945		assert!(
946			row_var
947				== row![
948					(A[(0, 0)] - row_mean[0]).abs2()
949						+ (A[(1, 0)] - row_mean[0]).abs2(),
950					(A[(0, 1)] - row_mean[1]).abs2()
951						+ (A[(1, 1)] - row_mean[1]).abs2(),
952				]
953		);
954		assert!(
955			col_mean
956				== col![
957					(A[(0, 0)] + A[(0, 1)]) / 2.0,
958					(A[(1, 0)] + A[(1, 1)]) / 2.0,
959				]
960		);
961		assert!(
962			col_var
963				== col![
964					(A[(0, 0)] - col_mean[0]).abs2()
965						+ (A[(0, 1)] - col_mean[0]).abs2(),
966					(A[(1, 0)] - col_mean[1]).abs2()
967						+ (A[(1, 1)] - col_mean[1]).abs2(),
968				]
969		);
970	}
971	#[test]
972	fn test_meanvar_ignore_nan_yesnan_c64() {
973		let c64 = c64::new;
974		let nan = f64::NAN;
975		let A = mat![
976			[c64(1.2, nan), c64(3.4, 1.2)],
977			[c64(1.7, -1.0), c64(-3.8, 1.95)],
978		];
979		let mut row_mean = Row::zeros(A.ncols());
980		let mut row_var = Row::zeros(A.ncols());
981		super::row_mean(row_mean.as_mut(), A.as_ref(), NanHandling::Ignore);
982		super::row_varm(
983			row_var.as_mut(),
984			A.as_ref(),
985			row_mean.as_ref(),
986			NanHandling::Ignore,
987		);
988		let mut col_mean = Col::zeros(A.nrows());
989		let mut col_var = Col::zeros(A.nrows());
990		super::col_mean(col_mean.as_mut(), A.as_ref(), NanHandling::Ignore);
991		super::col_varm(
992			col_var.as_mut(),
993			A.as_ref(),
994			col_mean.as_ref(),
995			NanHandling::Ignore,
996		);
997		assert!(
998			row_mean == row![A[(1, 0)] / 1.0, (A[(0, 1)] + A[(1, 1)]) / 2.0,]
999		);
1000		assert!(
1001			row_var
1002				== row![
1003					(A[(1, 0)] - row_mean[0]).abs2(),
1004					(A[(0, 1)] - row_mean[1]).abs2()
1005						+ (A[(1, 1)] - row_mean[1]).abs2(),
1006				]
1007		);
1008		assert!(
1009			col_mean == col![A[(0, 1)] / 1.0, (A[(1, 0)] + A[(1, 1)]) / 2.0,]
1010		);
1011		assert!(
1012			col_var
1013				== col![
1014					(A[(0, 1)] - col_mean[0]).abs2(),
1015					(A[(1, 0)] - col_mean[1]).abs2()
1016						+ (A[(1, 1)] - col_mean[1]).abs2(),
1017				]
1018		);
1019	}
1020}
faer/stats/meanvar.rs

faer/stats/
meanvar.rs
