Skip to content

Commit 48606c6

Browse files
andrei-papoubluss
authored andcommitted
Renamed stack to concatenate and added numpy-compliant stack
function
1 parent bdc0724 commit 48606c6

File tree

5 files changed

+96
-23
lines changed

5 files changed

+96
-23
lines changed

src/doc/ndarray_for_numpy_users/mod.rs

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -532,7 +532,8 @@
532532
//! ------|-----------|------
533533
//! `a[:] = 3.` | [`a.fill(3.)`][.fill()] | set all array elements to the same scalar value
534534
//! `a[:] = b` | [`a.assign(&b)`][.assign()] | copy the data from array `b` into array `a`
535-
//! `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
535+
//! `np.concatenate((a,b), axis=1)` | [`concatenate![Axis(1), a, b]`][concatenate!] or [`concatenate(Axis(1), &[a.view(), b.view()])`][concatenate()] | concatenate arrays `a` and `b` along axis 1
536+
//! `np.stack((a,b), axis=1)` | [`stack![Axis(1), a, b]`][stack!] or [`stack(Axis(1), vec![a.view(), b.view()])`][stack()] | stack arrays `a` and `b` along axis 1
536537
//! `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
537538
//! `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()`)
538539
//! `np.diag(a)` | [`a.diag()`][.diag()] | view the diagonal of `a`
@@ -640,6 +641,8 @@
640641
//! [.slice_move()]: ../../struct.ArrayBase.html#method.slice_move
641642
//! [.slice_mut()]: ../../struct.ArrayBase.html#method.slice_mut
642643
//! [.shape()]: ../../struct.ArrayBase.html#method.shape
644+
//! [concatenate!]: ../../macro.concatenate.html
645+
//! [concatenate()]: ../../fn.concatenate.html
643646
//! [stack!]: ../../macro.stack.html
644647
//! [stack()]: ../../fn.stack.html
645648
//! [.strides()]: ../../struct.ArrayBase.html#method.strides

src/impl_methods.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ use crate::iter::{
2828
IndexedIter, IndexedIterMut, Iter, IterMut, Lanes, LanesMut, Windows,
2929
};
3030
use crate::slice::MultiSlice;
31-
use crate::stacking::stack;
31+
use crate::stacking::concatenate;
3232
use crate::{NdIndex, Slice, SliceInfo, SliceOrIndex};
3333

3434
/// # Methods For All Array Types
@@ -840,7 +840,7 @@ where
840840
dim.set_axis(axis, 0);
841841
unsafe { Array::from_shape_vec_unchecked(dim, vec![]) }
842842
} else {
843-
stack(axis, &subs).unwrap()
843+
concatenate(axis, &subs).unwrap()
844844
}
845845
}
846846

