scirs2-spatial 0.4.2

//! Inverse Distance Weighting interpolation
//!
//! This module provides Inverse Distance Weighting (IDW) interpolation, a
//! simple and efficient method for interpolating scattered data.
//!
//! IDW interpolation works by weighting neighboring points by the inverse of
//! their distance raised to a power. The power parameter controls the smoothness
//! of the interpolation, with higher values giving more weight to close points.
//!
//! The method is fast but can produce "bull's-eye" patterns around sample points,
//! especially with high power values.

use crate::distance::EuclideanDistance;
use crate::error::{SpatialError, SpatialResult};
use crate::kdtree::KDTree;
use scirs2_core::ndarray::{Array1, Array2, ArrayView1, ArrayView2};

/// Inverse Distance Weighting interpolator for scattered data
///
/// # Examples
///
/// ```
/// use scirs2_spatial::interpolate::IDWInterpolator;
/// use scirs2_core::ndarray::array;
///
/// // Create sample points and values
/// let points = array![
///     [0.0, 0.0],
///     [1.0, 0.0],
///     [0.0, 1.0],
///     [1.0, 1.0],
/// ];
/// let values = array![0.0, 1.0, 2.0, 3.0];
///
/// // Create interpolator with power=2
/// let interp = IDWInterpolator::new(&points.view(), &values.view(), 2.0, None).expect("Operation failed");
///
/// // Interpolate at a point
/// let query_point = array![0.5, 0.5];
/// let result = interp.interpolate(&query_point.view()).expect("Operation failed");
///
/// // Should be close to 1.5
/// assert!((result - 1.5).abs() < 0.1);
/// ```
#[derive(Debug, Clone)]
pub struct IDWInterpolator {
    /// Input points (N x D)
    points: Array2<f64>,
    /// Input values (N)
    values: Array1<f64>,
    /// Dimensionality of the input points
    dim: usize,
    /// Number of input points
    n_points: usize,
    /// Power parameter (p)
    power: f64,
    /// Number of neighbors to use (None means use all points)
    n_neighbors: Option<usize>,
    /// KD-tree for fast nearest neighbor lookup
    kdtree: KDTree<f64, EuclideanDistance<f64>>,
}

impl IDWInterpolator {
    /// Create a new IDW interpolator
    ///
    /// # Arguments
    ///
    /// * `points` - Input points with shape (n_samples, n_dims)
    /// * `values` - Input values with shape (n_samples,)
    /// * `power` - Power parameter (p), controls the importance of nearby points
    /// * `n_neighbors` - Number of neighbors to use (None = use all points)
    ///
    /// # Returns
    ///
    /// A new IDWInterpolator
    ///
    /// # Errors
    ///
    /// * If points and values have different lengths
    /// * If power is negative
    /// * If n_neighbors is 0 or greater than n_points
    pub fn new(
        points: &ArrayView2<'_, f64>,
        values: &ArrayView1<f64>,
        power: f64,
        n_neighbors: Option<usize>,
    ) -> SpatialResult<Self> {
        // Check input dimensions
        let n_points = points.nrows();
        let dim = points.ncols();

        if n_points != values.len() {
            return Err(SpatialError::DimensionError(format!(
                "Number of points ({}) must match number of values ({})",
                n_points,
                values.len()
            )));
        }

        if power < 0.0 {
            return Err(SpatialError::ValueError(format!(
                "Power parameter must be non-negative, got {power}"
            )));
        }

        if let Some(k) = n_neighbors {
            if k == 0 {
                return Err(SpatialError::ValueError(
                    "Number of _neighbors must be positive".to_string(),
                ));
            }
            if k > n_points {
                return Err(SpatialError::ValueError(format!(
                    "Number of _neighbors ({k}) cannot exceed number of points ({n_points})"
                )));
            }
        }

        // Build KD-tree for fast nearest neighbor lookups
        let kdtree = KDTree::new(&points.to_owned())?;

        Ok(Self {
            points: points.to_owned(),
            values: values.to_owned(),
            dim,
            n_points,
            power,
            n_neighbors,
            kdtree,
        })
    }

