ndarray/iterators/
windows.rs1use super::ElementsBase;
2use crate::imp_prelude::*;
3use crate::IntoDimension;
4use crate::Layout;
5use crate::NdProducer;
6
7pub 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 *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
87pub 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}