Skip to main content

ferray_ma/
lib.rs

1// ferray-ma: Masked arrays with mask propagation
2//
3// This crate implements `numpy.ma`-style masked arrays for the ferray workspace.
4// A `MaskedArray<T, D>` pairs a data array with a boolean mask array where
5// `true` = masked/invalid. All operations (arithmetic, reductions, ufuncs)
6// respect the mask by skipping masked elements.
7//
8// # Modules
9// - `masked_array`: The core `MaskedArray<T, D>` type
10// - `reductions`: Masked mean, sum, min, max, var, std, count
11// - `constructors`: masked_where, masked_invalid, masked_equal, etc.
12// - `arithmetic`: Masked binary ops with mask union
13// - `ufunc_support`: Wrapper functions for ufunc operations on MaskedArrays
14// - `sorting`: Masked sort, argsort
15// - `mask_ops`: harden_mask, soften_mask, getmask, getdata, is_masked, count_masked
16// - `filled`: filled, compressed
17
18// Masked reductions divide running sums by valid-element counts and
19// truncate `f64` results to integer index types in argmin/argmax. Float
20// equality is also intrinsic to `masked_equal` and `getdata` semantics.
21#![allow(
22    clippy::cast_possible_truncation,
23    clippy::cast_possible_wrap,
24    clippy::cast_precision_loss,
25    clippy::cast_sign_loss,
26    clippy::cast_lossless,
27    clippy::float_cmp,
28    clippy::missing_errors_doc,
29    clippy::missing_panics_doc,
30    clippy::many_single_char_names,
31    clippy::similar_names,
32    clippy::items_after_statements,
33    clippy::option_if_let_else,
34    clippy::too_long_first_doc_paragraph,
35    clippy::needless_pass_by_value,
36    clippy::match_same_arms
37)]
38
39pub mod arithmetic;
40pub mod constructors;
41pub mod filled;
42pub mod interop;
43/// Binary I/O (save/load) for `MaskedArray` via ferray-io (#509).
44///
45/// Gated behind the `io` cargo feature so callers who don't need
46/// disk I/O don't have to pull in the zip + binary reader dependency
47/// tree through ferray-io.
48#[cfg(feature = "io")]
49pub mod io;
50pub mod manipulation;
51pub mod mask_ops;
52pub mod masked_array;
53pub mod reductions;
54pub mod sorting;
55pub mod ufunc_support;
56
57// Re-export the primary type at crate root
58pub use masked_array::MaskedArray;
59
60// Re-export masking constructors
61pub use constructors::{
62    fix_invalid, masked_equal, masked_greater, masked_greater_equal, masked_inside, masked_invalid,
63    masked_less, masked_less_equal, masked_not_equal, masked_outside, masked_where,
64};
65
66// Re-export arithmetic operations
67pub use arithmetic::{
68    masked_add, masked_add_array, masked_div, masked_div_array, masked_mul, masked_mul_array,
69    masked_sub, masked_sub_array,
70};
71
72// Re-export mask manipulation functions
73pub use mask_ops::{count_masked, getdata, getmask, is_masked};
74
75// Re-export MaskAware trait (#505) for downstream code that wants
76// to write functions polymorphic over Array and MaskedArray.
77pub use interop::{MaskAware, ma_apply_unary};
78
79// Re-export generic ufunc helpers (#513) — the escape hatch for
80// ufuncs that don't have a dedicated named wrapper. Users with an
81// arbitrary `Fn(T) -> T` / `Fn(T, T) -> T` closure can call
82// `ferray_ma::masked_unary(ma, my_fn)` directly.
83pub use ufunc_support::{masked_binary, masked_unary};
84
85// Domain-aware ufunc wrappers (#503) — auto-mask out-of-domain
86// inputs so the result mask carries a "safe to use" contract.
87pub use ufunc_support::{
88    arccos_domain, arccosh_domain, arcsin_domain, arctanh_domain, divide_domain, log_domain,
89    log2_domain, log10_domain, masked_binary_domain, masked_unary_domain, sqrt_domain,
90};
91
92#[cfg(test)]
93mod tests {
94    use super::*;
95    use ferray_core::Array;
96    use ferray_core::dimension::Ix1;
97
98    // -----------------------------------------------------------------------
99    // AC-1: MaskedArray::new([1,2,3,4,5], [false,false,true,false,false]).mean() == 3.0
100    // -----------------------------------------------------------------------
101    #[test]
102    fn ac1_masked_mean_skips_masked() {
103        let data =
104            Array::<f64, Ix1>::from_vec(Ix1::new([5]), vec![1.0, 2.0, 3.0, 4.0, 5.0]).unwrap();
105        let mask =
106            Array::<bool, Ix1>::from_vec(Ix1::new([5]), vec![false, false, true, false, false])
107                .unwrap();
108        let ma = MaskedArray::new(data, mask).unwrap();
109        let mean = ma.mean().unwrap();
110        // (1 + 2 + 4 + 5) / 4 = 3.0
111        assert!((mean - 3.0).abs() < 1e-10);
112    }
113
114    // -----------------------------------------------------------------------
115    // AC-2: filled(0.0) replaces masked elements with 0.0
116    // -----------------------------------------------------------------------
117    #[test]
118    fn ac2_filled_replaces_masked() {
119        let data = Array::<f64, Ix1>::from_vec(Ix1::new([4]), vec![1.0, 2.0, 3.0, 4.0]).unwrap();
120        let mask =
121            Array::<bool, Ix1>::from_vec(Ix1::new([4]), vec![false, true, false, true]).unwrap();
122        let ma = MaskedArray::new(data, mask).unwrap();
123        let filled = ma.filled(0.0).unwrap();
124        assert_eq!(filled.as_slice().unwrap(), &[1.0, 0.0, 3.0, 0.0]);
125    }
126
127    // -----------------------------------------------------------------------
128    // AC-3: compressed() returns only unmasked elements as 1D
129    // -----------------------------------------------------------------------
130    #[test]
131    fn ac3_compressed_returns_unmasked() {
132        let data =
133            Array::<f64, Ix1>::from_vec(Ix1::new([5]), vec![10.0, 20.0, 30.0, 40.0, 50.0]).unwrap();
134        let mask =
135            Array::<bool, Ix1>::from_vec(Ix1::new([5]), vec![false, true, false, true, false])
136                .unwrap();
137        let ma = MaskedArray::new(data, mask).unwrap();
138        let compressed = ma.compressed().unwrap();
139        assert_eq!(compressed.as_slice().unwrap(), &[10.0, 30.0, 50.0]);
140    }
141
142    // -----------------------------------------------------------------------
143    // AC-4: masked_invalid masks NaN and Inf
144    // -----------------------------------------------------------------------
145    #[test]
146    fn ac4_masked_invalid_nan_inf() {
147        let data =
148            Array::<f64, Ix1>::from_vec(Ix1::new([4]), vec![1.0, f64::NAN, 3.0, f64::INFINITY])
149                .unwrap();
150        let ma = masked_invalid(&data).unwrap();
151        let mask_vals: Vec<bool> = ma.mask().iter().copied().collect();
152        assert_eq!(mask_vals, vec![false, true, false, true]);
153    }
154
155    // -----------------------------------------------------------------------
156    // AC-5: ma1 + ma2 produces correct mask union and correct values
157    // -----------------------------------------------------------------------
158    #[test]
159    fn ac5_add_mask_union() {
160        let d1 = Array::<f64, Ix1>::from_vec(Ix1::new([4]), vec![1.0, 2.0, 3.0, 4.0]).unwrap();
161        let m1 =
162            Array::<bool, Ix1>::from_vec(Ix1::new([4]), vec![false, true, false, false]).unwrap();
163        let ma1 = MaskedArray::new(d1, m1).unwrap();
164
165        let d2 = Array::<f64, Ix1>::from_vec(Ix1::new([4]), vec![10.0, 20.0, 30.0, 40.0]).unwrap();
166        let m2 =
167            Array::<bool, Ix1>::from_vec(Ix1::new([4]), vec![false, false, true, false]).unwrap();
168        let ma2 = MaskedArray::new(d2, m2).unwrap();
169
170        let result = masked_add(&ma1, &ma2).unwrap();
171        let mask_vals: Vec<bool> = result.mask().iter().copied().collect();
172        // Mask union: [false, true, true, false]
173        assert_eq!(mask_vals, vec![false, true, true, false]);
174        // Unmasked values: 1+10=11, 4+40=44; masked get 0.0
175        let data_vals: Vec<f64> = result.data().iter().copied().collect();
176        assert!((data_vals[0] - 11.0).abs() < 1e-10);
177        assert!((data_vals[3] - 44.0).abs() < 1e-10);
178    }
179
180    #[test]
181    fn operator_add_matches_masked_add() {
182        let d1 = Array::<f64, Ix1>::from_vec(Ix1::new([3]), vec![1.0, 2.0, 3.0]).unwrap();
183        let m1 = Array::<bool, Ix1>::from_vec(Ix1::new([3]), vec![false, true, false]).unwrap();
184        let ma1 = MaskedArray::new(d1, m1).unwrap();
185
186        let d2 = Array::<f64, Ix1>::from_vec(Ix1::new([3]), vec![10.0, 20.0, 30.0]).unwrap();
187        let m2 = Array::<bool, Ix1>::from_vec(Ix1::new([3]), vec![false, false, true]).unwrap();
188        let ma2 = MaskedArray::new(d2, m2).unwrap();
189
190        // Use operator syntax
191        let result = (&ma1 + &ma2).unwrap();
192        let mask_vals: Vec<bool> = result.mask().iter().copied().collect();
193        assert_eq!(mask_vals, vec![false, true, true]);
194        let data_vals: Vec<f64> = result.data().iter().copied().collect();
195        assert!((data_vals[0] - 11.0).abs() < 1e-10);
196    }
197
198    // -----------------------------------------------------------------------
199    // AC-7: sin(masked_array) returns same mask, correct values
200    // -----------------------------------------------------------------------
201    #[test]
202    fn ac7_ufunc_sin_masked() {
203        use std::f64::consts::FRAC_PI_2;
204        let data =
205            Array::<f64, Ix1>::from_vec(Ix1::new([3]), vec![0.0, FRAC_PI_2, FRAC_PI_2]).unwrap();
206        let mask = Array::<bool, Ix1>::from_vec(Ix1::new([3]), vec![false, true, false]).unwrap();
207        let ma = MaskedArray::new(data, mask).unwrap();
208        let result = ufunc_support::sin(&ma).unwrap();
209        let mask_vals: Vec<bool> = result.mask().iter().copied().collect();
210        assert_eq!(mask_vals, vec![false, true, false]);
211        let data_vals: Vec<f64> = result.data().iter().copied().collect();
212        // sin(0) = 0, masked = 0.0 (skipped), sin(pi/2) = 1.0
213        assert!((data_vals[0] - 0.0).abs() < 1e-10);
214        assert!((data_vals[2] - 1.0).abs() < 1e-10);
215    }
216
217    // -----------------------------------------------------------------------
218    // AC-8: sort places masked at end; harden_mask prevents clearing
219    // -----------------------------------------------------------------------
220    #[test]
221    fn ac8_sort_masked_at_end() {
222        let data =
223            Array::<f64, Ix1>::from_vec(Ix1::new([5]), vec![5.0, 1.0, 3.0, 2.0, 4.0]).unwrap();
224        let mask =
225            Array::<bool, Ix1>::from_vec(Ix1::new([5]), vec![false, false, true, false, false])
226                .unwrap();
227        let ma = MaskedArray::new(data, mask).unwrap();
228        let sorted = ma.sort().unwrap();
229        let data_vals: Vec<f64> = sorted.data().iter().copied().collect();
230        let mask_vals: Vec<bool> = sorted.mask().iter().copied().collect();
231        // Unmasked [5, 1, 2, 4] sorted = [1, 2, 4, 5], then masked [3]
232        assert_eq!(data_vals, vec![1.0, 2.0, 4.0, 5.0, 3.0]);
233        assert_eq!(mask_vals, vec![false, false, false, false, true]);
234    }
235
236    #[test]
237    fn ac8_harden_mask_prevents_clearing() {
238        let data = Array::<f64, Ix1>::from_vec(Ix1::new([3]), vec![1.0, 2.0, 3.0]).unwrap();
239        let mask = Array::<bool, Ix1>::from_vec(Ix1::new([3]), vec![false, true, false]).unwrap();
240        let mut ma = MaskedArray::new(data, mask).unwrap();
241
242        ma.harden_mask().unwrap();
243        assert!(ma.is_hard_mask());
244
245        // Try to clear the mask at index 1 — should be silently ignored
246        ma.set_mask_flat(1, false).unwrap();
247        let mask_vals: Vec<bool> = ma.mask().iter().copied().collect();
248        assert_eq!(mask_vals, vec![false, true, false]);
249
250        // Setting a mask bit to true should still work
251        ma.set_mask_flat(0, true).unwrap();
252        let mask_vals2: Vec<bool> = ma.mask().iter().copied().collect();
253        assert_eq!(mask_vals2, vec![true, true, false]);
254
255        // Soften and then clearing should work
256        ma.soften_mask().unwrap();
257        assert!(!ma.is_hard_mask());
258        ma.set_mask_flat(1, false).unwrap();
259        let mask_vals3: Vec<bool> = ma.mask().iter().copied().collect();
260        assert_eq!(mask_vals3, vec![true, false, false]);
261    }
262
263    // -----------------------------------------------------------------------
264    // AC-9: is_masked returns true/false correctly
265    // -----------------------------------------------------------------------
266    #[test]
267    fn ac9_is_masked() {
268        let data1 = Array::<f64, Ix1>::from_vec(Ix1::new([3]), vec![1.0, 2.0, 3.0]).unwrap();
269        let mask1 = Array::<bool, Ix1>::from_vec(Ix1::new([3]), vec![false, true, false]).unwrap();
270        let ma1 = MaskedArray::new(data1, mask1).unwrap();
271        assert!(is_masked(&ma1).unwrap());
272
273        let data2 = Array::<f64, Ix1>::from_vec(Ix1::new([3]), vec![1.0, 2.0, 3.0]).unwrap();
274        let mask2 = Array::<bool, Ix1>::from_vec(Ix1::new([3]), vec![false, false, false]).unwrap();
275        let ma2 = MaskedArray::new(data2, mask2).unwrap();
276        assert!(!is_masked(&ma2).unwrap());
277    }
278
279    // -----------------------------------------------------------------------
280    // Additional tests
281    // -----------------------------------------------------------------------
282
283    #[test]
284    fn shape_mismatch_error() {
285        let data = Array::<f64, Ix1>::from_vec(Ix1::new([3]), vec![1.0, 2.0, 3.0]).unwrap();
286        let mask = Array::<bool, Ix1>::from_vec(Ix1::new([2]), vec![false, true]).unwrap();
287        assert!(MaskedArray::new(data, mask).is_err());
288    }
289
290    #[test]
291    fn from_data_no_mask() {
292        let data = Array::<f64, Ix1>::from_vec(Ix1::new([3]), vec![1.0, 2.0, 3.0]).unwrap();
293        let ma = MaskedArray::from_data(data).unwrap();
294        assert!(!is_masked(&ma).unwrap());
295        assert_eq!(ma.count().unwrap(), 3);
296    }
297
298    #[test]
299    fn sum_skips_masked() {
300        let data = Array::<f64, Ix1>::from_vec(Ix1::new([4]), vec![1.0, 2.0, 3.0, 4.0]).unwrap();
301        let mask =
302            Array::<bool, Ix1>::from_vec(Ix1::new([4]), vec![false, true, false, true]).unwrap();
303        let ma = MaskedArray::new(data, mask).unwrap();
304        assert!((ma.sum().unwrap() - 4.0).abs() < 1e-10);
305    }
306
307    #[test]
308    fn min_max_skip_masked() {
309        let data =
310            Array::<f64, Ix1>::from_vec(Ix1::new([5]), vec![5.0, 1.0, 3.0, 2.0, 4.0]).unwrap();
311        let mask =
312            Array::<bool, Ix1>::from_vec(Ix1::new([5]), vec![false, true, false, false, false])
313                .unwrap();
314        let ma = MaskedArray::new(data, mask).unwrap();
315        assert!((ma.min().unwrap() - 2.0).abs() < 1e-10);
316        assert!((ma.max().unwrap() - 5.0).abs() < 1e-10);
317    }
318
319    #[test]
320    fn var_std_skip_masked() {
321        // values: [2, 4, 6] (mask out index 1 and 4)
322        let data =
323            Array::<f64, Ix1>::from_vec(Ix1::new([5]), vec![2.0, 99.0, 4.0, 6.0, 99.0]).unwrap();
324        let mask =
325            Array::<bool, Ix1>::from_vec(Ix1::new([5]), vec![false, true, false, false, true])
326                .unwrap();
327        let ma = MaskedArray::new(data, mask).unwrap();
328        let mean = ma.mean().unwrap();
329        assert!((mean - 4.0).abs() < 1e-10);
330        // var = ((2-4)^2 + (4-4)^2 + (6-4)^2) / 3 = 8/3
331        let v = ma.var().unwrap();
332        assert!((v - 8.0 / 3.0).abs() < 1e-10);
333        let s = ma.std().unwrap();
334        assert!((s - (8.0_f64 / 3.0).sqrt()).abs() < 1e-10);
335    }
336
337    #[test]
338    fn count_elements() {
339        let data = Array::<f64, Ix1>::from_vec(Ix1::new([5]), vec![1.0; 5]).unwrap();
340        let mask =
341            Array::<bool, Ix1>::from_vec(Ix1::new([5]), vec![false, true, true, false, false])
342                .unwrap();
343        let ma = MaskedArray::new(data, mask).unwrap();
344        assert_eq!(ma.count().unwrap(), 3);
345    }
346
347    #[test]
348    fn masked_equal_test() {
349        let data =
350            Array::<f64, Ix1>::from_vec(Ix1::new([5]), vec![1.0, 2.0, 3.0, 2.0, 1.0]).unwrap();
351        let ma = masked_equal(&data, 2.0).unwrap();
352        let mask_vals: Vec<bool> = ma.mask().iter().copied().collect();
353        assert_eq!(mask_vals, vec![false, true, false, true, false]);
354    }
355
356    #[test]
357    fn masked_greater_test() {
358        let data = Array::<f64, Ix1>::from_vec(Ix1::new([4]), vec![1.0, 2.0, 3.0, 4.0]).unwrap();
359        let ma = masked_greater(&data, 2.0).unwrap();
360        let mask_vals: Vec<bool> = ma.mask().iter().copied().collect();
361        assert_eq!(mask_vals, vec![false, false, true, true]);
362    }
363
364    #[test]
365    fn masked_less_test() {
366        let data = Array::<f64, Ix1>::from_vec(Ix1::new([4]), vec![1.0, 2.0, 3.0, 4.0]).unwrap();
367        let ma = masked_less(&data, 3.0).unwrap();
368        let mask_vals: Vec<bool> = ma.mask().iter().copied().collect();
369        assert_eq!(mask_vals, vec![true, true, false, false]);
370    }
371
372    #[test]
373    fn masked_not_equal_test() {
374        let data = Array::<f64, Ix1>::from_vec(Ix1::new([3]), vec![1.0, 2.0, 3.0]).unwrap();
375        let ma = masked_not_equal(&data, 2.0).unwrap();
376        let mask_vals: Vec<bool> = ma.mask().iter().copied().collect();
377        assert_eq!(mask_vals, vec![true, false, true]);
378    }
379
380    #[test]
381    fn masked_greater_equal_test() {
382        let data = Array::<f64, Ix1>::from_vec(Ix1::new([4]), vec![1.0, 2.0, 3.0, 4.0]).unwrap();
383        let ma = masked_greater_equal(&data, 3.0).unwrap();
384        let mask_vals: Vec<bool> = ma.mask().iter().copied().collect();
385        assert_eq!(mask_vals, vec![false, false, true, true]);
386    }
387
388    #[test]
389    fn masked_less_equal_test() {
390        let data = Array::<f64, Ix1>::from_vec(Ix1::new([4]), vec![1.0, 2.0, 3.0, 4.0]).unwrap();
391        let ma = masked_less_equal(&data, 2.0).unwrap();
392        let mask_vals: Vec<bool> = ma.mask().iter().copied().collect();
393        assert_eq!(mask_vals, vec![true, true, false, false]);
394    }
395
396    #[test]
397    fn masked_inside_test() {
398        let data =
399            Array::<f64, Ix1>::from_vec(Ix1::new([5]), vec![1.0, 2.0, 3.0, 4.0, 5.0]).unwrap();
400        let ma = masked_inside(&data, 2.0, 4.0).unwrap();
401        let mask_vals: Vec<bool> = ma.mask().iter().copied().collect();
402        assert_eq!(mask_vals, vec![false, true, true, true, false]);
403    }
404
405    #[test]
406    fn masked_outside_test() {
407        let data =
408            Array::<f64, Ix1>::from_vec(Ix1::new([5]), vec![1.0, 2.0, 3.0, 4.0, 5.0]).unwrap();
409        let ma = masked_outside(&data, 2.0, 4.0).unwrap();
410        let mask_vals: Vec<bool> = ma.mask().iter().copied().collect();
411        assert_eq!(mask_vals, vec![true, false, false, false, true]);
412    }
413
414    #[test]
415    fn masked_where_test() {
416        let data = Array::<f64, Ix1>::from_vec(Ix1::new([4]), vec![1.0, 2.0, 3.0, 4.0]).unwrap();
417        let cond =
418            Array::<bool, Ix1>::from_vec(Ix1::new([4]), vec![true, false, true, false]).unwrap();
419        let ma = masked_where(&cond, &data).unwrap();
420        let mask_vals: Vec<bool> = ma.mask().iter().copied().collect();
421        assert_eq!(mask_vals, vec![true, false, true, false]);
422    }
423
424    #[test]
425    fn argsort_test() {
426        let data =
427            Array::<f64, Ix1>::from_vec(Ix1::new([5]), vec![5.0, 1.0, 3.0, 2.0, 4.0]).unwrap();
428        let mask =
429            Array::<bool, Ix1>::from_vec(Ix1::new([5]), vec![false, false, true, false, false])
430                .unwrap();
431        let ma = MaskedArray::new(data, mask).unwrap();
432        let indices = ma.argsort().unwrap();
433        let idx_vals: Vec<usize> = indices;
434        // Unmasked: index 1 (1.0), 3 (2.0), 4 (4.0), 0 (5.0); masked: 2
435        assert_eq!(idx_vals, vec![1, 3, 4, 0, 2]);
436    }
437
438    #[test]
439    fn getmask_getdata_test() {
440        let data = Array::<f64, Ix1>::from_vec(Ix1::new([3]), vec![1.0, 2.0, 3.0]).unwrap();
441        let mask = Array::<bool, Ix1>::from_vec(Ix1::new([3]), vec![false, true, false]).unwrap();
442        let ma = MaskedArray::new(data.clone(), mask.clone()).unwrap();
443
444        let got_mask = getmask(&ma).unwrap();
445        let got_data = getdata(&ma).unwrap();
446
447        assert_eq!(got_mask.as_slice().unwrap(), mask.as_slice().unwrap());
448        assert_eq!(got_data.as_slice().unwrap(), data.as_slice().unwrap());
449    }
450
451    #[test]
452    fn count_masked_test() {
453        let data = Array::<f64, Ix1>::from_vec(Ix1::new([5]), vec![1.0; 5]).unwrap();
454        let mask =
455            Array::<bool, Ix1>::from_vec(Ix1::new([5]), vec![true, false, true, true, false])
456                .unwrap();
457        let ma = MaskedArray::new(data, mask).unwrap();
458        assert_eq!(count_masked(&ma, None).unwrap(), 3);
459    }
460
461    #[test]
462    fn masked_add_array_test() {
463        let data = Array::<f64, Ix1>::from_vec(Ix1::new([3]), vec![1.0, 2.0, 3.0]).unwrap();
464        let mask = Array::<bool, Ix1>::from_vec(Ix1::new([3]), vec![false, true, false]).unwrap();
465        let ma = MaskedArray::new(data, mask).unwrap();
466        let arr = Array::<f64, Ix1>::from_vec(Ix1::new([3]), vec![10.0, 20.0, 30.0]).unwrap();
467        let result = masked_add_array(&ma, &arr).unwrap();
468        let data_vals: Vec<f64> = result.data().iter().copied().collect();
469        let mask_vals: Vec<bool> = result.mask().iter().copied().collect();
470        assert_eq!(mask_vals, vec![false, true, false]);
471        assert!((data_vals[0] - 11.0).abs() < 1e-10);
472        assert!((data_vals[2] - 33.0).abs() < 1e-10);
473    }
474
475    #[test]
476    fn all_masked_mean_is_nan() {
477        let data = Array::<f64, Ix1>::from_vec(Ix1::new([3]), vec![1.0, 2.0, 3.0]).unwrap();
478        let mask = Array::<bool, Ix1>::from_vec(Ix1::new([3]), vec![true, true, true]).unwrap();
479        let ma = MaskedArray::new(data, mask).unwrap();
480        assert!(ma.mean().unwrap().is_nan());
481    }
482
483    #[test]
484    fn all_masked_min_errors() {
485        let data = Array::<f64, Ix1>::from_vec(Ix1::new([3]), vec![1.0, 2.0, 3.0]).unwrap();
486        let mask = Array::<bool, Ix1>::from_vec(Ix1::new([3]), vec![true, true, true]).unwrap();
487        let ma = MaskedArray::new(data, mask).unwrap();
488        assert!(ma.min().is_err());
489    }
490
491    #[test]
492    fn ufunc_exp_masked() {
493        let data = Array::<f64, Ix1>::from_vec(Ix1::new([3]), vec![0.0, 1.0, 2.0]).unwrap();
494        let mask = Array::<bool, Ix1>::from_vec(Ix1::new([3]), vec![false, true, false]).unwrap();
495        let ma = MaskedArray::new(data, mask).unwrap();
496        let result = ufunc_support::exp(&ma).unwrap();
497        let data_vals: Vec<f64> = result.data().iter().copied().collect();
498        let mask_vals: Vec<bool> = result.mask().iter().copied().collect();
499        assert_eq!(mask_vals, vec![false, true, false]);
500        assert!((data_vals[0] - 1.0).abs() < 1e-10); // exp(0) = 1
501        assert!((data_vals[2] - 2.0_f64.exp()).abs() < 1e-10);
502    }
503
504    #[test]
505    fn ufunc_sqrt_masked() {
506        let data = Array::<f64, Ix1>::from_vec(Ix1::new([4]), vec![4.0, 9.0, 16.0, 25.0]).unwrap();
507        let mask =
508            Array::<bool, Ix1>::from_vec(Ix1::new([4]), vec![false, true, false, true]).unwrap();
509        let ma = MaskedArray::new(data, mask).unwrap();
510        let result = ufunc_support::sqrt(&ma).unwrap();
511        let data_vals: Vec<f64> = result.data().iter().copied().collect();
512        assert!((data_vals[0] - 2.0).abs() < 1e-10);
513        assert!((data_vals[2] - 4.0).abs() < 1e-10);
514    }
515
516    #[test]
517    fn set_mask_hardened() {
518        let data = Array::<f64, Ix1>::from_vec(Ix1::new([3]), vec![1.0, 2.0, 3.0]).unwrap();
519        let mask = Array::<bool, Ix1>::from_vec(Ix1::new([3]), vec![false, true, false]).unwrap();
520        let mut ma = MaskedArray::new(data, mask).unwrap();
521        ma.harden_mask().unwrap();
522
523        // set_mask with all-false should not clear the existing true
524        let new_mask =
525            Array::<bool, Ix1>::from_vec(Ix1::new([3]), vec![false, false, false]).unwrap();
526        ma.set_mask(new_mask).unwrap();
527        let mask_vals: Vec<bool> = ma.mask().iter().copied().collect();
528        // Hard mask: union of old [false, true, false] and new [false, false, false] = [false, true, false]
529        assert_eq!(mask_vals, vec![false, true, false]);
530    }
531
532    #[test]
533    fn masked_sub_test() {
534        let d1 = Array::<f64, Ix1>::from_vec(Ix1::new([3]), vec![10.0, 20.0, 30.0]).unwrap();
535        let m1 = Array::<bool, Ix1>::from_vec(Ix1::new([3]), vec![false, false, true]).unwrap();
536        let ma1 = MaskedArray::new(d1, m1).unwrap();
537
538        let d2 = Array::<f64, Ix1>::from_vec(Ix1::new([3]), vec![1.0, 2.0, 3.0]).unwrap();
539        let m2 = Array::<bool, Ix1>::from_vec(Ix1::new([3]), vec![false, true, false]).unwrap();
540        let ma2 = MaskedArray::new(d2, m2).unwrap();
541
542        let result = masked_sub(&ma1, &ma2).unwrap();
543        let mask_vals: Vec<bool> = result.mask().iter().copied().collect();
544        assert_eq!(mask_vals, vec![false, true, true]);
545        let data_vals: Vec<f64> = result.data().iter().copied().collect();
546        assert!((data_vals[0] - 9.0).abs() < 1e-10);
547    }
548
549    #[test]
550    fn masked_mul_test() {
551        let d1 = Array::<f64, Ix1>::from_vec(Ix1::new([3]), vec![2.0, 3.0, 4.0]).unwrap();
552        let m1 = Array::<bool, Ix1>::from_vec(Ix1::new([3]), vec![false, true, false]).unwrap();
553        let ma1 = MaskedArray::new(d1, m1).unwrap();
554
555        let d2 = Array::<f64, Ix1>::from_vec(Ix1::new([3]), vec![5.0, 6.0, 7.0]).unwrap();
556        let m2 = Array::<bool, Ix1>::from_vec(Ix1::new([3]), vec![false, false, false]).unwrap();
557        let ma2 = MaskedArray::new(d2, m2).unwrap();
558
559        let result = masked_mul(&ma1, &ma2).unwrap();
560        let mask_vals: Vec<bool> = result.mask().iter().copied().collect();
561        assert_eq!(mask_vals, vec![false, true, false]);
562        let data_vals: Vec<f64> = result.data().iter().copied().collect();
563        assert!((data_vals[0] - 10.0).abs() < 1e-10);
564        assert!((data_vals[2] - 28.0).abs() < 1e-10);
565    }
566
567    #[test]
568    fn masked_div_test() {
569        let d1 = Array::<f64, Ix1>::from_vec(Ix1::new([3]), vec![10.0, 20.0, 30.0]).unwrap();
570        let m1 = Array::<bool, Ix1>::from_vec(Ix1::new([3]), vec![false, false, true]).unwrap();
571        let ma1 = MaskedArray::new(d1, m1).unwrap();
572
573        let d2 = Array::<f64, Ix1>::from_vec(Ix1::new([3]), vec![2.0, 5.0, 6.0]).unwrap();
574        let m2 = Array::<bool, Ix1>::from_vec(Ix1::new([3]), vec![false, false, false]).unwrap();
575        let ma2 = MaskedArray::new(d2, m2).unwrap();
576
577        let result = masked_div(&ma1, &ma2).unwrap();
578        let data_vals: Vec<f64> = result.data().iter().copied().collect();
579        assert!((data_vals[0] - 5.0).abs() < 1e-10);
580        assert!((data_vals[1] - 4.0).abs() < 1e-10);
581    }
582
583    #[test]
584    fn masked_invalid_negative_inf() {
585        let data =
586            Array::<f64, Ix1>::from_vec(Ix1::new([3]), vec![1.0, f64::NEG_INFINITY, 3.0]).unwrap();
587        let ma = masked_invalid(&data).unwrap();
588        let mask_vals: Vec<bool> = ma.mask().iter().copied().collect();
589        assert_eq!(mask_vals, vec![false, true, false]);
590    }
591
592    #[test]
593    fn empty_array_operations() {
594        let data = Array::<f64, Ix1>::from_vec(Ix1::new([0]), vec![]).unwrap();
595        let mask = Array::<bool, Ix1>::from_vec(Ix1::new([0]), vec![]).unwrap();
596        let ma = MaskedArray::new(data, mask).unwrap();
597        assert_eq!(ma.count().unwrap(), 0);
598        assert!(ma.mean().unwrap().is_nan());
599        let compressed = ma.compressed().unwrap();
600        assert_eq!(compressed.size(), 0);
601    }
602
603    #[test]
604    fn ndim_shape_size() {
605        let data = Array::<f64, Ix1>::from_vec(Ix1::new([5]), vec![1.0; 5]).unwrap();
606        let mask = Array::<bool, Ix1>::from_vec(Ix1::new([5]), vec![false; 5]).unwrap();
607        let ma = MaskedArray::new(data, mask).unwrap();
608        assert_eq!(ma.ndim(), 1);
609        assert_eq!(ma.shape(), &[5]);
610        assert_eq!(ma.size(), 5);
611    }
612
613    #[test]
614    fn ufunc_binary_power() {
615        let d1 = Array::<f64, Ix1>::from_vec(Ix1::new([3]), vec![2.0, 3.0, 4.0]).unwrap();
616        let m1 = Array::<bool, Ix1>::from_vec(Ix1::new([3]), vec![false, true, false]).unwrap();
617        let ma1 = MaskedArray::new(d1, m1).unwrap();
618
619        let d2 = Array::<f64, Ix1>::from_vec(Ix1::new([3]), vec![3.0, 2.0, 2.0]).unwrap();
620        let m2 = Array::<bool, Ix1>::from_vec(Ix1::new([3]), vec![false, false, false]).unwrap();
621        let ma2 = MaskedArray::new(d2, m2).unwrap();
622
623        let result = ufunc_support::power(&ma1, &ma2).unwrap();
624        let data_vals: Vec<f64> = result.data().iter().copied().collect();
625        let mask_vals: Vec<bool> = result.mask().iter().copied().collect();
626        assert_eq!(mask_vals, vec![false, true, false]);
627        assert!((data_vals[0] - 8.0).abs() < 1e-10); // 2^3 = 8
628        assert!((data_vals[2] - 16.0).abs() < 1e-10); // 4^2 = 16
629    }
630
631    #[test]
632    fn filled_with_custom_value() {
633        let data = Array::<f64, Ix1>::from_vec(Ix1::new([4]), vec![1.0, 2.0, 3.0, 4.0]).unwrap();
634        let mask =
635            Array::<bool, Ix1>::from_vec(Ix1::new([4]), vec![true, false, true, false]).unwrap();
636        let ma = MaskedArray::new(data, mask).unwrap();
637        let filled = ma.filled(-999.0).unwrap();
638        assert_eq!(filled.as_slice().unwrap(), &[-999.0, 2.0, -999.0, 4.0]);
639    }
640
641    // --- 2D masked array tests ---
642
643    #[test]
644    fn masked_2d_construction() {
645        use ferray_core::dimension::Ix2;
646        let data =
647            Array::<f64, Ix2>::from_vec(Ix2::new([2, 3]), vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0])
648                .unwrap();
649        let mask = Array::<bool, Ix2>::from_vec(
650            Ix2::new([2, 3]),
651            vec![false, true, false, false, false, true],
652        )
653        .unwrap();
654        let ma = MaskedArray::new(data, mask).unwrap();
655        assert_eq!(ma.ndim(), 2);
656        assert_eq!(ma.shape(), &[2, 3]);
657        assert_eq!(ma.size(), 6);
658        assert_eq!(ma.count().unwrap(), 4);
659    }
660
661    #[test]
662    fn masked_2d_mean() {
663        use ferray_core::dimension::Ix2;
664        let data =
665            Array::<f64, Ix2>::from_vec(Ix2::new([2, 3]), vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0])
666                .unwrap();
667        // Mask out 2.0 and 6.0
668        let mask = Array::<bool, Ix2>::from_vec(
669            Ix2::new([2, 3]),
670            vec![false, true, false, false, false, true],
671        )
672        .unwrap();
673        let ma = MaskedArray::new(data, mask).unwrap();
674        // mean of [1, 3, 4, 5] = 13/4 = 3.25
675        let m = ma.mean().unwrap();
676        assert!((m - 3.25).abs() < 1e-10);
677    }
678
679    #[test]
680    fn masked_2d_sum() {
681        use ferray_core::dimension::Ix2;
682        let data =
683            Array::<f64, Ix2>::from_vec(Ix2::new([2, 3]), vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0])
684                .unwrap();
685        let mask = Array::<bool, Ix2>::from_vec(
686            Ix2::new([2, 3]),
687            vec![false, true, false, false, false, true],
688        )
689        .unwrap();
690        let ma = MaskedArray::new(data, mask).unwrap();
691        // sum of [1, 3, 4, 5] = 13
692        assert!((ma.sum().unwrap() - 13.0).abs() < 1e-10);
693    }
694
695    #[test]
696    fn masked_2d_add_operator() {
697        use ferray_core::dimension::Ix2;
698        let d1 = Array::<f64, Ix2>::from_vec(Ix2::new([2, 2]), vec![1.0, 2.0, 3.0, 4.0]).unwrap();
699        let m1 = Array::<bool, Ix2>::from_vec(Ix2::new([2, 2]), vec![false, true, false, false])
700            .unwrap();
701        let ma1 = MaskedArray::new(d1, m1).unwrap();
702
703        let d2 =
704            Array::<f64, Ix2>::from_vec(Ix2::new([2, 2]), vec![10.0, 20.0, 30.0, 40.0]).unwrap();
705        let m2 = Array::<bool, Ix2>::from_vec(Ix2::new([2, 2]), vec![false, false, true, false])
706            .unwrap();
707        let ma2 = MaskedArray::new(d2, m2).unwrap();
708
709        let result = (&ma1 + &ma2).unwrap();
710        let mask_vals: Vec<bool> = result.mask().iter().copied().collect();
711        assert_eq!(mask_vals, vec![false, true, true, false]);
712        let data_vals: Vec<f64> = result.data().iter().copied().collect();
713        assert!((data_vals[0] - 11.0).abs() < 1e-10);
714        assert!((data_vals[3] - 44.0).abs() < 1e-10);
715    }
716
717    #[test]
718    fn masked_2d_compressed() {
719        use ferray_core::dimension::Ix2;
720        let data = Array::<f64, Ix2>::from_vec(Ix2::new([2, 2]), vec![1.0, 2.0, 3.0, 4.0]).unwrap();
721        let mask =
722            Array::<bool, Ix2>::from_vec(Ix2::new([2, 2]), vec![false, true, false, true]).unwrap();
723        let ma = MaskedArray::new(data, mask).unwrap();
724        let compressed = ma.compressed().unwrap();
725        assert_eq!(compressed.as_slice().unwrap(), &[1.0, 3.0]);
726    }
727}