foxhunt/crates/ml/tests/diffusion_integration.rs

//! Diffusion Model (DDPM/DDIM) Integration Tests
//!
//! Validates the Diffusion trainable adapter end-to-end:
//! construction, forward pass, training pipeline, checkpoint save/load.
//!
//! NOTE: Diffusion models generate noise targets internally during forward(),
//! so we test pipeline integrity rather than loss monotonicity.

use candle_core::{Device, Tensor};
use ml::diffusion::config::DiffusionConfig;
use ml::diffusion::trainable::DiffusionTrainableAdapter;
use ml::training::unified_trainer::UnifiedTrainable;

fn small_diffusion_config() -> DiffusionConfig {
    DiffusionConfig {
        num_timesteps: 50,
        sampling_steps: 5,
        seq_len: 8,
        feature_dim: 1,
        hidden_dim: 16,
        num_layers: 1,
        time_embed_dim: 8,
        learning_rate: 1e-3,
        weight_decay: 1e-5,
        grad_clip: 1.0,
        ..Default::default()
    }
}

#[test]
fn test_diffusion_construction() {
    let config = small_diffusion_config();
    let adapter = DiffusionTrainableAdapter::new(config, Device::Cpu);
    assert!(
        adapter.is_ok(),
        "Diffusion construction failed: {:?}",
        adapter.err()
    );
    let adapter = adapter.unwrap();
    assert_eq!(adapter.model_type(), "Diffusion");
    assert_eq!(adapter.get_step(), 0);
}

#[test]
fn test_diffusion_forward_pass() {
    let config = small_diffusion_config();
    let data_dim = config.data_dim(); // seq_len * feature_dim = 8
    let mut adapter = DiffusionTrainableAdapter::new(config, Device::Cpu).unwrap();

    // [batch=4, data_dim=8]
    let input = Tensor::randn(0f32, 1.0, (4, data_dim), &Device::Cpu).unwrap();
    let output = adapter.forward(&input);
    assert!(output.is_ok(), "Forward failed: {:?}", output.err());

    let output = output.unwrap();
    println!("Diffusion output shape: {:?}", output.dims());
    assert_eq!(output.dims()[0], 4, "Batch dimension should be 4");
    // Output is predicted noise, should be finite
    let sum = output
        .abs()
        .unwrap()
        .sum_all()
        .unwrap()
        .to_scalar::<f32>()
        .unwrap();
    assert!(sum.is_finite(), "Output contains NaN/Inf");
}

#[test]
fn test_diffusion_training_pipeline() {
    let config = small_diffusion_config();
    let data_dim = config.data_dim();
    let mut adapter = DiffusionTrainableAdapter::new(config, Device::Cpu).unwrap();

    let batch_size = 4;
    let input = Tensor::randn(0f32, 1.0, (batch_size, data_dim), &Device::Cpu).unwrap();

    // The diffusion forward returns predicted noise.
    // Use the input as a pseudo-target (just to exercise the pipeline).
    // Loss values won't be meaningful but should be finite.
    let mut all_losses = Vec::new();

    for epoch in 0..20 {
        let predictions = adapter.forward(&input).unwrap();

        // Use input as target (exercising compute_loss, not expecting meaningful loss)
        let loss = adapter.compute_loss(&predictions, &input).unwrap();
        let loss_val = loss.to_scalar::<f32>().unwrap();

        assert!(loss_val.is_finite(), "Loss is NaN/Inf at epoch {}", epoch);
        all_losses.push(loss_val);

        let grad_norm = adapter.backward(&loss).unwrap();
        assert!(
            grad_norm.is_finite(),
            "Grad norm is NaN/Inf at epoch {}",
            epoch
        );

        adapter.optimizer_step().unwrap();
        adapter.zero_grad().unwrap();

        if epoch % 5 == 0 {
            println!(
                "Diffusion epoch {}: loss={:.6}, grad_norm={:.6}",
                epoch, loss_val, grad_norm
            );
        }
    }

    // Verify we got through all epochs without crash
    assert_eq!(all_losses.len(), 20);
    assert_eq!(adapter.get_step(), 20);
}

