warp-lang 1.7.2rc1__py3-none-manylinux_2_34_aarch64.whl → 1.8.0__py3-none-manylinux_2_34_aarch64.whl

warp/sim/collide.py

@@ -232,11 +232,11 @@ def box_sdf_grad(upper: wp.vec3, p: wp.vec3):
     sz = wp.sign(p[2])
 
     # x projection
-    if qx > qy and qx > qz or qy == 0.0 and qz == 0.0:
+    if (qx > qy and qx > qz) or (qy == 0.0 and qz == 0.0):
         return wp.vec3(sx, 0.0, 0.0)
 
     # y projection
-    if qy > qx and qy > qz or qx == 0.0 and qz == 0.0:
+    if (qy > qx and qy > qz) or (qx == 0.0 and qz == 0.0):
         return wp.vec3(0.0, sy, 0.0)
 
     # z projection
@@ -341,9 +341,9 @@ def closest_point_box(upper: wp.vec3, point: wp.vec3):
         sy = wp.abs(wp.abs(point[1]) - upper[1])
         sz = wp.abs(wp.abs(point[2]) - upper[2])
         # return closest point on closest side, handle corner cases
-        if sx < sy and sx < sz or sy == 0.0 and sz == 0.0:
+        if (sx < sy and sx < sz) or (sy == 0.0 and sz == 0.0):
             x = wp.sign(point[0]) * upper[0]
-        elif sy < sx and sy < sz or sx == 0.0 and sz == 0.0:
+        elif (sy < sx and sy < sz) or (sx == 0.0 and sz == 0.0):
             y = wp.sign(point[1]) * upper[1]
         else:
             z = wp.sign(point[2]) * upper[2]
