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()`](ArrayBase::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.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()`](ArrayBase::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()`](ArrayBase::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(a: ArrayView<'a, A, D>, axis: Axis, window_size: usize) -> Self
145    {
146        let window_strides = a.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 unit_stride = D::zeros(ndim);
154        unit_stride.slice_mut().fill(1);
155
156        let base = build_base(a, window.clone(), unit_stride);
157        AxisWindows {
158            base,
159            axis_idx,
160            window,
161            strides: window_strides,
162        }
163    }
164}
165
166impl<'a, A, D: Dimension> NdProducer for AxisWindows<'a, A, D>
167{
168    type Item = ArrayView<'a, A, D>;
169    type Dim = Ix1;
170    type Ptr = *mut A;
171    type Stride = isize;
172
173    fn raw_dim(&self) -> Ix1
174    {
175        Ix1(self.base.raw_dim()[self.axis_idx])
176    }
177
178    fn layout(&self) -> Layout
179    {
180        self.base.layout()
181    }
182
183    fn as_ptr(&self) -> *mut A
184    {
185        self.base.as_ptr() as *mut _
186    }
187
188    fn contiguous_stride(&self) -> isize
189    {
190        self.base.contiguous_stride()
191    }
192
193    unsafe fn as_ref(&self, ptr: *mut A) -> Self::Item
194    {
195        ArrayView::new_(ptr, self.window.clone(), self.strides.clone())
196    }
197
198    unsafe fn uget_ptr(&self, i: &Self::Dim) -> *mut A
199    {
200        let mut d = D::zeros(self.base.ndim());
201        d[self.axis_idx] = i[0];
202        self.base.uget_ptr(&d)
203    }
204
205    fn stride_of(&self, axis: Axis) -> isize
206    {
207        assert_eq!(axis, Axis(0));
208        self.base.stride_of(Axis(self.axis_idx))
209    }
210
211    fn split_at(self, axis: Axis, index: usize) -> (Self, Self)
212    {
213        assert_eq!(axis, Axis(0));
214        let (a, b) = self.base.split_at(Axis(self.axis_idx), index);
215        (
216            AxisWindows {
217                base: a,
218                axis_idx: self.axis_idx,
219                window: self.window.clone(),
220                strides: self.strides.clone(),
221            },
222            AxisWindows {
223                base: b,
224                axis_idx: self.axis_idx,
225                window: self.window,
226                strides: self.strides,
227            },
228        )
229    }
230
231    private_impl!{}
232}
233
234impl<'a, A, D> IntoIterator for AxisWindows<'a, A, D>
235where
236    D: Dimension,
237    A: 'a,
238{
239    type Item = <Self::IntoIter as Iterator>::Item;
240    type IntoIter = WindowsIter<'a, A, D>;
241    fn into_iter(self) -> Self::IntoIter
242    {
243        WindowsIter {
244            iter: self.base.into_base_iter(),
245            life: PhantomData,
246            window: self.window,
247            strides: self.strides,
248        }
249    }
250}
251
252/// build the base array of the `Windows` and `AxisWindows` structs
253fn build_base<A, D>(a: ArrayView<A, D>, window: D, strides: D) -> ArrayView<A, D>
254where D: Dimension
255{
256    ndassert!(
257        a.ndim() == window.ndim(),
258        concat!(
259            "Window dimension {} does not match array dimension {} ",
260            "(with array of shape {:?})"
261        ),
262        window.ndim(),
263        a.ndim(),
264        a.shape()
265    );
266
267    ndassert!(
268        a.ndim() == strides.ndim(),
269        concat!(
270            "Stride dimension {} does not match array dimension {} ",
271            "(with array of shape {:?})"
272        ),
273        strides.ndim(),
274        a.ndim(),
275        a.shape()
276    );
277
278    let mut base = a;
279    base.slice_each_axis_inplace(|ax_desc| {
280        let len = ax_desc.len;
281        let wsz = window[ax_desc.axis.index()];
282        let stride = strides[ax_desc.axis.index()];
283
284        if len < wsz {
285            Slice::new(0, Some(0), 1)
286        } else {
287            Slice::new(0, Some((len - wsz + 1) as isize), stride as isize)
288        }
289    });
290    base
291}
ndarray/iterators/windows.rs

ndarray/iterators/
windows.rs
