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

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

#[test]
fn test_feature_selector_creation() {
    let selector = FeatureSelector::variance_threshold(0.1);
    assert!(!selector.is_fitted());
    assert_eq!(selector.method(), "variance_threshold");

    let selector = FeatureSelector::k_best(5);
    assert_eq!(selector.method(), "k_best");
    assert_eq!(selector.k(), 5);

    let selector = FeatureSelector::percentile(90.0);
    assert_eq!(selector.method(), "percentile");
    assert_abs_diff_eq!(selector.get_percentile(), 90.0, epsilon = 1e-5);
}

#[test]
fn test_feature_selector_invalid_params() {
    let result = std::panic::catch_unwind(|| FeatureSelector::variance_threshold(-0.1));
    assert!(result.is_err());

    let result = std::panic::catch_unwind(|| FeatureSelector::k_best(0));
    assert!(result.is_err());

    let result = std::panic::catch_unwind(|| FeatureSelector::percentile(101.0));
    assert!(result.is_err());
}

#[test]
fn test_feature_selector_fit_variance_threshold() {
    // Features with different variances
    let data = Tensor::from_slice(
        &[1.0, 1.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0],
        &[3, 3],
        &Device::cpu(),
    )
    .unwrap();
    let mut selector = FeatureSelector::variance_threshold(0.5);

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

    // First feature: [1, 1, 1] -> variance = 0 (below threshold)
    // Second feature: [2, 3, 4] -> variance > 0.5 (above threshold)
    // Third feature: [5, 6, 7] -> variance > 0.5 (above threshold)
    let selected_features = selector.selected_features();
    assert_eq!(selected_features.len(), 2); // Should select features 1 and 2
    assert!(!selected_features.contains(&0)); // First feature should be removed
    assert!(selected_features.contains(&1));
    assert!(selected_features.contains(&2));
}

#[test]
fn test_feature_selector_fit_k_best() {
    let features = Tensor::from_slice(
        &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0],
        &[3, 3],
        &Device::cpu(),
    )
    .unwrap();
    let targets = Tensor::from_slice(&[1.0, 2.0, 3.0], &[3, 1], &Device::cpu()).unwrap();
    let mut selector = FeatureSelector::k_best(2);

    let result = selector.fit_with_target(&features, &targets);
    assert!(result.is_ok());
    assert!(selector.is_fitted());

    let selected_features = selector.selected_features();
    assert_eq!(selected_features.len(), 2); // Should select top 2 features
}

#[test]
fn test_feature_selector_fit_percentile() {
    let features = Tensor::from_slice(
        &[
            1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0,
        ],
        &[3, 4],
        &Device::cpu(),
    )
    .unwrap();
    let targets = Tensor::from_slice(&[1.0, 2.0, 3.0], &[3, 1], &Device::cpu()).unwrap();
    let mut selector = FeatureSelector::percentile(50.0); // Select top 50%

    selector.fit_with_target(&features, &targets).unwrap();

    let selected_features = selector.selected_features();
    assert_eq!(selected_features.len(), 2); // 50% of 4 features = 2
}

#[test]
fn test_feature_selector_transform_variance_threshold() {
    let data = Tensor::from_slice(
        &[1.0, 1.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0],
        &[3, 3],
        &Device::cpu(),
    )
    .unwrap();
    let mut selector = FeatureSelector::variance_threshold(0.5);

    selector.fit(&data).unwrap();
    let selected = selector.transform(&data).unwrap();

    // Should have removed the first feature (constant)
    assert_eq!(selected.shape(), &[3, 2]);

    let values = selected
        .to_cpu()
        .unwrap()
        .iter()
        .copied()
        .collect::<Vec<f32>>();
    // Should contain features 1 and 2: [[2, 5], [3, 6], [4, 7]]
    let expected = vec![2.0, 5.0, 3.0, 6.0, 4.0, 7.0];
    for (actual, expected) in values.iter().zip(expected.iter()) {
        assert_abs_diff_eq!(actual, expected, epsilon = 1e-5);
    }
}

