1use ferray_core::Array;
25use ferray_core::dimension::{Dimension, IxDyn};
26use ferray_core::dtype::Element;
27use ferray_core::error::{FerrayError, FerrayResult};
28use num_traits::Float;
29
30use crate::MaskedArray;
31
32fn compute_strides(shape: &[usize]) -> Vec<usize> {
38 let n = shape.len();
39 let mut s = vec![1usize; n];
40 for i in (0..n.saturating_sub(1)).rev() {
41 s[i] = s[i + 1] * shape[i + 1];
42 }
43 s
44}
45
46fn increment_multi(multi: &mut [usize], shape: &[usize]) -> bool {
48 for d in (0..multi.len()).rev() {
49 multi[d] += 1;
50 if multi[d] < shape[d] {
51 return true;
52 }
53 multi[d] = 0;
54 }
55 false
56}
57
58fn reduce_axis<T, D, F>(
65 ma: &MaskedArray<T, D>,
66 axis: usize,
67 fill_value: T,
68 kernel: F,
69) -> FerrayResult<MaskedArray<T, IxDyn>>
70where
71 T: Element + Copy,
72 D: Dimension,
73 F: Fn(&[(T, bool)]) -> Option<T>,
74{
75 let ndim = ma.ndim();
76 if axis >= ndim {
77 return Err(FerrayError::axis_out_of_bounds(axis, ndim));
78 }
79 let shape = ma.shape();
80 let axis_len = shape[axis];
81
82 let out_shape: Vec<usize> = shape
84 .iter()
85 .enumerate()
86 .filter_map(|(i, &s)| if i == axis { None } else { Some(s) })
87 .collect();
88 let out_size: usize = if out_shape.is_empty() {
89 1
90 } else {
91 out_shape.iter().product()
92 };
93
94 let src_data: Vec<T> = ma.data().iter().copied().collect();
97 let src_mask: Vec<bool> = ma.mask().iter().copied().collect();
98 let strides = compute_strides(shape);
99
100 let mut out_data = Vec::with_capacity(out_size);
101 let mut out_mask = Vec::with_capacity(out_size);
102 let mut out_multi = vec![0usize; out_shape.len()];
103 let mut in_multi = vec![0usize; ndim];
104 let mut lane: Vec<(T, bool)> = Vec::with_capacity(axis_len);
105
106 for _ in 0..out_size {
107 let mut out_dim = 0;
110 for (d, idx) in in_multi.iter_mut().enumerate() {
111 if d == axis {
112 *idx = 0;
113 } else {
114 *idx = out_multi[out_dim];
115 out_dim += 1;
116 }
117 }
118
119 lane.clear();
120 for k in 0..axis_len {
121 in_multi[axis] = k;
122 let flat = in_multi
123 .iter()
124 .zip(strides.iter())
125 .map(|(i, s)| i * s)
126 .sum::<usize>();
127 lane.push((src_data[flat], src_mask[flat]));
128 }
129
130 if let Some(value) = kernel(&lane) {
131 out_data.push(value);
132 out_mask.push(false);
133 } else {
134 out_data.push(fill_value);
135 out_mask.push(true);
136 }
137
138 if !out_shape.is_empty() {
139 increment_multi(&mut out_multi, &out_shape);
140 }
141 }
142
143 let data_arr = Array::<T, IxDyn>::from_vec(IxDyn::new(&out_shape), out_data)?;
144 let mask_arr = Array::<bool, IxDyn>::from_vec(IxDyn::new(&out_shape), out_mask)?;
145 let mut result = MaskedArray::new(data_arr, mask_arr)?;
146 result.set_fill_value(fill_value);
147 Ok(result)
148}
149
150fn count_axis<T, D>(ma: &MaskedArray<T, D>, axis: usize) -> FerrayResult<Array<u64, IxDyn>>
153where
154 T: Element + Copy,
155 D: Dimension,
156{
157 let ndim = ma.ndim();
158 if axis >= ndim {
159 return Err(FerrayError::axis_out_of_bounds(axis, ndim));
160 }
161 let shape = ma.shape();
162 let axis_len = shape[axis];
163 let out_shape: Vec<usize> = shape
164 .iter()
165 .enumerate()
166 .filter_map(|(i, &s)| if i == axis { None } else { Some(s) })
167 .collect();
168 let out_size: usize = if out_shape.is_empty() {
169 1
170 } else {
171 out_shape.iter().product()
172 };
173
174 let src_mask: Vec<bool> = ma.mask().iter().copied().collect();
175 let strides = compute_strides(shape);
176 let mut out: Vec<u64> = Vec::with_capacity(out_size);
177 let mut out_multi = vec![0usize; out_shape.len()];
178 let mut in_multi = vec![0usize; ndim];
179
180 for _ in 0..out_size {
181 let mut out_dim = 0;
182 for (d, idx) in in_multi.iter_mut().enumerate() {
183 if d == axis {
184 *idx = 0;
185 } else {
186 *idx = out_multi[out_dim];
187 out_dim += 1;
188 }
189 }
190
191 let mut count: u64 = 0;
192 for k in 0..axis_len {
193 in_multi[axis] = k;
194 let flat = in_multi
195 .iter()
196 .zip(strides.iter())
197 .map(|(i, s)| i * s)
198 .sum::<usize>();
199 if !src_mask[flat] {
200 count += 1;
201 }
202 }
203 out.push(count);
204
205 if !out_shape.is_empty() {
206 increment_multi(&mut out_multi, &out_shape);
207 }
208 }
209
210 Array::<u64, IxDyn>::from_vec(IxDyn::new(&out_shape), out)
211}
212
213impl<T, D> MaskedArray<T, D>
214where
215 T: Element + Copy,
216 D: Dimension,
217{
218 pub fn count(&self) -> FerrayResult<usize> {
224 let n = self
225 .data()
226 .iter()
227 .zip(self.mask().iter())
228 .filter(|(_, m)| !**m)
229 .count();
230 Ok(n)
231 }
232
233 pub fn count_axis(&self, axis: usize) -> FerrayResult<Array<u64, IxDyn>> {
241 count_axis(self, axis)
242 }
243}
244
245impl<T, D> MaskedArray<T, D>
246where
247 T: Element + Float,
248 D: Dimension,
249{
250 pub fn sum(&self) -> FerrayResult<T> {
265 let zero = num_traits::zero::<T>();
266 let mut any = false;
267 let s = self
268 .data()
269 .iter()
270 .zip(self.mask().iter())
271 .filter(|(_, m)| !**m)
272 .fold(zero, |acc, (v, _)| {
273 any = true;
274 acc + *v
275 });
276 if any { Ok(s) } else { Ok(T::nan()) }
277 }
278
279 pub fn mean(&self) -> FerrayResult<T> {
286 let zero = num_traits::zero::<T>();
287 let (sum, count) = self
288 .data()
289 .iter()
290 .zip(self.mask().iter())
291 .filter(|(_, m)| !**m)
292 .fold((zero, 0usize), |(s, c), (v, _)| (s + *v, c + 1));
293 if count == 0 {
294 return Ok(T::nan());
295 }
296 let n = T::from(count).ok_or_else(|| {
301 FerrayError::invalid_value(format!(
302 "cannot convert unmasked count {count} to element type"
303 ))
304 })?;
305 Ok(sum / n)
306 }
307
308 pub fn min(&self) -> FerrayResult<T> {
315 self.data()
316 .iter()
317 .zip(self.mask().iter())
318 .filter(|(_, m)| !**m)
319 .map(|(v, _)| *v)
320 .fold(None, |acc: Option<T>, v| {
321 Some(match acc {
322 Some(a) => {
323 if a <= v {
325 a
326 } else if a > v {
327 v
328 } else {
329 a
330 }
331 }
332 None => v,
333 })
334 })
335 .ok_or_else(|| FerrayError::invalid_value("min: all elements are masked"))
336 }
337
338 pub fn max(&self) -> FerrayResult<T> {
345 self.data()
346 .iter()
347 .zip(self.mask().iter())
348 .filter(|(_, m)| !**m)
349 .map(|(v, _)| *v)
350 .fold(None, |acc: Option<T>, v| {
351 Some(match acc {
352 Some(a) => {
353 if a >= v {
354 a
355 } else if a < v {
356 v
357 } else {
358 a
359 }
360 }
361 None => v,
362 })
363 })
364 .ok_or_else(|| FerrayError::invalid_value("max: all elements are masked"))
365 }
366
367 pub fn var(&self) -> FerrayResult<T> {
374 self.var_ddof(0)
375 }
376
377 pub fn var_ddof(&self, ddof: usize) -> FerrayResult<T> {
389 let zero = num_traits::zero::<T>();
396 let mut mean = zero;
397 let mut m2 = zero;
398 let mut count = 0usize;
399
400 for (v, m) in self.data().iter().zip(self.mask().iter()) {
401 if *m {
402 continue;
403 }
404 count += 1;
405 let n_t = T::from(count).ok_or_else(|| {
406 FerrayError::invalid_value(format!("cannot convert count {count} to element type"))
407 })?;
408 let delta = *v - mean;
409 mean = mean + delta / n_t;
410 let delta2 = *v - mean;
411 m2 = m2 + delta * delta2;
412 }
413
414 if count == 0 {
415 return Ok(T::nan());
416 }
417 if count <= ddof {
418 return Ok(T::nan());
419 }
420 let n = T::from(count - ddof).ok_or_else(|| {
421 FerrayError::invalid_value(format!(
422 "cannot convert (count - ddof) = {} to element type",
423 count - ddof
424 ))
425 })?;
426 Ok(m2 / n)
427 }
428
429 pub fn std(&self) -> FerrayResult<T> {
436 Ok(self.var()?.sqrt())
437 }
438
439 pub fn std_ddof(&self, ddof: usize) -> FerrayResult<T> {
450 Ok(self.var_ddof(ddof)?.sqrt())
451 }
452
453 pub fn sum_axis(&self, axis: usize) -> FerrayResult<MaskedArray<T, IxDyn>> {
464 let zero = num_traits::zero::<T>();
465 let fill = self.fill_value();
466 reduce_axis(self, axis, fill, |lane| {
467 let mut acc = zero;
468 let mut any = false;
469 for &(v, m) in lane {
470 if !m {
471 acc = acc + v;
472 any = true;
473 }
474 }
475 if any { Some(acc) } else { None }
476 })
477 }
478
479 pub fn mean_axis(&self, axis: usize) -> FerrayResult<MaskedArray<T, IxDyn>> {
482 let zero = num_traits::zero::<T>();
483 let fill = self.fill_value();
484 reduce_axis(self, axis, fill, |lane| {
485 let mut acc = zero;
486 let mut count = 0usize;
487 for &(v, m) in lane {
488 if !m {
489 acc = acc + v;
490 count += 1;
491 }
492 }
493 if count == 0 {
494 None
495 } else {
496 T::from(count).map(|n| acc / n)
501 }
502 })
503 }
504
505 pub fn min_axis(&self, axis: usize) -> FerrayResult<MaskedArray<T, IxDyn>> {
507 let fill = self.fill_value();
508 reduce_axis(self, axis, fill, |lane| {
509 let mut acc: Option<T> = None;
510 for &(v, m) in lane {
511 if !m {
512 acc = Some(match acc {
513 Some(a) => {
514 if a <= v {
516 a
517 } else if a > v {
518 v
519 } else {
520 a
521 }
522 }
523 None => v,
524 });
525 }
526 }
527 acc
528 })
529 }
530
531 pub fn max_axis(&self, axis: usize) -> FerrayResult<MaskedArray<T, IxDyn>> {
533 let fill = self.fill_value();
534 reduce_axis(self, axis, fill, |lane| {
535 let mut acc: Option<T> = None;
536 for &(v, m) in lane {
537 if !m {
538 acc = Some(match acc {
539 Some(a) => {
540 if a >= v {
541 a
542 } else if a < v {
543 v
544 } else {
545 a
546 }
547 }
548 None => v,
549 });
550 }
551 }
552 acc
553 })
554 }
555
556 pub fn var_axis(&self, axis: usize) -> FerrayResult<MaskedArray<T, IxDyn>> {
558 self.var_axis_ddof(axis, 0)
559 }
560
561 pub fn var_axis_ddof(&self, axis: usize, ddof: usize) -> FerrayResult<MaskedArray<T, IxDyn>> {
567 let zero = num_traits::zero::<T>();
569 let fill = self.fill_value();
570 reduce_axis(self, axis, fill, |lane| {
571 let mut mean = zero;
572 let mut m2 = zero;
573 let mut count = 0usize;
574 for &(v, m) in lane {
575 if m {
576 continue;
577 }
578 count += 1;
579 let n_t = T::from(count)?;
580 let delta = v - mean;
581 mean = mean + delta / n_t;
582 let delta2 = v - mean;
583 m2 = m2 + delta * delta2;
584 }
585 if count <= ddof {
586 return None;
587 }
588 let n_var = T::from(count - ddof)?;
589 Some(m2 / n_var)
590 })
591 }
592
593 pub fn std_axis(&self, axis: usize) -> FerrayResult<MaskedArray<T, IxDyn>> {
595 self.std_axis_ddof(axis, 0)
596 }
597
598 pub fn std_axis_ddof(&self, axis: usize, ddof: usize) -> FerrayResult<MaskedArray<T, IxDyn>> {
600 let result = self.var_axis_ddof(axis, ddof)?;
601 let fill = self.fill_value();
602 let mask = result.mask().clone();
603 let new_data: Vec<T> = result
604 .data()
605 .iter()
606 .zip(result.mask().iter())
607 .map(|(v, m)| if *m { fill } else { v.sqrt() })
608 .collect();
609 let data_arr = Array::<T, IxDyn>::from_vec(IxDyn::new(result.shape()), new_data)?;
610 let mut out = MaskedArray::new(data_arr, mask)?;
611 out.set_fill_value(fill);
612 Ok(out)
613 }
614}
615
616#[cfg(test)]
617mod tests {
618 use super::*;
619 use ferray_core::dimension::{Ix1, Ix2};
620
621 fn ma2d(rows: usize, cols: usize, data: Vec<f64>, mask: Vec<bool>) -> MaskedArray<f64, Ix2> {
622 let d = Array::<f64, Ix2>::from_vec(Ix2::new([rows, cols]), data).unwrap();
623 let m = Array::<bool, Ix2>::from_vec(Ix2::new([rows, cols]), mask).unwrap();
624 MaskedArray::new(d, m).unwrap()
625 }
626
627 #[test]
630 fn sum_axis_drops_axis() {
631 let ma = ma2d(2, 3, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], vec![false; 6]);
633 let s0 = ma.sum_axis(0).unwrap();
634 assert_eq!(s0.shape(), &[3]);
635 let d0: Vec<f64> = s0.data().iter().copied().collect();
636 assert_eq!(d0, vec![5.0, 7.0, 9.0]);
637
638 let s1 = ma.sum_axis(1).unwrap();
639 assert_eq!(s1.shape(), &[2]);
640 let d1: Vec<f64> = s1.data().iter().copied().collect();
641 assert_eq!(d1, vec![6.0, 15.0]);
642 }
643
644 #[test]
645 fn sum_axis_skips_masked() {
646 let ma = ma2d(
648 2,
649 3,
650 vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0],
651 vec![false, true, false, false, false, false],
652 );
653 let s0 = ma.sum_axis(0).unwrap();
655 let d0: Vec<f64> = s0.data().iter().copied().collect();
656 assert_eq!(d0, vec![5.0, 5.0, 9.0]);
657 let m0: Vec<bool> = s0.mask().iter().copied().collect();
658 assert_eq!(m0, vec![false, false, false]);
659 }
660
661 #[test]
662 fn sum_axis_all_masked_lane_is_masked() {
663 let ma = ma2d(
665 2,
666 3,
667 vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0],
668 vec![false, true, false, false, true, false],
669 );
670 let s0 = ma.sum_axis(0).unwrap();
671 let m0: Vec<bool> = s0.mask().iter().copied().collect();
672 assert_eq!(m0, vec![false, true, false]);
673 }
674
675 #[test]
676 fn mean_axis_skips_masked() {
677 let ma = ma2d(
678 2,
679 3,
680 vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0],
681 vec![false, true, false, false, false, false],
682 );
683 let m1 = ma.mean_axis(1).unwrap();
685 let d: Vec<f64> = m1.data().iter().copied().collect();
686 assert_eq!(d, vec![2.0, 5.0]);
687 }
688
689 #[test]
690 fn min_max_axis() {
691 let ma = ma2d(2, 3, vec![3.0, 1.0, 5.0, 2.0, 4.0, 0.0], vec![false; 6]);
692 let mn = ma.min_axis(0).unwrap();
693 let mx = ma.max_axis(0).unwrap();
694 let mn_d: Vec<f64> = mn.data().iter().copied().collect();
695 let mx_d: Vec<f64> = mx.data().iter().copied().collect();
696 assert_eq!(mn_d, vec![2.0, 1.0, 0.0]);
697 assert_eq!(mx_d, vec![3.0, 4.0, 5.0]);
698 }
699
700 #[test]
701 fn count_axis_basic() {
702 let ma = ma2d(
704 2,
705 3,
706 vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0],
707 vec![false, true, false, false, false, true],
708 );
709 let c0 = ma.count_axis(0).unwrap();
710 let v: Vec<u64> = c0.iter().copied().collect();
711 assert_eq!(v, vec![2u64, 1, 1]);
712 }
713
714 #[test]
715 fn axis_out_of_bounds_errors() {
716 let ma = ma2d(2, 3, vec![0.0; 6], vec![false; 6]);
717 assert!(ma.sum_axis(2).is_err());
718 }
719
720 #[test]
721 fn var_std_axis() {
722 let data = vec![1.0, 2.0, 3.0, 4.0, 5.0, 1.0, 2.0, 3.0, 4.0, 5.0];
724 let ma = ma2d(2, 5, data, vec![false; 10]);
725 let v = ma.var_axis(1).unwrap();
726 let s = ma.std_axis(1).unwrap();
727 let v_d: Vec<f64> = v.data().iter().copied().collect();
728 let s_d: Vec<f64> = s.data().iter().copied().collect();
729 for &x in &v_d {
730 assert!((x - 2.0).abs() < 1e-12);
731 }
732 for &x in &s_d {
733 assert!((x - 2.0_f64.sqrt()).abs() < 1e-12);
734 }
735 }
736
737 #[test]
740 fn fill_value_default_is_numpy_default_filler_1e20() {
741 let d = Array::<f64, Ix1>::from_vec(Ix1::new([3]), vec![1.0, 2.0, 3.0]).unwrap();
744 let m = Array::<bool, Ix1>::from_vec(Ix1::new([3]), vec![false; 3]).unwrap();
745 let ma = MaskedArray::new(d, m).unwrap();
746 assert_eq!(ma.fill_value(), 1e20);
747 }
748
749 #[test]
750 fn with_fill_value_sets_field() {
751 let d = Array::<f64, Ix1>::from_vec(Ix1::new([3]), vec![1.0, 2.0, 3.0]).unwrap();
752 let m = Array::<bool, Ix1>::from_vec(Ix1::new([3]), vec![false; 3]).unwrap();
753 let ma = MaskedArray::new(d, m).unwrap().with_fill_value(99.0);
754 assert_eq!(ma.fill_value(), 99.0);
755 }
756
757 #[test]
758 fn filled_default_uses_stored_fill_value() {
759 let d = Array::<f64, Ix1>::from_vec(Ix1::new([4]), vec![1.0, 2.0, 3.0, 4.0]).unwrap();
760 let m =
761 Array::<bool, Ix1>::from_vec(Ix1::new([4]), vec![false, true, false, true]).unwrap();
762 let ma = MaskedArray::new(d, m).unwrap().with_fill_value(-1.0);
763 let filled = ma.filled_default().unwrap();
764 let v: Vec<f64> = filled.iter().copied().collect();
765 assert_eq!(v, vec![1.0, -1.0, 3.0, -1.0]);
766 }
767
768 #[test]
769 fn arithmetic_uses_fill_value() {
770 use crate::masked_add;
773 let d_a = Array::<f64, Ix1>::from_vec(Ix1::new([3]), vec![1.0, 2.0, 3.0]).unwrap();
774 let m_a = Array::<bool, Ix1>::from_vec(Ix1::new([3]), vec![false, true, false]).unwrap();
775 let d_b = Array::<f64, Ix1>::from_vec(Ix1::new([3]), vec![10.0, 20.0, 30.0]).unwrap();
776 let m_b = Array::<bool, Ix1>::from_vec(Ix1::new([3]), vec![false; 3]).unwrap();
777 let a = MaskedArray::new(d_a, m_a).unwrap().with_fill_value(-999.0);
778 let b = MaskedArray::new(d_b, m_b).unwrap();
779 let r = masked_add(&a, &b).unwrap();
780 let r_d: Vec<f64> = r.data().iter().copied().collect();
781 assert_eq!(r_d, vec![11.0, -999.0, 33.0]);
782 assert_eq!(r.fill_value(), -999.0);
783 }
784
785 #[test]
788 fn masked_add_broadcasts_within_same_rank() {
789 use crate::masked_add;
790 let d_a = Array::<f64, Ix2>::from_vec(Ix2::new([3, 1]), vec![1.0, 2.0, 3.0]).unwrap();
792 let m_a = Array::<bool, Ix2>::from_vec(Ix2::new([3, 1]), vec![false; 3]).unwrap();
793 let d_b =
794 Array::<f64, Ix2>::from_vec(Ix2::new([1, 4]), vec![10.0, 20.0, 30.0, 40.0]).unwrap();
795 let m_b = Array::<bool, Ix2>::from_vec(Ix2::new([1, 4]), vec![false; 4]).unwrap();
796 let a = MaskedArray::new(d_a, m_a).unwrap();
797 let b = MaskedArray::new(d_b, m_b).unwrap();
798 let r = masked_add(&a, &b).unwrap();
799 assert_eq!(r.shape(), &[3, 4]);
800 let r_d: Vec<f64> = r.data().iter().copied().collect();
801 assert_eq!(
802 r_d,
803 vec![
804 11.0, 21.0, 31.0, 41.0, 12.0, 22.0, 32.0, 42.0, 13.0, 23.0, 33.0, 43.0, ]
808 );
809 let r_m: Vec<bool> = r.mask().iter().copied().collect();
810 assert_eq!(r_m, vec![false; 12]);
811 }
812
813 #[test]
814 fn masked_sub_broadcasts_with_mask_union() {
815 use crate::masked_sub;
816 let d_a = Array::<f64, Ix2>::from_vec(Ix2::new([3, 1]), vec![10.0, 20.0, 30.0]).unwrap();
819 let m_a = Array::<bool, Ix2>::from_vec(Ix2::new([3, 1]), vec![false, true, false]).unwrap();
820 let d_b = Array::<f64, Ix2>::from_vec(Ix2::new([1, 4]), vec![1.0, 2.0, 3.0, 4.0]).unwrap();
821 let m_b = Array::<bool, Ix2>::from_vec(Ix2::new([1, 4]), vec![false; 4]).unwrap();
822 let a = MaskedArray::new(d_a, m_a).unwrap();
823 let b = MaskedArray::new(d_b, m_b).unwrap();
824 let r = masked_sub(&a, &b).unwrap();
825 let r_m: Vec<bool> = r.mask().iter().copied().collect();
826 assert_eq!(
828 r_m,
829 vec![
830 false, false, false, false, true, true, true, true, false, false, false, false, ]
834 );
835 }
836
837 fn all_masked_ma1d(n: usize) -> MaskedArray<f64, Ix1> {
849 let d = Array::<f64, Ix1>::from_vec(Ix1::new([n]), vec![1.0; n]).unwrap();
850 let m = Array::<bool, Ix1>::from_vec(Ix1::new([n]), vec![true; n]).unwrap();
851 MaskedArray::new(d, m).unwrap()
852 }
853
854 #[test]
855 fn sum_all_masked_returns_nan_not_zero() {
856 let ma = all_masked_ma1d(4);
859 assert!(ma.sum().unwrap_or(0.0).is_nan());
860 }
861
862 #[test]
863 fn mean_all_masked_returns_nan() {
864 let ma = all_masked_ma1d(4);
865 assert!(ma.mean().unwrap().is_nan());
866 }
867
868 #[test]
869 fn var_all_masked_returns_nan() {
870 let ma = all_masked_ma1d(4);
871 assert!(ma.var().unwrap().is_nan());
872 }
873
874 #[test]
875 fn std_all_masked_returns_nan() {
876 let ma = all_masked_ma1d(4);
877 assert!(ma.std().unwrap().is_nan());
878 }
879
880 #[test]
881 fn min_max_all_masked_error() {
882 let ma = all_masked_ma1d(4);
885 assert!(ma.min().is_err());
886 assert!(ma.max().is_err());
887 }
888
889 #[test]
890 fn sum_var_std_all_masked_2d_matches_1d() {
891 let d = Array::<f64, Ix2>::from_vec(Ix2::new([2, 3]), vec![9.0; 6]).unwrap();
893 let m = Array::<bool, Ix2>::from_vec(Ix2::new([2, 3]), vec![true; 6]).unwrap();
894 let ma = MaskedArray::new(d, m).unwrap();
895 assert!(ma.sum().unwrap_or(0.0).is_nan());
896 assert!(ma.var().unwrap_or(0.0).is_nan());
897 assert!(ma.std().unwrap_or(0.0).is_nan());
898 }
899
900 #[test]
901 fn masked_add_broadcast_incompatible_errors() {
902 use crate::masked_add;
903 let d_a = Array::<f64, Ix1>::from_vec(Ix1::new([3]), vec![1.0, 2.0, 3.0]).unwrap();
904 let m_a = Array::<bool, Ix1>::from_vec(Ix1::new([3]), vec![false; 3]).unwrap();
905 let d_b = Array::<f64, Ix1>::from_vec(Ix1::new([4]), vec![1.0, 2.0, 3.0, 4.0]).unwrap();
906 let m_b = Array::<bool, Ix1>::from_vec(Ix1::new([4]), vec![false; 4]).unwrap();
907 let a = MaskedArray::new(d_a, m_a).unwrap();
908 let b = MaskedArray::new(d_b, m_b).unwrap();
909 assert!(masked_add(&a, &b).is_err());
910 }
911
912 #[test]
915 fn var_ddof_zero_matches_default_var() {
916 let data =
917 Array::<f64, Ix1>::from_vec(Ix1::new([5]), vec![1.0, 2.0, 3.0, 4.0, 5.0]).unwrap();
918 let mask = Array::<bool, Ix1>::from_vec(Ix1::new([5]), vec![false; 5]).unwrap();
919 let ma = MaskedArray::new(data, mask).unwrap();
920 let v0 = ma.var().unwrap();
921 let v_explicit = ma.var_ddof(0).unwrap();
922 assert!((v0 - v_explicit).abs() < 1e-14);
923 }
924
925 #[test]
926 fn var_ddof_one_is_bessel_corrected() {
927 let data =
931 Array::<f64, Ix1>::from_vec(Ix1::new([5]), vec![1.0, 2.0, 3.0, 4.0, 5.0]).unwrap();
932 let mask = Array::<bool, Ix1>::from_vec(Ix1::new([5]), vec![false; 5]).unwrap();
933 let ma = MaskedArray::new(data, mask).unwrap();
934 let v0 = ma.var_ddof(0).unwrap();
935 let v1 = ma.var_ddof(1).unwrap();
936 assert!((v0 - 2.0).abs() < 1e-14, "ddof=0: expected 2.0, got {v0}");
937 assert!((v1 - 2.5).abs() < 1e-14, "ddof=1: expected 2.5, got {v1}");
938 }
939
940 #[test]
941 fn var_ddof_skips_masked_elements() {
942 let data =
945 Array::<f64, Ix1>::from_vec(Ix1::new([5]), vec![1.0, 2.0, 99.0, 4.0, 5.0]).unwrap();
946 let mask =
947 Array::<bool, Ix1>::from_vec(Ix1::new([5]), vec![false, false, true, false, false])
948 .unwrap();
949 let ma = MaskedArray::new(data, mask).unwrap();
950 let v0 = ma.var_ddof(0).unwrap();
951 let v1 = ma.var_ddof(1).unwrap();
952 assert!((v0 - 2.5).abs() < 1e-14);
953 assert!((v1 - 10.0 / 3.0).abs() < 1e-14);
954 }
955
956 #[test]
957 fn var_ddof_returns_nan_when_count_le_ddof() {
958 let data = Array::<f64, Ix1>::from_vec(Ix1::new([3]), vec![1.0, 2.0, 3.0]).unwrap();
961 let mask = Array::<bool, Ix1>::from_vec(Ix1::new([3]), vec![true, false, true]).unwrap();
962 let ma = MaskedArray::new(data, mask).unwrap();
963 let v = ma.var_ddof(1).unwrap();
964 assert!(v.is_nan(), "expected NaN, got {v}");
965 }
966
967 #[test]
968 fn std_ddof_is_sqrt_of_var_ddof() {
969 let data =
970 Array::<f64, Ix1>::from_vec(Ix1::new([5]), vec![1.0, 2.0, 3.0, 4.0, 5.0]).unwrap();
971 let mask = Array::<bool, Ix1>::from_vec(Ix1::new([5]), vec![false; 5]).unwrap();
972 let ma = MaskedArray::new(data, mask).unwrap();
973 let s1 = ma.std_ddof(1).unwrap();
974 let v1 = ma.var_ddof(1).unwrap();
975 assert!((s1 - v1.sqrt()).abs() < 1e-14);
976 }
977
978 #[test]
979 fn var_welford_stable_on_high_offset_data() {
980 let offset = 1e9_f64;
985 let data: Vec<f64> = (1..=5).map(|i| offset + i as f64).collect();
986 let arr = Array::<f64, Ix1>::from_vec(Ix1::new([5]), data).unwrap();
987 let mask = Array::<bool, Ix1>::from_vec(Ix1::new([5]), vec![false; 5]).unwrap();
988 let ma = MaskedArray::new(arr, mask).unwrap();
989 let v = ma.var().unwrap();
990 assert!(
993 (v - 2.0).abs() < 1e-9,
994 "var with offset 1e9: got {v}, expected 2.0"
995 );
996 }
997
998 #[test]
999 fn var_axis_ddof_one_per_row() {
1000 use ferray_core::dimension::Ix2;
1001 let data =
1005 Array::<f64, Ix2>::from_vec(Ix2::new([2, 3]), vec![1.0, 2.0, 3.0, 10.0, 20.0, 30.0])
1006 .unwrap();
1007 let mask = Array::<bool, Ix2>::from_vec(Ix2::new([2, 3]), vec![false; 6]).unwrap();
1008 let ma = MaskedArray::new(data, mask).unwrap();
1009 let v = ma.var_axis_ddof(1, 1).unwrap();
1010 let vs: Vec<f64> = v.data().iter().copied().collect();
1011 assert_eq!(vs.len(), 2);
1012 assert!((vs[0] - 1.0).abs() < 1e-12);
1013 assert!((vs[1] - 100.0).abs() < 1e-12);
1014 }
1015}