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

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

#[test]
fn test_normalizer_creation() {
    let normalizer = Normalizer::new();
    assert!(!normalizer.is_fitted());
    assert_eq!(normalizer.norm(), "l2");

    let normalizer = Normalizer::with_norm("l1");
    assert_eq!(normalizer.norm(), "l1");

    let normalizer = Normalizer::with_norm("max");
    assert_eq!(normalizer.norm(), "max");
}

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

#[test]
fn test_normalizer_fit_simple() {
    let data = Tensor::from_slice(&[3.0, 4.0], &[1, 2], &Device::cpu()).unwrap();
    let mut normalizer = Normalizer::new();

    // Normalizer doesn't need fitting, but should succeed
    let result = normalizer.fit(&data);
    assert!(result.is_ok());
    assert!(normalizer.is_fitted());
}

#[test]
fn test_normalizer_transform_l2() {
    let data = Tensor::from_slice(&[3.0, 4.0], &[1, 2], &Device::cpu()).unwrap();
    let mut normalizer = Normalizer::new(); // L2 by default

    normalizer.fit(&data).unwrap();
    let normalized = normalizer.transform(&data).unwrap();

    let values = normalized
        .to_cpu()
        .unwrap()
        .iter()
        .copied()
        .collect::<Vec<f32>>();
    // L2 norm of [3, 4] is 5, so normalized should be [0.6, 0.8]
    assert_abs_diff_eq!(values[0], 0.6, epsilon = 1e-5);
    assert_abs_diff_eq!(values[1], 0.8, epsilon = 1e-5);
}

#[test]
fn test_normalizer_transform_l1() {
    let data = Tensor::from_slice(&[3.0, 4.0], &[1, 2], &Device::cpu()).unwrap();
    let mut normalizer = Normalizer::with_norm("l1");

    normalizer.fit(&data).unwrap();
    let normalized = normalizer.transform(&data).unwrap();

    let values = normalized
        .to_cpu()
        .unwrap()
        .iter()
        .copied()
        .collect::<Vec<f32>>();
    // L1 norm of [3, 4] is 7, so normalized should be [3/7, 4/7]
    assert_abs_diff_eq!(values[0], 3.0 / 7.0, epsilon = 1e-5);
    assert_abs_diff_eq!(values[1], 4.0 / 7.0, epsilon = 1e-5);
}

#[test]
fn test_normalizer_transform_max() {
    let data = Tensor::from_slice(&[3.0, 4.0], &[1, 2], &Device::cpu()).unwrap();
    let mut normalizer = Normalizer::with_norm("max");

    normalizer.fit(&data).unwrap();
    let normalized = normalizer.transform(&data).unwrap();

    let values = normalized
        .to_cpu()
        .unwrap()
        .iter()
        .copied()
        .collect::<Vec<f32>>();
    // Max norm of [3, 4] is 4, so normalized should be [0.75, 1.0]
    assert_abs_diff_eq!(values[0], 0.75, epsilon = 1e-5);
    assert_abs_diff_eq!(values[1], 1.0, epsilon = 1e-5);
}

#[test]
fn test_normalizer_multiple_samples() {
    let data = Tensor::from_slice(&[3.0, 4.0, 6.0, 8.0], &[2, 2], &Device::cpu()).unwrap();
    let mut normalizer = Normalizer::new();

    normalizer.fit(&data).unwrap();
    let normalized = normalizer.transform(&data).unwrap();

    let values = normalized
        .to_cpu()
        .unwrap()
        .iter()
        .copied()
        .collect::<Vec<f32>>();

    // First sample [3, 4] -> L2 norm = 5 -> [0.6, 0.8]
    assert_abs_diff_eq!(values[0], 0.6, epsilon = 1e-5);
    assert_abs_diff_eq!(values[1], 0.8, epsilon = 1e-5);

    // Second sample [6, 8] -> L2 norm = 10 -> [0.6, 0.8]
    assert_abs_diff_eq!(values[2], 0.6, epsilon = 1e-5);
    assert_abs_diff_eq!(values[3], 0.8, epsilon = 1e-5);
}

#[test]
fn test_normalizer_zero_norm() {
    let data = Tensor::from_slice(&[0.0, 0.0], &[1, 2], &Device::cpu()).unwrap();
    let mut normalizer = Normalizer::new();

    normalizer.fit(&data).unwrap();
    let result = normalizer.transform(&data);

    // Should handle zero norm gracefully
    match result {
        Ok(normalized) => {
            let values = normalized
                .to_cpu()
                .unwrap()
                .iter()
                .copied()
                .collect::<Vec<f32>>();
            // Could be [0, 0] or could error, depending on implementation
            assert!(values[0] == 0.0 && values[1] == 0.0);
        }
        Err(e) => {
            assert!(matches!(e, PreprocessingError::NumericalError { .. }));
        }
    }
}

#[test]
fn test_normalizer_reset() {
    let data = Tensor::from_slice(&[3.0, 4.0], &[1, 2], &Device::cpu()).unwrap();
    let mut normalizer = Normalizer::new();

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

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