gbrt_rs/tree/

splitter.rs

//! Split finding algorithms for decision tree construction.
//!
//! This module provides the core logic for identifying optimal feature splits
//! in gradient boosting trees. It defines the [`Splitter`] trait for different
//! split-finding strategies and implements [`BestSplitter`] as the default
//! exhaustive search algorithm.
//!
//! # Split Finding Strategies
//!
//! The module supports two main approaches:
//!
//! - **Exact splits**: Evaluate every possible threshold between sorted feature values.
//!   Most accurate but O(n) per feature. Suitable for small to medium datasets.
//!
//! - **Approximate splits**: Use histogram binning to discretize continuous features.
//!   O(bins) per feature, much faster for large datasets with negligible accuracy loss.
//!
//! # Performance Optimizations
//!
//! [`BestSplitter`] implements several key optimizations:
//! - Precomputes cumulative sums for O(n) gain calculation
//! - Merges small histogram bins to respect `min_samples_leaf`
//! - Early termination for degenerate cases (constant features, insufficient samples)
//! - Feature subsampling via `feature_indices` parameter
//!
//! # Split Candidate Evaluation
//!
//! Splits are evaluated using the gain formula from the objective function's
//! second-order Taylor approximation. The gain measures improvement over the parent node:
//!
//! ```text
//! Gain = LeftGain + RightGain - ParentGain
//! ```
//!
//! where `NodeGain = (sum_gradients)² / (sum_hessians + lambda)`
//!
//! [`Splitter`]: trait.Splitter.html
//! [`BestSplitter`]: struct.BestSplitter.html

use crate::data::FeatureMatrix;
use thiserror::Error;

/// Errors that can occur during split finding.
///
/// These errors cover invalid inputs, insufficient samples, and data integrity
/// issues that prevent finding valid splits.
#[derive(Error, Debug)]
pub enum SplitterError {
    /// Requested feature index is out of bounds.
    #[error("Invalid feature index: {index} (max: {max})")]
    InvalidFeatureIndex { index: usize, max: usize },
    
    /// Too few samples to perform a split meeting minimum leaf size requirements.
    #[error("Insufficient samples for splitting: {samples} (min: {min})")]
    InsufficientSamples { samples: usize, min: usize },
    
    /// A feature had no valid split points (e.g., all values identical).
    #[error("No valid splits found for feature {feature}")]
    NoValidSplits { feature: usize },
    
    /// Underlying data access error from feature matrix.
    #[error("Data error: {0}")]
    DataError(#[from] crate::data::DataError),
   
    /// Sample array is empty after filtering.
    #[error("Empty samples array")]
    EmptySamples,
}

/// Represents a candidate split with precomputed statistics.
///
/// This struct stores all information needed to evaluate and apply a split,
/// including the feature, threshold, gain, and sample indices for each child.
#[derive(Debug, Clone)]
pub struct SplitCandidate {
    /// Index of the feature to split on.
    pub feature_index: usize,
    /// Threshold value for the split (samples <= go left, > go right).
    pub split_value: f64,
    /// Improvement in loss from this split.
    pub gain: f64,
    /// Indices of samples assigned to the left child.
    pub left_indices: Vec<usize>,
    /// Indices of samples assigned to the right child.
    pub right_indices: Vec<usize>,
    /// Sum of gradients in the left child.
    pub left_grad_sum: f64,
    /// Sum of gradients in the right child.
    pub right_grad_sum: f64,
    /// Sum of hessians in the left child.
    pub left_hess_sum: f64,
    /// Sum of hessians in the right child.
    pub right_hess_sum: f64,
}

impl SplitCandidate {
    /// Creates a new split candidate with all required statistics.
    ///
    /// # Arguments
    ///
    /// * `feature_index` - Feature to split on
    /// * `split_value` - Threshold value
    /// * `gain` - Split improvement score
    /// * `left_indices` - Samples in left child
    /// * `right_indices` - Samples in right child
    /// * `left_grad_sum` - Gradient sum for left child
    /// * `right_grad_sum` - Gradient sum for right child
    /// * `left_hess_sum` - Hessian sum for left child
    /// * `right_hess_sum` - Hessian sum for right child
    pub fn new(
        feature_index: usize,
        split_value: f64,
        gain: f64,
        left_indices: Vec<usize>,
        right_indices: Vec<usize>,
        left_grad_sum: f64,
        right_grad_sum: f64,
        left_hess_sum: f64,
        right_hess_sum: f64,
    ) -> Self {
        Self {
            feature_index,
            split_value,
            gain,
            left_indices,
            right_indices,
            left_grad_sum,
            right_grad_sum,
            left_hess_sum,
            right_hess_sum,
        }
    }
   
