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