crater/analysis/rays/

algo.rs

//! Ray-surface intersection algorithms.
//!
//! This module implements various algorithms for finding intersections between rays and
//! implicit surfaces. For the mathematical theory and algorithm descriptions, see the
//! [Ray Casting](../../book/theory/ray-casting.md) chapter.

use super::prelude::*;
use crate::csg::prelude::*;
use burn::prelude::*;
use burn::tensor::{
    Tensor,
    backend::{AutodiffBackend, Backend},
};
use log::info;

pub const RAY_FAIL_VALUE: f32 = f32::MAX;

/// Ray-surface intersection algorithms for implicit surface ray casting.
///
/// This enum provides multiple algorithmic approaches for solving ray-surface intersection
/// problems, each optimized for different scenarios in safety-critical applications where
/// precision and reliability are paramount.
///
/// # Algorithm Selection Guidelines
///
/// * **Use [`Analytical`](RayCastAlgorithm::Analytical)** for:
///   - Simple primitive shapes (spheres, planes, cylinders)
///   - Maximum accuracy and performance requirements
///   - Real-time applications with tight computational budgets
///
/// * **Use [`Newton`](RayCastAlgorithm::Newton)** for:
///   - Smooth, well-conditioned implicit surfaces
///   - When autodifferentiation is available
///   - Applications requiring fast convergence with few iterations
///
/// * **Use [`BracketAndBisect`](RayCastAlgorithm::BracketAndBisect)** for:
///   - Complex CSG constructions with sharp edges
///   - Robustness over speed requirements
///   - Surfaces with potential numerical instabilities
///
/// For theoretical background, see the [Ray Casting](../../book/theory/ray-casting.md) chapter.
#[derive(Debug, Clone, Copy)]
pub enum RayCastAlgorithm<const N: usize> {
    /// Robust numerical method using bracketing followed by bisection.
    ///
    /// This algorithm first advances along each ray to bracket sign changes in the
    /// scalar field, then uses bisection to refine intersection points to high precision.
    /// It's the most reliable method for complex CSG constructions.
    ///
    /// # Parameters
    ///
    /// * `lambda` - Maximum ray extension distance for bracketing phase
    /// * `d_lambda` - Step size during bracketing search
    /// * `max_bisection_iterations` - Maximum refinement iterations (typically 20-50)
    ///
    /// # Convergence
    ///
    /// Guaranteed to converge for continuous scalar fields with precision ε ≈ λ/2ⁿ
    /// where n is the number of bisection iterations.
    BracketAndBisect {
        /// Maximum distance to extend rays during bracketing
        lambda: f32,
        /// Step size for bracketing phase
        d_lambda: f32,
        /// Maximum iterations for bisection refinement
        max_bisection_iterations: usize,
    },

    /// Direct analytical solutions for primitive geometric shapes.
    ///
    /// Computes exact ray-surface intersections using closed-form mathematical
    /// formulas. This is the fastest and most accurate method when applicable,
    /// but is limited to simple primitives with known intersection equations.
    ///
    /// # Supported Primitives
    ///
    /// - Spheres: Quadratic equation solutions
    /// - Planes: Linear ray-plane intersection
    /// - Cylinders: Analytical cylinder intersection
    /// - Cones: Quadratic cone intersection formulas
    ///
    /// # Limitations
    ///
    /// Cannot be used with:
    /// - Arbitrary scalar field expressions
    /// - CSG operations beyond simple unions
    /// - Transformed or deformed primitives
    Analytical,

    /// Iterative Newton-Raphson root finding using automatic differentiation.
    ///
    /// Uses gradient information to rapidly converge to surface intersections.
    /// Exhibits quadratic convergence near solutions but requires smooth,
    /// differentiable scalar fields for optimal performance.
    ///
    /// # Parameters
    ///
    /// * `max_iterations` - Maximum Newton iterations (typically 5-20)
    /// * `nudge_distance` - Perturbation for degenerate gradient cases
    /// * `step_size` - Relaxation factor for Newton steps (usually 0.5-1.0)
    ///
    /// # Convergence Properties
    ///
    /// - **Quadratic**: Error ∝ ε² per iteration near solutions
    /// - **Sensitive**: May fail near sharp edges or gradient discontinuities
    /// - **Fast**: Typically converges in 3-10 iterations for smooth surfaces
    Newton {
        /// Maximum number of Newton iterations
        max_iterations: usize,
        /// Distance to perturb points with zero gradients
        nudge_distance: f32,
        /// Step size relaxation factor (0.0 to 1.0)
        step_size: f32,
    },
}

