gbrt_rs/

lib.rs

#![allow(unused)]

//! Gradient Boosted Regression Trees (GBRT) in Rust
//!
//! A high-performance, feature-rich implementation of gradient boosting
//! for machine learning tasks including regression and binary classification.
//!
//! This library provides a complete gradient boosting framework with:
//! - **Multiple objective functions** (MSE, MAE, Huber, Log Loss)
//! - **Categorical feature encoding** with auto-detection
//! - **Early stopping** with validation support
//! - **Cross-validation** with proper categorical handling
//! - **Feature importance** analysis
//! - **Model serialization** with metadata
//! - **Comprehensive error handling** with unified error types
//!
//! # Architecture
//!
//! The library is organized into modular components:
//!
//! - [`core`]: Core configurations and loss functions
//! - [`boosting`]: Gradient boosting algorithms and training loop
//! - [`tree`]: Decision tree construction with pluggable split criteria
//! - [`data`]: Dataset management and feature preprocessing
//! - [`objective`]: Objective functions for optimization
//! - [`utils`]: Mathematical utilities, validation, and serialization
//! - [`io`]: Data loading and model persistence
//!
//! # Examples
//!
//! Basic regression example:
//!
//! ```
//! use gbrt_rs::{
//!     GradientBooster, BoosterFactory, Dataset, FeatureMatrix,
//!     DataLoader, ObjectiveType
//! };
//! use ndarray::Array2;
//!
//! // Load data from CSV
//! let data_loader = DataLoader::new()?;
//! let dataset = data_loader.load_csv("data/training.csv")?;
//!
//! // Train regression model
//! let mut booster = BoosterFactory::create_regression_booster()?;
//! booster.fit(&dataset, None)?;
//!
//! // Make predictions
//! let predictions = booster.predict(dataset.features())?;
//! # Ok::<(), Box<dyn std::error::Error>>(())
//! ```
//!
//! Training with validation and early stopping:
//!
//! ```
//! use gbrt_rs::{
//!     GBRTModel, Dataset, DataLoader, GBRTConfig,
//!     LossFunction, ObjectiveType
//! };
//!
//! // Configure model
//! let config = GBRTConfig {
//!     n_estimators: 100,
//!     learning_rate: 0.1,
//!     max_depth: 6,
//!     early_stopping_rounds: Some(10),
//!     ..Default::default()
//! };
//!
//! let mut model = GBRTModel::with_config(config)?;
//!
//! // Split data for validation
//! let data_loader = DataLoader::new()?;
//! let dataset = data_loader.load_csv("data.csv")?;
//! let (train_data, val_data) = data_loader.split_data(&dataset, 0.2, true, None)?;
//!
//! // Train with early stopping
//! model.fit_with_validation(&train_data, &val_data)?;
//!
//! // Evaluate
//! let predictions = model.predict(val_data.features())?;
//! # Ok::<(), Box<dyn std::error::Error>>(())
//! ```

// Core modules
pub mod core;
/// Gradient boosting algorithms and ensemble training
pub mod boosting;
/// Decision tree construction and split criteria
pub mod tree;
/// Dataset handling, feature matrices, and preprocessing
pub mod data;
/// Loss functions and objective functions for training
pub mod objective;
/// Mathematical utilities and validation helpers
pub mod utils;
/// Data loading and model serialization
pub mod io;

/// Data preprocessing utilities
pub mod preprocessing {
    pub use crate::data::preprocessing::*;
}

use std::collections::HashMap;

// Re-export main types for easy access
pub use boosting::{
    GradientBooster, BoosterFactory, create_booster,
    BoostingError, BoostingResult, BoostingType,
    TrainingState, IterationState
};
pub use data::{
    Dataset, FeatureMatrix, DataError, DataResult,
    Scaler, StandardScaler, MinMaxScaler
};
pub use objective::{
    Objective, RegressionObjective, BinaryClassificationObjective,
    ObjectiveFactory, create_objective, ObjectiveType,
    ObjectiveError, ObjectiveResult, ObjectiveConfig,
    MSEObjective, MAEObjective, HuberObjective, LogLossObjective
};
pub use tree::{
    DecisionTree, TreeBuilder, TreeError, TreeResult,
    Splitter, BestSplitter, SplitCriterion, MSECriterion
};
pub use utils::{
    MathError, MathResult, Statistics, VectorMath,
    ValidationError, ValidationResult, DataValidator,
    SerializationError, SerializationResult, ModelSerializer
};
pub use io::{
    DataLoader, DataLoaderError, DataLoaderResult, DataFormat,
    ModelIO, ModelIOError, ModelIOResult, ModelFormat,
    LoadOptions as DataLoadOptions, SaveOptions as ModelSaveOptions
};

