bs_trace/linalg/
vector.rs

1use crate::internal::uninit;
2use crate::linalg::matrix::Matrix;
3use num_traits::{ConstOne, ConstZero, Float, One, Zero};
4use std::iter::Sum;
5use std::ops::{
6    Add, AddAssign, Div, DivAssign, Index, IndexMut, Mul, MulAssign, Neg, Sub, SubAssign,
7};
8
9#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
10pub struct Vector<T, const N: usize> {
11    data: [T; N],
12}
13
14impl<T, const N: usize> Vector<T, N> {
15    pub const fn new(data: [T; N]) -> Self {
16        Self { data }
17    }
18
19    pub fn pure(x: T) -> Self
20    where
21        T: Clone,
22    {
23        let mut data = uninit::new_uninit_array::<T, N>();
24        for datum in data.iter_mut() {
25            datum.write(x.clone());
26        }
27        let data = unsafe { uninit::array_assume_init(data) };
28        Self::new(data)
29    }
30
31    pub const fn const_pure(x: T) -> Self
32    where
33        T: Copy,
34    {
35        Self::new([x; N])
36    }
37
38    pub fn gen<F>(f: F) -> Self
39    where
40        F: Fn(usize) -> T,
41    {
42        let mut data = uninit::new_uninit_array::<T, N>();
43        for (i, datum) in data.iter_mut().enumerate() {
44            datum.write(f(i));
45        }
46        let data = unsafe { uninit::array_assume_init(data) };
47        Self::new(data)
48    }
49
50    pub fn dot(self, rhs: Vector<T, N>) -> T
51    where
52        T: Mul<T, Output = T> + Sum,
53    {
54        (self * rhs).into_iter().sum()
55    }
56
57    pub fn mag_sq(self) -> T
58    where
59        T: Mul<T, Output = T> + Sum + Clone,
60    {
61        self.clone().dot(self)
62    }
63
64    pub fn mag(self) -> T
65    where
66        T: Float + Sum,
67    {
68        self.mag_sq().sqrt()
69    }
70
71    pub fn normalised(self) -> Self
72    where
73        T: Float + Sum,
74    {
75        self / self.mag()
76    }
77
78    pub fn map<F, U>(self, f: F) -> Vector<U, N>
79    where
80        F: Fn(T) -> U,
81    {
82        let mut new_data = uninit::new_uninit_array::<U, N>();
83        for (datum, lhs) in new_data.iter_mut().zip(self.data.into_iter()) {
84            datum.write(f(lhs));
85        }
86        let new_data = unsafe { uninit::array_assume_init(new_data) };
87        Vector::new(new_data)
88    }
89
90    pub fn map_mut<F>(&mut self, f: F)
91    where
92        F: Fn(&mut T),
93    {
94        for datum in self.data.iter_mut() {
95            f(datum);
96        }
97    }
98
99    pub fn apply<F, U, V>(self, f: F, rhs: Vector<U, N>) -> Vector<V, N>
100    where
101        F: Fn(T, U) -> V,
102    {
103        let mut new_data = uninit::new_uninit_array::<V, N>();
104        for (datum, (lhs, rhs)) in new_data
105            .iter_mut()
106            .zip(self.data.into_iter().zip(rhs.data.into_iter()))
107        {
108            datum.write(f(lhs, rhs));
109        }
110        let new_data = unsafe { uninit::array_assume_init(new_data) };
111        Vector::new(new_data)
112    }
113
114    pub fn apply_mut<F, U>(&mut self, f: F, rhs: Vector<U, N>)
115    where
116        F: Fn(&mut T, U),
117    {
118        for (datum, rhs) in self.data.iter_mut().zip(rhs.data.into_iter()) {
119            f(datum, rhs);
120        }
121    }
122
123    pub(super) fn into_matrix(self) -> Matrix<T, N, 1> {
124        Matrix::new([self])
125    }
126
127    pub fn into_iter(self) -> <Self as IntoIterator>::IntoIter {
128        <Self as IntoIterator>::into_iter(self)
129    }
130
131    pub fn iter(&self) -> <&Self as IntoIterator>::IntoIter {
132        <&Self as IntoIterator>::into_iter(self)
133    }
134
135    pub fn iter_mut(&mut self) -> <&mut Self as IntoIterator>::IntoIter {
136        <&mut Self as IntoIterator>::into_iter(self)
137    }
138}
139
140impl<T> Vector<T, 1> {
141    pub const fn new1(x: T) -> Self {
142        Vector::new([x])
143    }
144
145    pub const fn x(&self) -> &T {
146        &self.data[0]
147    }
148}
149
150impl<T> Vector<T, 2> {
151    pub const fn new2(x: T, y: T) -> Self {
152        Vector::new([x, y])
153    }
154
155    pub const fn x(&self) -> &T {
156        &self.data[0]
157    }
158
159    pub const fn y(&self) -> &T {
160        &self.data[1]
161    }
162}
163
164impl<T> Vector<T, 3> {
165    pub const fn new3(x: T, y: T, z: T) -> Self {
166        Vector::new([x, y, z])
167    }
168
169    pub const fn x(&self) -> &T {
170        &self.data[0]
171    }
172
173    pub const fn y(&self) -> &T {
174        &self.data[1]
175    }
176
177    pub const fn z(&self) -> &T {
178        &self.data[2]
179    }
180
181    pub fn cross(self, rhs: Vector<T, 3>) -> Self
182    where
183        T: Mul<T, Output = T> + Add<T, Output = T> + Sub<T, Output = T> + Neg<Output = T> + Clone,
184    {
185        let x = self.y().clone() * rhs.z().clone() - self.z().clone() * rhs.y().clone();
186        let y = self.z().clone() * rhs.x().clone() - self.x().clone() * rhs.z().clone();
187        let z = self.x().clone() * rhs.y().clone() - self.y().clone() * rhs.x().clone();
188        Vector::new3(x, y, z)
189    }
190}
191
192impl<T> Vector<T, 4> {
193    pub const fn new4(x: T, y: T, z: T, w: T) -> Self {
194        Vector::new([x, y, z, w])
195    }
196
197    pub const fn x(&self) -> &T {
198        &self.data[0]
199    }
200
201    pub const fn y(&self) -> &T {
202        &self.data[1]
203    }
204
205    pub const fn z(&self) -> &T {
206        &self.data[2]
207    }
208
209    pub const fn w(&self) -> &T {
210        &self.data[3]
211    }
212}
213
214impl<T, const N: usize> Neg for Vector<T, N>
215where
216    T: Neg<Output = T>,
217{
218    type Output = Vector<T, N>;
219
220    fn neg(self) -> Self::Output {
221        self.map(T::neg)
222    }
223}
224
225impl<T, const N: usize> Add<Vector<T, N>> for Vector<T, N>
226where
227    T: Add<T, Output = T>,
228{
229    type Output = Vector<T, N>;
230
231    fn add(self, rhs: Vector<T, N>) -> Self::Output {
232        self.apply(T::add, rhs)
233    }
234}
235
236impl<T, const N: usize> AddAssign<Vector<T, N>> for Vector<T, N>
237where
238    T: AddAssign<T>,
239{
240    fn add_assign(&mut self, rhs: Vector<T, N>) {
241        self.apply_mut(T::add_assign, rhs);
242    }
243}
244
245impl<T, const N: usize> Sub<Vector<T, N>> for Vector<T, N>
246where
247    T: Sub<T, Output = T>,
248{
249    type Output = Vector<T, N>;
250
251    fn sub(self, rhs: Vector<T, N>) -> Self::Output {
252        self.apply(T::sub, rhs)
253    }
254}
255
256impl<T, const N: usize> SubAssign<Vector<T, N>> for Vector<T, N>
257where
258    T: SubAssign<T>,
259{
260    fn sub_assign(&mut self, rhs: Vector<T, N>) {
261        self.apply_mut(T::sub_assign, rhs);
262    }
263}
264
265impl<T, const N: usize> Mul<Vector<T, N>> for Vector<T, N>
266where
267    T: Mul<T, Output = T>,
268{
269    type Output = Vector<T, N>;
270
271    fn mul(self, rhs: Vector<T, N>) -> Self::Output {
272        self.apply(T::mul, rhs)
273    }
274}
275
276impl<T, const N: usize> Mul<T> for Vector<T, N>
277where
278    T: Mul<T, Output = T> + Clone,
279{
280    type Output = Vector<T, N>;
281
282    fn mul(self, rhs: T) -> Self::Output {
283        self.map(|x| x * rhs.clone())
284    }
285}
286
287impl<T, const N: usize> MulAssign<Vector<T, N>> for Vector<T, N>
288where
289    T: MulAssign,
290{
291    fn mul_assign(&mut self, rhs: Vector<T, N>) {
292        self.apply_mut(T::mul_assign, rhs);
293    }
294}
295
296impl<T, const N: usize> MulAssign<T> for Vector<T, N>
297where
298    T: MulAssign + Clone,
299{
300    fn mul_assign(&mut self, rhs: T) {
301        self.map_mut(|x| *x *= rhs.clone());
302    }
303}
304
305impl<T, const N: usize> Div<Vector<T, N>> for Vector<T, N>
306where
307    T: Div<T, Output = T>,
308{
309    type Output = Vector<T, N>;
310
311    fn div(self, rhs: Vector<T, N>) -> Self::Output {
312        self.apply(T::div, rhs)
313    }
314}
315
316impl<T, const N: usize> Div<T> for Vector<T, N>
317where
318    T: Div<T, Output = T> + Clone,
319{
320    type Output = Vector<T, N>;
321
322    fn div(self, rhs: T) -> Self::Output {
323        self.map(|x| x / rhs.clone())
324    }
325}
326
327impl<T, const N: usize> DivAssign<Vector<T, N>> for Vector<T, N>
328where
329    T: DivAssign<T>,
330{
331    fn div_assign(&mut self, rhs: Vector<T, N>) {
332        self.apply_mut(T::div_assign, rhs);
333    }
334}
335
336impl<T, const N: usize> DivAssign<T> for Vector<T, N>
337where
338    T: DivAssign<T> + Clone,
339{
340    fn div_assign(&mut self, rhs: T) {
341        self.map_mut(|x| *x /= rhs.clone());
342    }
343}
344
345impl<T, const N: usize> Zero for Vector<T, N>
346where
347    T: Zero + Clone,
348{
349    fn zero() -> Self {
350        Self::pure(T::zero())
351    }
352
353    fn is_zero(&self) -> bool {
354        self.data.iter().all(T::is_zero)
355    }
356}
357
358impl<T, const N: usize> ConstZero for Vector<T, N>
359where
360    T: ConstZero + Copy,
361{
362    const ZERO: Self = Self::const_pure(T::ZERO);
363}
364
365impl<T, const N: usize> One for Vector<T, N>
366where
367    T: One + Clone,
368{
369    fn one() -> Self {
370        Vector::pure(T::one())
371    }
372}
373
374impl<T, const N: usize> ConstOne for Vector<T, N>
375where
376    T: ConstOne + Copy,
377{
378    const ONE: Self = Vector::const_pure(T::ONE);
379}
380
381impl<T, const N: usize> Sum for Vector<T, N>
382where
383    T: Add<T, Output = T> + Zero + Clone,
384{
385    fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
386        iter.fold(Vector::zero(), |acc, x| acc + x)
387    }
388}
389
390impl<T, const N: usize> IntoIterator for Vector<T, N> {
391    type Item = T;
392    type IntoIter = std::array::IntoIter<T, N>;
393
394    fn into_iter(self) -> Self::IntoIter {
395        self.data.into_iter()
396    }
397}
398
399impl<'a, T, const N: usize> IntoIterator for &'a Vector<T, N> {
400    type Item = &'a T;
401    type IntoIter = std::slice::Iter<'a, T>;
402
403    fn into_iter(self) -> Self::IntoIter {
404        self.data.iter()
405    }
406}
407
408impl<'a, T, const N: usize> IntoIterator for &'a mut Vector<T, N> {
409    type Item = &'a mut T;
410    type IntoIter = std::slice::IterMut<'a, T>;
411
412    fn into_iter(self) -> Self::IntoIter {
413        self.data.iter_mut()
414    }
415}
416
417impl<T, const N: usize> Index<usize> for Vector<T, N> {
418    type Output = T;
419
420    fn index(&self, index: usize) -> &Self::Output {
421        &self.data[index]
422    }
423}
424
425impl<T, const N: usize> IndexMut<usize> for Vector<T, N> {
426    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
427        &mut self.data[index]
428    }
429}
430
431impl<T, const N: usize> From<Matrix<T, N, 1>> for Vector<T, N> {
432    fn from(value: Matrix<T, N, 1>) -> Self {
433        value.into_vector()
434    }
435}
436
437#[cfg(test)]
438mod tests {
439    use super::*;
440
441    #[test]
442    fn check_add() {
443        let lhs = Vector::<isize, 4>::new([2, 3, 5, 7]);
444        let rhs = Vector::<isize, 4>::new([11, 13, 17, 19]);
445        let expected_sum = Vector::<isize, 4>::new([13, 16, 22, 26]);
446        let actual_sum = lhs + rhs;
447        assert_eq!(actual_sum, expected_sum);
448    }
449
450    #[test]
451    fn check_add_assign() {
452        let mut lhs = Vector::<isize, 5>::new([23, 29, 31, 37, 41]);
453        let rhs = Vector::<isize, 5>::new([43, 47, 53, 59, 61]);
454        let expected_sum = Vector::<isize, 5>::new([66, 76, 84, 96, 102]);
455        lhs += rhs;
456        assert_eq!(lhs, expected_sum);
457    }
458
459    #[test]
460    fn check_sub() {
461        let lhs = Vector::<isize, 3>::new([67, 71, 73]);
462        let rhs = Vector::<isize, 3>::new([79, 83, 89]);
463        let expected_diff = Vector::<isize, 3>::new([-12, -12, -16]);
464        let actual_diff = lhs - rhs;
465        assert_eq!(actual_diff, expected_diff);
466    }
467
468    #[test]
469    fn check_sub_assign() {
470        let mut lhs = Vector::<isize, 3>::new([97, 101, 103]);
471        let rhs = Vector::<isize, 3>::new([107, 109, 113]);
472        let expected_diff = Vector::<isize, 3>::new([-10, -8, -10]);
473        lhs -= rhs;
474        assert_eq!(lhs, expected_diff);
475    }
476
477    #[test]
478    fn check_mul_componentwise() {
479        let lhs = Vector::<isize, 2>::new([127, 131]);
480        let rhs = Vector::<isize, 2>::new([137, 139]);
481        let expected_prod = Vector::<isize, 2>::new([17399, 18209]);
482        let actual_prod = lhs * rhs;
483        assert_eq!(actual_prod, expected_prod);
484    }
485
486    #[test]
487    fn check_mul_scalar() {
488        let lhs = Vector::<isize, 4>::new([151, 157, 163, 167]);
489        let rhs = 149isize;
490        let expected_prod = Vector::<isize, 4>::new([22499, 23393, 24287, 24883]);
491        let actual_prod = lhs * rhs;
492        assert_eq!(actual_prod, expected_prod);
493    }
494
495    #[test]
496    fn check_mul_assign_componentwise() {
497        let mut lhs = Vector::<isize, 3>::new([173, 179, 181]);
498        let rhs = Vector::<isize, 3>::new([191, 193, 197]);
499        let expected_prod = Vector::<isize, 3>::new([33043, 34547, 35657]);
500        lhs *= rhs;
501        assert_eq!(lhs, expected_prod);
502    }
503
504    #[test]
505    fn check_mul_assign_scalar() {
506        let mut lhs = Vector::<isize, 2>::new([199, 211]);
507        let rhs = 223isize;
508        let expected_prod = Vector::<isize, 2>::new([44377, 47053]);
509        lhs *= rhs;
510        assert_eq!(lhs, expected_prod);
511    }
512
513    #[test]
514    fn check_div_componentwise() {
515        let lhs = Vector::<isize, 3>::new([8, 24, 33]);
516        let rhs = Vector::<isize, 3>::new([2, 8, 12]);
517        let expected_quot = Vector::<isize, 3>::new([4, 3, 2]);
518        let actual_quot = lhs / rhs;
519        assert_eq!(actual_quot, expected_quot);
520    }
521
522    #[test]
523    fn check_div_scalar() {
524        let lhs = Vector::<isize, 4>::new([11, 14, 15, 16]);
525        let rhs = 3isize;
526        let expected_quot = Vector::<isize, 4>::new([3, 4, 5, 5]);
527        let actual_quot = lhs / rhs;
528        assert_eq!(actual_quot, expected_quot);
529    }
530
531    #[test]
532    fn check_div_assign_componentwise() {
533        let mut lhs = Vector::<isize, 3>::new([44, 55, 66]);
534        let rhs = Vector::<isize, 3>::new([5, 6, 7]);
535        let expected_quot = Vector::<isize, 3>::new([8, 9, 9]);
536        lhs /= rhs;
537        assert_eq!(lhs, expected_quot);
538    }
539
540    #[test]
541    fn check_div_assign_scalar() {
542        let mut lhs = Vector::<isize, 2>::new([400, 450]);
543        let rhs = 50isize;
544        let expected_quot = Vector::<isize, 2>::new([8, 9]);
545        lhs /= rhs;
546        assert_eq!(lhs, expected_quot);
547    }
548}