impl<const N: usize> RayCastAlgorithm<N> {
    /// Computes ray-surface intersections using the specified algorithm.
    ///
    /// This method finds the points where rays intersect with the boundary of an implicit
    /// surface region. It supports multiple algorithmic approaches, each with different
    /// trade-offs between accuracy, performance, and applicability.
    ///
    /// # Mathematical Problem
    ///
    /// Given rays **r**(t) = **o** + t**d** and a region R with characteristic function f(x),
    /// find parameter values t where f(**r**(t)) = 0, indicating surface intersection.
    ///
    /// # Algorithm Variants
    ///
    /// * **Analytical**: Direct closed-form solutions for primitive shapes (spheres, planes, etc.)
    ///   - Fastest and most accurate when applicable
    ///   - Limited to simple geometric primitives with known ray intersection formulas
    ///
    /// * **Newton**: Iterative root-finding using gradient-based optimization
    ///   - Quadratic convergence near solutions
    ///   - Requires autodifferentiable scalar fields
    ///   - Best for smooth, well-conditioned surfaces
    ///
    /// * **BracketAndBisect**: Robust bracketing followed by bisection refinement
    ///   - Guaranteed convergence for continuous functions
    ///   - Works with any scalar field implementation
    ///   - Slower but most reliable for complex CSG constructions
    ///
    /// # Parameters
    ///
    /// * `rays` - Collection of rays to cast against the surface
    /// * `region` - Implicit surface region to intersect with
    /// * `algebra` - CSG algebra for evaluating boolean operations in composite regions
    ///
    /// # Returns
    ///
    /// [`RayCastResult`] containing intersection points, distances, field values, and metadata
    /// for each input ray. Rays that miss the surface are marked with sentinel values.
    ///
    /// # Performance Notes
    ///
    /// - **Analytical**: O(1) per ray for primitive surfaces
    /// - **Newton**: O(k) where k is iteration count (typically 3-10)
    /// - **BracketAndBisect**: O(log ε⁻¹) where ε is desired precision
    ///
    /// For complex CSG regions, cost scales with the number of constituent half-spaces.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use crater::analysis::prelude::*;
    /// use crater::csg::prelude::*;
    /// use burn::prelude::*;
    /// use burn::backend::ndarray::{NdArrayDevice, NdArray};
    /// use burn::backend::Autodiff;
    ///
    /// // Create a sphere region
    /// let sphere = Field3D::<Autodiff<NdArray>>::sphere(1.0, NdArrayDevice::Cpu).into_isosurface(0.0);
    /// let region = Region::HalfSpace(sphere, Side::Negative);
    /// let algebra = Algebra::<Autodiff<NdArray>>::default();
    ///
    /// // Create rays pointing at the sphere
    /// let origins = Tensor::from_data([[-2.0, 0.0, 0.0], [0.0, -2.0, 0.0]], &NdArrayDevice::Cpu);
    /// let directions = Tensor::from_data([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], &NdArrayDevice::Cpu);
    /// let rays = Rays::<Autodiff<NdArray>, 3>::new(origins, directions);
    ///
    /// // Use analytical algorithm for primitive shapes
    /// let algorithm = RayCastAlgorithm::<3>::Analytical;
    /// // let result = algorithm.solve(rays, &region, &algebra);
    ///
    /// // Check intersection distances (should be ≈ 1.0 for unit sphere)
    /// // let distances = result.distances();
    /// ```
    ///
    /// ```rust
    /// use crater::analysis::prelude::*;
    /// use crater::csg::prelude::*;
    /// use burn::prelude::*;
    /// use burn::backend::wgpu::WgpuDevice;
    ///
    /// // Complex CSG region requiring numerical methods
    /// let sphere = Field3D::<burn::backend::wgpu::Wgpu>::sphere(1.0, WgpuDevice::default()).into_isosurface(0.0);
    /// let cube = Field3D::<burn::backend::wgpu::Wgpu>::sphere(1.5, WgpuDevice::default()).into_isosurface(0.0);
    ///
    /// let sphere_region = Region::HalfSpace(sphere, Side::Negative);
    /// let cube_region = Region::HalfSpace(cube, Side::Negative);
    /// let intersection = Region::Intersection(Box::new(sphere_region), Box::new(cube_region));
    ///
    /// // Use bracket-and-bisect for robust intersection with CSG
    /// let algorithm = RayCastAlgorithm::<3>::BracketAndBisect {
    ///     lambda: 10.0,
    ///     d_lambda: 0.01,
    ///     max_bisection_iterations: 50,
    /// };
    ///
    /// // let result = algorithm.solve(rays, &intersection, &algebra);
    /// ```
    pub fn solve<B: AutodiffBackend>(
        &self,
        rays: Rays<B, N>,
        region: &Region<N, B>,
        algebra: &impl CSGAlgebra<B>,
    ) -> RayCastResult<B, N> {
        match self {
            Self::BracketAndBisect { .. } => self.bracket_and_bisect(rays, region, algebra),
            Self::Analytical => self.analytical(rays, region, algebra),
            Self::Newton { .. } => self.newton(rays, region, algebra),
        }
    }

