loess-rs 0.9.0

//! Memory management and buffer recycling for LOESS operations.
//!
//! This module provides centralized, reusable workspaces to minimize dynamic memory
//! allocations during LOESS fitting. By allocating buffers once and recycling them
//! across multiple query points or robustness iterations, we significantly reduce
//! allocator pressure and improve cache locality.

// Feature-gated dependencies
#[cfg(not(feature = "std"))]
use alloc::collections::BinaryHeap;
#[cfg(not(feature = "std"))]
use alloc::vec::Vec;
#[cfg(feature = "std")]
use std::collections::BinaryHeap;
#[cfg(feature = "std")]
use std::vec::Vec;

// External dependencies
use core::fmt::Debug;
use core::ops::{Deref, DerefMut};
use num_traits::Float;

// A reusable vector slot with automatic capacity management.
#[derive(Debug, Clone)]
pub struct Slot<T>(Vec<T>);

impl<T> Slot<T> {
    // Create a new slot with the given initial capacity.
    #[inline]
    pub fn new(capacity: usize) -> Self {
        Self(Vec::with_capacity(capacity))
    }

    // Ensure the slot has at least the given capacity.
    // Grows the underlying vector if needed; never shrinks.
    #[inline]
    pub fn ensure_capacity(&mut self, capacity: usize) {
        if self.0.capacity() < capacity {
            self.0.reserve(capacity - self.0.capacity());
        }
    }

    // Clear the slot (sets length to 0, preserves capacity).
    #[inline]
    pub fn clear(&mut self) {
        self.0.clear();
    }

    // Get a reference to the underlying vector.
    #[inline]
    pub fn as_vec(&self) -> &Vec<T> {
        &self.0
    }

    // Get a mutable reference to the underlying vector.
    #[inline]
    pub fn as_vec_mut(&mut self) -> &mut Vec<T> {
        &mut self.0
    }

    // Consume the slot and return the underlying vector.
    #[inline]
    pub fn into_inner(self) -> Vec<T> {
        self.0
    }
}

impl<T> Default for Slot<T> {
    fn default() -> Self {
        Self(Vec::new())
    }
}

impl<T> Deref for Slot<T> {
    type Target = Vec<T>;
    #[inline]
    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl<T> DerefMut for Slot<T> {
    #[inline]
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.0
    }
}

impl<T> From<Vec<T>> for Slot<T> {
    fn from(v: Vec<T>) -> Self {
        Self(v)
    }
}

// A struct containing pre-allocated buffers for LOESS operations.
pub struct LoessBuffer<T: Float, N, NH> {
    // Buffer for KD-tree search state.
    pub search_buffer: NeighborhoodSearchBuffer<N>,
    // Buffer for neighbor indices and distances.
    pub neighborhood: NH,
    // Buffer for regression fitting (WLS matrices).
    pub fitting_buffer: FittingBuffer<T>,
    // Buffer for global executor state.
    pub executor_buffer: ExecutorBuffer<T>,
    // Buffer for cross-validation.
    pub cv_buffer: CVBuffer<T>,
}

impl<T: Float + Debug + Send + Sync + 'static, N, NH> LoessBuffer<T, N, NH>
where
    N: Ord,
    NH: NeighborhoodStorage,
{
    // Create a new workspace with capacities matching the expected fit parameters.
    //
    // - `n`: Total number of points.
    // - `dims`: Predictor dimensions.
    // - `k`: Expected number of neighbors (window size).
    // - `n_coeffs`: Expected number of polynomial coefficients.
    pub fn new(n: usize, dims: usize, k: usize, n_coeffs: usize) -> Self {
        Self {
            search_buffer: NeighborhoodSearchBuffer::new(k),
            neighborhood: NH::with_capacity(k),
            fitting_buffer: FittingBuffer::<T>::new(k, n_coeffs),
            executor_buffer: ExecutorBuffer::<T>::new(n, dims),
            cv_buffer: CVBuffer::<T>::new(n, dims),
        }
    }

    // Ensure all buffers have enough capacity for the given problem size.
    pub fn ensure_capacity(&mut self, n_total: usize, dims: usize, k: usize, n_coeffs: usize) {
        if self.neighborhood.capacity() < k {
            self.neighborhood = NH::with_capacity(k);
            self.search_buffer = NeighborhoodSearchBuffer::new(k);
        }
        if self.fitting_buffer.weights.capacity() < k
            || self.fitting_buffer.xtw_x.capacity() < n_coeffs * n_coeffs
        {
            self.fitting_buffer = FittingBuffer::new(k, n_coeffs);
        }

        self.executor_buffer.ensure_capacity(n_total, dims);
        self.cv_buffer.ensure_capacity(n_total, dims);
    }
}

// Trait for neighborhood storage that can be injected into the workspace.
pub trait NeighborhoodStorage {
    // Create a new neighborhood storage with given capacity.
    fn with_capacity(k: usize) -> Self;
    // Get the current capacity of the storage.
    fn capacity(&self) -> usize;
}

// Persistent buffers for KD-tree search to avoid allocations.
pub struct NeighborhoodSearchBuffer<N> {
    pub(crate) heap: BinaryHeap<N>,
    pub(crate) stack: Vec<usize>,
}

