foxhunt/ml/tests/huber_loss_test.rs

//! Huber Loss Tests
//!
//! Test suite for Huber loss implementation in DQN.
//! Huber loss = MSE for small errors, L1 for large errors (robust to outliers).
//!
//! Formula:
//! ```
//! L(x) = {
//!     0.5 * x²                    if |x| <= delta
//!     delta * (|x| - 0.5 * delta) otherwise
//! }
//! ```

use anyhow::Result;
use candle_core::{Device, Tensor};

/// Huber loss: quadratic for small errors, linear for large errors
/// More robust to outliers than MSE
fn huber_loss(predictions: &Tensor, targets: &Tensor, delta: f32) -> Result<Tensor> {
    let errors = (predictions - targets)?;
    let abs_errors = errors.abs()?;

    // Create mask for small errors (|error| <= delta)
    let small_errors_mask = abs_errors.le(delta)?;

    // Quadratic loss for small errors: 0.5 * error²
    let quadratic_loss = (errors.sqr()? * 0.5)?;

    // Linear loss for large errors: delta * (|error| - 0.5 * delta)
    // Use affine() to avoid scalar multiplication issues: affine(x, a, b) = a*x + b
    let abs_errors_scaled = abs_errors.affine(delta as f64, -(0.5 * delta * delta) as f64)?;

    // Select based on mask
    let loss = small_errors_mask.where_cond(&quadratic_loss, &abs_errors_scaled)?;

    Ok(loss.mean_all()?)
}

#[test]
fn test_huber_small_error_quadratic() -> Result<()> {
    // Test 1: Small error (0.5, delta=1.0) → quadratic behavior
    // Expected: 0.5 * 0.5² = 0.5 * 0.25 = 0.125
    let device = Device::cuda_if_available(0)?;
    let predictions = Tensor::new(&[1.5f32], &device)?;
    let targets = Tensor::new(&[1.0f32], &device)?;
    let delta = 1.0;

    let loss = huber_loss(&predictions, &targets, delta)?;
    let loss_value = loss.to_scalar::<f32>()?;

    let expected = 0.125f32;
    assert!(
        (loss_value - expected).abs() < 1e-5,
        "Small error loss incorrect: got {}, expected {}",
        loss_value,
        expected
    );

    println!(
        "✅ Test 1 passed: Small error (0.5) → quadratic behavior ({:.6})",
        loss_value
    );
    Ok(())
}

#[test]
fn test_huber_large_error_linear() -> Result<()> {
    // Test 2: Large error (5.0, delta=1.0) → linear behavior
    // Expected: 1.0 * (5.0 - 0.5 * 1.0) = 1.0 * 4.5 = 4.5
    let device = Device::cuda_if_available(0)?;
    let predictions = Tensor::new(&[6.0f32], &device)?;
    let targets = Tensor::new(&[1.0f32], &device)?;
    let delta = 1.0;

    let loss = huber_loss(&predictions, &targets, delta)?;
    let loss_value = loss.to_scalar::<f32>()?;

    let expected = 4.5f32;
    assert!(
        (loss_value - expected).abs() < 1e-5,
        "Large error loss incorrect: got {}, expected {}",
        loss_value,
        expected
    );

    println!(
        "✅ Test 2 passed: Large error (5.0) → linear behavior ({:.6})",
        loss_value
    );
    Ok(())
}

#[test]
fn test_huber_threshold_smooth_transition() -> Result<()> {
    // Test 3: Error at threshold (1.0, delta=1.0) → smooth transition
    // Quadratic: 0.5 * 1.0² = 0.5
    // Linear: 1.0 * (1.0 - 0.5 * 1.0) = 1.0 * 0.5 = 0.5
    // Both formulas should give same result at threshold
    let device = Device::cuda_if_available(0)?;
    let predictions = Tensor::new(&[2.0f32], &device)?;
    let targets = Tensor::new(&[1.0f32], &device)?;
    let delta = 1.0;

    let loss = huber_loss(&predictions, &targets, delta)?;
    let loss_value = loss.to_scalar::<f32>()?;

    let expected = 0.5f32;
    assert!(
        (loss_value - expected).abs() < 1e-5,
        "Threshold error loss incorrect: got {}, expected {}",
        loss_value,
        expected
    );

    println!(
        "✅ Test 3 passed: Error at threshold (1.0) → smooth transition ({:.6})",
        loss_value
    );
    Ok(())
}

