Implemented BVH

Implemented the BVH algorithm to allow for much faster collision tests at the cost of having to compile the mesh

Raytrace/BVHMesh.py

@@ -0,0 +1,178 @@
+from typing import Self, Literal
+import numpy as np
+from Raytrace.TriangleMesh import TriangleMesh, Ray
+
+# Adapted from https://jacco.ompf2.com/2022/04/13/how-to-build-a-bvh-part-1-basics/
+class BVHAABB:
+    min_pos:np.ndarray
+    max_pos:np.ndarray
+    def __init__(self):
+        self.min_pos = np.array([np.nan]*3)
+        self.max_pos = np.array([np.nan]*3)
+    @staticmethod
+    def from_array(triangles:np.ndarray) -> 'BVHAABB':
+        out = BVHAABB()
+        if triangles.size == 0: return out
+        triangles = triangles.reshape((-1, 3))
+        out.min_pos = triangles.min(axis = 0)
+        out.max_pos = triangles.max(axis = 0)
+        return out
+    def raytrace(self, ray:Ray) -> float:
+        old_error_state = np.seterr(divide='ignore')
+        t1 = (self.min_pos - ray.origin) * ray.inverse_direction
+        t2 = (self.max_pos - ray.origin) * ray.inverse_direction
+        np.seterr(**old_error_state)
+        t = np.stack([t1, t2])
+        tmin:float = t.min(axis=0).max()
+        tmax:float = t.max(axis=0).min()
+        return tmin if tmax >= tmin and tmax > 0 else np.inf
+    def grow(self, point:np.ndarray):
+        self.min_pos = np.nanmin(np.stack([point, self.min_pos]), axis=0)
+        self.max_pos = np.nanmax(np.stack([point, self.max_pos]), axis=0)
+    def grow_aabb(self, other:Self):
+        self.grow(other.min_pos)
+        self.grow(other.max_pos)
+    def area(self) -> float:
+        e = self.max_pos - self.min_pos
+        return e[0] * e[1] + e[1] * e[2] + e[2] * e[0]
+    def __str__(self):
+        return f'({float(self.min_pos[0])}, {float(self.min_pos[1])}, {float(self.min_pos[2])})#({float(self.max_pos[0])}, {float(self.max_pos[1])}, {float(self.max_pos[2])})'
+class BVHNode:
+    aabb:BVHAABB
+    left:Self
+    right:Self
+    start_index:int
+    tri_count:int
+    def get_subarray(self, triangles:np.ndarray) -> np.ndarray:
+        return triangles[self.start_index:self.start_index + self.tri_count]
+    def update_bounds(self, triangles:np.ndarray):
+        self.aabb = BVHAABB.from_array(self.get_subarray(triangles))
+    def is_leaf(self) -> bool:
+        return self.tri_count > 0
+class BVHMesh(TriangleMesh):
+    centroids:np.ndarray
+    root:BVHNode
+    node_count:int
+    def __init__(self, *args, min_node_size = 2, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self.build_BVH(min_node_size)
+    def build_BVH(self, min_node_size = 100) -> None:
+        # calculate triangle centroids for partitioning
+        self.centroids = self.triangles.sum(axis = 1) / 3
+        # assign all triangles to root node
+        self.node_count = 1
+        self.root = root = BVHNode()
+        root.start_index = 0
+        root.tri_count = self.triangles.shape[0]
+        if root.tri_count == 0: return
+        root.update_bounds(self.triangles)
+        # subdivide recursively
+        self.subdivide(root, min_node_size)
+    def subdivide(self,  node:BVHNode, min_node_size = 100) -> None:
+        # determine split axis using SAH
+        print(self.node_count, node.tri_count,end='   \r')
+        if node.tri_count < min_node_size: return
+        best_cost, axis, split_pos = self.find_best_split(node)
+        if best_cost >= node.tri_count * node.aabb.area(): return
+        # in-place partition
+        i = node.start_index
+        j = i + node.tri_count - 1
+        while i <= j:
+            if self.centroids[i, axis] < split_pos: i += 1
+            else:
+                self.triangles[[i,j]] = self.triangles[[j,i]]
+                self.centroids[[i,j]] = self.centroids[[j,i]]
+                j -= 1
+        # abort split if one of the sides is empty
+        left_count = i - node.start_index
+        if left_count == 0 or left_count == node.tri_count: return
+        self.node_count += 2
+        # create child nodes
+        node.left = left = BVHNode()
+        left.start_index = node.start_index
+        left.tri_count = left_count
+
+        node.right = right = BVHNode()
+        right.start_index = i
+        right.tri_count = node.tri_count - left_count
+
+        node.tri_count = 0
+        left.update_bounds(self.triangles)
+        right.update_bounds(self.triangles)
+        # recurse
+        self.subdivide(left)
+        self.subdivide(right)
+
+    def raytrace(self, ray:Ray) -> float:
+        return self.BVH_raytrace(ray)
+
+    def BVH_raytrace(self, ray:Ray) -> float:
+        node:BVHNode = self.root
+        stack:list[BVHNode] = []
+        best:float = np.inf
+        while 1:
+            if node.is_leaf():
+                intersections = self.batch_triangle_ray_intersection(node.get_subarray(self.triangles), ray)
+                intersections[intersections < 0] = np.inf
+                best = min(best, np.min(intersections))
+                if len(stack) == 0: break
+                node = stack.pop()
+                continue
+            child1, child2 = node.left, node.right
+            dist1 = node.left.aabb.raytrace(ray)
+            dist2 = node.right.aabb.raytrace(ray)
+            if dist1 > dist2:
+                dist1, dist2 = dist2, dist1
+                child1, child2 = child2, child1
+            if not np.isfinite(dist1):
+                if len(stack) == 0: break
+                node = stack.pop()
+            else:
+                node = child1
+                if np.isfinite(dist2):
+                    stack.append(child2)
+        return best
+    def find_best_split(self, node:BVHNode) -> tuple[float, int, float]:
+        BINS = 8
+        axis:int = -1
+        split_pos:float = 0
+        best_cost:float = np.inf
+        triangles:np.ndarray[tuple[int, Literal[3], Literal[3]], np.dtype[np.float32]] = node.get_subarray(self.triangles)
+        centroids:np.ndarray[tuple[int, Literal[3]], np.dtype[np.float32]] = node.get_subarray(self.centroids)
+        for axis_i in range(3):
+            bounds_min = float(np.min(centroids[:,axis_i]))
+            bounds_max = float(np.max(centroids[:,axis_i]))
+
+            if bounds_min == bounds_max: continue
+            # populate the bins
+            scale:float = BINS / (bounds_max - bounds_min)
+            bin_idx = ((centroids[:,axis_i] - bounds_min) * scale).astype(int)
+            bin_idx[bin_idx > BINS - 1] = BINS - 1
+            bin_bounds = [BVHAABB.from_array(triangles[bin_idx == n]) for n in range(BINS)]
+            bin_counts = [np.count_nonzero(bin_idx == n) for n in range(BINS)]
+            # gather data for the 7 planes between the 8 bins
+            left_area = np.zeros(BINS - 1, float)
+            right_area = np.zeros(BINS - 1, float)
+            left_count = np.zeros(BINS - 1, int)
+            right_count = np.zeros(BINS - 1, int)
+            left_aabb = BVHAABB()
+            right_aabb = BVHAABB()
+            left_sum = right_sum = 0
+            for i in range(BINS - 1):
+                left_sum += bin_counts[i]
+                left_count[i] = left_sum
+                left_aabb.grow_aabb(bin_bounds[i])
+                left_area[i] = left_aabb.area()
+                right_sum += bin_counts[BINS - 1 - i]
+                right_count[BINS - 2 - i] = right_sum
+                right_aabb.grow_aabb(bin_bounds[BINS - 1 - i])
+                right_area[BINS - 2 - i] = right_aabb.area()
+            # calculate SAH cost for the 7 planes
+            scale = 1. / scale
+            for i in range(BINS - 1):
+                cost:float = left_count[i] * left_area[i] + right_count[i] * right_area[i]
+                if cost < best_cost:
+                    split_pos = bounds_min + scale * (i + 1)
+                    axis = axis_i
+                    best_cost = cost
+        return best_cost, axis, split_pos