rust-ndarray · andrei-papou · Apr 3, 2019 · May 15, 2019 · May 31, 2019 · Oct 1, 2019
diff --git a/src/doc/ndarray_for_numpy_users/mod.rs b/src/doc/ndarray_for_numpy_users/mod.rs
@@ -533,6 +533,7 @@
 //! `a[:] = 3.` | [`a.fill(3.)`][.fill()] | set all array elements to the same scalar value
 //! `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.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`
@@ -642,6 +643,8 @@
 //! [.shape()]: ../../struct.ArrayBase.html#method.shape
 //! [stack!]: ../../macro.stack.html
 //! [stack()]: ../../fn.stack.html
+//! [stack_new_axis!]: ../../macro.stack_new_axis.html
+//! [stack_new_axis()]: ../../fn.stack_new_axis.html
 //! [.strides()]: ../../struct.ArrayBase.html#method.strides
 //! [.index_axis()]: ../../struct.ArrayBase.html#method.index_axis
 //! [.sum_axis()]: ../../struct.ArrayBase.html#method.sum_axis

diff --git a/src/lib.rs b/src/lib.rs
@@ -131,7 +131,7 @@ use crate::iterators::{ElementsBase, ElementsBaseMut, Iter, IterMut, Lanes, Lane
 
 pub use crate::arraytraits::AsArray;
 pub use crate::linalg_traits::{LinalgScalar, NdFloat};
-pub use crate::stacking::stack;
+pub use crate::stacking::{stack, stack_new_axis};
 
 pub use crate::impl_views::IndexLonger;
 pub use crate::shape_builder::ShapeBuilder;

diff --git a/src/slice.rs b/src/slice.rs
@@ -543,62 +543,62 @@ impl_slicenextdim_larger!((), Slice);
 #[macro_export]
 macro_rules! s(
     // convert a..b;c into @convert(a..b, c), final item
-    (@parse $dim:expr, [$($stack:tt)*] $r:expr;$s:expr) => {
+    (@parse $dim:expr, [$($concatenate:tt)*] $r:expr;$s:expr) => {
         match $r {
             r => {
                 let out_dim = $crate::SliceNextDim::next_dim(&r, $dim);
                 #[allow(unsafe_code)]
                 unsafe {
                     $crate::SliceInfo::new_unchecked(
-                        [$($stack)* $crate::s!(@convert r, $s)],
+                        [$($concatenate)* $crate::s!(@convert r, $s)],
                         out_dim,
                     )
                 }
             }
         }
     };
     // convert a..b into @convert(a..b), final item
-    (@parse $dim:expr, [$($stack:tt)*] $r:expr) => {
+    (@parse $dim:expr, [$($concatenate:tt)*] $r:expr) => {
         match $r {
             r => {
                 let out_dim = $crate::SliceNextDim::next_dim(&r, $dim);
                 #[allow(unsafe_code)]
                 unsafe {
                     $crate::SliceInfo::new_unchecked(
-                        [$($stack)* $crate::s!(@convert r)],
+                        [$($concatenate)* $crate::s!(@convert r)],
                         out_dim,
                     )
                 }
             }
         }
     };
     // convert a..b;c into @convert(a..b, c), final item, trailing comma
-    (@parse $dim:expr, [$($stack:tt)*] $r:expr;$s:expr ,) => {
-        $crate::s![@parse $dim, [$($stack)*] $r;$s]
+    (@parse $dim:expr, [$($concatenate:tt)*] $r:expr;$s:expr ,) => {
+        $crate::s![@parse $dim, [$($concatenate)*] $r;$s]
     };
     // convert a..b into @convert(a..b), final item, trailing comma
-    (@parse $dim:expr, [$($stack:tt)*] $r:expr ,) => {
-        $crate::s![@parse $dim, [$($stack)*] $r]
+    (@parse $dim:expr, [$($concatenate:tt)*] $r:expr ,) => {
+        $crate::s![@parse $dim, [$($concatenate)*] $r]
     };
     // convert a..b;c into @convert(a..b, c)
-    (@parse $dim:expr, [$($stack:tt)*] $r:expr;$s:expr, $($t:tt)*) => {
+    (@parse $dim:expr, [$($concatenate:tt)*] $r:expr;$s:expr, $($t:tt)*) => {
         match $r {
             r => {
                 $crate::s![@parse
                    $crate::SliceNextDim::next_dim(&r, $dim),
-                   [$($stack)* $crate::s!(@convert r, $s),]
+                   [$($concatenate)* $crate::s!(@convert r, $s),]
                    $($t)*
                 ]
             }
         }
     };
     // convert a..b into @convert(a..b)
-    (@parse $dim:expr, [$($stack:tt)*] $r:expr, $($t:tt)*) => {
+    (@parse $dim:expr, [$($concatenate:tt)*] $r:expr, $($t:tt)*) => {
         match $r {
             r => {
                 $crate::s![@parse
                    $crate::SliceNextDim::next_dim(&r, $dim),
-                   [$($stack)* $crate::s!(@convert r),]
+                   [$($concatenate)* $crate::s!(@convert r),]
                    $($t)*
                 ]
             }

diff --git a/src/stacking.rs b/src/stacking.rs
@@ -29,10 +29,7 @@ use crate::imp_prelude::*;
 ///                  [3., 3.]]))
 /// );
 /// ```
-pub fn stack<'a, A, D>(
-    axis: Axis,
-    arrays: &[ArrayView<'a, A, D>],
-) -> Result<Array<A, D>, ShapeError>
+pub fn stack<A, D>(axis: Axis, arrays: &[ArrayView<A, D>]) -> Result<Array<A, D>, ShapeError>
 where
     A: Copy,
     D: RemoveAxis,
@@ -76,14 +73,57 @@ where
     Ok(res)
 }
 
-/// Stack arrays along the given axis.
+pub fn stack_new_axis<A, D>(
+    axis: Axis,
+    arrays: Vec<ArrayView<A, D>>,
+) -> Result<Array<A, D::Larger>, ShapeError>
+where
+    A: Copy,
+    D: Dimension,
+    D::Larger: RemoveAxis,
+{
+    if arrays.is_empty() {
+        return Err(from_kind(ErrorKind::Unsupported));
+    }
+    let common_dim = arrays[0].raw_dim();
+    // Avoid panic on `insert_axis` call, return an Err instead of it.
+    if axis.index() > common_dim.ndim() {
+        return Err(from_kind(ErrorKind::OutOfBounds));
+    }
+    let mut res_dim = common_dim.insert_axis(axis);
+
+    if arrays.iter().any(|a| a.raw_dim() != common_dim) {
+        return Err(from_kind(ErrorKind::IncompatibleShape));
+    }
+
+    res_dim.set_axis(axis, arrays.len());
+
+    // we can safely use uninitialized values here because they are Copy
+    // and we will only ever write to them
+    let size = res_dim.size();
+    let mut v = Vec::with_capacity(size);
+    unsafe {
+        v.set_len(size);
+    }
+    let mut res = Array::from_shape_vec(res_dim, v)?;
+
+    res.axis_iter_mut(axis)
+        .zip(arrays.into_iter())
+        .for_each(|(mut assign_view, array)| {
+            assign_view.assign(&array);
+        });
+
+    Ok(res)
+}
+
+/// Concatenate arrays along the given axis.
 ///
 /// Uses the [`stack`][1] function, calling `ArrayView::from(&a)` on each
 /// argument `a`.
 ///
 /// [1]: fn.stack.html
 ///
-/// ***Panics*** if the `stack` function would return an error.
+/// ***Panics*** if the `concatenate` function would return an error.
 ///
 /// ```
 /// extern crate ndarray;
@@ -109,3 +149,37 @@ macro_rules! stack {
         $crate::stack($axis, &[ $($crate::ArrayView::from(&$array) ),* ]).unwrap()
     }
 }
+
+/// Stack arrays along the new axis.
+///
+/// Uses the [`stack_new_axis`][1] function, calling `ArrayView::from(&a)` on each
+/// argument `a`.
+///
+/// [1]: fn.stack_new_axis.html
+///
+/// ***Panics*** if the `stack` function would return an error.
+///
+/// ```
+/// extern crate ndarray;
+///
+/// use ndarray::{arr2, arr3, stack_new_axis, Axis};
+///
+/// # fn main() {
+///
+/// let a = arr2(&[[2., 2.],
+///                [3., 3.]]);
+/// assert!(
+///     stack_new_axis![Axis(0), a, a]
+///     == arr3(&[[[2., 2.],
+///                [3., 3.]],
+///               [[2., 2.],
+///                [3., 3.]]])
+/// );
+/// # }
+/// ```
+#[macro_export]
+macro_rules! stack_new_axis {
+    ($axis:expr, $( $array:expr ),+ ) => {
+        $crate::stack_new_axis($axis, vec![ $($crate::ArrayView::from(&$array) ),* ]).unwrap()
+    }
+}
diff --git a/tests/stacking.rs b/tests/stacking.rs
@@ -1,7 +1,7 @@
-use ndarray::{arr2, aview1, stack, Array2, Axis, ErrorKind};
+use ndarray::{arr2, arr3, aview1, stack, Array2, Axis, ErrorKind, Ix1};
 
 #[test]
-fn stacking() {
+fn concatenating() {
     let a = arr2(&[[2., 2.], [3., 3.]]);
     let b = ndarray::stack(Axis(0), &[a.view(), a.view()]).unwrap();
     assert_eq!(b, arr2(&[[2., 2.], [3., 3.], [2., 2.], [3., 3.]]));
@@ -24,3 +24,20 @@ fn stacking() {
     let res: Result<Array2<f64>, _> = ndarray::stack(Axis(0), &[]);
     assert_eq!(res.unwrap_err().kind(), ErrorKind::Unsupported);
 }
+
+#[test]
+fn stacking() {
+    let a = arr2(&[[2., 2.], [3., 3.]]);
+    let b = ndarray::stack_new_axis(Axis(0), vec![a.view(), a.view()]).unwrap();
+    assert_eq!(b, arr3(&[[[2., 2.], [3., 3.]], [[2., 2.], [3., 3.]]]));
+
+    let c = arr2(&[[3., 2., 3.], [2., 3., 2.]]);
+    let res = ndarray::stack_new_axis(Axis(1), vec![a.view(), c.view()]);
+    assert_eq!(res.unwrap_err().kind(), ErrorKind::IncompatibleShape);
+
+    let res = ndarray::stack_new_axis(Axis(3), vec![a.view(), a.view()]);
+    assert_eq!(res.unwrap_err().kind(), ErrorKind::OutOfBounds);
+
+    let res: Result<Array2<f64>, _> = ndarray::stack_new_axis::<_, Ix1>(Axis(0), vec![]);
+    assert_eq!(res.unwrap_err().kind(), ErrorKind::Unsupported);
+}