#[test]
fn test_diffusion_checkpoint_roundtrip() {
    let config = small_diffusion_config();
    let data_dim = config.data_dim();
    let mut adapter = DiffusionTrainableAdapter::new(config.clone(), Device::Cpu).unwrap();

    // Do a few forward passes
    let input = Tensor::randn(0f32, 1.0, (4, data_dim), &Device::Cpu).unwrap();
    for _ in 0..3 {
        let pred = adapter.forward(&input).unwrap();
        let loss = adapter.compute_loss(&pred, &input).unwrap();
        adapter.backward(&loss).unwrap();
        adapter.optimizer_step().unwrap();
    }

    // Save - Diffusion uses directory-based checkpoints
    let tmp_dir = std::env::temp_dir().join("diffusion_test_checkpoint");
    std::fs::create_dir_all(&tmp_dir).unwrap();
    let save_result = adapter.save_checkpoint(tmp_dir.to_str().unwrap());
    assert!(
        save_result.is_ok(),
        "Save failed: {:?}",
        save_result.err()
    );

    // Load
    let mut adapter2 = DiffusionTrainableAdapter::new(config, Device::Cpu).unwrap();
    let load_result = adapter2.load_checkpoint(tmp_dir.to_str().unwrap());
    assert!(
        load_result.is_ok(),
        "Load failed: {:?}",
        load_result.err()
    );

    // Cleanup
    let _ = std::fs::remove_dir_all(&tmp_dir);
}

#[test]
fn test_diffusion_3d_input() {
    let config = small_diffusion_config();
    let mut adapter = DiffusionTrainableAdapter::new(config.clone(), Device::Cpu).unwrap();

    // [batch=4, seq_len=8, feature_dim=1] — 3D input should be flattened internally
    let input = Tensor::randn(
        0f32,
        1.0,
        (4, config.seq_len, config.feature_dim),
        &Device::Cpu,
    )
    .unwrap();
    let output = adapter.forward(&input);
    assert!(output.is_ok(), "3D forward failed: {:?}", output.err());

    let output = output.unwrap();
    println!("Diffusion 3D output shape: {:?}", output.dims());
    assert!(output.elem_count() > 0);
}

#[test]
fn test_diffusion_metrics_collection() {
    let config = small_diffusion_config();
    let data_dim = config.data_dim();
    let mut adapter = DiffusionTrainableAdapter::new(config, Device::Cpu).unwrap();

    let input = Tensor::randn(0f32, 1.0, (4, data_dim), &Device::Cpu).unwrap();
    let pred = adapter.forward(&input).unwrap();
    let loss = adapter.compute_loss(&pred, &input).unwrap();
    adapter.backward(&loss).unwrap();
    adapter.optimizer_step().unwrap();

    let metrics = adapter.collect_metrics();
    assert!(metrics.loss.is_finite(), "Metrics loss should be finite");
    assert!(metrics.learning_rate > 0.0);
}

#[test]
fn test_diffusion_validation() {
    let config = small_diffusion_config();
    let data_dim = config.data_dim();
    let mut adapter = DiffusionTrainableAdapter::new(config, Device::Cpu).unwrap();

    let val_data: Vec<(Tensor, Tensor)> = (0..5)
        .map(|_| {
            let input = Tensor::randn(0f32, 1.0, (4, data_dim), &Device::Cpu).unwrap();
            let target = Tensor::randn(0f32, 1.0, (4, data_dim), &Device::Cpu).unwrap();
            (input, target)
        })
        .collect();

    let val_loss = adapter.validate(&val_data);
    assert!(val_loss.is_ok(), "Validation failed: {:?}", val_loss.err());
    let loss_val = val_loss.unwrap();
    assert!(
        loss_val.is_finite(),
        "Validation loss is not finite: {}",
        loss_val
    );
    println!("Diffusion validation loss: {:.6}", loss_val);
}