    fn analytical<B: Backend>(
        &self,
        rays: Rays<B, N>,
        region: &Region<N, B>,
        algebra: &impl CSGAlgebra<B>,
    ) -> RayCastResult<B, N> {
        info!("Analytical ray casting: {}", rays);
        // First, iterate over all halfspaces, and evaluate the ray field
        let halfspace_evaluations = region
            .halfspaces()
            .enumerate()
            .map(|(i, halfspace)| {
                info!("Evaluating halfspace {}: {:?}", i, halfspace);
                let isosurface = match halfspace {
                    Region::HalfSpace(surf, _) => surf,
                    _ => panic!(),
                };

                assert!(
                    isosurface.ray_field.is_some(),
                    "Cannot use analytical ray casting with non-analytic isosurfaces"
                );

                // Evaluate ray field at the isosurface
                let result = isosurface
                    .ray_field
                    .as_ref()
                    .unwrap()
                    .solve(rays.clone(), isosurface.constant.clone());

                result
            })
            .collect::<Vec<_>>();

        // Next, at least one of the halfspaces must be on the surface of the REGION as well
        let mut intersections =
            Tensor::<B, 2, Float>::ones([rays.ray_count(), N], &rays.device()) * RAY_FAIL_VALUE;
        let mut distance_to_hit =
            Tensor::<B, 2, Float>::ones([rays.ray_count(), 1], &rays.device()) * RAY_FAIL_VALUE;

        halfspace_evaluations.iter().for_each(|result| {
            // True if the ray hits the surface of the REGION at this point

            let is_on_mask = region.evaluate(result.extensions(), algebra).is_on_mask();

            // If no intersection with this halfspace, skip
            if !is_on_mask.clone().any().into_scalar() {
                return;
            }

            let is_closer = distance_to_hit
                .clone()
                .greater_equal(result.distances().clone().unsqueeze::<2>().transpose());

            // Create a mask for rays that hit the surface AND are closer than current hit
            let should_replace_mask = is_on_mask
                .clone()
                .int()
                .mul(is_closer.clone().squeeze(1).int())
                .bool();

            // If no rays should be replaced, skip
            if !should_replace_mask.clone().any().into_scalar() {
                return;
            }

            // Use scatter_replace with boolean masks to replace entire rows
            intersections = intersections
                .scatter_replace(should_replace_mask.clone(), result.extensions().clone());
            distance_to_hit = distance_to_hit.scatter_replace(
                should_replace_mask.clone(),
                result.distances().clone().unsqueeze::<2>().transpose(),
            );
        });
        RayCastResult::new(
            rays.clone(),
            intersections.clone(),
            region.evaluate(intersections.clone(), algebra),
            Tensor::<B, 1, Float>::zeros([rays.ray_count()], &rays.device()),
        )
    }

