gbrt_rs/data/

preprocessing.rs

//! Data Preprocessing for Machine Learning
//! 
//! This module provides preprocessing utilities for preparing raw data
//! for gradient boosting models:
//! 
//! - **Feature Scaling**: Standardization and min-max normalization
//! - **Categorical Encoding**: Convert string categories to numerical values
//! - **Automatic Detection**: Identify categorical columns heuristically
//! - **Preprocessor Pipeline**: Combine multiple preprocessing steps
//! 
//! # Scaling Strategies
//! 
//! - [`StandardScaler`]: Zero mean, unit variance (z-score normalization)
//! - [`MinMaxScaler`]: Scale to specified range (e.g., [0, 1])
//! - Custom scalers via [`Scaler`] trait
//! 
//! # Categorical Encoding
//! 
//! [`CategoricalEncoder`] maps string categories to numeric codes. Unknown
//! categories are mapped to a default value (default: -1.0).
//! 
use super::{FeatureMatrix, DataError, DataResult};
use ndarray::{Array1, Array2};
use serde::{Serialize, Deserialize};
use crate::utils::Statistics;
use std::collections::{HashMap, HashSet};
use thiserror::Error;

/// Preprocessing-specific error types.
#[derive(Error, Debug)]
pub enum PreprocessingError {
    /// Errors during categorical encoding (e.g., parsing failures).
    #[error("Categorical encoding error: {0}")]
    EncodingError(String),
    
    /// Invalid data format or consistency issues.
    #[error("Invalid data: {0}")]
    InvalidData(String),
    
    /// Scaler used before fitting.    
    #[error("Scaler not fitted")]
    NotFitted,
}

/// Trait for feature scaling operations.
/// 
/// Scalers learn parameters from training data (via [`fit()`]) and apply
/// consistent transformations to new data (via [`transform()`]).
/// 
/// # Contract
/// 
/// - `fit()` must be called before `transform()`
/// - Transformations must be deterministic and reproducible
/// - Scaling parameters must be serializable for persistence
pub trait Scaler: Send + Sync {
    /// Learns scaling parameters from the training data.
    /// 
    /// # Parameters
    /// - `data`: Training feature matrix
    /// 
    /// # Returns
    /// `Ok(())` on successful fitting
    /// 
    /// # Errors
    /// - `DataError::PreprocessingError` if fitting fails
    fn fit(&mut self, data: &FeatureMatrix) -> DataResult<()>;
    
    /// Applies learned transformation to data.
    /// 
    /// # Parameters
    /// - `data`: Feature matrix to transform
    /// 
    /// # Returns
    /// Transformed feature matrix
    /// 
    /// # Errors
    /// - `DataError::PreprocessingError` if scaler not fitted
    /// - `DataError::PreprocessingError` if feature count mismatch
    fn transform(&self, data: &FeatureMatrix) -> DataResult<FeatureMatrix>;
    
    /// Fits and transforms in a single operation.
    /// 
    /// Equivalent to calling `fit()` then `transform()`.
    /// 
    /// # Returns
    /// Fitted and transformed feature matrix
    fn fit_transform(&mut self, data: &FeatureMatrix) -> DataResult<FeatureMatrix> {
        self.fit(data)?;
        self.transform(data)
    }
}

// ============================================================================
// StandardScaler - standardize features by removing mean and scaling to unit variance
// ============================================================================

/// Standardizes features by removing the mean and scaling to unit variance.
/// 
/// # Formula
/// 
/// For each feature column `x`:
/// 
/// ```
/// z = (x - μ) / σ
/// ```
/// 
/// where:
/// - `μ` is the mean of the feature
/// - `σ` is the standard deviation (σ² + ε)¹ᐟ²
/// - `ε` is a small constant for numerical stability
/// 
/// # When to Use
/// 
/// - When features have different scales
/// - For algorithms sensitive to feature magnitude (not tree-based)
/// - When you want zero-centered data
/// 
/// # Notes
/// 
/// Trees are scale-invariant, so scaling is optional for gradient boosting.
/// However, it can help with regularization and numerical stability.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StandardScaler {
    mean: Option<Array1<f64>>,
    scale: Option<Array1<f64>>,
    epsilon: f64,
}

impl Default for StandardScaler {
    fn default() -> Self {
        Self {
            mean: None,
            scale: None,
            epsilon: 1e-8,
        }
    }
}

impl StandardScaler {
    /// Creates a new unsc fitted StandardScaler with default epsilon.
    pub fn new() -> Self {
        Self::default()
    }
   
