1use super::tensor::Tensor;
2
3pub struct Grid<const N: usize> {
12 pub min: [f64; N],
13 pub dx: [f64; N],
14 pub tensor_info: Tensor<N>,
15}
16
17impl<const N: usize> Grid<N> {
18 pub fn new(shape: [usize; N], min: [f64; N], max: [f64; N]) -> Self {
19 let mut dx = [0.0; N];
20 for i in 0..N {
21 if shape[i] > 1 {
22 dx[i] = (max[i] - min[i]) / ((shape[i] - 1) as f64);
23 } else {
24 dx[i] = 0.0;
25 }
26 }
27
28 let tensor_info = Tensor::new(shape, 0.0);
29
30 Self {
31 min,
32 dx,
33 tensor_info,
34 }
35 }
36
37 pub fn total_size(&self) -> usize {
38 self.tensor_info.data.len()
39 }
40
41 pub fn get_closest_index(&self, state: &[f64; N]) -> usize {
42 let mut coords = [0; N];
43 for i in 0..N {
44 if self.dx[i] == 0.0 {
45 coords[i] = 0;
46 } else {
47 let c = ((state[i] - self.min[i]) / self.dx[i]).round();
48 coords[i] = (c as usize).clamp(0, self.tensor_info.shape[i] - 1);
49 }
50 }
51 self.tensor_info.get_idx(&coords)
52 }
53}