Redo parallel implementation of summator_incompr

This is a speed-up due to the call to `with_min_len` which however requires
indexed parallel iterators which `ndarray` provides only via `axis_iter` and
`axis_chunk_iter`, so this workaround is necessary until [1] is merged.

[1] rust-ndarray/ndarray#1081

Author: adamreichold

src/field.rs

@@ -1,4 +1,5 @@
-use ndarray::{Array1, Array2, ArrayView1, ArrayView2, Zip};
+use ndarray::{Array1, Array2, ArrayView1, ArrayView2, Axis, Zip};
+use rayon::prelude::*;
 
 pub fn summator(
     cov_samples: ArrayView2<'_, f64>,
@@ -38,34 +39,47 @@ pub fn summator_incompr(
     assert_eq!(cov_samples.dim().1, z1.dim());
     assert_eq!(cov_samples.dim().1, z2.dim());
 
-    let mut summed_modes = Array2::<f64>::zeros(pos.dim());
-
     // unit vector in x dir.
     let mut e1 = Array1::<f64>::zeros(pos.dim().0);
     e1[0] = 1.0;
 
-    Zip::from(cov_samples.columns())
-        .and(z1)
-        .and(z2)
-        .for_each(|cov_samples, z1, z2| {
-            let k_2 = cov_samples[0] / cov_samples.dot(&cov_samples);
-
-            Zip::from(pos.columns())
-                .and(summed_modes.columns_mut())
-                .par_for_each(|pos, mut summed_modes| {
-                    let phase = cov_samples.dot(&pos);
-                    let z12 = z1 * phase.cos() + z2 * phase.sin();
-
-                    Zip::from(&mut summed_modes)
-                        .and(&e1)
-                        .and(cov_samples)
-                        .for_each(|sum, e1, cs| {
-                            *sum += (*e1 - cs * k_2) * z12;
-                        });
-                });
-        });
-
-    summed_modes
+    // FIXME: Use Zip::from(cov_samples.columns()).and(z1).and(z1) after
+    // https://github.com/rust-ndarray/ndarray/pull/1081 is merged.
+    cov_samples
+        .axis_iter(Axis(1))
+        .into_par_iter()
+        .zip(z1.axis_iter(Axis(0)))
+        .zip(z2.axis_iter(Axis(0)))
+        .with_min_len(100)
+        .fold(
+            || Array2::<f64>::zeros(pos.dim()),
+            |mut summed_modes, ((cov_samples, z1), z2)| {
+                let k_2 = cov_samples[0] / cov_samples.dot(&cov_samples);
+                let z1 = z1.into_scalar();
+                let z2 = z2.into_scalar();
+
+                Zip::from(pos.columns())
+                    .and(summed_modes.columns_mut())
+                    .par_for_each(|pos, mut summed_modes| {
+                        let phase = cov_samples.dot(&pos);
+                        let z12 = z1 * phase.cos() + z2 * phase.sin();
+
+                        Zip::from(&mut summed_modes)
+                            .and(&e1)
+                            .and(cov_samples)
+                            .for_each(|sum, e1, cs| {
+                                *sum += (*e1 - cs * k_2) * z12;
+                            });
+                    });
+
+                summed_modes
+            },
+        )
+        .reduce_with(|mut lhs, rhs| {
+            lhs += &rhs;
+            lhs
+        })
+        .unwrap()
 }
 
 #[cfg(test)]