rustytorch/examples/pinn_mre_helmholtz/tests/validate_training.rs

//! Validation Tests for PINN Training
//!
//! These tests verify that the training loop actually updates weights
//! and that loss decreases over training.

use pinn_mre_helmholtz::{GpuAdam, GradientWorkspace};
use rtx_tensor::{Tensor, Device};

/// Test that GpuAdam actually modifies parameters
#[test]
fn test_gpu_adam_modifies_weights() {
    let device = Device::Cpu;

    // Create simple parameter and gradient
    let original_val = 1.0f32;
    let grad_val = 0.5f32;

    let shapes = vec![vec![1]];
    let mut opt = GpuAdam::with_defaults(&shapes, &device).unwrap();

    let mut params = vec![
        Tensor::from_vec(vec![original_val], &[1], &device).unwrap()
    ];
    let grads = vec![
        Tensor::from_vec(vec![grad_val], &[1], &device).unwrap()
    ];

    // Copy original value
    let original = params[0].to_cpu().unwrap()[0];

    // Perform optimization step
    opt.step(&mut params, &grads).unwrap();

    // Check that parameter changed
    let updated = params[0].to_cpu().unwrap()[0];
    assert!(
        (updated - original).abs() > 1e-8,
        "Parameter should have changed: original={}, updated={}",
        original,
        updated
    );

    // With positive gradient, parameter should decrease (gradient descent)
    assert!(
        updated < original,
        "With positive gradient, parameter should decrease: original={}, updated={}",
        original,
        updated
    );
}

/// Test that multiple Adam steps continue to modify parameters
#[test]
fn test_gpu_adam_multiple_steps() {
    let device = Device::Cpu;

    let shapes = vec![vec![2, 2]];
    // Use larger learning rate for faster convergence in test
    let mut opt = GpuAdam::new(&shapes, &device, 0.1, 0.9, 0.999, 1e-8).unwrap();

    let mut params = vec![
        Tensor::from_vec(vec![1.0f32, 2.0, 3.0, 4.0], &[2, 2], &device).unwrap()
    ];

    // Track parameter norm over steps
    let mut norms = Vec::new();
    let initial_norm: f32 = params[0].to_cpu().unwrap().iter().map(|x| x * x).sum::<f32>().sqrt();
    norms.push(initial_norm);

    // Perform 100 steps with constant gradient pointing toward zero
    for _ in 0..100 {
        let param_data = params[0].to_cpu().unwrap();
        // Gradient = parameter (so we're doing gradient descent on ||params||^2)
        let grads = vec![
            Tensor::from_vec(param_data, &[2, 2], &device).unwrap()
        ];

        opt.step(&mut params, &grads).unwrap();

        let norm: f32 = params[0].to_cpu().unwrap().iter().map(|x| x * x).sum::<f32>().sqrt();
        norms.push(norm);
    }

    // Check that the first 10 steps show decreasing norm
    for i in 1..10.min(norms.len()) {
        assert!(
            norms[i] < norms[i - 1],
            "Norm should decrease: step {} norm = {}, step {} norm = {}",
            i - 1,
            norms[i - 1],
            i,
            norms[i]
        );
    }

    // Final norm should be significantly smaller than initial
    assert!(
        *norms.last().unwrap() < *norms.first().unwrap() * 0.5,
        "Final norm should be at least 50% smaller than initial: initial={}, final={}",
        norms.first().unwrap(),
        norms.last().unwrap()
    );
}

/// Test GradientWorkspace creation
#[test]
fn test_gradient_workspace_creation() {
    let device = Device::Cpu;

    let batch_size = 32;
    let ff_dim = 16;
    let hidden_dim = 64;
    let num_hidden_layers = 3;

    let ws = GradientWorkspace::new(batch_size, ff_dim, hidden_dim, num_hidden_layers, &device).unwrap();

    // Check dimensions
    assert_eq!(ws.num_layers(), 4); // 3 hidden + 1 output

    // Check z/h shapes
    assert_eq!(ws.z.len(), 4);
    assert_eq!(ws.h.len(), 4);

    // Check dW/db shapes
    assert_eq!(ws.dW.len(), 4);
    assert_eq!(ws.db.len(), 4);

    // First layer: [hidden_dim, ff_dim * 2]
    assert_eq!(ws.dW[0].shape(), &[64, 32]);
    assert_eq!(ws.db[0].shape(), &[64]);

    // Output layer: [2, hidden_dim]
    assert_eq!(ws.dW[3].shape(), &[2, 64]);
    assert_eq!(ws.db[3].shape(), &[2]);
}

/// Test MSE backward gradient
#[test]
fn test_mse_backward_gradient() {
    use pinn_mre_helmholtz::mse_backward;

    let device = Device::Cpu;

    // Predictions and targets
    let pred = Tensor::from_vec(vec![1.0f32, 2.0, 3.0, 4.0], &[2, 2], &device).unwrap();
    let target = Tensor::from_vec(vec![0.0f32, 0.0, 0.0, 0.0], &[2, 2], &device).unwrap();

    let grad = mse_backward(&pred, &target).unwrap();

    // dL/dpred = (2/N) * (pred - target) = (2/2) * pred = pred
    let grad_cpu = grad.to_cpu().unwrap();
    assert!((grad_cpu[0] - 1.0).abs() < 1e-5, "Expected 1.0, got {}", grad_cpu[0]);
    assert!((grad_cpu[1] - 2.0).abs() < 1e-5, "Expected 2.0, got {}", grad_cpu[1]);
    assert!((grad_cpu[2] - 3.0).abs() < 1e-5, "Expected 3.0, got {}", grad_cpu[2]);
    assert!((grad_cpu[3] - 4.0).abs() < 1e-5, "Expected 4.0, got {}", grad_cpu[3]);
}

