ndarray/dimension/
axes.rs

1use crate::{Axis, Dimension, 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 D: Dimension
6{
7    Axes {
8        dim: d,
9        strides,
10        start: 0,
11        end: d.ndim(),
12    }
13}
14
15/// An iterator over the length and stride of each axis of an array.
16///
17/// This iterator is created from the array method
18/// [`.axes()`](crate::ArrayBase::axes).
19///
20/// Iterator element type is [`AxisDescription`].
21///
22/// # Examples
23///
24/// ```
25/// use ndarray::Array3;
26/// use ndarray::Axis;
27///
28/// let a = Array3::<f32>::zeros((3, 5, 4));
29///
30/// // find the largest axis in the array
31/// // check the axis index and its length
32///
33/// let largest_axis = a.axes()
34///                     .max_by_key(|ax| ax.len)
35///                     .unwrap();
36/// assert_eq!(largest_axis.axis, Axis(1));
37/// assert_eq!(largest_axis.len, 5);
38/// ```
39#[derive(Debug)]
40pub struct Axes<'a, D>
41{
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{
52    /// Axis identifier (index)
53    pub axis: Axis,
54    /// Length in count of elements of the current axis
55    pub len: usize,
56    /// Stride in count of elements of the current axis
57    pub stride: isize,
58}
59
60copy_and_clone!(AxisDescription);
61copy_and_clone!(['a, D] Axes<'a, D>);
62
63impl<'a, D> Iterator for Axes<'a, D>
64where D: Dimension
65{
66    /// Description of the axis, its length and its stride.
67    type Item = AxisDescription;
68
69    fn next(&mut self) -> Option<Self::Item>
70    {
71        if self.start < self.end {
72            let i = self.start.post_inc();
73            Some(AxisDescription {
74                axis: Axis(i),
75                len: self.dim[i],
76                stride: self.strides[i] as Ixs,
77            })
78        } else {
79            None
80        }
81    }
82
83    fn fold<B, F>(self, init: B, f: F) -> B
84    where F: FnMut(B, AxisDescription) -> B
85    {
86        (self.start..self.end)
87            .map(move |i| AxisDescription {
88                axis: Axis(i),
89                len: self.dim[i],
90                stride: self.strides[i] as isize,
91            })
92            .fold(init, f)
93    }
94
95    fn size_hint(&self) -> (usize, Option<usize>)
96    {
97        let len = self.end - self.start;
98        (len, Some(len))
99    }
100}
101
102impl<'a, D> DoubleEndedIterator for Axes<'a, D>
103where D: Dimension
104{
105    fn next_back(&mut self) -> Option<Self::Item>
106    {
107        if self.start < self.end {
108            let i = self.end.pre_dec();
109            Some(AxisDescription {
110                axis: Axis(i),
111                len: self.dim[i],
112                stride: self.strides[i] as Ixs,
113            })
114        } else {
115            None
116        }
117    }
118}
119
120trait IncOps: Copy
121{
122    fn post_inc(&mut self) -> Self;
123    fn pre_dec(&mut self) -> Self;
124}
125
126impl IncOps for usize
127{
128    #[inline(always)]
129    fn post_inc(&mut self) -> Self
130    {
131        let x = *self;
132        *self += 1;
133        x
134    }
135    #[inline(always)]
136    fn pre_dec(&mut self) -> Self
137    {
138        *self -= 1;
139        *self
140    }
141}
ndarray/dimension/axes.rs

ndarray/dimension/
axes.rs
