Repository: rust-ndarray/ndarray (rust)

Issue:
How to do parallel array `.sum`?
Hi thanks for the lib! It seems that `array.sum()` is a bit slow. So I wonder whether there is a way to parallelly sum it? I have lookat at [parallel](https://docs.rs/ndarray/0.15.3/ndarray/parallel/index.html) but it seems only to support `map`, and does not support things like `.fold()`.

Pull request:
Add a par_fold method to Zip to improve the discoverability of Rayon's fold-reduce idiom.
Fixes #1094 

Patch:
diff --git a/README.rst b/README.rst
index 50302488..a6ed9e90 100644
--- a/README.rst
+++ b/README.rst
@@ -15,7 +15,7 @@ __ https://docs.rs/ndarray/
    :alt: CI build status
 .. _build_status: https://github.com/rust-ndarray/ndarray/actions
 
-.. |crates| image:: http://meritbadge.herokuapp.com/ndarray
+.. |crates| image:: https://img.shields.io/crates/v/ndarray.svg
    :alt: ndarray at crates.io
 .. _crates: https://crates.io/crates/ndarray
 
diff --git a/src/impl_constructors.rs b/src/impl_constructors.rs
index e800a17e..df5e295c 100644
--- a/src/impl_constructors.rs
+++ b/src/impl_constructors.rs
@@ -230,6 +230,29 @@ where
         arr.diag_mut().assign(diag);
         arr
     }
+
+    /// Create a square 2D matrix of the specified size, with the specified
+    /// element along the diagonal and zeros elsewhere.
+    ///
+    /// **Panics** if `n * n` would overflow `isize`.
+    ///
+    /// ```rust
+    /// use ndarray::{array, Array2};
+    ///
+    /// let array = Array2::from_diag_elem(2, 5.);
+    /// assert_eq!(array, array![[5., 0.], [0., 5.]]);
+    /// ```
+    pub fn from_diag_elem(n: usize, elem: A) -> Self
+    where
+        S: DataMut,
+        A: Clone + Zero,
+    {
+        let mut eye = Self::zeros((n, n));
+        for a_ii in eye.diag_mut() {
+            *a_ii = elem.clone();
+        }
+        eye
+    }
 }
 
 #[cfg(not(debug_assertions))]
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 */
diff --git a/src/parallel/impl_par_methods.rs b/src/parallel/impl_par_methods.rs
index 33d39334..381b86b5 100644
--- a/src/parallel/impl_par_methods.rs
+++ b/src/parallel/impl_par_methods.rs
@@ -177,6 +177,56 @@ macro_rules! zip_impl {
                 self.par_map_assign_into(into, f)
             }
 
+            /// Parallel version of `fold`.
+            ///
+            /// Splits the producer in multiple tasks which each accumulate a single value
+            /// using the `fold` closure. Those tasks are executed in parallel and their results
+            /// are then combined to a single value using the `reduce` closure.
+            ///
+            /// The `identity` closure provides the initial values for each of the tasks and
+            /// for the final reduction.
+            ///
+            /// This is a shorthand for calling `self.into_par_iter().fold(...).reduce(...)`.
+            ///
+            /// Note that it is often more efficient to parallelize not per-element but rather
+            /// based on larger chunks of an array like generalized rows and operating on each chunk
+            /// using a sequential variant of the accumulation.
+            /// For example, sum each row sequentially and in parallel, taking advatange of locality
+            /// and vectorization within each task, and then reduce their sums to the sum of the matrix.
+            ///
+            /// Also note that the splitting of the producer into multiple tasks is _not_ deterministic
+            /// which needs to be considered when the accuracy of such an operation is analyzed.
+            ///
+            /// ## Examples
+            ///
+            /// ```rust
+            /// use ndarray::{Array, Zip};
+            ///
+            /// let a = Array::<usize, _>::ones((128, 1024));
+            /// let b = Array::<usize, _>::ones(128);
+            ///
+            /// let weighted_sum = Zip::from(a.rows()).and(&b).par_fold(
+            ///     || 0,
+            ///     |sum, row, factor| sum + row.sum() * factor,
+            ///     |sum, other_sum| sum + other_sum,
+            /// );
+            ///
+            /// assert_eq!(weighted_sum, a.len());
+            /// ```
+            pub fn par_fold<ID, F, R, T>(self, identity: ID, fold: F, reduce: R) -> T
+            where
+                ID: Fn() -> T + Send + Sync + Clone,
+                F: Fn(T, $($p::Item),*) -> T + Send + Sync,
+                R: Fn(T, T) -> T + Send + Sync,
+                T: Send
+            {
+                self.into_par_iter()
+                    .fold(identity.clone(), move |accumulator, ($($p,)*)| {
+                        fold(accumulator, $($p),*)
+                    })
+                    .reduce(identity, reduce)
+            }
+
             );
         }
         )+


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();