337 lines
11 KiB
Rust
337 lines
11 KiB
Rust
//! Decision Tree Tests - TDD Red Phase
|
|
//!
|
|
//! These tests define the expected behavior of the decision tree implementation.
|
|
//! They will fail initially and drive the implementation.
|
|
|
|
use approx::assert_abs_diff_eq;
|
|
use rtx_ml_classic::trees::{DecisionTree, SplitCriterion};
|
|
use rtx_ml_classic::{MLError, Result};
|
|
use rtx_tensor::{DType, Device, Tensor};
|
|
|
|
/// Test fixture for decision tree tests
|
|
struct DecisionTreeTestFixture {
|
|
device: Device,
|
|
x_binary: Tensor, // Simple binary classification data
|
|
y_binary: Tensor, // Binary labels
|
|
x_multiclass: Tensor, // Multiclass classification data
|
|
y_multiclass: Tensor, // Multiclass labels
|
|
x_regression: Tensor, // Regression data
|
|
y_regression: Tensor, // Regression targets
|
|
}
|
|
|
|
impl DecisionTreeTestFixture {
|
|
fn new() -> Result<Self> {
|
|
let device = Device::cpu();
|
|
|
|
// Create simple linearly separable binary classification data
|
|
// Features: [x1, x2] where x1 + x2 > 0.5 -> class 1, else class 0
|
|
let x_binary = Tensor::from_slice(
|
|
&[
|
|
0.1f32, 0.1, 0.9, 0.1, 0.1, 0.9, 0.9, 0.9, 0.2, 0.2, 0.8, 0.8, 0.3, 0.1, 0.1, 0.8,
|
|
],
|
|
&[8, 2],
|
|
&device,
|
|
)?;
|
|
let y_binary =
|
|
Tensor::from_slice(&[0.0f32, 1.0, 1.0, 1.0, 0.0, 1.0, 0.0, 1.0], &[8], &device)?;
|
|
|
|
// Create 3-class classification data
|
|
let x_multiclass = Tensor::from_slice(
|
|
&[
|
|
0.1f32, 0.1, 0.9, 0.9, 0.1, 0.9, 0.9, 0.1, 0.5, 0.1, 0.5, 0.9, 0.1, 0.5, 0.9, 0.5,
|
|
],
|
|
&[8, 2],
|
|
&device,
|
|
)?;
|
|
let y_multiclass =
|
|
Tensor::from_slice(&[0.0f32, 2.0, 1.0, 1.0, 0.0, 2.0, 1.0, 2.0], &[8], &device)?;
|
|
|
|
// Create simple regression data: y = x1 + 2*x2 + noise
|
|
let x_regression = Tensor::from_slice(
|
|
&[
|
|
1.0f32, 2.0, 2.0, 1.0, 3.0, 3.0, 1.0, 1.0, 2.0, 2.0, 3.0, 1.0, 1.0, 3.0, 4.0, 2.0,
|
|
],
|
|
&[8, 2],
|
|
&device,
|
|
)?;
|
|
let y_regression =
|
|
Tensor::from_slice(&[5.1f32, 4.1, 9.0, 3.0, 6.1, 5.0, 7.1, 8.0], &[8], &device)?;
|
|
|
|
Ok(Self {
|
|
device,
|
|
x_binary,
|
|
y_binary,
|
|
x_multiclass,
|
|
y_multiclass,
|
|
x_regression,
|
|
y_regression,
|
|
})
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_decision_tree_creation() {
|
|
let tree = DecisionTree::new();
|
|
|
|
// Test default parameters
|
|
assert_eq!(tree.get_max_depth(), None);
|
|
assert_eq!(tree.get_min_samples_split(), 2);
|
|
assert_eq!(tree.get_min_samples_leaf(), 1);
|
|
assert_eq!(tree.get_criterion(), &SplitCriterion::Gini);
|
|
assert!(tree.is_classifier());
|
|
assert!(!tree.is_fitted());
|
|
}
|
|
|
|
#[test]
|
|
fn test_decision_tree_builder_pattern() {
|
|
let tree = DecisionTree::new()
|
|
.max_depth(5)
|
|
.min_samples_split(4)
|
|
.min_samples_leaf(2)
|
|
.criterion("entropy")
|
|
.task("regression");
|
|
|
|
assert_eq!(tree.get_max_depth(), Some(5));
|
|
assert_eq!(tree.get_min_samples_split(), 4);
|
|
assert_eq!(tree.get_min_samples_leaf(), 2);
|
|
assert_eq!(tree.get_criterion(), &SplitCriterion::Entropy);
|
|
assert!(!tree.is_classifier()); // Should be regressor
|
|
}
|
|
|
|
#[test]
|
|
fn test_split_criterion_parsing() {
|
|
let tree_gini = DecisionTree::new().criterion("gini");
|
|
assert_eq!(tree_gini.get_criterion(), &SplitCriterion::Gini);
|
|
|
|
let tree_entropy = DecisionTree::new().criterion("ENTROPY");
|
|
assert_eq!(tree_entropy.get_criterion(), &SplitCriterion::Entropy);
|
|
|
|
let tree_log_loss = DecisionTree::new().criterion("log_loss");
|
|
assert_eq!(tree_log_loss.get_criterion(), &SplitCriterion::LogLoss);
|
|
|
|
// Invalid criterion should default to Gini
|
|
let tree_invalid = DecisionTree::new().criterion("invalid");
|
|
assert_eq!(tree_invalid.get_criterion(), &SplitCriterion::Gini);
|
|
}
|
|
|
|
#[test]
|
|
fn test_prediction_before_fitting_fails() {
|
|
let fixture = DecisionTreeTestFixture::new().unwrap();
|
|
let tree = DecisionTree::new();
|
|
|
|
let result = tree.predict(&fixture.x_binary);
|
|
assert!(matches!(result, Err(MLError::ModelNotFitted)));
|
|
|
|
let result = tree.predict_proba(&fixture.x_binary);
|
|
assert!(matches!(result, Err(MLError::ModelNotFitted)));
|
|
|
|
let result = tree.feature_importances();
|
|
assert!(matches!(result, Err(MLError::ModelNotFitted)));
|
|
}
|
|
|
|
#[test]
|
|
fn test_binary_classification_gini() {
|
|
let fixture = DecisionTreeTestFixture::new().unwrap();
|
|
let mut tree = DecisionTree::new().max_depth(3).criterion("gini");
|
|
|
|
// Fit should succeed
|
|
tree.fit(&fixture.x_binary, &fixture.y_binary).unwrap();
|
|
assert!(tree.is_fitted());
|
|
|
|
// Predictions should have correct shape
|
|
let predictions = tree.predict(&fixture.x_binary).unwrap();
|
|
assert_eq!(predictions.shape().dims(), &[8]); // 8 samples
|
|
|
|
// All predictions should be 0 or 1
|
|
let pred_data = predictions.data().unwrap();
|
|
for pred in pred_data {
|
|
let pred_int = pred.round() as i32;
|
|
assert!(pred_int == 0 || pred_int == 1);
|
|
}
|
|
|
|
// Probability predictions should sum to 1
|
|
let proba = tree.predict_proba(&fixture.x_binary).unwrap();
|
|
assert_eq!(proba.shape().dims(), &[8, 2]); // 8 samples, 2 classes
|
|
|
|
let proba_data = proba.data().unwrap();
|
|
for i in 0..8 {
|
|
let row_sum = proba_data[i * 2] + proba_data[i * 2 + 1];
|
|
assert_abs_diff_eq!(row_sum, 1.0, epsilon = 1e-6);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_binary_classification_entropy() {
|
|
let fixture = DecisionTreeTestFixture::new().unwrap();
|
|
let mut tree = DecisionTree::new().max_depth(3).criterion("entropy");
|
|
|
|
tree.fit(&fixture.x_binary, &fixture.y_binary).unwrap();
|
|
let predictions = tree.predict(&fixture.x_binary).unwrap();
|
|
|
|
// Should achieve perfect accuracy on this simple dataset
|
|
let pred_data = predictions.data().unwrap();
|
|
let y_data = fixture.y_binary.data().unwrap();
|
|
|
|
let mut correct = 0;
|
|
for (pred, true_val) in pred_data.iter().zip(y_data.iter()) {
|
|
if (pred.round() - true_val).abs() < 0.1 {
|
|
correct += 1;
|
|
}
|
|
}
|
|
|
|
let accuracy = correct as f32 / y_data.len() as f32;
|
|
assert!(
|
|
accuracy >= 0.75,
|
|
"Accuracy {} should be at least 75%",
|
|
accuracy
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_multiclass_classification() {
|
|
let fixture = DecisionTreeTestFixture::new().unwrap();
|
|
let mut tree = DecisionTree::new().max_depth(5).criterion("gini");
|
|
|
|
tree.fit(&fixture.x_multiclass, &fixture.y_multiclass)
|
|
.unwrap();
|
|
|
|
let predictions = tree.predict(&fixture.x_multiclass).unwrap();
|
|
let pred_data = predictions.data().unwrap();
|
|
|
|
// All predictions should be valid class indices (0, 1, or 2)
|
|
for pred in pred_data {
|
|
let pred_int = pred.round() as i32;
|
|
assert!(pred_int >= 0 && pred_int <= 2);
|
|
}
|
|
|
|
// Probability predictions for 3 classes
|
|
let proba = tree.predict_proba(&fixture.x_multiclass).unwrap();
|
|
assert_eq!(proba.shape().dims(), &[8, 3]);
|
|
|
|
let proba_data = proba.data().unwrap();
|
|
for i in 0..8 {
|
|
let row_sum = proba_data[i * 3] + proba_data[i * 3 + 1] + proba_data[i * 3 + 2];
|
|
assert_abs_diff_eq!(row_sum, 1.0, epsilon = 1e-6);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_regression() {
|
|
let fixture = DecisionTreeTestFixture::new().unwrap();
|
|
let mut tree = DecisionTree::new().task("regression").max_depth(4);
|
|
|
|
tree.fit(&fixture.x_regression, &fixture.y_regression)
|
|
.unwrap();
|
|
|
|
let predictions = tree.predict(&fixture.x_regression).unwrap();
|
|
assert_eq!(predictions.shape().dims(), &[8]);
|
|
|
|
// predict_proba should fail for regression
|
|
let result = tree.predict_proba(&fixture.x_regression);
|
|
assert!(matches!(result, Err(MLError::InvalidParameter { .. })));
|
|
|
|
// Calculate MSE - should be reasonable for this simple dataset
|
|
let pred_data = predictions.data().unwrap();
|
|
let y_data = fixture.y_regression.data().unwrap();
|
|
|
|
let mse: f32 = pred_data
|
|
.iter()
|
|
.zip(y_data.iter())
|
|
.map(|(pred, actual)| (pred - actual).powi(2))
|
|
.sum::<f32>()
|
|
/ y_data.len() as f32;
|
|
|
|
assert!(mse < 4.0, "MSE {} should be reasonable", mse);
|
|
}
|
|
|
|
#[test]
|
|
fn test_feature_importances() {
|
|
let fixture = DecisionTreeTestFixture::new().unwrap();
|
|
let mut tree = DecisionTree::new().max_depth(3);
|
|
|
|
tree.fit(&fixture.x_binary, &fixture.y_binary).unwrap();
|
|
|
|
let importances = tree.feature_importances().unwrap();
|
|
assert_eq!(importances.shape().dims(), &[2]); // 2 features
|
|
|
|
let imp_data = importances.data().unwrap();
|
|
|
|
// All importances should be non-negative
|
|
for imp in &imp_data {
|
|
assert!(*imp >= 0.0);
|
|
}
|
|
|
|
// Importances should sum to 1.0
|
|
let total: f32 = imp_data.iter().sum();
|
|
assert_abs_diff_eq!(total, 1.0, epsilon = 1e-6);
|
|
}
|
|
|
|
#[test]
|
|
fn test_min_samples_constraints() {
|
|
let fixture = DecisionTreeTestFixture::new().unwrap();
|
|
|
|
// Tree with high min_samples_split should create simpler trees
|
|
let mut tree_high_split = DecisionTree::new()
|
|
.min_samples_split(10) // More than our dataset size
|
|
.max_depth(10);
|
|
|
|
// Should fail because min_samples_split > n_samples
|
|
let result = tree_high_split.fit(&fixture.x_binary, &fixture.y_binary);
|
|
assert!(result.is_err());
|
|
|
|
// Let's test a valid case where min_samples_split causes simpler trees but doesn't fail
|
|
let mut tree_simple = DecisionTree::new()
|
|
.min_samples_split(4) // Reasonable value for 8 samples
|
|
.max_depth(10);
|
|
|
|
tree_simple
|
|
.fit(&fixture.x_binary, &fixture.y_binary)
|
|
.unwrap();
|
|
let predictions = tree_simple.predict(&fixture.x_binary).unwrap();
|
|
assert_eq!(predictions.shape().dims(), &[8]);
|
|
}
|
|
|
|
#[test]
|
|
fn test_max_depth_constraint() {
|
|
let fixture = DecisionTreeTestFixture::new().unwrap();
|
|
|
|
// Very shallow tree
|
|
let mut tree_shallow = DecisionTree::new().max_depth(1);
|
|
|
|
tree_shallow
|
|
.fit(&fixture.x_binary, &fixture.y_binary)
|
|
.unwrap();
|
|
let predictions_shallow = tree_shallow.predict(&fixture.x_binary).unwrap();
|
|
|
|
// Deeper tree
|
|
let mut tree_deep = DecisionTree::new().max_depth(5);
|
|
|
|
tree_deep.fit(&fixture.x_binary, &fixture.y_binary).unwrap();
|
|
let predictions_deep = tree_deep.predict(&fixture.x_binary).unwrap();
|
|
|
|
// Both should produce valid predictions
|
|
assert_eq!(predictions_shallow.shape().dims(), &[8]);
|
|
assert_eq!(predictions_deep.shape().dims(), &[8]);
|
|
}
|
|
|
|
#[test]
|
|
fn test_invalid_input_dimensions() {
|
|
let device = Device::cpu();
|
|
let mut tree = DecisionTree::new();
|
|
|
|
// Mismatched X and y dimensions
|
|
let x = Tensor::zeros([10, 3], &device).unwrap();
|
|
let y = Tensor::zeros([5], &device).unwrap(); // Wrong size
|
|
|
|
let result = tree.fit(&x, &y);
|
|
assert!(result.is_err());
|
|
|
|
// Empty data - use 1 sample of zeros instead since tensor doesn't allow zero dims
|
|
let x_single = Tensor::zeros([1, 3], &device).unwrap();
|
|
let y_single = Tensor::zeros([1], &device).unwrap();
|
|
|
|
let result = tree.fit(&x_single, &y_single);
|
|
// Should fail because insufficient data
|
|
assert!(result.is_err());
|
|
}
|