imbalanced_metrics/

lib.rs

//! Performance metrics for imbalanced learning
//!
//! This crate provides metrics specifically useful for evaluating
//! models on imbalanced datasets.

#![deny(missing_docs)]
#![warn(clippy::all, clippy::pedantic)]
#![allow(clippy::module_name_repetitions)]

use ndarray::ArrayView1;
use rayon::prelude::*;
use std::collections::HashMap;

/// Classification report containing precision, recall, F1-score, and support
#[derive(Debug, Clone)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct ClassificationReport {
    /// Per-class metrics
    pub class_metrics: HashMap<usize, ClassMetrics>,
    /// Macro-averaged metrics
    pub macro_avg: AggregateMetrics,
    /// Weighted-averaged metrics
    pub weighted_avg: AggregateMetrics,
    /// Overall accuracy
    pub accuracy: f64,
}

/// Metrics for a single class
#[derive(Debug, Clone, Copy)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct ClassMetrics {
    /// Precision: TP / (TP + FP)
    pub precision: f64,
    /// Recall: TP / (TP + FN)
    pub recall: f64,
    /// F1-score: 2 * (precision * recall) / (precision + recall)
    pub f1_score: f64,
    /// Number of true instances for this class
    pub support: usize,
}

/// Aggregate metrics (macro/weighted averages)
#[derive(Debug, Clone, Copy)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct AggregateMetrics {
    /// Average precision
    pub precision: f64,
    /// Average recall
    pub recall: f64,
    /// Average F1-score
    pub f1_score: f64,
}

/// Compute confusion matrix
pub fn confusion_matrix(y_true: ArrayView1<usize>, y_pred: ArrayView1<usize>) -> Array2<usize> {
    use ndarray::Array2;
    
    assert_eq!(y_true.len(), y_pred.len(), "Arrays must have same length");
    
    let n_classes = y_true.iter().chain(y_pred.iter()).max().unwrap() + 1;
    let mut matrix = Array2::zeros((n_classes, n_classes));
    
    for (&true_label, &pred_label) in y_true.iter().zip(y_pred.iter()) {
        matrix[[true_label, pred_label]] += 1;
    }
    
    matrix
}

/// Generate classification report
pub fn classification_report(
    y_true: ArrayView1<usize>, 
    y_pred: ArrayView1<usize>
) -> ClassificationReport {
    let cm = confusion_matrix(y_true, y_pred);
    let n_classes = cm.nrows();
    
    let class_metrics: HashMap<usize, ClassMetrics> = (0..n_classes)
        .into_par_iter()
        .map(|class| {
            let tp = cm[[class, class]] as f64;
            let fp = cm.column(class).sum() as f64 - tp;
            let false_neg = cm.row(class).sum() as f64 - tp;
            
            let precision = if tp + fp > 0.0 { tp / (tp + fp) } else { 0.0 };
            let recall = if tp + false_neg > 0.0 { tp / (tp + false_neg) } else { 0.0 };
            let f1_score = if precision + recall > 0.0 {
                2.0 * precision * recall / (precision + recall)
            } else {
                0.0
            };
            
            let support = cm.row(class).sum();
            
            (class, ClassMetrics { precision, recall, f1_score, support })
        })
        .collect();
    
    // Compute macro averages
    let macro_avg = AggregateMetrics {
        precision: class_metrics.values().map(|m| m.precision).sum::<f64>() / n_classes as f64,
        recall: class_metrics.values().map(|m| m.recall).sum::<f64>() / n_classes as f64,
        f1_score: class_metrics.values().map(|m| m.f1_score).sum::<f64>() / n_classes as f64,
    };
    
    // Compute weighted averages
    let total_support = class_metrics.values().map(|m| m.support).sum::<usize>() as f64;
    let weighted_avg = AggregateMetrics {
        precision: class_metrics.values()
            .map(|m| m.precision * m.support as f64)
            .sum::<f64>() / total_support,
        recall: class_metrics.values()
            .map(|m| m.recall * m.support as f64)
            .sum::<f64>() / total_support,
        f1_score: class_metrics.values()
            .map(|m| m.f1_score * m.support as f64)
            .sum::<f64>() / total_support,
    };
    
    let accuracy = cm.diag().sum() as f64 / cm.sum() as f64;
    
    ClassificationReport {
        class_metrics,
        macro_avg,
        weighted_avg,
        accuracy,
    }
}

/// Calculate F1 score for binary or multiclass classification
pub fn f1_score(y_true: ArrayView1<usize>, y_pred: ArrayView1<usize>) -> f64 {
    let report = classification_report(y_true, y_pred);
    report.weighted_avg.f1_score
}

/// Calculate balanced accuracy
pub fn balanced_accuracy(y_true: ArrayView1<usize>, y_pred: ArrayView1<usize>) -> f64 {
    let report = classification_report(y_true, y_pred);
    report.macro_avg.recall  // Balanced accuracy is macro-averaged recall
}

/// Display classification report
impl std::fmt::Display for ClassificationReport {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        writeln!(f, "\n              precision    recall  f1-score   support\n")?;
        
        let mut classes: Vec<_> = self.class_metrics.keys().copied().collect();
        classes.sort_unstable();
        
        for class in classes {
            let metrics = &self.class_metrics[&class];
            writeln!(
                f,
                "   class {:2}      {:5.2}      {:5.2}      {:5.2}      {:4}",
                class, metrics.precision, metrics.recall, metrics.f1_score, metrics.support
            )?;
        }
        
        writeln!(f, "\n   accuracy                          {:5.2}", self.accuracy)?;
        writeln!(
            f,
            "  macro avg      {:5.2}      {:5.2}      {:5.2}",
            self.macro_avg.precision, self.macro_avg.recall, self.macro_avg.f1_score
        )?;
        writeln!(
            f,
            "weighted avg      {:5.2}      {:5.2}      {:5.2}",
            self.weighted_avg.precision, self.weighted_avg.recall, self.weighted_avg.f1_score
        )?;
        
        Ok(())
    }
}

/// Prelude module for convenient imports
pub mod prelude {
    pub use crate::{
        classification_report, confusion_matrix, f1_score, balanced_accuracy,
        ClassificationReport, ClassMetrics, AggregateMetrics,
    };
}

// Missing import
use ndarray::Array2;