Skip to main content

ferray_ma/
reductions.rs

1// ferray-ma: Masked reductions (REQ-4)
2//
3// mean, sum, min, max, var, std, count — all skip masked elements.
4//
5// Whole-array variants (`sum`, `mean`, ...) return a scalar `T`.
6// Per-axis variants (`sum_axis`, `mean_axis`, ...) return a `MaskedArray<T, IxDyn>`
7// where each output position holds the reduction of one lane along the
8// chosen axis. Lanes that contain only masked elements are themselves
9// masked in the output (and hold the source array's `fill_value`).
10//
11// ## REQ status
12//
13// REQ-4 SHIPPED — audited, green. All reductions skip masked elements; the
14// all-masked case surfaces as `NaN` (ferray's analog of numpy's `masked`
15// singleton, `numpy/ma/core.py:5249` `elif newmask: result = masked`).
16//
17// | REQ | Status | Evidence |
18// |-----|--------|----------|
19// | REQ-4 (mask-skipping reductions) | SHIPPED | Whole-array `count`/`sum`/`mean`/`min`/`max`/`var`(`_ddof`)/`std`(`_ddof`) and the per-axis `count_axis`/`sum_axis`/`mean_axis`/`min_axis`/`max_axis`/`var_axis`(`_ddof`)/`std_axis`(`_ddof`) (this file) all divide by the unmasked count and mask all-masked lanes. Non-test production consumers: the `PyMaskedArray` reduction methods `sum`/`mean`/`var`/`std`/`count`/`min`/`max` and the `sum_axis`/`mean_axis` axis methods in `ferray-python/src/ma.rs`. All-masked `sum` returning `NaN` (not `0.0`) is additionally pinned by the divergence test `divergence_all_masked_sum_is_masked_not_zero`. |
20//
21// `argmin`/`argmax`/`prod`/`median`/`ptp` live in `extras.rs` (their REQ rows
22// are documented there); this file owns the original REQ-4 reduction family.
23
24use 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
32// ---------------------------------------------------------------------------
33// Internal helpers for axis-aware reductions
34// ---------------------------------------------------------------------------
35
36/// Compute row-major strides for `shape`.
37fn 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
46/// Increment a multi-index in row-major order. Returns false on overflow.
47fn 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
58/// Apply a per-lane masked reduction along `axis`.
59///
60/// `kernel` receives a `&[(T, bool)]` slice (data + mask) for each lane and
61/// returns either `Some(value)` (the reduction result) or `None` if every
62/// element in the lane was masked. Masked output positions are filled with
63/// `fill_value`.
64fn 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    // Output shape: drop the reduced axis.
83    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    // Materialize source data + mask in row-major order so we can index by
95    // computed flat indices regardless of the source memory layout.
96    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        // Map output multi-index back into the input multi-index by inserting
108        // a placeholder at `axis`.
109        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
150/// Per-axis count of unmasked elements. Returns a plain `Array<u64, IxDyn>`
151/// (not masked, since count is always defined) with the reduced axis dropped.
152fn 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    /// Count the number of unmasked (valid) elements.
219    ///
220    /// # Errors
221    /// This function does not currently error but returns `Result` for API
222    /// consistency.
223    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    /// Count the number of unmasked elements per lane along `axis`.
234    ///
235    /// Returns a plain `Array<u64, IxDyn>` (not masked, since count is
236    /// always defined) with the reduced axis removed.
237    ///
238    /// # Errors
239    /// Returns `FerrayError::AxisOutOfBounds` if `axis >= ndim`.
240    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    /// Compute the sum of unmasked elements.
251    ///
252    /// When **every** element is masked the result is `NaN`, not the bare
253    /// additive identity `0.0`. numpy's `MaskedArray.sum` returns the
254    /// `masked` singleton in this case (`numpy/ma/core.py:5249`
255    /// `elif newmask: result = masked`) — an all-masked sum has no defined
256    /// numeric value, so `0.0` (the fold's neutral element) would be a
257    /// wrong observable. ferray has no Python `masked` singleton, so we
258    /// surface the "no defined value" status as `NaN`, consistent with
259    /// [`MaskedArray::mean`]'s all-masked behavior. Live oracle (numpy
260    /// 2.4.5): `ma.array([1.,2.,3.,4.], mask=[1,1,1,1]).sum() is ma.masked`.
261    ///
262    /// # Errors
263    /// Returns an error only for internal failures.
264    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    /// Compute the mean of unmasked elements.
280    ///
281    /// Returns `NaN` if no elements are unmasked.
282    ///
283    /// # Errors
284    /// Returns an error only for internal failures.
285    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        // #267: T::from(count).unwrap_or(one) silently returned the
297        // sum as the "mean" if the conversion ever failed. Surface a
298        // typed error instead so a downstream NaN/garbage isn't
299        // misattributed to upstream data.
300        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    /// Compute the minimum of unmasked elements.
309    ///
310    /// NaN values in unmasked elements are propagated (returns NaN), matching `NumPy`.
311    ///
312    /// # Errors
313    /// Returns `FerrayError::InvalidValue` if no elements are unmasked.
314    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                        // NaN-propagating: if comparison is unordered, propagate NaN
324                        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    /// Compute the maximum of unmasked elements.
339    ///
340    /// NaN values in unmasked elements are propagated (returns NaN), matching `NumPy`.
341    ///
342    /// # Errors
343    /// Returns `FerrayError::InvalidValue` if no elements are unmasked.
344    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    /// Compute the variance of unmasked elements (population variance, ddof=0).
368    ///
369    /// Returns `NaN` if no elements are unmasked.
370    ///
371    /// # Errors
372    /// Returns an error only for internal failures.
373    pub fn var(&self) -> FerrayResult<T> {
374        self.var_ddof(0)
375    }
376
377    /// Compute the variance of unmasked elements with a delta degrees-of-freedom
378    /// adjustment.
379    ///
380    /// `ddof = 0` is the population variance (divides by `n`).
381    /// `ddof = 1` is Bessel's correction for sample variance (divides
382    /// by `n - 1`). Matches `numpy.ma.var(ddof=...)` (#270).
383    ///
384    /// Returns `NaN` if `count <= ddof` (insufficient unmasked elements).
385    ///
386    /// # Errors
387    /// Returns an error only for internal failures.
388    pub fn var_ddof(&self, ddof: usize) -> FerrayResult<T> {
389        // Single-pass Welford accumulator (#274). Replaces the previous
390        // two-pass mean-then-sum-of-squares form: one walk over the
391        // array updates the running mean and the m2 sum-of-squared-
392        // deviations simultaneously. Numerically more stable than the
393        // two-pass form for large datasets (less catastrophic
394        // cancellation when mean is far from zero).
395        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    /// Compute the standard deviation of unmasked elements (population, ddof=0).
430    ///
431    /// Returns `NaN` if no elements are unmasked.
432    ///
433    /// # Errors
434    /// Returns an error only for internal failures.
435    pub fn std(&self) -> FerrayResult<T> {
436        Ok(self.var()?.sqrt())
437    }
438
439    /// Compute the standard deviation of unmasked elements with a delta
440    /// degrees-of-freedom adjustment.
441    ///
442    /// `ddof = 0` is the population std; `ddof = 1` is the sample std
443    /// (Bessel's correction). Matches `numpy.ma.std(ddof=...)` (#270).
444    ///
445    /// Returns `NaN` if `count <= ddof`.
446    ///
447    /// # Errors
448    /// Returns an error only for internal failures.
449    pub fn std_ddof(&self, ddof: usize) -> FerrayResult<T> {
450        Ok(self.var_ddof(ddof)?.sqrt())
451    }
452
453    // -----------------------------------------------------------------------
454    // Per-axis reductions (issue #500)
455    //
456    // Each lane along `axis` is reduced independently. Lanes containing only
457    // masked elements produce a masked output position holding `fill_value`.
458    // -----------------------------------------------------------------------
459
460    /// Sum unmasked elements along `axis`. Returns a masked array with the
461    /// reduced axis removed; lanes that are entirely masked produce a
462    /// masked output position holding `fill_value`.
463    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    /// Mean of unmasked elements along `axis`. All-masked lanes are masked
480    /// in the output.
481    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                // #267: a failed count→T conversion would silently
497                // divide by 1, returning the sum as the "mean". Mask
498                // the cell instead — for f32/f64 the conversion never
499                // fails, so this only triggers on pathological types.
500                T::from(count).map(|n| acc / n)
501            }
502        })
503    }
504
505    /// Min of unmasked elements along `axis`. NaN-propagating per `NumPy`.
506    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                            // NaN-propagating: if comparison is unordered, return NaN
515                            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    /// Max of unmasked elements along `axis`. NaN-propagating per `NumPy`.
532    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    /// Population variance (ddof=0) of unmasked elements along `axis`.
557    pub fn var_axis(&self, axis: usize) -> FerrayResult<MaskedArray<T, IxDyn>> {
558        self.var_axis_ddof(axis, 0)
559    }
560
561    /// Variance with `ddof` adjustment along `axis` (#270).
562    ///
563    /// `ddof = 1` matches numpy's sample variance (Bessel's correction).
564    /// Lanes with `count <= ddof` produce a masked output position
565    /// holding `fill_value`.
566    pub fn var_axis_ddof(&self, axis: usize, ddof: usize) -> FerrayResult<MaskedArray<T, IxDyn>> {
567        // Single-pass Welford accumulator per lane (#274).
568        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    /// Population standard deviation (ddof=0) of unmasked elements along `axis`.
594    pub fn std_axis(&self, axis: usize) -> FerrayResult<MaskedArray<T, IxDyn>> {
595        self.std_axis_ddof(axis, 0)
596    }
597
598    /// Standard deviation with `ddof` adjustment along `axis` (#270).
599    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    // ---- #500: per-axis reductions ----
628
629    #[test]
630    fn sum_axis_drops_axis() {
631        // 2x3 array, no masks. axis=0 sums columns, axis=1 sums rows.
632        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        // 2x3 array. Mask out (0, 1) so column 1 and row 0 lose one element.
647        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        // axis=0 (per-column): col0 = 1+4=5, col1 = 5 (only row 1), col2 = 3+6=9
654        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        // 2x3 array, mask out entire column 1.
664        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        // mean axis 1 (per-row): row 0 = (1+3)/2 = 2.0, row 1 = (4+5+6)/3 = 5.0
684        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        // 2x3 array. Mask: (0,1) and (1,2). col0:2, col1:1, col2:1.
703        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        // Two rows of [1, 2, 3, 4, 5] — variance along axis 1 should be 2.0 each.
723        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    // ---- #501: fill_value ----
738
739    #[test]
740    fn fill_value_default_is_numpy_default_filler_1e20() {
741        // numpy/ma/core.py:166 `default_filler['f'] = 1.e20`; live oracle
742        // numpy 2.4.5: `ma.array([1.,2.,3.]).fill_value == np.float64(1e+20)`.
743        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        // (Adding two masked arrays) — result data at masked positions should
771        // be the receiver's fill_value, not zero.
772        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    // ---- #504: broadcasting in masked arithmetic ----
786
787    #[test]
788    fn masked_add_broadcasts_within_same_rank() {
789        use crate::masked_add;
790        // (3, 1) + (1, 4) -> (3, 4) — both Ix2.
791        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, // row 0 = 1 + {10,20,30,40}
805                12.0, 22.0, 32.0, 42.0, // row 1 = 2 + ...
806                13.0, 23.0, 33.0, 43.0, // row 2 = 3 + ...
807            ]
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        // Mask one element in `a`. After broadcasting (3,1) -> (3,4),
817        // the masked row becomes a fully-masked row in the result.
818        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        // Row 1 is fully masked (4 positions).
827        assert_eq!(
828            r_m,
829            vec![
830                false, false, false, false, // row 0
831                true, true, true, true, // row 1 (masked)
832                false, false, false, false, // row 2
833            ]
834        );
835    }
836
837    // ---- #276: all-masked whole-array reductions ----
838    //
839    // Semantics when every element is masked (matching numpy.ma):
840    //   sum   -> NaN (numpy returns the `masked` singleton, core.py:5249;
841    //            NOT the additive identity 0.0)
842    //   mean  -> NaN
843    //   var   -> NaN
844    //   std   -> NaN
845    //   min   -> error (FerrayError::InvalidValue)
846    //   max   -> error (FerrayError::InvalidValue)
847
848    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        // numpy/ma/core.py:5249: an all-masked sum is the `masked`
857        // singleton, NOT 0.0. ferray surfaces that as NaN.
858        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        // Documenting the asymmetry: sum/mean/var/std fall through with
883        // 0/NaN sentinels, but min/max have no neutral element and error.
884        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        // Same semantics hold for multi-dimensional whole-array reductions.
892        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    // ----- ddof=1 sample variance/std (#270) -----------------------------
913
914    #[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        // [1,2,3,4,5] mean = 3; squared deviations sum = 10.
928        // ddof=0: 10/5 = 2.0
929        // ddof=1: 10/4 = 2.5 (Bessel)
930        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        // [1,2,_,4,5] with index 2 masked: mean = 3, sq deviations = 10, count = 4.
943        // ddof=0: 10/4 = 2.5; ddof=1: 10/3.
944        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        // 1 unmasked element + ddof=1 → division by zero would occur;
959        // numpy returns NaN.
960        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        // #274: classic two-pass var loses precision when the data is
981        // a small variance riding on a huge offset. Welford handles
982        // it cleanly. Use values [1e9 + 1, 1e9 + 2, 1e9 + 3, 1e9 + 4,
983        // 1e9 + 5]; ddof=0 var should be exactly 2.0.
984        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        // Welford should be within ~ULPs of 2.0; the previous two-pass
991        // form was within ~1e-7 at this scale.
992        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        // Two rows: [1,2,3] and [10,20,30]; both unmasked.
1002        // Row 0: mean=2, sum_sq = 1+0+1 = 2; ddof=1 → 2.
1003        // Row 1: mean=20, sum_sq = 100+0+100 = 200; ddof=1 → 100.
1004        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}