Skip to main content

solver/core/
grid.rs

1use super::tensor::Tensor;
2
3/// Represents a N-dimensional Cartesian discretization grid.
4///
5/// Contains parameters for converting between discrete indices and continuous state space values.
6///
7/// # Fields
8/// * `min` - The lower bound coordinate for each dimension.
9/// * `dx` - The step size for each dimension: $(Max - Min) / (Shape - 1)$.
10/// * `tensor_info` - Underlying tensor structure holding topology (shape, strides) but not necessarily data values.
11pub 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}