solver/lookup/

adaptive.rs

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

/// Result of recursive bisection: a set of segments with refined points.
#[derive(Clone, Debug)]
pub struct AdaptiveSegments {
    /// Ordered list of grid points (strictly increasing)
    pub points: Vec<f64>,
}

impl AdaptiveSegments {
    /// Create from a set of points (will be sorted and deduplicated)
    pub fn from_points(mut points: Vec<f64>) -> Self {
        points.sort_by(|a, b| a.partial_cmp(b).unwrap());
        points.dedup_by(|a, b| (*a - *b).abs() < 1e-14);
        Self { points }
    }

    /// Length of the grid
    pub fn len(&self) -> usize {
        self.points.len()
    }

    /// Check if empty
    pub fn is_empty(&self) -> bool {
        self.points.is_empty()
    }
}

/// Adaptive grid generator using recursive bisection.
///
/// This generator recursively subdivides intervals in 1D, refining only where
/// the underlying function exhibits high variance (error exceeds tolerance).
///
/// **Algorithm:**
/// 1. Start with an interval [a, b]
/// 2. Evaluate the function at a, midpoint m, and b
/// 3. Compute the linear interpolation at m from endpoints
/// 4. If |f(m) - interpolated(m)| > tolerance, bisect into [a, m] and [m, b] and recurse
/// 5. Otherwise, mark the interval as sufficiently smooth
///
/// This guarantees adaptive refinement: dense grids form only near sharp features,
/// while smooth regions remain coarse.
pub struct AdaptiveGrid1D {
    start: f64,
    end: f64,
    tolerance: f64,
    max_depth: usize,
}

impl AdaptiveGrid1D {
    /// Create a new adaptive 1D grid generator.
    ///
    /// # Arguments
    /// * `start` - Left boundary
    /// * `end` - Right boundary
    /// * `tolerance` - Error threshold for bisection
    /// * `max_depth` - Maximum recursion depth (prevents pathological refinement)
    pub fn new(start: f64, end: f64, tolerance: f64, max_depth: usize) -> Self {
        assert!(start < end, "start must be less than end");
        assert!(tolerance > 0.0, "tolerance must be positive");
        Self {
            start,
            end,
            tolerance,
            max_depth,
        }
    }

    /// Generate an adaptive grid by evaluating `f` at points determined by tolerance.
    ///
    /// Returns a strictly increasing list of grid points.
    pub fn generate<F>(&self, f: F) -> AdaptiveSegments
    where
        F: Fn(f64) -> f64,
    {
        let mut points = Vec::new();
        self.bisect(self.start, self.end, &f, 0, &mut points);
        points.push(self.end); // Always include right boundary
        AdaptiveSegments::from_points(points)
    }

    /// Recursive bisection helper.
    fn bisect<F>(&self, a: f64, b: f64, f: &F, depth: usize, points: &mut Vec<f64>)
    where
        F: Fn(f64) -> f64,
    {
        // Add left boundary only once
        if points.is_empty() || (points.last().unwrap() - a).abs() > 1e-14 {
            points.push(a);
        }

        let m = (a + b) / 2.0;
        let fa = f(a);
        let fm = f(m);
        let fb = f(b);

        // Linear interpolation at midpoint
        let fm_interp = (fa + fb) / 2.0;
        let error = (fm - fm_interp).abs();

        if error > self.tolerance && depth < self.max_depth {
            // Error too large: refine by bisecting
            self.bisect(a, m, f, depth + 1, points);
            self.bisect(m, b, f, depth + 1, points);
        } else {
            // Error acceptable: mark this interval as done, add midpoint if needed
            if points.is_empty() || (points.last().unwrap() - m).abs() > 1e-14 {
                points.push(m);
            }
        }
    }
}

/// A two-tier indirection table for deterministic O(1) lookup.
///
/// **Architecture:**
/// - **Level 0 (Coarse Array):** A flat, uniform array spanning the entire domain at low resolution.
/// - **Level 1 (High-Resolution Array):** Packed blocks of fine-grained data for refined regions.
///
/// Each entry in Level 0 contains either:
/// - A **Direct Value:** The actual lookup result (if no refinement needed)
/// - A **Pointer:** An offset into Level 1 pointing to the start of a high-res block
///
/// A bit-flag (sign bit) distinguishes the two cases.
#[derive(Clone, Serialize, Deserialize, Debug)]
pub struct IndirectionTable {
    /// Coarse uniform grid (Level 0). Entries are either direct values or pointers to Level 1.
    /// Stored as (value_or_offset, is_pointer) pairs for clarity in serialization.
    coarse: Vec<(f64, bool)>,
    /// High-resolution packed blocks (Level 1). Contains refined values for adaptive regions.
    fine: Vec<f64>,
    /// Boundaries of the coarse grid
    min: f64,
    max: f64,
    /// Coarse grid resolution
    coarse_resolution: usize,
}

