/rust/registry/src/index.crates.io-1949cf8c6b5b557f/ndarray-0.17.2/src/dimension/axes.rs
Line | Count | Source |
1 | | use crate::{Axis, Dimension, Ixs}; |
2 | | |
3 | | /// Create a new Axes iterator |
4 | 0 | pub(crate) fn axes_of<'a, D>(d: &'a D, strides: &'a D) -> Axes<'a, D> |
5 | 0 | where D: Dimension |
6 | | { |
7 | 0 | Axes { |
8 | 0 | dim: d, |
9 | 0 | strides, |
10 | 0 | start: 0, |
11 | 0 | end: d.ndim(), |
12 | 0 | } |
13 | 0 | } |
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)] |
40 | | pub 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)] |
50 | | pub 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 | | |
60 | | copy_and_clone!(AxisDescription); |
61 | | copy_and_clone!(['a, D] Axes<'a, D>); |
62 | | |
63 | | impl<D> Iterator for Axes<'_, D> |
64 | | where D: Dimension |
65 | | { |
66 | | /// Description of the axis, its length and its stride. |
67 | | type Item = AxisDescription; |
68 | | |
69 | 0 | fn next(&mut self) -> Option<Self::Item> |
70 | | { |
71 | 0 | if self.start < self.end { |
72 | 0 | let i = self.start.post_inc(); |
73 | 0 | Some(AxisDescription { |
74 | 0 | axis: Axis(i), |
75 | 0 | len: self.dim[i], |
76 | 0 | stride: self.strides[i] as Ixs, |
77 | 0 | }) |
78 | | } else { |
79 | 0 | None |
80 | | } |
81 | 0 | } |
82 | | |
83 | 0 | fn fold<B, F>(self, init: B, f: F) -> B |
84 | 0 | where F: FnMut(B, AxisDescription) -> B |
85 | | { |
86 | 0 | (self.start..self.end) |
87 | 0 | .map(move |i| AxisDescription { |
88 | 0 | axis: Axis(i), |
89 | 0 | len: self.dim[i], |
90 | 0 | stride: self.strides[i] as isize, |
91 | 0 | }) |
92 | 0 | .fold(init, f) |
93 | 0 | } |
94 | | |
95 | 0 | fn size_hint(&self) -> (usize, Option<usize>) |
96 | | { |
97 | 0 | let len = self.end - self.start; |
98 | 0 | (len, Some(len)) |
99 | 0 | } |
100 | | } |
101 | | |
102 | | impl<D> DoubleEndedIterator for Axes<'_, D> |
103 | | where D: Dimension |
104 | | { |
105 | 0 | fn next_back(&mut self) -> Option<Self::Item> |
106 | | { |
107 | 0 | if self.start < self.end { |
108 | 0 | let i = self.end.pre_dec(); |
109 | 0 | Some(AxisDescription { |
110 | 0 | axis: Axis(i), |
111 | 0 | len: self.dim[i], |
112 | 0 | stride: self.strides[i] as Ixs, |
113 | 0 | }) |
114 | | } else { |
115 | 0 | None |
116 | | } |
117 | 0 | } |
118 | | } |
119 | | |
120 | | trait IncOps: Copy |
121 | | { |
122 | | fn post_inc(&mut self) -> Self; |
123 | | fn pre_dec(&mut self) -> Self; |
124 | | } |
125 | | |
126 | | impl IncOps for usize |
127 | | { |
128 | | #[inline(always)] |
129 | 0 | fn post_inc(&mut self) -> Self |
130 | | { |
131 | 0 | let x = *self; |
132 | 0 | *self += 1; |
133 | 0 | x |
134 | 0 | } |
135 | | #[inline(always)] |
136 | 0 | fn pre_dec(&mut self) -> Self |
137 | | { |
138 | 0 | *self -= 1; |
139 | 0 | *self |
140 | 0 | } |
141 | | } |