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// ## REQ status
19//
20// This file is INFRASTRUCTURE: the crate-root re-export and
21// `numpy.ma`-surface registration layer. It owns no `REQ-N` of its own —
22// each `REQ-N` is shipped by the module it re-exports below, and the
23// per-module `## REQ status` tables carry the quoted-code evidence. This
24// table records the crate-level surface status.
25//
26// | Surface | Status | Evidence |
27// |---------|--------|----------|
28// | Crate-root re-export | SHIPPED | `pub use masked_array::MaskedArray` plus the `pub use` blocks for `constructors`, `arithmetic`, `mask_ops` (`count_masked`/`getdata`/`getmask`/`is_masked`), `put::PutMode`, `interop::{MaskAware, ma_apply_unary}`, `ufunc_support::{masked_binary, masked_unary, …}`, `algorithms::{ma_where, ma_choose, …}`, and the `extras::{…}` block (this file). |
29// | `numpy.ma` callable surface | SHIPPED | The re-exported library functions back the `ferray.ma` Python module: every name is wired through a `#[pyfunction]`/`PyMaskedArray` shim in `ferray-python/src/ma.rs` and registered by `register_ma_module` in `ferray-python/src/lib.rs`. Per `.design/ferray-ma.md`, the surface covers 219/220 `numpy.ma` callables (the sole exclusion is `np.ma.test`, numpy's own pytest entry point). |
30// | Feature-gated `io` | SHIPPED (optional) | `#[cfg(feature = "io")] pub mod io` re-exports `save_masked`/`load_masked`; library-only (no binding consumer) — see `io.rs` REQ status. |
31//
32// REQ-1..REQ-39 are classified SHIPPED/NOT-STARTED in their owning modules:
33// REQ-1/2/3/5 → `masked_array.rs`; REQ-4 → `reductions.rs`; REQ-12 →
34// `ufunc_support.rs`; REQ-16 → `interop.rs`; REQ-36/37 → `put.rs`; the
35// extras family (REQ-18..REQ-33, REQ-38/39) → `extras.rs`/`algorithms.rs`.
36// No REQ owned by these routed modules is NOT-STARTED.
37
38// Masked reductions divide running sums by valid-element counts and
39// truncate `f64` results to integer index types in argmin/argmax. Float
40// equality is also intrinsic to `masked_equal` and `getdata` semantics.
41#![allow(
42    clippy::cast_possible_truncation,
43    clippy::cast_possible_wrap,
44    clippy::cast_precision_loss,
45    clippy::cast_sign_loss,
46    clippy::cast_lossless,
47    clippy::float_cmp,
48    clippy::missing_errors_doc,
49    clippy::missing_panics_doc,
50    clippy::many_single_char_names,
51    clippy::similar_names,
52    clippy::items_after_statements,
53    clippy::option_if_let_else,
54    clippy::too_long_first_doc_paragraph,
55    clippy::needless_pass_by_value,
56    clippy::match_same_arms
57)]
58
59pub mod algorithms;
60pub mod arithmetic;
61pub mod constructors;
62pub mod extras;
63pub mod filled;
64pub mod interop;
65/// Binary I/O (save/load) for `MaskedArray` via ferray-io (#509).
66///
67/// Gated behind the `io` cargo feature so callers who don't need
68/// disk I/O don't have to pull in the zip + binary reader dependency
69/// tree through ferray-io.
70#[cfg(feature = "io")]
71pub mod io;
72pub mod manipulation;
73pub mod mask_ops;
74pub mod masked_array;
75/// In-place flat assignment ops (`put` / `putmask`, #835 REQ-36/37) that
76/// mutate a `MaskedArray`'s data + mask honouring the hard-mask flag.
77pub mod put;
78pub mod reductions;
79pub mod sorting;
80pub mod ufunc_support;
81
82// Re-export the primary type at crate root
83pub use masked_array::MaskedArray;
84
85// Re-export masking constructors
86pub use constructors::{
87    fix_invalid, masked_equal, masked_greater, masked_greater_equal, masked_inside, masked_invalid,
88    masked_less, masked_less_equal, masked_not_equal, masked_outside, masked_where,
89};
90
91// Re-export arithmetic operations
92pub use arithmetic::{
93    masked_add, masked_add_array, masked_div, masked_div_array, masked_mul, masked_mul_array,
94    masked_sub, masked_sub_array,
95};
96
97// Re-export mask manipulation functions
98pub use mask_ops::{count_masked, count_masked_axis, getdata, getmask, is_masked};
99
100// Re-export the in-place flat assignment mode enum (#835 REQ-36): the
101// `MaskedArray::put`/`putmask` methods are inherent methods on the type, so
102// only the `PutMode` mode selector needs a crate-root re-export.
103pub use put::PutMode;
104
105// Re-export MaskAware trait (#505) for downstream code that wants
106// to write functions polymorphic over Array and MaskedArray.
107pub use interop::{MaskAware, ma_apply_unary};
108
109// Re-export generic ufunc helpers (#513) — the escape hatch for
110// ufuncs that don't have a dedicated named wrapper. Users with an
111// arbitrary `Fn(T) -> T` / `Fn(T, T) -> T` closure can call
112// `ferray_ma::masked_unary(ma, my_fn)` directly.
113pub use ufunc_support::{masked_binary, masked_unary};
114
115// Domain-aware ufunc wrappers (#503) — auto-mask out-of-domain
116// inputs so the result mask carries a "safe to use" contract.
117pub use ufunc_support::{
118    arccos_domain, arccosh_domain, arcsin_domain, arctanh_domain, divide_domain, log_domain,
119    log2_domain, log10_domain, masked_binary_domain, masked_unary_domain, sqrt_domain,
120};
121
122// numpy.ma specialized algorithms (#835): masked where/choose/diff/
123// ediff1d/nonzero — each mirrors its unmasked numpy counterpart with
124// explicit mask propagation. See algorithms.rs.
125pub use algorithms::{ma_choose, ma_diff, ma_ediff1d, ma_nonzero, ma_where};
126
127// numpy.ma extras: full reductions, constructors, mask manipulation,
128// linalg-lite, set ops, fill-value protocol, comparison/logical ufuncs,
129// and class helpers. See extras.rs for the catalogue.
130pub use extras::{
131    NOMASK, clump_masked, clump_unmasked, common_fill_value, default_fill_value_bool,
132    default_fill_value_f32, default_fill_value_f64, default_fill_value_i64,
133    flatnotmasked_contiguous, flatnotmasked_edges, getmaskarray, ids, is_ma, is_masked_array,
134    ma_apply_along_axis, ma_apply_over_axes, ma_compress_cols, ma_compress_rowcols,
135    ma_compress_rows, ma_concatenate, ma_corrcoef, ma_cov, ma_equal, ma_greater, ma_greater_equal,
136    ma_in1d, ma_intersect1d, ma_isin, ma_less, ma_less_equal, ma_logical_and, ma_logical_not,
137    ma_logical_or, ma_logical_xor, ma_mask_rowcols, ma_median_axis, ma_not_equal, ma_setdiff1d,
138    ma_setxor1d, ma_union1d, ma_unique, ma_unique_masked, ma_vander, make_mask, make_mask_none,
139    mask_or, masked_all, masked_all_like, masked_values, maximum_fill_value, minimum_fill_value,
140    notmasked_contiguous_axis, notmasked_edges, notmasked_edges_axis2,
141};
142
143#[cfg(test)]
144mod tests {
145    use super::*;
146    use ferray_core::Array;
147    use ferray_core::dimension::Ix1;
148
149    // -----------------------------------------------------------------------
150    // AC-1: MaskedArray::new([1,2,3,4,5], [false,false,true,false,false]).mean() == 3.0
151    // -----------------------------------------------------------------------
152    #[test]
153    fn ac1_masked_mean_skips_masked() {
154        let data =
155            Array::<f64, Ix1>::from_vec(Ix1::new([5]), vec![1.0, 2.0, 3.0, 4.0, 5.0]).unwrap();
156        let mask =
157            Array::<bool, Ix1>::from_vec(Ix1::new([5]), vec![false, false, true, false, false])
158                .unwrap();
159        let ma = MaskedArray::new(data, mask).unwrap();
160        let mean = ma.mean().unwrap();
161        // (1 + 2 + 4 + 5) / 4 = 3.0
162        assert!((mean - 3.0).abs() < 1e-10);
163    }
164
165    // -----------------------------------------------------------------------
166    // AC-2: filled(0.0) replaces masked elements with 0.0
167    // -----------------------------------------------------------------------
168    #[test]
169    fn ac2_filled_replaces_masked() {
170        let data = Array::<f64, Ix1>::from_vec(Ix1::new([4]), vec![1.0, 2.0, 3.0, 4.0]).unwrap();
171        let mask =
172            Array::<bool, Ix1>::from_vec(Ix1::new([4]), vec![false, true, false, true]).unwrap();
173        let ma = MaskedArray::new(data, mask).unwrap();
174        let filled = ma.filled(0.0).unwrap();
175        assert_eq!(filled.as_slice().unwrap(), &[1.0, 0.0, 3.0, 0.0]);
176    }
177
178    // -----------------------------------------------------------------------
179    // AC-3: compressed() returns only unmasked elements as 1D
180    // -----------------------------------------------------------------------
181    #[test]
182    fn ac3_compressed_returns_unmasked() {
183        let data =
184            Array::<f64, Ix1>::from_vec(Ix1::new([5]), vec![10.0, 20.0, 30.0, 40.0, 50.0]).unwrap();
185        let mask =
186            Array::<bool, Ix1>::from_vec(Ix1::new([5]), vec![false, true, false, true, false])
187                .unwrap();
188        let ma = MaskedArray::new(data, mask).unwrap();
189        let compressed = ma.compressed().unwrap();
190        assert_eq!(compressed.as_slice().unwrap(), &[10.0, 30.0, 50.0]);
191    }
192
193    // -----------------------------------------------------------------------
194    // AC-4: masked_invalid masks NaN and Inf
195    // -----------------------------------------------------------------------
196    #[test]
197    fn ac4_masked_invalid_nan_inf() {
198        let data =
199            Array::<f64, Ix1>::from_vec(Ix1::new([4]), vec![1.0, f64::NAN, 3.0, f64::INFINITY])
200                .unwrap();
201        let ma = masked_invalid(&data).unwrap();
202        let mask_vals: Vec<bool> = ma.mask().iter().copied().collect();
203        assert_eq!(mask_vals, vec![false, true, false, true]);
204    }
205
206    // -----------------------------------------------------------------------
207    // AC-5: ma1 + ma2 produces correct mask union and correct values
208    // -----------------------------------------------------------------------
209    #[test]
210    fn ac5_add_mask_union() {
211        let d1 = Array::<f64, Ix1>::from_vec(Ix1::new([4]), vec![1.0, 2.0, 3.0, 4.0]).unwrap();
212        let m1 =
213            Array::<bool, Ix1>::from_vec(Ix1::new([4]), vec![false, true, false, false]).unwrap();
214        let ma1 = MaskedArray::new(d1, m1).unwrap();
215
216        let d2 = Array::<f64, Ix1>::from_vec(Ix1::new([4]), vec![10.0, 20.0, 30.0, 40.0]).unwrap();
217        let m2 =
218            Array::<bool, Ix1>::from_vec(Ix1::new([4]), vec![false, false, true, false]).unwrap();
219        let ma2 = MaskedArray::new(d2, m2).unwrap();
220
221        let result = masked_add(&ma1, &ma2).unwrap();
222        let mask_vals: Vec<bool> = result.mask().iter().copied().collect();
223        // Mask union: [false, true, true, false]
224        assert_eq!(mask_vals, vec![false, true, true, false]);
225        // Unmasked values: 1+10=11, 4+40=44; masked get 0.0
226        let data_vals: Vec<f64> = result.data().iter().copied().collect();
227        assert!((data_vals[0] - 11.0).abs() < 1e-10);
228        assert!((data_vals[3] - 44.0).abs() < 1e-10);
229    }
230
231    #[test]
232    fn operator_add_matches_masked_add() {
233        let d1 = Array::<f64, Ix1>::from_vec(Ix1::new([3]), vec![1.0, 2.0, 3.0]).unwrap();
234        let m1 = Array::<bool, Ix1>::from_vec(Ix1::new([3]), vec![false, true, false]).unwrap();
235        let ma1 = MaskedArray::new(d1, m1).unwrap();
236
237        let d2 = Array::<f64, Ix1>::from_vec(Ix1::new([3]), vec![10.0, 20.0, 30.0]).unwrap();
238        let m2 = Array::<bool, Ix1>::from_vec(Ix1::new([3]), vec![false, false, true]).unwrap();
239        let ma2 = MaskedArray::new(d2, m2).unwrap();
240
241        // Use operator syntax
242        let result = (&ma1 + &ma2).unwrap();
243        let mask_vals: Vec<bool> = result.mask().iter().copied().collect();
244        assert_eq!(mask_vals, vec![false, true, true]);
245        let data_vals: Vec<f64> = result.data().iter().copied().collect();
246        assert!((data_vals[0] - 11.0).abs() < 1e-10);
247    }
248
249    // -----------------------------------------------------------------------
250    // AC-7: sin(masked_array) returns same mask, correct values
251    // -----------------------------------------------------------------------
252    #[test]
253    fn ac7_ufunc_sin_masked() {
254        use std::f64::consts::FRAC_PI_2;
255        let data =
256            Array::<f64, Ix1>::from_vec(Ix1::new([3]), vec![0.0, FRAC_PI_2, FRAC_PI_2]).unwrap();
257        let mask = Array::<bool, Ix1>::from_vec(Ix1::new([3]), vec![false, true, false]).unwrap();
258        let ma = MaskedArray::new(data, mask).unwrap();
259        let result = ufunc_support::sin(&ma).unwrap();
260        let mask_vals: Vec<bool> = result.mask().iter().copied().collect();
261        assert_eq!(mask_vals, vec![false, true, false]);
262        let data_vals: Vec<f64> = result.data().iter().copied().collect();
263        // sin(0) = 0, masked = 0.0 (skipped), sin(pi/2) = 1.0
264        assert!((data_vals[0] - 0.0).abs() < 1e-10);
265        assert!((data_vals[2] - 1.0).abs() < 1e-10);
266    }
267
268    // -----------------------------------------------------------------------
269    // AC-8: sort places masked at end; harden_mask prevents clearing
270    // -----------------------------------------------------------------------
271    #[test]
272    fn ac8_sort_masked_at_end() {
273        let data =
274            Array::<f64, Ix1>::from_vec(Ix1::new([5]), vec![5.0, 1.0, 3.0, 2.0, 4.0]).unwrap();
275        let mask =
276            Array::<bool, Ix1>::from_vec(Ix1::new([5]), vec![false, false, true, false, false])
277                .unwrap();
278        let ma = MaskedArray::new(data, mask).unwrap();
279        let sorted = ma.sort().unwrap();
280        let data_vals: Vec<f64> = sorted.data().iter().copied().collect();
281        let mask_vals: Vec<bool> = sorted.mask().iter().copied().collect();
282        // Unmasked [5, 1, 2, 4] sorted = [1, 2, 4, 5], then masked [3]
283        assert_eq!(data_vals, vec![1.0, 2.0, 4.0, 5.0, 3.0]);
284        assert_eq!(mask_vals, vec![false, false, false, false, true]);
285    }
286
287    #[test]
288    fn ac8_harden_mask_prevents_clearing() {
289        let data = Array::<f64, Ix1>::from_vec(Ix1::new([3]), vec![1.0, 2.0, 3.0]).unwrap();
290        let mask = Array::<bool, Ix1>::from_vec(Ix1::new([3]), vec![false, true, false]).unwrap();
291        let mut ma = MaskedArray::new(data, mask).unwrap();
292
293        ma.harden_mask().unwrap();
294        assert!(ma.is_hard_mask());
295
296        // Try to clear the mask at index 1 — should be silently ignored
297        ma.set_mask_flat(1, false).unwrap();
298        let mask_vals: Vec<bool> = ma.mask().iter().copied().collect();
299        assert_eq!(mask_vals, vec![false, true, false]);
300
301        // Setting a mask bit to true should still work
302        ma.set_mask_flat(0, true).unwrap();
303        let mask_vals2: Vec<bool> = ma.mask().iter().copied().collect();
304        assert_eq!(mask_vals2, vec![true, true, false]);
305
306        // Soften and then clearing should work
307        ma.soften_mask().unwrap();
308        assert!(!ma.is_hard_mask());
309        ma.set_mask_flat(1, false).unwrap();
310        let mask_vals3: Vec<bool> = ma.mask().iter().copied().collect();
311        assert_eq!(mask_vals3, vec![true, false, false]);
312    }
313
314    // -----------------------------------------------------------------------
315    // AC-9: is_masked returns true/false correctly
316    // -----------------------------------------------------------------------
317    #[test]
318    fn ac9_is_masked() {
319        let data1 = Array::<f64, Ix1>::from_vec(Ix1::new([3]), vec![1.0, 2.0, 3.0]).unwrap();
320        let mask1 = Array::<bool, Ix1>::from_vec(Ix1::new([3]), vec![false, true, false]).unwrap();
321        let ma1 = MaskedArray::new(data1, mask1).unwrap();
322        assert!(is_masked(&ma1).unwrap());
323
324        let data2 = Array::<f64, Ix1>::from_vec(Ix1::new([3]), vec![1.0, 2.0, 3.0]).unwrap();
325        let mask2 = Array::<bool, Ix1>::from_vec(Ix1::new([3]), vec![false, false, false]).unwrap();
326        let ma2 = MaskedArray::new(data2, mask2).unwrap();
327        assert!(!is_masked(&ma2).unwrap());
328    }
329
330    // -----------------------------------------------------------------------
331    // Additional tests
332    // -----------------------------------------------------------------------
333
334    #[test]
335    fn shape_mismatch_error() {
336        let data = Array::<f64, Ix1>::from_vec(Ix1::new([3]), vec![1.0, 2.0, 3.0]).unwrap();
337        let mask = Array::<bool, Ix1>::from_vec(Ix1::new([2]), vec![false, true]).unwrap();
338        assert!(MaskedArray::new(data, mask).is_err());
339    }
340
341    #[test]
342    fn from_data_no_mask() {
343        let data = Array::<f64, Ix1>::from_vec(Ix1::new([3]), vec![1.0, 2.0, 3.0]).unwrap();
344        let ma = MaskedArray::from_data(data).unwrap();
345        assert!(!is_masked(&ma).unwrap());
346        assert_eq!(ma.count().unwrap(), 3);
347    }
348
349    #[test]
350    fn sum_skips_masked() {
351        let data = Array::<f64, Ix1>::from_vec(Ix1::new([4]), vec![1.0, 2.0, 3.0, 4.0]).unwrap();
352        let mask =
353            Array::<bool, Ix1>::from_vec(Ix1::new([4]), vec![false, true, false, true]).unwrap();
354        let ma = MaskedArray::new(data, mask).unwrap();
355        assert!((ma.sum().unwrap() - 4.0).abs() < 1e-10);
356    }
357
358    #[test]
359    fn min_max_skip_masked() {
360        let data =
361            Array::<f64, Ix1>::from_vec(Ix1::new([5]), vec![5.0, 1.0, 3.0, 2.0, 4.0]).unwrap();
362        let mask =
363            Array::<bool, Ix1>::from_vec(Ix1::new([5]), vec![false, true, false, false, false])
364                .unwrap();
365        let ma = MaskedArray::new(data, mask).unwrap();
366        assert!((ma.min().unwrap() - 2.0).abs() < 1e-10);
367        assert!((ma.max().unwrap() - 5.0).abs() < 1e-10);
368    }
369
370    #[test]
371    fn var_std_skip_masked() {
372        // values: [2, 4, 6] (mask out index 1 and 4)
373        let data =
374            Array::<f64, Ix1>::from_vec(Ix1::new([5]), vec![2.0, 99.0, 4.0, 6.0, 99.0]).unwrap();
375        let mask =
376            Array::<bool, Ix1>::from_vec(Ix1::new([5]), vec![false, true, false, false, true])
377                .unwrap();
378        let ma = MaskedArray::new(data, mask).unwrap();
379        let mean = ma.mean().unwrap();
380        assert!((mean - 4.0).abs() < 1e-10);
381        // var = ((2-4)^2 + (4-4)^2 + (6-4)^2) / 3 = 8/3
382        let v = ma.var().unwrap();
383        assert!((v - 8.0 / 3.0).abs() < 1e-10);
384        let s = ma.std().unwrap();
385        assert!((s - (8.0_f64 / 3.0).sqrt()).abs() < 1e-10);
386    }
387
388    #[test]
389    fn count_elements() {
390        let data = Array::<f64, Ix1>::from_vec(Ix1::new([5]), vec![1.0; 5]).unwrap();
391        let mask =
392            Array::<bool, Ix1>::from_vec(Ix1::new([5]), vec![false, true, true, false, false])
393                .unwrap();
394        let ma = MaskedArray::new(data, mask).unwrap();
395        assert_eq!(ma.count().unwrap(), 3);
396    }
397
398    #[test]
399    fn masked_equal_test() {
400        let data =
401            Array::<f64, Ix1>::from_vec(Ix1::new([5]), vec![1.0, 2.0, 3.0, 2.0, 1.0]).unwrap();
402        let ma = masked_equal(&data, 2.0).unwrap();
403        let mask_vals: Vec<bool> = ma.mask().iter().copied().collect();
404        assert_eq!(mask_vals, vec![false, true, false, true, false]);
405    }
406
407    #[test]
408    fn masked_greater_test() {
409        let data = Array::<f64, Ix1>::from_vec(Ix1::new([4]), vec![1.0, 2.0, 3.0, 4.0]).unwrap();
410        let ma = masked_greater(&data, 2.0).unwrap();
411        let mask_vals: Vec<bool> = ma.mask().iter().copied().collect();
412        assert_eq!(mask_vals, vec![false, false, true, true]);
413    }
414
415    #[test]
416    fn masked_less_test() {
417        let data = Array::<f64, Ix1>::from_vec(Ix1::new([4]), vec![1.0, 2.0, 3.0, 4.0]).unwrap();
418        let ma = masked_less(&data, 3.0).unwrap();
419        let mask_vals: Vec<bool> = ma.mask().iter().copied().collect();
420        assert_eq!(mask_vals, vec![true, true, false, false]);
421    }
422
423    #[test]
424    fn masked_not_equal_test() {
425        let data = Array::<f64, Ix1>::from_vec(Ix1::new([3]), vec![1.0, 2.0, 3.0]).unwrap();
426        let ma = masked_not_equal(&data, 2.0).unwrap();
427        let mask_vals: Vec<bool> = ma.mask().iter().copied().collect();
428        assert_eq!(mask_vals, vec![true, false, true]);
429    }
430
431    #[test]
432    fn masked_greater_equal_test() {
433        let data = Array::<f64, Ix1>::from_vec(Ix1::new([4]), vec![1.0, 2.0, 3.0, 4.0]).unwrap();
434        let ma = masked_greater_equal(&data, 3.0).unwrap();
435        let mask_vals: Vec<bool> = ma.mask().iter().copied().collect();
436        assert_eq!(mask_vals, vec![false, false, true, true]);
437    }
438
439    #[test]
440    fn masked_less_equal_test() {
441        let data = Array::<f64, Ix1>::from_vec(Ix1::new([4]), vec![1.0, 2.0, 3.0, 4.0]).unwrap();
442        let ma = masked_less_equal(&data, 2.0).unwrap();
443        let mask_vals: Vec<bool> = ma.mask().iter().copied().collect();
444        assert_eq!(mask_vals, vec![true, true, false, false]);
445    }
446
447    #[test]
448    fn masked_inside_test() {
449        let data =
450            Array::<f64, Ix1>::from_vec(Ix1::new([5]), vec![1.0, 2.0, 3.0, 4.0, 5.0]).unwrap();
451        let ma = masked_inside(&data, 2.0, 4.0).unwrap();
452        let mask_vals: Vec<bool> = ma.mask().iter().copied().collect();
453        assert_eq!(mask_vals, vec![false, true, true, true, false]);
454    }
455
456    #[test]
457    fn masked_outside_test() {
458        let data =
459            Array::<f64, Ix1>::from_vec(Ix1::new([5]), vec![1.0, 2.0, 3.0, 4.0, 5.0]).unwrap();
460        let ma = masked_outside(&data, 2.0, 4.0).unwrap();
461        let mask_vals: Vec<bool> = ma.mask().iter().copied().collect();
462        assert_eq!(mask_vals, vec![true, false, false, false, true]);
463    }
464
465    #[test]
466    fn masked_where_test() {
467        let data = Array::<f64, Ix1>::from_vec(Ix1::new([4]), vec![1.0, 2.0, 3.0, 4.0]).unwrap();
468        let cond =
469            Array::<bool, Ix1>::from_vec(Ix1::new([4]), vec![true, false, true, false]).unwrap();
470        let ma = masked_where(&cond, &data).unwrap();
471        let mask_vals: Vec<bool> = ma.mask().iter().copied().collect();
472        assert_eq!(mask_vals, vec![true, false, true, false]);
473    }
474
475    #[test]
476    fn argsort_test() {
477        let data =
478            Array::<f64, Ix1>::from_vec(Ix1::new([5]), vec![5.0, 1.0, 3.0, 2.0, 4.0]).unwrap();
479        let mask =
480            Array::<bool, Ix1>::from_vec(Ix1::new([5]), vec![false, false, true, false, false])
481                .unwrap();
482        let ma = MaskedArray::new(data, mask).unwrap();
483        let indices = ma.argsort().unwrap();
484        // Unmasked: index 1 (1.0), 3 (2.0), 4 (4.0), 0 (5.0); masked: 2
485        assert_eq!(indices.shape(), &[5]);
486        assert_eq!(indices.as_slice().unwrap(), &[1u64, 3, 4, 0, 2]);
487    }
488
489    #[test]
490    fn getmask_getdata_test() {
491        let data = Array::<f64, Ix1>::from_vec(Ix1::new([3]), vec![1.0, 2.0, 3.0]).unwrap();
492        let mask = Array::<bool, Ix1>::from_vec(Ix1::new([3]), vec![false, true, false]).unwrap();
493        let ma = MaskedArray::new(data.clone(), mask.clone()).unwrap();
494
495        let got_mask = getmask(&ma).unwrap();
496        let got_data = getdata(&ma).unwrap();
497
498        assert_eq!(got_mask.as_slice().unwrap(), mask.as_slice().unwrap());
499        assert_eq!(got_data.as_slice().unwrap(), data.as_slice().unwrap());
500    }
501
502    #[test]
503    fn count_masked_test() {
504        let data = Array::<f64, Ix1>::from_vec(Ix1::new([5]), vec![1.0; 5]).unwrap();
505        let mask =
506            Array::<bool, Ix1>::from_vec(Ix1::new([5]), vec![true, false, true, true, false])
507                .unwrap();
508        let ma = MaskedArray::new(data, mask).unwrap();
509        assert_eq!(count_masked(&ma).unwrap(), 3);
510    }
511
512    #[test]
513    fn count_masked_axis_2d_along_rows() {
514        // #268: per-row masked counts on a 2x3 array.
515        // Mask:
516        //   [[T, F, T],
517        //    [F, F, T]]
518        use ferray_core::dimension::Ix2;
519        let data = Array::<f64, Ix2>::from_vec(Ix2::new([2, 3]), vec![1.0; 6]).unwrap();
520        let mask = Array::<bool, Ix2>::from_vec(
521            Ix2::new([2, 3]),
522            vec![true, false, true, false, false, true],
523        )
524        .unwrap();
525        let ma = MaskedArray::new(data, mask).unwrap();
526        // axis=1 reduces columns -> per-row counts: [2, 1].
527        let counts = count_masked_axis(&ma, 1).unwrap();
528        assert_eq!(counts.shape(), &[2]);
529        assert_eq!(counts.as_slice().unwrap(), &[2u64, 1]);
530    }
531
532    #[test]
533    fn count_masked_axis_2d_along_cols() {
534        // axis=0 reduces rows -> per-column counts: [1, 0, 2].
535        use ferray_core::dimension::Ix2;
536        let data = Array::<f64, Ix2>::from_vec(Ix2::new([2, 3]), vec![1.0; 6]).unwrap();
537        let mask = Array::<bool, Ix2>::from_vec(
538            Ix2::new([2, 3]),
539            vec![true, false, true, false, false, true],
540        )
541        .unwrap();
542        let ma = MaskedArray::new(data, mask).unwrap();
543        let counts = count_masked_axis(&ma, 0).unwrap();
544        assert_eq!(counts.shape(), &[3]);
545        assert_eq!(counts.as_slice().unwrap(), &[1u64, 0, 2]);
546    }
547
548    #[test]
549    fn count_masked_axis_rejects_out_of_bounds() {
550        use ferray_core::dimension::Ix2;
551        let data = Array::<f64, Ix2>::from_vec(Ix2::new([2, 3]), vec![1.0; 6]).unwrap();
552        let mask = Array::<bool, Ix2>::from_vec(Ix2::new([2, 3]), vec![false; 6]).unwrap();
553        let ma = MaskedArray::new(data, mask).unwrap();
554        assert!(count_masked_axis(&ma, 2).is_err());
555    }
556
557    #[test]
558    fn sort_axis_2d_per_row() {
559        // #271: sort along axis=1 (columns) should sort each row
560        // independently. Row 0: [3, 1, _] (mask 2) → unmasked sorted
561        // ascending [1, 3] then masked → [1, 3, _].
562        // Row 1: [_, 4, 2] (mask 0) → [2, 4, _].
563        use ferray_core::dimension::Ix2;
564        let data =
565            Array::<f64, Ix2>::from_vec(Ix2::new([2, 3]), vec![3.0, 1.0, 99.0, 99.0, 4.0, 2.0])
566                .unwrap();
567        let mask = Array::<bool, Ix2>::from_vec(
568            Ix2::new([2, 3]),
569            vec![false, false, true, true, false, false],
570        )
571        .unwrap();
572        let ma = MaskedArray::new(data, mask).unwrap();
573        let sorted = ma.sort_axis(1).unwrap();
574        assert_eq!(sorted.shape(), &[2, 3]);
575        let d: Vec<f64> = sorted.data().iter().copied().collect();
576        let m: Vec<bool> = sorted.mask().iter().copied().collect();
577        // Row 0
578        assert!((d[0] - 1.0).abs() < 1e-12);
579        assert!((d[1] - 3.0).abs() < 1e-12);
580        assert!(m[2], "row 0 col 2 should be masked");
581        // Row 1
582        assert!((d[3] - 2.0).abs() < 1e-12);
583        assert!((d[4] - 4.0).abs() < 1e-12);
584        assert!(m[5], "row 1 col 2 should be masked");
585    }
586
587    #[test]
588    fn sort_axis_2d_per_column() {
589        // axis=0 sorts each column. Column 0: [3, 1] both unmasked →
590        // [1, 3]. Column 1: [2, _] (row 1 masked) → [2, _].
591        use ferray_core::dimension::Ix2;
592        let data =
593            Array::<f64, Ix2>::from_vec(Ix2::new([2, 2]), vec![3.0, 2.0, 1.0, 99.0]).unwrap();
594        let mask = Array::<bool, Ix2>::from_vec(Ix2::new([2, 2]), vec![false, false, false, true])
595            .unwrap();
596        let ma = MaskedArray::new(data, mask).unwrap();
597        let sorted = ma.sort_axis(0).unwrap();
598        let d: Vec<f64> = sorted.data().iter().copied().collect();
599        let m: Vec<bool> = sorted.mask().iter().copied().collect();
600        // Column 0: index [0,0]=1, [1,0]=3.
601        assert!((d[0] - 1.0).abs() < 1e-12);
602        assert!((d[2] - 3.0).abs() < 1e-12);
603        // Column 1: index [0,1]=2 (unmasked), [1,1]=99 (masked).
604        assert!((d[1] - 2.0).abs() < 1e-12);
605        assert!(m[3]);
606    }
607
608    #[test]
609    fn sort_axis_rejects_out_of_bounds() {
610        use ferray_core::dimension::Ix2;
611        let data = Array::<f64, Ix2>::from_vec(Ix2::new([2, 3]), vec![1.0; 6]).unwrap();
612        let mask = Array::<bool, Ix2>::from_vec(Ix2::new([2, 3]), vec![false; 6]).unwrap();
613        let ma = MaskedArray::new(data, mask).unwrap();
614        assert!(ma.sort_axis(2).is_err());
615    }
616
617    #[test]
618    fn data_mut_only_exposes_element_slice() {
619        // #273: data_mut returns &mut [T], not &mut Array — callers can
620        // update values but cannot reshape or resize.
621        let data = Array::<f64, Ix1>::from_vec(Ix1::new([4]), vec![1.0, 2.0, 3.0, 4.0]).unwrap();
622        let mask = Array::<bool, Ix1>::from_vec(Ix1::new([4]), vec![false; 4]).unwrap();
623        let mut ma = MaskedArray::new(data, mask).unwrap();
624        // Mutate values in place.
625        if let Some(s) = ma.data_mut() {
626            s[2] = 99.0;
627        }
628        assert_eq!(ma.shape(), &[4]);
629        let vals: Vec<f64> = ma.data().iter().copied().collect();
630        assert_eq!(vals, vec![1.0, 2.0, 99.0, 4.0]);
631        // Mask stays the same length and value.
632        assert_eq!(ma.mask().shape(), &[4]);
633    }
634
635    #[test]
636    fn masked_add_array_test() {
637        let data = Array::<f64, Ix1>::from_vec(Ix1::new([3]), vec![1.0, 2.0, 3.0]).unwrap();
638        let mask = Array::<bool, Ix1>::from_vec(Ix1::new([3]), vec![false, true, false]).unwrap();
639        let ma = MaskedArray::new(data, mask).unwrap();
640        let arr = Array::<f64, Ix1>::from_vec(Ix1::new([3]), vec![10.0, 20.0, 30.0]).unwrap();
641        let result = masked_add_array(&ma, &arr).unwrap();
642        let data_vals: Vec<f64> = result.data().iter().copied().collect();
643        let mask_vals: Vec<bool> = result.mask().iter().copied().collect();
644        assert_eq!(mask_vals, vec![false, true, false]);
645        assert!((data_vals[0] - 11.0).abs() < 1e-10);
646        assert!((data_vals[2] - 33.0).abs() < 1e-10);
647    }
648
649    #[test]
650    fn all_masked_mean_is_nan() {
651        let data = Array::<f64, Ix1>::from_vec(Ix1::new([3]), vec![1.0, 2.0, 3.0]).unwrap();
652        let mask = Array::<bool, Ix1>::from_vec(Ix1::new([3]), vec![true, true, true]).unwrap();
653        let ma = MaskedArray::new(data, mask).unwrap();
654        assert!(ma.mean().unwrap().is_nan());
655    }
656
657    #[test]
658    fn all_masked_min_errors() {
659        let data = Array::<f64, Ix1>::from_vec(Ix1::new([3]), vec![1.0, 2.0, 3.0]).unwrap();
660        let mask = Array::<bool, Ix1>::from_vec(Ix1::new([3]), vec![true, true, true]).unwrap();
661        let ma = MaskedArray::new(data, mask).unwrap();
662        assert!(ma.min().is_err());
663    }
664
665    #[test]
666    fn ufunc_exp_masked() {
667        let data = Array::<f64, Ix1>::from_vec(Ix1::new([3]), vec![0.0, 1.0, 2.0]).unwrap();
668        let mask = Array::<bool, Ix1>::from_vec(Ix1::new([3]), vec![false, true, false]).unwrap();
669        let ma = MaskedArray::new(data, mask).unwrap();
670        let result = ufunc_support::exp(&ma).unwrap();
671        let data_vals: Vec<f64> = result.data().iter().copied().collect();
672        let mask_vals: Vec<bool> = result.mask().iter().copied().collect();
673        assert_eq!(mask_vals, vec![false, true, false]);
674        assert!((data_vals[0] - 1.0).abs() < 1e-10); // exp(0) = 1
675        assert!((data_vals[2] - 2.0_f64.exp()).abs() < 1e-10);
676    }
677
678    #[test]
679    fn ufunc_sqrt_masked() {
680        let data = Array::<f64, Ix1>::from_vec(Ix1::new([4]), vec![4.0, 9.0, 16.0, 25.0]).unwrap();
681        let mask =
682            Array::<bool, Ix1>::from_vec(Ix1::new([4]), vec![false, true, false, true]).unwrap();
683        let ma = MaskedArray::new(data, mask).unwrap();
684        let result = ufunc_support::sqrt(&ma).unwrap();
685        let data_vals: Vec<f64> = result.data().iter().copied().collect();
686        assert!((data_vals[0] - 2.0).abs() < 1e-10);
687        assert!((data_vals[2] - 4.0).abs() < 1e-10);
688    }
689
690    #[test]
691    fn set_mask_hardened() {
692        let data = Array::<f64, Ix1>::from_vec(Ix1::new([3]), vec![1.0, 2.0, 3.0]).unwrap();
693        let mask = Array::<bool, Ix1>::from_vec(Ix1::new([3]), vec![false, true, false]).unwrap();
694        let mut ma = MaskedArray::new(data, mask).unwrap();
695        ma.harden_mask().unwrap();
696
697        // set_mask with all-false should not clear the existing true
698        let new_mask =
699            Array::<bool, Ix1>::from_vec(Ix1::new([3]), vec![false, false, false]).unwrap();
700        ma.set_mask(new_mask).unwrap();
701        let mask_vals: Vec<bool> = ma.mask().iter().copied().collect();
702        // Hard mask: union of old [false, true, false] and new [false, false, false] = [false, true, false]
703        assert_eq!(mask_vals, vec![false, true, false]);
704    }
705
706    #[test]
707    fn masked_sub_test() {
708        let d1 = Array::<f64, Ix1>::from_vec(Ix1::new([3]), vec![10.0, 20.0, 30.0]).unwrap();
709        let m1 = Array::<bool, Ix1>::from_vec(Ix1::new([3]), vec![false, false, true]).unwrap();
710        let ma1 = MaskedArray::new(d1, m1).unwrap();
711
712        let d2 = Array::<f64, Ix1>::from_vec(Ix1::new([3]), vec![1.0, 2.0, 3.0]).unwrap();
713        let m2 = Array::<bool, Ix1>::from_vec(Ix1::new([3]), vec![false, true, false]).unwrap();
714        let ma2 = MaskedArray::new(d2, m2).unwrap();
715
716        let result = masked_sub(&ma1, &ma2).unwrap();
717        let mask_vals: Vec<bool> = result.mask().iter().copied().collect();
718        assert_eq!(mask_vals, vec![false, true, true]);
719        let data_vals: Vec<f64> = result.data().iter().copied().collect();
720        assert!((data_vals[0] - 9.0).abs() < 1e-10);
721    }
722
723    #[test]
724    fn masked_mul_test() {
725        let d1 = Array::<f64, Ix1>::from_vec(Ix1::new([3]), vec![2.0, 3.0, 4.0]).unwrap();
726        let m1 = Array::<bool, Ix1>::from_vec(Ix1::new([3]), vec![false, true, false]).unwrap();
727        let ma1 = MaskedArray::new(d1, m1).unwrap();
728
729        let d2 = Array::<f64, Ix1>::from_vec(Ix1::new([3]), vec![5.0, 6.0, 7.0]).unwrap();
730        let m2 = Array::<bool, Ix1>::from_vec(Ix1::new([3]), vec![false, false, false]).unwrap();
731        let ma2 = MaskedArray::new(d2, m2).unwrap();
732
733        let result = masked_mul(&ma1, &ma2).unwrap();
734        let mask_vals: Vec<bool> = result.mask().iter().copied().collect();
735        assert_eq!(mask_vals, vec![false, true, false]);
736        let data_vals: Vec<f64> = result.data().iter().copied().collect();
737        assert!((data_vals[0] - 10.0).abs() < 1e-10);
738        assert!((data_vals[2] - 28.0).abs() < 1e-10);
739    }
740
741    #[test]
742    fn masked_div_test() {
743        let d1 = Array::<f64, Ix1>::from_vec(Ix1::new([3]), vec![10.0, 20.0, 30.0]).unwrap();
744        let m1 = Array::<bool, Ix1>::from_vec(Ix1::new([3]), vec![false, false, true]).unwrap();
745        let ma1 = MaskedArray::new(d1, m1).unwrap();
746
747        let d2 = Array::<f64, Ix1>::from_vec(Ix1::new([3]), vec![2.0, 5.0, 6.0]).unwrap();
748        let m2 = Array::<bool, Ix1>::from_vec(Ix1::new([3]), vec![false, false, false]).unwrap();
749        let ma2 = MaskedArray::new(d2, m2).unwrap();
750
751        let result = masked_div(&ma1, &ma2).unwrap();
752        let data_vals: Vec<f64> = result.data().iter().copied().collect();
753        assert!((data_vals[0] - 5.0).abs() < 1e-10);
754        assert!((data_vals[1] - 4.0).abs() < 1e-10);
755    }
756
757    #[test]
758    fn masked_invalid_negative_inf() {
759        let data =
760            Array::<f64, Ix1>::from_vec(Ix1::new([3]), vec![1.0, f64::NEG_INFINITY, 3.0]).unwrap();
761        let ma = masked_invalid(&data).unwrap();
762        let mask_vals: Vec<bool> = ma.mask().iter().copied().collect();
763        assert_eq!(mask_vals, vec![false, true, false]);
764    }
765
766    #[test]
767    fn empty_array_operations() {
768        let data = Array::<f64, Ix1>::from_vec(Ix1::new([0]), vec![]).unwrap();
769        let mask = Array::<bool, Ix1>::from_vec(Ix1::new([0]), vec![]).unwrap();
770        let ma = MaskedArray::new(data, mask).unwrap();
771        assert_eq!(ma.count().unwrap(), 0);
772        assert!(ma.mean().unwrap().is_nan());
773        let compressed = ma.compressed().unwrap();
774        assert_eq!(compressed.size(), 0);
775    }
776
777    #[test]
778    fn ndim_shape_size() {
779        let data = Array::<f64, Ix1>::from_vec(Ix1::new([5]), vec![1.0; 5]).unwrap();
780        let mask = Array::<bool, Ix1>::from_vec(Ix1::new([5]), vec![false; 5]).unwrap();
781        let ma = MaskedArray::new(data, mask).unwrap();
782        assert_eq!(ma.ndim(), 1);
783        assert_eq!(ma.shape(), &[5]);
784        assert_eq!(ma.size(), 5);
785    }
786
787    #[test]
788    fn ufunc_binary_power() {
789        let d1 = Array::<f64, Ix1>::from_vec(Ix1::new([3]), vec![2.0, 3.0, 4.0]).unwrap();
790        let m1 = Array::<bool, Ix1>::from_vec(Ix1::new([3]), vec![false, true, false]).unwrap();
791        let ma1 = MaskedArray::new(d1, m1).unwrap();
792
793        let d2 = Array::<f64, Ix1>::from_vec(Ix1::new([3]), vec![3.0, 2.0, 2.0]).unwrap();
794        let m2 = Array::<bool, Ix1>::from_vec(Ix1::new([3]), vec![false, false, false]).unwrap();
795        let ma2 = MaskedArray::new(d2, m2).unwrap();
796
797        let result = ufunc_support::power(&ma1, &ma2).unwrap();
798        let data_vals: Vec<f64> = result.data().iter().copied().collect();
799        let mask_vals: Vec<bool> = result.mask().iter().copied().collect();
800        assert_eq!(mask_vals, vec![false, true, false]);
801        assert!((data_vals[0] - 8.0).abs() < 1e-10); // 2^3 = 8
802        assert!((data_vals[2] - 16.0).abs() < 1e-10); // 4^2 = 16
803    }
804
805    #[test]
806    fn filled_with_custom_value() {
807        let data = Array::<f64, Ix1>::from_vec(Ix1::new([4]), vec![1.0, 2.0, 3.0, 4.0]).unwrap();
808        let mask =
809            Array::<bool, Ix1>::from_vec(Ix1::new([4]), vec![true, false, true, false]).unwrap();
810        let ma = MaskedArray::new(data, mask).unwrap();
811        let filled = ma.filled(-999.0).unwrap();
812        assert_eq!(filled.as_slice().unwrap(), &[-999.0, 2.0, -999.0, 4.0]);
813    }
814
815    // --- 2D masked array tests ---
816
817    #[test]
818    fn masked_2d_construction() {
819        use ferray_core::dimension::Ix2;
820        let data =
821            Array::<f64, Ix2>::from_vec(Ix2::new([2, 3]), vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0])
822                .unwrap();
823        let mask = Array::<bool, Ix2>::from_vec(
824            Ix2::new([2, 3]),
825            vec![false, true, false, false, false, true],
826        )
827        .unwrap();
828        let ma = MaskedArray::new(data, mask).unwrap();
829        assert_eq!(ma.ndim(), 2);
830        assert_eq!(ma.shape(), &[2, 3]);
831        assert_eq!(ma.size(), 6);
832        assert_eq!(ma.count().unwrap(), 4);
833    }
834
835    #[test]
836    fn masked_2d_mean() {
837        use ferray_core::dimension::Ix2;
838        let data =
839            Array::<f64, Ix2>::from_vec(Ix2::new([2, 3]), vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0])
840                .unwrap();
841        // Mask out 2.0 and 6.0
842        let mask = Array::<bool, Ix2>::from_vec(
843            Ix2::new([2, 3]),
844            vec![false, true, false, false, false, true],
845        )
846        .unwrap();
847        let ma = MaskedArray::new(data, mask).unwrap();
848        // mean of [1, 3, 4, 5] = 13/4 = 3.25
849        let m = ma.mean().unwrap();
850        assert!((m - 3.25).abs() < 1e-10);
851    }
852
853    #[test]
854    fn masked_2d_sum() {
855        use ferray_core::dimension::Ix2;
856        let data =
857            Array::<f64, Ix2>::from_vec(Ix2::new([2, 3]), vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0])
858                .unwrap();
859        let mask = Array::<bool, Ix2>::from_vec(
860            Ix2::new([2, 3]),
861            vec![false, true, false, false, false, true],
862        )
863        .unwrap();
864        let ma = MaskedArray::new(data, mask).unwrap();
865        // sum of [1, 3, 4, 5] = 13
866        assert!((ma.sum().unwrap() - 13.0).abs() < 1e-10);
867    }
868
869    #[test]
870    fn masked_2d_add_operator() {
871        use ferray_core::dimension::Ix2;
872        let d1 = Array::<f64, Ix2>::from_vec(Ix2::new([2, 2]), vec![1.0, 2.0, 3.0, 4.0]).unwrap();
873        let m1 = Array::<bool, Ix2>::from_vec(Ix2::new([2, 2]), vec![false, true, false, false])
874            .unwrap();
875        let ma1 = MaskedArray::new(d1, m1).unwrap();
876
877        let d2 =
878            Array::<f64, Ix2>::from_vec(Ix2::new([2, 2]), vec![10.0, 20.0, 30.0, 40.0]).unwrap();
879        let m2 = Array::<bool, Ix2>::from_vec(Ix2::new([2, 2]), vec![false, false, true, false])
880            .unwrap();
881        let ma2 = MaskedArray::new(d2, m2).unwrap();
882
883        let result = (&ma1 + &ma2).unwrap();
884        let mask_vals: Vec<bool> = result.mask().iter().copied().collect();
885        assert_eq!(mask_vals, vec![false, true, true, false]);
886        let data_vals: Vec<f64> = result.data().iter().copied().collect();
887        assert!((data_vals[0] - 11.0).abs() < 1e-10);
888        assert!((data_vals[3] - 44.0).abs() < 1e-10);
889    }
890
891    #[test]
892    fn masked_2d_compressed() {
893        use ferray_core::dimension::Ix2;
894        let data = Array::<f64, Ix2>::from_vec(Ix2::new([2, 2]), vec![1.0, 2.0, 3.0, 4.0]).unwrap();
895        let mask =
896            Array::<bool, Ix2>::from_vec(Ix2::new([2, 2]), vec![false, true, false, true]).unwrap();
897        let ma = MaskedArray::new(data, mask).unwrap();
898        let compressed = ma.compressed().unwrap();
899        assert_eq!(compressed.as_slice().unwrap(), &[1.0, 3.0]);
900    }
901}
ferray_ma/lib.rs

ferray_ma/
lib.rs