/// Unified error type for the entire crate.
///
/// This enum aggregates all possible errors that can occur during
/// model training, prediction, or data processing.
///
/// # Variants
///
/// * `BoostingError` - Errors during gradient boosting
/// * `DataError` - Dataset loading or preprocessing errors
/// * `TreeError` - Decision tree construction errors
/// * `ObjectiveError` - Loss function computation errors
/// * `MathError` - Mathematical operation errors
/// * `ValidationError` - Data validation failures
/// * `SerializationError` - Model serialization/deserialization errors
/// * `DataLoaderError` - Data loading errors
/// * `ModelIOError` - Model I/O errors
/// * `IoError` - General I/O errors
/// * `CsvError` - CSV parsing errors
/// * `JsonError` - JSON serialization errors
/// * `ConfigError` - Invalid configuration
/// * `TrainingError` - Training process errors
/// * `PredictionError` - Prediction errors
/// * `NotTrained` - Model not yet trained
#[derive(thiserror::Error, Debug)]
pub enum GBRTError {
    #[error("Boosting error: {0}")]
    BoostingError(#[from] BoostingError),
    
    #[error("Data error: {0}")]
    DataError(#[from] DataError),
    
    #[error("Tree error: {0}")]
    TreeError(#[from] TreeError),
    
    #[error("Objective error: {0}")]
    ObjectiveError(#[from] ObjectiveError),
    
    #[error("Math error: {0}")]
    MathError(#[from] MathError),
    
    #[error("Validation error: {0}")]
    ValidationError(#[from] ValidationError),
    
    #[error("Serialization error: {0}")]
    SerializationError(#[from] SerializationError),
    
    #[error("Data loader error: {0}")]
    DataLoaderError(#[from] DataLoaderError),
    
    #[error("Model IO error: {0}")]
    ModelIOError(#[from] ModelIOError),
    
    #[error("IO error: {0}")]
    IoError(#[from] std::io::Error),
        
    #[error("CSV error: {0}")]
    CsvError(#[from] csv::Error),

    #[error("JSON error: {0}")]
    JsonError(#[from] serde_json::Error),

    #[error("Configuration error: {0}")]
    ConfigError(String),
    
    #[error("Training error: {0}")]
    TrainingError(String),
    
    #[error("Prediction error: {0}")]
    PredictionError(String),
    
    #[error("Model not trained")]
    NotTrained,
}

/// Result type alias for the crate using `GBRTError`
pub type Result<T> = std::result::Result<T, GBRTError>;

/// Main GBRT model struct that provides a high-level API for training and prediction.
///
/// This struct encapsulates the gradient boosting algorithm and provides
/// a simple interface for regression and classification tasks.
///
/// # Examples
///
/// ```
/// # use gbrt_rs::{GBRTModel, Dataset, FeatureMatrix};
/// # use ndarray::Array2;
/// # fn example() -> Result<(), Box<dyn std::error::Error>> {
/// let features = FeatureMatrix::new(Array2::ones((10, 3)))?;
/// let targets = vec![1.0; 10];
/// let dataset = Dataset::new(features, targets)?;
///
/// let mut model = GBRTModel::new()?;
/// model.fit(&dataset)?;
/// # Ok(())
/// # }
/// ```
pub struct GBRTModel {
    // The underlying gradient booster implementation
    booster: GradientBooster,
    // Optional feature names for better interpretability
    feature_names: Option<Vec<String>>,
    // Optional target name for better interpretability
    target_name: Option<String>,
    // Minimum samples required to split a node (default: 2)
    pub min_samples_split: usize,
    // Minimum samples required in a leaf node (default: 1)
    pub min_samples_leaf: usize,
    // L2 regularization parameter (default: 0.0)
    pub lambda_l2: f64,          
}

impl GBRTModel {
    /// Creates a new GBRT model with default configuration for regression.
    ///
    /// # Returns
    ///
    /// A `Result` containing the configured model or an error if creation fails.
    ///
    /// # Errors
    ///
    /// Returns an error if the underlying booster cannot be created.
    pub fn new() -> Result<Self> {
        let booster = BoosterFactory::create_regression_booster()?;
        Ok(Self {
            booster,
            feature_names: None,
            target_name: None,
            min_samples_split: 2,  // ✅ Default: split nodes with ≥2 samples
            min_samples_leaf: 1,   // ✅ Default: allow leaves with 1 sample
            lambda_l2: 0.0,        // ✅ Default: no L2 regularization
        })
    }
    
    /// Creates a new GBRT model for binary classification.
    ///
    /// # Returns
    ///
    /// A `Result` containing the configured model or an error if creation fails.
    ///
    /// # Errors
    ///
    /// Returns an error if the underlying booster cannot be created.
    pub fn new_classifier() -> Result<Self> {
        let booster = BoosterFactory::create_classification_booster()?;
        Ok(Self {
            booster,
            feature_names: None,
            target_name: None,
            min_samples_split: 2,
            min_samples_leaf: 1,
            lambda_l2: 0.0,
        })
    }
    
    /// Creates a new GBRT model with custom configuration.
    ///
    /// # Parameters
    ///
    /// * `config` - Custom gradient boosting configuration
    ///
    /// # Returns
    ///
    /// A `Result` containing the configured model or an error if creation fails.
    ///
    /// # Errors
    ///
    /// Returns an error if the booster cannot be created with the given config.
    pub fn with_config(config: boosting::GBRTConfig) -> Result<Self> {
        let booster = GradientBooster::new(config.clone())?;
        Ok(Self {
            booster,
            feature_names: None,
            target_name: None,
            min_samples_split: config.tree_config.min_samples_split,
            min_samples_leaf: config.tree_config.min_samples_leaf,
            lambda_l2: config.tree_config.lambda,
        })
    }
    
    /// Trains the model on a dataset.
    ///
    /// # Parameters
    ///
    /// * `dataset` - The training dataset containing features and targets
    ///
    /// # Returns
    ///
    /// A `Result` indicating success or failure.
    ///
    /// # Errors
    ///
    /// Returns an error if training fails or the dataset is invalid.
    pub fn fit(&mut self, dataset: &Dataset) -> Result<()> {
        self.booster.fit(dataset, None)?;
        Ok(())
    }

    /// Trains the model with validation data for early stopping.
    ///
    /// # Parameters
    ///
    /// * `train_data` - The training dataset
    /// * `val_data` - The validation dataset for monitoring overfitting
    ///
    /// # Returns
    ///
    /// A `Result` indicating success or failure.
    ///
    /// # Errors
    ///
    /// Returns an error if training fails or either dataset is invalid.
    pub fn fit_with_validation(&mut self, train_data: &Dataset, val_data: &Dataset) -> Result<()> {
        self.booster.fit(train_data, Some(val_data))?;
        Ok(())
    }
    
    /// Makes predictions for a feature matrix.
    ///
    /// # Parameters
    ///
    /// * `features` - The feature matrix to predict on
    ///
    /// # Returns
    ///
    /// A `Result` containing a vector of predictions.
    ///
    /// # Errors
    ///
    /// Returns an error if the model is not trained or prediction fails.
    pub fn predict(&self, features: &FeatureMatrix) -> Result<Vec<f64>> {
        self.booster.predict(features).map_err(GBRTError::from)
    }
    
    /// Makes prediction for a single sample.
    ///
    /// # Parameters
    ///
    /// * `features` - A slice of feature values for one sample
    ///
    /// # Returns
    ///
    /// A `Result` containing the predicted value.
    ///
    /// # Errors
    ///
    /// Returns an error if the model is not trained or the feature length is invalid.
    pub fn predict_single(&self, features: &[f64]) -> Result<f64> {
        self.booster.predict_single(features).map_err(GBRTError::from)
    }
    
    /// Gets feature importance scores.
    ///
    /// # Returns
    ///
    /// A vector of importance scores, one per feature.
    pub fn feature_importance(&self) -> Vec<f64> {
        self.booster.feature_importance().to_vec()
    }

    /// Gets training history and state.
    ///
    /// # Returns
    ///
    /// An optional reference to the training state.
    pub fn training_history(&self) -> Option<&TrainingState> {
        self.booster.training_state()
    }
    
    /// Checks if the model has been trained.
    ///
    /// # Returns
    ///
    /// `true` if the model has been trained, `false` otherwise. 
    pub fn is_trained(&self) -> bool {
        self.booster.is_trained()
    }
    
    /// Gets the number of trees in the ensemble.
    ///
    /// # Returns
    ///
    /// The number of trees, or 0 if untrained.
    pub fn n_trees(&self) -> usize {
        self.booster.n_trees()
    }
    
    /// Sets feature names for better interpretability.
    ///
    /// # Parameters
    ///
    /// * `names` - Vector of feature names in order
    pub fn set_feature_names(&mut self, names: Vec<String>) {
        self.feature_names = Some(names);
    }

    /// Gets the feature names if they have been set.
    ///
    /// # Returns
    ///
    /// An optional reference to the feature names vector.
    pub fn feature_names(&self) -> Option<&Vec<String>> {
        self.feature_names.as_ref()
    }

    /// Sets the target name for better interpretability.
    ///
    /// # Parameters
    ///
    /// * `name` - The target variable name
    pub fn set_target_name(&mut self, name: String) {
        self.target_name = Some(name);
    }

    /// Gets a reference to the underlying booster.
    ///
    /// # Returns
    ///
    /// Reference to the gradient booster.
    pub fn booster(&self) -> &GradientBooster {
        &self.booster
    }
}

impl Default for GBRTModel {
    /// Creates a default GBRT model for regression.
    fn default() -> Self {
        Self::new().unwrap()
    }
}

/// Convenience functions for common tasks

/// Trains a GBRT model on data from a CSV file.
///
/// # Parameters
///
/// * `train_path` - Path to the CSV training data file
/// * `target_column` - Optional name of the target column (defaults to "target")
/// * `model_type` - Type of model to train (regression or binary classification)
///
/// # Returns
///
/// A `Result` containing the trained model.
///
/// # Errors
///
/// Returns an error if data loading or training fails.
pub fn train_from_csv(
    train_path: &str,
    target_column: Option<&str>,
    model_type: ObjectiveType,
) -> Result<GBRTModel> {
    let data_loader = DataLoader::new()?;
    let options = DataLoadOptions::default()
        .with_target_column(target_column.unwrap_or("target"));
    
    let dataset = data_loader.load_csv(train_path)?;
    
    let mut model = match model_type {
        ObjectiveType::Regression => GBRTModel::new()?,
        ObjectiveType::BinaryClassification => GBRTModel::new_classifier()?,
    };
    
    model.fit(&dataset)?;
    Ok(model)
}

/// Trains a GBRT model with train/validation split.
///
/// # Parameters
///
/// * `data_path` - Path to the CSV data file
/// * `target_column` - Optional name of the target column (defaults to "target")
/// * `test_size` - Fraction of data to use for validation (e.g., 0.2 for 20%)
/// * `model_type` - Type of model to train
///
/// # Returns
///
/// A `Result` containing the trained model with validation.
///
/// # Errors
///
/// Returns an error if data loading, splitting, or training fails.
pub fn train_with_validation_split(
    data_path: &str,
    target_column: Option<&str>,
    test_size: f64,
    model_type: ObjectiveType,
) -> Result<GBRTModel> {
    let data_loader = DataLoader::new()?;
    let options = DataLoadOptions::default()
        .with_target_column(target_column.unwrap_or("target"));
    
    let dataset = data_loader.load_csv(data_path)?;
    let (train_data, val_data) = data_loader.split_data(&dataset, test_size, true, None)?;
    
    let mut model = match model_type {
        ObjectiveType::Regression => GBRTModel::new()?,
        ObjectiveType::BinaryClassification => GBRTModel::new_classifier()?,
    };
    
    model.fit_with_validation(&train_data, &val_data)?;
    Ok(model)
}

/// Performs k-fold cross-validation for model evaluation.
///
/// # Parameters
///
/// * `data_path` - Path to the CSV data file
/// * `target_column` - Optional name of the target column (defaults to "target")
/// * `n_splits` - Number of cross-validation folds
/// * `model_type` - Type of model to evaluate
/// * `categorical_columns` - Optional comma-separated list of categorical column names
/// * `categorical_threshold` - Threshold for auto-detecting categorical features
///
/// # Returns
///
/// A `Result` containing a vector of R² scores for each fold.
///
/// # Errors
///
/// Returns an error if data loading, model training, or evaluation fails.
pub fn cross_validate(
    data_path: &str,
    target_column: Option<&str>,
    n_splits: usize,
    model_type: ObjectiveType,
    categorical_columns: Option<&str>, // ✅ NEW: Explicit categorical columns
    categorical_threshold: f64,        // ✅ NEW: Auto-detect threshold
) -> Result<Vec<f64>> {
    let data_loader = DataLoader::new()?;

    // ✅ Parse CSV headers to get column name to index mapping (same logic as train_model)
    let mut temp_reader = csv::Reader::from_path(data_path)?;
    let headers: Vec<String> = temp_reader
        .headers()?
        .iter()
        .map(|s| s.to_string())
        .collect();

    let header_map: HashMap<_, _> = headers.iter()
        .enumerate()
        .map(|(i, name)| (name.as_str(), i))
        .collect();

    // ✅ Parse categorical column names to indices (same logic as train_model)
    let categorical_indices = if let Some(col_names) = categorical_columns {
        col_names.split(',')
            .map(|s| s.trim())
            .map(|name| {
                header_map.get(name)
                    .copied()
                    .ok_or_else(|| {
                        GBRTError::ConfigError(format!(
                            "Unknown categorical column '{}'. Available columns: {:?}",
                            name, headers
                        ))
                    })
            })
            .collect::<Result<Vec<_>>>()?
    } else {
        vec![] // Auto-detect will be handled by load_csv_with_categorical
    };

    // ✅ Load data WITH proper categorical encoding (critical fix)
    let dataset = data_loader.load_csv_with_categorical(
        data_path,
        target_column.unwrap_or("target"),
        if categorical_columns.is_some() {
            Some(&categorical_indices)
        } else {
            None
        },
        categorical_threshold,
    )?;

    // Create cross-validation splits
    let splits = data_loader.cross_validation_splits(&dataset, n_splits, true, None)?;

    let mut scores = Vec::new();

    for (fold_idx, (train_data, test_data)) in splits.into_iter().enumerate() {
        let mut model = match model_type {
            ObjectiveType::Regression => GBRTModel::new()?,
            ObjectiveType::BinaryClassification => GBRTModel::new_classifier()?,
        };

        model.fit(&train_data)?;
        let predictions = model.predict(test_data.features())?;

        // ✅ Calculate R² (coefficient of determination) instead of RMSE
        let targets = test_data.targets().as_slice().unwrap();
        let score = calculate_r2(&predictions, targets)?;

        scores.push(score);
    }

    Ok(scores)
}

/// Computes the R² (coefficient of determination) metric.
///
/// # Parameters
///
/// * `predictions` - Slice of predicted values
/// * `targets` - Slice of true target values
///
/// # Returns
///
/// A `Result` containing the R² score.
///
/// # Errors
///
/// Returns an error if the input slices have different lengths.
///
/// # Formula
///
/// R² = 1 - (Σ(y_true - y_pred)² / Σ(y_true - y_mean)²)
fn calculate_r2(predictions: &[f64], targets: &[f64]) -> Result<f64> {
    if predictions.len() != targets.len() {
        return Err(GBRTError::ValidationError(
            ValidationError::ValidationFailed("Predictions and targets length mismatch".to_string())
        ));
    }

    let y_mean = targets.iter().sum::<f64>() / targets.len() as f64;
    let total_ss: f64 = targets.iter()
        .map(|&y| (y - y_mean).powi(2))
        .sum();
    let residual_ss: f64 = targets.iter()
        .zip(predictions.iter())
        .map(|(&true_val, &pred_val)| (true_val - pred_val).powi(2))
        .sum();

    let r2 = if total_ss == 0.0 {
        1.0 // Perfect prediction for constant target
    } else {
        1.0 - (residual_ss / total_ss)
    };

    Ok(r2)
}

/// Container for common model evaluation metrics.
pub struct ModelMetrics {
    /// Root Mean Squared Error
    pub rmse: f64,
    /// Mean Absolute Error
    pub mae: f64,
    /// R² (coefficient of determination)
    pub r2: f64,
}

impl ModelMetrics {
    /// Calculates regression metrics from predictions and true values.
    ///
    /// # Parameters
    ///
    /// * `y_true` - Slice of true target values
    /// * `y_pred` - Slice of predicted values
    ///
    /// # Returns
    ///
    /// A `Result` containing the computed metrics.
    ///
    /// # Errors
    ///
    /// Returns an error if the input slices have different lengths.
    pub fn regression_metrics(y_true: &[f64], y_pred: &[f64]) -> Result<Self> {
        if y_true.len() != y_pred.len() {
            return Err(GBRTError::ValidationError(
                ValidationError::ValidationFailed("True and predicted values must have same length".to_string())
            ));
        }
        
        let rmse = utils::rmse(y_pred, y_true)?;
        let mae = utils::compute_mean(
            &y_true.iter()
                .zip(y_pred.iter())
                .map(|(&true_val, &pred_val)| (true_val - pred_val).abs())
                .collect::<Vec<f64>>()
        )?;
        
        // Calculate R²
        let y_mean = utils::compute_mean(y_true)?;
        let total_ss: f64 = y_true.iter()
            .map(|&y| (y - y_mean).powi(2))
            .sum();
        let residual_ss: f64 = y_true.iter()
            .zip(y_pred.iter())
            .map(|(&true_val, &pred_val)| (true_val - pred_val).powi(2))
            .sum();
        
        let r2 = if total_ss == 0.0 {
            1.0
        } else {
            1.0 - (residual_ss / total_ss)
        };
        
        Ok(Self { rmse, mae, r2 })
    }
}

/// Returns the crate version string.
///
/// # Returns
///
/// The version from Cargo.toml
pub fn version() -> &'static str {
    env!("CARGO_PKG_VERSION")
}

/// Feature importance analysis with named features.
pub struct FeatureAnalysis {
    /// Sorted vector of (feature_name, importance_score) pairs
    pub importance: Vec<(String, f64)>,
    /// Summary statistics of feature importances
    pub summary: FeatureSummary,
}

/// Summary statistics for feature importance analysis.
pub struct FeatureSummary {
    /// Total number of features
    pub n_features: usize,
    /// Top 10 most important features
    pub top_features: Vec<(String, f64)>,
    /// Sum of all importance scores
    pub total_importance: f64,
}

impl FeatureAnalysis {
    /// Creates a feature analysis from a trained model.
    ///
    /// # Parameters
    ///
    /// * `model` - Reference to a trained gradient booster
    /// * `feature_names` - Optional vector of feature names
    ///
    /// # Returns
    ///
    /// A new `FeatureAnalysis` instance.
     pub fn from_model(model: &GradientBooster, feature_names: Option<Vec<String>>) -> Self {
        let importance_scores = model.feature_importance();
        let names = feature_names.unwrap_or_else(|| {
            (0..importance_scores.len())
                .map(|i| format!("feature_{}", i))
                .collect()
        });

        let mut importance: Vec<(String, f64)> = names
            .into_iter()
            .zip(importance_scores.iter().copied()) // Use .copied() to get f64 values from &[f64]
            .collect();

        // Sort by importance (descending)
        importance.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());

        let total_importance: f64 = importance.iter().map(|(_, imp)| imp).sum();
        let top_features = importance.iter().take(10).cloned().collect();

        let summary = FeatureSummary {
            n_features: importance.len(),
            top_features,
            total_importance,
        };

        Self { importance, summary }
    }
}