gbrt_rs/data/

feature_matrix.rs

//! Feature Matrix for Efficient Column-Oriented Access
//! 
//! This module provides [`FeatureMatrix`], a matrix structure optimized for
//! feature-wise operations common in machine learning, particularly decision
//! tree algorithms. Unlike typical row-major storage, this uses column-major
//! layout for efficient feature access during split finding.
//! 
//! # Design Decisions
//! 
//! - **Column-major storage**: [`ndarray::Array2`] stores data column-wise,
//!   making feature extraction (columns) cache-friendly during tree building
//! - **Named features**: Each column can have a human-readable name
//! - **Type safety**: Custom errors for index bounds and shape validation
//! - **Parallel-ready**: Integrates with Rayon for parallel feature processing
//! 
//! # When to Use
//! 
//! This structure is ideal when you need frequent column-wise access,
//! such as:
//! - Decision tree split finding
//! - Feature importance computation
//! - Feature scaling and normalization
//! - Feature selection operations
//!

use ndarray::{Array2, Array1, Axis, ArrayView1, ArrayView2};
use serde::{Serialize, Deserialize};
use std::path::Path;
use rayon::prelude::*;

/// Errors that can occur when working with feature matrices.
#[derive(thiserror::Error, Debug)]
pub enum FeatureMatrixError {
    /// Invalid matrix shape (zero rows or columns).
    #[error("Invalid shape: expected at least 1 row and 1 column, got {rows}×{cols}")]
    InvalidShape { rows: usize, cols: usize },
    
    /// Feature column index out of valid range.
    #[error("Feature index out of bounds: {index} (max: {max})")]
    FeatureIndexOutOfBounds { index: usize, max: usize },
    
    /// Sample row index out of valid range.    
    #[error("Sample index out of bounds: {index} (max: {max})")]
    SampleIndexOutOfBounds { index: usize, max: usize },

    /// CSV parsing error.    
    #[error("CSV error: {0}")]
    CsvError(#[from] csv::Error),

    /// File I/O error.    
    #[error("IO error: {0}")]
    IoError(#[from] std::io::Error),
}

/// A feature matrix optimized for column-wise access in machine learning.
/// 
/// [`FeatureMatrix`] stores data in column-major format using [`ndarray::Array2`],
/// making feature extraction efficient for algorithms that process columns
/// repeatedly. This is particularly beneficial for decision trees during split
/// finding.
/// 
/// # Invariants
/// 
/// - `n_samples > 0 && n_features > 0`
/// - `feature_names.len() == n_features`
/// - All data values are finite (validated at construction)
/// 
/// # Thread Safety
/// 
/// Cloning creates a new copy of the data; the struct is not designed for
/// shared mutable access. Use immutable references or ownership passing.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FeatureMatrix {
    /// Underlying column-major data storage
    data: Array2<f64>,
    /// Feature/column names for human readability
    feature_names: Vec<String>,
}

impl FeatureMatrix {
    /// Creates a new feature matrix with auto-generated names.
    /// 
    /// Generates feature names in the format "feature_0", "feature_1", etc.
    /// 
    /// # Parameters
    /// - `data`: 2D array with shape (n_samples, n_features)
    /// 
    /// # Returns
    /// `Ok(FeatureMatrix)` if validation passes
    /// 
    /// # Errors
    /// - `FeatureMatrixError::InvalidShape` if data has zero rows or columns
    pub fn new(data: Array2<f64>) -> std::result::Result<Self, FeatureMatrixError> {
        let (n_samples, n_features) = data.dim();
        
        if n_samples == 0 || n_features == 0 {
            return Err(FeatureMatrixError::InvalidShape { 
                rows: n_samples, 
                cols: n_features 
            });
        }
        
        let feature_names = (0..n_features)
            .map(|i| format!("feature_{}", i))
            .collect();
            
        Ok(FeatureMatrix { data, feature_names })
    }
    
    /// Creates a new feature matrix with custom feature names.
    /// 
    /// # Parameters
    /// - `data`: 2D array with shape (n_samples, n_features)
    /// - `feature_names`: Vector of names, one per column
    /// 
    /// # Returns
    /// `Ok(FeatureMatrix)` if validation passes
    /// 
    /// # Errors
    /// - `FeatureMatrixError::InvalidShape` if data has zero rows or columns
    /// - `FeatureMatrixError::InvalidShape` if feature_names length doesn't match columns
    pub fn with_feature_names(
        data: Array2<f64>, 
        feature_names: Vec<String>
    ) -> std::result::Result<Self, FeatureMatrixError> {
        let (n_samples, n_features) = data.dim();
        
        if n_samples == 0 || n_features == 0 {
            return Err(FeatureMatrixError::InvalidShape { 
                rows: n_samples, 
                cols: n_features 
            });
        }
        
        if feature_names.len() != n_features {
            return Err(FeatureMatrixError::InvalidShape {
                rows: n_samples,
                cols: n_features,
            });
        }
        
        Ok(FeatureMatrix { data, feature_names })
    }
    
