[FEATURE] Enhance SAP coupler for coupling between Rigid body and Fem object.

genesis/engine/bvh.py

@@ -277,8 +277,8 @@ def _kernel_radix_sort_morton_codes_one_round(self, i: int):
 
         # Reorder morton codes
         for i_b, i_a in ti.ndrange(self.n_batches, self.n_aabbs):
-            code = (self.morton_codes[i_b, i_a][1 - (i // 4)] >> ((i % 4) * 8)) & 0xFF
-            idx = ti.i32(self.offset[i_b, i_a] + self.prefix_sum[i_b, ti.i32(code)])
+            code = ti.i32((self.morton_codes[i_b, i_a][1 - (i // 4)] >> ((i % 4) * 8)) & 0xFF)
+            idx = ti.i32(self.offset[i_b, i_a] + self.prefix_sum[i_b, code])
             self.tmp_morton_codes[i_b, idx] = self.morton_codes[i_b, i_a]
 
         # Swap the temporary and original morton codes
@@ -299,7 +299,7 @@ def _kernel_radix_sort_morton_codes_one_round_group(self, i: int):
                 (i_g + 1) * self.group_size,
             )
             for i_a in range(start, end):
-                code = (self.morton_codes[i_b, i_a][1 - (i // 4)] >> ((i % 4) * 8)) & 0xFF
+                code = ti.i32((self.morton_codes[i_b, i_a][1 - (i // 4)] >> ((i % 4) * 8)) & 0xFF)
                 self.offset[i_b, i_a] = self.hist_group[i_b, i_g, code]
                 self.hist_group[i_b, i_g, code] = self.hist_group[i_b, i_g, code] + 1
 
@@ -319,7 +319,7 @@ def _kernel_radix_sort_morton_codes_one_round_group(self, i: int):
 
         # Reorder morton codes
         for i_b, i_a in ti.ndrange(self.n_batches, self.n_aabbs):
-            code = (self.morton_codes[i_b, i_a][1 - (i // 4)] >> ((i % 4) * 8)) & 0xFF
+            code = ti.i32((self.morton_codes[i_b, i_a][1 - (i // 4)] >> ((i % 4) * 8)) & 0xFF)
             i_g = ti.min(i_a // self.group_size, self.n_radix_sort_groups - 1)
             idx = ti.i32(self.prefix_sum[i_b, code] + self.prefix_sum_group[i_b, i_g, code] + self.offset[i_b, i_a])
             # Use the group prefix sum to find the correct index
@@ -441,14 +441,15 @@ def _kernel_compute_bounds_one_layer(self) -> ti.i32:
 
         return is_done
 
-    @ti.kernel
+    @ti.func
     def query(self, aabbs: ti.template()):
         """
         Query the BVH for intersections with the given AABBs.
 
         The results are stored in the query_result field.
         """
         self.query_result_count[None] = 0
+        overflow = False
 
         n_querys = aabbs.shape[1]
         for i_b, i_q in ti.ndrange(self.n_batches, n_querys):
@@ -470,6 +471,8 @@ def query(self, aabbs: ti.template()):
                         idx = ti.atomic_add(self.query_result_count[None], 1)
                         if idx < self.max_n_query_results:
                             self.query_result[idx] = gs.ti_ivec3(i_b, i_a, i_q)  # Store the AABB index
+                        else:
+                            overflow = True
                     else:
                         # Push children onto the stack
                         if node.right != -1:
@@ -479,12 +482,13 @@ def query(self, aabbs: ti.template()):
                             query_stack[stack_depth] = node.left
                             stack_depth += 1
 
+        return overflow
+
 
 @ti.data_oriented
 class FEMSurfaceTetLBVH(LBVH):
     """
     FEMSurfaceTetLBVH is a specialized Linear BVH for FEM surface tetrahedrals.
-
     It extends the LBVH class to support filtering based on FEM surface tetrahedral elements.
     """
 
@@ -495,14 +499,11 @@ def __init__(self, fem_solver, aabb: AABB, max_n_query_result_per_aabb: int = 8,
     @ti.func
     def filter(self, i_a, i_q):
         """
-        Filter function for FEM surface tets. Filter out tet that share vertices.
-
+        Filter function for FEM surface tets. Filter out tet that share vertices.\
         This is used to avoid self-collisions in FEM surface tets.
-
         i_a: index of the found AABB
         i_q: index of the query AABB
         """
-
         result = i_a >= i_q
         i_av = self.fem_solver.elements_i[self.fem_solver.surface_elements[i_a]].el2v
         i_qv = self.fem_solver.elements_i[self.fem_solver.surface_elements[i_q]].el2v