gbrt_rs/boosting/

gradient_booster.rs

//! Core Gradient Boosting Implementation
//! 
//! This module provides the main [`GradientBooster`] algorithm implementation with support for:
//! 
//! - **Multiple loss functions**: MSE, MAE, Huber, and LogLoss
//! - **Stochastic gradient boosting**: Subsampling for improved generalization
//! - **Early stopping**: With configurable patience and tolerance
//! - **Validation monitoring**: Optional validation dataset for loss tracking
//! - **Feature importance**: Computed from tree gain contributions
//! - **Robust training**: Handles edge cases like NaN loss values
//! 
//! # Training Process
//! 
//! The booster trains an ensemble of decision trees iteratively:
//! 1. Initialize predictions with the optimal constant value
//! 2. For each iteration:
//!    - Compute gradients and hessians of the loss function
//!    - Sample data (if subsampling < 1.0)
//!    - Fit a new tree to the negative gradients
//!    - Update predictions with shrinkage (learning rate)
//!    - Monitor validation loss for early stopping
use crate::core::{GBRTConfig, LossFunction};
use crate::tree::Tree;
use crate::data::{Dataset, FeatureMatrix};
use crate::core::{GradientLoss, create_loss};
use crate::tree::DecisionTree;
use serde::{Deserialize, Serialize};
use thiserror::Error;
use rand::seq::SliceRandom;

/// Errors that can occur during gradient boosting operations.
///
/// This error type covers all failure modes: configuration issues,
/// training errors, prediction failures, and data validation problems.
#[derive(Error, Debug)]
pub enum BoostingError {
    #[error("Invalid input data: {0}")]
    InvalidInput(String),

    #[error("Training error: {0}")]
    TrainingError(String),

    #[error("Prediction error: {0}")]
    PredictionError(String),

    #[error("Configuration error: {0}")]
    ConfigError(String),

    #[error("Tree building error: {0}")]
    TreeError(String),

    #[error("Loss function error: {0}")]
    LossError(String),

    #[error("Serialization error: {0}")]
    SerializationError(String),
}

/// Result type for all gradient boosting operations.
pub type BoostingResult<T> = std::result::Result<T, BoostingError>;



/// State of the boosting process at a specific training iteration.
/// 
/// Captures performance metrics and model statistics for monitoring
/// training progress and implementing early stopping.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct IterationState {
    /// Current iteration number (0-indexed).
    pub iteration: usize,
    /// Training loss at this iteration.
    pub train_loss: f64,
    /// Validation loss if validation data is provided.
    pub validation_loss: Option<f64>,
    /// Number of trees in the ensemble after this iteration.
    pub n_trees: usize,
    /// Number of leaves in the tree added this iteration.
    pub n_leaves: usize,
}


/// Complete training history and state.
/// 
/// Stores all iteration states and tracks the best performing iteration
/// for early stopping and model selection.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TrainingState {
    /// Per-iteration training state history.
    pub iterations: Vec<IterationState>,
    /// Iteration with the best validation loss (if early stopping is used).
    pub best_iteration: Option<usize>,
    /// Best validation loss achieved during training.
    pub best_validation_loss: Option<f64>,
}

/// Main gradient boosting implementation.
/// 
/// [`GradientBooster`] is the core ensemble model that trains a sequence
/// of decision trees to minimize a given loss function. It supports
/// various training strategies and provides comprehensive monitoring capabilities.
/// 
/// # Key Features
/// 
/// - **Loss Functions**: MSE, MAE, Huber (robust regression), LogLoss (classification)
/// - **Regularization**: Learning rate shrinkage, L2 regularization (`lambda`)
/// - **Stochastic Boosting**: Row subsampling for variance reduction
/// - **Early Stopping**: Configurable patience and improvement tolerance
/// - **Feature Importance**: Gain-based importance scores
/// 
/// # Serialization
/// 
/// This struct derives `Serialize` and `Deserialize`, allowing trained models
/// to be saved and loaded using serde-compatible formats.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GradientBooster {
    /// Model configuration and hyperparameters.
    config: GBRTConfig,
    /// Sequence of trained decision trees.
    trees: Vec<Tree>,
    /// Initial prediction (optimal constant model before boosting iterations).
    initial_prediction: f64,
    /// Training history and validation monitoring state.
    training_state: Option<TrainingState>,
    /// Feature importance scores (normalized gain).
    feature_importance: Vec<f64>,
    /// Whether the model has been trained.
    is_trained: bool,
}

