rustytorch/crates/specialized/rtx-cfd/tests/kernels_tests.rs

//! Tests for GPU kernels module
//!
//! This module tests the real CUDA implementations of CFD kernels including:
//! - Advection schemes (upwind, central differencing)
//! - Diffusion schemes (implicit/explicit)
//! - Pressure Poisson solver
//! - Matrix operations for CFD

use approx::assert_relative_eq;
use rtx_cfd::{CfdConfig, CfdResult};

#[cfg(feature = "cuda")]
mod cuda_tests {
    use super::*;
    use rtx_cfd::kernels::{
        AdvectionKernel, AdvectionScheme, CudaKernelManager, DiffusionKernel, DiffusionScheme,
        MatrixOpsKernel, PoissonKernel,
    };

    #[tokio::test]
    async fn test_advection_kernel_upwind() -> CfdResult<()> {
        let config = CfdConfig::new().with_gpu(true);
        let kernel_manager = std::sync::Arc::new(CudaKernelManager::new(&config)?);
        let advection_kernel = AdvectionKernel::new(&kernel_manager, AdvectionScheme::Upwind)?;

        // Test data: 1D advection with known analytical solution
        let nx = 128;
        let dx = 1.0 / (nx as f64);
        let dt = 0.001;
        let velocity = 1.0;

        // Initial condition: Gaussian pulse
        let mut phi = vec![0.0f32; nx];
        let center = nx / 4;
        let sigma = 5.0;
        for i in 0..nx {
            let x = (i as f64 - center as f64) * dx;
            phi[i] = (-(x * x) / (2.0 * sigma * sigma)).exp() as f32;
        }

        // Allocate GPU memory
        let mut d_phi = kernel_manager.allocate_f32(nx)?;
        let mut d_phi_new = kernel_manager.allocate_f32(nx)?;

        // Copy to GPU
        d_phi = kernel_manager.copy_to_device(&phi)?;

        // Run advection kernel
        // CFL = v dt / dx ≈ 0.1: march 50 steps so the pulse moves several cells.
        for _ in 0..50 {
            advection_kernel.apply(
                &d_phi,
                &mut d_phi_new,
                velocity as f32,
                dt as f32,
                dx as f32,
            )?;
            std::mem::swap(&mut d_phi, &mut d_phi_new);
        }

        // Copy result back
        let mut result = vec![0.0f32; nx];
        result = kernel_manager.copy_from_device(&d_phi)?;

        // Verify that the pulse has moved (mass conservation)
        let initial_mass: f32 = phi.iter().sum();
        let final_mass: f32 = result.iter().sum();
        assert_relative_eq!(initial_mass, final_mass, epsilon = 1e-6);

        // Verify that maximum has shifted in the correct direction
        let initial_max_idx = phi
            .iter()
            .enumerate()
            .max_by(|a, b| a.1.total_cmp(b.1))
            .unwrap()
            .0;
        let final_max_idx = result
            .iter()
            .enumerate()
            .max_by(|a, b| a.1.total_cmp(b.1))
            .unwrap()
            .0;

        assert!(
            final_max_idx > initial_max_idx,
            "Pulse should move in positive direction"
        );

        Ok(())
    }

    #[tokio::test]
    async fn test_advection_kernel_central() -> CfdResult<()> {
        let config = CfdConfig::new().with_gpu(true);
        let kernel_manager = std::sync::Arc::new(CudaKernelManager::new(&config)?);
        let advection_kernel = AdvectionKernel::new(&kernel_manager, AdvectionScheme::Central)?;

        // Test data: smooth sinusoidal wave
        let nx = 256;
        let dx = 2.0 * std::f64::consts::PI / (nx as f64);
        let dt = 0.001;
        let velocity = 1.0;

        let mut phi = vec![0.0f32; nx];
        for i in 0..nx {
            let x = i as f64 * dx;
            phi[i] = (2.0 * x).sin() as f32;
        }

        let mut d_phi = kernel_manager.allocate_f32(nx)?;
        let mut d_phi_new = kernel_manager.allocate_f32(nx)?;

        d_phi = kernel_manager.copy_to_device(&phi)?;
        advection_kernel.apply(
            &d_phi,
            &mut d_phi_new,
            velocity as f32,
            dt as f32,
            dx as f32,
        )?;

        let mut result = vec![0.0f32; nx];
        result = kernel_manager.copy_from_device(&d_phi_new)?;

        // For central scheme, verify mass conservation and smoothness
        let initial_mass: f32 = phi.iter().sum();
        let final_mass: f32 = result.iter().sum();
        assert_relative_eq!(initial_mass, final_mass, epsilon = 1e-5);

        Ok(())
    }

