Repository: rust-ndarray/ndarray (rust)

Issue:
Add axis_windows_iter
We should add an `axis_windows_iter` method:

```rust
impl<A, S, D> ArrayBase<S, D>
where
    S: RawData<Elem = A>,
    D: Dimension,
{
    /// Returns an iterator which traverses over all windows of a given length along an axis.
    ///
    /// The windows are all distinct, possibly-overlapping views. The shape of each window
    /// is the shape of `self`, with the length of `axis` replaced with `window_size`.
    ///
    /// **Panics** if `axis` is out-of-bounds or if `window_size` is zero.
    ///
    /// ```
    /// use ndarray::Array3;
    ///
    /// let arr = Array3::from_shape_fn([4, 5, 2], |(i, j, k)| i * 100 + j * 10 + k);
    /// let correct = vec![
    ///     arr.slice(s![.., 0..3, ..]),
    ///     arr.slice(s![.., 1..4, ..]),
    ///     arr.slice(s![.., 2..5, ..]),
    /// ];
    /// for (window, correct) in arr.axis_windows_iter(Axis(1), 3).zip(&correct) {
    ///     assert_eq!(window, correct);
    ///     assert_eq!(window.shape(), &[4, 3, 2]);
    /// }
    /// ```
    pub fn axis_windows_iter(&self, axis: Axis, window_size: usize) -> AxisWindowsIter<'_, A, D>
    where
        S: Data,
    { ... }
}
```

Optionally, we could use `NonZeroUsize` instead of `usize` as the type for `window_size`.

Existing, related methods include `axis_chunks_iter` and `windows`.

Pull request description:
Add `axis_windows` method
Closes #1041.

I have deviated a bit on the name, for consistency with [windows](https://docs.rs/ndarray/0.15.3/ndarray/struct.ArrayBase.html#method.windows) method.

I have not used `NonZeroUsize`, mainly because of the additional effort required from the User.

The doc comments have been taken verbatim from @jturner314's issue comments.


Patch:
diff --git a/src/impl_methods.rs b/src/impl_methods.rs
index 27abe689..cdba1791 100644
--- a/src/impl_methods.rs
+++ b/src/impl_methods.rs
@@ -1399,6 +1399,50 @@ where
         Windows::new(self.view(), window_size)
     }
 
+    /// Returns a producer which traverses over all windows of a given length along an axis.
+    ///
+    /// The windows are all distinct, possibly-overlapping views. The shape of each window
+    /// is the shape of `self`, with the length of `axis` replaced with `window_size`.
+    ///
+    /// **Panics** if `axis` is out-of-bounds or if `window_size` is zero.
+    ///
+    /// ```
+    /// use ndarray::{Array3, Axis, s};
+    ///
+    /// let arr = Array3::from_shape_fn([4, 5, 2], |(i, j, k)| i * 100 + j * 10 + k);
+    /// let correct = vec![
+    ///     arr.slice(s![.., 0..3, ..]),
+    ///     arr.slice(s![.., 1..4, ..]),
+    ///     arr.slice(s![.., 2..5, ..]),
+    /// ];
+    /// for (window, correct) in arr.axis_windows(Axis(1), 3).into_iter().zip(&correct) {
+    ///     assert_eq!(window, correct);
+    ///     assert_eq!(window.shape(), &[4, 3, 2]);
+    /// }
+    /// ```
+    pub fn axis_windows(&self, axis: Axis, window_size: usize) -> Windows<'_, A, D>
+    where
+        S: Data,
+    {
+        let axis_index = axis.index();
+
+        ndassert!(
+            axis_index < self.ndim(),
+            concat!(
+                "Window axis {} does not match array dimension {} ",
+                "(with array of shape {:?})"
+            ),
+            axis_index,
+            self.ndim(),
+            self.shape()
+        );
+
+        let mut size = self.raw_dim();
+        size[axis_index] = window_size;
+
+        Windows::new(self.view(), size)
+    }
+
     // Return (length, stride) for diagonal
     fn diag_params(&self) -> (Ix, Ixs) {
         /* empty shape has len 1 */


Test patch:
diff --git a/tests/windows.rs b/tests/windows.rs
index b0482e4b..664616f6 100644
--- a/tests/windows.rs
+++ b/tests/windows.rs
@@ -6,7 +6,7 @@
 )]
 
 use ndarray::prelude::*;
