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