    /// Interpolate at a single point
    ///
    /// # Arguments
    ///
    /// * `point` - Query point with shape (n_dims,)
    ///
    /// # Returns
    ///
    /// Interpolated value at the query point
    ///
    /// # Errors
    ///
    /// * If the point dimensions don't match the interpolator
    pub fn interpolate(&self, point: &ArrayView1<f64>) -> SpatialResult<f64> {
        // Check dimension
        if point.len() != self.dim {
            return Err(SpatialError::DimensionError(format!(
                "Query point has dimension {}, expected {}",
                point.len(),
                self.dim
            )));
        }

        // Handle exact matches first
        for i in 0..self.n_points {
            let data_point = self.points.row(i);
            if Self::is_same_point(&data_point, point) {
                return Ok(self.values[i]);
            }
        }

        // Get the neighbors to use
        let (indices, distances) = match self.n_neighbors {
            Some(k) => {
                // Use k nearest neighbors
                self.kdtree
                    .query(point.as_slice().expect("Operation failed"), k)?
            }
            None => {
                // Use all points
                let mut indices = Vec::with_capacity(self.n_points);
                let mut distances = Vec::with_capacity(self.n_points);

                for i in 0..self.n_points {
                    let data_point = self.points.row(i);
                    let dist_sq = Self::squared_distance(&data_point, point);
                    indices.push(i);
                    distances.push(dist_sq);
                }

                (indices, distances)
            }
        };

        // Calculate IDW weights and interpolated value
        let mut weighted_sum = 0.0;
        let mut weight_sum = 0.0;

        for i in 0..indices.len() {
            let dist_sq = distances[i];

            // Handle zero distance case (coincident point)
            if dist_sq < 1e-10 {
                return Ok(self.values[indices[i]]);
            }

            // Calculate weight
            let weight = 1.0 / dist_sq.powf(self.power / 2.0);

            weighted_sum += weight * self.values[indices[i]];
            weight_sum += weight;
        }

        if weight_sum > 0.0 {
            Ok(weighted_sum / weight_sum)
        } else {
            // This should not happen with valid data
            Err(SpatialError::ComputationError(
                "Zero weight sum in IDW interpolation".to_string(),
            ))
        }
    }

    /// Interpolate at multiple points
    ///
    /// # Arguments
    ///
    /// * `points` - Query points with shape (n_queries, n_dims)
    ///
    /// # Returns
    ///
    /// Interpolated values with shape (n_queries,)
    ///
    /// # Errors
    ///
    /// * If the points dimensions don't match the interpolator
    pub fn interpolate_many(&self, points: &ArrayView2<'_, f64>) -> SpatialResult<Array1<f64>> {
        // Check dimensions
        if points.ncols() != self.dim {
            return Err(SpatialError::DimensionError(format!(
                "Query _points have dimension {}, expected {}",
                points.ncols(),
                self.dim
            )));
        }

        let n_queries = points.nrows();
        let mut results = Array1::zeros(n_queries);

        // Interpolate each point
        for i in 0..n_queries {
            let point = points.row(i);
            results[i] = self.interpolate(&point)?;
        }

        Ok(results)
    }

    /// Get the power parameter used by this interpolator
    pub fn power(&self) -> f64 {
        self.power
    }

    /// Get the number of neighbors used by this interpolator
    pub fn n_neighbors(&self) -> Option<usize> {
        self.n_neighbors
    }

    /// Set the power parameter
    ///
    /// # Arguments
    ///
    /// * `power` - New power parameter
    ///
    /// # Errors
    ///
    /// * If power is negative
    pub fn set_power(&mut self, power: f64) -> SpatialResult<()> {
        if power < 0.0 {
            return Err(SpatialError::ValueError(format!(
                "Power parameter must be non-negative, got {power}"
            )));
        }

        self.power = power;
        Ok(())
    }