-use ndarray::Zip;
+use ndarray::{arr3, Zip};
 
 // Edge Cases for Windows iterator:
 //
@@ -79,7 +79,6 @@ fn windows_iterator_2d() {
 /// Simple test for iterating 3d-arrays via `Windows`.
 #[test]
 fn windows_iterator_3d() {
-    use ndarray::arr3;
     let a = Array::from_iter(10..37).into_shape((3, 3, 3)).unwrap();
     itertools::assert_equal(
         a.windows(Dim((2, 2, 2))),
@@ -117,6 +116,87 @@ fn test_window_zip() {
     }
 }
 
+/// Test verifies that non existant Axis results in panic
+#[test]
+#[should_panic]
+fn axis_windows_outofbound() {
+    let a = Array::from_iter(10..37).into_shape((3, 3, 3)).unwrap();
+    a.axis_windows(Axis(4), 2);
+}
+
+/// Test verifies that zero sizes results in panic
+#[test]
+#[should_panic]
+fn axis_windows_zero_size() {
+    let a = Array::from_iter(10..37).into_shape((3, 3, 3)).unwrap();
+    a.axis_windows(Axis(0), 0);
+}
+
+/// Test verifies that over sized windows yield nothing
+#[test]
+fn axis_windows_oversized() {
+    let a = Array::from_iter(10..37).into_shape((3, 3, 3)).unwrap();
+    let mut iter = a.axis_windows(Axis(2), 4).into_iter();
+    assert_eq!(iter.next(), None);
+}
+
+/// Simple test for iterating 1d-arrays via `Axis Windows`.
+#[test]
+fn test_axis_windows_1d() {
+    let a = Array::from_iter(10..20).into_shape(10).unwrap();
+
+    itertools::assert_equal(
+        a.axis_windows(Axis(0), 5),
+        vec![
+            arr1(&[10, 11, 12, 13, 14]),
+            arr1(&[11, 12, 13, 14, 15]),
+            arr1(&[12, 13, 14, 15, 16]),
+            arr1(&[13, 14, 15, 16, 17]),
+            arr1(&[14, 15, 16, 17, 18]),
+            arr1(&[15, 16, 17, 18, 19]),
+        ],
+    );
+}
+
+/// Simple test for iterating 2d-arrays via `Axis Windows`.
+#[test]
+fn test_axis_windows_2d() {
+    let a = Array::from_iter(10..30).into_shape((5, 4)).unwrap();
+
+    itertools::assert_equal(
+        a.axis_windows(Axis(0), 2),
+        vec![
+            arr2(&[[10, 11, 12, 13], [14, 15, 16, 17]]),
+            arr2(&[[14, 15, 16, 17], [18, 19, 20, 21]]),
+            arr2(&[[18, 19, 20, 21], [22, 23, 24, 25]]),
+            arr2(&[[22, 23, 24, 25], [26, 27, 28, 29]]),
+        ],
+    );
+}
+
+/// Simple test for iterating 3d-arrays via `Axis Windows`.
+#[test]
+fn test_axis_windows_3d() {
+    let a = Array::from_iter(0..27).into_shape((3, 3, 3)).unwrap();
+
+    itertools::assert_equal(
+        a.axis_windows(Axis(1), 2),
+        vec![
+            arr3(&[
+                [[0, 1, 2], [3, 4, 5]],
+                [[9, 10, 11], [12, 13, 14]],
+                [[18, 19, 20], [21, 22, 23]],
+            ]),
+            arr3(&[
+                [[3, 4, 5], [6, 7, 8]],
+                [[12, 13, 14], [15, 16, 17]],
+                [[21, 22, 23], [24, 25, 26]],
+            ]),
+        ],
+    );
+}
+
+
 #[test]
 fn test_window_neg_stride() {
     let array = Array::from_iter(1..10).into_shape((3, 3)).unwrap();