impl IndirectionTable {
    /// Create an indirection table from coarse and fine-grained data.
    ///
    /// # Arguments
    /// * `coarse` - Level 0 entries: (value_or_offset, is_pointer) tuples
    /// * `fine` - Level 1: packed high-resolution blocks
    /// * `min` - Domain minimum
    /// * `max` - Domain maximum
    /// * `coarse_resolution` - Number of cells in coarse grid
    pub fn new(
        coarse: Vec<(f64, bool)>,
        fine: Vec<f64>,
        min: f64,
        max: f64,
        coarse_resolution: usize,
    ) -> Self {
        assert_eq!(
            coarse.len(),
            coarse_resolution,
            "coarse length must match coarse_resolution"
        );
        assert!(min < max, "min must be less than max");
        Self {
            coarse,
            fine,
            min,
            max,
            coarse_resolution,
        }
    }

    /// Branchless $O(1)$ lookup at a single point.
    ///
    /// Returns either the direct value (single fetch) or looks up in Level 1 (dual fetch).
    pub fn lookup(&self, x: f64) -> f64 {
        if x <= self.min {
            return self.coarse[0].0;
        }
        if x >= self.max {
            return self.coarse[self.coarse_resolution - 1].0;
        }

        // Direct index calculation for coarse grid (O(1))
        let normalized = (x - self.min) / (self.max - self.min);
        let coarse_idx =
            ((normalized * self.coarse_resolution as f64) as usize).min(self.coarse_resolution - 1);

        let (payload, is_pointer) = self.coarse[coarse_idx];

        if !is_pointer {
            // Direct value: single memory fetch complete
            payload
        } else {
            // Pointer into fine-grained array: second fetch required
            let offset = payload as usize;
            if offset < self.fine.len() {
                self.fine[offset]
            } else {
                payload // Fallback to payload if offset invalid
            }
        }
    }

    /// Save the indirection table to disk.
    pub fn save(&self, path: &Path) -> bincode::Result<()> {
        let file = File::create(path)?;
        let writer = BufWriter::new(file);
        bincode::serialize_into(writer, self)
    }

    /// Load an indirection table from disk.
    pub fn load(path: &Path) -> bincode::Result<Self> {
        let file = File::open(path)?;
        let reader = BufReader::new(file);
        bincode::deserialize_from(reader)
    }

    /// Get statistics about the table structure.
    pub fn stats(&self) -> IndirectionStats {
        let num_direct = self.coarse.iter().filter(|(_, is_ptr)| !is_ptr).count();
        let num_pointers = self.coarse.iter().filter(|(_, is_ptr)| *is_ptr).count();
        let coarse_bytes = std::mem::size_of::<f64>() * 2 * self.coarse.len(); // (f64, bool) per entry
        let fine_bytes = std::mem::size_of::<f64>() * self.fine.len();

        IndirectionStats {
            num_direct_values: num_direct,
            num_pointers,
            fine_array_size: self.fine.len(),
            coarse_bytes,
            fine_bytes,
            total_bytes: coarse_bytes + fine_bytes,
        }
    }
}