impl GradientBooster {
    /// Creates a new gradient booster with the given configuration.
    /// 
    /// Validates the configuration before returning. The booster is untrained
    /// and ready for [`fit()`] to be called.
    /// 
    /// # Parameters
    /// - `config`: [`GBRTConfig`] with hyperparameters and training options.
    /// 
    /// # Errors
    /// Returns `BoostingError::ConfigError` if configuration validation fails.
    pub fn new(config: GBRTConfig) -> BoostingResult<Self> {
        config.validate()
            .map_err(|e| BoostingError::ConfigError(e))?;

        Ok(Self {
            config,
            trees: Vec::new(),
            initial_prediction: 0.0,
            training_state: None,
            feature_importance: Vec::new(),
            is_trained: false,
        })
    }

    /// Trains the model on a dataset with optional validation data.
    /// 
    /// This is the main training entry point. It performs the complete
    /// gradient boosting algorithm including:
    /// - Data validation and initialization
    /// - Iterative tree building with gradient optimization
    /// - Optional stochastic subsampling
    /// - Validation loss monitoring and early stopping
    /// - Feature importance computation
    /// 
    /// # Parameters
    /// - `train_data`: Training dataset with features and targets.
    /// - `validation_data`: Optional validation dataset for early stopping and monitoring.
    /// 
    /// # Errors
    /// Returns errors for:
    /// - Invalid training data (empty, mismatched dimensions)
    /// - Configuration issues detected during training
    /// - Tree building failures
    /// - Loss computation problems (e.g., NaN values)
    /// 
    /// # Early Stopping
    /// If `config.early_stopping_rounds` is set and `validation_data` is provided,
    /// training stops when validation loss doesn't improve significantly
    /// (by more than `config.early_stopping_tolerance`) for the specified number of rounds.
    pub fn fit(
        &mut self,
        train_data: &Dataset,
        validation_data: Option<&Dataset>,
    ) -> BoostingResult<()> {
        self.validate_training_data(train_data, validation_data)?;

        // Reset state
        self.trees.clear();
        self.initial_prediction = 0.0;
        self.is_trained = false;
        self.training_state = Some(TrainingState {
            iterations: Vec::new(),
            best_iteration: None,
            best_validation_loss: None,
        });

        // Initialize predictions
        self.initial_prediction = self.compute_initial_prediction(train_data.targets().as_slice().unwrap())?;
        let mut predictions = vec![self.initial_prediction; train_data.n_samples()];

        // Create loss function
        let loss_fn = create_loss(&self.config.loss);

        // Prepare validation data
        let (val_features, val_targets, mut val_predictions) = 
            self.prepare_validation_data(validation_data, train_data.n_samples())?;

        // Early stopping state
        let mut best_val_loss = f64::INFINITY;
        let mut no_improvement_count = 0;
        let mut best_iteration = 0;

        // Main boosting loop
        for iteration in 0..self.config.n_estimators {
            // Compute gradients and hessians
            let (gradients, hessians) = loss_fn.gradient_hessian(
                train_data.targets().as_slice().unwrap(), 
                &predictions
            );

            // Convert ndarray results to slices
            let gradients_slice = gradients.as_slice().unwrap();
            let hessians_slice = hessians.as_slice().unwrap();

            // Sample data for stochastic gradient boosting
            let (sampled_features, sampled_gradients, sampled_hessians, sample_indices) = 
                self.sample_data(train_data.features(), gradients_slice, hessians_slice)?;

            // Fit a new tree to the negative gradients
            let tree = self.fit_tree(&sampled_features, &sampled_gradients, &sampled_hessians)?;

            // Update predictions with learning rate
            self.update_predictions(train_data.features(), &tree, &mut predictions, sample_indices.as_ref())?;

            // Update training state and check early stopping
            let should_stop = self.update_training_state(
                iteration,
                train_data,
                validation_data,
                &predictions,
                &mut val_predictions,
                &tree,
                &mut best_val_loss,
                &mut no_improvement_count,
                &mut best_iteration,
                &loss_fn,
            )?;

            self.trees.push(tree);

            if should_stop {
                println!("Early stopping at iteration {}. Best iteration: {}", iteration, best_iteration);
                break;
            }
        }

        // Finalize training
        self.compute_feature_importance(train_data.n_features());
        self.is_trained = true;

        Ok(())
    }

