Merge pull request #291 from bluss/whole-chunks-mut

Add .whole_chunks_mut()

Author: bluss

benches/bench1.rs

@@ -348,6 +348,45 @@ fn add_2d_zip_cutout(bench: &mut test::Bencher)
     });
 }
 
+#[bench]
+fn add_2d_cutouts_by_4(bench: &mut test::Bencher)
+{
+    let mut a = Array::<i32, _>::zeros((64 * 1, 64 * 1));
+    let b = Array::<i32, _>::zeros((64 * 1, 64 * 1));
+    let chunksz = (4, 4);
+    bench.iter(|| {
+        Zip::from(a.whole_chunks_mut(chunksz))
+            .and(b.whole_chunks(chunksz))
+            .apply(|mut a, b| a += &b);
+    });
+}
+
+#[bench]
+fn add_2d_cutouts_by_16(bench: &mut test::Bencher)
+{
+    let mut a = Array::<i32, _>::zeros((64 * 1, 64 * 1));
+    let b = Array::<i32, _>::zeros((64 * 1, 64 * 1));
+    let chunksz = (16, 16);
+    bench.iter(|| {
+        Zip::from(a.whole_chunks_mut(chunksz))
+            .and(b.whole_chunks(chunksz))
+            .apply(|mut a, b| a += &b);
+    });
+}
+
+#[bench]
+fn add_2d_cutouts_by_32(bench: &mut test::Bencher)
+{
+    let mut a = Array::<i32, _>::zeros((64 * 1, 64 * 1));
+    let b = Array::<i32, _>::zeros((64 * 1, 64 * 1));
+    let chunksz = (32, 32);
+    bench.iter(|| {
+        Zip::from(a.whole_chunks_mut(chunksz))
+            .and(b.whole_chunks(chunksz))
+            .apply(|mut a, b| a += &b);
+    });
+}
+
 #[bench]
 fn add_2d_broadcast_1_to_2(bench: &mut test::Bencher)
 {

benches/chunks.rs

@@ -0,0 +1,35 @@
+#![feature(test)]
+
+extern crate test;
+use test::Bencher;
+
+#[macro_use(azip)]
+extern crate ndarray;
+use ndarray::prelude::*;
+use ndarray::NdProducer;
+
+#[bench]
+fn chunk2x2_sum(bench: &mut Bencher)
+{
+    let a = Array::<f32, _>::zeros((256, 256));
+    let chunksz = (2, 2);
+    let mut sum = Array::zeros(a.whole_chunks(chunksz).raw_dim());
+    bench.iter(|| {
+        azip!(ref a (a.whole_chunks(chunksz)), mut sum in {
+            *sum = a.iter().sum::<f32>();
+        });
+    });
+}
+
+#[bench]
+fn chunk2x2_scalar_sum(bench: &mut Bencher)
+{
+    let a = Array::<f32, _>::zeros((256, 256));
+    let chunksz = (2, 2);
+    let mut sum = Array::zeros(a.whole_chunks(chunksz).raw_dim());
+    bench.iter(|| {
+        azip!(ref a (a.whole_chunks(chunksz)), mut sum in {
+            *sum = a.scalar_sum();
+        });
+    });
+}

src/arrayformat.rs

@@ -10,6 +10,7 @@ use super::{
     ArrayBase,
     Data,
     Dimension,
+    NdProducer,
 };
 use dimension::IntoDimension;
 
@@ -109,7 +110,8 @@ impl<'a, A: fmt::Debug, S, D: Dimension> fmt::Debug for ArrayBase<S, D>
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         // Add extra information for Debug
         try!(format_array(self, f, <_>::fmt));
-        try!(write!(f, " shape={:?}, strides={:?}", self.shape(), self.strides()));
+        try!(write!(f, " shape={:?}, strides={:?}, layout={:?}",
+                    self.shape(), self.strides(), layout=self.view().layout()));
         Ok(())
     }
 }

src/impl_methods.rs

@@ -22,10 +22,11 @@ use super::zipsl;
 use super::ZipExt;
 use dimension::IntoDimension;
 use dimension::{axes_of, Axes, merge_axes, stride_offset};