    /// Returns the number of samples (rows) in the matrix.
    pub fn n_samples(&self) -> usize {
        self.data.nrows()
    }
    
    /// Returns the number of features (columns) in the matrix.
    pub fn n_features(&self) -> usize {
        self.data.ncols()
    }
    
    /// Returns the shape as (n_samples, n_features).
    pub fn shape(&self) -> (usize, usize) {
        self.data.dim()
    }
    
    /// Returns a reference to the underlying data array.
    pub fn data(&self) -> &Array2<f64> {
        &self.data
    }
    
    /// Returns a slice of feature names.
    pub fn feature_names(&self) -> &[String] {
        &self.feature_names
    }
    
    /// Gets a column view for a specific feature.
    /// 
    /// # Parameters
    /// - `feature_idx`: Zero-based column index
    /// 
    /// # Returns
    /// Array view of shape (n_samples,)
    /// 
    /// # Errors
    /// - `FeatureMatrixError::FeatureIndexOutOfBounds` if index >= n_features
    pub fn get_feature(&self, feature_idx: usize) -> std::result::Result<ArrayView1<'_, f64>, FeatureMatrixError> {
        if feature_idx >= self.n_features() {
            return Err(FeatureMatrixError::FeatureIndexOutOfBounds { 
                index: feature_idx, 
                max: self.n_features() - 1 
            });
        }
        Ok(self.data.column(feature_idx))
    }
    