    /// Validates that this split meets minimum leaf size requirements.
    ///
    /// # Arguments
    ///
    /// * `min_samples_leaf` - Minimum samples required in each child
    ///
    /// # Returns
    ///
    /// `true` if both children have at least `min_samples_leaf` samples
    pub fn is_valid(&self, min_samples_leaf: usize) -> bool {
        self.left_indices.len() >= min_samples_leaf && 
        self.right_indices.len() >= min_samples_leaf
    }
}

/// Trait for pluggable split-finding algorithms.
///
/// Implementations define different strategies for finding optimal feature splits,
/// enabling experimentation with exact search, approximate methods, or random splits.
pub trait Splitter: Send + Sync {
    /// Finds the best split across the specified features.
    ///
    /// This method searches for the optimal (feature, threshold) pair that maximizes
    /// the split gain while respecting `min_samples_leaf` and other constraints.
    ///
    /// # Arguments
    ///
    /// * `features` - Training feature matrix
    /// * `gradients` - First derivatives from the objective function
    /// * `hessians` - Second derivatives from the objective function
    /// * `feature_indices` - Subset of features to consider (enables feature sampling)
    /// * `min_samples_leaf` - Minimum samples in each child
    /// * `lambda` - L2 regularization parameter
    ///
    /// # Returns
    ///
    /// `Ok(Some(candidate))` if a valid split is found, `Ok(None)` if no valid split exists,
    /// or a [`SplitterError`] if computation fails.
    fn find_best_split(
        &self,
        features: &FeatureMatrix,
        gradients: &[f64],
        hessians: &[f64],
        feature_indices: &[usize],
        min_samples_leaf: usize,
        lambda: f64,
    ) -> Result<Option<SplitCandidate>, SplitterError>;
}

/// Default splitter that exhaustively searches for the optimal split.
///
/// `BestSplitter` evaluates all valid split points for each feature, computing
/// the exact gain using cumulative sums. For large datasets, it can use histogram
/// approximation to trade a small amount of accuracy for significant speedup.
///
/// # Configuration
///
/// - `n_bins`: Number of histogram bins for approximate splitting (default: 256)
/// - `use_exact_splits`: Force exact evaluation even for large features
///
/// # Algorithm
///
/// For each feature:
/// 1. Sort samples by feature value (O(n log n))
/// 2. Precompute cumulative gradient/hessian sums (O(n))
/// 3. Evaluate all splits meeting `min_samples_leaf` (O(n))
/// 4. Track the split with maximum gain
///
/// The overall complexity is O(k × n log n) where k is the number of features considered.
#[derive(Debug, Clone)]
pub struct BestSplitter {
    /// Number of bins for histogram approximation (None = exact).
    pub n_bins: Option<usize>,
    /// Whether to disable approximation and use exact splits.
    pub use_exact_splits: bool,
}

impl Default for BestSplitter {
    fn default() -> Self {
        Self {
            n_bins: Some(256),
            use_exact_splits: false,
        }
    }
}

impl BestSplitter {
    /// Creates a new `BestSplitter` with default settings (256 bins, approximate).
    pub fn new() -> Self {
        Self::default()
    }
   
    /// Sets the number of histogram bins for approximate splitting.
    ///
    /// More bins = more accurate but slower. Fewer bins = faster but coarser approximations.
    ///
    /// # Arguments
    ///
    /// * `n_bins` - Number of histogram bins (must be ≥ 2)
    ///
    /// # Returns
    ///
    /// A new `BestSplitter` instance
    pub fn with_n_bins(n_bins: usize) -> Self {
        Self {
            n_bins: Some(n_bins),
            use_exact_splits: false,
        }
    }
   
    /// Enables or disables exact split evaluation.
    ///
    /// When `true`, all splits are evaluated exactly regardless of dataset size.
    /// When `false`, histogram approximation is used for large features.
    ///
    /// # Arguments
    ///
    /// * `use_exact` - Whether to force exact evaluation
    ///
    /// # Returns
    ///
    /// Self with updated setting (builder pattern)
    pub fn use_exact_splits(mut self, use_exact: bool) -> Self {
        self.use_exact_splits = use_exact;
        self
    }
   