    /// Makes predictions for a batch of samples.
    /// 
    /// Calculates predictions by summing contributions from all trees
    /// with shrinkage applied. Applies loss-specific transformation
    /// (e.g., sigmoid for LogLoss).
    /// 
    /// # Parameters
    /// - `features`: Feature matrix with shape `(n_samples, n_features)`.
    /// 
    /// # Returns
    /// Vector of predictions of length `n_samples`.
    /// 
    /// # Errors
    /// - `PredictionError::ModelNotTrained` if called before [`fit()`]
    /// - `PredictionError::FeatureMismatch` if `features.n_features()` doesn't match training data
    pub fn predict(&self, features: &FeatureMatrix) -> BoostingResult<Vec<f64>> {
        if !self.is_trained {
            return Err(BoostingError::PredictionError("Model not trained".to_string()));
        }

        if features.n_features() != self.feature_importance.len() {
            return Err(BoostingError::PredictionError(
                format!("Expected {} features, got {}", self.feature_importance.len(), features.n_features())
            ));
        }

        let mut predictions = vec![self.initial_prediction; features.n_samples()];

        for tree in &self.trees {
            for (i, pred) in predictions.iter_mut().enumerate() {
                let sample = features.get_sample(i)
                    .map_err(|e| BoostingError::PredictionError(e.to_string()))?;
                *pred += self.config.learning_rate * tree.predict(&sample.to_vec());
            }
        }

        // Apply loss-specific transformation
        let transformed_predictions = self.apply_prediction_transform(&predictions);

        Ok(transformed_predictions)
    }

    /// Makes a prediction for a single sample.
    /// 
    /// More efficient than `predict()` for single-sample inference.
    /// 
    /// # Parameters
    /// - `features`: Slice of feature values of length `n_features`.
    /// 
    /// # Returns
    /// Single prediction value.
    /// 
    /// # Errors
    /// - `PredictionError::ModelNotTrained` if called before [`fit()`]
    /// - `PredictionError::FeatureMismatch` if `features.len()` doesn't match training data
    pub fn predict_single(&self, features: &[f64]) -> BoostingResult<f64> {
        if !self.is_trained {
            return Err(BoostingError::PredictionError("Model not trained".to_string()));
        }

        if features.len() != self.feature_importance.len() {
            return Err(BoostingError::PredictionError(
                format!("Expected {} features, got {}", self.feature_importance.len(), features.len())
            ));
        }

        let mut prediction = self.initial_prediction;

        for tree in &self.trees {
            prediction += self.config.learning_rate * tree.predict(features);
        }

        let transformed_prediction = self.apply_single_prediction_transform(prediction);

        Ok(transformed_prediction)
    }

    /// Returns feature importance scores from the trained model.
    /// 
    /// Importance is computed as the normalized total gain contributed
    /// by splits on each feature across all trees. Scores sum to 1.0.
    /// 
    /// # Returns
    /// Slice of importance scores of length `n_features`.
    /// Returns zeros if `config.compute_feature_importance` is `false`.
    /// 
    /// # Panics
    /// Will panic if called before any model has been trained (no features are known).
    pub fn feature_importance(&self) -> &[f64] {
        &self.feature_importance
    }

