ndarray/iterators/
windows.rs

1use super::ElementsBase;
2use crate::imp_prelude::*;
3use crate::IntoDimension;
4use crate::Layout;
5use crate::NdProducer;
6
7/// Window producer and iterable
8///
9/// See [`.windows()`](ArrayBase::windows) for more
10/// information.
11pub struct Windows<'a, A, D> {
12    base: ArrayView<'a, A, D>,
13    window: D,
14    strides: D,
15}
16
17impl<'a, A, D: Dimension> Windows<'a, A, D> {
18    pub(crate) fn new<E>(a: ArrayView<'a, A, D>, window_size: E) -> Self
19    where
20        E: IntoDimension<Dim = D>,
21    {
22        let window = window_size.into_dimension();
23        ndassert!(
24            a.ndim() == window.ndim(),
25            concat!(
26                "Window dimension {} does not match array dimension {} ",
27                "(with array of shape {:?})"
28            ),
29            window.ndim(),
30            a.ndim(),
31            a.shape()
32        );
33        let mut size = a.dim;
34        for (sz, &ws) in size.slice_mut().iter_mut().zip(window.slice()) {
35            assert_ne!(ws, 0, "window-size must not be zero!");
36            // cannot use std::cmp::max(0, ..) since arithmetic underflow panics
37            *sz = if *sz < ws { 0 } else { *sz - ws + 1 };
38        }
39
40        let window_strides = a.strides.clone();
41
42        unsafe {
43            Windows {
44                base: ArrayView::new(a.ptr, size, a.strides),
45                window,
46                strides: window_strides,
47            }
48        }
49    }
50}
51
52impl_ndproducer! {
53    ['a, A, D: Dimension]
54    [Clone => 'a, A, D: Clone ]
55    Windows {
56        base,
57        window,
58        strides,
59    }
60    Windows<'a, A, D> {
61        type Item = ArrayView<'a, A, D>;
62        type Dim = D;
63
64        unsafe fn item(&self, ptr) {
65            ArrayView::new_(ptr, self.window.clone(),
66                            self.strides.clone())
67        }
68    }
69}
70
71impl<'a, A, D> IntoIterator for Windows<'a, A, D>
72where
73    D: Dimension,
74    A: 'a,
75{
76    type Item = <Self::IntoIter as Iterator>::Item;
77    type IntoIter = WindowsIter<'a, A, D>;
78    fn into_iter(self) -> Self::IntoIter {
79        WindowsIter {
80            iter: self.base.into_elements_base(),
81            window: self.window,
82            strides: self.strides,
83        }
84    }
85}
86
87/// Window iterator.
88///
89/// See [`.windows()`](ArrayBase::windows) for more
90/// information.
91pub struct WindowsIter<'a, A, D> {
92    iter: ElementsBase<'a, A, D>,
93    window: D,
94    strides: D,
95}
96
97impl_iterator! {
98    ['a, A, D: Dimension]
99    [Clone => 'a, A, D: Clone]
100    WindowsIter {
101        iter,
102        window,
103        strides,
104    }
105    WindowsIter<'a, A, D> {
106        type Item = ArrayView<'a, A, D>;
107
108        fn item(&mut self, elt) {
109            unsafe {
110                ArrayView::new_(
111                    elt,
112                    self.window.clone(),
113                    self.strides.clone())
114            }
115        }
116    }
117}