    /// Creates a new StandardScaler with custom epsilon.
    /// 
    /// # Parameters
    /// - `epsilon`: Small constant added to denominator for numerical stability.
    ///   Set to 1e-8 by default. Increase for very small features.
    pub fn with_epsilon(epsilon: f64) -> Self {
        Self {
            mean: None,
            scale: None,
            epsilon,
        }
    }

    /// Returns true if the scaler has been fitted.
    pub fn is_fitted(&self) -> bool {
        self.mean.is_some() && self.scale.is_some()
    }
}

impl Scaler for StandardScaler {
    fn fit(&mut self, data: &FeatureMatrix) -> DataResult<()> {
        let n_features = data.n_features();
        
        let mean: Array1<f64> = (0..n_features)
            .map(|i| {
                let feature = data.get_feature(i).unwrap();
                feature.mean().unwrap_or(0.0)
            })
            .collect();
            
        let scale: Array1<f64> = (0..n_features)
            .map(|i| {
                let feature = data.get_feature(i).unwrap();
                let variance = feature.variance().ok().unwrap_or(0.0);
                (variance + self.epsilon).sqrt()
            })
            .collect();
            
        self.mean = Some(mean);
        self.scale = Some(scale);
        
        Ok(())
    }
    
    fn transform(&self, data: &FeatureMatrix) -> DataResult<FeatureMatrix> {
        let mean = self.mean.as_ref()
            .ok_or_else(|| DataError::PreprocessingError("Scaler not fitted".to_string()))?;
        let scale = self.scale.as_ref()
            .ok_or_else(|| DataError::PreprocessingError("Scaler not fitted".to_string()))?;
            
        if data.n_features() != mean.len() {
            return Err(DataError::PreprocessingError(
                format!("Expected {} features, got {}", mean.len(), data.n_features())
            ));
        }
        
        let mut transformed_data = data.data().clone();
        
        for i in 0..data.n_features() {
            let mut column = transformed_data.column_mut(i);
            for j in 0..column.len() {
                column[j] = (column[j] - mean[i]) / scale[i];
            }
        }
        
        FeatureMatrix::with_feature_names(
            transformed_data,
            data.feature_names().to_vec()
        ).map_err(Into::into)
    }
}

// ============================================================================
// MinMaxScaler - scale features to a given range (default [0, 1])
// ============================================================================

/// Scales features to a given range (default [0, 1]).
///
/// # Formula
///
/// For each feature column `x`:
///
/// ```
/// z = (x - x_min) / (x_max - x_min) * (max_range - min_range) + min_range
/// ```
///
/// where:
/// - `x_min` is the minimum value seen during fitting
/// - `x_max` is the maximum value seen during fitting
/// - `max_range` and `min_range` are the target range bounds
///
/// # When to Use
///
/// - When you need bounded feature ranges
/// - For neural networks (input normalization)
/// - When preserving zero values is important
/// - For visualization purposes
///
/// # Handling Constant Features
///
/// If a feature has zero range (all values identical), the scaler
/// sets the scale factor to 1.0 to avoid division by zero.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MinMaxScaler {
    min: Option<Array1<f64>>,
    scale: Option<Array1<f64>>,
    feature_range: (f64, f64),
}

impl Default for MinMaxScaler {
    fn default() -> Self {
        Self {
            min: None,
            scale: None,
            feature_range: (0.0, 1.0),
        }
    }
}

impl MinMaxScaler {
    /// Creates a new unsc fitted MinMaxScaler with default [0, 1] range.
    pub fn new() -> Self {
        Self::default()
    }
   
    /// Creates a new MinMaxScaler with custom output range.
    /// 
    /// # Parameters
    /// - `min`: Lower bound of output range
    /// - `max`: Upper bound of output range
    /// 
    /// # Panics
    /// Will cause transformation errors if `min >= max`.
    pub fn with_range(min: f64, max: f64) -> Self {
        Self {
            min: None,
            scale: None,
            feature_range: (min, max),
        }
    }
   
    /// Returns true if the scaler has been fitted.
    pub fn is_fitted(&self) -> bool {
        self.min.is_some() && self.scale.is_some()
    }
}

