rustytorch/demos/rtx-neural-operator-demo/examples/benchmark_solvers.rs

//! PDE Solver Benchmark Demo
//!
//! This example demonstrates the benchmark system comparing:
//! - FNO (Fourier Neural Operator) - learned PDE solver
//! - FDM (Finite Difference Method) - classical numerical method
//! - FEM (Finite Element Method) - classical numerical method
//!
//! Usage:
//! ```bash
//! cargo run --package rtx-neural-operator-demo --example benchmark_solvers
//! cargo run --package rtx-neural-operator-demo --example benchmark_solvers -- --quick
//! cargo run --package rtx-neural-operator-demo --example benchmark_solvers -- --comprehensive
//! ```

use rtx_neural_operator_demo::{BenchmarkConfig, BenchmarkRunner};
use std::env;

fn main() {
    println!("\n{:=<80}", "");
    println!("Neural Operator PDE Benchmark Demo");
    println!("Comparing FNO, FDM, and FEM solvers");
    println!("{:=<80}\n", "");

    // Parse command line arguments
    let args: Vec<String> = env::args().collect();
    let config = if args.contains(&"--quick".to_string()) {
        println!("Running quick benchmark (small resolutions, few trials)...\n");
        BenchmarkConfig::quick()
    } else if args.contains(&"--comprehensive".to_string()) {
        println!("Running comprehensive benchmark (all resolutions, many trials)...\n");
        BenchmarkConfig::comprehensive()
    } else {
        println!("Running standard benchmark...");
        println!("Use --quick for faster testing or --comprehensive for full analysis\n");
        BenchmarkConfig::default()
    };

    // Display configuration
    println!("Configuration:");
    println!("  PDE Type: {}", config.pde_type.name());
    println!("  Resolutions: {:?}", config.resolutions);
    println!("  Problems per resolution: {}", config.n_problems);
    println!("  Reference resolution: {}", config.reference_resolution);
    println!();

    // Run benchmarks
    println!("Running benchmarks...\n");
    let runner = BenchmarkRunner::new(config);
    let results = runner.run_classical_benchmarks();

    // Display results
    BenchmarkRunner::print_results(&results);

    // Compute and display summary statistics
    println!("\nRunning 3 trials for statistical analysis...\n");
    let summaries = runner.run_with_statistics(3);
    BenchmarkRunner::print_summary(&summaries);

    // Analyze speedup between methods
    print_speedup_analysis(&summaries);

    // Show accuracy comparison
    print_accuracy_analysis(&results);

    // Memory efficiency analysis
    print_memory_analysis(&results);

    println!("\n{:=<80}", "");
    println!("Benchmark Complete!");
    println!("{:=<80}\n", "");

    println!("Note: To benchmark FNO, train a model first and use the");
    println!("      benchmark_fno() method with a trained model.");
}

fn print_speedup_analysis(summaries: &[rtx_neural_operator_demo::BenchmarkSummary]) {
    println!("\n{:=<80}", "");
    println!("Speedup Analysis");
    println!("{:=<80}\n", "");

    // Group by resolution
    let mut by_resolution: std::collections::HashMap<
        usize,
        Vec<&rtx_neural_operator_demo::BenchmarkSummary>,
    > = std::collections::HashMap::new();

    for summary in summaries {
        by_resolution
            .entry(summary.resolution)
            .or_default()
            .push(summary);
    }

    for (resolution, res_summaries) in by_resolution.iter() {
        println!("Resolution: {}x{}", resolution, resolution);
        println!("{:-<40}", "");

        let fdm_summary = res_summaries.iter().find(|s| s.method.contains("FDM"));
        let fem_summary = res_summaries.iter().find(|s| s.method.contains("FEM"));

        if let (Some(fdm), Some(fem)) = (fdm_summary, fem_summary) {
            let speedup = fem.speedup_vs(fdm);
            if speedup > 1.0 {
                println!(
                    "  FDM is {:.2}x faster than FEM ({:.2}ms vs {:.2}ms)",
                    speedup, fdm.avg_time_ms, fem.avg_time_ms
                );
            } else {
                println!(
                    "  FEM is {:.2}x faster than FDM ({:.2}ms vs {:.2}ms)",
                    1.0 / speedup,
                    fem.avg_time_ms,
                    fdm.avg_time_ms
                );
            }
        }

        println!();
    }
}

fn print_accuracy_analysis(results: &[rtx_neural_operator_demo::BenchmarkResult]) {
    println!("\n{:=<80}", "");
    println!("Accuracy Analysis (L2 Error vs Analytical Solution)");
    println!("{:=<80}\n", "");

    // Group by resolution
    let mut by_resolution: std::collections::HashMap<
        usize,
        Vec<&rtx_neural_operator_demo::BenchmarkResult>,
    > = std::collections::HashMap::new();

    for result in results {
        by_resolution
            .entry(result.resolution)
            .or_default()
            .push(result);
    }

    for (resolution, res_results) in by_resolution.iter() {
        println!("Resolution: {}x{}", resolution, resolution);
        println!("{:-<40}", "");

        for result in res_results {
            if let Some(l2_error) = result.l2_error {
                println!("  {:<12} L2 error: {:.6e}", result.method, l2_error);
            }
        }

        println!();
    }

    println!("Note: Lower L2 error indicates better accuracy.");
    println!("      Both FDM and FEM converge to the exact solution as resolution increases.");
}

fn print_memory_analysis(results: &[rtx_neural_operator_demo::BenchmarkResult]) {
    println!("\n{:=<80}", "");
    println!("Memory Efficiency Analysis");
    println!("{:=<80}\n", "");

    // Group by method
    let mut by_method: std::collections::HashMap<
        String,
        Vec<&rtx_neural_operator_demo::BenchmarkResult>,
    > = std::collections::HashMap::new();

    for result in results {
        by_method
            .entry(result.method.clone())
            .or_default()
            .push(result);
    }

    for (method, method_results) in by_method.iter() {
        println!("Method: {}", method);
        println!("{:-<40}", "");

        let mut sorted_results = method_results.to_vec();
        sorted_results.sort_by_key(|r| r.resolution);

        for result in &sorted_results {
            println!(
                "  {}x{}: {:.2} MB",
                result.resolution, result.resolution, result.memory_mb
            );
        }

        // Compute scaling factor
        if sorted_results.len() >= 2 {
            let first = sorted_results[0];
            let last = sorted_results[sorted_results.len() - 1];
            let res_ratio = (last.resolution as f64 / first.resolution as f64).powi(2);
            let mem_ratio = last.memory_mb / first.memory_mb;

            println!(
                "  Scaling: {:.2}x memory for {:.1}x grid points (O(N²) expected)",
                mem_ratio, res_ratio
            );
        }

        println!();
    }

    println!("Note: Memory should scale as O(N²) for 2D problems.");
    println!("      FEM uses more memory than FDM due to CG solver vectors.");
}