crater/analysis/rays/

bench_utils.rs

//! Shared utilities for ray casting benchmarks and examples

use crate::analysis::rays::prelude::*;
use crate::csg::prelude::*;
use burn::prelude::*;
use burn::tensor::cast::ToElement;

use burn::tensor::backend::AutodiffBackend;
use rand::rngs::StdRng;
use rand::{Rng, SeedableRng};
use std::fmt;

/// Results from a ray casting benchmark
#[derive(Debug, Clone)]
pub struct BenchmarkResult {
    pub algorithm: String,
    pub sphere_count: usize,
    pub ray_count: usize,
    pub successful_hits: usize,
    pub hit_rate: f32,
    pub average_distance: f32,
    pub ray_cast_time: std::time::Duration,
    pub rays_per_second: f64,
}

impl BenchmarkResult {
    /// Create a new benchmark result from ray cast data
    pub fn new<B: Backend>(
        result: &RayCastResult<B, 3>,
        elapsed: std::time::Duration,
        sphere_count: usize,
        ray_count: usize,
        algorithm: &str,
    ) -> Self {
        // Count successful hits (rays that didn't fail)
        let successful_mask = result.distances().lower_elem(f32::MAX);
        let successful_hits: usize = successful_mask.clone().int().sum().into_scalar().to_usize();

        // Calculate average distance for successful hits only
        let distances = result.distances();
        let valid_distances = distances.clone().mask_fill(successful_mask.bool_not(), 0.0);
        let average_distance = if successful_hits > 0 {
            valid_distances.sum().into_scalar().to_f32() / successful_hits as f32
        } else {
            0.0
        };

        let hit_rate = (successful_hits as f32 / ray_count as f32) * 100.0;
        let rays_per_second = ray_count as f64 / elapsed.as_secs_f64();

        Self {
            algorithm: algorithm.to_string(),
            sphere_count,
            ray_count,
            successful_hits,
            hit_rate,
            average_distance,
            ray_cast_time: elapsed,
            rays_per_second,
        }
    }
}

impl fmt::Display for BenchmarkResult {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        writeln!(
            f,
            "Results for {} algorithm with {} spheres and {} rays:",
            self.algorithm, self.sphere_count, self.ray_count
        )?;
        writeln!(f, "  Total rays: {}", self.ray_count)?;
        writeln!(f, "  Successful hits: {}", self.successful_hits)?;
        writeln!(f, "  Hit rate: {:.2}%", self.hit_rate)?;
        writeln!(f, "  Average distance: {:.4}", self.average_distance)?;
        writeln!(f, "  Ray-cast time: {:.2?}", self.ray_cast_time)?;
        writeln!(f, "  Rays per second: {:.0}", self.rays_per_second)?;
        Ok(())
    }
}

/// Create a region consisting of multiple random spheres
pub fn make_random_spheres<B: Backend>(
    num_spheres: usize,
    seed: u64,
    device: &B::Device,
) -> Region<3, B> {
    let mut rng = StdRng::seed_from_u64(seed);
    let mut spheres = Vec::new();

    for _ in 0..num_spheres {
        // Generate random center in unit cube [-1, 1]^3
        let center = [
            rng.random_range(-0.8..=0.8),
            rng.random_range(-0.8..=0.8),
            rng.random_range(-0.8..=0.8),
        ];

        // Generate random radius between 0.1 and 0.3
        let radius = rng.random_range(0.1..=0.3);

        // Create sphere and translate it to the random center
        let sphere = Field3D::<B>::sphere(radius, device.clone())
            .into_isosurface(0.0)
            .transform(Translate(center))
            .region();

        spheres.push(sphere);
    }

    // Create union of all spheres
    let mut union = spheres.remove(0);
    for sphere in spheres {
        union = &union | &sphere;
    }

    union
}