    /// Finds the best split for a single feature.
    ///
    /// This internal method delegates to either exact or approximate search
    /// based on the splitter's configuration.
    ///
    /// # Arguments
    ///
    /// * `features` - Feature matrix
    /// * `gradients` - Gradient values
    /// * `hessians` - Hessian values
    /// * `feature_index` - Feature to evaluate
    /// * `min_samples_leaf` - Minimum leaf size
    /// * `lambda` - L2 regularization
    ///
    /// # Returns
    ///
    /// The best split candidate for this feature, or `None` if no valid split exists
    fn find_best_split_for_feature(
        &self,
        features: &FeatureMatrix,
        gradients: &[f64],
        hessians: &[f64],
        feature_index: usize,
        min_samples_leaf: usize,
        lambda: f64,
    ) -> Result<Option<SplitCandidate>, SplitterError> {
        let n_samples = features.n_samples();
        
        if n_samples < min_samples_leaf * 2 {
            return Ok(None);
        }
        
        // Get feature values with error handling
        let mut samples: Vec<(f64, f64, f64, usize)> = (0..n_samples)
            .filter_map(|i| {
                match features.get(i, feature_index) {
                    Ok(feature_val) => Some((feature_val, gradients[i], hessians[i], i)),
                    Err(_) => None,
                }
            })
            .collect();
        
        if samples.is_empty() {
            return Err(SplitterError::EmptySamples);
        }
        
        // Sort by feature value with safe comparison
        samples.sort_by(|a, b| {
            a.0.partial_cmp(&b.0).unwrap_or(std::cmp::Ordering::Equal)
        });
        
        // Remove duplicate feature values at boundaries to avoid degenerate splits
        samples.dedup_by(|a, b| (a.0 - b.0).abs() < 1e-10);
        
        if samples.len() < min_samples_leaf * 2 {
            return Ok(None);
        }
        
        if self.use_exact_splits {
            self.find_best_exact_split(&samples, feature_index, min_samples_leaf, lambda)
        } else {
            self.find_best_approximate_split(&samples, feature_index, min_samples_leaf, lambda)
        }
    }
    
    /// Exhaustively evaluates all possible split points.
    ///
    /// # Arguments
    ///
    /// * `samples` - Sorted samples: (feature_value, gradient, hessian, index)
    /// * `feature_index` - Feature being evaluated
    /// * `min_samples_leaf` - Minimum samples per child
    /// * `lambda` - L2 regularization
    ///
    /// # Returns
    ///
    /// Best split candidate found via exhaustive search
    fn find_best_exact_split(
        &self,
        samples: &[(f64, f64, f64, usize)],
        feature_index: usize,
        min_samples_leaf: usize,
        lambda: f64,
    ) -> Result<Option<SplitCandidate>, SplitterError> {
        let n_samples = samples.len();
        let mut best_gain = -f64::INFINITY;
        let mut best_split: Option<SplitCandidate> = None;
        
        // Precompute cumulative sums
        let mut grad_prefix = Vec::with_capacity(n_samples + 1);
        let mut hess_prefix = Vec::with_capacity(n_samples + 1);
        
        grad_prefix.push(0.0);
        hess_prefix.push(0.0);
        
        for i in 0..n_samples {
            grad_prefix.push(grad_prefix[i] + samples[i].1);
            hess_prefix.push(hess_prefix[i] + samples[i].2);
        }
        
        let total_grad = grad_prefix[n_samples];
        let total_hess = hess_prefix[n_samples];
        
        if total_hess + lambda <= 0.0 {
            return Ok(None);
        }
        
        let parent_gain = self.compute_gain(total_grad, total_hess, lambda);
        
        // Try all possible splits with bounds checking
        let max_i = n_samples.saturating_sub(min_samples_leaf);
        for i in min_samples_leaf..=max_i {
            // Ensure we don't go out of bounds
            if i >= n_samples {
                break;
            }
            
            // Skip if feature value is too close to next (degenerate)
            if i < n_samples - 1 {
                let diff = (samples[i].0 - samples[i + 1].0).abs();
                if diff < 1e-10 {
                    continue;
                }
            }
            
            let left_grad = grad_prefix[i];
            let left_hess = hess_prefix[i];
            let right_grad = total_grad - left_grad;
            let right_hess = total_hess - left_hess;
            
            // Avoid division by zero
            if left_hess + lambda <= 0.0 || right_hess + lambda <= 0.0 {
                continue;
            }
            
            let left_gain = self.compute_gain(left_grad, left_hess, lambda);
            let right_gain = self.compute_gain(right_grad, right_hess, lambda);
            let gain = left_gain + right_gain - parent_gain;
            
            // Only consider positive gains
            if gain > best_gain && gain > 1e-10 {
                best_gain = gain;
                
                let split_value = if i < n_samples - 1 {
                    (samples[i].0 + samples[i + 1].0) / 2.0
                } else {
                    samples[i].0
                };
                
                let left_indices: Vec<usize> = samples[0..i].iter().map(|&(_, _, _, idx)| idx).collect();
                let right_indices: Vec<usize> = samples[i..].iter().map(|&(_, _, _, idx)| idx).collect();
                
                // Validate split sizes
                if left_indices.len() >= min_samples_leaf && right_indices.len() >= min_samples_leaf {
                    best_split = Some(SplitCandidate::new(
                        feature_index,
                        split_value,
                        gain,
                        left_indices,
                        right_indices,
                        left_grad,
                        right_grad,
                        left_hess,
                        right_hess,
                    ));
                }
            }
        }
        
        Ok(best_split)
    }
    