    #[tokio::test]
    async fn test_diffusion_kernel_explicit() -> CfdResult<()> {
        let config = CfdConfig::new().with_gpu(true);
        let kernel_manager = std::sync::Arc::new(CudaKernelManager::new(&config)?);
        let diffusion_kernel = DiffusionKernel::new(&kernel_manager, DiffusionScheme::Explicit)?;

        // Test 1D heat equation with analytical solution
        let nx = 128;
        let dx = 1.0 / (nx as f64);
        let dt = 0.0001; // Small dt for stability
        let alpha = 0.1; // Thermal diffusivity

        // Initial condition: step function
        let mut temperature = vec![0.0f32; nx];
        for i in nx / 4..3 * nx / 4 {
            temperature[i] = 1.0;
        }

        let mut d_temp = kernel_manager.allocate_f32(nx)?;
        let mut d_temp_new = kernel_manager.allocate_f32(nx)?;

        d_temp = kernel_manager.copy_to_device(&temperature)?;
        diffusion_kernel.apply(&d_temp, &mut d_temp_new, alpha as f32, dt as f32, dx as f32)?;

        let mut result = vec![0.0f32; nx];
        result = kernel_manager.copy_from_device(&d_temp_new)?;

        // Verify diffusion: edges should be smoother, total heat conserved
        let initial_total: f32 = temperature.iter().sum();
        let final_total: f32 = result.iter().sum();
        assert_relative_eq!(initial_total, final_total, epsilon = 1e-6);

        // Verify smoothing: gradient at edges should be reduced
        let initial_gradient = (temperature[nx / 4] - temperature[nx / 4 - 1]).abs();
        let final_gradient = (result[nx / 4] - result[nx / 4 - 1]).abs();
        assert!(
            final_gradient < initial_gradient,
            "Diffusion should smooth gradients"
        );

        Ok(())
    }

    #[tokio::test]
    async fn test_poisson_kernel_2d() -> CfdResult<()> {
        let config = CfdConfig::new().with_gpu(true);
        let kernel_manager = std::sync::Arc::new(CudaKernelManager::new(&config)?);
        let poisson_kernel = PoissonKernel::new(&kernel_manager)?;

        // Test 2D Poisson equation: ∇²φ = f
        // Analytical solution: φ(x,y) = sin(πx)sin(πy), f = -2π²sin(πx)sin(πy)
        let nx = 64;
        let ny = 64;
        let dx = 1.0 / (nx as f64);
        let dy = 1.0 / (ny as f64);

        // Setup source term f
        let mut source = vec![0.0f32; nx * ny];
        for j in 0..ny {
            for i in 0..nx {
                let x = i as f64 * dx;
                let y = j as f64 * dy;
                let idx = j * nx + i;
                source[idx] = (-2.0
                    * std::f64::consts::PI.powi(2)
                    * (std::f64::consts::PI * x).sin()
                    * (std::f64::consts::PI * y).sin()) as f32;
            }
        }

        // Initial guess (zeros)
        let phi = vec![0.0f32; nx * ny];

        let mut d_phi = kernel_manager.allocate_f32(nx * ny)?;
        let mut d_source = kernel_manager.allocate_f32(nx * ny)?;

        d_phi = kernel_manager.copy_to_device(&phi)?;
        d_source = kernel_manager.copy_to_device(&source)?;

        // Solve Poisson equation
        // Jacobi needs O(n²) sweeps on this grid and its residual floor in f32 sits
        // above 1e-6; the accuracy check below is the pin, the stop is loose.
        let max_iterations = 20_000;
        let tolerance = 1e-2; // absolute, on a source of size 2π²
        let iterations = poisson_kernel.solve_2d(
            &mut d_phi,
            &d_source,
            nx,
            ny,
            dx as f32,
            dy as f32,
            max_iterations,
            tolerance,
        )?;

        let mut result = vec![0.0f32; nx * ny];
        result = kernel_manager.copy_from_device(&d_phi)?;

        // Verify against analytical solution
        let mut max_error = 0.0f32;
        for j in 1..ny - 1 {
            for i in 1..nx - 1 {
                let x = i as f64 * dx;
                let y = j as f64 * dy;
                let idx = j * nx + i;
                let analytical =
                    (std::f64::consts::PI * x).sin() * (std::f64::consts::PI * y).sin();
                let error = (result[idx] - analytical as f32).abs();
                max_error = max_error.max(error);
            }
        }

        assert!(iterations < max_iterations, "Solver should converge");
        assert!(max_error < 0.1, "Solution should be reasonably accurate");

        Ok(())
    }

