gbrt_rs/tree/

criterion.rs

//! Split criteria for decision tree construction in gradient boosting.
//!
//! This module defines traits and implementations for evaluating the quality of potential
//! splits when building decision trees. Split criteria determine which feature and threshold
//! provide the best separation of data to minimize the loss function.
//!
//! # Key Concepts
//!
//! In gradient boosting, split criteria operate on **gradients** and **hessians** (first and
//! second derivatives of the loss function) rather than raw target values. This allows the
//! same tree-building algorithm to work with any differentiable loss function.
//!
//! # Available Criteria
//!
//! - [`MSECriterion`]: Basic mean squared error criterion with L2 regularization
//! - [`FriedmanMSECriterion`]: Enhanced MSE criterion with L1/L2 regularization and
//!   minimum gain threshold (used in XGBoost and LightGBM)
//!
//! # Regularization
//!
//! Both criteria support regularization parameters to prevent overfitting:
//! - `lambda`: L2 regularization on leaf weights
//! - `alpha`: L1 regularization (Friedman only)
//! - `gamma`: Minimum gain required to make a split (Friedman only)
//! - `min_hessian`: Minimum sum of hessians required for a valid split

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

/// Errors that can occur during split criterion computation.
///
/// These errors cover invalid inputs, insufficient samples, and numerical issues
/// that prevent proper calculation of split quality.
#[derive(Error, Debug)]
pub enum CriterionError {
    /// Input data is invalid for computation.
    ///
    /// This includes mismatched array lengths, non-finite values, or other
    /// data integrity issues that violate the criterion's preconditions.
    #[error("Invalid input data: {0}")]
    InvalidInput(String),
   
    /// Not enough samples to perform a meaningful split.
    ///
    /// This error is returned when a node or split candidate has fewer samples
    /// than the minimum required for stable statistical estimation.
    #[error("Insufficient samples: {samples} (min: {min})")]
    InsufficientSamples { samples: usize, min: usize },
}

/// Core trait for evaluating split quality in decision trees.
///
/// Split criteria determine how well a candidate split separates the data
/// in terms of reducing the loss function. They operate on gradients and hessians
/// from the objective function's derivatives.
///
/// # Type Safety
///
/// Implementations must be thread-safe (`Send + Sync`) as they may be used
/// in parallel tree construction.
pub trait SplitCriterion: Send + Sync {
    /// Computes the improvement (gain) from making a split.
    ///
    /// The gain measures how much the loss is reduced by splitting a parent node
    /// into left and right children. Higher gain indicates a better split.
    ///
    /// # Arguments
    ///
    /// * `features` - Feature matrix of the training data
    /// * `gradients` - First derivatives of the loss w.r.t predictions (one per sample)
    /// * `hessians` - Second derivatives of the loss w.r.t predictions (one per sample)
    /// * `left_indices` - Indices of samples assigned to the left child node
    /// * `right_indices` - Indices of samples assigned to the right child node
    ///
    /// # Returns
    ///
    /// The split gain (improvement score). Should be non-negative.
    ///
    /// # Errors
    ///
    /// Returns [`CriterionError::InvalidInput`] if gradients and hessians have different lengths
    /// Returns [`CriterionError::InsufficientSamples`] if either child has too few samples
    fn compute_gain(
        &self,
        features: &FeatureMatrix,
        gradients: &[f64],
        hessians: &[f64],
        left_indices: &[usize],
        right_indices: &[usize],
    ) -> Result<f64, CriterionError>;
    
    /// Computes the optimal prediction value for a leaf node.
    ///
    /// The leaf value minimizes the loss for the samples reaching that leaf.
    /// For MSE-based criteria, this is typically `-sum(gradients) / sum(hessians + lambda)`.
    ///
    /// # Arguments
    ///
    /// * `gradients` - First derivatives of samples in the leaf
    /// * `hessians` - Second derivatives of samples in the leaf
    ///
    /// # Returns
    ///
    /// The optimal leaf prediction value
    ///
    /// # Errors
    ///
    /// Returns [`CriterionError::InsufficientSamples`] if the leaf is empty
    fn compute_leaf_value(
        &self,
        gradients: &[f64],
        hessians: &[f64],
    ) -> Result<f64, CriterionError>;
    
