ndarray/iterators/
windows.rs

1use std::marker::PhantomData;
2
3use super::Baseiter;
4use crate::imp_prelude::*;
5use crate::IntoDimension;
6use crate::Layout;
7use crate::NdProducer;
8use crate::Slice;
9
10/// Window producer and iterable
11///
12/// See [`.windows()`](crate::ArrayRef::windows) for more
13/// information.
14pub struct Windows<'a, A, D>
15{
16    base: RawArrayView<A, D>,
17    life: PhantomData<&'a A>,
18    window: D,
19    strides: D,
20}
21
22impl<'a, A, D: Dimension> Windows<'a, A, D>
23{
24    pub(crate) fn new<E>(a: ArrayView<'a, A, D>, window_size: E) -> Self
25    where E: IntoDimension<Dim = D>
26    {
27        let window = window_size.into_dimension();
28        let ndim = window.ndim();
29
30        let mut unit_stride = D::zeros(ndim);
31        unit_stride.slice_mut().fill(1);
32
33        Windows::new_with_stride(a, window, unit_stride)
34    }
35
36    pub(crate) fn new_with_stride<E>(a: ArrayView<'a, A, D>, window_size: E, axis_strides: E) -> Self
37    where E: IntoDimension<Dim = D>
38    {
39        let window = window_size.into_dimension();
40
41        let strides = axis_strides.into_dimension();
42        let window_strides = a.parts.strides.clone();
43
44        let base = build_base(a, window.clone(), strides);
45        Windows {
46            base: base.into_raw_view(),
47            life: PhantomData,
48            window,
49            strides: window_strides,
50        }
51    }
52}
53
54impl_ndproducer! {
55    ['a, A, D: Dimension]
56    [Clone => 'a, A, D: Clone ]
57    Windows {
58        base,
59        life,
60        window,
61        strides,
62    }
63    Windows<'a, A, D> {
64        type Item = ArrayView<'a, A, D>;
65        type Dim = D;
66
67        unsafe fn item(&self, ptr) {
68            ArrayView::new_(ptr, self.window.clone(),
69                            self.strides.clone())
70        }
71    }
72}
73
74impl<'a, A, D> IntoIterator for Windows<'a, A, D>
75where
76    D: Dimension,
77    A: 'a,
78{
79    type Item = <Self::IntoIter as Iterator>::Item;
80    type IntoIter = WindowsIter<'a, A, D>;
81    fn into_iter(self) -> Self::IntoIter
82    {
83        WindowsIter {
84            iter: self.base.into_base_iter(),
85            life: self.life,
86            window: self.window,
87            strides: self.strides,
88        }
89    }
90}
91
92/// Window iterator.
93///
94/// See [`.windows()`](crate::ArrayRef::windows) for more
95/// information.
96pub struct WindowsIter<'a, A, D>
97{
98    iter: Baseiter<A, D>,
99    life: PhantomData<&'a A>,
100    window: D,
101    strides: D,
102}
103
104impl_iterator! {
105    ['a, A, D: Dimension]
106    [Clone => 'a, A, D: Clone]
107    WindowsIter {
108        iter,
109        life,
110        window,
111        strides,
112    }
113    WindowsIter<'a, A, D> {
114        type Item = ArrayView<'a, A, D>;
115
116        fn item(&mut self, ptr) {
117            unsafe {
118                ArrayView::new(
119                    ptr,
120                    self.window.clone(),
121                    self.strides.clone())
122            }
123        }
124    }
125}
126
127send_sync_read_only!(Windows);
128send_sync_read_only!(WindowsIter);
129
130/// Window producer and iterable
131///
132/// See [`.axis_windows()`](crate::ArrayRef::axis_windows) for more
133/// information.
134pub struct AxisWindows<'a, A, D>
135{
136    base: ArrayView<'a, A, D>,
137    axis_idx: usize,
138    window: D,
139    strides: D,
140}
141
142impl<'a, A, D: Dimension> AxisWindows<'a, A, D>
143{
144    pub(crate) fn new_with_stride(a: ArrayView<'a, A, D>, axis: Axis, window_size: usize, stride_size: usize) -> Self
145    {
146        let window_strides = a.parts.strides.clone();
147        let axis_idx = axis.index();
148
149        let mut window = a.raw_dim();
150        window[axis_idx] = window_size;
151
152        let ndim = window.ndim();
153        let mut stride = D::zeros(ndim);
154        stride.slice_mut().fill(1);
155        stride[axis_idx] = stride_size;
156
157        let base = build_base(a, window.clone(), stride);
158        AxisWindows {
159            base,
160            axis_idx,
161            window,
162            strides: window_strides,
163        }
164    }
165}
166
167impl<'a, A, D: Dimension> NdProducer for AxisWindows<'a, A, D>
168{
169    type Item = ArrayView<'a, A, D>;
170    type Dim = Ix1;
171    type Ptr = *mut A;
172    type Stride = isize;
173
174    fn raw_dim(&self) -> Ix1
175    {
176        Ix1(self.base.raw_dim()[self.axis_idx])
177    }
178
179    fn layout(&self) -> Layout
180    {
181        self.base.layout()
182    }
183
184    fn as_ptr(&self) -> *mut A
185    {
186        self.base.as_ptr() as *mut _
187    }
188
189    fn contiguous_stride(&self) -> isize
190    {
191        self.base.contiguous_stride()
192    }
193
194    unsafe fn as_ref(&self, ptr: *mut A) -> Self::Item
195    {
196        ArrayView::new_(ptr, self.window.clone(), self.strides.clone())
197    }
198
199    unsafe fn uget_ptr(&self, i: &Self::Dim) -> *mut A
200    {
201        let mut d = D::zeros(self.base.ndim());
202        d[self.axis_idx] = i[0];
203        self.base.uget_ptr(&d)
204    }
205
206    fn stride_of(&self, axis: Axis) -> isize
207    {
208        assert_eq!(axis, Axis(0));
209        self.base.stride_of(Axis(self.axis_idx))
210    }
211
212    fn split_at(self, axis: Axis, index: usize) -> (Self, Self)
213    {
214        assert_eq!(axis, Axis(0));
215        let (a, b) = self.base.split_at(Axis(self.axis_idx), index);
216        (
217            AxisWindows {
218                base: a,
219                axis_idx: self.axis_idx,
220                window: self.window.clone(),
221                strides: self.strides.clone(),
222            },
223            AxisWindows {
224                base: b,
225                axis_idx: self.axis_idx,
226                window: self.window,
227                strides: self.strides,
228            },
229        )
230    }
231
232    private_impl!{}
233}
234
235impl<'a, A, D> IntoIterator for AxisWindows<'a, A, D>
236where
237    D: Dimension,
238    A: 'a,
239{
240    type Item = <Self::IntoIter as Iterator>::Item;
241    type IntoIter = WindowsIter<'a, A, D>;
242    fn into_iter(self) -> Self::IntoIter
243    {
244        WindowsIter {
245            iter: self.base.into_base_iter(),
246            life: PhantomData,
247            window: self.window,
248            strides: self.strides,
249        }
250    }
251}
252
253/// build the base array of the `Windows` and `AxisWindows` structs
254fn build_base<A, D>(a: ArrayView<A, D>, window: D, strides: D) -> ArrayView<A, D>
255where D: Dimension
256{
257    ndassert!(
258        a.ndim() == window.ndim(),
259        concat!(
260            "Window dimension {} does not match array dimension {} ",
261            "(with array of shape {:?})"
262        ),
263        window.ndim(),
264        a.ndim(),
265        a.shape()
266    );
267
268    ndassert!(
269        a.ndim() == strides.ndim(),
270        concat!(
271            "Stride dimension {} does not match array dimension {} ",
272            "(with array of shape {:?})"
273        ),
274        strides.ndim(),
275        a.ndim(),
276        a.shape()
277    );
278
279    let mut base = a;
280    base.slice_each_axis_inplace(|ax_desc| {
281        let len = ax_desc.len;
282        let wsz = window[ax_desc.axis.index()];
283        let stride = strides[ax_desc.axis.index()];
284
285        if len < wsz {
286            Slice::new(0, Some(0), 1)
287        } else {
288            Slice::new(0, Some((len - wsz + 1) as isize), stride as isize)
289        }
290    });
291    base
292}
ndarray/iterators/windows.rs

ndarray/iterators/
windows.rs