impl Scaler for MinMaxScaler {
    fn fit(&mut self, data: &FeatureMatrix) -> DataResult<()> {
        // Get min values for each feature using fold_axis
        let data_min = data.data().fold_axis(
            ndarray::Axis(0),
            f64::INFINITY,
            |&min, &x| min.min(x)
        );
        
        // Get max values for each feature using fold_axis
        let data_max = data.data().fold_axis(
            ndarray::Axis(0),
            f64::NEG_INFINITY,
            |&max, &x| max.max(x)
        );
        
        // Calculate range (max - min) for each feature
        let data_range = &data_max - &data_min;
        
        // Handle constant features to avoid division by zero
        let scale: Array1<f64> = data_range.iter()
            .map(|&range| {
                if range == 0.0 {
                    1.0
                } else {
                    (self.feature_range.1 - self.feature_range.0) / range
                }
            })
            .collect();
            
        self.min = Some(data_min);
        self.scale = Some(scale);
        
        Ok(())
    }
    
    fn transform(&self, data: &FeatureMatrix) -> DataResult<FeatureMatrix> {
        let data_min = self.min.as_ref()
            .ok_or_else(|| DataError::PreprocessingError("Scaler not fitted".to_string()))?;
        let scale = self.scale.as_ref()
            .ok_or_else(|| DataError::PreprocessingError("Scaler not fitted".to_string()))?;
            
        if data.n_features() != data_min.len() {
            return Err(DataError::PreprocessingError(
                format!("Expected {} features, got {}", data_min.len(), data.n_features())
            ));
        }
        
        let mut transformed_data = data.data().clone();
        
        for i in 0..data.n_features() {
            let mut column = transformed_data.column_mut(i);
            for j in 0..column.len() {
                column[j] = (column[j] - data_min[i]) * scale[i] + self.feature_range.0;
            }
        }
        
        FeatureMatrix::with_feature_names(
            transformed_data,
            data.feature_names().to_vec()
        ).map_err(Into::into)
    }
}

// ============================================================================
// CategoricalEncoder - encode string categories as numerical values
// ============================================================================

/// Encodes string categorical features into numerical values.
/// 
/// # Encoding Strategy
/// 
/// - Each unique string gets an integer code starting from 0
/// - Unknown categories (not seen during fitting) get a default value
/// - Numerical columns (parseable as f64) are passed through unchanged
/// 
/// # Usage Pattern
/// 
/// 1. `fit_transform()`: Learn mappings and encode training data
/// 2. `transform()`: Encode new data using existing mappings
/// 
/// # Unknown Categories
/// 
/// Categories not present during fitting are mapped to `default_value`
/// (default: -1.0). This prevents errors but may affect model performance.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CategoricalEncoder {
    /// Mapping: column_index -> (category_string -> numeric_code)
    mappings: HashMap<String, HashMap<String, f64>>,
    /// Value for unknown categories
    default_value: f64,
}

impl Default for CategoricalEncoder {
    fn default() -> Self {
        Self::new()
    }
}

impl CategoricalEncoder {
    /// Creates a new encoder with default value -1.0 for unknown categories.
    pub fn new() -> Self {
        Self {
            mappings: HashMap::new(),
            default_value: -1.0, // Use -1.0 for unknown categories
        }
    }
    
    /// Creates a new encoder with custom default value.
    /// 
    /// # Parameters
    /// - `default_value`: Value to use for categories not seen during fitting
    pub fn with_default(default_value: f64) -> Self {
        Self {
            mappings: HashMap::new(),
            default_value,
        }
    }
    

