ferray_ma/
masked_array.rs

1// ferray-ma: MaskedArray<T, D> type (REQ-1, REQ-2, REQ-3)
2
3use ferray_core::Array;
4use ferray_core::dimension::Dimension;
5use ferray_core::dtype::Element;
6use ferray_core::error::{FerrayError, FerrayResult};
7
8/// A masked array that pairs data with a boolean mask.
9///
10/// Each element position has a corresponding mask bit:
11/// - `true` means the element is **masked** (invalid / missing)
12/// - `false` means the element is valid
13///
14/// All operations (arithmetic, reductions, ufuncs) respect the mask by
15/// skipping masked elements.
16#[derive(Debug, Clone)]
17pub struct MaskedArray<T: Element, D: Dimension> {
18    /// The underlying data array.
19    data: Array<T, D>,
20    /// Boolean mask (`true` = masked/invalid).
21    mask: Array<bool, D>,
22    /// Whether the mask is hardened (cannot be cleared by assignment).
23    pub(crate) hard_mask: bool,
24}
25
26impl<T: Element, D: Dimension> MaskedArray<T, D> {
27    /// Create a new masked array from data and mask arrays.
28    ///
29    /// # Errors
30    /// Returns `FerrayError::ShapeMismatch` if data and mask shapes differ.
31    pub fn new(data: Array<T, D>, mask: Array<bool, D>) -> FerrayResult<Self> {
32        if data.shape() != mask.shape() {
33            return Err(FerrayError::shape_mismatch(format!(
34                "MaskedArray::new: data shape {:?} does not match mask shape {:?}",
35                data.shape(),
36                mask.shape()
37            )));
38        }
39        Ok(Self {
40            data,
41            mask,
42            hard_mask: false,
43        })
44    }
45
46    /// Create a masked array with no masked elements (all-false mask).
47    ///
48    /// # Errors
49    /// Returns an error if the mask array cannot be created.
50    pub fn from_data(data: Array<T, D>) -> FerrayResult<Self> {
51        let mask = Array::<bool, D>::from_elem(data.dim().clone(), false)?;
52        Ok(Self {
53            data,
54            mask,
55            hard_mask: false,
56        })
57    }
58
59    /// Return a reference to the underlying data array.
60    #[inline]
61    pub fn data(&self) -> &Array<T, D> {
62        &self.data
63    }
64
65    /// Return a reference to the mask array.
66    #[inline]
67    pub fn mask(&self) -> &Array<bool, D> {
68        &self.mask
69    }
70
71    /// Return a mutable reference to the underlying data array.
72    #[inline]
73    pub fn data_mut(&mut self) -> &mut Array<T, D> {
74        &mut self.data
75    }
76
77    /// Return the shape of the masked array.
78    #[inline]
79    pub fn shape(&self) -> &[usize] {
80        self.data.shape()
81    }
82
83    /// Return the number of dimensions.
84    #[inline]
85    pub fn ndim(&self) -> usize {
86        self.data.ndim()
87    }
88
89    /// Return the total number of elements (including masked).
90    #[inline]
91    pub fn size(&self) -> usize {
92        self.data.size()
93    }
94
95    /// Return the dimension descriptor.
96    #[inline]
97    pub fn dim(&self) -> &D {
98        self.data.dim()
99    }
100
101    /// Return whether the mask is hardened.
102    #[inline]
103    pub fn is_hard_mask(&self) -> bool {
104        self.hard_mask
105    }
106
107    /// Set a mask value at a flat index.
108    ///
109    /// If the mask is hardened, only `true` (masking) is allowed; attempts to
110    /// clear a mask bit are silently ignored.
111    ///
112    /// # Errors
113    /// Returns `FerrayError::IndexOutOfBounds` if `flat_idx >= size`.
114    pub fn set_mask_flat(&mut self, flat_idx: usize, value: bool) -> FerrayResult<()> {
115        let size = self.size();
116        if flat_idx >= size {
117            return Err(FerrayError::index_out_of_bounds(flat_idx as isize, 0, size));
118        }
119        if self.hard_mask && !value {
120            // Hard mask: cannot clear mask bits
121            return Ok(());
122        }
123        // Set via iter_mut at the flat index
124        if let Some(m) = self.mask.iter_mut().nth(flat_idx) {
125            *m = value;
126        }
127        Ok(())
128    }
129
130    /// Replace the mask with a new one.
131    ///
132    /// If the mask is hardened, only bits that are `true` in both the old and
133    /// new masks (or newly set to `true`) are allowed; cleared bits are ignored.
134    ///
135    /// # Errors
136    /// Returns `FerrayError::ShapeMismatch` if shapes differ.
137    pub fn set_mask(&mut self, new_mask: Array<bool, D>) -> FerrayResult<()> {
138        if self.mask.shape() != new_mask.shape() {
139            return Err(FerrayError::shape_mismatch(format!(
140                "set_mask: mask shape {:?} does not match array shape {:?}",
141                new_mask.shape(),
142                self.mask.shape()
143            )));
144        }
145        if self.hard_mask {
146            // Union: keep old trues, add new trues, but never clear
147            let merged: Vec<bool> = self
148                .mask
149                .iter()
150                .zip(new_mask.iter())
151                .map(|(old, new)| *old || *new)
152                .collect();
153            self.mask = Array::from_vec(self.mask.dim().clone(), merged)?;
154        } else {
155            self.mask = new_mask;
156        }
157        Ok(())
158    }
159}
ferray_ma/masked_array.rs

ferray_ma/
masked_array.rs
