Add support for inserting new axes while slicing

Author: jturner314

src/dimension/mod.rs

@@ -599,7 +599,11 @@ pub fn slices_intersect<D: Dimension>(
     indices2: &impl CanSlice<D>,
 ) -> bool {
     debug_assert_eq!(indices1.in_ndim(), indices2.in_ndim());
-    for (&axis_len, &si1, &si2) in izip!(dim.slice(), indices1.as_ref(), indices2.as_ref()) {
+    for (&axis_len, &si1, &si2) in izip!(
+        dim.slice(),
+        indices1.as_ref().iter().filter(|si| !si.is_new_axis()),
+        indices2.as_ref().iter().filter(|si| !si.is_new_axis()),
+    ) {
         // The slices do not intersect iff any pair of `AxisSliceInfo` does not intersect.
         match (si1, si2) {
             (
@@ -647,6 +651,7 @@ pub fn slices_intersect<D: Dimension>(
                     return false;
                 }
             }
+            (AxisSliceInfo::NewAxis, _) | (_, AxisSliceInfo::NewAxis) => unreachable!(),
         }
     }
     true
@@ -688,7 +693,7 @@ mod test {
     };
     use crate::error::{from_kind, ErrorKind};
     use crate::slice::Slice;
-    use crate::{Dim, Dimension, Ix0, Ix1, Ix2, Ix3, IxDyn};
+    use crate::{Dim, Dimension, Ix0, Ix1, Ix2, Ix3, IxDyn, NewAxis};
     use num_integer::gcd;
     use quickcheck::{quickcheck, TestResult};
 
@@ -962,17 +967,45 @@ mod test {
 
     #[test]
     fn slices_intersect_true() {
-        assert!(slices_intersect(&Dim([4, 5]), s![.., ..], s![.., ..]));
-        assert!(slices_intersect(&Dim([4, 5]), s![0, ..], s![0, ..]));
-        assert!(slices_intersect(&Dim([4, 5]), s![..;2, ..], s![..;3, ..]));
-        assert!(slices_intersect(&Dim([4, 5]), s![.., ..;2], s![.., 1..;3]));
+        assert!(slices_intersect(
+            &Dim([4, 5]),
+            s![NewAxis, .., NewAxis, ..],
+            s![.., NewAxis, .., NewAxis]
+        ));
+        assert!(slices_intersect(
+            &Dim([4, 5]),
+            s![NewAxis, 0, ..],
+            s![0, ..]
+        ));
+        assert!(slices_intersect(
+            &Dim([4, 5]),
+            s![..;2, ..],
+            s![..;3, NewAxis, ..]
+        ));
+        assert!(slices_intersect(
+            &Dim([4, 5]),
+            s![.., ..;2],
+            s![.., 1..;3, NewAxis]
+        ));
         assert!(slices_intersect(&Dim([4, 10]), s![.., ..;9], s![.., 3..;6]));
     }
 
     #[test]
     fn slices_intersect_false() {
-        assert!(!slices_intersect(&Dim([4, 5]), s![..;2, ..], s![1..;2, ..]));
-        assert!(!slices_intersect(&Dim([4, 5]), s![..;2, ..], s![1..;3, ..]));
-        assert!(!slices_intersect(&Dim([4, 5]), s![.., ..;9], s![.., 3..;6]));
+        assert!(!slices_intersect(
+            &Dim([4, 5]),
+            s![..;2, ..],
+            s![NewAxis, 1..;2, ..]
+        ));
+        assert!(!slices_intersect(
+            &Dim([4, 5]),
+            s![..;2, NewAxis, ..],
+            s![1..;3, ..]
+        ));
+        assert!(!slices_intersect(
+            &Dim([4, 5]),
+            s![.., ..;9],
+            s![.., 3..;6, NewAxis]
+        ));
     }
 }

src/doc/ndarray_for_numpy_users/mod.rs

@@ -532,7 +532,7 @@
 //! `a[:] = b` | [`a.assign(&b)`][.assign()] | copy the data from array `b` into array `a`
 //! `np.concatenate((a,b), axis=1)` | [`stack![Axis(1), a, b]`][stack!] or [`stack(Axis(1), &[a.view(), b.view()])`][stack()] | concatenate arrays `a` and `b` along axis 1
 //! `np.stack((a,b), axis=1)` | [`stack_new_axis![Axis(1), a, b]`][stack_new_axis!] or [`stack_new_axis(Axis(1), vec![a.view(), b.view()])`][stack_new_axis()] | stack arrays `a` and `b` along axis 1
-//! `a[:,np.newaxis]` or `np.expand_dims(a, axis=1)` | [`a.insert_axis(Axis(1))`][.insert_axis()] | create an array from `a`, inserting a new axis 1
+//! `a[:,np.newaxis]` or `np.expand_dims(a, axis=1)` | [`a.slice(s![.., NewAxis])`][.slice()] or [`a.insert_axis(Axis(1))`][.insert_axis()] | create an view of 1-D array `a`, inserting a new axis 1
 //! `a.transpose()` or `a.T` | [`a.t()`][.t()] or [`a.reversed_axes()`][.reversed_axes()] | transpose of array `a` (view for `.t()` or by-move for `.reversed_axes()`)
 //! `np.diag(a)` | [`a.diag()`][.diag()] | view the diagonal of `a`
 //! `a.flatten()` | [`use std::iter::FromIterator; Array::from_iter(a.iter().cloned())`][::from_iter()] | create a 1-D array by flattening `a`

src/impl_methods.rs

@@ -435,6 +435,12 @@ where
                     // Skip the old axis since it should be removed.
                     old_axis += 1;
                 }
+                AxisSliceInfo::NewAxis => {
+                    // Set the dim and stride of the new axis.
+                    new_dim[new_axis] = 1;
+                    new_strides[new_axis] = 0;
+                    new_axis += 1;
+                }
             });
             debug_assert_eq!(old_axis, self.ndim());
             debug_assert_eq!(new_axis, out_ndim);