    fn newton<B: AutodiffBackend>(
        &self,
        rays: Rays<B, N>,
        region: &Region<N, B>,
        algebra: &impl CSGAlgebra<B>,
    ) -> RayCastResult<B, N> {
        if let Self::Newton {
            max_iterations,
            nudge_distance,
            step_size,
        } = self
        {
            let device = rays.device();
            let origins = rays.origins();
            let batch_size = origins.dims()[0];

            let zeros = Tensor::<B, 1>::zeros([batch_size], &device);
            // Start with a mask of all rays active
            let mut m_i = Tensor::<B, 1>::ones([batch_size], &device);
            // Start from ray origins
            let mut points = origins.clone();

            for _ in 0..*max_iterations {
                // Compute field values at the current points
                let field_values = region.evaluate(points.clone(), algebra);

                // For points that are determined to be ON the surface, we set
                // their mask positions to 0.
                let is_on_mask = field_values.is_on_mask();
                m_i = m_i.clone().mul(is_on_mask.clone().bool_not().float());

                // Then, if m_i == 1, it means we should take a gradient step
                // If m_i is all 0, then we can stop
                if m_i.clone().any().into_scalar() {
                    let mut directional_derivatives = region.directional_derivative(
                        Rays::new(points.clone(), rays.directions()),
                        algebra,
                    ); // [batch_size, 1]

                    // Degenerate indices are those where the directional derivative is 0
                    let degenerate_mask = directional_derivatives.clone().equal_elem(0.0);

                    // Set the directional derivative to 1 for degenerate indices
                    if degenerate_mask.clone().any().into_scalar() {
                        directional_derivatives = directional_derivatives
                            .clone()
                            .unsqueeze::<2>()
                            .transpose()
                            .scatter_replace(
                                degenerate_mask.clone(),
                                Tensor::ones([degenerate_mask.clone().dims()[0], 1], &device)
                                    .mul_scalar(*nudge_distance),
                            )
                            .squeeze(1);
                    }

                    let step_size = -field_values
                        .clone()
                        .inner()
                        .div(directional_derivatives)
                        .mul_scalar(*step_size)
                        .unsqueeze::<2>()
                        .transpose();

                    // Take a gradient step using directional derivative
                    let step = rays
                        .directions()
                        .clone()
                        .inner()
                        .mul(step_size.clone())
                        .mul(m_i.clone().inner().unsqueeze::<2>().transpose());

                    points = points.add(Tensor::<B, 2>::from_inner(step));
                } else {
                    println!("No rays are still active, stopping");
                    break;
                }
            }

            // Ensure that rays are not moving backwards
            let displacement = points.clone() - rays.origins().clone();
            let backwards_mask = displacement
                .clone()
                .mul(rays.directions().clone())
                .sum_dim(1)
                .squeeze::<1>(1)
                .lower_elem(0.0);

            if backwards_mask.clone().any().into_scalar() {
                points = points.scatter_replace(
                    backwards_mask.clone(),
                    Tensor::<B, 2>::ones([backwards_mask.dims()[0], N], &device) * RAY_FAIL_VALUE,
                );
            }

            // For all rays that never had m_i == 1, set their points to RAY_FAIL_VALUE
            let failed_ray_mask = m_i.clone().bool();
            if failed_ray_mask.clone().any().into_scalar() {
                points = points.scatter_replace(
                    failed_ray_mask.clone(),
                    Tensor::<B, 2>::ones([failed_ray_mask.dims()[0], N], &device) * RAY_FAIL_VALUE,
                );
            }

            RayCastResult::new(
                rays.clone(),
                points.clone(),
                region.evaluate(points, algebra),
                zeros,
            )
        } else {
            panic!("Expected GradientDescent variant")
        }
    }