    /// Returns the training history and validation state.
    /// 
    /// Contains per-iteration metrics and the best iteration if
    /// early stopping was used. This is useful for plotting learning curves
    /// or analyzing training dynamics.
    /// 
    /// # Returns
    /// `Some(TrainingState)` after training, `None` before training.
    pub fn training_state(&self) -> Option<&TrainingState> {
        self.training_state.as_ref()
    }

    /// Returns the number of trees in the ensemble.
    /// 
    /// Note: This may be fewer than `config.n_estimators` if early stopping triggered.
    /// 
    /// # Returns
    /// Number of trees after training.
    pub fn n_trees(&self) -> usize {
        self.trees.len()
    }

    /// Returns the configuration used to create this booster.
    pub fn config(&self) -> &GBRTConfig {
        &self.config
    }

    /// Checks whether the model has been trained.
    pub fn is_trained(&self) -> bool {
        self.is_trained
    }

    // Private helper methods

    /// Validates training and validation datasets before fitting.
    fn validate_training_data(
        &self,
        train_data: &Dataset,
        validation_data: Option<&Dataset>,
    ) -> BoostingResult<()> {
        if train_data.n_samples() == 0 {
            return Err(BoostingError::InvalidInput("Training dataset is empty".to_string()));
        }

        if let Some(val_data) = validation_data {
            if val_data.n_features() != train_data.n_features() {
                return Err(BoostingError::InvalidInput(
                    "Validation data has different number of features".to_string()
                ));
            }
        }

        Ok(())
    }

    /// Computes the optimal initial prediction before boosting iterations.
    /// 
    /// For MSE/Huber: mean of targets. For MAE: median. For LogLoss: log-odds.
    fn compute_initial_prediction(&self, targets: &[f64]) -> BoostingResult<f64> {
        match self.config.loss {
            LossFunction::MSE | LossFunction::Huber(_) => {
                Ok(targets.iter().sum::<f64>() / targets.len() as f64)
            }
            LossFunction::MAE => {
                let mut sorted = targets.to_vec();
                sorted.sort_by(|a, b| a.partial_cmp(b).unwrap());
                Ok(sorted[sorted.len() / 2]) // median
            }
            LossFunction::LogLoss => {
                // For classification, use log-odds of positive class proportion
                let mean = targets.iter().sum::<f64>() / targets.len() as f64;
                let mean_clamped = mean.max(1e-15).min(1.0 - 1e-15);
                Ok((mean_clamped / (1.0 - mean_clamped)).ln())
            }
        }
    }

    /// Prepares validation data structures for loss monitoring.
    fn prepare_validation_data(
        &self,
        validation_data: Option<&Dataset>,
        n_train_samples: usize,
    ) -> BoostingResult<(Option<FeatureMatrix>, Option<Vec<f64>>, Vec<f64>)> {
        match validation_data {
            Some(val_data) => {
                let val_predictions = vec![self.initial_prediction; val_data.n_samples()];
                Ok((
                    Some(val_data.features().clone()),
                    Some(val_data.targets().as_slice().unwrap().to_vec()),
                    val_predictions,
                ))
            }
            None => {
                Ok((None, None, Vec::new()))
            }
        }
    }

