add axis_iter_mut

Author: vbarrielle

src/iterators.rs

@@ -630,3 +630,14 @@ pub fn new_outer_iter_mut<A, D>(v: ArrayViewMut<A, D>) -> OuterIterMut<A, D::Sma
         life: PhantomData,
     }
 }
+
+pub fn new_axis_iter_mut<A, D>(v: ArrayViewMut<A, D>,
+                               axis: usize
+                              ) -> OuterIterMut<A, D::Smaller>
+    where D: RemoveAxis,
+{
+    OuterIterMut {
+        iter: new_outer_core(v, axis),
+        life: PhantomData,
+    }
+}

src/lib.rs

@@ -1247,6 +1247,22 @@ impl<A, S, D> ArrayBase<S, D> where S: Data<Elem=A>, D: Dimension
         iterators::new_outer_iter_mut(self.view_mut())
     }
 
+    /// Return an iterator that traverses over the `axis` dimension
+    /// and yields each mutable subview.
+    ///
+    /// Iterator element is `ArrayViewMut<A, D::Smaller>`
+    /// (read-write array view).
+    ///
+    /// # Panics
+    ///
+    /// If axis is out of bounds.
+    pub fn axis_iter_mut(&mut self, axis: usize) -> OuterIterMut<A, D::Smaller>
+        where S: DataMut,
+              D: RemoveAxis,
+    {
+        iterators::new_axis_iter_mut(self.view_mut(), axis)
+    }
+
     // Return (length, stride) for diagonal
     fn diag_params(&self) -> (Ix, Ixs)
     {

tests/iterators.rs

@@ -8,6 +8,7 @@ use ndarray::{
     Data,
     Dimension,
     aview1,
+    arr3,
 };
 
 use itertools::assert_equal;
@@ -250,6 +251,31 @@ fn outer_iter_mut() {
     assert_equal(a.inner_iter(), found_rows);
 }
 
+#[test]
+fn axis_iter_mut() {
+    let a = Array::from_iter(0..12);
+    let a = a.reshape((2, 3, 2));
+    // [[[0, 1],
+    //   [2, 3],
+    //   [4, 5]],
+    //  [[6, 7],
+    //   [8, 9],
+    //    ...
+    let mut a = a.to_owned();
+
+    for mut subview in a.axis_iter_mut(1) {
+        subview[[0, 0]] = 42;
+    }
+
+    let b = arr3(&[[[42, 1],
+                    [42, 3],
+                    [42, 5]],
+                   [[6, 7],
+                    [8, 9],
+                    [10, 11]]]);
+    assert_eq!(a, b);
+}
+
 #[test]
 fn outer_iter_size_hint() {
     // Check that the size hint is correctly computed