// ferray-ma: Interop with ferray-core's Array<T, D>
//
// In NumPy, MaskedArray subclasses ndarray, so anywhere ndarray is accepted
// MaskedArray works too — with mask propagation handled automatically by
// __array_ufunc__. Rust has no inheritance, so we provide the interop in
// three layers:
//
// 1. `AsRef<Array<T, D>>` for low-level access — pass `&masked.as_ref()`
// where `&Array<T, D>` is expected, accepting that the mask is dropped.
// Useful when you need the data quickly and the mask is irrelevant.
//
// 2. `From<MaskedArray<T, D>>` for `Array<T, D>` — owned data extraction
// that consumes the mask. Use when you want to convert and forget.
//
// 3. `apply_unary` / `apply_binary` — mask-aware adapters that take any
// Array→Array function (e.g. `ferray_ufunc::sin`) and propagate the
// mask through to the result. The function operates on every element
// of the underlying data array (including masked positions, which
// hold whatever value is there); the mask is then re-attached to the
// result so that downstream consumers see the correct invalidity.
//
// See: https://github.com/dollspace-gay/ferray/issues/505
//
// ## REQ status
//
// REQ-16 SHIPPED — audited, green. This module provides the
// `MaskedArray <-> Array` conversion layer and the mask-aware adapters that
// back `numpy.ma.getdata`/`getmask`/`getmaskarray`.
//
// | REQ | Status | Evidence |
// |-----|--------|----------|
// | REQ-16 (`getmask`/`getdata`) | SHIPPED | The `AsRef<Array<T, D>>` impl + `From<MaskedArray<T, D>> for Array<T, D>` + `into_data` (this file) extract the underlying data; `getdata`/`getmask` live in `mask_ops.rs` and `getmaskarray` in `extras.rs`, but the data-view machinery they read is defined here. Non-test production consumers: `fma::getdata` is called throughout `ferray-python/src/ma.rs` (data egress for repr/median/cov/etc.), and `fma::getmaskarray` is the mask-extraction primitive for the binary/comparison shims there. |
//
// Also SHIPPED here: the `MaskAware<T, D>` trait (uniform mask access over
// `Array` and `MaskedArray`) and the `apply_unary`/`apply_binary`/
// `apply_unary_to`/`ma_apply_unary` adapters — `ma_apply_unary` (re-exported
// from `ferray-ma/src/lib.rs`) is the mask-propagating bridge consumed by the
// `numpy.ma` unary shims; the default fill from REQ-5 fills masked slots here.
use ferray_core::Array;
use ferray_core::dimension::Dimension;
use ferray_core::dtype::Element;
use ferray_core::error::{FerrayError, FerrayResult};
use crate::MaskedArray;
// ---------------------------------------------------------------------------
// AsRef / From — passive interop
// ---------------------------------------------------------------------------
impl<T: Element, D: Dimension> AsRef<Array<T, D>> for MaskedArray<T, D> {
/// Borrow the underlying data array, dropping the mask.
///
/// This lets you pass a `&MaskedArray<T, D>` to any function that takes
/// `&Array<T, D>`, but the mask is **not** consulted — masked positions
/// will be processed like normal data. Use [`MaskedArray::apply_unary`]
/// or [`MaskedArray::apply_binary`] when you want mask propagation.
///
/// # Example
/// ```
/// # use ferray_core::{Array, dimension::Ix1};
/// # use ferray_ma::MaskedArray;
/// # let data = Array::<f64, Ix1>::from_vec(Ix1::new([3]), vec![1.0, 2.0, 3.0]).unwrap();
/// # let mask = Array::<bool, Ix1>::from_vec(Ix1::new([3]), vec![false, true, false]).unwrap();
/// let ma = MaskedArray::new(data, mask).unwrap();
/// // Pass through a function that operates on `&Array<T, D>`:
/// let arr_ref: &Array<f64, Ix1> = ma.as_ref();
/// assert_eq!(arr_ref.shape(), &[3]);
/// ```
fn as_ref(&self) -> &Array<T, D> {
self.data()
}
}
impl<T: Element + Copy, D: Dimension> From<MaskedArray<T, D>> for Array<T, D> {
/// Consume a `MaskedArray` and return its underlying data array,
/// **discarding** the mask. Equivalent to calling `ma.into_data()`.
///
/// Requires `T: Copy` because the underlying data buffer is cloned;
/// use [`MaskedArray::filled`] / [`MaskedArray::filled_default`] for a
/// mask-aware materialization with custom fill semantics.
fn from(ma: MaskedArray<T, D>) -> Self {
ma.into_data()
}
}
// ---------------------------------------------------------------------------
// Active interop — apply functions with mask propagation
// ---------------------------------------------------------------------------
impl<T, D> MaskedArray<T, D>
where
T: Element + Copy,
D: Dimension,
{
/// Consume the masked array and return its underlying data array,
/// dropping the mask.
///
/// Use [`MaskedArray::filled_default`] (or [`MaskedArray::filled`]) if
/// you want masked positions replaced by a sentinel value before
/// dropping the mask.
pub fn into_data(self) -> Array<T, D> {
// We can't move out of a struct with non-Copy fields directly
// because of `data_mut()` borrowing semantics, so destructure
// the unsafe-but-safe internal `data` field via the public getter.
// The clone here is unavoidable without a fully private accessor.
self.data().clone()
}
/// Apply a unary function to the underlying data and re-attach the
/// mask, propagating it to the result.
///
/// The function `f` is called on the **entire** data array — masked
/// positions are processed alongside unmasked ones, but their values
/// in the result are immediately overwritten with the masked array's
/// `fill_value`. This matches `NumPy`'s `__array_ufunc__` semantics where
/// ufuncs run over the raw data and the mask is propagated separately.
///
/// # Example
/// ```ignore
/// // Apply ferray-ufunc::sin to a masked f64 array:
/// let result = ma.apply_unary(|arr| ferray_ufunc::sin(arr))?;
/// ```
///
/// # Errors
/// Forwards any error from `f`.
pub fn apply_unary<F>(&self, f: F) -> FerrayResult<Self>
where
F: FnOnce(&Array<T, D>) -> FerrayResult<Array<T, D>>,
{
let data_out = f(self.data())?;
if data_out.shape() != self.shape() {
return Err(FerrayError::shape_mismatch(format!(
"apply_unary: function changed shape from {:?} to {:?}",
self.shape(),
data_out.shape()
)));
}
let fill = self.fill_value();
// Replace masked positions in the result with fill_value to keep
// operations like log/sqrt from leaving misleading data behind.
let masked_data: Vec<T> = data_out
.iter()
.zip(self.mask().iter())
.map(|(v, m)| if *m { fill } else { *v })
.collect();
let final_data = Array::from_vec(self.dim().clone(), masked_data)?;
let mut result = Self::new(final_data, self.mask().clone())?;
result.set_fill_value(fill);
Ok(result)
}
/// Apply a unary function that maps `T -> U`, propagating the mask.
///
/// This is the type-changing variant of [`MaskedArray::apply_unary`],
/// useful for predicates like `isnan` that return `Array<bool, D>` from
/// `Array<T, D>`. Masked positions in the result hold the explicitly
/// supplied `default_for_masked` value.
///
/// # Errors
/// Forwards any error from `f`. Returns `FerrayError::ShapeMismatch` if
/// `f` produces an array with a different shape.
pub fn apply_unary_to<U, F>(
&self,
f: F,
default_for_masked: U,
) -> FerrayResult<MaskedArray<U, D>>
where
U: Element + Copy,
F: FnOnce(&Array<T, D>) -> FerrayResult<Array<U, D>>,
{
let data_out = f(self.data())?;
if data_out.shape() != self.shape() {
return Err(FerrayError::shape_mismatch(format!(
"apply_unary_to: function changed shape from {:?} to {:?}",
self.shape(),
data_out.shape()
)));
}
let masked_data: Vec<U> = data_out
.iter()
.zip(self.mask().iter())
.map(|(v, m)| if *m { default_for_masked } else { *v })
.collect();
let final_data = Array::from_vec(self.dim().clone(), masked_data)?;
let mut result = MaskedArray::new(final_data, self.mask().clone())?;
result.set_fill_value(default_for_masked);
Ok(result)
}
/// Apply a binary function to two masked arrays, propagating the mask
/// union. Both inputs must have the same shape.
///
/// The function `f` is called on the underlying data of both inputs
/// (no broadcasting — use [`crate::masked_add`]-style functions for
/// that). The result mask is the OR of the two input masks; masked
/// positions in the data are overwritten with the receiver's `fill_value`.
///
/// # Example
/// ```ignore
/// let result = a.apply_binary(&b, |x, y| ferray_ufunc::power(x, y))?;
/// ```
///
/// # Errors
/// Returns `FerrayError::ShapeMismatch` if shapes differ. Forwards any
/// error from `f`.
pub fn apply_binary<F>(&self, other: &Self, f: F) -> FerrayResult<Self>
where
F: FnOnce(&Array<T, D>, &Array<T, D>) -> FerrayResult<Array<T, D>>,
{
if self.shape() != other.shape() {
return Err(FerrayError::shape_mismatch(format!(
"apply_binary: shapes {:?} and {:?} differ",
self.shape(),
other.shape()
)));
}
let data_out = f(self.data(), other.data())?;
if data_out.shape() != self.shape() {
return Err(FerrayError::shape_mismatch(format!(
"apply_binary: function changed shape from {:?} to {:?}",
self.shape(),
data_out.shape()
)));
}
// Mask union.
let union_data: Vec<bool> = self
.mask()
.iter()
.zip(other.mask().iter())
.map(|(a, b)| *a || *b)
.collect();
let union_mask = Array::from_vec(self.dim().clone(), union_data)?;
let fill = self.fill_value();
let masked_data: Vec<T> = data_out
.iter()
.zip(union_mask.iter())
.map(|(v, m)| if *m { fill } else { *v })
.collect();
let final_data = Array::from_vec(self.dim().clone(), masked_data)?;
let mut result = Self::new(final_data, union_mask)?;
result.set_fill_value(fill);
Ok(result)
}
}
#[cfg(test)]
mod tests {
use super::*;
use ferray_core::dimension::Ix1;
fn ma1(data: Vec<f64>, mask: Vec<bool>) -> MaskedArray<f64, Ix1> {
let n = data.len();
let d = Array::<f64, Ix1>::from_vec(Ix1::new([n]), data).unwrap();
let m = Array::<bool, Ix1>::from_vec(Ix1::new([n]), mask).unwrap();
MaskedArray::new(d, m).unwrap()
}
#[test]
fn as_ref_returns_underlying_data() {
let ma = ma1(vec![1.0, 2.0, 3.0], vec![false, true, false]);
let arr_ref: &Array<f64, Ix1> = ma.as_ref();
assert_eq!(arr_ref.shape(), &[3]);
// The data is the unmasked-equivalent, mask is dropped.
let v: Vec<f64> = arr_ref.iter().copied().collect();
assert_eq!(v, vec![1.0, 2.0, 3.0]);
}
#[test]
fn from_masked_to_array_drops_mask() {
let ma = ma1(vec![1.0, 2.0, 3.0], vec![false, true, false]);
let arr: Array<f64, Ix1> = ma.into();
assert_eq!(arr.shape(), &[3]);
}
#[test]
fn into_data_method() {
let ma = ma1(vec![1.0, 2.0, 3.0], vec![false, true, false]);
let arr = ma.into_data();
assert_eq!(arr.shape(), &[3]);
}
#[test]
fn apply_unary_propagates_mask() {
let ma =
ma1(vec![1.0, 4.0, 9.0, 16.0], vec![false, false, true, false]).with_fill_value(-1.0);
let result = ma
.apply_unary(|arr| {
// Squaring closure as a stand-in for any ferray-ufunc function
let data: Vec<f64> = arr.iter().map(|&x| x.sqrt()).collect();
Array::<f64, Ix1>::from_vec(Ix1::new([arr.size()]), data)
})
.unwrap();
// Masked position (index 2) holds fill_value; others hold sqrt.
let d: Vec<f64> = result.data().iter().copied().collect();
assert_eq!(d, vec![1.0, 2.0, -1.0, 4.0]);
// Mask is preserved.
let m: Vec<bool> = result.mask().iter().copied().collect();
assert_eq!(m, vec![false, false, true, false]);
// Fill value is preserved.
assert_eq!(result.fill_value(), -1.0);
}
#[test]
fn apply_unary_forwards_error() {
let ma = ma1(vec![1.0, 2.0], vec![false, false]);
let result: FerrayResult<MaskedArray<f64, Ix1>> =
ma.apply_unary(|_| Err(FerrayError::invalid_value("simulated failure")));
assert!(result.is_err());
}
#[test]
fn apply_unary_rejects_shape_change() {
let ma = ma1(vec![1.0, 2.0, 3.0], vec![false, false, false]);
let result = ma.apply_unary(|_| Array::<f64, Ix1>::from_vec(Ix1::new([2]), vec![1.0, 2.0]));
assert!(result.is_err());
}
#[test]
fn apply_unary_to_changes_type_with_mask_default() {
// Apply a "is positive" predicate that returns bool.
let ma = ma1(vec![1.0, -2.0, 3.0, -4.0], vec![false, false, true, false]);
let result = ma
.apply_unary_to(
|arr| {
let data: Vec<bool> = arr.iter().map(|&x| x > 0.0).collect();
Array::<bool, Ix1>::from_vec(Ix1::new([arr.size()]), data)
},
false, // default for masked positions
)
.unwrap();
let d: Vec<bool> = result.data().iter().copied().collect();
// Index 2 is masked → false (the default); others reflect the predicate.
assert_eq!(d, vec![true, false, false, false]);
let m: Vec<bool> = result.mask().iter().copied().collect();
assert_eq!(m, vec![false, false, true, false]);
}
#[test]
fn apply_binary_unions_masks() {
let a = ma1(vec![10.0, 20.0, 30.0], vec![false, true, false]).with_fill_value(-1.0);
let b = ma1(vec![1.0, 2.0, 3.0], vec![false, false, true]);
let result = a
.apply_binary(&b, |x, y| {
let data: Vec<f64> = x.iter().zip(y.iter()).map(|(&a, &b)| a + b).collect();
Array::<f64, Ix1>::from_vec(Ix1::new([x.size()]), data)
})
.unwrap();
let d: Vec<f64> = result.data().iter().copied().collect();
// Indices 1 and 2 are masked (union of the two input masks); index 0
// gets the actual sum.
assert_eq!(d, vec![11.0, -1.0, -1.0]);
let m: Vec<bool> = result.mask().iter().copied().collect();
assert_eq!(m, vec![false, true, true]);
assert_eq!(result.fill_value(), -1.0);
}
#[test]
fn apply_binary_rejects_shape_mismatch() {
let a = ma1(vec![1.0, 2.0, 3.0], vec![false; 3]);
let b = ma1(vec![1.0, 2.0], vec![false; 2]);
let result = a.apply_binary(&b, |x, _y| Ok(x.clone()));
assert!(result.is_err());
}
/// Demonstrates the canonical interop pattern with a real ferray-stats
/// call: pass `&MaskedArray` through `as_ref()` to a function that
/// expects `&Array`. This loses the mask (per the `AsRef` contract) but
/// is the cheapest way to bridge.
#[test]
fn as_ref_works_with_array_consuming_function() {
// A simple Array -> Array function (any closure works as a stand-in).
fn double(arr: &Array<f64, Ix1>) -> FerrayResult<Array<f64, Ix1>> {
let data: Vec<f64> = arr.iter().map(|&x| x * 2.0).collect();
Array::<f64, Ix1>::from_vec(Ix1::new([arr.size()]), data)
}
let ma = ma1(vec![1.0, 2.0, 3.0], vec![false, true, false]);
// Direct call via AsRef — the mask is dropped:
let result = double(ma.as_ref()).unwrap();
let v: Vec<f64> = result.iter().copied().collect();
assert_eq!(v, vec![2.0, 4.0, 6.0]);
// For mask-preserving usage, route through apply_unary instead:
let masked_result = ma.apply_unary(double).unwrap();
let m: Vec<bool> = masked_result.mask().iter().copied().collect();
assert_eq!(m, vec![false, true, false]);
}
}
// ---------------------------------------------------------------------------
// MaskAware trait: common interface for Array and MaskedArray (#505)
//
// NumPy's MaskedArray subclasses ndarray so any ndarray-accepting function
// automatically accepts MaskedArray (with __array_ufunc__ handling mask
// propagation). Rust doesn't have inheritance, so the ferray equivalent is a
// trait both types implement: mask-aware functions take
// `&impl MaskAware<T, D>` and dispatch via the trait methods.
//
// Array<T, D> is treated as "always fully unmasked" — `mask_opt()` returns
// None and `fill_value()` falls back to `T::zero()`. MaskedArray<T, D>
// delegates to its actual accessors. This lets callers write one function
// that works on both and still propagates masks correctly.
// ---------------------------------------------------------------------------
/// Shared view contract for functions that want to accept either an
/// `Array` or a `MaskedArray` (#505).
///
/// Implementations:
/// - `Array<T, D>`: `data()` returns `self`, `mask_opt()` returns `None`
/// (no mask), `fill_value()` returns `T::zero()`. The array is treated
/// as fully unmasked.
/// - `MaskedArray<T, D>`: delegates to the existing accessors.
///
/// Downstream code that wants to write "one function, works on both"
/// should take `&impl MaskAware<T, D>` and consult `mask_opt()` to
/// decide whether to do mask propagation.
pub trait MaskAware<T: Element, D: Dimension> {
/// Return a reference to the underlying data array.
fn data(&self) -> &Array<T, D>;
/// Return the mask array if one is explicitly present, or `None`
/// when the input carries no mask (treated as fully unmasked).
///
/// For `Array<T, D>` this always returns `None`. For
/// `MaskedArray<T, D>` it returns `Some` when a real mask has
/// been explicitly set and `None` when the array is in the
/// nomask-sentinel state (#506).
fn mask_opt(&self) -> Option<&Array<bool, D>>;
/// Return the fill value to use for masked positions in
/// derived results.
fn fill_value(&self) -> T
where
T: Copy;
/// Return the shape of the underlying data.
fn shape(&self) -> &[usize] {
self.data().shape()
}
}
impl<T: Element, D: Dimension> MaskAware<T, D> for Array<T, D> {
#[inline]
fn data(&self) -> &Self {
self
}
/// A plain `Array<T, D>` has no mask — always returns `None`.
#[inline]
fn mask_opt(&self) -> Option<&Array<bool, D>> {
None
}
/// A plain `Array<T, D>` has no fill value; returns `T::zero()`.
#[inline]
fn fill_value(&self) -> T
where
T: Copy,
{
T::zero()
}
}
impl<T: Element, D: Dimension> MaskAware<T, D> for MaskedArray<T, D> {
#[inline]
fn data(&self) -> &Array<T, D> {
Self::data(self)
}
#[inline]
fn mask_opt(&self) -> Option<&Array<bool, D>> {
Self::mask_opt(self)
}
#[inline]
fn fill_value(&self) -> T
where
T: Copy,
{
Self::fill_value(self)
}
}
/// Apply a unary function to any `MaskAware` input, propagating the
/// mask if one is present.
///
/// When the input is a plain `Array<T, D>` (or a nomask-sentinel
/// `MaskedArray`), the function is applied directly and the result
/// is returned as a nomask `MaskedArray`. When the input has a real
/// mask, this delegates to the existing [`MaskedArray::apply_unary`]
/// path so masked positions are overwritten with the fill value.
///
/// Use this to write "one function, works on both" adapters:
///
/// ```ignore
/// fn my_op<X: MaskAware<f64, Ix1>>(x: &X) -> FerrayResult<MaskedArray<f64, Ix1>> {
/// ma_apply_unary(x, |a| ferray_ufunc::sin(a))
/// }
/// ```
///
/// # Errors
/// Forwards any error from `f`, plus shape-mismatch errors if `f`
/// returns a differently-shaped array.
pub fn ma_apply_unary<T, D, X, F>(input: &X, f: F) -> FerrayResult<MaskedArray<T, D>>
where
T: Element + Copy,
D: Dimension,
X: MaskAware<T, D>,
F: FnOnce(&Array<T, D>) -> FerrayResult<Array<T, D>>,
{
let data_out = f(input.data())?;
if data_out.shape() != input.shape() {
return Err(FerrayError::shape_mismatch(format!(
"ma_apply_unary: function changed shape from {:?} to {:?}",
input.shape(),
data_out.shape()
)));
}
match input.mask_opt() {
None => {
// No mask — wrap the result in a nomask-sentinel
// MaskedArray so the caller gets a uniform return type.
let mut out = MaskedArray::from_data(data_out)?;
out.set_fill_value(input.fill_value());
Ok(out)
}
Some(mask) => {
// Overwrite masked positions with fill_value so downstream
// operations can't see stale data at masked slots.
let fill = input.fill_value();
let masked_data: Vec<T> = data_out
.iter()
.zip(mask.iter())
.map(|(v, m)| if *m { fill } else { *v })
.collect();
let final_data = Array::from_vec(input.data().dim().clone(), masked_data)?;
let mut result = MaskedArray::new(final_data, mask.clone())?;
result.set_fill_value(fill);
Ok(result)
}
}
}
#[cfg(test)]
mod mask_aware_tests {
use super::*;
use ferray_core::dimension::Ix1;
fn arr_f64(data: Vec<f64>) -> Array<f64, Ix1> {
let n = data.len();
Array::<f64, Ix1>::from_vec(Ix1::new([n]), data).unwrap()
}
fn ma_f64(data: Vec<f64>, mask: Vec<bool>) -> MaskedArray<f64, Ix1> {
let d = arr_f64(data);
let n = d.size();
let m = Array::<bool, Ix1>::from_vec(Ix1::new([n]), mask).unwrap();
MaskedArray::new(d, m).unwrap()
}
// ---- MaskAware trait impls (#505) ----
#[test]
fn array_implements_mask_aware_with_none_mask() {
let a = arr_f64(vec![1.0, 2.0, 3.0]);
// Plain Array carries no mask.
assert!(<Array<f64, Ix1> as MaskAware<f64, Ix1>>::mask_opt(&a).is_none());
assert_eq!(
<Array<f64, Ix1> as MaskAware<f64, Ix1>>::fill_value(&a),
0.0
);
assert_eq!(<Array<f64, Ix1> as MaskAware<f64, Ix1>>::shape(&a), &[3]);
}
#[test]
fn masked_array_implements_mask_aware_with_real_mask() {
let ma = ma_f64(vec![1.0, 2.0, 3.0], vec![false, true, false]);
let via_trait = <MaskedArray<f64, Ix1> as MaskAware<f64, Ix1>>::mask_opt(&ma);
assert!(via_trait.is_some());
assert_eq!(
via_trait.unwrap().iter().copied().collect::<Vec<_>>(),
vec![false, true, false]
);
}
#[test]
fn nomask_sentinel_masked_array_reports_none_via_trait() {
// A from_data-constructed MaskedArray (nomask sentinel) should
// report None through the MaskAware trait, matching the
// behavior of a plain Array.
let ma = MaskedArray::from_data(arr_f64(vec![1.0, 2.0, 3.0])).unwrap();
let via_trait = <MaskedArray<f64, Ix1> as MaskAware<f64, Ix1>>::mask_opt(&ma);
assert!(via_trait.is_none());
}
#[test]
fn ma_apply_unary_on_plain_array_returns_nomask_result() {
let a = arr_f64(vec![1.0, 2.0, 3.0]);
let result = ma_apply_unary(&a, |x| {
let data: Vec<f64> = x.iter().map(|v| v * 2.0).collect();
Array::from_vec(x.dim().clone(), data)
})
.unwrap();
assert_eq!(
result.data().iter().copied().collect::<Vec<_>>(),
vec![2.0, 4.0, 6.0]
);
// Plain-Array input → nomask-sentinel result.
assert!(!result.has_real_mask());
}
#[test]
fn ma_apply_unary_on_masked_array_propagates_mask() {
let ma = ma_f64(vec![1.0, 2.0, 3.0], vec![false, true, false]);
let result = ma_apply_unary(&ma, |x| {
let data: Vec<f64> = x.iter().map(|v| v * 2.0).collect();
Array::from_vec(x.dim().clone(), data)
})
.unwrap();
// Mask survives the operation.
assert!(result.has_real_mask());
assert_eq!(
result.mask().iter().copied().collect::<Vec<_>>(),
vec![false, true, false]
);
// Masked position was overwritten with fill_value (numpy default
// 1e20, default_filler['f'] core.py:166).
let d: Vec<f64> = result.data().iter().copied().collect();
assert_eq!(d[0], 2.0);
assert_eq!(d[1], 1e20); // masked → numpy default fill value
assert_eq!(d[2], 6.0);
}
#[test]
fn ma_apply_unary_generic_over_both_types() {
// Write a helper that works on both Array and MaskedArray.
fn double_it<T, D, X>(x: &X) -> FerrayResult<MaskedArray<T, D>>
where
T: Element + Copy + std::ops::Mul<Output = T> + num_traits::FromPrimitive,
D: Dimension,
X: MaskAware<T, D>,
{
let two = T::from_f64(2.0).unwrap();
ma_apply_unary(x, move |a| {
let data: Vec<T> = a.iter().map(|v| *v * two).collect();
Array::from_vec(a.dim().clone(), data)
})
}
let plain = arr_f64(vec![1.0, 2.0, 3.0]);
let masked = ma_f64(vec![1.0, 2.0, 3.0], vec![false, true, false]);
// Both inputs go through the same helper.
let r_plain = double_it(&plain).unwrap();
let r_masked = double_it(&masked).unwrap();
// Plain result: no mask, all values doubled.
assert!(!r_plain.has_real_mask());
assert_eq!(
r_plain.data().iter().copied().collect::<Vec<_>>(),
vec![2.0, 4.0, 6.0]
);
// Masked result: mask preserved, masked position holds fill.
assert!(r_masked.has_real_mask());
assert_eq!(
r_masked.mask().iter().copied().collect::<Vec<_>>(),
vec![false, true, false]
);
}
#[test]
fn ma_apply_unary_rejects_shape_changing_function() {
let a = arr_f64(vec![1.0, 2.0, 3.0]);
let result = ma_apply_unary(&a, |_| {
// Return a wrong-shape result deliberately.
Ok(arr_f64(vec![1.0, 2.0]))
});
assert!(result.is_err());
}
}