    /// Samples data for stochastic gradient boosting.
    /// 
    /// Returns sampled features, gradients, hessians, and original indices.
    fn sample_data(
        &self,
        features: &FeatureMatrix,
        gradients: &[f64],
        hessians: &[f64],
    ) -> BoostingResult<(FeatureMatrix, Vec<f64>, Vec<f64>, Option<Vec<usize>>)> {
        if self.config.subsample >= 1.0 {
            // No sampling
            return Ok((
                features.clone(),
                gradients.to_vec(),
                hessians.to_vec(),
                None,
            ));
        }

        // Stochastic gradient boosting: sample a subset of data
        let n_samples = (features.n_samples() as f64 * self.config.subsample) as usize;
        let n_samples = n_samples.max(1).min(features.n_samples());

        let mut rng = rand::thread_rng();
        let mut indices: Vec<usize> = (0..features.n_samples()).collect();
        indices.shuffle(&mut rng);
        indices.truncate(n_samples);

        let sampled_features = features.select_samples(&indices)
            .map_err(|e| BoostingError::TrainingError(e.to_string()))?;
        let sampled_gradients: Vec<f64> = indices.iter().map(|&i| gradients[i]).collect();
        let sampled_hessians: Vec<f64> = indices.iter().map(|&i| hessians[i]).collect();

        Ok((sampled_features, sampled_gradients, sampled_hessians, Some(indices)))
    }

    /// Trains a single decision tree on the gradient data.
    fn fit_tree(
        &self,
        features: &FeatureMatrix,
        gradients: &[f64],
        hessians: &[f64],
    ) -> BoostingResult<Tree> {
        let tree_config = &self.config.tree_config;

        let tree_builder = DecisionTree::new(
            tree_config.max_depth,
            tree_config.min_samples_split,
            tree_config.min_samples_leaf,
            tree_config.min_impurity_decrease,
            tree_config.max_features,
            tree_config.lambda,
        );

        tree_builder.fit(features, gradients, hessians)
            .map_err(|e| BoostingError::TreeError(format!("Tree fitting failed: {}", e)))
    }

    /// Updates predictions after fitting a new tree.
    fn update_predictions(
        &self,
        features: &FeatureMatrix,
        tree: &Tree,
        predictions: &mut Vec<f64>,
        sample_indices: Option<&Vec<usize>>,
    ) -> BoostingResult<()> {
        match sample_indices {
            Some(indices) => {
                // Update only sampled indices for stochastic gradient boosting
                for &idx in indices {
                    let sample = features.get_sample(idx)
                        .map_err(|e| BoostingError::TrainingError(e.to_string()))?;
                    predictions[idx] += self.config.learning_rate * tree.predict(&sample.to_vec());
                }
            }
            None => {
                // Update all predictions
                for (i, pred) in predictions.iter_mut().enumerate() {
                    let sample = features.get_sample(i)
                        .map_err(|e| BoostingError::TrainingError(e.to_string()))?;
                    *pred += self.config.learning_rate * tree.predict(&sample.to_vec());
                }
            }
        }

        Ok(())
    }

   
    /// Updates training state and implements early stopping logic.
    /// 
    /// Returns `true` if training should stop, `false` otherwise.
    fn update_training_state(
        &mut self,
        iteration: usize,
        train_data: &Dataset,
        validation_data: Option<&Dataset>,
        predictions: &[f64],
        val_predictions: &mut Vec<f64>,
        tree: &Tree,
        best_val_loss: &mut f64,
        no_improvement_count: &mut usize,
        best_iteration: &mut usize,
        loss_fn: &Box<dyn GradientLoss>,
    ) -> BoostingResult<bool> {
        let mut iteration_state = IterationState {
            iteration,
            train_loss: 0.0,
            validation_loss: None,
            n_trees: self.trees.len() + 1,
            n_leaves: tree.n_leaves(),
        };

        // Compute training loss
        iteration_state.train_loss = loss_fn.loss(
            train_data.targets().as_slice().unwrap(), 
            predictions
        );

        let should_stop = if let (Some(val_features), Some(val_targets)) = (
            validation_data.map(|d| d.features()), 
            validation_data.map(|d| d.targets())
        ) {
            // Update validation predictions
            for (i, pred) in val_predictions.iter_mut().enumerate() {
                let sample = val_features.get_sample(i)
                    .map_err(|e| BoostingError::TrainingError(e.to_string()))?;
                *pred += self.config.learning_rate * tree.predict(&sample.to_vec());
            }

            let current_val_loss = loss_fn.loss(
                val_targets.as_slice().unwrap(), 
                val_predictions
            );

            // PROTECT AGAINST NAN/INF
            if !current_val_loss.is_finite() {
                eprintln!("Warning: Validation loss became NaN/Inf at iteration {}", iteration);
                return Ok(true); // Stop immediately
            }

            iteration_state.validation_loss = Some(current_val_loss);

            // EARLY STOPPING WITH TOLERANCE
            let improvement_threshold = *best_val_loss * (1.0 + self.config.early_stopping_tolerance);
            
            if current_val_loss < improvement_threshold {
                // Significant improvement
                *best_val_loss = current_val_loss;
                *no_improvement_count = 0;
                *best_iteration = iteration;
            } else {
                // No significant improvement
                *no_improvement_count += 1;
            }

            // Check patience
            if let Some(patience) = self.config.early_stopping_rounds {
                if *no_improvement_count >= patience {
                    println!("Early stopping triggered at iteration {}. Best validation loss: {:.6} at iteration {}", 
                             iteration, *best_val_loss, *best_iteration);
                    if let Some(state) = &mut self.training_state {
                        state.best_iteration = Some(*best_iteration);
                        state.best_validation_loss = Some(*best_val_loss);
                    }
                    return Ok(true); // Signal to stop
                }
            }

            false
        } else {
            false
        };

        // Update training state
        if let Some(state) = &mut self.training_state {
            state.iterations.push(iteration_state);
        }

        Ok(should_stop)
    }