    /// Fits encoder to data and transforms it in one pass.
    /// 
    /// # Parameters
    /// - `data`: Dataset as strings (CSV-like structure)
    /// - `categorical_columns`: Column indices to treat as categorical
    /// 
    /// # Returns
    /// Encoded numerical data
    /// 
    /// # Errors
    /// - `PreprocessingError::InvalidData` if data is empty or inconsistent
    /// - `PreprocessingError::EncodingError` if numerical parsing fails
    pub fn fit_transform(
        &mut self, 
        data: &[Vec<String>], 
        categorical_columns: &[usize]
    ) -> Result<Vec<Vec<f64>>, PreprocessingError> {
        if data.is_empty() {
            return Err(PreprocessingError::InvalidData("Empty dataset".to_string()));
        }
        
        let n_features = data[0].len();
        let n_samples = data.len();
        let mut encoded_data = Vec::with_capacity(n_samples);
        
        // Validate categorical column indices
        for &col_idx in categorical_columns {
            if col_idx >= n_features {
                return Err(PreprocessingError::InvalidData(
                    format!("Column index {} out of bounds (max: {})", col_idx, n_features - 1)
                ));
            }
        }
        
        // Initialize mappings for each categorical column
        for &col_idx in categorical_columns {
            self.mappings.insert(col_idx.to_string(), HashMap::new());
        }
        
        // First pass: collect unique values for each categorical column
        for &col_idx in categorical_columns {
            let mut unique_values = Vec::new();
            for row in data {
                let value = &row[col_idx];
                if !unique_values.contains(value) {
                    unique_values.push(value.clone());
                }
            }
            
            // Sort for consistent encoding
            unique_values.sort();
            
            // Create mapping (value -> index)
            let mapping = self.mappings.get_mut(&col_idx.to_string()).unwrap();
            for (idx, value) in unique_values.iter().enumerate() {
                mapping.insert(value.clone(), idx as f64);
            }
        }
        
        // Second pass: encode data
        for row in data {
            let mut encoded_row = Vec::with_capacity(n_features);
            
            for (col_idx, value) in row.iter().enumerate() {
                if self.mappings.contains_key(&col_idx.to_string()) {
                    // Categorical column - use encoding
                    let mapping = self.mappings.get(&col_idx.to_string()).unwrap();
                    let encoded_value = mapping.get(value).unwrap_or(&self.default_value);
                    encoded_row.push(*encoded_value);
                } else {
                    // Numerical column - try to parse as f64
                    match value.parse::<f64>() {
                        Ok(num) => encoded_row.push(num),
                        Err(_) => {
                            return Err(PreprocessingError::EncodingError(
                                format!("Failed to parse numerical value '{}' in column {}", value, col_idx)
                            ));
                        }
                    }
                }
            }
            
            encoded_data.push(encoded_row);
        }
        
        Ok(encoded_data)
    }
    
    /// Transforms new data using existing mappings.
    /// 
    /// # Parameters
    /// - `data`: New dataset as strings
    /// 
    /// # Returns
    /// Encoded numerical data
    /// 
    /// # Errors
    /// - `PreprocessingError::InvalidData` if row lengths are inconsistent
    /// - `PreprocessingError::EncodingError` if numerical parsing fails
    /// 
    /// # Unknown Categories
    /// Categories not seen during fitting are mapped to `default_value`.
    pub fn transform(&self, data: &[Vec<String>]) -> Result<Vec<Vec<f64>>, PreprocessingError> {
        if data.is_empty() {
            return Ok(Vec::new());
        }
        
        let n_features = data[0].len();
        let mut encoded_data = Vec::with_capacity(data.len());
        
        for row in data {
            if row.len() != n_features {
                return Err(PreprocessingError::InvalidData(
                    format!("Inconsistent row length: expected {}, got {}", n_features, row.len())
                ));
            }
            
            let mut encoded_row = Vec::with_capacity(n_features);
            
            for (col_idx, value) in row.iter().enumerate() {
                if let Some(mapping) = self.mappings.get(&col_idx.to_string()) {
                    // Categorical column - use existing mapping
                    let encoded_value = mapping.get(value).unwrap_or(&self.default_value);
                    encoded_row.push(*encoded_value);
                } else {
                    // Numerical column - try to parse as f64
                    match value.parse::<f64>() {
                        Ok(num) => encoded_row.push(num),
                        Err(_) => {
                            return Err(PreprocessingError::EncodingError(
                                format!("Failed to parse numerical value '{}' in column {}", value, col_idx)
                            ));
                        }
                    }
                }
            }
            
            encoded_data.push(encoded_row);
        }
        
        Ok(encoded_data)
    }
    
    /// Gets the mapping for a specific column.
    /// 
    /// # Parameters
    /// - `column`: Column index
    /// 
    /// # Returns
    /// `Some(&HashMap)` if column was fitted as categorical, `None` otherwise
    pub fn get_mapping(&self, column: usize) -> Option<&HashMap<String, f64>> {
        self.mappings.get(&column.to_string())
    }
    
    /// Gets the number of unique categories for a column.
    /// 
    /// # Parameters
    /// - `column`: Column index
    /// 
    /// # Returns
    /// `Some(count)` if column was fitted as categorical, `None` otherwise
    pub fn n_categories(&self, column: usize) -> Option<usize> {
        self.get_mapping(column).map(|m| m.len())
    }
}

// ============================================================================
// Helper Functions
// ============================================================================

