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

ferray_ma/
lib.rs