    /// Uses histogram binning for faster approximate split finding.
    ///
    /// # Arguments
    ///
    /// * `samples` - Sorted samples: (feature_value, gradient, hessian, index)
    /// * `feature_index` - Feature being evaluated
    /// * `min_samples_leaf` - Minimum samples per child
    /// * `lambda` - L2 regularization
    ///
    /// # Returns
    ///
    /// Best split candidate found via histogram approximation
    fn find_best_approximate_split(
        &self,
        samples: &[(f64, f64, f64, usize)],
        feature_index: usize,
        min_samples_leaf: usize,
        lambda: f64,
    ) -> Result<Option<SplitCandidate>, SplitterError> {
        let n_samples = samples.len();
        let n_bins = self.n_bins.unwrap_or(256).min(n_samples);

        if n_bins < 2 {
            return self.find_best_exact_split(samples, feature_index, min_samples_leaf, lambda);
        }

        // Create bins
        let min_val = samples[0].0;
        let max_val = samples[n_samples - 1].0;
        if (max_val - min_val).abs() < 1e-10 {
            return Ok(None);
        }

        let bin_size = (max_val - min_val) / (n_bins as f64);
        let mut bins: Vec<(f64, f64, f64, f64, Vec<usize>)> = Vec::with_capacity(n_bins);

        for _ in 0..n_bins {
            bins.push((f64::INFINITY, f64::NEG_INFINITY, 0.0, 0.0, Vec::new()));
        }

        // Accumulate into bins
        for &(feature_val, grad, hess, idx) in samples {
            let bin_idx = (((feature_val - min_val) / bin_size) as usize).min(n_bins - 1);
            let bin = &mut bins[bin_idx];
            bin.0 = bin.0.min(feature_val);
            bin.1 = bin.1.max(feature_val);
            bin.2 += grad;
            bin.3 += hess;
            bin.4.push(idx);
        }

        // Filter empty bins
        let non_empty_bins: Vec<(f64, f64, f64, f64, Vec<usize>)> = bins
            .into_iter()
            .filter(|bin| !bin.4.is_empty())
            .collect();

        if non_empty_bins.len() < 2 {
            return Ok(None);
        }

        // Merge small bins
        let mut merged_bins = Vec::new();
        let mut current_bin = (f64::INFINITY, f64::NEG_INFINITY, 0.0, 0.0, Vec::new());

        for bin in non_empty_bins {
            if current_bin.4.len() + bin.4.len() < min_samples_leaf && !current_bin.4.is_empty() {
                current_bin.0 = current_bin.0.min(bin.0);
                current_bin.1 = current_bin.1.max(bin.1);
                current_bin.2 += bin.2;
                current_bin.3 += bin.3;
                current_bin.4.extend(bin.4);
            } else {
                if !current_bin.4.is_empty() {
                    merged_bins.push(current_bin);
                }
                current_bin = bin;
            }
        }

        if !current_bin.4.is_empty() {
            merged_bins.push(current_bin);
        }

        if merged_bins.len() < 2 {
            return Ok(None);
        }

        // Find best split among bins
        let mut best_gain = -f64::INFINITY;
        let mut best_split: Option<SplitCandidate> = None;

        let total_grad: f64 = merged_bins.iter().map(|bin| bin.2).sum();
        let total_hess: f64 = merged_bins.iter().map(|bin| bin.3).sum();
        
        if total_hess + lambda <= 0.0 {
            return Ok(None);
        }
        
        let parent_gain = self.compute_gain(total_grad, total_hess, lambda);

        let mut left_grad = 0.0;
        let mut left_hess = 0.0;
        let mut left_indices = Vec::new();

        for i in 0..(merged_bins.len() - 1) {
            left_grad += merged_bins[i].2;
            left_hess += merged_bins[i].3;
            left_indices.extend(&merged_bins[i].4);

            if left_indices.len() < min_samples_leaf {
                continue;
            }

            let right_grad = total_grad - left_grad;
            let right_hess = total_hess - left_hess;
            let right_indices: Vec<usize> = merged_bins[i + 1..]
                .iter()
                .flat_map(|bin| bin.4.iter().cloned())
                .collect();

            if right_indices.len() < min_samples_leaf {
                continue;
            }

            if left_hess + lambda <= 0.0 || right_hess + lambda <= 0.0 {
                continue;
            }
            
            let left_gain = self.compute_gain(left_grad, left_hess, lambda);
            let right_gain = self.compute_gain(right_grad, right_hess, lambda);
            let gain = left_gain + right_gain - parent_gain;

            if gain > best_gain && gain > 1e-10 {
                best_gain = gain;

                let split_value = (merged_bins[i].1 + merged_bins[i + 1].0) / 2.0;

                best_split = Some(SplitCandidate::new(
                    feature_index,
                    split_value,
                    gain,
                    left_indices.clone(),
                    right_indices,
                    left_grad,
                    right_grad,
                    left_hess,
                    right_hess,
                ));
            }
        }

        Ok(best_split)
    }
    
