Mon, 24 Oct 2022 09:41:43 +0300
Allow step closure of AlgIterators to indicate succesfull termination or failure.
0 | 1 | /// Production of equally spaced nodes within intervals. |
2 | ||
3 | use crate::types::*; | |
4 | use crate::loc::Loc; | |
5 | use crate::sets::Cube; | |
6 | use crate::iter::{RestartableIterator, StatefulIterator}; | |
7 | use crate::maputil::{map2, map4}; | |
8 | use serde::Serialize; | |
9 | ||
10 | // TODO: rewrite this using crate::sets::Cube. | |
11 | ||
12 | /// The interval `[start, end]` divided into `count` nodes. | |
13 | /// Implementes `IntoIterator` to iterate over the nodes of the `LinSpace`. | |
14 | /// It is similar to [`num::range_step_inclusive`], but [restartable][RestartableIterator] and | |
15 | /// parametrised by the number of nodes instead of a step. This way it can be ensured that the | |
16 | /// last item produced is equal to `end`. To use other types than `float` and `f32`, create | |
17 | /// the `LinSpace` struct directly. | |
18 | #[derive(Clone,Copy,Debug,Serialize,Eq,PartialEq)] | |
19 | pub struct LinSpace<F, I> { | |
20 | pub start : F, | |
21 | pub end : F, | |
22 | pub count : I, | |
23 | } | |
24 | ||
25 | #[allow(type_alias_bounds)] // Need it to access F::CompatibleSize. | |
26 | pub type LinGrid<F : Float, const N : usize> = LinSpace<Loc<F, N>, [usize; N]>; | |
27 | ||
28 | /// Creates a [`LinSpace`] on the real line. | |
29 | pub fn linspace<F : Float>(start : F, end : F, count : usize) -> LinSpace<F, usize> { | |
30 | LinSpace{ start : start, end : end, count : count } | |
31 | } | |
32 | ||
33 | /// Create a multi-dimensional linear grid. | |
34 | /// | |
35 | /// The first and last point in each dimension are the boundaries of the corresponding | |
36 | /// dimensions of `cube`, and there are `count` nodes along each dimension. | |
37 | pub fn lingrid<F : Float, const N : usize>( | |
38 | cube : &Cube<F, N>, | |
39 | count : &[usize; N] | |
40 | ) -> LinSpace<Loc<F, N>, [usize; N]> { | |
41 | LinSpace{ start : cube.span_start(), end : cube.span_end(), count : *count } | |
42 | } | |
43 | ||
44 | /// Create a multi-dimensional linear grid with centered nodes. | |
45 | /// | |
46 | /// There are `count` along each dimension and each node has equally-sized subcube surrounding it | |
47 | /// inside `cube`. Thus, if $w_i$ is the width of the cube along dimension $i$, and $n_i$ the number | |
48 | /// of nodes, the width of the subcube along this dimension is $h_i = w_i/(n_i+1)$, and the first | |
49 | /// and last nodes are at a distance $h_i/2$ from the closest boundary. | |
50 | pub fn lingrid_centered<F : Float, const N : usize>( | |
51 | cube : &Cube<F, N>, | |
52 | count : &[usize; N] | |
53 | ) -> LinSpace<Loc<F, N>, [usize; N]> { | |
54 | let h_div_2 = map2(cube.width(), count, |w, &n| w / F::cast_from(2 * (n + 1))); | |
55 | let span_start = map2(cube.span_start(), &h_div_2, |a, &t| a + t).into(); | |
56 | let span_end = map2(cube.span_end(), &h_div_2, |b, &t| b - t).into(); | |
57 | LinSpace{ start : span_start, end : span_end, count : *count } | |
58 | } | |
59 | ||
60 | ||
61 | /// Iterator over a `LinSpace`. | |
62 | #[derive(Clone, Debug)] | |
63 | pub struct LinSpaceIterator<F, I> { | |
64 | lingrid : LinSpace<F,I>, | |
65 | current : Option<I>, | |
66 | } | |
67 | ||
68 | pub trait Grid<F, I> { | |
69 | fn entry_linear_unchecked(&self, i : usize) -> F; | |
70 | //fn entry(&self, i : I) -> Option<F>; | |
71 | fn entry_unchecked(&self, i : &I) -> F; | |
72 | } | |
73 | ||
74 | pub trait GridIteration<F, I> { | |
75 | fn next_index(&mut self) -> Option<I>; | |
76 | } | |
77 | ||
78 | impl<F : Float + CastFrom<I>, I : Unsigned> Grid<F, I> for LinSpace<F, I> { | |
79 | /*fn entry(&self, i : I) -> Option<F> { | |
80 | if i < self.count { | |
81 | Some(self.entry_unchecked(i)) | |
82 | } else { | |
83 | None | |
84 | } | |
85 | }*/ | |
86 | ||
87 | #[inline] | |
88 | fn entry_linear_unchecked(&self, i : usize) -> F { | |
89 | self.entry_unchecked(&I::cast_from(i)) | |
90 | } | |
91 | ||
92 | #[inline] | |
93 | fn entry_unchecked(&self, i : &I) -> F { | |
94 | let idx = F::cast_from(*i); | |
95 | let scale = F::cast_from(self.count-I::ONE); | |
96 | self.start + (self.end-self.start)*idx/scale | |
97 | } | |
98 | } | |
99 | ||
100 | impl<F : Float + CastFrom<I>, I : Unsigned> GridIteration<F, I> | |
101 | for LinSpaceIterator<F, I> { | |
102 | #[inline] | |
103 | fn next_index(&mut self) -> Option<I> { | |
104 | match self.current { | |
105 | None if I::ZERO < self.lingrid.count => { self.current = Some(I::ZERO); self.current } | |
106 | Some(v) if v+I::ONE < self.lingrid.count => { self.current = Some(v+I::ONE); self.current } | |
107 | _ => { None } | |
108 | } | |
109 | } | |
110 | } | |
111 | ||
112 | impl<F : Float + CastFrom<I>, I : Unsigned, const N : usize> Grid<Loc<F,N>, [I; N]> | |
113 | for LinSpace<Loc<F,N>, [I; N]> { | |
114 | #[inline] | |
115 | fn entry_linear_unchecked(&self, i_ : usize) -> Loc<F, N> { | |
116 | let mut i = I::cast_from(i_); | |
117 | let mut tmp = [I::ZERO; N]; | |
118 | for k in 0..N { | |
119 | tmp[k] = i % self.count[k]; | |
120 | i /= self.count[k]; | |
121 | } | |
122 | self.entry_unchecked(&tmp) | |
123 | } | |
124 | ||
125 | #[inline] | |
126 | fn entry_unchecked(&self, i : &[I; N]) -> Loc<F, N> { | |
127 | let LinSpace{ start, end, count } = self; | |
128 | map4(i, start, end, count, |&ik, &sk, &ek, &ck| { | |
129 | let idx = F::cast_from(ik); | |
130 | let scale = F::cast_from(ck-I::ONE); | |
131 | sk + (ek - sk) * idx / scale | |
132 | }).into() | |
133 | } | |
134 | } | |
135 | ||
136 | impl<F : Float + CastFrom<I>, I : Unsigned, const N : usize> GridIteration<Loc<F,N>, [I; N]> | |
137 | for LinSpaceIterator<Loc<F,N>, [I; N]> { | |
138 | ||
139 | #[inline] | |
140 | fn next_index(&mut self) -> Option<[I; N]> { | |
141 | match self.current { | |
142 | None if self.lingrid.count.iter().all(|v| I::ZERO < *v) => { | |
143 | self.current = Some([I::ZERO; N]); | |
144 | self.current | |
145 | }, | |
146 | Some(ref mut v) => { | |
147 | for k in 0..N { | |
148 | let a = v[k] + I::ONE; | |
149 | if a < self.lingrid.count[k] { | |
150 | v[k] = a; | |
151 | return self.current | |
152 | } else { | |
153 | v[k] = I::ZERO; | |
154 | } | |
155 | } | |
156 | None | |
157 | }, | |
158 | _ => None | |
159 | } | |
160 | } | |
161 | } | |
162 | ||
163 | impl<F, I> IntoIterator for LinSpace<F,I> | |
164 | where LinSpace<F, I> : Grid<F, I>, | |
165 | LinSpaceIterator<F, I> : GridIteration<F, I> { | |
166 | type Item = F; | |
167 | type IntoIter = LinSpaceIterator<F,I>; | |
168 | ||
169 | #[inline] | |
170 | fn into_iter(self) -> Self::IntoIter { | |
171 | LinSpaceIterator { lingrid : self, current : None } | |
172 | } | |
173 | } | |
174 | ||
175 | impl<F, I> Iterator for LinSpaceIterator<F,I> | |
176 | where LinSpace<F, I> : Grid<F, I>, | |
177 | LinSpaceIterator<F, I> : GridIteration<F, I> { | |
178 | type Item = F; | |
179 | #[inline] | |
180 | fn next(&mut self) -> Option<F> { | |
181 | self.next_index().map(|v| self.lingrid.entry_unchecked(&v)) | |
182 | } | |
183 | } | |
184 | ||
185 | impl<F, I> StatefulIterator for LinSpaceIterator<F,I> | |
186 | where LinSpace<F, I> : Grid<F, I>, | |
187 | LinSpaceIterator<F, I> : GridIteration<F, I> { | |
188 | #[inline] | |
189 | fn current(&self) -> Option<F> { | |
190 | self.current.as_ref().map(|c| self.lingrid.entry_unchecked(c)) | |
191 | } | |
192 | } | |
193 | ||
194 | ||
195 | impl<F, I> RestartableIterator for LinSpaceIterator<F,I> | |
196 | where LinSpace<F, I> : Grid<F, I>, | |
197 | LinSpaceIterator<F, I> : GridIteration<F, I> { | |
198 | #[inline] | |
199 | fn restart(&mut self) -> Option<F> { | |
200 | self.current = None; | |
201 | self.next() | |
202 | } | |
203 | } |