    /// Set the number of neighbors
    ///
    /// # Arguments
    ///
    /// * `n_neighbors` - New number of neighbors (None = use all points)
    ///
    /// # Errors
    ///
    /// * If n_neighbors is 0 or greater than n_points
    pub fn set_n_neighbors(&mut self, _nneighbors: Option<usize>) -> SpatialResult<()> {
        if let Some(k) = _nneighbors {
            if k == 0 {
                return Err(SpatialError::ValueError(
                    "Number of _neighbors must be positive".to_string(),
                ));
            }
            if k > self.n_points {
                return Err(SpatialError::ValueError(format!(
                    "Number of _neighbors ({}) cannot exceed number of points ({})",
                    k, self.n_points
                )));
            }
        }

        self.n_neighbors = _nneighbors;
        Ok(())
    }

    /// Check if two points are the same (within a small tolerance)
    ///
    /// # Arguments
    ///
    /// * `p1` - First point
    /// * `p2` - Second point
    ///
    /// # Returns
    ///
    /// True if the points are considered the same
    fn is_same_point(p1: &ArrayView1<f64>, p2: &ArrayView1<f64>) -> bool {
        Self::squared_distance(p1, p2) < 1e-10
    }

    /// Compute the squared Euclidean distance between two points
    ///
    /// # Arguments
    ///
    /// * `p1` - First point
    /// * `p2` - Second point
    ///
    /// # Returns
    ///
    /// Squared Euclidean distance between the points
    fn squared_distance(p1: &ArrayView1<f64>, p2: &ArrayView1<f64>) -> f64 {
        let mut sum_sq = 0.0;
        for i in 0..p1.len().min(p2.len()) {
            let diff = p1[i] - p2[i];
            sum_sq += diff * diff;
        }
        sum_sq
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use approx::assert_relative_eq;
    use scirs2_core::ndarray::array;

    #[test]
    fn test_idw_interpolation_basic() {
        // Create a simple grid of points
        let points = array![
            [0.0, 0.0], // 0: bottom-left
            [1.0, 0.0], // 1: bottom-right
            [0.0, 1.0], // 2: top-left
            [1.0, 1.0], // 3: top-right
        ];

        // Set up a simple function z = x + y
        let values = array![0.0, 1.0, 1.0, 2.0];

        // Test with different power values
        for power in &[1.0, 2.0, 3.0] {
            // Create the interpolator
            let interp = IDWInterpolator::new(&points.view(), &values.view(), *power, None)
                .expect("Operation failed");

            // Test at the data points (should interpolate exactly)
            let val_00 = interp
                .interpolate(&array![0.0, 0.0].view())
                .expect("Operation failed");
            let val_10 = interp
                .interpolate(&array![1.0, 0.0].view())
                .expect("Operation failed");
            let val_01 = interp
                .interpolate(&array![0.0, 1.0].view())
                .expect("Operation failed");
            let val_11 = interp
                .interpolate(&array![1.0, 1.0].view())
                .expect("Operation failed");

            assert_relative_eq!(val_00, 0.0, epsilon = 1e-10);
            assert_relative_eq!(val_10, 1.0, epsilon = 1e-10);
            assert_relative_eq!(val_01, 1.0, epsilon = 1e-10);
            assert_relative_eq!(val_11, 2.0, epsilon = 1e-10);

            // Test at the center
            let val_center = interp
                .interpolate(&array![0.5, 0.5].view())
                .expect("Operation failed");
            assert_relative_eq!(val_center, 1.0, epsilon = 0.1);
        }
    }

    #[test]
    fn test_idw_with_neighbors() {
        // Create a more complex set of points
        let points = array![
            [0.0, 0.0], // 0
            [1.0, 0.0], // 1
            [0.0, 1.0], // 2
            [1.0, 1.0], // 3
            [0.5, 0.5], // 4
            [0.2, 0.8], // 5
            [0.8, 0.2], // 6
            [0.3, 0.3], // 7
            [0.7, 0.7], // 8
        ];

        // Function z = x + y
        let values = Array1::from_vec(
            points
                .rows()
                .into_iter()
                .map(|row| row[0] + row[1])
                .collect(),
        );

        // Create interpolator with different numbers of neighbors
        let interp_all = IDWInterpolator::new(&points.view(), &values.view(), 2.0, None)
            .expect("Operation failed");

        let interp_3 = IDWInterpolator::new(&points.view(), &values.view(), 2.0, Some(3))
            .expect("Operation failed");

        // Test at a new point
        let test_point = array![0.6, 0.4];

        let val_all = interp_all
            .interpolate(&test_point.view())
            .expect("Operation failed");
        let val_3 = interp_3
            .interpolate(&test_point.view())
            .expect("Operation failed");

        // Both should be close to x + y = 0.6 + 0.4 = 1.0
        assert_relative_eq!(val_all, 1.0, epsilon = 0.1);
        assert_relative_eq!(val_3, 1.0, epsilon = 0.1);

        // They might be slightly different, but not guaranteed in all implementations
        // Different implementations may produce very similar results
        // assert!(f64::abs(val_all - val_3) > 1e-6);
    }

    #[test]
    fn test_interpolate_many() {
        // Create a simple grid of points
        let points = array![[0.0, 0.0], [1.0, 0.0], [0.0, 1.0], [1.0, 1.0],];

        // Set up a simple function z = x + y
        let values = array![0.0, 1.0, 1.0, 2.0];

        // Create the interpolator
        let interp = IDWInterpolator::new(&points.view(), &values.view(), 2.0, None)
            .expect("Operation failed");

        // Test multiple points at once
        let query_points = array![[0.0, 0.0], [1.0, 0.0], [0.0, 1.0], [1.0, 1.0], [0.5, 0.5],];

        let results = interp
            .interpolate_many(&query_points.view())
            .expect("Operation failed");

        assert_eq!(results.len(), 5);
        assert_relative_eq!(results[0], 0.0, epsilon = 1e-10);
        assert_relative_eq!(results[1], 1.0, epsilon = 1e-10);
        assert_relative_eq!(results[2], 1.0, epsilon = 1e-10);
        assert_relative_eq!(results[3], 2.0, epsilon = 1e-10);
        assert_relative_eq!(results[4], 1.0, epsilon = 0.1);
    }

    #[test]
    fn test_setter_methods() {
        // Create interpolator
        let points = array![[0.0, 0.0], [1.0, 0.0], [0.0, 1.0], [1.0, 1.0],];

        let values = array![0.0, 1.0, 1.0, 2.0];

        let mut interp = IDWInterpolator::new(&points.view(), &values.view(), 2.0, None)
            .expect("Operation failed");

        // Test setter methods
        assert_eq!(interp.power(), 2.0);
        assert_eq!(interp.n_neighbors(), None);

        interp.set_power(3.0).expect("Operation failed");
        assert_eq!(interp.power(), 3.0);

        interp.set_n_neighbors(Some(2)).expect("Operation failed");
        assert_eq!(interp.n_neighbors(), Some(2));

        // Test error cases
        let result = interp.set_power(-1.0);
        assert!(result.is_err());

        let result = interp.set_n_neighbors(Some(0));
        assert!(result.is_err());

        let result = interp.set_n_neighbors(Some(10));
        assert!(result.is_err());
    }

    #[test]
    fn test_error_handling() {
        // Wrong dimensions
        let points = array![[0.0, 0.0], [1.0, 0.0], [0.0, 1.0]];
        let values = array![0.0, 1.0, 1.0];

        let interp = IDWInterpolator::new(&points.view(), &values.view(), 2.0, None)
            .expect("Operation failed");

        let result = interp.interpolate(&array![0.0].view());
        assert!(result.is_err());

        // Negative power
        let result = IDWInterpolator::new(&points.view(), &values.view(), -1.0, None);
        assert!(result.is_err());

        // Invalid neighbors
        let result = IDWInterpolator::new(&points.view(), &values.view(), 2.0, Some(0));
        assert!(result.is_err());

        let result = IDWInterpolator::new(&points.view(), &values.view(), 2.0, Some(10));
        assert!(result.is_err());
    }
}