    /// Returns the name of the criterion.
    ///
    /// Used for logging, debugging, and serialization.
    ///
    /// # Returns
    ///
    /// A static string slice identifying the criterion (e.g., "MSE", "FriedmanMSE") 
    fn name(&self) -> &str;
}


/// Standard Mean Squared Error split criterion with L2 regularization.
///
/// This criterion computes split gain based on the reduction in squared error,
/// treating gradients as residuals and hessians as constant weights (typically 1.0).
///
/// # Formula
///
/// - Node gain: `(sum_grad)² / (sum_hess + lambda)`
/// - Leaf value: `-sum_grad / (sum_hess + lambda)`
///
/// # Parameters
///
/// - `lambda`: L2 regularization term added to the denominator to prevent division by zero
///   and control leaf value magnitude
/// - `min_hessian`: Minimum sum of hessians required for a valid split (avoids unstable splits)
#[derive(Debug, Clone)]
pub struct MSECriterion {
    /// L2 regularization parameter (lambda) for leaf weights.
    ///
    /// Larger values shrink leaf predictions toward zero, preventing overfitting.
    pub lambda: f64,
    /// Minimum sum of hessians required to consider a split valid.
    ///
    /// This prevents splits on nodes with too few effective samples.
    pub min_hessian: f64,
}

impl Default for MSECriterion {
    fn default() -> Self {
        Self {
            lambda: 1.0,
            min_hessian: 1e-8,
        }
    }
}

impl MSECriterion {
    /// Creates a new MSE criterion with the specified L2 regularization.
    ///
    /// # Arguments
    ///
    /// * `lambda` - L2 regularization parameter (typically 1.0)
    ///
    /// # Returns
    ///
    /// A new `MSECriterion` instance with default `min_hessian = 1e-8`
    pub fn new(lambda: f64) -> Self {
        Self {
            lambda,
            min_hessian: 1e-8,
        }
    }
    
    /// Sets the minimum hessian threshold (builder pattern).
    ///
    /// # Arguments
    ///
    /// * `min_hessian` - Minimum sum of hessians for valid splits
    ///
    /// # Returns
    ///
    /// Self with the updated threshold
    pub fn with_min_hessian(mut self, min_hessian: f64) -> Self {
        self.min_hessian = min_hessian;
        self
    }
    
    /// Computes the gain for a specific node (without children).
    ///
    /// This helper method calculates the base gain used for both parent
    /// and child nodes in split evaluation.
    ///
    /// # Arguments
    ///
    /// * `sum_grad` - Sum of gradients in the node
    /// * `sum_hess` - Sum of hessians in the node
    ///
    /// # Returns
    ///
    /// Node gain score (non-negative) 
    fn compute_node_gain(&self, sum_grad: f64, sum_hess: f64) -> f64 {
        if sum_hess < self.min_hessian {
            return 0.0;
        }
        (sum_grad * sum_grad) / (sum_hess + self.lambda)
    }
}

impl SplitCriterion for MSECriterion {
    fn compute_gain(
        &self,
        _features: &FeatureMatrix,
        gradients: &[f64],
        hessians: &[f64],
        left_indices: &[usize],
        right_indices: &[usize],
    ) -> Result<f64, CriterionError> {
        if gradients.len() != hessians.len() {
            return Err(CriterionError::InvalidInput(
                "Gradients and hessians must have the same length".to_string()
            ));
        }
        
        if left_indices.is_empty() || right_indices.is_empty() {
            return Err(CriterionError::InsufficientSamples {
                samples: left_indices.len().min(right_indices.len()),
                min: 1,
            });
        }
        
        // Compute parent statistics
        let parent_grad: f64 = gradients.iter().sum();
        let parent_hess: f64 = hessians.iter().sum();
        let parent_gain = self.compute_node_gain(parent_grad, parent_hess);
        
        // Compute left child statistics
        let left_grad: f64 = left_indices.iter().map(|&i| gradients[i]).sum();
        let left_hess: f64 = left_indices.iter().map(|&i| hessians[i]).sum();
        let left_gain = self.compute_node_gain(left_grad, left_hess);
        
        // Compute right child statistics
        let right_grad: f64 = right_indices.iter().map(|&i| gradients[i]).sum();
        let right_hess: f64 = right_indices.iter().map(|&i| hessians[i]).sum();
        let right_gain = self.compute_node_gain(right_grad, right_hess);
        
        // Gain is the improvement over the parent
        Ok(left_gain + right_gain - parent_gain)
    }
    