    /// Computes normalized feature importance from all trees.
    fn compute_feature_importance(&mut self, n_features: usize) {
        if !self.config.compute_feature_importance {
            self.feature_importance = vec![0.0; n_features];
            return;
        }

        let mut importance = vec![0.0; n_features];
        let total_gain: f64 = self.trees.iter()
            .map(|tree| {
                tree.feature_importance()
                    .iter()
                    .map(|&(feature, gain)| {
                        importance[feature] += gain;
                        gain
                    })
                    .sum::<f64>()
            })
            .sum();

        // Normalize
        if total_gain > 0.0 {
            for imp in &mut importance {
                *imp /= total_gain;
            }
        }

        self.feature_importance = importance;
    }

    /// Applies loss-specific transformations to predictions.
    fn apply_prediction_transform(&self, predictions: &[f64]) -> Vec<f64> {
        if matches!(self.config.loss, LossFunction::LogLoss) {
            let loss_fn = create_loss(&self.config.loss);
            loss_fn.transform(predictions)
        } else {
            predictions.to_vec()
        }
    }

    /// Applies loss-specific transformation to a single prediction.
    fn apply_single_prediction_transform(&self, prediction: f64) -> f64 {
        if matches!(self.config.loss, LossFunction::LogLoss) {
            let loss_fn = create_loss(&self.config.loss);
            loss_fn.transform(&[prediction])[0]
        } else {
            prediction
        }
    }
}

impl std::fmt::Display for GradientBooster {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        writeln!(f, "GradientBooster")?;
        writeln!(f, "  Trained: {}", self.is_trained)?;
        writeln!(f, "  Trees: {}", self.trees.len())?;
        writeln!(f, "  Loss: {}", self.config.loss)?;
        writeln!(f, "  Learning Rate: {:.4}", self.config.learning_rate)?;
        writeln!(f, "  Subsampling: {:.2}", self.config.subsample)?;

        if let Some(state) = &self.training_state {
            if let Some(iter_state) = state.iterations.last() {
                writeln!(f, "  Final Training Loss: {:.6}", iter_state.train_loss)?;
                if let Some(val_loss) = iter_state.validation_loss {
                    writeln!(f, "  Final Validation Loss: {:.6}", val_loss)?;
                }
            }
        }

        Ok(())
    }
}