    fn bracket_and_bisect<B: Backend>(
        &self,
        rays: Rays<B, N>,
        region: &Region<N, B>,
        algebra: &impl CSGAlgebra<B>,
    ) -> RayCastResult<B, N> {
        if let Self::BracketAndBisect {
            lambda,
            d_lambda,
            max_bisection_iterations,
        } = self
        {
            let device = rays.device();
            let origins = rays.origins();
            let directions = rays.directions();
            let batch_size = origins.dims()[0];

            let zeros = Tensor::<B, 1>::zeros([batch_size], &device);
            let ones = Tensor::<B, 1>::ones([batch_size], &device);
            let lambda_0 = LAMBDA;

            let mut m_i = ones.clone();
            let mut l_i = zeros.clone();
            let l_0 = l_i.clone();
            let mut r_i = l_i.clone() + lambda_0;
            let mut f_l = region.evaluate(
                extend_rays::<B, N>(origins.clone(), directions.clone(), l_i.clone()),
                algebra,
            );

            // Bracketing phase:
            // Find intervals [l_i, r_i] along each ray that contain the isosurface
            // by advancing until the field value changes sign or we exceed max_distance
            loop {
                // Compute the right side of the bracket
                let f_r = region.evaluate(
                    extend_rays::<B, N>(origins.clone(), directions.clone(), r_i.clone()),
                    algebra,
                );

                // Update the mask to reflect those brackets that now bound a root
                m_i = m_i.clone().mul(
                    ones.clone()
                        .sub(contains_root_mask(f_l.clone().mul(f_r.clone()))),
                );

                // Check if we've exceeded max_distance for any active rays
                let exceeded_max = (r_i.clone() - l_0.clone()).greater_elem(*lambda);

                if exceeded_max.clone().any().into_scalar() {
                    // For rays that exceeded max_distance, mark them as not intersecting
                    m_i = m_i.clone().mul(ones.clone().sub(exceeded_max.float()));
                }

                // If no rays are still active, we can stop
                if !m_i.clone().any().into_scalar() {
                    break;
                }

                // Update the left and right sides of the bracket
                // Step fixed distance
                let step_size = m_i.clone().mul_scalar(*d_lambda);
                l_i = l_i.clone().add(step_size.clone());
                r_i = r_i.clone().add(step_size.clone());
                f_l = f_r;
            }

            // Bisection phase:
            // Refine the brackets to precisely locate the intersection points
            // using the bisection method for root finding
            (0..*max_bisection_iterations).for_each(|_| {
                // Compute the midpoint of the bracket
                let c_i = l_i.clone().add(r_i.clone()).div_scalar(2.0);
                // Evaluate the field at left, center, and right points
                let f_c = region.evaluate(
                    extend_rays::<B, N>(origins.clone(), directions.clone(), c_i.clone()),
                    algebra,
                );
                let f_l = region.evaluate(
                    extend_rays::<B, N>(origins.clone(), directions.clone(), l_i.clone()),
                    algebra,
                );
                // Create masks for brackets:
                // - is_bracketed_left: true if the root is in [l_i, c_i]
                // - is_bracketed_right: true if the root is in [c_i, r_i]
                // Note: there is a chance that the root is in both brackets (i.e., c_i = 0), but we
                // just break ties by assuming anything that is bracketed left must also not be bracketed right
                let is_bracketed_left = contains_root_mask(f_c.clone().mul(f_l.clone()));
                let is_bracketed_right = ones.clone() - is_bracketed_left.clone();

                // Update brackets based on which half contains the root
                l_i = c_i.clone().mul(ones.clone().sub(is_bracketed_left.clone()))
                    + l_i.clone().mul(is_bracketed_left.clone());
                r_i = c_i
                    .clone()
                    .mul(ones.clone().sub(is_bracketed_right.clone()))
                    + r_i.clone().mul(is_bracketed_right.clone());
            });
            let mut intersections =
                extend_rays::<B, N>(origins.clone(), directions.clone(), r_i.clone());

            // If the ray's origin is on the surface of the region, overwrite
            let is_on_mask = region.evaluate(origins.clone(), algebra).is_on_mask();

            if is_on_mask.clone().any().into_scalar() {
                intersections = intersections.scatter_replace(is_on_mask.clone(), origins.clone());
            }

            // If we did not find an extension that is on the surface of the region,
            // provide sentinel values
            let is_not_on_mask = region
                .evaluate(intersections.clone(), algebra)
                .is_on_mask()
                .bool_not();
            if is_not_on_mask.clone().any().into_scalar() {
                intersections = intersections.scatter_replace(
                    is_not_on_mask.clone(),
                    Tensor::<B, 2, Float>::ones([is_not_on_mask.clone().dims()[0], N], &device)
                        * RAY_FAIL_VALUE,
                );
            }

            // Return the final ray intersection points and their field values
            let final_field_values = region.evaluate(intersections.clone(), algebra);
            RayCastResult::new(rays, intersections, final_field_values, zeros)
        } else {
            panic!("Expected MarchAndBisect variant")
        }
    }
}