#[test]
fn test_huber_negative_errors() -> Result<()> {
    // Test 4: Negative errors handled correctly (symmetry)
    // Error of -5.0 should give same loss as +5.0
    let device = Device::cuda_if_available(0)?;

    // Positive error
    let pred_pos = Tensor::new(&[6.0f32], &device)?;
    let target_pos = Tensor::new(&[1.0f32], &device)?;
    let loss_pos = huber_loss(&pred_pos, &target_pos, 1.0)?;
    let loss_pos_value = loss_pos.to_scalar::<f32>()?;

    // Negative error
    let pred_neg = Tensor::new(&[-4.0f32], &device)?;
    let target_neg = Tensor::new(&[1.0f32], &device)?;
    let loss_neg = huber_loss(&pred_neg, &target_neg, 1.0)?;
    let loss_neg_value = loss_neg.to_scalar::<f32>()?;

    assert!(
        (loss_pos_value - loss_neg_value).abs() < 1e-5,
        "Negative error loss asymmetric: pos={}, neg={}",
        loss_pos_value,
        loss_neg_value
    );

    println!(
        "✅ Test 4 passed: Negative errors handled correctly (pos={:.6}, neg={:.6})",
        loss_pos_value, loss_neg_value
    );
    Ok(())
}

#[test]
fn test_huber_batch_mixed_errors() -> Result<()> {
    // Test 5: Batch of errors (mixed small/large)
    // Errors: [0.5, 2.0, 5.0, 0.1] with delta=1.0
    // Expected losses:
    //   0.5: 0.5 * 0.5² = 0.125 (quadratic)
    //   2.0: 1.0 * (2.0 - 0.5) = 1.5 (linear)
    //   5.0: 1.0 * (5.0 - 0.5) = 4.5 (linear)
    //   0.1: 0.5 * 0.1² = 0.005 (quadratic)
    // Average: (0.125 + 1.5 + 4.5 + 0.005) / 4 = 6.13 / 4 = 1.5325
    let device = Device::cuda_if_available(0)?;
    let predictions = Tensor::new(&[1.5f32, 3.0, 6.0, 1.1], &device)?;
    let targets = Tensor::new(&[1.0f32, 1.0, 1.0, 1.0], &device)?;
    let delta = 1.0;

    let loss = huber_loss(&predictions, &targets, delta)?;
    let loss_value = loss.to_scalar::<f32>()?;

    let expected = 1.5325f32;
    assert!(
        (loss_value - expected).abs() < 1e-3,
        "Batch loss incorrect: got {}, expected {}",
        loss_value,
        expected
    );

    println!(
        "✅ Test 5 passed: Batch of mixed errors → average loss ({:.6})",
        loss_value
    );
    Ok(())
}

#[test]
fn test_huber_gradient_bounded() -> Result<()> {
    // Test 6: Gradient is bounded for large errors
    // For MSE: gradient = 2 * error (unbounded, grows linearly)
    // For Huber: gradient = delta for |error| > delta (bounded)
    //
    // Error of 100.0 with delta=1.0:
    // MSE gradient would be 200 (2 * 100)
    // Huber gradient is bounded by delta=1.0
    //
    // We verify this by checking that large errors don't cause
    // disproportionately large losses (which would indicate large gradients)
    let device = Device::cuda_if_available(0)?;

    // Small outlier: error=10
    let pred_small = Tensor::new(&[11.0f32], &device)?;
    let target_small = Tensor::new(&[1.0f32], &device)?;
    let loss_small = huber_loss(&pred_small, &target_small, 1.0)?;
    let loss_small_value = loss_small.to_scalar::<f32>()?;

    // Large outlier: error=100
    let pred_large = Tensor::new(&[101.0f32], &device)?;
    let target_large = Tensor::new(&[1.0f32], &device)?;
    let loss_large = huber_loss(&pred_large, &target_large, 1.0)?;
    let loss_large_value = loss_large.to_scalar::<f32>()?;

    // Huber loss should grow linearly with error size (not quadratically)
    // loss_large / loss_small should be approximately 100/10 = 10 (linear growth)
    // For MSE it would be (100²)/(10²) = 100 (quadratic growth)
    let ratio = loss_large_value / loss_small_value;

    assert!(
        ratio > 8.0 && ratio < 12.0,
        "Gradient not bounded: loss ratio {} (expected ~10 for linear growth, ~100 for quadratic)",
        ratio
    );

    println!(
        "✅ Test 6 passed: Gradient bounded for large errors (ratio={:.2}, linear growth confirmed)",
        ratio
    );
    Ok(())
}

