Tue, 06 Dec 2022 08:29:13 +0200
cargo-d alias generation with a cargo-d script
5 | 1 | /*! |
2 | Linear grids. | |
3 | ||
8
4e09b7829b51
Multithreaded bisection tree operations
Tuomo Valkonen <tuomov@iki.fi>
parents:
5
diff
changeset
|
4 | These are multi-dimensional intervals $\prod\_{i=1}^N [a\_i, b\_i]$ divided along each dimension |
4e09b7829b51
Multithreaded bisection tree operations
Tuomo Valkonen <tuomov@iki.fi>
parents:
5
diff
changeset
|
5 | into $n\_i$ equally-spaced nodes, with $a\_i$ the first node and $b\_i$ the last node along each |
5 | 6 | dimension. |
7 | ||
8 | The [`LinSpace`]s provided by this module are similar to [`num::range_step_inclusive`], but as an | |
9 | iterator they are [restartable][RestartableIterator] and parametrised by the number of nodes | |
8
4e09b7829b51
Multithreaded bisection tree operations
Tuomo Valkonen <tuomov@iki.fi>
parents:
5
diff
changeset
|
10 | instead of a step. This way it can be ensured that $a\_i$ and $b\_i$ are the last and the first |
4e09b7829b51
Multithreaded bisection tree operations
Tuomo Valkonen <tuomov@iki.fi>
parents:
5
diff
changeset
|
11 | node. |
5 | 12 | |
13 | The starting points for the use of this module are the [`linspace`], [`lingrid`], and | |
14 | [`lingrid_centered`] functions. They return a [`LinSpace`]s that implements [`IntoIterator`] for | |
15 | iteration over the grid. Additional utility functions are in the [`Grid`] trait. | |
16 | */ | |
0 | 17 | |
18 | use crate::types::*; | |
19 | use crate::loc::Loc; | |
20 | use crate::sets::Cube; | |
21 | use crate::iter::{RestartableIterator, StatefulIterator}; | |
22 | use crate::maputil::{map2, map4}; | |
5 | 23 | use serde::{Serialize, Deserialize}; |
0 | 24 | |
25 | // TODO: rewrite this using crate::sets::Cube. | |
26 | ||
5 | 27 | /// An abstraction of possibly multi-dimensional linear grids. |
28 | /// | |
29 | /// `U` is typically a `F` for a `Float` `F` for one-dimensional grids created by `linspace`, | |
30 | /// or [`Loc`]`<F, N>` for multi-dimensional grids created by `lingrid`. | |
31 | /// In the first case `count` of nodes is `usize`, and in the second case `[usize; N]`. | |
32 | #[derive(Clone, Copy, Debug, Serialize, Deserialize, Eq, PartialEq)] | |
33 | pub struct LinSpace<U, I> { | |
34 | pub start : U, | |
35 | pub end : U, | |
0 | 36 | pub count : I, |
37 | } | |
38 | ||
5 | 39 | /// A `N`-dimensional interval divided into an indicated number of equally-spaced nodes along |
40 | /// each dimension. | |
0 | 41 | #[allow(type_alias_bounds)] // Need it to access F::CompatibleSize. |
42 | pub type LinGrid<F : Float, const N : usize> = LinSpace<Loc<F, N>, [usize; N]>; | |
43 | ||
44 | /// Creates a [`LinSpace`] on the real line. | |
45 | pub fn linspace<F : Float>(start : F, end : F, count : usize) -> LinSpace<F, usize> { | |
46 | LinSpace{ start : start, end : end, count : count } | |
47 | } | |
48 | ||
5 | 49 | /// Creates a multi-dimensional linear grid. |
0 | 50 | /// |
51 | /// The first and last point in each dimension are the boundaries of the corresponding | |
52 | /// dimensions of `cube`, and there are `count` nodes along each dimension. | |
53 | pub fn lingrid<F : Float, const N : usize>( | |
54 | cube : &Cube<F, N>, | |
55 | count : &[usize; N] | |
56 | ) -> LinSpace<Loc<F, N>, [usize; N]> { | |
57 | LinSpace{ start : cube.span_start(), end : cube.span_end(), count : *count } | |
58 | } | |
59 | ||
60 | /// Create a multi-dimensional linear grid with centered nodes. | |
61 | /// | |
62 | /// There are `count` along each dimension and each node has equally-sized subcube surrounding it | |
8
4e09b7829b51
Multithreaded bisection tree operations
Tuomo Valkonen <tuomov@iki.fi>
parents:
5
diff
changeset
|
63 | /// inside `cube`. Thus, if $w\_i$ is the width of the cube along dimension $i$, and $n_i$ the number |
4e09b7829b51
Multithreaded bisection tree operations
Tuomo Valkonen <tuomov@iki.fi>
parents:
5
diff
changeset
|
64 | /// of nodes, the width of the subcube along this dimension is $h_i = w\_i/(n\_i+1)$, and the first |
4e09b7829b51
Multithreaded bisection tree operations
Tuomo Valkonen <tuomov@iki.fi>
parents:
5
diff
changeset
|
65 | /// and last nodes are at a distance $h\_i/2$ from the closest boundary. |
0 | 66 | pub fn lingrid_centered<F : Float, const N : usize>( |
67 | cube : &Cube<F, N>, | |
68 | count : &[usize; N] | |
69 | ) -> LinSpace<Loc<F, N>, [usize; N]> { | |
70 | let h_div_2 = map2(cube.width(), count, |w, &n| w / F::cast_from(2 * (n + 1))); | |
71 | let span_start = map2(cube.span_start(), &h_div_2, |a, &t| a + t).into(); | |
72 | let span_end = map2(cube.span_end(), &h_div_2, |b, &t| b - t).into(); | |
73 | LinSpace{ start : span_start, end : span_end, count : *count } | |
74 | } | |
75 | ||
76 | ||
77 | /// Iterator over a `LinSpace`. | |
78 | #[derive(Clone, Debug)] | |
79 | pub struct LinSpaceIterator<F, I> { | |
80 | lingrid : LinSpace<F,I>, | |
81 | current : Option<I>, | |
82 | } | |
83 | ||
5 | 84 | /// Abstraction of a linear grid over space `U` with multi-dimensional index set `I`. |
85 | pub trait Grid<U, I> { | |
86 | /// Converts a linear index `i` into a grid point. | |
87 | fn entry_linear_unchecked(&self, i : usize) -> U; | |
88 | // Converts a multi-dimensional index `i` into a grid point. | |
89 | fn entry_unchecked(&self, i : &I) -> U; | |
90 | ||
91 | // fn entry(&self, i : I) -> Option<F> | |
0 | 92 | } |
93 | ||
5 | 94 | /// Helper trait for iteration of [`Grid`]s. |
0 | 95 | pub trait GridIteration<F, I> { |
5 | 96 | /// Returns the next multi-dimensional index (not yet converted into grid point). |
0 | 97 | fn next_index(&mut self) -> Option<I>; |
98 | } | |
99 | ||
100 | impl<F : Float + CastFrom<I>, I : Unsigned> Grid<F, I> for LinSpace<F, I> { | |
101 | /*fn entry(&self, i : I) -> Option<F> { | |
102 | if i < self.count { | |
103 | Some(self.entry_unchecked(i)) | |
104 | } else { | |
105 | None | |
106 | } | |
107 | }*/ | |
108 | ||
109 | #[inline] | |
110 | fn entry_linear_unchecked(&self, i : usize) -> F { | |
111 | self.entry_unchecked(&I::cast_from(i)) | |
112 | } | |
113 | ||
114 | #[inline] | |
115 | fn entry_unchecked(&self, i : &I) -> F { | |
116 | let idx = F::cast_from(*i); | |
117 | let scale = F::cast_from(self.count-I::ONE); | |
118 | self.start + (self.end-self.start)*idx/scale | |
119 | } | |
120 | } | |
121 | ||
122 | impl<F : Float + CastFrom<I>, I : Unsigned> GridIteration<F, I> | |
123 | for LinSpaceIterator<F, I> { | |
124 | #[inline] | |
125 | fn next_index(&mut self) -> Option<I> { | |
126 | match self.current { | |
5 | 127 | None if I::ZERO < self.lingrid.count |
128 | => { self.current = Some(I::ZERO); self.current } | |
129 | Some(v) if v+I::ONE < self.lingrid.count | |
130 | => { self.current = Some(v+I::ONE); self.current } | |
131 | _ | |
132 | => { None } | |
0 | 133 | } |
134 | } | |
135 | } | |
136 | ||
137 | impl<F : Float + CastFrom<I>, I : Unsigned, const N : usize> Grid<Loc<F,N>, [I; N]> | |
138 | for LinSpace<Loc<F,N>, [I; N]> { | |
139 | #[inline] | |
140 | fn entry_linear_unchecked(&self, i_ : usize) -> Loc<F, N> { | |
141 | let mut i = I::cast_from(i_); | |
142 | let mut tmp = [I::ZERO; N]; | |
143 | for k in 0..N { | |
144 | tmp[k] = i % self.count[k]; | |
145 | i /= self.count[k]; | |
146 | } | |
147 | self.entry_unchecked(&tmp) | |
148 | } | |
149 | ||
150 | #[inline] | |
151 | fn entry_unchecked(&self, i : &[I; N]) -> Loc<F, N> { | |
152 | let LinSpace{ start, end, count } = self; | |
153 | map4(i, start, end, count, |&ik, &sk, &ek, &ck| { | |
154 | let idx = F::cast_from(ik); | |
155 | let scale = F::cast_from(ck-I::ONE); | |
156 | sk + (ek - sk) * idx / scale | |
157 | }).into() | |
158 | } | |
159 | } | |
160 | ||
161 | impl<F : Float + CastFrom<I>, I : Unsigned, const N : usize> GridIteration<Loc<F,N>, [I; N]> | |
162 | for LinSpaceIterator<Loc<F,N>, [I; N]> { | |
163 | ||
164 | #[inline] | |
165 | fn next_index(&mut self) -> Option<[I; N]> { | |
166 | match self.current { | |
167 | None if self.lingrid.count.iter().all(|v| I::ZERO < *v) => { | |
168 | self.current = Some([I::ZERO; N]); | |
169 | self.current | |
170 | }, | |
171 | Some(ref mut v) => { | |
172 | for k in 0..N { | |
173 | let a = v[k] + I::ONE; | |
174 | if a < self.lingrid.count[k] { | |
175 | v[k] = a; | |
176 | return self.current | |
177 | } else { | |
178 | v[k] = I::ZERO; | |
179 | } | |
180 | } | |
181 | None | |
182 | }, | |
183 | _ => None | |
184 | } | |
185 | } | |
186 | } | |
187 | ||
188 | impl<F, I> IntoIterator for LinSpace<F,I> | |
189 | where LinSpace<F, I> : Grid<F, I>, | |
190 | LinSpaceIterator<F, I> : GridIteration<F, I> { | |
191 | type Item = F; | |
192 | type IntoIter = LinSpaceIterator<F,I>; | |
193 | ||
194 | #[inline] | |
195 | fn into_iter(self) -> Self::IntoIter { | |
196 | LinSpaceIterator { lingrid : self, current : None } | |
197 | } | |
198 | } | |
199 | ||
200 | impl<F, I> Iterator for LinSpaceIterator<F,I> | |
201 | where LinSpace<F, I> : Grid<F, I>, | |
202 | LinSpaceIterator<F, I> : GridIteration<F, I> { | |
203 | type Item = F; | |
204 | #[inline] | |
205 | fn next(&mut self) -> Option<F> { | |
206 | self.next_index().map(|v| self.lingrid.entry_unchecked(&v)) | |
207 | } | |
208 | } | |
209 | ||
210 | impl<F, I> StatefulIterator for LinSpaceIterator<F,I> | |
211 | where LinSpace<F, I> : Grid<F, I>, | |
212 | LinSpaceIterator<F, I> : GridIteration<F, I> { | |
213 | #[inline] | |
214 | fn current(&self) -> Option<F> { | |
215 | self.current.as_ref().map(|c| self.lingrid.entry_unchecked(c)) | |
216 | } | |
217 | } | |
218 | ||
219 | ||
220 | impl<F, I> RestartableIterator for LinSpaceIterator<F,I> | |
221 | where LinSpace<F, I> : Grid<F, I>, | |
222 | LinSpaceIterator<F, I> : GridIteration<F, I> { | |
223 | #[inline] | |
224 | fn restart(&mut self) -> Option<F> { | |
225 | self.current = None; | |
226 | self.next() | |
227 | } | |
228 | } |