/// Statistics about an indirection table.
#[derive(Debug, Clone)]
pub struct IndirectionStats {
    pub num_direct_values: usize,
    pub num_pointers: usize,
    pub fine_array_size: usize,
    pub coarse_bytes: usize,
    pub fine_bytes: usize,
    pub total_bytes: usize,
}

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

    #[test]
    fn test_adaptive_grid_1d_linear() {
        let grid_gen = AdaptiveGrid1D::new(0.0, 1.0, 0.1, 10);
        let grid = grid_gen.generate(|x| x); // Linear function: no refinement needed

        // Should have at least start and end
        assert!(grid.len() >= 2);
        assert!((grid.points[0] - 0.0).abs() < 1e-10);
        assert!((grid.points[grid.len() - 1] - 1.0).abs() < 1e-10);

        // Linear function should require minimal refinement
        assert!(grid.len() <= 10, "Linear function should not over-refine");
    }

    #[test]
    fn test_adaptive_grid_1d_quadratic() {
        // f(x) = x^2: smooth curve, moderate refinement
        let grid_gen = AdaptiveGrid1D::new(0.0, 1.0, 0.01, 10);
        let grid = grid_gen.generate(|x| x * x);

        assert!(grid.len() >= 2);
        assert!((grid.points[0] - 0.0).abs() < 1e-10);
        assert!((grid.points[grid.len() - 1] - 1.0).abs() < 1e-10);
    }

    #[test]
    fn test_adaptive_grid_1d_sharp_feature() {
        // f(x) = 1 / (1 + (x - 0.5)^2): bell curve with sharp peak
        // This is more amenable to bisection-based refinement
        let grid_gen = AdaptiveGrid1D::new(0.0, 1.0, 0.01, 12);
        let grid = grid_gen.generate(|x| {
            let dx = x - 0.5;
            1.0 / (1.0 + 100.0 * dx * dx)
        });

        // Should refine around the peak at x=0.5
        assert!(
            grid.len() >= 5,
            "Sharp features should trigger some refinement"
        );
    }

    #[test]
    fn test_adaptive_grid_monotonicity() {
        let grid_gen = AdaptiveGrid1D::new(0.0, 1.0, 0.01, 10);
        let grid = grid_gen.generate(|x| x.exp());

        // Points must be strictly increasing
        for i in 0..grid.len() - 1 {
            assert!(
                grid.points[i] < grid.points[i + 1],
                "Grid points must be strictly increasing"
            );
        }
    }

    #[test]
    fn test_indirection_table_direct_values() {
        // Create a simple indirection table with only direct values (no refinement)
        let coarse = vec![(1.0, false), (2.0, false), (3.0, false), (4.0, false)];
        let fine = vec![];
        let table = IndirectionTable::new(coarse, fine, 0.0, 4.0, 4);

        // Lookups in each coarse region
        assert_eq!(table.lookup(0.0), 1.0); // Left boundary
        assert_eq!(table.lookup(0.5), 1.0); // First coarse cell
        assert_eq!(table.lookup(1.0), 2.0); // Second coarse cell
        assert_eq!(table.lookup(2.5), 3.0); // Third coarse cell
        assert_eq!(table.lookup(4.0), 4.0); // Right boundary
    }

    #[test]
    fn test_indirection_table_with_pointers() {
        // Create an indirection table mixing direct values and pointers
        let coarse = vec![
            (1.0, false), // Direct value
            (0.0, true),  // Pointer to fine[0]
            (3.0, false), // Direct value
            (1.0, true),  // Pointer to fine[1]
        ];
        let fine = vec![1.5, 2.0, 3.5]; // Fine-grained values
        let table = IndirectionTable::new(coarse, fine, 0.0, 4.0, 4);

        // Lookups that hit direct values
        assert_eq!(table.lookup(0.1), 1.0);
        assert_eq!(table.lookup(2.1), 3.0);

        // Lookups that hit pointers (transition through fine array)
        assert_eq!(table.lookup(1.1), 1.5); // Coarse[1] -> fine[0]
        assert_eq!(table.lookup(3.1), 2.0); // Coarse[3] -> fine[1]
    }

    #[test]
    fn test_indirection_stats() {
        let coarse = vec![(1.0, false), (0.0, true), (3.0, false), (1.0, true)];
        let fine = vec![1.5, 2.0, 3.5];
        let table = IndirectionTable::new(coarse, fine, 0.0, 4.0, 4);

        let stats = table.stats();
        assert_eq!(stats.num_direct_values, 2);
        assert_eq!(stats.num_pointers, 2);
        assert_eq!(stats.fine_array_size, 3);
        assert!(stats.total_bytes > 0);
    }

    #[test]
    fn test_adaptive_grid_max_depth_limit() {
        // Verify max_depth prevents infinite refinement
        let grid_gen = AdaptiveGrid1D::new(0.0, 1.0, 0.0001, 5); // Very tight tolerance, shallow depth
        let grid = grid_gen.generate(|x| x.sin());

        // Should respect max_depth and not explode in size
        assert!(
            grid.len() < 1000,
            "max_depth should prevent pathological refinement"
        );
    }

    #[test]
    fn test_adaptive_grid_constant_function() {
        // Constant function should require minimal refinement (zero error everywhere)
        let grid_gen = AdaptiveGrid1D::new(0.0, 1.0, 0.1, 10);
        let grid = grid_gen.generate(|_x| 5.0);

        // Should be minimal: algorithm adds boundaries and at least one midpoint
        assert!(grid.len() <= 3, "Constant function should refine minimally");
    }

    #[test]
    fn test_indirection_table_boundary_clamp() {
        let coarse = vec![(1.0, false), (2.0, false), (3.0, false)];
        let fine = vec![];
        let table = IndirectionTable::new(coarse, fine, 1.0, 3.0, 3);

        // Out-of-bounds on left
        assert_eq!(table.lookup(-10.0), 1.0);
        // Out-of-bounds on right
        assert_eq!(table.lookup(100.0), 3.0);
    }
}