/// Automatically detects categorical columns based on data characteristics.
/// 
/// Uses two heuristics:
/// 1. **Data type**: Columns containing non-numeric values are categorical
/// 2. **Cardinality**: Numeric columns with low uniqueness ratio (< threshold) are categorical
/// 
/// # Parameters
/// - `data`: Dataset as strings (CSV-like structure)
/// - `threshold`: Ratio threshold (unique_values / total_samples). Default: 0.1
/// 
/// # Returns
/// Vector of column indices identified as categorical
pub fn detect_categorical_columns(data: &[Vec<String>], threshold: f64) -> Vec<usize> {
    if data.is_empty() {
        return Vec::new();
    }
    
    let n_samples = data.len();
    let n_features = data[0].len();
    let mut categorical_columns = Vec::new();
    
    for col_idx in 0..n_features {
        let mut unique_values = HashSet::new();
        let mut is_numeric = true;
        
        for row in data {
            let value = &row[col_idx];
            
            // Check if value can be parsed as f64
            if value.parse::<f64>().is_err() {
                is_numeric = false;
            }
            
            unique_values.insert(value.clone());
        }
        
        // If not numeric OR if numeric but has few unique values (could be categorical codes)
        let uniqueness_ratio = unique_values.len() as f64 / n_samples as f64;
        if !is_numeric || uniqueness_ratio < threshold {
            categorical_columns.push(col_idx);
        }
    }
    
    categorical_columns
}



// ============================================================================
// DataPreprocessor - Combine multiple preprocessing steps
// ============================================================================

/// High-level preprocessor that combines categorical encoding and scaling.
/// 
/// This convenience struct orchestrates multiple preprocessing steps:
/// 1. Categorical encoding (if encoder provided)
/// 2. Numerical scaling (if scaler provided)
/// 
pub struct DataPreprocessor {
    categorical_encoder: Option<CategoricalEncoder>,
    scaler: Option<Box<dyn Scaler>>,
}

impl DataPreprocessor {
    /// Creates a new preprocessor with no steps configured.   
    pub fn new() -> Self {
        Self {
            categorical_encoder: None,
            scaler: None,
        }
    }
   
    /// Adds categorical encoding step.
    /// 
    /// # Parameters
    /// - `encoder`: Categorical encoder instance
    /// 
    /// # Returns
    /// Self for method chaining
    pub fn with_categorical_encoder(mut self, encoder: CategoricalEncoder) -> Self {
        self.categorical_encoder = Some(encoder);
        self
    }
   
    /// Adds scaling step.
    /// 
    /// # Parameters
    /// - `scaler`: Boxed scaler trait object (e.g., Box::new(StandardScaler::new()))
    /// 
    /// # Returns
    /// Self for method chaining
    pub fn with_scaler(mut self, scaler: Box<dyn Scaler>) -> Self {
        self.scaler = Some(scaler);
        self
    }
    
    /// Fits encoder (if any) and transforms data, then applies scaler (if any).
    /// 
    /// # Parameters
    /// - `data`: Raw string data (CSV-like)
    /// - `categorical_columns`: Column indices to treat as categorical
    /// 
    /// # Returns
    /// Preprocessed feature matrix
    /// 
    /// # Errors
    /// - `PreprocessingError::InvalidData` if encoding fails
    /// - `PreprocessingError::EncodingError` if numerical parsing fails
    /// - `DataError` if scaler fitting/transforming fails
    pub fn fit_transform(
        &mut self,
        data: &[Vec<String>],
        categorical_columns: &[usize],
    ) -> Result<FeatureMatrix, PreprocessingError> {
        // Encode categorical data
        let mut encoder = CategoricalEncoder::new();
        let encoded_data = encoder.fit_transform(data, categorical_columns)?;
        
        if self.categorical_encoder.is_none() {
            self.categorical_encoder = Some(encoder);
        }
        
        // Convert to FeatureMatrix
        let n_samples = encoded_data.len();
        let n_features = encoded_data[0].len();
        
        let flat_data: Vec<f64> = encoded_data.into_iter().flatten().collect();
        let array = Array2::from_shape_vec((n_samples, n_features), flat_data)
            .map_err(|e| PreprocessingError::InvalidData(e.to_string()))?;
        
        let feature_matrix = FeatureMatrix::new(array)
            .map_err(|e| PreprocessingError::InvalidData(e.to_string()))?;
        
        // Apply scaler if specified
        if let Some(scaler) = &mut self.scaler {
            scaler.fit(&feature_matrix)
                .map_err(|e| PreprocessingError::InvalidData(e.to_string()))?;
            return scaler.transform(&feature_matrix)
                .map_err(|e| PreprocessingError::InvalidData(e.to_string()));
        }
        
        Ok(feature_matrix)
    }
}