1use std::ops::{Index, IndexMut};
2
3use num::Num;
4
5#[derive(Debug, Clone, Default)]
7pub struct Ndarray<T>
8where
9 T: Num + Clone,
10{
11 pub data: Vec<T>,
13
14 pub shape: Vec<usize>,
16
17 pub stride: Vec<usize>,
19}
20
21impl<T> Ndarray<T>
22where
23 T: Num + Clone,
24{
25 pub fn new(shape: &[usize], default: T) -> Self {
27 let d = shape.len();
28
29 let mut size = 1;
30 shape.iter().for_each(|x| size *= x);
31
32 let data = vec![default; size];
33
34 let mut stride = vec![0; d];
35
36 let mut s = 1;
37 for i in (0..d).rev() {
38 stride[i] = s;
39 s *= shape[i];
40 }
41
42 Self {
43 data,
44 shape: shape.to_vec(),
45 stride,
46 }
47 }
48
49 pub fn index(&self, coords: &[usize]) -> usize {
51 coords
52 .iter()
53 .zip(self.stride.iter())
54 .map(|(a, b)| a * b)
55 .sum()
56 }
57
58 pub fn contains(&self, coords: &[usize]) -> bool {
60 !coords.iter().zip(self.shape.iter()).any(|(&a, &b)| a >= b)
61 }
62}
63
64impl<T: Num + Clone> Index<&[usize]> for Ndarray<T> {
65 type Output = T;
66
67 fn index(&self, index: &[usize]) -> &Self::Output {
68 &self.data.get(self.index(index)).unwrap()
69 }
70}
71
72impl<T: Num + Clone> IndexMut<&[usize]> for Ndarray<T> {
73 fn index_mut(&mut self, index: &[usize]) -> &mut Self::Output {
74 let index = self.index(index);
75 self.data.get_mut(index).unwrap()
76 }
77}
78
79pub fn ndarray<T: Num + Clone>(shape: &[usize], default: T) -> Ndarray<T> {
81 Ndarray::new(shape, default)
82}