Skip to content

Commit

Permalink
Add slice_each_axis_inplace method
Browse files Browse the repository at this point in the history
  • Loading branch information
jturner314 committed Feb 5, 2021
1 parent a66f364 commit da71b98
Show file tree
Hide file tree
Showing 2 changed files with 54 additions and 0 deletions.
43 changes: 43 additions & 0 deletions src/impl_methods.rs
Expand Up @@ -25,6 +25,7 @@ use crate::error::{self, ErrorKind, ShapeError};
use crate::math_cell::MathCell;
use crate::itertools::zip;
use crate::zip::Zip;
use crate::AxisDescription;

use crate::iter::{
AxisChunksIter, AxisChunksIterMut, AxisIter, AxisIterMut, ExactChunks, ExactChunksMut,
Expand Down Expand Up @@ -511,6 +512,48 @@ where
debug_assert!(self.pointer_is_inbounds());
}

/// Slice the array in place, with a closure specifying the slice for each
/// axis.
///
/// This is especially useful for code which is generic over the
/// dimensionality of the array.
///
/// **Panics** if an index is out of bounds or step size is zero.
///
/// ```
/// use ndarray::{s, Array2, ArrayBase, ArrayView, Data, Dimension, Slice};
///
/// fn view_halved_axes<S, D>(arr: &ArrayBase<S, D>) -> ArrayView<'_, S::Elem, D>
/// where
/// S: Data,
/// D: Dimension,
/// {
///
/// let mut view = arr.view();
/// view.slice_each_axis_inplace(|ax| Slice::from(0..ax.len() / 2));
/// view
/// }
///
/// let a = Array2::<f32>::eye(8);
/// let v = view_halved_axes(&a);
/// assert_eq!(v, a.slice(s![..4, ..4]));
/// ```
pub fn slice_each_axis_inplace<F>(&mut self, mut f: F)
where
F: FnMut(AxisDescription) -> Slice,
{
(0..self.ndim()).for_each(|ax| {
self.slice_axis_inplace(
Axis(ax),
f(AxisDescription(
Axis(ax),
self.dim[ax],
self.strides[ax] as isize,
)),
)
})
}

/// Return a reference to the element at `index`, or return `None`
/// if the index is out of bounds.
///
Expand Down
11 changes: 11 additions & 0 deletions src/lib.rs
Expand Up @@ -503,6 +503,17 @@ pub type Ixs = isize;
/// [`.slice_move()`]: #method.slice_move
/// [`.slice_collapse()`]: #method.slice_collapse
///
/// When slicing arrays with generic dimensionality, creating an instance of
/// [`&SliceInfo`] to pass to the multi-axis slicing methods like [`.slice()`]
/// is awkward. In these cases, it's usually more convenient to create a view
/// and then slice individual axes of the view using methods such as
/// [`.slice_axis_inplace()`], [`.collapse_axis()`], and
/// [`.slice_each_axis_inplace()`].
///
/// [`.slice_axis_inplace()`]: #method.slice_axis_inplace
/// [`.collapse_axis()`]: #method.collapse_axis
/// [`.slice_each_axis_inplace()`]: #method.slice_each_axis_inplace
///
/// It's possible to take multiple simultaneous *mutable* slices with
/// [`.multi_slice_mut()`] or (for [`ArrayViewMut`] only)
/// [`.multi_slice_move()`].
Expand Down

0 comments on commit da71b98

Please sign in to comment.