/// Test layer backward computation
#[test]
fn test_layer_backward() {
    use pinn_mre_helmholtz::layer_backward;

    let device = Device::Cpu;

    // Simple case: 2 samples, 2 output units, 2 input units
    let d_l_dh = Tensor::from_vec(vec![1.0f32, 1.0, 1.0, 1.0], &[2, 2], &device).unwrap();
    let h = Tensor::from_vec(vec![0.5f32, 0.5, 0.5, 0.5], &[2, 2], &device).unwrap();
    let h_prev = Tensor::from_vec(vec![1.0f32, 0.5, 1.0, 0.5], &[2, 2], &device).unwrap();
    let w = Tensor::from_vec(vec![1.0f32, 0.0, 0.0, 1.0], &[2, 2], &device).unwrap();

    let (d_w, d_b, dh_prev) = layer_backward(&d_l_dh, &h, &h_prev, &w, true).unwrap();

    // Check shapes
    assert_eq!(d_w.shape(), &[2, 2]);
    assert_eq!(d_b.shape(), &[2]);
    assert_eq!(dh_prev.shape(), &[2, 2]);

    // With h=0.5, tanh'(z) = 1 - 0.5^2 = 0.75
    // dL_dz = dL_dh * 0.75 = 0.75 for all elements

    // db = sum(dL_dz, axis=0) = [1.5, 1.5]
    let db_cpu = d_b.to_cpu().unwrap();
    assert!((db_cpu[0] - 1.5).abs() < 1e-5, "db[0] should be 1.5, got {}", db_cpu[0]);
    assert!((db_cpu[1] - 1.5).abs() < 1e-5, "db[1] should be 1.5, got {}", db_cpu[1]);
}

/// Test that PINN solver training actually reduces loss
#[test]
fn test_pinn_training_reduces_loss() {
    use pinn_mre_helmholtz::{Config, Mre1DPinnSolver};

    // Create solver with small network for quick test
    let mut cfg = Config::default();
    cfg.n_data = 50;  // Smaller dataset for faster test
    cfg.u_layers = 2;  // Fewer layers
    cfg.u_hidden = 32;
    cfg.u_ff_dim = 16;
    cfg.epochs = 100;  // We'll control steps manually

    let mut solver = Mre1DPinnSolver::new(cfg).expect("Failed to create solver");

    // Run forward pass to get initial loss
    solver.train_step_data_only().expect("First training step failed");
    let initial_loss = solver.compute_data_loss_for_logging()
        .expect("Failed to compute initial loss");

    println!("Initial loss: {:.6e}", initial_loss);

    // Run 500 training steps
    for step in 1..=500 {
        solver.train_step_data_only().expect("Training step failed");

        // Log progress every 100 steps
        if step % 100 == 0 {
            let loss = solver.compute_data_loss_for_logging().unwrap();
            println!("Step {}: loss = {:.6e}", step, loss);
        }
    }

    // Compute final loss
    let final_loss = solver.compute_data_loss_for_logging()
        .expect("Failed to compute final loss");

    println!("Final loss: {:.6e}", final_loss);
    println!("Improvement: {:.1}x", initial_loss / final_loss);

    // The loss should decrease over training
    // We expect at least some improvement (2% reduction - conservative threshold)
    // Note: The network may converge quickly on this simple problem
    assert!(
        final_loss < initial_loss,
        "Training should reduce loss: initial={:.6e}, final={:.6e}",
        initial_loss,
        final_loss
    );
}

/// Profile training step to identify bottlenecks
#[test]
#[cfg(feature = "cuda")]
fn test_profile_training_step() {
    use pinn_mre_helmholtz::{Config, Mre1DPinnSolver};
    use std::time::Instant;

    let mut cfg = Config::default();
    cfg.n_data = 200;  // Same as benchmark
    cfg.u_layers = 3;
    cfg.u_hidden = 64;
    cfg.u_ff_dim = 64;

    let mut solver = Mre1DPinnSolver::new(cfg).expect("Failed to create solver");

    // Warm up
    for _ in 0..5 {
        solver.train_step_data_only().expect("Training step failed");
    }

    // Get CUDA context for sync
    let ctx = rtx_tensor::storage::cuda_manager::get_or_create_context(0).unwrap();
    ctx.synchronize().unwrap();

    // Time 10 complete training steps
    let start = Instant::now();
    for _ in 0..10 {
        solver.train_step_data_only().expect("Training step failed");
    }
    ctx.synchronize().unwrap();
    let elapsed = start.elapsed();

    let per_step = elapsed / 10;
    println!("\n=== TRAINING STEP PROFILING ===");
    println!("Total time for 10 steps: {:?}", elapsed);
    println!("Per step: {:?}", per_step);
    println!("Target: < 100µs");
    println!("Status: {}", if per_step.as_micros() < 100 { "✅ PASS" } else { "❌ FAIL - too slow!" });
}

/// Test the train_analytical convenience method
#[test]
fn test_train_analytical_method() {
    use pinn_mre_helmholtz::{Config, Mre1DPinnSolver};

    let mut cfg = Config::default();
    cfg.n_data = 30;
    cfg.u_layers = 2;
    cfg.u_hidden = 16;
    cfg.u_ff_dim = 8;

    let mut solver = Mre1DPinnSolver::new(cfg).expect("Failed to create solver");

    // Run 100 steps with logging
    let best_loss = solver.train_analytical(100, 50)
        .expect("train_analytical failed");

    println!("Best loss from train_analytical: {:.6e}", best_loss);

    // Should produce a finite, reasonable loss
    assert!(best_loss.is_finite(), "Loss should be finite");
    assert!(best_loss < 1.0, "Loss should be reasonable (< 1.0)");
}