ndarray/dimension/
axes.rs

1use crate::{Axis, Dimension, Ix, Ixs};
2
3/// Create a new Axes iterator
4pub(crate) fn axes_of<'a, D>(d: &'a D, strides: &'a D) -> Axes<'a, D>
5where
6    D: Dimension,
7{
8    Axes {
9        dim: d,
10        strides,
11        start: 0,
12        end: d.ndim(),
13    }
14}
15
16/// An iterator over the length and stride of each axis of an array.
17///
18/// This iterator is created from the array method
19/// [`.axes()`](crate::ArrayBase::axes).
20///
21/// Iterator element type is [`AxisDescription`].
22///
23/// # Examples
24///
25/// ```
26/// use ndarray::Array3;
27/// use ndarray::Axis;
28///
29/// let a = Array3::<f32>::zeros((3, 5, 4));
30///
31/// // find the largest axis in the array
32/// // check the axis index and its length
33///
34/// let largest_axis = a.axes()
35///                     .max_by_key(|ax| ax.len)
36///                     .unwrap();
37/// assert_eq!(largest_axis.axis, Axis(1));
38/// assert_eq!(largest_axis.len, 5);
39/// ```
40#[derive(Debug)]
41pub struct Axes<'a, D> {
42    dim: &'a D,
43    strides: &'a D,
44    start: usize,
45    end: usize,
46}
47
48/// Description of the axis, its length and its stride.
49#[derive(Debug)]
50pub struct AxisDescription {
51    pub axis: Axis,
52    pub len: usize,
53    pub stride: isize,
54}
55
56copy_and_clone!(AxisDescription);
57
58// AxisDescription can't really be empty
59// https://github.com/rust-ndarray/ndarray/pull/642#discussion_r296051702
60#[allow(clippy::len_without_is_empty)]
61impl AxisDescription {
62    /// Return axis
63    #[deprecated(note = "Use .axis field instead", since = "0.15.0")]
64    #[inline(always)]
65    pub fn axis(self) -> Axis {
66        self.axis
67    }
68    /// Return length
69    #[deprecated(note = "Use .len field instead", since = "0.15.0")]
70    #[inline(always)]
71    pub fn len(self) -> Ix {
72        self.len
73    }
74    /// Return stride
75    #[deprecated(note = "Use .stride field instead", since = "0.15.0")]
76    #[inline(always)]
77    pub fn stride(self) -> Ixs {
78        self.stride
79    }
80}
81
82copy_and_clone!(['a, D] Axes<'a, D>);
83
84impl<'a, D> Iterator for Axes<'a, D>
85where
86    D: Dimension,
87{
88    /// Description of the axis, its length and its stride.
89    type Item = AxisDescription;
90
91    fn next(&mut self) -> Option<Self::Item> {
92        if self.start < self.end {
93            let i = self.start.post_inc();
94            Some(AxisDescription {
95                axis: Axis(i),
96                len: self.dim[i],
97                stride: self.strides[i] as Ixs,
98            })
99        } else {
100            None
101        }
102    }
103
104    fn fold<B, F>(self, init: B, f: F) -> B
105    where
106        F: FnMut(B, AxisDescription) -> B,
107    {
108        (self.start..self.end)
109            .map(move |i| AxisDescription {
110                axis: Axis(i),
111                len: self.dim[i],
112                stride: self.strides[i] as isize,
113            })
114            .fold(init, f)
115    }
116
117    fn size_hint(&self) -> (usize, Option<usize>) {
118        let len = self.end - self.start;
119        (len, Some(len))
120    }
121}
122
123impl<'a, D> DoubleEndedIterator for Axes<'a, D>
124where
125    D: Dimension,
126{
127    fn next_back(&mut self) -> Option<Self::Item> {
128        if self.start < self.end {
129            let i = self.end.pre_dec();
130            Some(AxisDescription {
131                axis: Axis(i),
132                len: self.dim[i],
133                stride: self.strides[i] as Ixs,
134            })
135        } else {
136            None
137        }
138    }
139}
140
141trait IncOps: Copy {
142    fn post_inc(&mut self) -> Self;
143    fn post_dec(&mut self) -> Self;
144    fn pre_dec(&mut self) -> Self;
145}
146
147impl IncOps for usize {
148    #[inline(always)]
149    fn post_inc(&mut self) -> Self {
150        let x = *self;
151        *self += 1;
152        x
153    }
154    #[inline(always)]
155    fn post_dec(&mut self) -> Self {
156        let x = *self;
157        *self -= 1;
158        x
159    }
160    #[inline(always)]
161    fn pre_dec(&mut self) -> Self {
162        *self -= 1;
163        *self
164    }
165}
ndarray/dimension/axes.rs

ndarray/dimension/
axes.rs
