rustytorch/crates/training/rtx-preprocessing/tests/imputer_tests.rs

#![allow(clippy::approx_constant)]
use approx::assert_abs_diff_eq;
use rtx_preprocessing::{Imputer, PreprocessingError, Transformer};
use rtx_tensor::{Device, Tensor};

#[test]
fn test_imputer_creation() {
    let imputer = Imputer::new();
    assert!(!imputer.is_fitted());
    assert_eq!(imputer.strategy(), "mean");

    let imputer = Imputer::with_strategy("median");
    assert_eq!(imputer.strategy(), "median");

    let imputer = Imputer::with_strategy("most_frequent");
    assert_eq!(imputer.strategy(), "most_frequent");

    let imputer = Imputer::with_constant(42.0);
    assert_eq!(imputer.strategy(), "constant");
    assert_abs_diff_eq!(imputer.fill_value(), 42.0, epsilon = 1e-5);
}

#[test]
fn test_imputer_invalid_strategy() {
    let result = std::panic::catch_unwind(|| Imputer::with_strategy("invalid"));
    assert!(result.is_err());
}

#[test]
fn test_imputer_fit_mean_strategy() {
    let data = Tensor::from_slice(
        &[1.0f32, f32::NAN, 3.0, 4.0, f32::NAN],
        &[5, 1],
        &Device::cpu(),
    )
    .unwrap();
    let mut imputer = Imputer::new(); // Default is mean

    let result = imputer.fit(&data);
    assert!(result.is_ok());
    assert!(imputer.is_fitted());

    // Mean of [1, 3, 4] = 8/3 ≈ 2.667
    let statistics = imputer.statistics();
    assert_abs_diff_eq!(statistics[0], 8.0 / 3.0, epsilon = 1e-5);
}

#[test]
fn test_imputer_fit_median_strategy() {
    let data =
        Tensor::from_slice(&[1.0f32, f32::NAN, 3.0, 4.0, 5.0], &[5, 1], &Device::cpu()).unwrap();
    let mut imputer = Imputer::with_strategy("median");

    imputer.fit(&data).unwrap();

    // Median of [1, 3, 4, 5] = 3.5
    let statistics = imputer.statistics();
    assert_abs_diff_eq!(statistics[0], 3.5, epsilon = 1e-5);
}

#[test]
fn test_imputer_fit_most_frequent_strategy() {
    let data =
        Tensor::from_slice(&[1.0f32, 2.0, 1.0, f32::NAN, 1.0], &[5, 1], &Device::cpu()).unwrap();
    let mut imputer = Imputer::with_strategy("most_frequent");

    imputer.fit(&data).unwrap();

    // Most frequent in [1, 2, 1, 1] is 1
    let statistics = imputer.statistics();
    assert_abs_diff_eq!(statistics[0], 1.0, epsilon = 1e-5);
}

#[test]
fn test_imputer_fit_constant_strategy() {
    let data = Tensor::from_slice(&[1.0f32, f32::NAN, 3.0], &[3, 1], &Device::cpu()).unwrap();
    let mut imputer = Imputer::with_constant(99.0);

    imputer.fit(&data).unwrap();

    // Constant strategy should store the fill value
    let statistics = imputer.statistics();
    assert_abs_diff_eq!(statistics[0], 99.0, epsilon = 1e-5);
}

#[test]
fn test_imputer_fit_multi_feature() {
    let data = Tensor::from_slice(
        &[1.0f32, 10.0, f32::NAN, 20.0, 3.0, f32::NAN],
        &[3, 2],
        &Device::cpu(),
    )
    .unwrap();
    let mut imputer = Imputer::new();

    imputer.fit(&data).unwrap();

    let statistics = imputer.statistics();
    assert_eq!(statistics.len(), 2);

    // Feature 0: mean of [1, 3] = 2
    assert_abs_diff_eq!(statistics[0], 2.0, epsilon = 1e-5);
    // Feature 1: mean of [10, 20] = 15
    assert_abs_diff_eq!(statistics[1], 15.0, epsilon = 1e-5);
}

#[test]
fn test_imputer_fit_all_nan_column() {
    let data =
        Tensor::from_slice(&[f32::NAN, f32::NAN, f32::NAN], &[3, 1], &Device::cpu()).unwrap();
    let mut imputer = Imputer::new();

    let result = imputer.fit(&data);
    // Should handle all-NaN columns gracefully (might use 0 or error)
    match result {
        Ok(_) => {
            let statistics = imputer.statistics();
            // Could be 0 or NaN depending on implementation
            assert!(statistics[0].is_nan() || statistics[0] == 0.0);
        }
        Err(e) => assert!(matches!(
            e,
            PreprocessingError::InvalidInput { .. } | PreprocessingError::EmptyDataset
        )),
    }
}

