solver/lookup/

table.rs

use ndarray::ArrayD;
use serde::{Deserialize, Serialize};
use std::fs::File;
use std::io::{BufReader, BufWriter};
use std::path::Path;

#[derive(Clone, Serialize, Deserialize, Debug)]
pub struct LookupTable {
    /// The grid coordinates for each dimension.
    /// axes[0] corresponds to the first dimension of `data`, etc.
    pub axes: Vec<Vec<f64>>,
    /// The data stored in a dynamic-dimensional array.
    pub data: ArrayD<f64>,
    /// Cached flat strides for each dimension (C-order).
    /// strides[i] = product of axis lengths for dimensions i+1 .. ndim.
    /// Precomputed in `new` so `interpolate` avoids repeated multiplications.
    #[serde(skip)]
    strides: Vec<usize>,
}

impl LookupTable {
    pub fn new(axes: Vec<Vec<f64>>, data: ArrayD<f64>) -> Self {
        assert_eq!(
            axes.len(),
            data.ndim(),
            "Number of axes must match data dimensions"
        );
        for (i, axis) in axes.iter().enumerate() {
            assert_eq!(
                axis.len(),
                data.shape()[i],
                "Axis length must match data shape at dim {}",
                i
            );
            for j in 0..axis.len() - 1 {
                assert!(
                    axis[j] < axis[j + 1],
                    "Axes must be sorted strictly increasing"
                );
            }
        }

        // Precompute C-order strides
        let ndim = axes.len();
        let mut strides = vec![1usize; ndim];
        for i in (0..ndim.saturating_sub(1)).rev() {
            strides[i] = strides[i + 1] * axes[i + 1].len();
        }

        Self {
            axes,
            data,
            strides,
        }
    }

    pub fn save(&self, path: &Path) -> bincode::Result<()> {
        let file = File::create(path)?;
        let writer = BufWriter::new(file);
        bincode::serialize_into(writer, self)
    }

    pub fn load(path: &Path) -> bincode::Result<Self> {
        let file = File::open(path)?;
        let reader = BufReader::new(file);
        let mut t: Self = bincode::deserialize_from(reader)?;
        // Recompute derived fields that are not serialized
        let ndim = t.axes.len();
        let mut strides = vec![1usize; ndim];
        for i in (0..ndim.saturating_sub(1)).rev() {
            strides[i] = strides[i + 1] * t.axes[i + 1].len();
        }
        t.strides = strides;
        Ok(t)
    }

