From ba74db96f06d81fea3182aac0abc87ba0f4f0c7f Mon Sep 17 00:00:00 2001 From: Alex Selimov Date: Sat, 8 Mar 2025 23:09:15 -0500 Subject: [PATCH] Improved BVH using surface area heuristic* - Add custom error type - Update more mesh based calculations to f32 --- src/bvh.rs | 173 +++++++++++++++++++++++---- src/errors.rs | 19 +++ src/lib.rs | 1 + src/primitives.rs | 106 +++++++++------- src/utilities/intersection_checks.rs | 12 +- 5 files changed, 235 insertions(+), 76 deletions(-) create mode 100644 src/errors.rs diff --git a/src/bvh.rs b/src/bvh.rs index cd9ed0a..aa06a44 100644 --- a/src/bvh.rs +++ b/src/bvh.rs @@ -1,4 +1,17 @@ -use crate::{primitives::Box3, stl::parser::StlSolid}; +use core::f64; + +use crate::{errors::RoctreeError, primitives::Box3, stl::parser::StlSolid}; + +/// Number of potential split planes to check +const NBINS: usize = 8; + +/// Struct to hold bin information, this includes an axis-aligned bounding box as well as the count +/// of triangles with the +#[derive(Default, Clone)] +struct Bin { + aabb: Box3, + tri_count: usize, +} /// A node representing a node single in the bounding volume hierarchy #[derive(Default)] @@ -10,6 +23,16 @@ struct BvhNode { prim_count: usize, } +impl BvhNode { + /// This computes the computational cost associated with the specific node. + /// The cost is an aggregate that considers the size of the bounding box (how likely it will + /// intersect a ray) * the number of triangles (the number of ray-triangle checks we have to + /// perform if an intersect occurs) + fn calculate_node_cost(&self) -> f32 { + self.aabb.surface_area() * self.prim_count as f32 + } +} + /// Struct representing the total bounding volume hierarchy pub struct Bvh { nodes: Vec, @@ -17,14 +40,15 @@ pub struct Bvh { } /// Return type for get_subdivision_position_and_axis -struct SplitPlane { +struct BestSplitPlane { axis: usize, pos: f64, + cost: f64, } impl Bvh { /// Create a new Bounding volume hierarchy from an STL solid. - pub fn new(solid: &StlSolid) -> Self { + pub fn new(solid: &StlSolid) -> Result { // Initialize the root node let mut nodes = vec![BvhNode { prim_idx: 0, @@ -40,16 +64,22 @@ impl Bvh { tri_idx: Vec::from_iter(0..solid.triangles.len()), }; - bvh.subdivide(0, solid); - bvh + bvh.subdivide(0, solid)?; + Ok(bvh) } /// Subdivide a node, subdivision doesn't occur if the node has only 2 primitives - pub fn subdivide(&mut self, node_idx: usize, solid: &StlSolid) { - let SplitPlane { + pub fn subdivide(&mut self, node_idx: usize, solid: &StlSolid) -> Result<(), RoctreeError> { + let BestSplitPlane { + cost, axis: split_axis, pos: split_position, - } = self.get_split_plane(node_idx); + } = self.get_best_split(node_idx, solid)?; + + // If the split cost is greater than keeping the node intact just return + if cost > self.nodes[node_idx].calculate_node_cost() as f64 { + return Ok(()); + } let prim_count = self.nodes[node_idx].prim_count; let prim_idx = self.nodes[node_idx].prim_idx; @@ -74,7 +104,7 @@ impl Bvh { // If all primitives belong to either the left or right half exit because the parent node // is a leaf if left_count == 0 || left_count == prim_count { - return; + return Ok(()); } // Otherwise split the nodes into left and right child @@ -117,24 +147,117 @@ impl Bvh { self.nodes[node_idx].right_child = right_child; // Recurse to divide children - self.subdivide(left_child, solid); - self.subdivide(right_child, solid); + self.subdivide(left_child, solid)?; + self.subdivide(right_child, solid)?; + + Ok(()) } - /// Calculate the optimal split plane axis and position. - /// This function should use the surface area heuristic to select the optimal split - fn get_split_plane(&self, node_idx: usize) -> SplitPlane { - // Initial approach just splits halfway along the longest axis - // TODO: Improve this to use the surface area heuristic - let axis = (0..3).fold(0, |acc, a| { - if self.nodes[node_idx].aabb.extents[acc] < self.nodes[node_idx].aabb.extents[a] { - a - } else { - acc - } - }); + /// Get the best split plane based on the surface area heuristic + /// This procedure works by binning + fn get_best_split( + &self, + node_idx: usize, + solid: &StlSolid, + ) -> Result { + // Initialize the best split plane information with bad values so we know if no plane gets + // assigned + let mut best_cost = f64::INFINITY; + let mut best_pos = f64::NAN; - let pos = self.nodes[node_idx].aabb.extents[axis] / 2.0; - SplitPlane { axis, pos } + // Pick the longest axis as the splitting axis + let axis = self.nodes[node_idx] + .aabb + .get_extents() + .iter() + .enumerate() + .fold((0_usize, &f32::NEG_INFINITY), |acc, a| { + if acc.1 > a.1 { + acc + } else { + a + } + }) + .0; + + let node = &self.nodes[node_idx]; + // Get minimum and maximum bounds over the centroids off the triangles. This improves + // the binning compared to doing the binning over the vertices + let (min_bd, max_bd) = self.tri_idx[node.prim_idx..node.prim_idx + node.prim_count] + .iter() + .fold((f64::INFINITY, f64::NEG_INFINITY), |bounds, tri_idx| { + ( + bounds + .0 + .min(solid.triangle_centroids[*tri_idx][axis].into()), + bounds + .1 + .max(solid.triangle_centroids[*tri_idx][axis].into()), + ) + }); + if min_bd == max_bd { + return Err(RoctreeError::SplitPlaneError); + }; + + let mut bins = vec![Bin::default(); NBINS]; + let scale = NBINS as f64 / (max_bd - min_bd); + + // Assign the triangles to each of the bins + for tri_idx in self.tri_idx[node.prim_idx..node.prim_idx + node.prim_count].iter() { + let bin_index = usize::min( + NBINS - 1, + ((solid.triangle_centroids[*tri_idx][axis] as f64 - min_bd) * scale) as usize, + ); + + bins[bin_index].tri_count += 1; + bins[bin_index] + .aabb + .grow(solid.vertices[solid.triangles[*tri_idx][0]]); + bins[bin_index] + .aabb + .grow(solid.vertices[solid.triangles[*tri_idx][1]]); + bins[bin_index] + .aabb + .grow(solid.vertices[solid.triangles[*tri_idx][2]]); + } + + // Now get the left and right counts for each split plane (based on the bins) + let mut left_area = [0.0; NBINS - 1]; + let mut right_area = [0.0; NBINS - 1]; + let mut left_count = [0; NBINS - 1]; + let mut right_count = [0; NBINS - 1]; + let mut left_box = Box3::default(); + let mut right_box = Box3::default(); + let mut left_sum = 0; + let mut right_sum = 0; + for plane_idx in 0..NBINS - 1 { + // Calculate the left side of the plane + left_sum += bins[plane_idx].tri_count; + left_count[plane_idx] = left_sum; + left_box.grow_by_other_box(&bins[plane_idx].aabb); + left_area[plane_idx] = left_box.surface_area(); + + // Calculate the right side of the plane + right_sum += bins[NBINS - 1 - plane_idx].tri_count; + right_count[NBINS - 2 - plane_idx] = right_sum; + right_box.grow_by_other_box(&bins[NBINS - 1 - plane_idx].aabb); + right_area[NBINS - 2 - plane_idx] = right_box.surface_area(); + } + let scale = 1.0 / scale; + for plane_idx in 0..NBINS - 1 { + let plane_cost: f64 = (left_count[plane_idx] as f32 * left_area[plane_idx] + + right_count[plane_idx] as f32 * right_area[plane_idx]) + .into(); + if plane_cost < best_cost { + best_cost = plane_cost; + best_pos = scale * plane_idx as f64; + } + } + + Ok(BestSplitPlane { + axis, + pos: best_pos, + cost: best_cost, + }) } } diff --git a/src/errors.rs b/src/errors.rs new file mode 100644 index 0000000..140f012 --- /dev/null +++ b/src/errors.rs @@ -0,0 +1,19 @@ +use std::{error::Error, fmt::Display}; + +#[derive(Debug)] +pub enum RoctreeError { + SplitPlaneError, +} + +impl Error for RoctreeError {} + +impl Display for RoctreeError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + RoctreeError::SplitPlaneError => write!( + f, + "Failed to get split plane, centroids overlap along longest axis" + ), + } + } +} diff --git a/src/lib.rs b/src/lib.rs index d7364e1..d1da0a5 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,4 +1,5 @@ pub mod bvh; +pub mod errors; pub mod octree; pub mod primitives; pub mod stl; diff --git a/src/primitives.rs b/src/primitives.rs index 51779a5..0d0c57d 100644 --- a/src/primitives.rs +++ b/src/primitives.rs @@ -1,37 +1,36 @@ -use std::collections::BTreeSet; +use vecmath::vec3_sub; -use itertools::Itertools; -use vecmath::{vec3_add, vec3_scale}; - -#[derive(Default)] +#[derive(Default, Clone)] pub struct Box3 { - pub center: [f64; 3], - pub extents: [f64; 3], + pub min_bd: [f32; 3], + pub max_bd: [f32; 3], } impl Box3 { /// Create a new box from the bounding values - pub fn new_from_bounds(min_bound: [f32; 3], max_bound: [f32; 3]) -> Self { - let mut center = [0.0; 3]; - let mut extents = [0.0; 3]; - for (i, (min_val, max_val)) in min_bound.iter().zip(max_bound).enumerate() { - center[i] = ((min_val + max_val) / 2.0) as f64; - extents[i] = (max_val - min_val) as f64; - } - Box3 { center, extents } + pub fn new_from_bounds(min_bd: [f32; 3], max_bd: [f32; 3]) -> Self { + Box3 { min_bd, max_bd } } /// This function returns the 8 vertices of an axis-aligned bounding box. - pub fn get_vertices(&self) -> Vec<[f64; 3]> { - let half_extent = vec3_scale(self.extents, 0.5); - - // Define all eight corners relative to the center. - [-half_extent[0], half_extent[0]] - .into_iter() - .cartesian_product([-half_extent[1], half_extent[1]]) - .cartesian_product([-half_extent[2], half_extent[2]]) - .map(|((x, y), z)| vec3_add(self.center, [x, y, z])) - .collect() + /// These are ordered + /// 7----6 + /// | /| /| + /// | 4----5 + /// | | 3--| 2 + /// | |/ |/ + /// | 0----1 + pub fn get_vertices(&self) -> Vec<[f32; 3]> { + vec![ + self.min_bd, + [self.max_bd[0], self.min_bd[1], self.min_bd[2]], + [self.max_bd[0], self.max_bd[1], self.min_bd[2]], + [self.min_bd[0], self.max_bd[1], self.min_bd[2]], + [self.max_bd[0], self.min_bd[1], self.max_bd[2]], + [self.max_bd[0], self.max_bd[1], self.max_bd[2]], + self.max_bd, + [self.min_bd[0], self.max_bd[1], self.max_bd[2]], + ] } /// Expand a box to fit all primitive inside of it. @@ -43,31 +42,50 @@ impl Box3 { { // If the current primitive list is empty, zero bounding box and exit if primitives.is_empty() { - self.center = [0.0; 3]; - self.extents = [0.0; 3] + self.min_bd = [0.0; 3]; + self.max_bd = [0.0; 3]; + return; } - // Get all of the unique vertex indices so we don't have to check them multiple times - let unique_vertex_idxs: BTreeSet<&usize> = BTreeSet::from_iter(primitives.iter().flatten()); - // Now expand the box - let mut min_bd = [f32::INFINITY; 3]; - let mut max_bd = [f32::NEG_INFINITY; 3]; + for vertex_idx in primitives.iter().flatten() { + self.grow(vertices[*vertex_idx]); + } + } - for idx in unique_vertex_idxs.into_iter() { - for ((vertex_pos, min_pos), max_pos) in vertices[*idx] - .iter() - .zip(min_bd.iter_mut()) - .zip(max_bd.iter_mut()) - { - *min_pos = vertex_pos.min(*min_pos); - *max_pos = vertex_pos.max(*max_pos); - } + /// Compute the surface area of a box + pub fn surface_area(&self) -> f32 { + let extents = vec3_sub(self.max_bd, self.min_bd); + extents[0] * extents[1] + extents[1] * extents[2] + extents[2] * extents[0] + } + + /// Get the extents of the current box + #[inline(always)] + pub fn get_extents(&self) -> [f32; 3] { + vec3_sub(self.max_bd, self.min_bd) + } + + /// Grow the axis aligned bounding box to include a specific vertex + #[inline(always)] + pub fn grow(&mut self, vertex: [f32; 3]) { + for (axis, pos) in vertex.iter().enumerate() { + self.min_bd[axis] = self.min_bd[axis].min(*pos); + self.max_bd[axis] = self.max_bd[axis].min(*pos); + } + } + + /// Grow the axis aligned bounding box to include a min and max boundary + /// This is useful when trying to combine two bounding boxes + #[inline(always)] + pub fn grow_by_other_box(&mut self, other: &Self) { + for axis in 0..3 { + self.min_bd[axis] = self.min_bd[axis].min(other.min_bd[axis]); + self.max_bd[axis] = self.max_bd[axis].max(other.max_bd[axis]); } } } pub struct Triangle { - pub v1: [f64; 3], - pub v2: [f64; 3], - pub v3: [f64; 3], + pub v1: [f32; 3], + pub v2: [f32; 3], + pub v3: [f32; 3], } diff --git a/src/utilities/intersection_checks.rs b/src/utilities/intersection_checks.rs index 628f0e8..0439aa3 100644 --- a/src/utilities/intersection_checks.rs +++ b/src/utilities/intersection_checks.rs @@ -1,5 +1,3 @@ -use core::f64; - use vecmath::{vec3_cross, vec3_dot, vec3_sub}; use crate::primitives::{Box3, Triangle}; @@ -28,8 +26,8 @@ pub fn tri_aabb_intersect(bx: &Box3, tri: &Triangle) -> bool { let box_projection = project_pnts(&box_vertices, axis); let tri_projection = project_pnts(&[tri.v1, tri.v2, tri.v3], axis); - let start = f64::max(box_projection.0, tri_projection.0); - let end = f64::min(box_projection.1, tri_projection.1); + let start = f32::max(box_projection.0, tri_projection.0); + let end = f32::min(box_projection.1, tri_projection.1); if start < end { return false; } @@ -38,9 +36,9 @@ pub fn tri_aabb_intersect(bx: &Box3, tri: &Triangle) -> bool { } /// Project points unto an Axis and return the max and min -fn project_pnts(pnts: &[[f64; 3]], axis: [f64; 3]) -> (f64, f64) { - let mut min_proj = f64::INFINITY; - let mut max_proj = f64::NEG_INFINITY; +fn project_pnts(pnts: &[[f32; 3]], axis: [f32; 3]) -> (f32, f32) { + let mut min_proj = f32::INFINITY; + let mut max_proj = f32::NEG_INFINITY; for pnt in pnts { let projection = vec3_dot(axis, *pnt); min_proj = min_proj.min(projection);