#[test]
fn test_imputer_transform_mean() {
    let data =
        Tensor::from_slice(&[1.0f32, f32::NAN, 3.0, f32::NAN], &[4, 1], &Device::cpu()).unwrap();
    let mut imputer = Imputer::new();

    imputer.fit(&data).unwrap();
    let imputed = imputer.transform(&data).unwrap();

    let values: Vec<f32> = imputed.to_cpu().unwrap().iter().copied().collect();
    // Mean of [1, 3] = 2, so NaNs should be replaced with 2
    let expected = vec![1.0f32, 2.0, 3.0, 2.0];

    for (actual, expected) in values.iter().zip(expected.iter()) {
        assert_abs_diff_eq!(actual, expected, epsilon = 1e-5);
    }
}

#[test]
fn test_imputer_transform_constant() {
    let data =
        Tensor::from_slice(&[1.0f32, f32::NAN, 3.0, f32::NAN], &[4, 1], &Device::cpu()).unwrap();
    let mut imputer = Imputer::with_constant(-999.0);

    imputer.fit(&data).unwrap();
    let imputed = imputer.transform(&data).unwrap();

    let values: Vec<f32> = imputed.to_cpu().unwrap().iter().copied().collect();
    // NaNs should be replaced with -999
    let expected = vec![1.0f32, -999.0, 3.0, -999.0];

    for (actual, expected) in values.iter().zip(expected.iter()) {
        assert_abs_diff_eq!(actual, expected, epsilon = 1e-5);
    }
}

#[test]
fn test_imputer_transform_no_missing() {
    let data = Tensor::from_slice(&[1.0f32, 2.0, 3.0, 4.0], &[4, 1], &Device::cpu()).unwrap();
    let mut imputer = Imputer::new();

    imputer.fit(&data).unwrap();
    let imputed = imputer.transform(&data).unwrap();

    let original_values: Vec<f32> = data.to_cpu().unwrap().iter().copied().collect();
    let imputed_values: Vec<f32> = imputed.to_cpu().unwrap().iter().copied().collect();

    // Should be unchanged when no missing values
    for (orig, imp) in original_values.iter().zip(imputed_values.iter()) {
        assert_abs_diff_eq!(orig, imp, epsilon = 1e-5);
    }
}

#[test]
fn test_imputer_transform_not_fitted() {
    let data = Tensor::from_slice(&[1.0f32, f32::NAN, 3.0], &[3, 1], &Device::cpu()).unwrap();
    let imputer = Imputer::new();

    let result = imputer.transform(&data);
    assert!(result.is_err());
    assert!(matches!(result.unwrap_err(), PreprocessingError::NotFitted));
}

#[test]
fn test_imputer_transform_dimension_mismatch() {
    let fit_data = Tensor::from_slice(&[1.0f32, 2.0, 3.0], &[3, 1], &Device::cpu()).unwrap();
    let transform_data =
        Tensor::from_slice(&[1.0f32, 2.0, 3.0, 4.0], &[2, 2], &Device::cpu()).unwrap();
    let mut imputer = Imputer::new();

    imputer.fit(&fit_data).unwrap();

    let result = imputer.transform(&transform_data);
    assert!(result.is_err());
    assert!(matches!(
        result.unwrap_err(),
        PreprocessingError::DimensionMismatch { .. }
    ));
}

#[test]
fn test_imputer_fit_transform() {
    let data = Tensor::from_slice(&[1.0f32, f32::NAN, 3.0], &[3, 1], &Device::cpu()).unwrap();
    let mut imputer = Imputer::new();

    let result1 = imputer.fit_transform(&data).unwrap();

    let mut imputer2 = Imputer::new();
    imputer2.fit(&data).unwrap();
    let result2 = imputer2.transform(&data).unwrap();

    assert_eq!(result1.shape(), result2.shape());
    let values1: Vec<f32> = result1.to_cpu().unwrap().iter().copied().collect();
    let values2: Vec<f32> = result2.to_cpu().unwrap().iter().copied().collect();

    for (v1, v2) in values1.iter().zip(values2.iter()) {
        assert_abs_diff_eq!(v1, v2, epsilon = 1e-5);
    }
}

#[test]
fn test_imputer_reset() {
    let data = Tensor::from_slice(&[1.0f32, f32::NAN, 3.0], &[3, 1], &Device::cpu()).unwrap();
    let mut imputer = Imputer::new();

    imputer.fit(&data).unwrap();
    assert!(imputer.is_fitted());

    imputer.reset();
    assert!(!imputer.is_fitted());
}

#[test]
fn test_imputer_empty_data() {
    let empty_data = Tensor::zeros(&[0, 1], &Device::cpu()).unwrap();
    let mut imputer = Imputer::new();

    let result = imputer.fit(&empty_data);
    assert!(result.is_err());
    assert!(matches!(
        result.unwrap_err(),
        PreprocessingError::EmptyDataset
    ));
}