#[test]
fn test_huber_vs_mse_convergence() -> Result<()> {
    // Test 7: Huber converges faster than MSE on outlier data
    // Simulate training data with outliers: [1.0, 1.1, 0.9, 10.0, 1.05]
    // The outlier (10.0) should have less influence on Huber loss
    let device = Device::cuda_if_available(0)?;
    let predictions = Tensor::new(&[2.0f32, 2.1, 1.9, 11.0, 2.05], &device)?;
    let targets = Tensor::new(&[1.0f32, 1.1, 0.9, 10.0, 1.05], &device)?;
    let delta = 1.0;

    // Huber loss
    let huber = huber_loss(&predictions, &targets, delta)?;
    let huber_value = huber.to_scalar::<f32>()?;

    // MSE loss
    let errors = (predictions - &targets)?;
    let mse = errors.sqr()?.mean_all()?;
    let mse_value = mse.to_scalar::<f32>()?;

    // Huber should be significantly smaller than MSE due to outlier robustness
    // The outlier (error=1.0) contributes:
    //   MSE: 1.0² = 1.0
    //   Huber: 1.0 * (1.0 - 0.5) = 0.5
    // So Huber should be approximately half of MSE
    assert!(
        huber_value < mse_value,
        "Huber loss should be smaller than MSE for outlier data: huber={}, mse={}",
        huber_value,
        mse_value
    );

    let reduction = (mse_value - huber_value) / mse_value * 100.0;
    println!(
        "✅ Test 7 passed: Huber converges better on outliers (MSE={:.6}, Huber={:.6}, {:.1}% reduction)",
        mse_value,
        huber_value,
        reduction
    );
    Ok(())
}

#[test]
fn test_huber_different_deltas() -> Result<()> {
    // Additional test: Different delta values affect transition point
    let device = Device::cuda_if_available(0)?;
    let predictions = Tensor::new(&[3.0f32], &device)?;
    let targets = Tensor::new(&[1.0f32], &device)?; // Error = 2.0

    // With delta=1.0: error=2.0 is large (linear)
    let loss_delta1 = huber_loss(&predictions, &targets, 1.0)?;
    let loss1 = loss_delta1.to_scalar::<f32>()?;

    // With delta=3.0: error=2.0 is small (quadratic)
    let loss_delta3 = huber_loss(&predictions, &targets, 3.0)?;
    let loss3 = loss_delta3.to_scalar::<f32>()?;

    // delta=1.0: 1.0 * (2.0 - 0.5 * 1.0) = 1.5 (linear)
    // delta=3.0: 0.5 * 2.0² = 2.0 (quadratic)
    assert!(
        (loss1 - 1.5).abs() < 1e-5,
        "Delta=1.0 loss incorrect: got {}, expected 1.5",
        loss1
    );
    assert!(
        (loss3 - 2.0).abs() < 1e-5,
        "Delta=3.0 loss incorrect: got {}, expected 2.0",
        loss3
    );

    println!(
        "✅ Additional test passed: Different deltas work correctly (delta=1.0: {:.6}, delta=3.0: {:.6})",
        loss1,
        loss3
    );
    Ok(())
}