    fn compute_leaf_value(
        &self,
        gradients: &[f64],
        hessians: &[f64],
    ) -> Result<f64, CriterionError> {
        if gradients.is_empty() {
            return Err(CriterionError::InsufficientSamples {
                samples: 0,
                min: 1,
            });
        }
        
        let sum_grad: f64 = gradients.iter().sum();
        let sum_hess: f64 = hessians.iter().sum();
        
        if sum_hess.abs() < self.min_hessian {
            return Ok(0.0);
        }
        
        // For MSE loss, the optimal leaf value is -sum_grad / (sum_hess + lambda)
        Ok(-sum_grad / (sum_hess + self.lambda))
    }
    
    fn name(&self) -> &str {
        "MSE"
    }
}

/// Friedman's MSE criterion with regularization (XGBoost-style).
///
/// This enhanced criterion adds L1 regularization, minimum gain threshold,
/// and proper handling of regularized leaf values as used in modern
/// gradient boosting frameworks.
///
/// # Formula
///
/// - Node gain: `(reg_grad)² / (sum_hess + lambda)` where `reg_grad` is L1-regularized
/// - Leaf value: `-reg_grad / (sum_hess + lambda)`
/// - Split gain: `left_gain + right_gain - parent_gain - gamma`
///
/// # Regularization Parameters
///
/// - `lambda`: L2 regularization on leaf weights
/// - `alpha`: L1 regularization (shrinkage effect)
/// - `gamma`: Minimum gain required to make a split (pruning)
/// - `min_hessian`: Minimum hessian sum for numerical stability
#[derive(Debug, Clone)]
pub struct FriedmanMSECriterion {
    /// L2 regularization parameter for leaf weights.
    ///
    /// Controls the magnitude of leaf predictions to prevent overfitting.
    pub lambda: f64,
    /// L1 regularization parameter for sparse models.
    ///
    /// Encourages leaf values to be exactly zero, creating sparse trees.
    pub alpha: f64,
    /// Minimum sum of hessians required for a valid split.
    ///
    /// Prevents splits on nodes with insufficient statistical support.
    pub min_hessian: f64,
    /// Minimum gain required to create a split (gamma parameter).
    ///
    /// Acts as pre-pruning: splits with gain < gamma are rejected.
    pub gamma: f64,
}

impl Default for FriedmanMSECriterion {
    fn default() -> Self {
        Self {
            lambda: 1.0,
            alpha: 0.0,
            min_hessian: 1e-8,
            gamma: 0.0,
        }
    }
}

impl FriedmanMSECriterion {
    /// Creates a new Friedman MSE criterion with L2 and minimum gain regularization.
    ///
    /// # Arguments
    ///
    /// * `lambda` - L2 regularization parameter (typically 1.0-2.0)
    /// * `gamma` - Minimum gain threshold for splits (0.0 for no pruning)
    ///
    /// # Returns
    ///
    /// A new `FriedmanMSECriterion` with default `alpha = 0.0` and `min_hessian = 1e-8`
    pub fn new(lambda: f64, gamma: f64) -> Self {
        Self {
            lambda,
            alpha: 0.0,
            min_hessian: 1e-8,
            gamma,
        }
    }
   
    /// Sets the L1 regularization parameter (builder pattern).
    ///
    /// # Arguments
    ///
    /// * `alpha` - L1 regularization parameter (useful for sparse data)
    ///
    /// # Returns
    ///
    /// Self with updated alpha
    pub fn with_alpha(mut self, alpha: f64) -> Self {
        self.alpha = alpha;
        self
    }
    
    /// Computes the gain for a node with L1 regularization applied.
    ///
    /// L1 regularization modifies the gradient by shrinking it toward zero,
    /// which can produce sparse leaf values.
    ///
    /// # Arguments
    ///
    /// * `sum_grad` - Sum of gradients in the node
    /// * `sum_hess` - Sum of hessians in the node
    ///
    /// # Returns
    ///
    /// Regularized node gain 
    fn compute_node_gain(&self, sum_grad: f64, sum_hess: f64) -> f64 {
        if sum_hess < self.min_hessian {
            return 0.0;
        }
        
        // Apply L1 regularization (like in XGBoost)
        let reg_sum_grad = if sum_grad >= 0.0 {
            (sum_grad - self.alpha).max(0.0)
        } else {
            (sum_grad + self.alpha).min(0.0)
        };
        
        (reg_sum_grad * reg_sum_grad) / (sum_hess + self.lambda)
    }
    
