1use crate::traits::{Num, Real};
22use core::ops::{Add, Div, Mul, Sub};
23
24#[derive(Clone, Copy, PartialEq, Debug)]
41pub struct Array<T, const N: usize> {
42 data: [T; N],
43}
44
45impl<T: Num + Eq, const N: usize> Eq for Array<T, N> {}
46
47impl<T: Num, const N: usize> Default for Array<T, N> {
48 fn default() -> Self {
49 Self::zeros()
50 }
51}
52
53impl<T: Num, const N: usize> Array<T, N> {
54 #[must_use]
56 #[inline]
57 pub const fn new(data: [T; N]) -> Self {
58 Self { data }
59 }
60
61 #[must_use]
69 pub fn from_fn<F: FnMut(usize) -> T>(mut f: F) -> Self {
70 let mut data = [T::zero(); N];
71 for (i, slot) in data.iter_mut().enumerate() {
72 *slot = f(i);
73 }
74 Self { data }
75 }
76
77 #[must_use]
84 pub fn zeros() -> Self {
85 Self::from_fn(|_| T::zero())
86 }
87
88 #[must_use]
90 pub fn ones() -> Self {
91 Self::from_fn(|_| T::one())
92 }
93
94 pub fn as_slice(&self) -> &[T] {
96 &self.data
97 }
98
99 pub fn as_mut_slice(&mut self) -> &mut [T] {
101 &mut self.data
102 }
103
104 #[must_use]
110 #[inline]
111 pub fn get(&self, i: usize) -> T {
112 self.data[i]
113 }
114
115 pub fn set(&mut self, i: usize, value: T) {
121 self.data[i] = value;
122 }
123
124 #[must_use]
128 #[inline]
129 pub fn try_get(&self, i: usize) -> Option<T> {
130 self.data.get(i).copied()
131 }
132
133 #[must_use]
136 pub fn sum(&self) -> T {
137 let mut sum = T::zero();
138 let mut c = T::zero();
139 for &x in &self.data {
140 let y = x - c;
141 let t = sum + y;
142 c = (t - sum) - y;
143 sum = t;
144 }
145 sum
146 }
147
148 #[must_use]
158 pub fn dot(&self, other: &Self) -> T {
159 let mut sum = T::zero();
160 let mut c = T::zero();
161 for i in 0..N {
162 let x = self.data[i] * other.data[i];
163 let y = x - c;
164 let t = sum + y;
165 c = (t - sum) - y;
166 sum = t;
167 }
168 sum
169 }
170
171 #[must_use]
173 #[inline]
174 pub fn map<U: Num, F: FnMut(T) -> U>(&self, mut f: F) -> Array<U, N> {
175 Array::from_fn(|i| f(self.data[i]))
176 }
177
178 #[must_use]
180 #[inline]
181 pub fn zip_with<U: Num, V: Num, F: FnMut(T, U) -> V>(
182 &self,
183 other: &Array<U, N>,
184 mut f: F,
185 ) -> Array<V, N> {
186 Array::from_fn(|i| f(self.data[i], other.data[i]))
187 }
188
189 #[must_use]
191 #[inline]
192 pub fn component_mul(&self, other: &Self) -> Self {
193 self.zip_with(other, |x, y| x * y)
194 }
195
196 #[must_use]
198 #[inline]
199 pub fn component_div(&self, other: &Self) -> Self
200 where
201 T: Div<Output = T>,
202 {
203 self.zip_with(other, |x, y| x / y)
204 }
205
206 #[must_use]
214 pub fn cumulative_sum(&self) -> Self {
215 let mut acc = T::zero();
216 let mut out = [T::zero(); N];
217 for (i, slot) in out.iter_mut().enumerate() {
218 acc = acc + self.data[i];
219 *slot = acc;
220 }
221 Self { data: out }
222 }
223
224 #[must_use]
226 pub fn reversed(&self) -> Self {
227 let mut out = [T::zero(); N];
228 for (i, item) in out.iter_mut().enumerate() {
229 *item = self.data[N - 1 - i];
230 }
231 Self { data: out }
232 }
233 #[must_use]
235 #[inline]
236 pub fn all_eq(&self, other: &Self) -> bool
237 where
238 T: PartialEq,
239 {
240 self.data == other.data
241 }
242}
243
244impl<T: Num, const N: usize> Array<T, N> {
245 #[must_use]
247 #[inline]
248 pub fn add(&self, other: &Self) -> Self {
249 self.zip_with(other, |x, y| x + y)
250 }
251
252 #[must_use]
254 #[inline]
255 pub fn sub(&self, other: &Self) -> Self {
256 self.zip_with(other, |x, y| x - y)
257 }
258
259 #[must_use]
261 #[inline]
262 pub fn scalar_mul(&self, s: T) -> Self {
263 self.map(|x| x * s)
264 }
265
266 #[must_use]
268 #[inline]
269 pub fn scalar_add(&self, s: T) -> Self {
270 self.map(|x| x + s)
271 }
272
273 #[must_use]
275 #[inline]
276 pub fn scalar_div(&self, s: T) -> Self
277 where
278 T: Div<Output = T>,
279 {
280 self.map(|x| x / s)
281 }
282}
283
284impl<T: Real, const N: usize> Array<T, N> {
285 #[must_use]
293 #[inline]
294 pub fn l2_norm(&self) -> T {
295 self.dot(self).sqrt()
296 }
297
298 #[must_use]
300 pub fn l1_norm(&self) -> T {
301 let mut acc = T::zero();
302 for &x in &self.data {
303 acc = acc + x.abs();
304 }
305 acc
306 }
307
308 #[must_use]
310 pub fn max_norm(&self) -> T {
311 let mut acc = T::zero();
312 for &x in &self.data {
313 let ax = x.abs();
314 if ax > acc {
315 acc = ax;
316 }
317 }
318 acc
319 }
320
321 #[must_use]
323 pub fn try_normalize(&self) -> Option<Self> {
324 let n = self.l2_norm();
325 if n == T::zero() {
326 None
327 } else {
328 Some(self.scalar_div(n))
329 }
330 }
331
332 #[must_use]
334 #[inline]
335 pub fn normalize(&self) -> Self {
336 self.try_normalize().unwrap_or_else(Self::zeros)
337 }
338
339 #[must_use]
341 #[inline]
342 pub fn mean(&self) -> T {
343 self.sum() / T::from_f64(N as f64)
344 }
345}
346
347impl<T, const N: usize> From<[T; N]> for Array<T, N> {
348 fn from(data: [T; N]) -> Self {
349 Self { data }
350 }
351}
352
353impl<T, const N: usize> From<Array<T, N>> for [T; N] {
354 fn from(arr: Array<T, N>) -> Self {
355 arr.data
356 }
357}
358
359impl<T, const N: usize> AsRef<[T]> for Array<T, N> {
360 fn as_ref(&self) -> &[T] {
361 &self.data
362 }
363}
364
365impl<T, const N: usize> core::ops::Index<usize> for Array<T, N> {
366 type Output = T;
367
368 fn index(&self, index: usize) -> &Self::Output {
369 &self.data[index]
370 }
371}
372
373impl<T: Num, const N: usize> Add for Array<T, N> {
374 type Output = Self;
375
376 fn add(self, rhs: Self) -> Self::Output {
377 Self::add(&self, &rhs)
378 }
379}
380
381impl<T: Num, const N: usize> Sub for Array<T, N> {
382 type Output = Self;
383
384 fn sub(self, rhs: Self) -> Self::Output {
385 Self::sub(&self, &rhs)
386 }
387}
388
389impl<T: Num, const N: usize> Mul<T> for Array<T, N> {
390 type Output = Self;
391
392 fn mul(self, rhs: T) -> Self::Output {
393 self.scalar_mul(rhs)
394 }
395}
396
397impl<T: Num + Div<Output = T>, const N: usize> Div<T> for Array<T, N> {
398 type Output = Self;
399
400 fn div(self, rhs: T) -> Self::Output {
401 self.scalar_div(rhs)
402 }
403}
404
405#[cfg(test)]
406mod tests {
407 use super::*;
408
409 #[test]
410 fn basic_ops() {
411 let a = Array::from([1.0f64, 2.0, 3.0]);
412 let b = Array::from([4.0, 5.0, 6.0]);
413 assert_eq!(a + b, Array::from([5.0, 7.0, 9.0]));
414 assert_eq!(a.dot(&b), 32.0);
415 assert_eq!(a.sum(), 6.0);
416 }
417
418 #[test]
419 fn const_generic_sizes() {
420 let a: Array<u32, 4> = Array::ones();
421 assert_eq!(a.sum(), 4);
422 let b: Array<u32, 8> = Array::zeros();
423 assert_eq!(b.sum(), 0);
424 }
425
426 #[test]
427 fn norms() {
428 let a = Array::from([3.0f64, 4.0]);
429 assert!((a.l2_norm() - 5.0).abs() < 1e-12);
430 assert_eq!(a.l1_norm(), 7.0);
431 assert_eq!(a.max_norm(), 4.0);
432 assert!((a.try_normalize().unwrap().l2_norm() - 1.0).abs() < 1e-12);
433 }
434
435 #[test]
436 fn int_arrays() {
437 let a = Array::from([1i32, 2, 3]);
438 let b = Array::from([10, 20, 30]);
439 assert_eq!(a.component_mul(&b), Array::from([10, 40, 90]));
440 assert_eq!(a.cumulative_sum(), Array::from([1, 3, 6]));
441 }
442
443 #[test]
444 fn normalize_zero() {
445 let z: Array<f64, 3> = Array::zeros();
446 assert!(z.try_normalize().is_none());
447 assert_eq!(z.normalize(), z);
448 }
449
450 #[test]
451 fn float_array_compiles_and_eq() {
452 let a = Array::from([1.0f64, 2.0, 3.0]);
453 let b = Array::from([1.0f64, 2.0, 3.0]);
454 assert_eq!(a, b);
455 assert!(Array::<f64, 3>::eq(&a, &b));
456 }
457
458 #[test]
459 fn kahan_sum_precision() {
460 let vals: [f64; 1000] = [0.1; 1000];
463 let a = Array::<f64, 1000> { data: vals };
464 assert!((a.sum() - 100.0).abs() < 1e-10);
465 }
466
467 #[test]
468 fn kahan_dot_precision() {
469 let vals: [f64; 1000] = [0.1; 1000];
471 let a = Array::<f64, 1000> { data: vals };
472 let b = Array::<f64, 1000> { data: [1.0; 1000] };
473 assert!((a.dot(&b) - 100.0).abs() < 1e-10);
474 }
475}