Implemented voxel-mesh intersection

xlb/operator/boundary_masker/mesh_boundary_masker.py

@@ -59,6 +59,78 @@ def index_to_position(index: wp.vec3i):
             pos = ijk + wp.vec3(0.5, 0.5, 0.5)  # cell center
             return pos
 
+        # Function to precompute useful values per triangle, assuming spacing is (1,1,1)
+        @wp.func
+        def pre_compute( v: wp.mat33f, # triangle vertices
+                        normal: wp.vec3f ): # triangle normal
+            c = wp.vec3f(float(normal[0] > 0.0),float(normal[1] > 0.0),float(normal[2] > 0.0))
+
+            d1 = wp.dot(normal, c-v[0])
+            d2 = wp.dot(normal, wp.vec3f(1.,1.,1.) - c - v[0])
+
+            edges = wp.transpose(wp.mat33(v[1]-v[0],v[2]-v[1],v[0]-v[2]))
+            ne0 = wp.mat33f(0.0)
+            ne1 = wp.mat33f(0.0)
+            de = wp.mat33f(0.0)
+
+            for ax0 in range(0,3):
+                ax2 = ( ax0 + 2 ) % 3
+
+                sgn = 1.0
+                if ( normal[ax2] < 0.0 ):
+                    sgn = -1.0
+
+                for i in range(0,3):
+                    ne0[i][ax0] = -1.0 * sgn * edges[i][ax0]
+                    ne1[i][ax0] = sgn * edges[i][ax0]
+
+                    de[i][ax0] = -1. * ( ne0[i][ax0] * v[i][ax0] + ne1[i][ax0] * v[i][ax0] ) \
+                        + wp.max(0., ne0[i][ax0] ) + wp.max(0., ne1[i][ax0])
+
+            return d1, d2, ne0, ne1, de
+
+        # Check whether this triangle intersects the unit cube at position low
+        @wp.func
+        def triangle_box_intersect( low: wp.vec3f,
+                                    normal: wp.vec3f,
+                                    d1: wp.float32,
+                                    d2: wp.float32,
+                                    ne0: wp.mat33f,
+                                    ne1: wp.mat33f,
+                                    de: wp.mat33f ):
+            if ( ( wp.dot(normal, low ) + d1 ) * ( wp.dot( normal, low ) + d2 ) <= 0.0 ):
+                intersect = True
+                #  Loop over primary axis for projection
+                for ax0 in range(0,3):
+                    ax1 = ( ax0 + 1 ) % 3
+                    for i in range(0,3):
+                        intersect = intersect and ( ne0[i][ax0] * low[ax0] + ne1[i][ax0] * low[ax1] + de[i][ax0] >= 0.0 )
+
+                return intersect
+            else:
+                return False
+
+        @wp.func
+        def mesh_voxel_intersect( mesh_id: wp.uint64, low: wp.vec3 ):
+
+            query = wp.mesh_query_aabb(mesh_id, low, low + wp.vec3f(1.,1.,1.))
+
+            for f in query:
+                v0 = wp.mesh_eval_position(mesh_id, f, 1.0, 0.0)
+                v1 = wp.mesh_eval_position(mesh_id, f, 0.0, 1.0)
+                v2 = wp.mesh_eval_position(mesh_id, f, 0.0, 0.0)
+                normal = wp.mesh_eval_face_normal(mesh_id, f)
+
+                v = wp.transpose(wp.mat33f(v0,v1,v2))
+
+                # TODO: run this on triangles in advance
+                d1, d2, ne0, ne1, de = pre_compute( v= v, normal= normal )
+
+                if triangle_box_intersect(low=low, normal=normal, d1=d1, d2=d2, ne0=ne0, ne1=ne1, de=de):
+                    return True
+
+            return False
+
         # Construct the warp kernel
         # Do voxelization mesh query (warp.mesh_query_aabb) to find solid voxels
         #  - this gives an approximate 1 voxel thick surface around mesh
@@ -79,20 +151,16 @@ def kernel(
             pos_bc_cell = index_to_position(index)
             half = wp.vec3(0.5, 0.5, 0.5)
 
-            vox_query = wp.mesh_query_aabb(mesh_id, pos_bc_cell - half, pos_bc_cell + half)
-            face = wp.int32(0)
-            if wp.mesh_query_aabb_next(vox_query, face):
+            if mesh_voxel_intersect(mesh_id=mesh_id, low=pos_bc_cell - half):
                 # Make solid voxel
                 bc_mask[0, index[0], index[1], index[2]] = wp.uint8(255)
             else:
                 # Find the fractional distance to the mesh in each direction
                 for l in range(1, _q):
                     _dir = wp.vec3f(_c_float[0, l], _c_float[1, l], _c_float[2, l])
 
-                    # Check to see if this neighbor is solid
-                    vox_query_dir = wp.mesh_query_aabb(mesh_id, pos_bc_cell + _dir - half, pos_bc_cell + _dir + half)
-                    face = wp.int32(0)
-                    if wp.mesh_query_aabb_next(vox_query_dir, face):
+                    # Check to see if this neighbor is solid - this is super inefficient TODO: make it way better
+                    if mesh_voxel_intersect(mesh_id=mesh_id, low=pos_bc_cell + _dir - half):
                         # We know we have a solid neighbor
                         # Set the boundary id and missing_mask
                         bc_mask[0, index[0], index[1], index[2]] = wp.uint8(id_number)