@@ -448,6 +454,8 @@ where
 
     /// Slice the array in place without changing the number of dimensions.
     ///
+    /// Note that `NewAxis` elements in `info` are ignored.
+    ///
     /// See [*Slicing*](#slicing) for full documentation.
     ///
     /// **Panics** if an index is out of bounds or step size is zero.<br>
@@ -461,18 +469,20 @@ where
             self.ndim(),
             "The input dimension of `info` must match the array to be sliced.",
         );
-        info.as_ref()
-            .iter()
-            .enumerate()
-            .for_each(|(axis, &ax_info)| match ax_info {
+        let mut axis = 0;
+        info.as_ref().iter().for_each(|&ax_info| match ax_info {
                 AxisSliceInfo::Slice { start, end, step } => {
-                    self.slice_axis_inplace(Axis(axis), Slice { start, end, step })
+                    self.slice_axis_inplace(Axis(axis), Slice { start, end, step });
+                    axis += 1;
                 }
                 AxisSliceInfo::Index(index) => {
                     let i_usize = abs_index(self.len_of(Axis(axis)), index);
-                    self.collapse_axis(Axis(axis), i_usize)
+                    self.collapse_axis(Axis(axis), i_usize);
+                    axis += 1;
                 }
+                AxisSliceInfo::NewAxis => {}
             });
+        debug_assert_eq!(axis, self.ndim());
     }
 
     /// Return a view of the array, sliced along the specified axis.

src/lib.rs

@@ -140,7 +140,7 @@ pub use crate::dimension::IxDynImpl;
 pub use crate::dimension::NdIndex;
 pub use crate::error::{ErrorKind, ShapeError};
 pub use crate::indexes::{indices, indices_of};
-pub use crate::slice::{AxisSliceInfo, Slice, SliceInfo, SliceNextInDim, SliceNextOutDim};
+pub use crate::slice::{AxisSliceInfo, NewAxis, Slice, SliceInfo, SliceNextInDim, SliceNextOutDim};
 
 use crate::iterators::Baseiter;
 use crate::iterators::{ElementsBase, ElementsBaseMut, Iter, IterMut, Lanes, LanesMut};
@@ -494,14 +494,16 @@ pub type Ixs = isize;
 ///
 /// If a range is used, the axis is preserved. If an index is used, that index
 /// is selected and the axis is removed; this selects a subview. See
-/// [*Subviews*](#subviews) for more information about subviews. Note that
-/// [`.slice_collapse()`] behaves like [`.collapse_axis()`] by preserving
-/// the number of dimensions.
+/// [*Subviews*](#subviews) for more information about subviews. If a
+/// [`NewAxis`] instance is used, a new axis is inserted. Note that
+/// [`.slice_collapse()`] ignores `NewAxis` elements and behaves like
+/// [`.collapse_axis()`] by preserving the number of dimensions.
 ///
 /// [`.slice()`]: #method.slice
 /// [`.slice_mut()`]: #method.slice_mut
 /// [`.slice_move()`]: #method.slice_move
 /// [`.slice_collapse()`]: #method.slice_collapse
+/// [`NewAxis`]: struct.NewAxis.html
 ///
 /// It's possible to take multiple simultaneous *mutable* slices with
 /// [`.multi_slice_mut()`] or (for [`ArrayViewMut`] only)
@@ -512,7 +514,7 @@ pub type Ixs = isize;
 ///
 /// ```
 ///
-/// use ndarray::{arr2, arr3, s};
+/// use ndarray::{arr2, arr3, s, NewAxis};
 ///
 /// // 2 submatrices of 2 rows with 3 elements per row, means a shape of `[2, 2, 3]`.
 ///
@@ -547,16 +549,17 @@ pub type Ixs = isize;
 /// assert_eq!(d, e);
 /// assert_eq!(d.shape(), &[2, 1, 3]);
 ///
-/// // Let’s create a slice while selecting a subview with
+/// // Let’s create a slice while selecting a subview and inserting a new axis with
 /// //
 /// // - Both submatrices of the greatest dimension: `..`
 /// // - The last row in each submatrix, removing that axis: `-1`
 /// // - Row elements in reverse order: `..;-1`
-/// let f = a.slice(s![.., -1, ..;-1]);
-/// let g = arr2(&[[ 6,  5,  4],
-///                [12, 11, 10]]);
+/// // - A new axis at the end.
+/// let f = a.slice(s![.., -1, ..;-1, NewAxis]);
+/// let g = arr3(&[[ [6],  [5],  [4]],
+///                [[12], [11], [10]]]);
 /// assert_eq!(f, g);
-/// assert_eq!(f.shape(), &[2, 3]);
+/// assert_eq!(f.shape(), &[2, 3, 1]);
 ///
 /// // Let's take two disjoint, mutable slices of a matrix with
 /// //

src/prelude.rs

@@ -49,6 +49,9 @@ pub use crate::{array, azip, s};
 #[doc(no_inline)]
 pub use crate::ShapeBuilder;
 
+#[doc(no_inline)]
+pub use crate::NewAxis;
+
 #[doc(no_inline)]
 pub use crate::AsArray;