impl<N: Ord> NeighborhoodSearchBuffer<N> {
    // Create a new search buffer with capacity k.
    pub fn new(k: usize) -> Self {
        // Stack depth is bounded by tree height, typically O(log n).
        // Pre-allocate for ~1M points (log2(1M) ≈ 20).
        Self {
            heap: BinaryHeap::with_capacity(k),
            stack: Vec::with_capacity(32),
        }
    }

    // Clear all internal buffers for reuse.
    pub fn clear(&mut self) {
        self.heap.clear();
        self.stack.clear();
    }
}

// Persistent buffers for local regression to avoid allocations.
pub struct FittingBuffer<T> {
    // Weights for each neighbor.
    pub weights: Slot<T>,
    // Normal matrix X'WX.
    pub xtw_x: Slot<T>,
    // Normal vector X'WY.
    pub xtw_y: Slot<T>,
    // Column norms for equilibration.
    pub col_norms: Slot<T>,
}

impl<T> FittingBuffer<T> {
    // Create a new fitting buffer with estimated capacities.
    pub fn new(k: usize, n_coeffs: usize) -> Self {
        Self {
            weights: Slot::new(k),
            xtw_x: Slot::new(n_coeffs * n_coeffs),
            xtw_y: Slot::new(n_coeffs),
            col_norms: Slot::new(n_coeffs),
        }
    }
}

// Cache of pre-computed neighborhoods for each query point.
//
// Used to avoid repeated KD-tree searches during robustness iterations,
// since neighborhoods are invariant (only regression weights change).
#[derive(Debug, Clone)]
pub struct NeighborhoodCache<T> {
    // Cached neighborhoods: (indices, distances, max_distance) for each point.
    pub entries: Vec<CachedNeighborhood<T>>,
    // Whether the cache is populated and valid.
    pub is_valid: bool,
}

// A single cached neighborhood for one query point.
#[derive(Debug, Clone)]
pub struct CachedNeighborhood<T> {
    // Indices of the k nearest neighbors.
    pub indices: Vec<usize>,
    // Distances to each neighbor.
    pub distances: Vec<T>,
    // Maximum distance in the neighborhood (bandwidth).
    pub max_distance: T,
}

impl<T: Float> NeighborhoodCache<T> {
    // Create a new empty cache.
    pub fn new() -> Self {
        Self {
            entries: Vec::new(),
            is_valid: false,
        }
    }

    // Create a cache with pre-allocated capacity for n points.
    pub fn with_capacity(n: usize) -> Self {
        Self {
            entries: Vec::with_capacity(n),
            is_valid: false,
        }
    }
}

impl<T: Float> Default for NeighborhoodCache<T> {
    fn default() -> Self {
        Self::new()
    }
}

// Persistent buffers for global executor state.
pub struct ExecutorBuffer<T> {
    // Minimum values for each dimension.
    pub mins: Slot<T>,
    // Maximum values for each dimension.
    pub maxs: Slot<T>,
    // Normalization scales for each dimension.
    pub scales: Slot<T>,
    // Robustness weights for iterative refinement.
    pub robustness_weights: Slot<T>,
    // Residuals for iterative refinement.
    pub residuals: Slot<T>,
    // Sorted residuals for median computation.
    pub sorted_residuals: Slot<T>,
    // Cached neighborhoods to avoid repeated KD-tree searches.
    pub neighborhood_cache: NeighborhoodCache<T>,
}

impl<T: Float> ExecutorBuffer<T> {
    // Create a new executor buffer with given capacities.
    pub fn new(n: usize, dims: usize) -> Self {
        Self {
            mins: Slot::new(dims),
            maxs: Slot::new(dims),
            scales: Slot::new(dims),
            robustness_weights: Slot::new(n),
            residuals: Slot::new(n),
            sorted_residuals: Slot::new(n),
            neighborhood_cache: NeighborhoodCache::with_capacity(n),
        }
    }

    // Ensure buffers have enough capacity for given dimensions and points.
    pub fn ensure_capacity(&mut self, n: usize, dims: usize) {
        self.mins.ensure_capacity(dims);
        self.maxs.ensure_capacity(dims);
        self.scales.ensure_capacity(dims);
        self.robustness_weights.ensure_capacity(n);
        self.residuals.ensure_capacity(n);
        self.sorted_residuals.ensure_capacity(n);
    }
}

// Persistent buffers for cross-validation subsets.
pub struct CVBuffer<T> {
    // Training subset x-values.
    pub train_x: Slot<T>,
    // Training subset y-values.
    pub train_y: Slot<T>,
    // Test subset x-values.
    pub test_x: Slot<T>,
    // Test subset y-values.
    pub test_y: Slot<T>,
}

impl<T> CVBuffer<T> {
    // Create a new CV buffer with given capacities.
    pub fn new(n: usize, dims: usize) -> Self {
        Self {
            train_x: Slot::new(n * dims),
            train_y: Slot::new(n),
            test_x: Slot::new(n * dims),
            test_y: Slot::new(n),
        }
    }

    // Ensure buffers have enough capacity for given dimensions and points.
    pub fn ensure_capacity(&mut self, n: usize, dims: usize) {
        self.train_x.ensure_capacity(n * dims);
        self.train_y.ensure_capacity(n);
        self.test_x.ensure_capacity(n * dims);
        self.test_y.ensure_capacity(n);
    }
}