    /// Gets a row view for a specific sample.
    /// 
    /// # Parameters
    /// - `sample_idx`: Zero-based row index
    /// 
    /// # Returns
    /// Array view of shape (n_features,)
    /// 
    /// # Errors
    /// - `FeatureMatrixError::SampleIndexOutOfBounds` if index >= n_samples
    pub fn get_sample(&self, sample_idx: usize) -> std::result::Result<ArrayView1<'_, f64>, FeatureMatrixError> {
        if sample_idx >= self.n_samples() {
            return Err(FeatureMatrixError::SampleIndexOutOfBounds { 
                index: sample_idx, 
                max: self.n_samples() - 1 
            });
        }
        Ok(self.data.row(sample_idx))
    }
    
    /// Gets a single element at (sample_idx, feature_idx).
    /// 
    /// # Parameters
    /// - `sample_idx`: Row index
    /// - `feature_idx`: Column index
    /// 
    /// # Returns
    /// The scalar value at that position
    /// 
    /// # Errors
    /// - `FeatureMatrixError::SampleIndexOutOfBounds` if sample_idx >= n_samples
    /// - `FeatureMatrixError::FeatureIndexOutOfBounds` if feature_idx >= n_features
    pub fn get(&self, sample_idx: usize, feature_idx: usize) -> std::result::Result<f64, FeatureMatrixError> {
        if sample_idx >= self.n_samples() {
            return Err(FeatureMatrixError::SampleIndexOutOfBounds { 
                index: sample_idx, 
                max: self.n_samples() - 1 
            });
        }
        if feature_idx >= self.n_features() {
            return Err(FeatureMatrixError::FeatureIndexOutOfBounds { 
                index: feature_idx, 
                max: self.n_features() - 1 
            });
        }
        Ok(self.data[[sample_idx, feature_idx]])
    }
    
    /// Gets a view of the entire matrix.
    /// 
    /// Useful for passing to ndarray-based algorithms.
    pub fn view(&self) -> ArrayView2<'_, f64> {
        self.data.view()
    }
    
    /// Returns an iterator over samples (rows).
    /// 
    /// Each iteration yields a view of shape (n_features,).
    pub fn samples(&self) -> impl Iterator<Item = ArrayView1<'_, f64>> + '_ {
        self.data.rows().into_iter()
    }
    
    /// Returns an iterator over features (columns).
    /// 
    /// Each iteration yields a view of shape (n_samples,).
    /// This is cache-friendly due to column-major storage.
    pub fn features(&self) -> impl Iterator<Item = ArrayView1<'_, f64>> + '_ {
        self.data.columns().into_iter()
    }
    
    /// Computes the min and max range for each feature.
    /// 
    /// Useful for decision tree split finding and feature scaling.
    /// 
    /// # Returns
    /// Vector of (min, max) tuples, one per feature
    pub fn feature_ranges(&self) -> Vec<(f64, f64)> {
        (0..self.n_features())
            .map(|i| {
                let col = self.data.column(i);
                let min = col.fold(f64::INFINITY, |a, &b| a.min(b));
                let max = col.fold(f64::NEG_INFINITY, |a, &b| a.max(b));
                (min, max)
            })
            .collect()
    }
    
    /// Computes the mean of each feature.
    /// 
    /// Useful for imputation and model initialization.
    /// 
    /// # Returns
    /// Vector of means, one per feature
    pub fn feature_means(&self) -> Vec<f64> {
        (0..self.n_features())
            .map(|i| {
                let col = self.data.column(i);
                col.mean().unwrap_or(0.0)
            })
            .collect()
    }

    /// Creates a submatrix with selected samples (rows).
    /// 
    /// # Parameters
    /// - `sample_indices`: Row indices to include in the subset
    /// 
    /// # Returns
    /// New FeatureMatrix containing only selected rows
    /// 
    /// # Errors
    /// - `FeatureMatrixError::SampleIndexOutOfBounds` if any index >= n_samples
    pub fn select_samples(&self, sample_indices: &[usize]) -> std::result::Result<Self, FeatureMatrixError> {
        let mut selected_data = Vec::new();

        for &idx in sample_indices {
            if idx >= self.n_samples() {
                return Err(FeatureMatrixError::SampleIndexOutOfBounds {
                    index: idx,
                    max: self.n_samples() - 1
                });
            }
            selected_data.push(self.data.row(idx).to_owned());
        }

        // Stack rows into a new matrix
        let data = ndarray::stack(
            Axis(0),
            &selected_data.iter().map(|v| v.view()).collect::<Vec<_>>()
        ).map_err(|_| FeatureMatrixError::InvalidShape {
            rows: sample_indices.len(),
            cols: self.n_features(),
        })?;

        Ok(FeatureMatrix {
            data,
            feature_names: self.feature_names.clone(),
        })
    }
    
    /// Creates a submatrix with selected features (columns).
    /// 
    /// # Parameters
    /// - `feature_indices`: Column indices to include in the subset
    /// 
    /// # Returns
    /// New FeatureMatrix containing only selected columns
    /// 
    /// # Errors
    /// - `FeatureMatrixError::FeatureIndexOutOfBounds` if any index >= n_features
    pub fn select_features(&self, feature_indices: &[usize]) -> std::result::Result<Self, FeatureMatrixError> {
        let mut selected_data = Vec::new();
        let mut selected_names = Vec::new();
        
        for &idx in feature_indices {
            if idx >= self.n_features() {
                return Err(FeatureMatrixError::FeatureIndexOutOfBounds {
                    index: idx,
                    max: self.n_features() - 1,
                });
            }
            selected_data.push(self.data.column(idx).to_owned());
            selected_names.push(self.feature_names[idx].clone());
        }
        
        // Stack columns into a new matrix
        let data = ndarray::stack(
            Axis(1),
            &selected_data.iter().map(|v| v.view()).collect::<Vec<_>>()
        ).map_err(|_| FeatureMatrixError::InvalidShape {
            rows: self.n_samples(),
            cols: feature_indices.len(),
        })?;
        
        Ok(FeatureMatrix {
            data,
            feature_names: selected_names,
        })
    }
    
    /// Loads a feature matrix from a CSV file.
    /// 
    /// # Parameters
    /// - `path`: Path to CSV file with headers
    /// 
    /// # Returns
    /// FeatureMatrix with data and column names
    /// 
    /// # Errors
    /// - `FeatureMatrixError::CsvError` if CSV parsing fails
    /// - `FeatureMatrixError::IoError` if file cannot be read
    /// - `FeatureMatrixError::InvalidShape` if file is empty
    /// 
    /// # CSV Format
    /// - First row: column headers (become feature names)
    /// - Subsequent rows: numeric data (NaN for missing values)
    pub fn from_csv<P: AsRef<Path>>(path: P) -> std::result::Result<Self, FeatureMatrixError> {
        let mut reader = csv::ReaderBuilder::new()
            .has_headers(true)
            .from_path(path)?;
            
        let headers: Vec<String> = reader
            .headers()?
            .iter()
            .map(|s| s.to_string())
            .collect();
            
        let mut data = Vec::new();
        
        for result in reader.records() {
            let record = result?;
            let row: Vec<f64> = record
                .iter()
                .map(|s| s.parse().unwrap_or(f64::NAN))
                .collect();
            data.push(row);
        }
        
        let n_samples = data.len();
        let n_features = headers.len();
        
        if n_samples == 0 {
            return Err(FeatureMatrixError::InvalidShape { rows: 0, cols: n_features });
        }
        
        let flat_data: Vec<f64> = data.into_iter().flatten().collect();
        let array = Array2::from_shape_vec((n_samples, n_features), flat_data)
            .map_err(|_| FeatureMatrixError::InvalidShape { rows: n_samples, cols: n_features })?;
            
        Ok(FeatureMatrix {
            data: array,
            feature_names: headers,
        })
    }
}

/// Equality comparison for feature matrices.
///
/// Two matrices are equal if they have:
/// - Identical underlying data (element-wise)
/// - Same feature names (in same order)
impl PartialEq for FeatureMatrix {
    fn eq(&self, other: &Self) -> bool {
        self.data == other.data && self.feature_names == other.feature_names
    }
}