    /// Multilinear interpolation at the given point.
    /// `point` must have the same number of dimensions as `axes`.
    ///
    /// For uniform axes (produced by `linspace`) the bracket index is computed
    /// in O(1) via direct division instead of binary search, avoiding a
    /// per-call heap allocation.
    pub fn interpolate(&self, point: &[f64]) -> f64 {
        let ndim = self.axes.len();
        assert_eq!(point.len(), ndim, "Point dimension mismatch");
        let flat = self
            .data
            .as_slice_memory_order()
            .expect("LookupTable data must be contiguous in memory order");

        // Stack-allocated arrays (ndim <= 16 is always true for this use case)
        let mut indices = [0usize; 16];
        let mut weights = [0.0f64; 16];

        for i in 0..ndim {
            let axis = &self.axes[i];
            let n = axis.len();
            let val = point[i];

            if n == 1 {
                // indices[i] = 0, weights[i] = 0.0 already
                continue;
            }

            if val <= axis[0] {
                // weights[i] = 0.0 already
            } else if val >= axis[n - 1] {
                indices[i] = n - 2;
                weights[i] = 1.0;
            } else {
                // Fast path for uniform axes: O(1) direct computation.
                // A uniform axis has constant spacing; check if step is uniform
                // by comparing first and last intervals.
                let step0 = axis[1] - axis[0];
                let step_last = axis[n - 1] - axis[n - 2];
                let idx = if (step0 - step_last).abs() < step0 * 1e-9 {
                    // Uniform: direct index
                    let raw = (val - axis[0]) / step0;
                    (raw as usize).min(n - 2)
                } else {
                    // Non-uniform: binary search
                    let idx = match axis.binary_search_by(|x| x.partial_cmp(&val).unwrap()) {
                        Ok(exact) => exact,
                        Err(insert) => insert - 1,
                    };
                    idx.min(n - 2)
                };
                let w = (val - axis[idx]) / (axis[idx + 1] - axis[idx]);
                indices[i] = idx;
                weights[i] = w;
            }
        }

        // Iterate over 2^ndim corners using precomputed strides and flat buffer
        let num_corners = 1usize << ndim;
        let mut result = 0.0;

        for corner in 0..num_corners {
            let mut w = 1.0f64;
            let mut flat_idx = 0usize;

            for dim in 0..ndim {
                let is_upper = (corner >> dim) & 1 == 1;
                if is_upper {
                    w *= weights[dim];
                    flat_idx +=
                        (indices[dim] + 1).min(self.axes[dim].len() - 1) * self.strides[dim];
                } else {
                    w *= 1.0 - weights[dim];
                    flat_idx += indices[dim] * self.strides[dim];
                }
            }

            result += flat[flat_idx] * w;
        }

        result
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use ndarray::IxDyn;

    /// Helper to create a 1D linear table for testing interpolation
    fn create_linear_1d() -> LookupTable {
        let axes = vec![vec![0.0, 1.0, 2.0, 3.0, 4.0]];
        let data = ArrayD::from_shape_vec(IxDyn(&[5]), vec![0.0, 1.0, 2.0, 3.0, 4.0]).unwrap();
        LookupTable::new(axes, data)
    }

    /// Helper to create a 2D table: f(x, y) = x + 2*y
    fn create_2d_linear() -> LookupTable {
        let x_axis = vec![0.0, 1.0, 2.0];
        let y_axis = vec![0.0, 1.0, 2.0];
        let mut data_vec = Vec::new();
        // C-order: last index changes fastest
        // For [nx, ny] shape, data[i*ny + j] = f(x[i], y[j])
        for i in 0..3 {
            for j in 0..3 {
                let x = x_axis[i];
                let y = y_axis[j];
                data_vec.push(x + 2.0 * y);
            }
        }
        let data = ArrayD::from_shape_vec(IxDyn(&[3, 3]), data_vec).unwrap();
        LookupTable::new(vec![x_axis, y_axis], data)
    }

    /// Helper to create a 3D table: f(x, y, z) = x + y + z
    fn create_3d_linear() -> LookupTable {
        let x_axis = vec![0.0, 1.0, 2.0];
        let y_axis = vec![0.0, 1.0, 2.0];
        let z_axis = vec![0.0, 1.0, 2.0];
        let mut data_vec = Vec::new();
        // C-order: for [nx, ny, nz], data[i*ny*nz + j*nz + k] = f(x[i], y[j], z[k])
        for i in 0..3 {
            for j in 0..3 {
                for k in 0..3 {
                    let x = x_axis[i];
                    let y = y_axis[j];
                    let z = z_axis[k];
                    data_vec.push(x + y + z);
                }
            }
        }
        let data = ArrayD::from_shape_vec(IxDyn(&[3, 3, 3]), data_vec).unwrap();
        LookupTable::new(vec![x_axis, y_axis, z_axis], data)
    }

    /// Helper to create a 4D table: f(a, b, c, d) = a + b + c + d
    fn create_4d_linear() -> LookupTable {
        let axes = vec![
            vec![0.0, 1.0, 2.0],
            vec![0.0, 1.0, 2.0],
            vec![0.0, 1.0, 2.0],
            vec![0.0, 1.0, 2.0],
        ];
        let mut data_vec = Vec::new();
        // C-order: for [na, nb, nc, nd]
        for a in 0..3 {
            for b in 0..3 {
                for c in 0..3 {
                    for d in 0..3 {
                        let val = axes[0][a] + axes[1][b] + axes[2][c] + axes[3][d];
                        data_vec.push(val);
                    }
                }
            }
        }
        let data = ArrayD::from_shape_vec(IxDyn(&[3, 3, 3, 3]), data_vec).unwrap();
        LookupTable::new(axes, data)
    }

    #[test]
    fn test_1d_exact_points() {
        let table = create_linear_1d();
        // Test at exact grid points
        assert_eq!(table.interpolate(&[0.0]), 0.0);
        assert_eq!(table.interpolate(&[1.0]), 1.0);
        assert_eq!(table.interpolate(&[2.0]), 2.0);
        assert_eq!(table.interpolate(&[3.0]), 3.0);
        assert_eq!(table.interpolate(&[4.0]), 4.0);
    }

    #[test]
    fn test_1d_interpolation() {
        let table = create_linear_1d();
        // Test interpolation at midpoints
        assert!((table.interpolate(&[0.5]) - 0.5).abs() < 1e-10);
        assert!((table.interpolate(&[1.5]) - 1.5).abs() < 1e-10);
        assert!((table.interpolate(&[2.5]) - 2.5).abs() < 1e-10);
        assert!((table.interpolate(&[3.5]) - 3.5).abs() < 1e-10);
    }

    #[test]
    fn test_1d_boundary_clamp() {
        let table = create_linear_1d();
        // Values beyond boundaries should be clamped
        assert_eq!(table.interpolate(&[-1.0]), 0.0);
        assert_eq!(table.interpolate(&[5.0]), 4.0);
    }

    #[test]
    fn test_2d_exact_points() {
        let table = create_2d_linear();
        // f(x, y) = x + 2*y
        assert_eq!(table.interpolate(&[0.0, 0.0]), 0.0);
        assert_eq!(table.interpolate(&[1.0, 0.0]), 1.0);
        assert_eq!(table.interpolate(&[0.0, 1.0]), 2.0);
        assert_eq!(table.interpolate(&[1.0, 1.0]), 3.0);
        assert_eq!(table.interpolate(&[2.0, 2.0]), 6.0);
    }

    #[test]
    fn test_2d_interpolation() {
        let table = create_2d_linear();
        // f(x, y) = x + 2*y
        // At (0.5, 0.5): should be 0.5 + 2*0.5 = 1.5
        assert!((table.interpolate(&[0.5, 0.5]) - 1.5).abs() < 1e-10);
        // At (1.5, 0.5): should be 1.5 + 2*0.5 = 2.5
        assert!((table.interpolate(&[1.5, 0.5]) - 2.5).abs() < 1e-10);
        // At (1.0, 1.5): should be 1.0 + 2*1.5 = 4.0
        assert!((table.interpolate(&[1.0, 1.5]) - 4.0).abs() < 1e-10);
    }

    #[test]
    fn test_3d_exact_points() {
        let table = create_3d_linear();
        // f(x, y, z) = x + y + z
        assert_eq!(table.interpolate(&[0.0, 0.0, 0.0]), 0.0);
        assert_eq!(table.interpolate(&[1.0, 1.0, 1.0]), 3.0);
        assert_eq!(table.interpolate(&[2.0, 2.0, 2.0]), 6.0);
        assert_eq!(table.interpolate(&[1.0, 2.0, 0.0]), 3.0);
    }

    #[test]
    fn test_3d_interpolation() {
        let table = create_3d_linear();
        // f(x, y, z) = x + y + z
        // At (0.5, 0.5, 0.5): should be 1.5
        assert!((table.interpolate(&[0.5, 0.5, 0.5]) - 1.5).abs() < 1e-10);
        // At (1.5, 1.5, 1.5): should be 4.5
        assert!((table.interpolate(&[1.5, 1.5, 1.5]) - 4.5).abs() < 1e-10);
    }

    #[test]
    fn test_4d_exact_points() {
        let table = create_4d_linear();
        // f(a, b, c, d) = a + b + c + d
        assert_eq!(table.interpolate(&[0.0, 0.0, 0.0, 0.0]), 0.0);
        assert_eq!(table.interpolate(&[1.0, 1.0, 1.0, 1.0]), 4.0);
        assert_eq!(table.interpolate(&[2.0, 2.0, 2.0, 2.0]), 8.0);
    }

    #[test]
    fn test_4d_interpolation() {
        let table = create_4d_linear();
        // f(a, b, c, d) = a + b + c + d
        // At (0.5, 0.5, 0.5, 0.5): should be 2.0
        assert!((table.interpolate(&[0.5, 0.5, 0.5, 0.5]) - 2.0).abs() < 1e-10);
        // At (1.5, 1.5, 1.5, 1.5): should be 6.0
        assert!((table.interpolate(&[1.5, 1.5, 1.5, 1.5]) - 6.0).abs() < 1e-10);
    }

    #[test]
    fn test_nonuniform_axes() {
        // Create table with non-uniform axis spacing
        let x_axis = vec![0.0, 0.1, 0.5, 2.0];
        let data = ArrayD::from_shape_vec(IxDyn(&[4]), vec![0.0, 0.1, 0.5, 2.0]).unwrap();
        let table = LookupTable::new(vec![x_axis], data);

        // Test at exact points
        assert_eq!(table.interpolate(&[0.0]), 0.0);
        assert_eq!(table.interpolate(&[0.1]), 0.1);
        assert_eq!(table.interpolate(&[0.5]), 0.5);
        assert_eq!(table.interpolate(&[2.0]), 2.0);

        // Test interpolation between first two points
        let mid_01 = table.interpolate(&[0.05]);
        assert!((mid_01 - 0.05).abs() < 1e-10);
    }

    #[test]
    fn test_single_element_axis() {
        // Test with a table that has only one element in a dimension
        let axes = vec![
            vec![0.0, 1.0],
            vec![5.0], // Single element
        ];
        let data = ArrayD::from_shape_vec(IxDyn(&[2, 1]), vec![10.0, 20.0]).unwrap();
        let table = LookupTable::new(axes, data);

        // Should return the single value for any query in that dimension
        assert_eq!(table.interpolate(&[0.0, 5.0]), 10.0);
        assert_eq!(table.interpolate(&[0.5, 5.0]), 15.0);
        assert_eq!(table.interpolate(&[1.0, 5.0]), 20.0);
        // Out-of-bounds in single-element dimension should still work
        assert_eq!(table.interpolate(&[0.5, 100.0]), 15.0);
    }

    #[test]
    fn test_boundary_clamping() {
        let table = create_2d_linear();
        // Query outside bounds should clamp
        let result_below = table.interpolate(&[-1.0, -1.0]);
        assert_eq!(result_below, table.interpolate(&[0.0, 0.0]));

        let result_above = table.interpolate(&[10.0, 10.0]);
        assert_eq!(result_above, table.interpolate(&[2.0, 2.0]));

        // Mixed: below in one dim, above in another
        let result_mixed = table.interpolate(&[-1.0, 10.0]);
        assert_eq!(result_mixed, table.interpolate(&[0.0, 2.0]));
    }

    #[test]
    fn test_save_load_round_trip() {
        use tempfile::TempDir;

        let original = create_3d_linear();

        let temp_dir = TempDir::new().expect("Failed to create temp dir");
        let temp_file = temp_dir.path().join("test_lookup_table.bin");
        original.save(&temp_file).expect("Save failed");
        assert!(temp_file.exists(), "File should exist after save");

        let loaded = LookupTable::load(&temp_file).expect("Load failed");
        assert_eq!(loaded.axes.len(), original.axes.len());
        for i in 0..loaded.axes.len() {
            assert_eq!(loaded.axes[i], original.axes[i]);
        }

        // Test that loaded table produces same interpolation results
        assert_eq!(loaded.interpolate(&[0.0, 0.0, 0.0]), 0.0);
        assert_eq!(loaded.interpolate(&[1.5, 1.5, 1.5]), 4.5);
        assert!((loaded.interpolate(&[0.5, 0.5, 0.5]) - 1.5).abs() < 1e-10);
    }

    #[test]
    fn test_axes_validation() {
        // Test that constructor validates axes are strictly increasing
        let axes_bad = vec![vec![0.0, 1.0, 1.0, 2.0]]; // Not strictly increasing!
        let data = ArrayD::from_shape_vec(IxDyn(&[4]), vec![0.0, 1.0, 1.0, 2.0]).unwrap();
        let result = std::panic::catch_unwind(|| LookupTable::new(axes_bad, data));
        assert!(
            result.is_err(),
            "Should panic on non-strictly-increasing axis"
        );
    }

    #[test]
    fn test_dimension_mismatch() {
        let axes = vec![vec![0.0, 1.0], vec![0.0, 1.0]];
        let data = ArrayD::from_shape_vec(IxDyn(&[3, 3]), vec![0.0; 9]).unwrap(); // Wrong shape!
        let result = std::panic::catch_unwind(|| LookupTable::new(axes, data));
        assert!(result.is_err(), "Should panic on dimension mismatch");
    }

    #[test]
    fn test_interpolate_dimension_mismatch() {
        let table = create_2d_linear();
        let result = std::panic::catch_unwind(|| table.interpolate(&[0.5])); // Only 1 arg for 2D table
        assert!(result.is_err(), "Should panic on point dimension mismatch");
    }

    #[test]
    fn test_large_grid_interpolation() {
        // Test with larger grid to ensure performance isn't issues
        let x_axis: Vec<f64> = (0..100).map(|i| i as f64 / 99.0).collect();
        let data = ArrayD::from_shape_vec(IxDyn(&[100]), x_axis.clone()).unwrap();
        let table = LookupTable::new(vec![x_axis], data);

        // Spot checks
        assert!((table.interpolate(&[0.0]) - 0.0).abs() < 1e-10);
        assert!((table.interpolate(&[1.0]) - 1.0).abs() < 1e-10);
        assert!((table.interpolate(&[0.5]) - 0.5).abs() < 1e-3); // Looser tolerance for large grid
    }
}