/// Min step distance in bracketing phase
/// Use non-rounded epsilon to avoid precision issues
const LAMBDA: f32 = f32::EPSILON * 1.0e6;

/// Extend a collection of rays by a vector of scalars
pub fn extend_rays<B: Backend, const N: usize>(
    positions: Origins<B, N>,
    directions: Directions<B, N>,
    lambda_i: Tensor<B, 1>,
) -> Origins<B, N> {
    // Compute the extended points
    positions.clone().add(
        lambda_i
            .clone()
            .unsqueeze::<2>()
            .transpose()
            .mul(directions),
    )
}

/// Creates a mask tensor where negative values become 1.0 and non-negative values become 0.0
fn contains_root_mask<B: Backend>(tensor: Tensor<B, 1>) -> Tensor<B, 1> {
    tensor.clone().lower_equal_elem(0.0).float()
}

impl<const N: usize, B: AutodiffBackend> Region<N, B> {
    /// Solve the ray casting problem for a given region and algebra
    pub fn ray_cast(
        &self,
        rays: Rays<B, N>,
        algebra: &impl CSGAlgebra<B>,
        algorithm: RayCastAlgorithm<N>,
    ) -> RayCastResult<B, N> {
        algorithm.solve(rays, self, algebra)
    }
    /// Returns the directional derivative along a collection of rays
    ///
    /// This is the scalar quantity:
    ///
    /// $$ \frac{\partial f}{\partial \mathbf{d}} = \nabla f \cdot \mathbf{d} $$
    pub fn directional_derivative(
        &self,
        rays: Rays<B, N>,
        algebra: &impl CSGAlgebra<B>,
    ) -> Scalars<B::InnerBackend> {
        let points = rays.origins();
        let grads = self.gradient(points, algebra);
        // Project onto the ray direction
        grads
            .mul(rays.directions().require_grad().inner())
            .sum_dim(1)
            .squeeze(1)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::primitives::prelude::*;
    use crate::test_utils::assert_tensor_almost_eq;
    use backend_macro::with_backend;
    use burn::tensor::Tensor;

    #[with_backend]
    #[test]
    fn test_directional_derivative_3d_sphere() {
        // Create a simple sphere field
        let sphere = Field3D::<Backend>::sphere(1.0, device()).into_isosurface(0.0);
        let region = sphere.region();
        let algebra = Algebra::default();

        // Create rays pointing in different directions
        let origins = Tensor::<Backend, 2>::from_data(
            [
                [0.0, 0.0, 0.0], // Origin at center
                [0.5, 0.0, 0.0], // Origin inside sphere
                [2.0, 0.0, 1.0], // Origin outside sphere
                [1.0, 0.0, 0.0], // Origin on surface
                [0.1, 0.1, 0.1], // Origin inside sphere
            ],
            &device(),
        );

        let directions = Tensor::<Backend, 2>::from_data(
            [
                [1.0, 0.0, 0.0],             // Pointing right
                [1.0, 0.0, 0.0],             // Pointing right
                [0.0, 0.0, 1.0],             // Pointing up
                [1.0, 0.0, 0.0],             // Pointing right
                normalize(&[1.0, 1.0, 1.0]), // Pointing along diagonal
            ],
            &device(),
        );
        let expected_derivatives =
            Tensor::<Backend, 1>::from_data([0.0, 1.0, 2.0, 2.0, 0.2 * 3.0_f32.sqrt()], &device());

        let rays = Rays::new(origins.clone(), directions.clone());

        // Compute directional derivatives
        let derivatives = region.directional_derivative(rays.clone(), &algebra);
        println!("derivatives: {}", derivatives);
        // Check that we get the expected shape
        assert_eq!(derivatives.dims(), [5]);

        // Check that the derivatives are correct
        assert_tensor_almost_eq(derivatives, expected_derivatives.inner(), Some(EPSILON));
    }
}