Skip to content

Commit 859cc2b

Browse files
committed
move core logic of creating the Windows logic into seperate fn
1 parent 51bccf9 commit 859cc2b

File tree

1 file changed

+50
-49
lines changed

1 file changed

+50
-49
lines changed

src/iterators/windows.rs

+50-49
Original file line numberDiff line numberDiff line change
@@ -38,41 +38,7 @@ impl<'a, A, D: Dimension> Windows<'a, A, D> {
3838
let strides = axis_strides.into_dimension();
3939
let window_strides = a.strides.clone();
4040

41-
ndassert!(
42-
a.ndim() == window.ndim(),
43-
concat!(
44-
"Window dimension {} does not match array dimension {} ",
45-
"(with array of shape {:?})"
46-
),
47-
window.ndim(),
48-
a.ndim(),
49-
a.shape()
50-
);
51-
52-
ndassert!(
53-
a.ndim() == strides.ndim(),
54-
concat!(
55-
"Stride dimension {} does not match array dimension {} ",
56-
"(with array of shape {:?})"
57-
),
58-
strides.ndim(),
59-
a.ndim(),
60-
a.shape()
61-
);
62-
63-
let mut base = a;
64-
base.slice_each_axis_inplace(|ax_desc| {
65-
let len = ax_desc.len;
66-
let wsz = window[ax_desc.axis.index()];
67-
let stride = strides[ax_desc.axis.index()];
68-
69-
if len < wsz {
70-
Slice::new(0, Some(0), 1)
71-
} else {
72-
Slice::new(0, Some((len - wsz + 1) as isize), stride as isize)
73-
}
74-
});
75-
41+
let base = build_base(a, window.clone(), strides);
7642
Windows {
7743
base,
7844
window,
@@ -162,28 +128,22 @@ pub struct AxisWindows<'a, A, D>{
162128
impl<'a, A, D: Dimension> AxisWindows<'a, A, D> {
163129
pub(crate) fn new(a: ArrayView<'a, A, D>, axis: Axis, window_size: usize) -> Self
164130
{
165-
let strides = a.strides.clone();
166-
let mut base = a;
131+
let window_strides = a.strides.clone();
167132
let axis_idx = axis.index();
168-
let mut window = base.raw_dim();
169-
window[axis_idx] = window_size;
170133

171-
base.slice_each_axis_inplace(|ax_desc| {
172-
let len = ax_desc.len;
173-
let wsz = window[ax_desc.axis.index()];
134+
let mut window = a.raw_dim();
135+
window[axis_idx] = window_size;
174136

175-
if len < wsz {
176-
Slice::new(0, Some(0), 1)
177-
} else {
178-
Slice::new(0, Some((len - wsz + 1) as isize), 1)
179-
}
180-
});
137+
let ndim = window.ndim();
138+
let mut unit_stride = D::zeros(ndim);
139+
unit_stride.slice_mut().fill(1);
181140

141+
let base = build_base(a, window.clone(), unit_stride);
182142
AxisWindows {
183143
base,
184144
axis_idx,
185145
window,
186-
strides,
146+
strides: window_strides,
187147
}
188148
}
189149
}
@@ -263,3 +223,44 @@ where
263223
}
264224
}
265225
}
226+
227+
/// build the base array of the `Windows` and `AxisWindows` structs
228+
fn build_base<'a, A, D>(a: ArrayView<'a, A, D>, window: D, strides: D) -> ArrayView<'a, A, D>
229+
where D:Dimension
230+
{
231+
ndassert!(
232+
a.ndim() == window.ndim(),
233+
concat!(
234+
"Window dimension {} does not match array dimension {} ",
235+
"(with array of shape {:?})"
236+
),
237+
window.ndim(),
238+
a.ndim(),
239+
a.shape()
240+
);
241+
242+
ndassert!(
243+
a.ndim() == strides.ndim(),
244+
concat!(
245+
"Stride dimension {} does not match array dimension {} ",
246+
"(with array of shape {:?})"
247+
),
248+
strides.ndim(),
249+
a.ndim(),
250+
a.shape()
251+
);
252+
253+
let mut base = a;
254+
base.slice_each_axis_inplace(|ax_desc| {
255+
let len = ax_desc.len;
256+
let wsz = window[ax_desc.axis.index()];
257+
let stride = strides[ax_desc.axis.index()];
258+
259+
if len < wsz {
260+
Slice::new(0, Some(0), 1)
261+
} else {
262+
Slice::new(0, Some((len - wsz + 1) as isize), stride as isize)
263+
}
264+
});
265+
base
266+
}

0 commit comments

Comments
 (0)