concision_core/nn/mask/
mask.rs

1/*
2    Appellation: mask <module>
3    Contrib: FL03 <jo3mccain@icloud.com>
4*/
5use nd::iter::{Iter, IterMut};
6use nd::prelude::*;
7use nd::{Data, DataMut, OwnedRepr, RawData, RawDataClone};
8
9pub struct Mask<S = OwnedRepr<bool>, D = Ix2>(ArrayBase<S, D>)
10where
11    D: Dimension,
12    S: RawData<Elem = bool>;
13
14impl<S, D> Mask<S, D>
15where
16    D: Dimension,
17    S: RawData<Elem = bool>,
18{
19    pub fn from_arr(data: ArrayBase<S, D>) -> Self {
20        Self(data)
21    }
22
23    pub fn apply<A, T, F>(&mut self, data: &ArrayBase<T, D>, fill: A) -> ArrayBase<T, D>
24    where
25        A: Clone,
26        S: Data,
27        T: DataMut<Elem = A> + RawDataClone,
28    {
29        let mut res = data.clone();
30        res.zip_mut_with(self.as_mut(), |x, &m| {
31            if m {
32                *x = fill.clone();
33            }
34        });
35        res
36    }
37
38    pub fn mask_inplace<'a, A, T, F>(
39        &mut self,
40        data: &'a mut ArrayBase<T, D>,
41        fill: A,
42    ) -> &'a mut ArrayBase<T, D>
43    where
44        A: Clone,
45        S: Data,
46        T: DataMut<Elem = A>,
47    {
48        data.zip_mut_with(&mut self.0, |x, &m| {
49            if m {
50                *x = fill.clone();
51            }
52        });
53        data
54    }
55
56    pub fn as_slice(&self) -> &[bool]
57    where
58        S: Data,
59    {
60        self.get().as_slice().unwrap()
61    }
62
63    pub fn as_mut_slice(&mut self) -> &mut [bool]
64    where
65        S: DataMut,
66    {
67        self.get_mut().as_slice_mut().unwrap()
68    }
69
70    pub fn dim(&self) -> D::Pattern {
71        self.get().dim()
72    }
73
74    pub fn iter(&self) -> Iter<'_, bool, D>
75    where
76        S: Data,
77    {
78        self.get().iter()
79    }
80
81    pub fn iter_mut(&mut self) -> IterMut<'_, bool, D>
82    where
83        S: DataMut,
84    {
85        self.get_mut().iter_mut()
86    }
87
88    pub fn get(&self) -> &ArrayBase<S, D> {
89        &self.0
90    }
91
92    pub fn get_mut(&mut self) -> &mut ArrayBase<S, D> {
93        &mut self.0
94    }
95
96    pub fn into_inner(self) -> ArrayBase<S, D> {
97        self.0
98    }
99
100    pub fn ndim(&self) -> usize {
101        self.get().ndim()
102    }
103
104    pub fn raw_dim(&self) -> D {
105        self.get().raw_dim()
106    }
107
108    pub fn set(&mut self, data: ArrayBase<S, D>) {
109        self.0 = data;
110    }
111
112    pub fn shape(&self) -> D {
113        self.get().raw_dim()
114    }
115}
116
117/*
118 ************* Implementations *************
119*/
120mod impls {
121    use super::Mask;
122    use core::borrow::{Borrow, BorrowMut};
123    use core::ops::{Deref, DerefMut, Index, IndexMut};
124    use nd::{ArrayBase, Data, DataMut, Dimension, NdIndex, RawData};
125
126    impl<S, D> AsRef<ArrayBase<S, D>> for Mask<S, D>
127    where
128        D: Dimension,
129        S: RawData<Elem = bool>,
130    {
131        fn as_ref(&self) -> &ArrayBase<S, D> {
132            &self.0
133        }
134    }
135
136    impl<S, D> AsMut<ArrayBase<S, D>> for Mask<S, D>
137    where
138        D: Dimension,
139        S: RawData<Elem = bool>,
140    {
141        fn as_mut(&mut self) -> &mut ArrayBase<S, D> {
142            &mut self.0
143        }
144    }
145
146    impl<S, D> Borrow<ArrayBase<S, D>> for Mask<S, D>
147    where
148        D: Dimension,
149        S: RawData<Elem = bool>,
150    {
151        fn borrow(&self) -> &ArrayBase<S, D> {
152            &self.0
153        }
154    }
155
156    impl<S, D> BorrowMut<ArrayBase<S, D>> for Mask<S, D>
157    where
158        D: Dimension,
159        S: RawData<Elem = bool>,
160    {
161        fn borrow_mut(&mut self) -> &mut ArrayBase<S, D> {
162            &mut self.0
163        }
164    }
165
166    impl<S, D> Deref for Mask<S, D>
167    where
168        D: Dimension,
169        S: RawData<Elem = bool>,
170    {
171        type Target = ArrayBase<S, D>;
172
173        fn deref(&self) -> &Self::Target {
174            &self.0
175        }
176    }
177
178    impl<S, D> DerefMut for Mask<S, D>
179    where
180        D: Dimension,
181        S: RawData<Elem = bool>,
182    {
183        fn deref_mut(&mut self) -> &mut Self::Target {
184            &mut self.0
185        }
186    }
187
188    impl<S, D, I> Index<I> for Mask<S, D>
189    where
190        D: Dimension,
191        I: NdIndex<D>,
192        S: Data<Elem = bool>,
193    {
194        type Output = <ArrayBase<S, D> as Index<I>>::Output;
195
196        fn index(&self, index: I) -> &Self::Output {
197            &self.0[index]
198        }
199    }
200
201    impl<S, D, I> IndexMut<I> for Mask<S, D>
202    where
203        D: Dimension,
204        I: NdIndex<D>,
205        S: DataMut<Elem = bool>,
206    {
207        fn index_mut(&mut self, index: I) -> &mut Self::Output {
208            &mut self.0[index]
209        }
210    }
211}
212
213mod impl_from {
214    use super::Mask;
215    use nd::{ArrayBase, Dimension, RawData};
216
217    impl<S, D> From<ArrayBase<S, D>> for Mask<S, D>
218    where
219        D: Dimension,
220        S: RawData<Elem = bool>,
221    {
222        fn from(mask: ArrayBase<S, D>) -> Self {
223            Mask(mask)
224        }
225    }
226
227    impl<S, D> From<Mask<S, D>> for ArrayBase<S, D>
228    where
229        D: Dimension,
230        S: RawData<Elem = bool>,
231    {
232        fn from(mask: Mask<S, D>) -> Self {
233            mask.0
234        }
235    }
236}