    /// Computes the optimal leaf value with L1/L2 regularization.
    ///
    /// # Arguments
    ///
    /// * `sum_grad` - Sum of gradients in the leaf
    /// * `sum_hess` - Sum of hessians in the leaf
    ///
    /// # Returns
    ///
    /// Regularized leaf prediction value    
    fn compute_leaf_value_with_reg(&self, sum_grad: f64, sum_hess: f64) -> f64 {
        if sum_hess.abs() < self.min_hessian {
            return 0.0;
        }
        
        // Apply L1 regularization
        let reg_sum_grad = if sum_grad >= 0.0 {
            (sum_grad - self.alpha).max(0.0)
        } else {
            (sum_grad + self.alpha).min(0.0)
        };
        
        -reg_sum_grad / (sum_hess + self.lambda)
    }
}

impl SplitCriterion for FriedmanMSECriterion {
    fn compute_gain(
        &self,
        _features: &FeatureMatrix,
        gradients: &[f64],
        hessians: &[f64],
        left_indices: &[usize],
        right_indices: &[usize],
    ) -> Result<f64, CriterionError> {
        if gradients.len() != hessians.len() {
            return Err(CriterionError::InvalidInput(
                "Gradients and hessians must have the same length".to_string()
            ));
        }
        
        if left_indices.is_empty() || right_indices.is_empty() {
            return Err(CriterionError::InsufficientSamples {
                samples: left_indices.len().min(right_indices.len()),
                min: 1,
            });
        }
        
        // Compute parent statistics
        let parent_grad: f64 = gradients.iter().sum();
        let parent_hess: f64 = hessians.iter().sum();
        let parent_gain = self.compute_node_gain(parent_grad, parent_hess);
        
        // Compute left child statistics
        let left_grad: f64 = left_indices.iter().map(|&i| gradients[i]).sum();
        let left_hess: f64 = left_indices.iter().map(|&i| hessians[i]).sum();
        let left_gain = self.compute_node_gain(left_grad, left_hess);
        
        // Compute right child statistics
        let right_grad: f64 = right_indices.iter().map(|&i| gradients[i]).sum();
        let right_hess: f64 = right_indices.iter().map(|&i| hessians[i]).sum();
        let right_gain = self.compute_node_gain(right_grad, right_hess);
        
        // Gain is the improvement over the parent, minus gamma (minimum gain required)
        let gain = left_gain + right_gain - parent_gain - self.gamma;
        Ok(gain.max(0.0))
    }
    
    fn compute_leaf_value(
        &self,
        gradients: &[f64],
        hessians: &[f64],
    ) -> Result<f64, CriterionError> {
        if gradients.is_empty() {
            return Err(CriterionError::InsufficientSamples {
                samples: 0,
                min: 1,
            });
        }
        
        let sum_grad: f64 = gradients.iter().sum();
        let sum_hess: f64 = hessians.iter().sum();
        
        Ok(self.compute_leaf_value_with_reg(sum_grad, sum_hess))
    }
    
    fn name(&self) -> &str {
        "FriedmanMSE"
    }
}

/// Factory function to create split criterion instances.
///
/// This convenience function creates criterion objects from string names,
/// supporting both built-in criteria and custom regularization parameters.
///
/// # Arguments
///
/// * `name` - Name of the criterion: "friedman_mse" or "mse"
/// * `lambda` - L2 regularization parameter
/// * `gamma` - Minimum gain threshold (only used by FriedmanMSECriterion)
///
/// # Returns
///
/// A boxed trait object implementing [`SplitCriterion`]
pub fn create_criterion(name: &str, lambda: f64, gamma: f64) -> Box<dyn SplitCriterion> {
    match name {
        "friedman_mse" => Box::new(FriedmanMSECriterion::new(lambda, gamma)),
        "mse" | _ => Box::new(MSECriterion::new(lambda)),
    }
}