@@ -1516,11 +1516,8 @@ def handle_contact_pairs(
         diff = p_a_world - p_b_world
 
         # if the plane is infinite or the point is within the plane we fix the normal to prevent intersections
-        if (
-            geo_scale_b[0] == 0.0
-            and geo_scale_b[1] == 0.0
-            or wp.abs(query_b[0]) < geo_scale_b[0]
-            and wp.abs(query_b[2]) < geo_scale_b[1]
+        if (geo_scale_b[0] == 0.0 and geo_scale_b[1] == 0.0) or (
+            wp.abs(query_b[0]) < geo_scale_b[0] and wp.abs(query_b[2]) < geo_scale_b[1]
         ):
             normal = wp.transform_vector(X_ws_b, wp.vec3(0.0, 1.0, 0.0))
             distance = wp.dot(diff, normal)
@@ -1618,7 +1615,7 @@ def collide(
                 device=model.device,
             )
 
-        if model.shape_contact_pair_count or model.ground and model.shape_ground_contact_pair_count:
+        if model.shape_contact_pair_count or (model.ground and model.shape_ground_contact_pair_count):
             # clear old count
             model.rigid_contact_count.zero_()
 
@@ -1683,7 +1680,7 @@ def collide(
                 record_tape=False,
             )
 
-        if model.shape_contact_pair_count or model.ground and model.shape_ground_contact_pair_count:
+        if model.shape_contact_pair_count or (model.ground and model.shape_ground_contact_pair_count):
             if requires_grad:
                 model.rigid_contact_point0 = wp.empty_like(model.rigid_contact_point0)
                 model.rigid_contact_point1 = wp.empty_like(model.rigid_contact_point1)
@@ -1850,7 +1847,7 @@ def vertex_triangle_collision_detection_kernel(
     """
     This function applies discrete collision detection between vertices and triangles. It uses pre-allocated spaces to
     record the collision data. This collision detector works both ways, i.e., it records vertices' colliding triangles to
-    vertex_colliding_triangles, and records each triangles colliding vertices to triangle_colliding_vertices.
+    `vertex_colliding_triangles`, and records each triangles colliding vertices to `triangle_colliding_vertices`.
 
     This function assumes that all the vertices are on triangles, and can be indexed from the pos argument.
 
@@ -1870,7 +1867,8 @@ def vertex_triangle_collision_detection_kernel(
         vertex_colliding_triangles_offsets (array): where each vertex' collision buffer starts
         vertex_colliding_triangles_buffer_sizes (array): size of each vertex' collision buffer, will be modified if resizing is needed
         vertex_colliding_triangles_min_dist (array): each vertex' min distance to all (non-neighbor) triangles
-        triangle_colliding_vertices (array): positions of all the triangles' collision vertices
+        triangle_colliding_vertices (array): positions of all the triangles' collision vertices, every two elements
+            records the vertex index and a triangle index it collides to
         triangle_colliding_vertices_count (array): number of triangles each vertex collides
         triangle_colliding_vertices_offsets (array): where each triangle's collision buffer starts
         triangle_colliding_vertices_buffer_sizes (array): size of each triangle's collision buffer, will be modified if resizing is needed
@@ -1880,6 +1878,8 @@ def vertex_triangle_collision_detection_kernel(
 
     v_index = wp.tid()
     v = pos[v_index]
+    vertex_buffer_offset = vertex_colliding_triangles_offsets[v_index]
+    vertex_buffer_size = vertex_colliding_triangles_offsets[v_index + 1] - vertex_buffer_offset
 
     lower = wp.vec3(v[0] - query_radius, v[1] - query_radius, v[2] - query_radius)
     upper = wp.vec3(v[0] + query_radius, v[1] + query_radius, v[2] + query_radius)
@@ -1905,13 +1905,11 @@ def vertex_triangle_collision_detection_kernel(
         dist = wp.length(closest_p - v)
 
         if dist < query_radius:
-            vertex_buffer_offset = vertex_colliding_triangles_offsets[v_index]
-            vertex_buffer_size = vertex_colliding_triangles_offsets[v_index + 1] - vertex_buffer_offset
-
             # record v-f collision to vertex
             min_dis_to_tris = wp.min(min_dis_to_tris, dist)
             if vertex_num_collisions < vertex_buffer_size:
-                vertex_colliding_triangles[vertex_buffer_offset + vertex_num_collisions] = tri_index
+                vertex_colliding_triangles[2 * (vertex_buffer_offset + vertex_num_collisions)] = v_index
+                vertex_colliding_triangles[2 * (vertex_buffer_offset + vertex_num_collisions) + 1] = tri_index
             else:
                 resize_flags[VERTEX_COLLISION_BUFFER_OVERFLOW_INDEX] = 1
 
@@ -1932,6 +1930,94 @@ def vertex_triangle_collision_detection_kernel(
     vertex_colliding_triangles_min_dist[v_index] = min_dis_to_tris
 
 
+@wp.kernel
+def vertex_triangle_collision_detection_no_triangle_buffers_kernel(
+    query_radius: float,
+    bvh_id: wp.uint64,
+    pos: wp.array(dtype=wp.vec3),
+    tri_indices: wp.array(dtype=wp.int32, ndim=2),
+    vertex_colliding_triangles_offsets: wp.array(dtype=wp.int32),
+    vertex_colliding_triangles_buffer_sizes: wp.array(dtype=wp.int32),
+    # outputs
+    vertex_colliding_triangles: wp.array(dtype=wp.int32),
+    vertex_colliding_triangles_count: wp.array(dtype=wp.int32),
+    vertex_colliding_triangles_min_dist: wp.array(dtype=float),
+    triangle_colliding_vertices_min_dist: wp.array(dtype=float),
+    resize_flags: wp.array(dtype=wp.int32),
+):
+    """
+    This function applies discrete collision detection between vertices and triangles. It uses pre-allocated spaces to
+    record the collision data. Unlike `vertex_triangle_collision_detection_kernel`, this collision detection kernel
+    works only in one way, i.e., it only records vertices' colliding triangles to `vertex_colliding_triangles`.
+
+    This function assumes that all the vertices are on triangles, and can be indexed from the pos argument.
+
+    Note:
+
+        The collision date buffer is pre-allocated and cannot be changed during collision detection, therefore, the space
+        may not be enough. If the space is not enough to record all the collision information, the function will set a
+        certain element in resized_flag to be true. The user can reallocate the buffer based on vertex_colliding_triangles_count
+        and vertex_colliding_triangles_count.
+
+    Attributes:
+        bvh_id (int): the bvh id you want to collide with
+        query_radius (float): the contact radius. vertex-triangle pairs whose distance are less than this will get detected
+        pos (array): positions of all the vertices that make up triangles
+        vertex_colliding_triangles (array): flattened buffer of vertices' collision triangles, every two elements records
+            the vertex index and a triangle index it collides to
+        vertex_colliding_triangles_count (array): number of triangles each vertex collides
+        vertex_colliding_triangles_offsets (array): where each vertex' collision buffer starts
+        vertex_colliding_triangles_buffer_sizes (array): size of each vertex' collision buffer, will be modified if resizing is needed
+        vertex_colliding_triangles_min_dist (array): each vertex' min distance to all (non-neighbor) triangles
+        triangle_colliding_vertices_min_dist (array): each triangle's min distance to all (non-self) vertices
+        resized_flag (array): size == 3, (vertex_buffer_resize_required, triangle_buffer_resize_required, edge_buffer_resize_required)
+    """
+
+    v_index = wp.tid()
+    v = pos[v_index]
+    vertex_buffer_offset = vertex_colliding_triangles_offsets[v_index]
+    vertex_buffer_size = vertex_colliding_triangles_offsets[v_index + 1] - vertex_buffer_offset
+
+    lower = wp.vec3(v[0] - query_radius, v[1] - query_radius, v[2] - query_radius)
+    upper = wp.vec3(v[0] + query_radius, v[1] + query_radius, v[2] + query_radius)
+
+    query = wp.bvh_query_aabb(bvh_id, lower, upper)
+
+    tri_index = wp.int32(0)
+    vertex_num_collisions = wp.int32(0)
+    min_dis_to_tris = query_radius
+    while wp.bvh_query_next(query, tri_index):
+        t1 = tri_indices[tri_index, 0]
+        t2 = tri_indices[tri_index, 1]
+        t3 = tri_indices[tri_index, 2]
+        if vertex_adjacent_to_triangle(v_index, t1, t2, t3):
+            continue
+
+        u1 = pos[t1]
+        u2 = pos[t2]
+        u3 = pos[t3]
+
+        closest_p, bary, feature_type = triangle_closest_point(u1, u2, u3, v)
+
+        dist = wp.length(closest_p - v)
+
+        if dist < query_radius:
+            # record v-f collision to vertex
+            min_dis_to_tris = wp.min(min_dis_to_tris, dist)
+            if vertex_num_collisions < vertex_buffer_size:
+                vertex_colliding_triangles[2 * (vertex_buffer_offset + vertex_num_collisions)] = v_index
+                vertex_colliding_triangles[2 * (vertex_buffer_offset + vertex_num_collisions) + 1] = tri_index
+            else:
+                resize_flags[VERTEX_COLLISION_BUFFER_OVERFLOW_INDEX] = 1
+
+            vertex_num_collisions = vertex_num_collisions + 1
+
+            wp.atomic_min(triangle_colliding_vertices_min_dist, tri_index, dist)
+
+    vertex_colliding_triangles_count[v_index] = vertex_num_collisions
+    vertex_colliding_triangles_min_dist[v_index] = min_dis_to_tris
+
+
 @wp.kernel
 def edge_colliding_edges_detection_kernel(
     query_radius: float,
@@ -2001,7 +2087,8 @@ def edge_colliding_edges_detection_kernel(
             # record e-e collision to e0, and leave e1; e1 will detect this collision from its own thread
             min_dis_to_edges = wp.min(min_dis_to_edges, dist)
             if edge_num_collisions < edge_buffer_size:
-                edge_colliding_edges[edge_buffer_offset + edge_num_collisions] = colliding_edge_index
+                edge_colliding_edges[2 * (edge_buffer_offset + edge_num_collisions)] = e_index
+                edge_colliding_edges[2 * (edge_buffer_offset + edge_num_collisions) + 1] = colliding_edge_index
             else:
                 resize_flags[EDGE_COLLISION_BUFFER_OVERFLOW_INDEX] = 1
 
@@ -2068,6 +2155,9 @@ def triangle_triangle_collision_detection_kernel(
 
 @wp.struct
 class TriMeshCollisionInfo:
+    vertex_indices: wp.array(dtype=wp.int32)
+    # size: 2 x sum(vertex_colliding_triangles_buffer_sizes)
+    # every two elements records the vertex index and a triangle index it collides to
     vertex_colliding_triangles: wp.array(dtype=wp.int32)
     vertex_colliding_triangles_offsets: wp.array(dtype=wp.int32)
     vertex_colliding_triangles_buffer_sizes: wp.array(dtype=wp.int32)
@@ -2080,6 +2170,8 @@ class TriMeshCollisionInfo:
     triangle_colliding_vertices_count: wp.array(dtype=wp.int32)
     triangle_colliding_vertices_min_dist: wp.array(dtype=float)
 
+    # size: 2 x sum(edge_colliding_edges_buffer_sizes)
+    # every two elements records the edge index and an edge index it collides to
     edge_colliding_edges: wp.array(dtype=wp.int32)
     edge_colliding_edges_offsets: wp.array(dtype=wp.int32)
     edge_colliding_edges_buffer_sizes: wp.array(dtype=wp.int32)
@@ -2095,7 +2187,13 @@ def get_vertex_colliding_triangles_count(col_info: TriMeshCollisionInfo, v: int)
 @wp.func
 def get_vertex_colliding_triangles(col_info: TriMeshCollisionInfo, v: int, i_collision: int):
     offset = col_info.vertex_colliding_triangles_offsets[v]
-    return col_info.vertex_colliding_triangles[offset + i_collision]
+    return col_info.vertex_colliding_triangles[2 * (offset + i_collision) + 1]
+
+
+@wp.func
+def get_vertex_collision_buffer_vertex_index(col_info: TriMeshCollisionInfo, v: int, i_collision: int):
+    offset = col_info.vertex_colliding_triangles_offsets[v]
+    return col_info.vertex_colliding_triangles[2 * (offset + i_collision)]
 
 
 @wp.func
@@ -2119,13 +2217,20 @@ def get_edge_colliding_edges_count(col_info: TriMeshCollisionInfo, e: int):
 @wp.func
 def get_edge_colliding_edges(col_info: TriMeshCollisionInfo, e: int, i_collision: int):
     offset = col_info.edge_colliding_edges_offsets[e]
-    return col_info.edge_colliding_edges[offset + i_collision]
+    return col_info.edge_colliding_edges[2 * (offset + i_collision) + 1]
+
+
+@wp.func
+def get_edge_collision_buffer_edge_index(col_info: TriMeshCollisionInfo, e: int, i_collision: int):
+    offset = col_info.edge_colliding_edges_offsets[e]
+    return col_info.edge_colliding_edges[2 * (offset + i_collision)]
 
 
 class TriMeshCollisionDetector:
     def __init__(
         self,
         model: Model,
+        record_triangle_contacting_vertices=False,
         vertex_positions=None,
         vertex_collision_buffer_pre_alloc=8,
         vertex_collision_buffer_max_alloc=256,
@@ -2138,6 +2243,7 @@ class TriMeshCollisionDetector:
         edge_edge_parallel_epsilon=1e-5,
     ):
         self.model = model
+        self.record_triangle_contacting_vertices = record_triangle_contacting_vertices
         self.vertex_positions = model.particle_q if vertex_positions is None else vertex_positions
         self.device = model.device
         self.vertex_collision_buffer_pre_alloc = vertex_collision_buffer_pre_alloc
@@ -2166,7 +2272,9 @@ class TriMeshCollisionDetector:
 
         # vertex collision buffers
         self.vertex_colliding_triangles = wp.zeros(
-            shape=(model.particle_count * self.vertex_collision_buffer_pre_alloc,), dtype=wp.int32, device=self.device
+            shape=(2 * model.particle_count * self.vertex_collision_buffer_pre_alloc,),
+            dtype=wp.int32,
+            device=self.device,
         )
         self.vertex_colliding_triangles_count = wp.array(
             shape=(model.particle_count,), dtype=wp.int32, device=self.device
@@ -2187,26 +2295,39 @@ class TriMeshCollisionDetector:
             self.vertex_colliding_triangles_buffer_sizes, self.vertex_colliding_triangles_offsets
         )
 
-        # triangle collision buffers
-        self.triangle_colliding_vertices = wp.zeros(
-            shape=(model.tri_count * self.triangle_collision_buffer_pre_alloc,), dtype=wp.int32, device=self.device
-        )
-        self.triangle_colliding_vertices_count = wp.zeros(shape=(model.tri_count,), dtype=wp.int32, device=self.device)
-        self.triangle_colliding_vertices_buffer_sizes = wp.full(
-            shape=(model.tri_count,), value=self.triangle_collision_buffer_pre_alloc, dtype=wp.int32, device=self.device
-        )
-        self.triangle_colliding_vertices_min_dist = wp.array(shape=(model.tri_count,), dtype=float, device=self.device)
+        if record_triangle_contacting_vertices:
+            # triangle collision buffers
+            self.triangle_colliding_vertices = wp.zeros(
+                shape=(model.tri_count * self.triangle_collision_buffer_pre_alloc,), dtype=wp.int32, device=self.device
+            )
+            self.triangle_colliding_vertices_count = wp.zeros(
+                shape=(model.tri_count,), dtype=wp.int32, device=self.device
+            )
+            self.triangle_colliding_vertices_buffer_sizes = wp.full(
+                shape=(model.tri_count,),
+                value=self.triangle_collision_buffer_pre_alloc,
+                dtype=wp.int32,
+                device=self.device,
+            )
 
-        self.triangle_colliding_vertices_offsets = wp.array(
-            shape=(model.tri_count + 1,), dtype=wp.int32, device=self.device
-        )
-        self.compute_collision_buffer_offsets(
-            self.triangle_colliding_vertices_buffer_sizes, self.triangle_colliding_vertices_offsets
-        )
+            self.triangle_colliding_vertices_offsets = wp.array(
+                shape=(model.tri_count + 1,), dtype=wp.int32, device=self.device
+            )
+            self.compute_collision_buffer_offsets(
+                self.triangle_colliding_vertices_buffer_sizes, self.triangle_colliding_vertices_offsets
+            )
+        else:
+            self.triangle_colliding_vertices = None
+            self.triangle_colliding_vertices_count = None
+            self.triangle_colliding_vertices_buffer_sizes = None
+            self.triangle_colliding_vertices_offsets = None
+
+        # this is need regardless of whether we record triangle contacting vertices
+        self.triangle_colliding_vertices_min_dist = wp.array(shape=(model.tri_count,), dtype=float, device=self.device)
 
         # edge collision buffers
         self.edge_colliding_edges = wp.zeros(
-            shape=(model.edge_count * self.edge_collision_buffer_pre_alloc,), dtype=wp.int32, device=self.device
+            shape=(2 * model.edge_count * self.edge_collision_buffer_pre_alloc,), dtype=wp.int32, device=self.device
         )
         self.edge_colliding_edges_count = wp.zeros(shape=(model.edge_count,), dtype=wp.int32, device=self.device)
         self.edge_colliding_edges_buffer_sizes = wp.full(
@@ -2248,10 +2369,12 @@ class TriMeshCollisionDetector:
         collision_info.vertex_colliding_triangles_count = self.vertex_colliding_triangles_count
         collision_info.vertex_colliding_triangles_min_dist = self.vertex_colliding_triangles_min_dist
 
-        collision_info.triangle_colliding_vertices = self.triangle_colliding_vertices
-        collision_info.triangle_colliding_vertices_offsets = self.triangle_colliding_vertices_offsets
-        collision_info.triangle_colliding_vertices_buffer_sizes = self.triangle_colliding_vertices_buffer_sizes
-        collision_info.triangle_colliding_vertices_count = self.triangle_colliding_vertices_count
+        if self.record_triangle_contacting_vertices:
+            collision_info.triangle_colliding_vertices = self.triangle_colliding_vertices
+            collision_info.triangle_colliding_vertices_offsets = self.triangle_colliding_vertices_offsets
+            collision_info.triangle_colliding_vertices_buffer_sizes = self.triangle_colliding_vertices_buffer_sizes
+            collision_info.triangle_colliding_vertices_count = self.triangle_colliding_vertices_count
+
         collision_info.triangle_colliding_vertices_min_dist = self.triangle_colliding_vertices_min_dist
 
         collision_info.edge_colliding_edges = self.edge_colliding_edges
@@ -2272,6 +2395,31 @@ class TriMeshCollisionDetector:
 
         offsets.assign(offsets_np)
 
+    def rebuild(self, new_pos=None):
+        if new_pos is not None:
+            self.vertex_positions = new_pos
+
+        wp.launch(
+            kernel=compute_tri_aabbs,
+            inputs=[
+                self.vertex_positions,
+                self.model.tri_indices,
+            ],
+            outputs=[self.lower_bounds_tris, self.upper_bounds_tris],
+            dim=self.model.tri_count,
+            device=self.model.device,
+        )
+        self.bvh_tris = wp.Bvh(self.lower_bounds_tris, self.upper_bounds_tris)
+
+        wp.launch(
+            kernel=compute_edge_aabbs,
+            inputs=[self.vertex_positions, self.model.edge_indices],
+            outputs=[self.lower_bounds_edges, self.upper_bounds_edges],
+            dim=self.model.edge_count,
+            device=self.model.device,
+        )
+        self.bvh_edges = wp.Bvh(self.lower_bounds_edges, self.upper_bounds_edges)
+
     def refit(self, new_pos=None):
         if new_pos is not None:
             self.vertex_positions = new_pos
@@ -2298,46 +2446,72 @@ class TriMeshCollisionDetector:
         self.bvh_edges.refit()
 
     def vertex_triangle_collision_detection(self, query_radius):
-        wp.launch(
-            kernel=init_triangle_collision_data_kernel,
-            inputs=[
-                query_radius,
-            ],
-            outputs=[
-                self.triangle_colliding_vertices_count,
-                self.triangle_colliding_vertices_min_dist,
-                self.resize_flags,
-            ],
-            dim=self.model.tri_count,
-            device=self.model.device,
-        )
+        self.vertex_colliding_triangles.fill_(-1)
 
-        wp.launch(
-            kernel=vertex_triangle_collision_detection_kernel,
-            inputs=[
-                query_radius,
-                self.bvh_tris.id,
-                self.vertex_positions,
-                self.model.tri_indices,
-                self.vertex_colliding_triangles_offsets,
-                self.vertex_colliding_triangles_buffer_sizes,
-                self.triangle_colliding_vertices_offsets,
-                self.triangle_colliding_vertices_buffer_sizes,
-            ],
-            outputs=[
-                self.vertex_colliding_triangles,
-                self.vertex_colliding_triangles_count,
-                self.vertex_colliding_triangles_min_dist,
-                self.triangle_colliding_vertices,
-                self.triangle_colliding_vertices_count,
-                self.triangle_colliding_vertices_min_dist,
-                self.resize_flags,
-            ],
-            dim=self.model.particle_count,
-            device=self.model.device,
-        )
+        if self.record_triangle_contacting_vertices:
+            wp.launch(
+                kernel=init_triangle_collision_data_kernel,
+                inputs=[
+                    query_radius,
+                ],
+                outputs=[
+                    self.triangle_colliding_vertices_count,
+                    self.triangle_colliding_vertices_min_dist,
+                    self.resize_flags,
+                ],
+                dim=self.model.tri_count,
+                device=self.model.device,
+            )
+
+            wp.launch(
+                kernel=vertex_triangle_collision_detection_kernel,
+                inputs=[
+                    query_radius,
+                    self.bvh_tris.id,
+                    self.vertex_positions,
+                    self.model.tri_indices,
+                    self.vertex_colliding_triangles_offsets,
+                    self.vertex_colliding_triangles_buffer_sizes,
+                    self.triangle_colliding_vertices_offsets,
+                    self.triangle_colliding_vertices_buffer_sizes,
+                ],
+                outputs=[
+                    self.vertex_colliding_triangles,
+                    self.vertex_colliding_triangles_count,
+                    self.vertex_colliding_triangles_min_dist,
+                    self.triangle_colliding_vertices,
+                    self.triangle_colliding_vertices_count,
+                    self.triangle_colliding_vertices_min_dist,
+                    self.resize_flags,
+                ],
+                dim=self.model.particle_count,
+                device=self.model.device,
+            )
+        else:
+            self.triangle_colliding_vertices_min_dist.fill_(query_radius)
+            wp.launch(
+                kernel=vertex_triangle_collision_detection_no_triangle_buffers_kernel,
+                inputs=[
+                    query_radius,
+                    self.bvh_tris.id,
+                    self.vertex_positions,
+                    self.model.tri_indices,
+                    self.vertex_colliding_triangles_offsets,
+                    self.vertex_colliding_triangles_buffer_sizes,
+                ],
+                outputs=[
+                    self.vertex_colliding_triangles,
+                    self.vertex_colliding_triangles_count,
+                    self.vertex_colliding_triangles_min_dist,
+                    self.triangle_colliding_vertices_min_dist,
+                    self.resize_flags,
+                ],
+                dim=self.model.particle_count,
+                device=self.model.device,
+            )
 
     def edge_edge_collision_detection(self, query_radius):
+        self.edge_colliding_edges.fill_(-1)
         wp.launch(
             kernel=edge_colliding_edges_detection_kernel,
             inputs=[

warp/sim/graph_coloring.py

@@ -178,6 +178,13 @@ def color_trimesh(
         algorithm: the parameter passed to `color_graph`, see `color_graph`'s document
 
     """
+    if num_nodes == 0:
+        return []
+
+    if trimesh_edge_indices.shape[0] == 0:
+        # no edge, all the particle can have same color
+        return [np.arange(0, num_nodes, dtype=int)]
+
     if include_bending_energy:
         graph_edge_indices = construct_trimesh_graph_edges(trimesh_edge_indices, return_wp_array=True)
     else:
@@ -218,7 +225,7 @@ def color_graph(
         Ordered Greedy: Ton-That, Q. M., Kry, P. G., & Andrews, S. (2023). Parallel block Neo-Hookean XPBD using graph clustering. Computers & Graphics, 110, 1-10.
     """
     if num_nodes == 0:
-        return
+        return []
 
     particle_colors = wp.empty(shape=(num_nodes), dtype=wp.int32, device="cpu")
 

warp/sim/import_mjcf.py

@@ -13,11 +13,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import math
 import os
 import re
 import xml.etree.ElementTree as ET
-from typing import Union
 
 import numpy as np
 
@@ -30,7 +31,7 @@ def parse_mjcf(
     builder,
     xform=None,
     floating=False,
-    base_joint: Union[dict, str, None] = None,
+    base_joint: dict | str | None = None,
     density=1000.0,
     stiffness=100.0,
     damping=10.0,
@@ -447,7 +448,7 @@ def parse_mjcf(
 
         return shapes
 
-    def parse_body(body, parent, incoming_defaults: dict, childclass: str = None):
+    def parse_body(body, parent, incoming_defaults: dict, childclass: str | None = None):
         body_class = body.get("class")
         if body_class is None:
             body_class = childclass

warp/sim/import_usd.py

@@ -13,6 +13,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import re
 
 import numpy as np
@@ -810,17 +812,19 @@ def parse_usd(
     }
 
 
-def resolve_usd_from_url(url: str, target_folder_name: str = None, export_usda: bool = False):
-    """
-    Downloads a USD file from a URL and resolves all references to other USD files to be downloaded to the given target folder.
+def resolve_usd_from_url(url: str, target_folder_name: str | None = None, export_usda: bool = False) -> str:
+    """Download a USD file from a URL and resolves all references to other USD files to be downloaded to the given target folder.
 
     Args:
-        url (str): URL to the USD file.
-        target_folder_name (str): Target folder name. If None, a timestamped folder will be created in the current directory.
-        export_usda (bool): If True, converts each downloaded USD file to USDA and saves the additional USDA file in the target folder with the same base name as the original USD file.
+        url: URL to the USD file.
+        target_folder_name: Target folder name. If ``None``, a time-stamped
+          folder will be created in the current directory.
+        export_usda: If ``True``, converts each downloaded USD file to USDA and
+          saves the additional USDA file in the target folder with the same
+          base name as the original USD file.
 
     Returns:
-        str: File path to the downloaded USD file.
+        File path to the downloaded USD file.
     """
     import datetime
     import os

warp/sim/integrator.py

@@ -32,11 +32,14 @@ def integrate_particles(
     v_new: wp.array(dtype=wp.vec3),
 ):
     tid = wp.tid()
+    x0 = x[tid]
+    v0 = v[tid]
+
     if (particle_flags[tid] & PARTICLE_FLAG_ACTIVE) == 0:
+        x_new[tid] = x0
+        v_new[tid] = wp.vec3(0.0)
         return
 
-    x0 = x[tid]
-    v0 = v[tid]
     f0 = f[tid]
 
     inv_mass = w[tid]

warp/sim/integrator_euler.py

@@ -1781,7 +1781,7 @@ def eval_tetrahedral_forces(model: Model, state: State, control: Control, partic
 
 def eval_body_contact_forces(model: Model, state: State, particle_f: wp.array, friction_smoothing: float = 1.0):
     if model.rigid_contact_max and (
-        model.ground and model.shape_ground_contact_pair_count or model.shape_contact_pair_count
+        (model.ground and model.shape_ground_contact_pair_count) or model.shape_contact_pair_count
     ):
         wp.launch(
             kernel=eval_rigid_contacts,

warp/sim/integrator_featherstone.py

@@ -1825,7 +1825,7 @@ class FeatherstoneIntegrator(Integrator):
                 )
 
                 if model.rigid_contact_max and (
-                    model.ground and model.shape_ground_contact_pair_count or model.shape_contact_pair_count
+                    (model.ground and model.shape_ground_contact_pair_count) or model.shape_contact_pair_count
                 ):
                     wp.launch(
                         kernel=eval_rigid_contacts,