/// Generate random rays with the given seed
pub fn random_rays<B: Backend>(
    batch_size: usize,
    seed: u64,
    device: &B::Device,
) -> (Tensor<B, 2, Float>, Tensor<B, 2, Float>) {
    let mut rng = StdRng::seed_from_u64(seed);

    // Generate random directions (unit vectors)
    let mut directions = Vec::with_capacity(batch_size * 3);
    for _ in 0..batch_size {
        let mut v = [0.0f32; 3];
        loop {
            for val in &mut v {
                *val = rng.random_range(-1.0..=1.0);
            }
            let norm = (v[0] * v[0] + v[1] * v[1] + v[2] * v[2]).sqrt();
            if norm > 1e-6 {
                for val in &v {
                    directions.push(val / norm);
                }
                break;
            }
        }
    }

    // Generate random origins inside the unit cube
    let mut origins = Vec::with_capacity(batch_size * 3);
    for _ in 0..batch_size {
        let x = rng.random_range(-1.0..=1.0);
        let y = rng.random_range(-1.0..=1.0);
        let z = rng.random_range(-1.0..=1.0);
        origins.push(x);
        origins.push(y);
        origins.push(z);
    }

    let directions =
        Tensor::<B, 1, Float>::from_data(directions.as_slice(), device).reshape([batch_size, 3]);
    let origins =
        Tensor::<B, 1, Float>::from_data(origins.as_slice(), device).reshape([batch_size, 3]);
    (origins, directions)
}

/// Parse algorithm string into RayCastAlgorithm
pub fn parse_algorithm(algorithm: &str) -> RayCastAlgorithm<3> {
    match algorithm.to_lowercase().as_str() {
        "analytical" => RayCastAlgorithm::<3>::Analytical,
        "march_and_bisect" => RayCastAlgorithm::<3>::BracketAndBisect {
            lambda: 10.0,
            d_lambda: f32::EPSILON * 1e6,
            max_bisection_iterations: 30,
        },
        "newton" => RayCastAlgorithm::<3>::Newton {
            max_iterations: 250,
            nudge_distance: 0.01,
            step_size: 0.1,
        },
        _ => {
            eprintln!("Unknown algorithm: {}. Using analytical.", algorithm);
            RayCastAlgorithm::<3>::Analytical
        }
    }
}

/// Run a ray casting benchmark with the given parameters
pub fn run_raycast_benchmark<B: AutodiffBackend>(
    sphere_count: usize,
    ray_count: usize,
    algorithm: &str,
    sphere_seed: u64,
    ray_seed: u64,
    device: &B::Device,
) -> BenchmarkResult {
    let region = make_random_spheres::<B>(sphere_count, sphere_seed, device);
    let (origins, directions) = random_rays::<B>(ray_count, ray_seed, device);
    let algebra = Algebra::default();
    let rays = Rays::new(origins, directions);
    let parsed_algorithm = parse_algorithm(algorithm);

    let start_time = std::time::Instant::now();
    let result = region.ray_cast(rays, &algebra, parsed_algorithm);
    let elapsed = start_time.elapsed();

    BenchmarkResult::new(&result, elapsed, sphere_count, ray_count, algorithm)
}

/// Run a ray casting benchmark and return both the benchmark result and the ray cast result
pub fn run_raycast_benchmark_with_result<B: AutodiffBackend>(
    sphere_count: usize,
    ray_count: usize,
    algorithm: &str,
    sphere_seed: u64,
    ray_seed: u64,
    device: &B::Device,
) -> (BenchmarkResult, RayCastResult<B, 3>) {
    let region = make_random_spheres::<B>(sphere_count, sphere_seed, device);
    let (origins, directions) = random_rays::<B>(ray_count, ray_seed, device);
    let algebra = Algebra::default();
    let rays = Rays::new(origins, directions);
    let parsed_algorithm = parse_algorithm(algorithm);

    let start_time = std::time::Instant::now();
    let result = region.ray_cast(rays, &algebra, parsed_algorithm);
    let elapsed = start_time.elapsed();

    let benchmark_result =
        BenchmarkResult::new(&result, elapsed, sphere_count, ray_count, algorithm);
    (benchmark_result, result)
}