    #[tokio::test]
    async fn test_matrix_ops_kernel() -> CfdResult<()> {
        let config = CfdConfig::new().with_gpu(true);
        let kernel_manager = std::sync::Arc::new(CudaKernelManager::new(&config)?);
        let matrix_ops = MatrixOpsKernel::new(&kernel_manager)?;

        // Test sparse matrix-vector multiplication (typical in CFD)
        let n = 1000;
        let mut diagonal = vec![2.0f32; n];
        let mut off_diagonal = vec![-1.0f32; n - 1];
        let mut x = vec![1.0f32; n];

        // Create tridiagonal matrix (common in CFD discretizations)
        let mut d_diag = kernel_manager.allocate_f32(n)?;
        let mut d_off_diag = kernel_manager.allocate_f32(n - 1)?;
        let mut d_x = kernel_manager.allocate_f32(n)?;
        let mut d_y = kernel_manager.allocate_f32(n)?;

        d_diag = kernel_manager.copy_to_device(&diagonal)?;
        d_off_diag = kernel_manager.copy_to_device(&off_diagonal)?;
        d_x = kernel_manager.copy_to_device(&x)?;

        // Perform A*x = y operation
        matrix_ops.tridiagonal_matvec(&d_diag, &d_off_diag, &d_x, &mut d_y)?;

        let mut result = vec![0.0f32; n];
        result = kernel_manager.copy_from_device(&d_y)?;

        // Verify result manually for first few elements
        assert_relative_eq!(result[0], 2.0 * x[0] - x[1], epsilon = 1e-6);
        assert_relative_eq!(result[1], -x[0] + 2.0 * x[1] - x[2], epsilon = 1e-6);
        assert_relative_eq!(result[n - 1], -x[n - 2] + 2.0 * x[n - 1], epsilon = 1e-6);

        Ok(())
    }
}

#[cfg(not(feature = "cuda"))]
mod cpu_fallback_tests {
    use super::*;

    #[test]
    fn test_cpu_fallback_advection() {
        // Basic test to ensure CPU fallback works
        let config = CfdConfig::new().with_gpu(false);
        // Implementation will provide CPU fallback
        assert!(config.validate().is_ok());
    }
}

#[tokio::test]
async fn test_kernel_manager_initialization() -> CfdResult<()> {
    let config = CfdConfig::new();

    #[cfg(feature = "cuda")]
    {
        // Should successfully create kernel manager
        let mut result = rtx_cfd::kernels::CudaKernelManager::new(&config);
        // May fail if no CUDA device available - that's expected in CI
        match result {
            Ok(_) => println!("CUDA kernel manager created successfully"),
            Err(e) => println!("CUDA not available (expected in CI): {}", e),
        }
    }

    #[cfg(not(feature = "cuda"))]
    {
        println!("CUDA feature not enabled - using CPU fallback");
    }

    Ok(())
}