140 lines
4.7 KiB
Rust
140 lines
4.7 KiB
Rust
#![allow(clippy::approx_constant)]
|
|
use approx::assert_abs_diff_eq;
|
|
use rtx_preprocessing::{OneHotEncoder, PreprocessingError, Transformer};
|
|
use rtx_tensor::{Device, Tensor};
|
|
|
|
#[test]
|
|
fn test_onehot_encoder_creation() {
|
|
let encoder = OneHotEncoder::new();
|
|
assert!(!encoder.is_fitted());
|
|
assert!(encoder.sparse());
|
|
assert!(encoder.drop_first());
|
|
|
|
let encoder = OneHotEncoder::with_params(false, false, None);
|
|
assert!(!encoder.sparse());
|
|
assert!(!encoder.drop_first());
|
|
}
|
|
|
|
#[test]
|
|
fn test_onehot_encoder_fit_simple() {
|
|
let data = Tensor::from_slice(&[0.0, 1.0, 2.0, 1.0, 0.0], &[5, 1], &Device::cpu()).unwrap();
|
|
let mut encoder = OneHotEncoder::new();
|
|
|
|
let result = encoder.fit(&data);
|
|
assert!(result.is_ok());
|
|
assert!(encoder.is_fitted());
|
|
assert_eq!(encoder.categories()[0].len(), 3); // [0, 1, 2]
|
|
}
|
|
|
|
#[test]
|
|
fn test_onehot_encoder_transform_simple() {
|
|
let data = Tensor::from_slice(&[0.0, 1.0, 2.0], &[3, 1], &Device::cpu()).unwrap();
|
|
let mut encoder = OneHotEncoder::with_params(false, false, None); // Dense, no drop
|
|
|
|
encoder.fit(&data).unwrap();
|
|
let encoded = encoder.transform(&data).unwrap();
|
|
|
|
assert_eq!(encoded.shape(), &[3, 3]); // 3 samples, 3 categories
|
|
let values: Vec<f32> = encoded.to_cpu().unwrap().iter().copied().collect();
|
|
|
|
// Should be identity matrix
|
|
let expected = vec![1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0];
|
|
for (actual, expected) in values.iter().zip(expected.iter()) {
|
|
assert_abs_diff_eq!(actual, expected, epsilon = 1e-5);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_onehot_encoder_transform_drop_first() {
|
|
let data = Tensor::from_slice(&[0.0, 1.0, 2.0], &[3, 1], &Device::cpu()).unwrap();
|
|
let mut encoder = OneHotEncoder::with_params(false, true, None); // Dense, drop first
|
|
|
|
encoder.fit(&data).unwrap();
|
|
let encoded = encoder.transform(&data).unwrap();
|
|
|
|
assert_eq!(encoded.shape(), &[3, 2]); // 3 samples, 2 categories (dropped first)
|
|
let values: Vec<f32> = encoded.to_cpu().unwrap().iter().copied().collect();
|
|
|
|
// Should drop first column
|
|
let expected = vec![0.0, 0.0, 1.0, 0.0, 0.0, 1.0];
|
|
for (actual, expected) in values.iter().zip(expected.iter()) {
|
|
assert_abs_diff_eq!(actual, expected, epsilon = 1e-5);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_onehot_encoder_unknown_category() {
|
|
let data = Tensor::from_slice(&[0.0, 1.0, 2.0], &[3, 1], &Device::cpu()).unwrap();
|
|
let mut encoder = OneHotEncoder::new();
|
|
|
|
encoder.fit(&data).unwrap();
|
|
|
|
let unknown_data = Tensor::from_slice(&[3.0], &[1, 1], &Device::cpu()).unwrap(); // Category 3 not seen
|
|
let result = encoder.transform(&unknown_data);
|
|
|
|
// Should handle unknown categories based on handle_unknown parameter
|
|
// Default might be 'error' or 'ignore'
|
|
match result {
|
|
Ok(_) => (), // If 'ignore' is default
|
|
Err(e) => assert!(matches!(e, PreprocessingError::InvalidInput { .. })),
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_onehot_encoder_multi_feature() {
|
|
let data =
|
|
Tensor::from_slice(&[0.0, 10.0, 1.0, 20.0, 0.0, 10.0], &[3, 2], &Device::cpu()).unwrap();
|
|
let mut encoder = OneHotEncoder::with_params(false, false, None);
|
|
|
|
encoder.fit(&data).unwrap();
|
|
let encoded = encoder.transform(&data).unwrap();
|
|
|
|
// Feature 0: [0, 1] -> 2 categories
|
|
// Feature 1: [10, 20] -> 2 categories
|
|
// Total: 4 columns
|
|
assert_eq!(encoded.shape(), &[3, 4]);
|
|
}
|
|
|
|
#[test]
|
|
fn test_onehot_encoder_inverse_transform() {
|
|
let data = Tensor::from_slice(&[0.0, 1.0, 2.0, 1.0], &[4, 1], &Device::cpu()).unwrap();
|
|
let mut encoder = OneHotEncoder::with_params(false, false, None);
|
|
|
|
encoder.fit(&data).unwrap();
|
|
let encoded = encoder.transform(&data).unwrap();
|
|
let decoded = encoder.inverse_transform(&encoded).unwrap();
|
|
|
|
let original_values: Vec<f32> = data.to_cpu().unwrap().iter().copied().collect();
|
|
let decoded_values: Vec<f32> = decoded.to_cpu().unwrap().iter().copied().collect();
|
|
|
|
for (orig, dec) in original_values.iter().zip(decoded_values.iter()) {
|
|
assert_abs_diff_eq!(orig, dec, epsilon = 1e-5);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_onehot_encoder_fit_transform() {
|
|
let data = Tensor::from_slice(&[0.0, 1.0, 2.0], &[3, 1], &Device::cpu()).unwrap();
|
|
let mut encoder = OneHotEncoder::new();
|
|
|
|
let result1 = encoder.fit_transform(&data).unwrap();
|
|
|
|
let mut encoder2 = OneHotEncoder::new();
|
|
encoder2.fit(&data).unwrap();
|
|
let result2 = encoder2.transform(&data).unwrap();
|
|
|
|
assert_eq!(result1.shape(), result2.shape());
|
|
}
|
|
|
|
#[test]
|
|
fn test_onehot_encoder_reset() {
|
|
let data = Tensor::from_slice(&[0.0, 1.0, 2.0], &[3, 1], &Device::cpu()).unwrap();
|
|
let mut encoder = OneHotEncoder::new();
|
|
|
|
encoder.fit(&data).unwrap();
|
|
assert!(encoder.is_fitted());
|
|
|
|
encoder.reset();
|
|
assert!(!encoder.is_fitted());
|
|
}
|