#[test]
fn test_feature_selector_transform_k_best() {
    let features =
        Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[2, 3], &Device::cpu()).unwrap();
    let targets = Tensor::from_slice(&[10.0, 20.0], &[2, 1], &Device::cpu()).unwrap();
    let mut selector = FeatureSelector::k_best(2);

    selector.fit_with_target(&features, &targets).unwrap();
    let selected = selector.transform(&features).unwrap();

    assert_eq!(selected.shape(), &[2, 2]); // 2 samples, 2 selected features
    assert!(selected.shape()[1] <= features.shape()[1]); // Fewer or equal features
}

#[test]
fn test_feature_selector_transform_all_low_variance() {
    // All features have low variance
    let data =
        Tensor::from_slice(&[1.0, 2.0, 1.0, 2.0, 1.0, 2.0], &[3, 2], &Device::cpu()).unwrap();
    let mut selector = FeatureSelector::variance_threshold(1.0);

    selector.fit(&data).unwrap();
    let selected = selector.transform(&data).unwrap();

    // Might select 0 features or keep at least 1
    assert!(selected.shape()[1] <= data.shape()[1]);
}

#[test]
fn test_feature_selector_transform_not_fitted() {
    let data = Tensor::from_slice(&[1.0, 2.0, 3.0], &[1, 3], &Device::cpu()).unwrap();
    let selector = FeatureSelector::variance_threshold(0.1);

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

#[test]
fn test_feature_selector_transform_dimension_mismatch() {
    let fit_data = Tensor::from_slice(&[1.0, 2.0, 3.0], &[1, 3], &Device::cpu()).unwrap();
    let transform_data = Tensor::from_slice(&[1.0, 2.0], &[1, 2], &Device::cpu()).unwrap();
    let mut selector = FeatureSelector::variance_threshold(0.1);

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

    let result = selector.transform(&transform_data);
    assert!(result.is_err());
    assert!(matches!(
        result.unwrap_err(),
        PreprocessingError::DimensionMismatch {
            expected: 3,
            actual: 2
        }
    ));
}

#[test]
fn test_feature_selector_inverse_transform() {
    let data = Tensor::from_slice(
        &[1.0, 1.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0],
        &[3, 3],
        &Device::cpu(),
    )
    .unwrap();
    let mut selector = FeatureSelector::variance_threshold(0.5);

    selector.fit(&data).unwrap();
    let selected = selector.transform(&data).unwrap();
    let reconstructed = selector.inverse_transform(&selected).unwrap();

    // Should have same number of features as original
    assert_eq!(reconstructed.shape()[1], data.shape()[1]);

    // Removed features should be filled with zeros or some default value
    let recon_values = reconstructed
        .to_cpu()
        .unwrap()
        .iter()
        .copied()
        .collect::<Vec<f32>>();

    // First column should be zeros (was removed)
    assert_abs_diff_eq!(recon_values[0], 0.0, epsilon = 1e-5);
    assert_abs_diff_eq!(recon_values[3], 0.0, epsilon = 1e-5);
    assert_abs_diff_eq!(recon_values[6], 0.0, epsilon = 1e-5);

    // Other columns should match selected features
    assert_abs_diff_eq!(recon_values[1], 2.0, epsilon = 1e-5);
    assert_abs_diff_eq!(recon_values[2], 5.0, epsilon = 1e-5);
}

#[test]
fn test_feature_selector_get_support() {
    let data = Tensor::from_slice(
        &[1.0, 1.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0],
        &[3, 3],
        &Device::cpu(),
    )
    .unwrap();
    let mut selector = FeatureSelector::variance_threshold(0.5);

    selector.fit(&data).unwrap();
    let support = selector.get_support();

    assert_eq!(support.len(), 3); // Same as number of input features
    assert!(!support[0]); // First feature should not be supported
    assert!(support[1]); // Second feature should be supported
    assert!(support[2]); // Third feature should be supported
}

#[test]
fn test_feature_selector_reset() {
    let data = Tensor::from_slice(&[1.0, 2.0, 3.0], &[1, 3], &Device::cpu()).unwrap();
    let mut selector = FeatureSelector::variance_threshold(0.1);

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

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

#[test]
fn test_feature_selector_empty_data() {
    let empty_data = Tensor::zeros(&[0, 3], &Device::cpu()).unwrap();
    let mut selector = FeatureSelector::variance_threshold(0.1);

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