src/lib.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -131,7 +131,7 @@ use crate::iterators::{ElementsBase, ElementsBaseMut, Iter, IterMut, Lanes, Lane
131131

132132
pub use crate::arraytraits::AsArray;
133133
pub use crate::linalg_traits::{LinalgScalar, NdFloat};
134-
pub use crate::stacking::stack;
134+
pub use crate::stacking::{concatenate, stack};
135135

136136
pub use crate::impl_views::IndexLonger;
137137
pub use crate::shape_builder::ShapeBuilder;

src/stacking.rs

Lines changed: 64 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -17,21 +17,21 @@ use crate::imp_prelude::*;
1717
/// if the result is larger than is possible to represent.
1818
///
1919
/// ```
20-
/// use ndarray::{arr2, Axis, stack};
20+
/// use ndarray::{arr2, Axis, concatenate};
2121
///
2222
/// let a = arr2(&[[2., 2.],
2323
/// [3., 3.]]);
2424
/// assert!(
25-
/// stack(Axis(0), &[a.view(), a.view()])
25+
/// concatenate(Axis(0), &[a.view(), a.view()])
2626
/// == Ok(arr2(&[[2., 2.],
2727
/// [3., 3.],
2828
/// [2., 2.],
2929
/// [3., 3.]]))
3030
/// );
3131
/// ```
32-
pub fn stack<'a, A, D>(
32+
pub fn concatenate<A, D>(
3333
axis: Axis,
34-
arrays: &[ArrayView<'a, A, D>],
34+
arrays: &[ArrayView<A, D>],
3535
) -> Result<Array<A, D>, ShapeError>
3636
where
3737
A: Copy,
@@ -76,26 +76,45 @@ where
7676
Ok(res)
7777
}
7878

79-
/// Stack arrays along the given axis.
79+
pub fn stack<A, D>(
80+
axis: Axis,
81+
arrays: Vec<ArrayView<A, D>>,
82+
) -> Result<Array<A, D::Larger>, ShapeError>
83+
where
84+
A: Copy,
85+
D: Dimension,
86+
D::Larger: RemoveAxis,
87+
{
88+
if let Some(ndim) = D::NDIM {
89+
if axis.index() > ndim {
90+
return Err(from_kind(ErrorKind::OutOfBounds));
91+
}
92+
}
93+
let arrays: Vec<ArrayView<A, D::Larger>> = arrays.into_iter()
94+
.map(|a| a.insert_axis(axis)).collect();
95+
concatenate(axis, &arrays)
96+
}
97+
98+
/// Concatenate arrays along the given axis.
8099
///
81-
/// Uses the [`stack`][1] function, calling `ArrayView::from(&a)` on each
100+
/// Uses the [`concatenate`][1] function, calling `ArrayView::from(&a)` on each
82101
/// argument `a`.
83102
///
84-
/// [1]: fn.stack.html
103+
/// [1]: fn.concatenate.html
85104
///
86-
/// ***Panics*** if the `stack` function would return an error.
105+
/// ***Panics*** if the `concatenate` function would return an error.
87106
///
88107
/// ```
89108
/// extern crate ndarray;
90109
///
91-
/// use ndarray::{arr2, stack, Axis};
110+
/// use ndarray::{arr2, concatenate, Axis};
92111
///
93112
/// # fn main() {
94113
///
95114
/// let a = arr2(&[[2., 2.],
96115
/// [3., 3.]]);
97116
/// assert!(
98-
/// stack![Axis(0), a, a]
117+
/// concatenate![Axis(0), a, a]
99118
/// == arr2(&[[2., 2.],
100119
/// [3., 3.],
101120
/// [2., 2.],
@@ -104,8 +123,42 @@ where
104123
/// # }
105124
/// ```
106125
#[macro_export]
126+
macro_rules! concatenate {
127+
($axis:expr, $( $array:expr ),+ ) => {
128+
$crate::concatenate($axis, &[ $($crate::ArrayView::from(&$array) ),* ]).unwrap()
129+
}
130+
}
131+
132+
/// Stack arrays along the new axis.
133+
///
134+
/// Uses the [`stack`][1] function, calling `ArrayView::from(&a)` on each
135+
/// argument `a`.
136+
///
137+
/// [1]: fn.concatenate.html
138+
///
139+
/// ***Panics*** if the `stack` function would return an error.
140+
///
141+
/// ```
142+
/// extern crate ndarray;
143+
///
144+
/// use ndarray::{arr2, arr3, stack, Axis};
145+
///
146+
/// # fn main() {
147+
///
148+
/// let a = arr2(&[[2., 2.],
149+
/// [3., 3.]]);
150+
/// assert!(
151+
/// stack![Axis(0), a, a]
152+
/// == arr3(&[[[2., 2.],
153+
/// [3., 3.]],
154+
/// [[2., 2.],
155+
/// [3., 3.]]])
156+
/// );
157+
/// # }
158+
/// ```
159+
#[macro_export]
107160
macro_rules! stack {
108161
($axis:expr, $( $array:expr ),+ ) => {
109-
$crate::stack($axis, &[ $($crate::ArrayView::from(&$array) ),* ]).unwrap()
162+
$crate::stack($axis, vec![ $($crate::ArrayView::from(&$array) ),* ]).unwrap()
110163
}
111164
}

tests/stacking.rs

Lines changed: 25 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,26 +1,43 @@
1-
use ndarray::{arr2, aview1, stack, Array2, Axis, ErrorKind};
1+
use ndarray::{arr2, arr3, aview1, concatenate, Array2, Axis, ErrorKind, Ix1};
22

33
#[test]
4-
fn stacking() {
4+
fn concatenating() {
55
let a = arr2(&[[2., 2.], [3., 3.]]);
6-
let b = ndarray::stack(Axis(0), &[a.view(), a.view()]).unwrap();
6+
let b = ndarray::concatenate(Axis(0), &[a.view(), a.view()]).unwrap();
77
assert_eq!(b, arr2(&[[2., 2.], [3., 3.], [2., 2.], [3., 3.]]));
88

9-
let c = stack![Axis(0), a, b];
9+
let c = concatenate![Axis(0), a, b];
1010
assert_eq!(
1111
c,
1212
arr2(&[[2., 2.], [3., 3.], [2., 2.], [3., 3.], [2., 2.], [3., 3.]])
1313
);
1414

15-
let d = stack![Axis(0), a.row(0), &[9., 9.]];
15+
let d = concatenate![Axis(0), a.row(0), &[9., 9.]];
1616
assert_eq!(d, aview1(&[2., 2., 9., 9.]));
1717

18-
let res = ndarray::stack(Axis(1), &[a.view(), c.view()]);
18+
let res = ndarray::concatenate(Axis(1), &[a.view(), c.view()]);
19+
assert_eq!(res.unwrap_err().kind(), ErrorKind::IncompatibleShape);
20+
21+
let res = ndarray::concatenate(Axis(2), &[a.view(), c.view()]);
22+
assert_eq!(res.unwrap_err().kind(), ErrorKind::OutOfBounds);
23+
24+
let res: Result<Array2<f64>, _> = ndarray::concatenate(Axis(0), &[]);
25+
assert_eq!(res.unwrap_err().kind(), ErrorKind::Unsupported);
26+
}
27+
28+
#[test]
29+
fn stacking() {
30+
let a = arr2(&[[2., 2.], [3., 3.]]);
31+
let b = ndarray::stack(Axis(0), vec![a.view(), a.view()]).unwrap();
32+
assert_eq!(b, arr3(&[[[2., 2.], [3., 3.]], [[2., 2.], [3., 3.]]]));
33+
34+
let c = arr2(&[[3., 2., 3.], [2., 3., 2.]]);
35+
let res = ndarray::stack(Axis(1), vec![a.view(), c.view()]);
1936
assert_eq!(res.unwrap_err().kind(), ErrorKind::IncompatibleShape);
2037

21-
let res = ndarray::stack(Axis(2), &[a.view(), c.view()]);
38+
let res = ndarray::stack(Axis(3), vec![a.view(), a.view()]);
2239
assert_eq!(res.unwrap_err().kind(), ErrorKind::OutOfBounds);
2340

24-
let res: Result<Array2<f64>, _> = ndarray::stack(Axis(0), &[]);
41+
let res: Result<Array2<f64>, _> = ndarray::stack::<_, Ix1>(Axis(0), vec![]);
2542
assert_eq!(res.unwrap_err().kind(), ErrorKind::Unsupported);
2643
}

0 commit comments

Comments
 (0)