    /// Computes the gain for a node given its gradient/hessian sums.
    ///
    /// # Formula
    ///
    /// `gain = (sum_grad)² / (sum_hess + lambda)`
    ///
    /// # Arguments
    ///
    /// * `sum_grad` - Sum of gradients in the node
    /// * `sum_hess` - Sum of hessians in the node
    /// * `lambda` - L2 regularization parameter
    ///
    /// # Returns
    ///
    /// Node gain, or `-∞` if the denominator would be non-positive
    fn compute_gain(&self, sum_grad: f64, sum_hess: f64, lambda: f64) -> f64 {
        if sum_hess + lambda <= 0.0 {
            -f64::INFINITY
        } else {
            (sum_grad * sum_grad) / (sum_hess + lambda)
        }
    }
}

impl Splitter for BestSplitter {
    /// Finds the globally optimal split across all specified features.
    ///
    /// This method iterates through `feature_indices`, finds the best split for each,
    /// and returns the candidate with maximum gain. It respects all regularization
    /// and sampling constraints.
    ///
    /// # Arguments
    ///
    /// * `features` - Feature matrix
    /// * `gradients` - Gradient values
    /// * `hessians` - Hessian values
    /// * `feature_indices` - Features to evaluate
    /// * `min_samples_leaf` - Minimum samples per child
    /// * `lambda` - L2 regularization
    ///
    /// # Returns
    ///
    /// Best split across all features, or `None` if no valid split exists
    fn find_best_split(
        &self,
        features: &FeatureMatrix,
        gradients: &[f64],
        hessians: &[f64],
        feature_indices: &[usize],
        min_samples_leaf: usize,
        lambda: f64,
    ) -> Result<Option<SplitCandidate>, SplitterError> {
        if features.n_samples() != gradients.len() || gradients.len() != hessians.len() {
            return Err(SplitterError::InsufficientSamples {
                samples: gradients.len(),
                min: features.n_samples(),
            });
        }
        
        if features.n_samples() < min_samples_leaf * 2 {
            return Ok(None);
        }
        
        let mut best_candidate: Option<SplitCandidate> = None;
        
        // Add progress tracking for large datasets
        if features.n_features() > 100 && features.n_samples() > 10000 {
            // Potentially log progress here
        }
        
        for (idx, &feature_idx) in feature_indices.iter().enumerate() {
            // Log progress every 10 features for large datasets
            if idx > 0 && idx % 10 == 0 && features.n_features() > 50 {
                // Progress log
            }
            
            if feature_idx >= features.n_features() {
                return Err(SplitterError::InvalidFeatureIndex {
                    index: feature_idx,
                    max: features.n_features() - 1,
                });
            }
            
            if let Some(candidate) = self.find_best_split_for_feature(
                features,
                gradients,
                hessians,
                feature_idx,
                min_samples_leaf,
                lambda,
            )? {
                if best_candidate.as_ref().map(|c| c.gain < candidate.gain).unwrap_or(true) {
                    best_candidate = Some(candidate);
                }
            }
        }
        
        Ok(best_candidate)
    }
}