-use iterators::whole_chunks_of;
 use iterators::{
     new_inner_iter_smaller,
     new_inner_iter_smaller_mut,
+    whole_chunks_of,
+    whole_chunks_mut_of,
 };
 
 use {
@@ -41,6 +42,7 @@ use {
     AxisIter,
     AxisIterMut,
     WholeChunks,
+    WholeChunksMut,
 };
 use stacking::stack;
 
@@ -619,15 +621,57 @@ impl<A, S, D> ArrayBase<S, D> where S: Data<Elem=A>, D: Dimension
     /// It produces the whole chunks of a given n-dimensional chunk size,
     /// skipping the remainder along each dimension that doesn't fit evenly.
     ///
-    /// Iterator element is `ArrayView<A, D>`
+    /// The produced element is a `ArrayView<A, D>` with exactly the dimension
+    /// `chunk_size`.
     ///
-    /// **Panics** if any dimension of `chunk_size` is zero
+    /// **Panics** if any dimension of `chunk_size` is zero<br>
+    /// (**Panics** if `D` is `IxDyn` and `chunk_size` does not match the
+    /// number of array axes.)
     pub fn whole_chunks<E>(&self, chunk_size: E) -> WholeChunks<A, D> 
         where E: IntoDimension<Dim=D>,
     {
         whole_chunks_of(self.view(), chunk_size)
     }
 
+    /// Return a whole chunks producer (and iterable).
+    ///
+    /// It produces the whole chunks of a given n-dimensional chunk size,
+    /// skipping the remainder along each dimension that doesn't fit evenly.
+    ///
+    /// The produced element is a `ArrayViewMut<A, D>` with exactly
+    /// the dimension `chunk_size`.
+    ///
+    /// **Panics** if any dimension of `chunk_size` is zero<br>
+    /// (**Panics** if `D` is `IxDyn` and `chunk_size` does not match the
+    /// number of array axes.)
+    ///
+    /// ```rust
+    /// use ndarray::Array;
+    /// use ndarray::arr2;
+    /// let mut a = Array::zeros((6, 7));
+    ///
+    /// // Fill each 2 × 2 chunk with the index of where it appeared in iteration
+    /// for (i, mut chunk) in a.whole_chunks_mut((2, 2)).into_iter().enumerate() {
+    ///     chunk.fill(i);
+    /// }
+    ///
+    /// // The resulting array is:
+    /// assert_eq!(
+    ///   a,
+    ///   arr2(&[[0, 0, 1, 1, 2, 2, 0],
+    ///          [0, 0, 1, 1, 2, 2, 0],
+    ///          [3, 3, 4, 4, 5, 5, 0],
+    ///          [3, 3, 4, 4, 5, 5, 0],
+    ///          [6, 6, 7, 7, 8, 8, 0],
+    ///          [6, 6, 7, 7, 8, 8, 0]]));
+    /// ```
+    pub fn whole_chunks_mut<E>(&mut self, chunk_size: E) -> WholeChunksMut<A, D> 
+        where E: IntoDimension<Dim=D>,
+              S: DataMut
+    {
+        whole_chunks_mut_of(self.view_mut(), chunk_size)
+    }
+
     // Return (length, stride) for diagonal
     fn diag_params(&self) -> (Ix, Ixs) {
         /* empty shape has len 1 */

src/impl_views.rs

@@ -75,7 +75,7 @@ impl<'a, A, D> ArrayBase<ViewRepr<&'a A>, D>
         ArrayView::new_(ptr, dim, strides)
     }
 
-    /// Split the array along `axis` and return one view strictly before the
+    /// Split the array view along `axis` and return one view strictly before the
     /// split and one view after the split.
     ///
     /// **Panics** if `axis` or `index` is out of bounds.
@@ -190,7 +190,7 @@ impl<'a, A, D> ArrayBase<ViewRepr<&'a mut A>, D>
         ArrayViewMut::new_(ptr, dim, strides)
     }
 
-    /// Split the array along `axis` and return one mutable view strictly
+    /// Split the array view along `axis` and return one mutable view strictly
     /// before the split and one mutable view after the split.
     ///
     /// **Panics** if `axis` or `index` is out of bounds.