Fri, 18 Nov 2022 10:29:50 +0200
Improvements and minor fixes to bisection tree refinement.
0 | 1 | /*! |
5 | 2 | Array containers that support vector space operations on floats. |
3 | For working with small vectors in $ℝ^2$ or $ℝ^3$. | |
0 | 4 | */ |
5 | ||
6 | use std::ops::{Add,Sub,AddAssign,SubAssign,Mul,Div,MulAssign,DivAssign,Neg,Index,IndexMut}; | |
7 | use std::slice::{Iter,IterMut}; | |
8 | use crate::types::{Float,Num,SignedNum}; | |
9 | use crate::maputil::{FixedLength,FixedLengthMut,map1,map2,map1_mut,map2_mut}; | |
5 | 10 | use crate::euclidean::*; |
0 | 11 | use crate::norms::*; |
12 | use crate::linops::AXPY; | |
13 | use serde::ser::{Serialize, Serializer, SerializeSeq}; | |
14 | ||
5 | 15 | /// A container type for (short) `N`-dimensional vectors of element type `F`. |
16 | /// | |
17 | /// Supports basic operations of an [`Euclidean`] space, several [`Norm`]s, and | |
18 | /// fused [`AXPY`] operations, among others. | |
0 | 19 | #[derive(Copy,Clone,Debug,PartialEq,Eq)] |
5 | 20 | pub struct Loc<F, const N : usize>( |
21 | /// An array of the elements of the vector | |
22 | pub [F; N] | |
23 | ); | |
0 | 24 | |
25 | // Need to manually implement as [F; N] serialisation is provided only for some N. | |
26 | impl<F, const N : usize> Serialize for Loc<F, N> | |
27 | where | |
28 | F: Serialize, | |
29 | { | |
30 | fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error> | |
31 | where | |
32 | S: Serializer, | |
33 | { | |
34 | let mut seq = serializer.serialize_seq(Some(N))?; | |
35 | for e in self.iter() { | |
36 | seq.serialize_element(e)?; | |
37 | } | |
38 | seq.end() | |
39 | } | |
40 | } | |
41 | ||
42 | impl<F, const N : usize> Loc<F, N> { | |
5 | 43 | /// Creates a new `Loc` vector from an array. |
0 | 44 | #[inline] |
45 | pub fn new(arr : [F; N]) -> Self { | |
46 | Loc(arr) | |
47 | } | |
5 | 48 | |
49 | /// Returns an iterator over the elements of the vector | |
0 | 50 | #[inline] |
51 | pub fn iter(&self) -> Iter<'_, F> { | |
52 | self.0.iter() | |
53 | } | |
54 | ||
5 | 55 | /// Returns an iterator over mutable references to the elements of the vector |
0 | 56 | #[inline] |
57 | pub fn iter_mut(&mut self) -> IterMut<'_, F> { | |
58 | self.0.iter_mut() | |
59 | } | |
60 | } | |
61 | ||
62 | impl<F : Copy, const N : usize> Loc<F, N> { | |
5 | 63 | /// Maps `g` over the elements of the vector, returning a new [`Loc`] vector |
0 | 64 | #[inline] |
5 | 65 | pub fn map<H>(&self, g : impl Fn(F) -> H) -> Loc<H, N> { |
0 | 66 | Loc::new(map1(self, |u| g(*u))) |
67 | } | |
68 | ||
5 | 69 | /// Maps `g` over pairs of elements of two vectors, retuning a new one. |
0 | 70 | #[inline] |
5 | 71 | pub fn map2<H>(&self, other : &Self, g : impl Fn(F, F) -> H) -> Loc<H, N> { |
0 | 72 | Loc::new(map2(self, other, |u, v| g(*u, *v))) |
73 | } | |
74 | ||
5 | 75 | /// Maps `g` over mutable references to elements of the vector. |
0 | 76 | #[inline] |
5 | 77 | pub fn map_mut(&mut self, g : impl Fn(&mut F)) { |
0 | 78 | map1_mut(self, g) |
79 | } | |
80 | ||
5 | 81 | /// Maps `g` over pairs of mutable references to elements of `self, and elements |
82 | /// of `other` vector. | |
0 | 83 | #[inline] |
5 | 84 | pub fn map2_mut(&mut self, other : &Self, g : impl Fn(&mut F, F)) { |
0 | 85 | map2_mut(self, other, |u, v| g(u, *v)) |
86 | } | |
87 | ||
5 | 88 | /// Maps `g` over the elements of `self` and returns the product of the results. |
0 | 89 | #[inline] |
5 | 90 | pub fn product_map<A : Num>(&self, g : impl Fn(F) -> A) -> A { |
0 | 91 | match N { |
5 | 92 | 1 => g(unsafe { *self.0.get_unchecked(0) }), |
93 | 2 => g(unsafe { *self.0.get_unchecked(0) }) * | |
94 | g(unsafe { *self.0.get_unchecked(1) }), | |
95 | 3 => g(unsafe { *self.0.get_unchecked(0) }) * | |
96 | g(unsafe { *self.0.get_unchecked(1) }) * | |
97 | g(unsafe { *self.0.get_unchecked(2) }), | |
98 | _ => self.iter().fold(A::ONE, |m, &x| m * g(x)) | |
0 | 99 | } |
100 | } | |
101 | } | |
102 | ||
5 | 103 | /// Construct a [`Loc`]. |
104 | /// | |
105 | /// Use as | |
106 | /// ``` | |
107 | /// # use alg_tools::loc::Loc; | |
108 | /// # use alg_tools::loc; | |
109 | /// let x = loc![1.0, 2.0]; | |
110 | /// ``` | |
0 | 111 | #[macro_export] |
112 | macro_rules! loc { | |
113 | ($($x:expr),+ $(,)?) => { Loc::new([$($x),+]) } | |
114 | } | |
115 | ||
116 | ||
117 | impl<F, const N : usize> From<[F; N]> for Loc<F, N> { | |
118 | #[inline] | |
119 | fn from(other: [F; N]) -> Loc<F, N> { | |
120 | Loc(other) | |
121 | } | |
122 | } | |
123 | ||
124 | /*impl<F : Copy, const N : usize> From<&[F; N]> for Loc<F, N> { | |
125 | #[inline] | |
126 | fn from(other: &[F; N]) -> Loc<F, N> { | |
127 | Loc(*other) | |
128 | } | |
129 | }*/ | |
130 | ||
131 | impl<F> From<F> for Loc<F, 1> { | |
132 | #[inline] | |
133 | fn from(other: F) -> Loc<F, 1> { | |
134 | Loc([other]) | |
135 | } | |
136 | } | |
137 | ||
138 | impl<F, const N : usize> From<Loc<F, N>> for [F; N] { | |
139 | #[inline] | |
140 | fn from(other : Loc<F, N>) -> [F; N] { | |
141 | other.0 | |
142 | } | |
143 | } | |
144 | ||
145 | /*impl<F : Copy, const N : usize> From<&Loc<F, N>> for [F; N] { | |
146 | #[inline] | |
147 | fn from(other : &Loc<F, N>) -> [F; N] { | |
148 | other.0 | |
149 | } | |
150 | }*/ | |
151 | ||
152 | ||
153 | impl<F, const N : usize> IntoIterator for Loc<F, N> { | |
154 | type Item = <[F; N] as IntoIterator>::Item; | |
155 | type IntoIter = <[F; N] as IntoIterator>::IntoIter; | |
156 | ||
157 | #[inline] | |
158 | fn into_iter(self) -> Self::IntoIter { | |
159 | self.0.into_iter() | |
160 | } | |
161 | } | |
162 | ||
163 | // Indexing | |
164 | ||
165 | impl<F, Ix, const N : usize> Index<Ix> for Loc<F,N> | |
166 | where [F; N] : Index<Ix> { | |
167 | type Output = <[F; N] as Index<Ix>>::Output; | |
168 | ||
169 | #[inline] | |
170 | fn index(&self, ix : Ix) -> &Self::Output { | |
171 | self.0.index(ix) | |
172 | } | |
173 | } | |
174 | ||
175 | impl<F, Ix, const N : usize> IndexMut<Ix> for Loc<F,N> | |
176 | where [F; N] : IndexMut<Ix> { | |
177 | #[inline] | |
178 | fn index_mut(&mut self, ix : Ix) -> &mut Self::Output { | |
179 | self.0.index_mut(ix) | |
180 | } | |
181 | } | |
182 | ||
183 | // Arithmetic | |
184 | ||
185 | macro_rules! make_binop { | |
186 | ($trait:ident, $fn:ident, $trait_assign:ident, $fn_assign:ident) => { | |
187 | impl<F : Num, const N : usize> $trait<Loc<F,N>> for Loc<F, N> { | |
188 | type Output = Loc<F, N>; | |
189 | #[inline] | |
190 | fn $fn(mut self, other : Loc<F, N>) -> Self::Output { | |
191 | self.$fn_assign(other); | |
192 | self | |
193 | } | |
194 | } | |
195 | ||
196 | impl<'a, F : Num, const N : usize> $trait<&'a Loc<F,N>> for Loc<F, N> { | |
197 | type Output = Loc<F, N>; | |
198 | #[inline] | |
199 | fn $fn(mut self, other : &'a Loc<F, N>) -> Self::Output { | |
200 | self.$fn_assign(other); | |
201 | self | |
202 | } | |
203 | } | |
204 | ||
205 | impl<'b, F : Num, const N : usize> $trait<Loc<F,N>> for &'b Loc<F, N> { | |
206 | type Output = Loc<F, N>; | |
207 | #[inline] | |
208 | fn $fn(self, other : Loc<F, N>) -> Self::Output { | |
209 | self.map2(&other, |a, b| a.$fn(b)) | |
210 | } | |
211 | } | |
212 | ||
213 | impl<'a, 'b, F : Num, const N : usize> $trait<&'a Loc<F,N>> for &'b Loc<F, N> { | |
214 | type Output = Loc<F, N>; | |
215 | #[inline] | |
216 | fn $fn(self, other : &'a Loc<F, N>) -> Self::Output { | |
217 | self.map2(other, |a, b| a.$fn(b)) | |
218 | } | |
219 | } | |
220 | ||
221 | impl<F : Num, const N : usize> $trait_assign<Loc<F,N>> for Loc<F, N> { | |
222 | #[inline] | |
223 | fn $fn_assign(&mut self, other : Loc<F, N>) { | |
224 | self.map2_mut(&other, |a, b| a.$fn_assign(b)) | |
225 | } | |
226 | } | |
227 | ||
228 | impl<'a, F : Num, const N : usize> $trait_assign<&'a Loc<F,N>> for Loc<F, N> { | |
229 | #[inline] | |
230 | fn $fn_assign(&mut self, other : &'a Loc<F, N>) { | |
231 | self.map2_mut(other, |a, b| a.$fn_assign(b)) | |
232 | } | |
233 | } | |
234 | } | |
235 | } | |
236 | ||
237 | make_binop!(Add, add, AddAssign, add_assign); | |
238 | make_binop!(Sub, sub, SubAssign, sub_assign); | |
239 | ||
240 | macro_rules! make_scalarop_rhs { | |
241 | ($trait:ident, $fn:ident, $trait_assign:ident, $fn_assign:ident) => { | |
242 | impl<F : Num, const N : usize> $trait<F> for Loc<F, N> { | |
243 | type Output = Loc<F, N>; | |
244 | #[inline] | |
245 | fn $fn(self, b : F) -> Self::Output { | |
246 | self.map(|a| a.$fn(b)) | |
247 | } | |
248 | } | |
249 | ||
250 | impl<'a, F : Num, const N : usize> $trait<&'a F> for Loc<F, N> { | |
251 | type Output = Loc<F, N>; | |
252 | #[inline] | |
253 | fn $fn(self, b : &'a F) -> Self::Output { | |
254 | self.map(|a| a.$fn(*b)) | |
255 | } | |
256 | } | |
257 | ||
258 | impl<'b, F : Num, const N : usize> $trait<F> for &'b Loc<F, N> { | |
259 | type Output = Loc<F, N>; | |
260 | #[inline] | |
261 | fn $fn(self, b : F) -> Self::Output { | |
262 | self.map(|a| a.$fn(b)) | |
263 | } | |
264 | } | |
265 | ||
266 | impl<'a, 'b, F : Float, const N : usize> $trait<&'a F> for &'b Loc<F, N> { | |
267 | type Output = Loc<F, N>; | |
268 | #[inline] | |
269 | fn $fn(self, b : &'a F) -> Self::Output { | |
270 | self.map(|a| a.$fn(*b)) | |
271 | } | |
272 | } | |
273 | ||
274 | impl<F : Num, const N : usize> $trait_assign<F> for Loc<F, N> { | |
275 | #[inline] | |
276 | fn $fn_assign(&mut self, b : F) { | |
277 | self.map_mut(|a| a.$fn_assign(b)); | |
278 | } | |
279 | } | |
280 | ||
281 | impl<'a, F : Num, const N : usize> $trait_assign<&'a F> for Loc<F, N> { | |
282 | #[inline] | |
283 | fn $fn_assign(&mut self, b : &'a F) { | |
284 | self.map_mut(|a| a.$fn_assign(*b)); | |
285 | } | |
286 | } | |
287 | } | |
288 | } | |
289 | ||
290 | ||
291 | make_scalarop_rhs!(Mul, mul, MulAssign, mul_assign); | |
292 | make_scalarop_rhs!(Div, div, DivAssign, div_assign); | |
293 | ||
294 | macro_rules! make_unaryop { | |
295 | ($trait:ident, $fn:ident) => { | |
296 | impl<F : SignedNum, const N : usize> $trait for Loc<F, N> { | |
297 | type Output = Loc<F, N>; | |
298 | #[inline] | |
299 | fn $fn(mut self) -> Self::Output { | |
300 | self.map_mut(|a| *a = (*a).$fn()); | |
301 | self | |
302 | } | |
303 | } | |
304 | ||
305 | impl<'a, F : SignedNum, const N : usize> $trait for &'a Loc<F, N> { | |
306 | type Output = Loc<F, N>; | |
307 | #[inline] | |
308 | fn $fn(self) -> Self::Output { | |
309 | self.map(|a| a.$fn()) | |
310 | } | |
311 | } | |
312 | } | |
313 | } | |
314 | ||
315 | make_unaryop!(Neg, neg); | |
316 | ||
317 | macro_rules! make_scalarop_lhs { | |
318 | ($trait:ident, $fn:ident; $($f:ident)+) => { $( | |
319 | impl<const N : usize> $trait<Loc<$f,N>> for $f { | |
320 | type Output = Loc<$f, N>; | |
321 | #[inline] | |
322 | fn $fn(self, v : Loc<$f,N>) -> Self::Output { | |
323 | v.map(|b| self.$fn(b)) | |
324 | } | |
325 | } | |
326 | ||
327 | impl<'a, const N : usize> $trait<&'a Loc<$f,N>> for $f { | |
328 | type Output = Loc<$f, N>; | |
329 | #[inline] | |
330 | fn $fn(self, v : &'a Loc<$f,N>) -> Self::Output { | |
331 | v.map(|b| self.$fn(b)) | |
332 | } | |
333 | } | |
334 | ||
335 | impl<'b, const N : usize> $trait<Loc<$f,N>> for &'b $f { | |
336 | type Output = Loc<$f, N>; | |
337 | #[inline] | |
338 | fn $fn(self, v : Loc<$f,N>) -> Self::Output { | |
339 | v.map(|b| self.$fn(b)) | |
340 | } | |
341 | } | |
342 | ||
343 | impl<'a, 'b, const N : usize> $trait<&'a Loc<$f,N>> for &'b $f { | |
344 | type Output = Loc<$f, N>; | |
345 | #[inline] | |
346 | fn $fn(self, v : &'a Loc<$f, N>) -> Self::Output { | |
347 | v.map(|b| self.$fn(b)) | |
348 | } | |
349 | } | |
350 | )+ } | |
351 | } | |
352 | ||
353 | make_scalarop_lhs!(Mul, mul; f32 f64 i8 i16 i32 i64 isize u8 u16 u32 u64 usize); | |
354 | make_scalarop_lhs!(Div, div; f32 f64 i8 i16 i32 i64 isize u8 u16 u32 u64 usize); | |
355 | ||
356 | // Norms | |
357 | ||
358 | macro_rules! domination { | |
359 | ($norm:ident, $dominates:ident) => { | |
360 | impl<F : Float, const N : usize> Dominated<F, $dominates, Loc<F, N>> for $norm { | |
361 | #[inline] | |
362 | fn norm_factor(&self, _p : $dominates) -> F { | |
363 | F::ONE | |
364 | } | |
365 | #[inline] | |
366 | fn from_norm(&self, p_norm : F, _p : $dominates) -> F { | |
367 | p_norm | |
368 | } | |
369 | } | |
370 | }; | |
371 | ($norm:ident, $dominates:ident, $fn:path) => { | |
372 | impl<F : Float, const N : usize> Dominated<F, $dominates, Loc<F, N>> for $norm { | |
373 | #[inline] | |
374 | fn norm_factor(&self, _p : $dominates) -> F { | |
375 | $fn(F::cast_from(N)) | |
376 | } | |
377 | } | |
378 | }; | |
379 | } | |
380 | ||
381 | domination!(L1, L1); | |
382 | domination!(L2, L2); | |
383 | domination!(Linfinity, Linfinity); | |
384 | ||
385 | domination!(L1, L2, F::sqrt); | |
386 | domination!(L2, Linfinity, F::sqrt); | |
387 | domination!(L1, Linfinity, std::convert::identity); | |
388 | ||
389 | domination!(Linfinity, L1); | |
390 | domination!(Linfinity, L2); | |
391 | domination!(L2, L1); | |
392 | ||
393 | impl<F : Num,const N : usize> Dot<Loc<F, N>,F> for Loc<F, N> { | |
394 | /// This implementation is not stabilised as it's meant to be used for very small vectors. | |
395 | /// Use [`nalgebra`] for larger vectors. | |
396 | #[inline] | |
397 | fn dot(&self, other : &Loc<F, N>) -> F { | |
398 | self.0.iter() | |
399 | .zip(other.0.iter()) | |
400 | .fold(F::ZERO, |m, (&v, &w)| m + v * w) | |
401 | } | |
402 | } | |
403 | ||
404 | impl<F : Float,const N : usize> Euclidean<F> for Loc<F, N> { | |
405 | type Output = Self; | |
406 | ||
407 | #[inline] | |
408 | fn similar_origin(&self) -> Self { | |
409 | Self::ORIGIN | |
410 | } | |
411 | ||
412 | /// This implementation is not stabilised as it's meant to be used for very small vectors. | |
413 | /// Use [`nalgebra`] for larger vectors. | |
414 | #[inline] | |
415 | fn norm2_squared(&self) -> F { | |
416 | self.iter().fold(F::ZERO, |m, &v| m + v * v) | |
417 | } | |
418 | ||
419 | fn dist2_squared(&self, other : &Self) -> F { | |
420 | self.iter() | |
421 | .zip(other.iter()) | |
422 | .fold(F::ZERO, |m, (&v, &w)| { let d = v - w; m + d * d }) | |
423 | } | |
424 | ||
425 | #[inline] | |
426 | fn norm2(&self) -> F { | |
427 | // Optimisation for N==1 that avoids squaring and square rooting. | |
428 | if N==1 { | |
429 | unsafe { self.0.get_unchecked(0) }.abs() | |
430 | } else { | |
431 | self.norm2_squared().sqrt() | |
432 | } | |
433 | } | |
434 | ||
435 | #[inline] | |
436 | fn dist2(&self, other : &Self) -> F { | |
437 | // Optimisation for N==1 that avoids squaring and square rooting. | |
438 | if N==1 { | |
439 | unsafe { *self.0.get_unchecked(0) - *other.0.get_unchecked(0) }.abs() | |
440 | } else { | |
441 | self.dist2_squared(other).sqrt() | |
442 | } | |
443 | } | |
444 | } | |
445 | ||
446 | impl<F : Num, const N : usize> Loc<F, N> { | |
447 | pub const ORIGIN : Self = Loc([F::ZERO; N]); | |
448 | } | |
449 | ||
450 | impl<F : Float,const N : usize> StaticEuclidean<F> for Loc<F, N> { | |
451 | #[inline] | |
452 | fn origin() -> Self { | |
453 | Self::ORIGIN | |
454 | } | |
455 | } | |
456 | ||
457 | impl<F : Float, const N : usize> Norm<F, L2> for Loc<F, N> { | |
458 | #[inline] | |
459 | fn norm(&self, _ : L2) -> F { self.norm2() } | |
460 | } | |
461 | ||
462 | impl<F : Float, const N : usize> Dist<F, L2> for Loc<F, N> { | |
463 | #[inline] | |
464 | fn dist(&self, other : &Self, _ : L2) -> F { self.dist2(other) } | |
465 | } | |
466 | ||
467 | impl<F : Float, const N : usize> Norm<F, L1> for Loc<F, N> { | |
468 | /// This implementation is not stabilised as it's meant to be used for very small vectors. | |
469 | /// Use [`nalgebra`] for larger vectors. | |
470 | #[inline] | |
471 | fn norm(&self, _ : L1) -> F { | |
472 | self.iter().fold(F::ZERO, |m, v| m + v.abs()) | |
473 | } | |
474 | } | |
475 | ||
476 | impl<F : Float, const N : usize> Dist<F, L1> for Loc<F, N> { | |
477 | #[inline] | |
478 | fn dist(&self, other : &Self, _ : L1) -> F { | |
479 | self.iter() | |
480 | .zip(other.iter()) | |
481 | .fold(F::ZERO, |m, (&v, &w)| m + (v-w).abs() ) | |
482 | } | |
483 | } | |
484 | ||
485 | impl<F : Float, const N : usize> Projection<F, Linfinity> for Loc<F, N> { | |
486 | #[inline] | |
487 | fn proj_ball_mut(&mut self, ρ : F, _ : Linfinity) { | |
488 | self.iter_mut().for_each(|v| *v = num_traits::clamp(*v, -ρ, ρ)) | |
489 | } | |
490 | } | |
491 | ||
492 | impl<F : Float, const N : usize> Norm<F, Linfinity> for Loc<F, N> { | |
493 | /// This implementation is not stabilised as it's meant to be used for very small vectors. | |
494 | /// Use [`nalgebra`] for larger vectors. | |
495 | #[inline] | |
496 | fn norm(&self, _ : Linfinity) -> F { | |
497 | self.iter().fold(F::ZERO, |m, v| m.max(v.abs())) | |
498 | } | |
499 | } | |
500 | ||
501 | impl<F : Float, const N : usize> Dist<F, Linfinity> for Loc<F, N> { | |
502 | #[inline] | |
503 | fn dist(&self, other : &Self, _ : Linfinity) -> F { | |
504 | self.iter() | |
505 | .zip(other.iter()) | |
506 | .fold(F::ZERO, |m, (&v, &w)| m.max((v-w).abs())) | |
507 | } | |
508 | } | |
509 | ||
510 | ||
511 | // Misc. | |
512 | ||
513 | impl<A, const N : usize> FixedLength<N> for Loc<A,N> { | |
514 | type Iter = std::array::IntoIter<A, N>; | |
515 | type Elem = A; | |
516 | #[inline] | |
517 | fn fl_iter(self) -> Self::Iter { | |
518 | self.into_iter() | |
519 | } | |
520 | } | |
521 | ||
522 | impl<A, const N : usize> FixedLengthMut<N> for Loc<A,N> { | |
523 | type IterMut<'a> = std::slice::IterMut<'a, A> where A : 'a; | |
524 | #[inline] | |
525 | fn fl_iter_mut(&mut self) -> Self::IterMut<'_> { | |
526 | self.iter_mut() | |
527 | } | |
528 | } | |
529 | ||
530 | impl<'a, A, const N : usize> FixedLength<N> for &'a Loc<A,N> { | |
531 | type Iter = std::slice::Iter<'a, A>; | |
532 | type Elem = &'a A; | |
533 | #[inline] | |
534 | fn fl_iter(self) -> Self::Iter { | |
535 | self.iter() | |
536 | } | |
537 | } | |
538 | ||
539 | impl<F : Num, const N : usize> AXPY<F, Loc<F, N>> for Loc<F, N> { | |
540 | ||
541 | #[inline] | |
542 | fn axpy(&mut self, α : F, x : &Loc<F, N>, β : F) { | |
543 | if β == F::ZERO { | |
544 | map2_mut(self, x, |yi, xi| { *yi = α * (*xi) }) | |
545 | } else { | |
546 | map2_mut(self, x, |yi, xi| { *yi = β * (*yi) + α * (*xi) }) | |
547 | } | |
548 | } | |
549 | ||
550 | #[inline] | |
551 | fn copy_from(&mut self, x : &Loc<F, N>) { | |
552 | map2_mut(self, x, |yi, xi| *yi = *xi ) | |
553 | } | |
554 | } |