Improved BVH using surface area heuristic*

- Add custom error type
- Update more mesh based calculations to f32

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<BvhNode>,
@@ -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<Bvh, RoctreeError> {
         // 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.subdivide(0, solid)?;
-        bvh
+        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) {
+    pub fn subdivide(&mut self, node_idx: usize, solid: &StlSolid) -> Result<(), RoctreeError> {
-        let SplitPlane {
+        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(left_child, solid)?;
-        self.subdivide(right_child, solid);
+        self.subdivide(right_child, solid)?;
+
+        Ok(())
     }
 
-    /// Calculate the optimal split plane axis and position.
+    /// Get the best split plane based on the surface area heuristic
-    /// This function should use the surface area heuristic to select the optimal split
+    /// This procedure works by binning
-    fn get_split_plane(&self, node_idx: usize) -> SplitPlane {
+    fn get_best_split(
-        // Initial approach just splits halfway along the longest axis
+        &self,
-        // TODO: Improve this to use the surface area heuristic
+        node_idx: usize,
-        let axis = (0..3).fold(0, |acc, a| {
+        solid: &StlSolid,
-            if self.nodes[node_idx].aabb.extents[acc] < self.nodes[node_idx].aabb.extents[a] {
+    ) -> Result<BestSplitPlane, RoctreeError> {
-                a
+        // Initialize the best split plane information with bad values so we know if no plane gets
-            } else {
+        // assigned
-                acc
+        let mut best_cost = f64::INFINITY;
-            }
+        let mut best_pos = f64::NAN;
-        });
 
-        let pos = self.nodes[node_idx].aabb.extents[axis] / 2.0;
+        // Pick the longest axis as the splitting axis
-        SplitPlane { axis, pos }
+        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,
+        })
     }
 }

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"
+            ),
+        }
+    }
+}

src/lib.rs

@@ -1,4 +1,5 @@
 pub mod bvh;
+pub mod errors;
 pub mod octree;
 pub mod primitives;
 pub mod stl;

src/primitives.rs

@@ -1,37 +1,36 @@
-use std::collections::BTreeSet;
+use vecmath::vec3_sub;
 
-use itertools::Itertools;
+#[derive(Default, Clone)]
-use vecmath::{vec3_add, vec3_scale};
-
-#[derive(Default)]
 pub struct Box3 {
-    pub center: [f64; 3],
+    pub min_bd: [f32; 3],
-    pub extents: [f64; 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 {
+    pub fn new_from_bounds(min_bd: [f32; 3], max_bd: [f32; 3]) -> Self {
-        let mut center = [0.0; 3];
+        Box3 { min_bd, max_bd }
-        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 }
     }
 
     /// This function returns the 8 vertices of an axis-aligned bounding box.
-    pub fn get_vertices(&self) -> Vec<[f64; 3]> {
+    /// These are ordered
-        let half_extent = vec3_scale(self.extents, 0.5);
+    ///     7----6
-
+    /// |  /|   /|
-        // Define all eight corners relative to the center.
+    /// | 4----5
-        [-half_extent[0], half_extent[0]]
+    /// | | 3--| 2
-            .into_iter()
+    /// | |/   |/
-            .cartesian_product([-half_extent[1], half_extent[1]])
+    /// | 0----1
-            .cartesian_product([-half_extent[2], half_extent[2]])
+    pub fn get_vertices(&self) -> Vec<[f32; 3]> {
-            .map(|((x, y), z)| vec3_add(self.center, [x, y, z]))
+        vec![
-            .collect()
+            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.min_bd = [0.0; 3];
-            self.extents = [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
+        for vertex_idx in primitives.iter().flatten() {
-        let mut min_bd = [f32::INFINITY; 3];
+            self.grow(vertices[*vertex_idx]);
-        let mut max_bd = [f32::NEG_INFINITY; 3];
+        }
+    }
 
-        for idx in unique_vertex_idxs.into_iter() {
+    /// Compute the surface area of a box
-            for ((vertex_pos, min_pos), max_pos) in vertices[*idx]
+    pub fn surface_area(&self) -> f32 {
-                .iter()
+        let extents = vec3_sub(self.max_bd, self.min_bd);
-                .zip(min_bd.iter_mut())
+        extents[0] * extents[1] + extents[1] * extents[2] + extents[2] * extents[0]
-                .zip(max_bd.iter_mut())
+    }
-            {
+
-                *min_pos = vertex_pos.min(*min_pos);
+    /// Get the extents of the current box
-                *max_pos = vertex_pos.max(*max_pos);
+    #[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 v1: [f32; 3],
-    pub v2: [f64; 3],
+    pub v2: [f32; 3],
-    pub v3: [f64; 3],
+    pub v3: [f32; 3],
 }

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 start = f32::max(box_projection.0, tri_projection.0);
-        let end = f64::min(box_projection.1, tri_projection.1);
+        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) {
+fn project_pnts(pnts: &[[f32; 3]], axis: [f32; 3]) -> (f32, f32) {
-    let mut min_proj = f64::INFINITY;
+    let mut min_proj = f32::INFINITY;
-    let mut max_proj = f64::NEG_INFINITY;
+    let mut max_proj = f32::NEG_INFINITY;
     for pnt in pnts {
         let projection = vec3_dot(axis, *pnt);
         min_proj = min_proj.min(projection);