foxhunt/crates/ml/tests/diffusion_integration.rs

#![allow(
    clippy::assertions_on_constants,
    clippy::assertions_on_result_states,
    clippy::clone_on_copy,
    clippy::decimal_literal_representation,
    clippy::doc_markdown,
    clippy::empty_line_after_doc_comments,
    clippy::field_reassign_with_default,
    clippy::get_unwrap,
    clippy::identity_op,
    clippy::inconsistent_digit_grouping,
    clippy::indexing_slicing,
    clippy::integer_division,
    clippy::len_zero,
    clippy::let_underscore_must_use,
    clippy::manual_div_ceil,
    clippy::manual_let_else,
    clippy::manual_range_contains,
    clippy::modulo_arithmetic,
    clippy::needless_range_loop,
    clippy::non_ascii_literal,
    clippy::redundant_clone,
    clippy::shadow_reuse,
    clippy::shadow_same,
    clippy::shadow_unrelated,
    clippy::single_match_else,
    clippy::str_to_string,
    clippy::string_slice,
    clippy::tests_outside_test_module,
    clippy::too_many_lines,
    clippy::unnecessary_wraps,
    clippy::unseparated_literal_suffix,
    clippy::use_debug,
    clippy::useless_vec,
    clippy::wildcard_enum_match_arm,
    clippy::else_if_without_else,
    clippy::expect_used,
    clippy::missing_const_for_fn,
    clippy::similar_names,
    clippy::type_complexity,
    clippy::collapsible_else_if,
    clippy::doc_lazy_continuation,
    clippy::items_after_test_module,
    clippy::map_clone,
    clippy::multiple_unsafe_ops_per_block,
    clippy::unwrap_or_default,
    clippy::assign_op_pattern,
    clippy::needless_borrow,
    clippy::println_empty_string,
    clippy::unnecessary_cast,
    clippy::used_underscore_binding,
    clippy::create_dir,
    clippy::implicit_saturating_sub,
    clippy::exit,
    clippy::expect_fun_call,
    clippy::too_many_arguments,
    clippy::unnecessary_map_or,
    clippy::unwrap_used,
    dead_code,
    unused_imports,
    unused_variables,
    clippy::cloned_ref_to_slice_refs,
    clippy::neg_multiply,
    clippy::while_let_loop,
    clippy::bool_assert_comparison,
    clippy::excessive_precision,
    clippy::trivially_copy_pass_by_ref,
    clippy::op_ref,
    clippy::redundant_closure,
    clippy::unnecessary_lazy_evaluations,
    clippy::if_then_some_else_none,
    clippy::unnecessary_to_owned,
    clippy::single_component_path_imports,
)]
//! 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.
//!
//! The adapter uses GPU-native StreamTensor (aliased as GpuTensor).
//! forward_gpu / compute_loss_gpu operate on GPU tensors directly.
//! The UnifiedTrainable trait's forward_loss takes &[f32] slices.

use ml::diffusion::config::DiffusionConfig;
use ml::diffusion::trainable::DiffusionTrainableAdapter;
use ml::training::unified_trainer::UnifiedTrainable;
use ml_supervised::gpu_tensor::GpuTensor;
use std::sync::Arc;
use tracing::info;

/// Create a CUDA stream for test tensor allocation.
fn test_stream() -> Arc<cudarc::driver::CudaStream> {
    let ctx = cudarc::driver::CudaContext::new(0).expect("CUDA device required");
    ctx.new_stream().expect("CUDA stream required")
}

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 stream = test_stream();
    let adapter = DiffusionTrainableAdapter::new(config, &stream);
    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 stream = test_stream();
    let mut adapter = DiffusionTrainableAdapter::new(config, &stream).unwrap();

    // [batch=4, data_dim=8]
    let input = GpuTensor::randn(&[4, data_dim], 1.0, &stream).unwrap();
    let output = adapter.forward_gpu(&input);
    assert!(output.is_ok(), "Forward failed: {:?}", output.err());

    let output = output.unwrap();
    info!(dims = ?output.shape, "Diffusion output shape");
    assert_eq!(output.shape[0], 4, "Batch dimension should be 4");
    // Output is predicted noise, should be finite
    let host = output.to_vec().unwrap();
    let sum: f32 = host.iter().map(|x| x.abs()).sum();
    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 stream = test_stream();
    let mut adapter = DiffusionTrainableAdapter::new(config, &stream).unwrap();

    let batch_size = 4;
    let input = GpuTensor::randn(&[batch_size, data_dim], 1.0, &stream).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_gpu(&input).unwrap();

        // Use input as target (exercising compute_loss_gpu, not expecting meaningful loss)
        let loss_val = adapter.compute_loss_gpu(&predictions, &input).unwrap();

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

        let grad_norm = adapter.backward(loss_val as f64).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 {
            info!(epoch, loss_val, grad_norm, "Diffusion epoch");
        }
    }

    // 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 stream = test_stream();
    let mut adapter = DiffusionTrainableAdapter::new(config.clone(), &stream).unwrap();

    let data_dim = config.data_dim();
    // Do a few forward passes
    let input = GpuTensor::randn(&[4, data_dim], 1.0, &stream).unwrap();
    for _ in 0..3 {
        let pred = adapter.forward_gpu(&input).unwrap();
        let loss = adapter.compute_loss_gpu(&pred, &input).unwrap();
        adapter.backward(loss as f64).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, &stream).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 stream = test_stream();
    let mut adapter = DiffusionTrainableAdapter::new(config.clone(), &stream).unwrap();

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

    let output = output.unwrap();
    info!(dims = ?output.shape, "Diffusion 3D output shape");
    assert!(output.numel() > 0);
}

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

    let input = GpuTensor::randn(&[4, data_dim], 1.0, &stream).unwrap();
    let pred = adapter.forward_gpu(&input).unwrap();
    let loss = adapter.compute_loss_gpu(&pred, &input).unwrap();
    adapter.backward(loss as f64).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 stream = test_stream();
    let mut adapter = DiffusionTrainableAdapter::new(config, &stream).unwrap();

    let val_data: Vec<(GpuTensor, GpuTensor)> = (0..5)
        .map(|_| {
            let input = GpuTensor::randn(&[4, data_dim], 1.0, &stream).unwrap();
            let target = GpuTensor::randn(&[4, data_dim], 1.0, &stream).unwrap();
            (input, target)
        })
        .collect();

    let val_loss = adapter.validate_gpu(&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
    );
    info!(loss_val, "Diffusion validation loss");
}