Skip to content

Commit da71b98

Browse files
committed
Add slice_each_axis_inplace method
1 parent a66f364 commit da71b98

File tree

2 files changed

+54
-0
lines changed

2 files changed

+54
-0
lines changed

src/impl_methods.rs

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ use crate::error::{self, ErrorKind, ShapeError};
2525
use crate::math_cell::MathCell;
2626
use crate::itertools::zip;
2727
use crate::zip::Zip;
28+
use crate::AxisDescription;
2829

2930
use crate::iter::{
3031
AxisChunksIter, AxisChunksIterMut, AxisIter, AxisIterMut, ExactChunks, ExactChunksMut,
@@ -511,6 +512,48 @@ where
511512
debug_assert!(self.pointer_is_inbounds());
512513
}
513514

515+
/// Slice the array in place, with a closure specifying the slice for each
516+
/// axis.
517+
///
518+
/// This is especially useful for code which is generic over the
519+
/// dimensionality of the array.
520+
///
521+
/// **Panics** if an index is out of bounds or step size is zero.
522+
///
523+
/// ```
524+
/// use ndarray::{s, Array2, ArrayBase, ArrayView, Data, Dimension, Slice};
525+
///
526+
/// fn view_halved_axes<S, D>(arr: &ArrayBase<S, D>) -> ArrayView<'_, S::Elem, D>
527+
/// where
528+
/// S: Data,
529+
/// D: Dimension,
530+
/// {
531+
///
532+
/// let mut view = arr.view();
533+
/// view.slice_each_axis_inplace(|ax| Slice::from(0..ax.len() / 2));
534+
/// view
535+
/// }
536+
///
537+
/// let a = Array2::<f32>::eye(8);
538+
/// let v = view_halved_axes(&a);
539+
/// assert_eq!(v, a.slice(s![..4, ..4]));
540+
/// ```
541+
pub fn slice_each_axis_inplace<F>(&mut self, mut f: F)
542+
where
543+
F: FnMut(AxisDescription) -> Slice,
544+
{
545+
(0..self.ndim()).for_each(|ax| {
546+
self.slice_axis_inplace(
547+
Axis(ax),
548+
f(AxisDescription(
549+
Axis(ax),
550+
self.dim[ax],
551+
self.strides[ax] as isize,
552+
)),
553+
)
554+
})
555+
}
556+
514557
/// Return a reference to the element at `index`, or return `None`
515558
/// if the index is out of bounds.
516559
///

src/lib.rs

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -503,6 +503,17 @@ pub type Ixs = isize;
503503
/// [`.slice_move()`]: #method.slice_move
504504
/// [`.slice_collapse()`]: #method.slice_collapse
505505
///
506+
/// When slicing arrays with generic dimensionality, creating an instance of
507+
/// [`&SliceInfo`] to pass to the multi-axis slicing methods like [`.slice()`]
508+
/// is awkward. In these cases, it's usually more convenient to create a view
509+
/// and then slice individual axes of the view using methods such as
510+
/// [`.slice_axis_inplace()`], [`.collapse_axis()`], and
511+
/// [`.slice_each_axis_inplace()`].
512+
///
513+
/// [`.slice_axis_inplace()`]: #method.slice_axis_inplace
514+
/// [`.collapse_axis()`]: #method.collapse_axis
515+
/// [`.slice_each_axis_inplace()`]: #method.slice_each_axis_inplace
516+
///
506517
/// It's possible to take multiple simultaneous *mutable* slices with
507518
/// [`.multi_slice_mut()`] or (for [`ArrayViewMut`] only)
508519
/// [`.multi_slice_move()`].

0 commit comments

Comments
 (0)