solver/core/
tensor.rs

1/// N-Dimensional Tensor structure for storing grid values.
2///
3/// Uses standard row-major storage to map N-dimensional coordinates to a flat index.
4///
5/// # Generic Parameters
6/// * `N` - The dimension of the tensor (e.g. 2 for Price-Inventory).
7#[derive(Clone, Debug)]
8pub struct Tensor<const N: usize> {
9    pub data: Vec<f64>,
10    pub shape: [usize; N],
11    pub strides: [usize; N], // Pre-computed strides for indexing efficiency
12}
13
14impl<const N: usize> Tensor<N> {
15    pub fn new(shape: [usize; N], initial_val: f64) -> Self {
16        let total_size = shape.iter().product();
17        let mut strides = [1; N];
18        // Calculate row-major strides
19        for i in (0..N - 1).rev() {
20            strides[i] = strides[i + 1] * shape[i + 1];
21        }
22
23        Self {
24            data: vec![initial_val; total_size],
25            shape,
26            strides,
27        }
28    }
29
30    #[inline(always)]
31    pub fn get_idx(&self, coords: &[usize; N]) -> usize {
32        let mut idx = 0;
33        for (i, &coord) in coords.iter().enumerate() {
34            idx += coord * self.strides[i];
35        }
36        idx
37    }
38
39    #[inline(always)]
40    pub fn get_coords(&self, mut idx: usize) -> [usize; N] {
41        let mut coords = [0; N];
42        for (i, coord) in coords.iter_mut().enumerate() {
43            *coord = idx / self.strides[i];
44            idx %= self.strides[i];
45        }
46        coords
47    }
48}
solver/core/tensor.rs

solver/core/
tensor.rs
