warp-lang 1.3.3__py3-none-macosx_10_13_universal2.whl → 1.4.1__py3-none-macosx_10_13_universal2.whl

warp/sim/integrator_xpbd.py

@@ -2808,12 +2808,8 @@ class XPBDIntegrator(Integrator):
 
         with wp.ScopedTimer("simulate", False):
             if model.particle_count:
-                if requires_grad:
-                    particle_q = state_out.particle_q
-                    particle_qd = state_out.particle_qd
-                else:
-                    particle_q = state_out.particle_q
-                    particle_qd = state_out.particle_qd
+                particle_q = state_out.particle_q
+                particle_qd = state_out.particle_qd
 
                 self.particle_q_init = wp.clone(state_in.particle_q)
                 if self.enable_restitution:

warp/sim/model.py

@@ -527,8 +527,9 @@ class Model:
         tri_poses (array): Triangle element rest pose, shape [tri_count, 2, 2], float
         tri_activations (array): Triangle element activations, shape [tri_count], float
         tri_materials (array): Triangle element materials, shape [tri_count, 5], float
+        tri_areas (array): Triangle element rest areas, shape [tri_count], float
 
-        edge_indices (array): Bending edge indices, shape [edge_count*4], int
+        edge_indices (array): Bending edge indices, shape [edge_count*4], int, each row is [o0, o1, v1, v2], where v1, v2 are on the edge
         edge_rest_angle (array): Bending edge rest angle, shape [edge_count], float
         edge_bending_properties (array): Bending edge stiffness and damping parameters, shape [edge_count, 2], float
 
@@ -558,6 +559,7 @@ class Model:
         joint_type (array): Joint type, shape [joint_count], int
         joint_parent (array): Joint parent body indices, shape [joint_count], int
         joint_child (array): Joint child body indices, shape [joint_count], int
+        joint_ancestor (array): Maps from joint index to the index of the joint that has the current joint parent body as child (-1 if no such joint ancestor exists), shape [joint_count], int
         joint_X_p (array): Joint transform in parent frame, shape [joint_count, 7], float
         joint_X_c (array): Joint mass frame in child frame, shape [joint_count, 7], float
         joint_axis (array): Joint axis in child frame, shape [joint_axis_count, 3], float
@@ -639,6 +641,8 @@ class Model:
         joint_dof_count (int): Total number of velocity degrees of freedom of all joints in the system
         joint_coord_count (int): Total number of position degrees of freedom of all joints in the system
 
+        particle_coloring (list of array): The coloring of all the particles, used for VBD's Gauss-Seidel iteration.
+
         device (wp.Device): Device on which the Model was allocated
 
     Note:
@@ -695,6 +699,7 @@ class Model:
         self.tri_poses = None
         self.tri_activations = None
         self.tri_materials = None
+        self.tri_areas = None
 
         self.edge_indices = None
         self.edge_rest_angle = None
@@ -727,6 +732,7 @@ class Model:
         self.joint_type = None
         self.joint_parent = None
         self.joint_child = None
+        self.joint_ancestor = None
         self.joint_X_p = None
         self.joint_X_c = None
         self.joint_axis = None
@@ -806,6 +812,8 @@ class Model:
         self.joint_dof_count = 0
         self.joint_coord_count = 0
 
+        self.particle_coloring = []
+
         self.device = wp.get_device(device)
 
     def state(self, requires_grad=None) -> State:
@@ -1174,6 +1182,7 @@ class ModelBuilder:
         self.tri_poses = []
         self.tri_activations = []
         self.tri_materials = []
+        self.tri_areas = []
 
         # edges (bending)
         self.edge_indices = []
@@ -1395,9 +1404,8 @@ class ModelBuilder:
             self.joint_X_p.extend(joint_X_p)
             self.joint_q.extend(joint_q)
 
-            self.add_articulation()
-
             # offset the indices
+            self.articulation_start.extend([a + self.joint_count for a in builder.articulation_start])
             self.joint_parent.extend([p + self.joint_count if p != -1 else -1 for p in builder.joint_parent])
             self.joint_child.extend([c + self.joint_count for c in builder.joint_child])
 
@@ -1493,6 +1501,7 @@ class ModelBuilder:
             "tri_poses",
             "tri_activations",
             "tri_materials",
+            "tri_areas",
             "tet_poses",
             "tet_activations",
             "tet_materials",
@@ -3529,6 +3538,7 @@ class ModelBuilder:
             self.tri_poses.append(inv_D.tolist())
             self.tri_activations.append(0.0)
             self.tri_materials.append((tri_ke, tri_ka, tri_kd, tri_drag, tri_lift))
+            self.tri_areas.append(area)
             return area
 
     def add_triangles(
@@ -3617,7 +3627,9 @@ class ModelBuilder:
                 np.array(tri_lift)[valid_inds],
             )
         )
-        return areas.tolist()
+        areas = areas.tolist()
+        self.tri_areas.extend(areas)
+        return areas
 
     def add_tetrahedron(
         self, i: int, j: int, k: int, l: int, k_mu: float = 1.0e3, k_lambda: float = 1.0e3, k_damp: float = 0.0
@@ -3685,10 +3697,10 @@ class ModelBuilder:
         by the `model.tri_kb` parameter.
 
         Args:
-            i: The index of the first particle
-            j: The index of the second particle
-            k: The index of the third particle
-            l: The index of the fourth particle
+            i: The index of the first particle, i.e., opposite vertex 0
+            j: The index of the second particle, i.e., opposite vertex 1
+            k: The index of the third particle, i.e., vertex 0
+            l: The index of the fourth particle, i.e., vertex 1
             rest: The rest angle across the edge in radians, if not specified it will be computed
 
         Note:
@@ -3736,10 +3748,10 @@ class ModelBuilder:
         by the `model.tri_kb` parameter.
 
         Args:
-            i: The indices of the first particle
-            j: The indices of the second particle
-            k: The indices of the third particle
-            l: The indices of the fourth particle
+            i: The index of the first particle, i.e., opposite vertex 0
+            j: The index of the second particle, i.e., opposite vertex 1
+            k: The index of the third particle, i.e., vertex 0
+            l: The index of the fourth particle, i.e., vertex 1
             rest: The rest angles across the edges in radians, if not specified they will be computed
 
         Note:
@@ -3849,16 +3861,22 @@ class ModelBuilder:
                 p = wp.quat_rotate(rot, g) + pos
                 m = mass
 
+                particle_flag = PARTICLE_FLAG_ACTIVE
+
                 if x == 0 and fix_left:
                     m = 0.0
+                    particle_flag = wp.uint32(int(particle_flag) & ~int(PARTICLE_FLAG_ACTIVE))
                 elif x == dim_x and fix_right:
                     m = 0.0
+                    particle_flag = wp.uint32(int(particle_flag) & ~int(PARTICLE_FLAG_ACTIVE))
                 elif y == 0 and fix_bottom:
                     m = 0.0
+                    particle_flag = wp.uint32(int(particle_flag) & ~int(PARTICLE_FLAG_ACTIVE))
                 elif y == dim_y and fix_top:
                     m = 0.0
+                    particle_flag = wp.uint32(int(particle_flag) & ~int(PARTICLE_FLAG_ACTIVE))
 
-                self.add_particle(p, vel, m)
+                self.add_particle(p, vel, m, flags=particle_flag)
 
                 if x > 0 and y > 0:
                     if reverse_winding:
@@ -3902,22 +3920,20 @@ class ModelBuilder:
         spring_indices = set()
 
         for _k, e in adj.edges.items():
-            # skip open edges
-            if e.f0 == -1 or e.f1 == -1:
-                continue
-
             self.add_edge(
                 e.o0, e.o1, e.v0, e.v1, edge_ke=edge_ke, edge_kd=edge_kd
             )  # opposite 0, opposite 1, vertex 0, vertex 1
 
-            spring_indices.add((min(e.o0, e.o1), max(e.o0, e.o1)))
-            spring_indices.add((min(e.o0, e.v0), max(e.o0, e.v0)))
-            spring_indices.add((min(e.o0, e.v1), max(e.o0, e.v1)))
+            # skip constraints open edges
+            if e.f0 != -1 and e.f1 != -1:
+                spring_indices.add((min(e.o0, e.o1), max(e.o0, e.o1)))
+                spring_indices.add((min(e.o0, e.v0), max(e.o0, e.v0)))
+                spring_indices.add((min(e.o0, e.v1), max(e.o0, e.v1)))
 
-            spring_indices.add((min(e.o1, e.v0), max(e.o1, e.v0)))
-            spring_indices.add((min(e.o1, e.v1), max(e.o1, e.v1)))
+                spring_indices.add((min(e.o1, e.v0), max(e.o1, e.v0)))
+                spring_indices.add((min(e.o1, e.v1), max(e.o1, e.v1)))
 
-            spring_indices.add((min(e.v0, e.v1), max(e.v0, e.v1)))
+                spring_indices.add((min(e.v0, e.v1), max(e.v0, e.v1)))
 
         if add_springs:
             for i, j in spring_indices:
@@ -4001,14 +4017,14 @@ class ModelBuilder:
         adj = wp.utils.MeshAdjacency(self.tri_indices[start_tri:end_tri], end_tri - start_tri)
 
         edgeinds = np.fromiter(
-            (x for e in adj.edges.values() if e.f0 != -1 and e.f1 != -1 for x in (e.o0, e.o1, e.v0, e.v1)),
+            (x for e in adj.edges.values() for x in (e.o0, e.o1, e.v0, e.v1)),
             int,
         ).reshape(-1, 4)
         self.add_edges(
             edgeinds[:, 0],
             edgeinds[:, 1],
             edgeinds[:, 2],
-            edgeinds[:, 0],
+            edgeinds[:, 3],
             edge_ke=[edge_ke] * len(edgeinds),
             edge_kd=[edge_kd] * len(edgeinds),
         )
@@ -4459,6 +4475,7 @@ class ModelBuilder:
             m.tri_poses = wp.array(self.tri_poses, dtype=wp.mat22, requires_grad=requires_grad)
             m.tri_activations = wp.array(self.tri_activations, dtype=wp.float32, requires_grad=requires_grad)
             m.tri_materials = wp.array(self.tri_materials, dtype=wp.float32, requires_grad=requires_grad)
+            m.tri_areas = wp.array(self.tri_areas, dtype=wp.float32, requires_grad=requires_grad)
 
             # ---------------------
             # edges
@@ -4514,6 +4531,14 @@ class ModelBuilder:
             m.joint_q = wp.array(self.joint_q, dtype=wp.float32, requires_grad=requires_grad)
             m.joint_qd = wp.array(self.joint_qd, dtype=wp.float32, requires_grad=requires_grad)
             m.joint_name = self.joint_name
+            # compute joint ancestors
+            child_to_joint = {}
+            for i, child in enumerate(self.joint_child):
+                child_to_joint[child] = i
+            parent_joint = []
+            for parent in self.joint_parent:
+                parent_joint.append(child_to_joint.get(parent, -1))
+            m.joint_ancestor = wp.array(parent_joint, dtype=wp.int32)
 
             # dynamics properties
             m.joint_armature = wp.array(self.joint_armature, dtype=wp.float32, requires_grad=requires_grad)

warp/sparse.py

@@ -106,7 +106,7 @@ class BsrMatrix(Generic[_BlockType]):
             return
 
         BsrMatrix.__setattr__(
-            self, "_nnz_buf", wp.zeros(dtype=int, shape=(1,), device="cpu", pinned=self.device.is_cuda)
+            self, "_nnz_buf", wp.empty(dtype=int, shape=(1,), device="cpu", pinned=self.device.is_cuda)
         )
         if self.device.is_cuda:
             BsrMatrix.__setattr__(self, "_nnz_event", wp.Event(self.device))
@@ -524,7 +524,7 @@ def _bsr_assign_split_blocks(
     if dest_block >= dest_offsets[dest_row_count]:
         return
 
-    dest_row = wp.lower_bound(dest_offsets, dest_block + 1) - 1
+    dest_row = wp.lower_bound(dest_offsets, 0, dest_row_count + 1, dest_block + 1) - 1
     src_row = dest_row // row_factor
 
     dest_col_in_row = dest_block - dest_offsets[dest_row]
@@ -566,7 +566,7 @@ def _bsr_assign_merge_row_col(
         dest_rows[block] = -1  # invalid
         dest_cols[block] = -1
     else:
-        row = wp.lower_bound(src_offsets, block + 1) - 1
+        row = wp.lower_bound(src_offsets, 0, src_row_count + 1, block + 1) - 1
         dest_rows[block] = row // row_factor
         dest_cols[block] = src_columns[block] // col_factor
 
@@ -589,7 +589,7 @@ def _bsr_assign_merge_blocks(
     if src_block >= src_offsets[src_row_count]:
         return
 
-    src_row = wp.lower_bound(src_offsets, src_block + 1) - 1
+    src_row = wp.lower_bound(src_offsets, 0, src_row_count + 1, src_block + 1) - 1
     src_col = src_columns[src_block]
 
     dest_row = src_row // row_factor
@@ -828,7 +828,7 @@ def bsr_copy(
         block_type=block_type,
         device=A.device,
     )
-    bsr_assign(dest=copy, src=A)
+    bsr_assign(dest=copy, src=A, structure_only=structure_only)
     return copy
 
 
@@ -1190,7 +1190,7 @@ def _bsr_get_block_row(dest_offset: int, row_count: int, bsr_offsets: wp.array(d
     if i >= bsr_offsets[row_count]:
         rows[dest_offset + i] = -1  # invalid
     else:
-        row = wp.lower_bound(bsr_offsets, i + 1) - 1
+        row = wp.lower_bound(bsr_offsets, 0, row_count + 1, i + 1) - 1
         rows[dest_offset + i] = row
 
 
@@ -1461,13 +1461,14 @@ def _bsr_mm_compute_values(
     y_offsets: wp.array(dtype=int),
     y_columns: wp.array(dtype=int),
     y_values: wp.array(dtype=Any),
+    mm_row_count: int,
     mm_offsets: wp.array(dtype=int),
     mm_cols: wp.array(dtype=int),
     mm_values: wp.array(dtype=Any),
 ):
     mm_block = wp.tid()
 
-    row = wp.lower_bound(mm_offsets, mm_block + 1) - 1
+    row = wp.lower_bound(mm_offsets, 0, mm_row_count + 1, mm_block + 1) - 1
     col = mm_cols[mm_block]
 
     mm_val = mm_values.dtype(type(alpha)(0.0))
@@ -1759,6 +1760,7 @@ def bsr_mm(
             work_arrays._old_z_offsets if y == z else y.offsets,
             work_arrays._old_z_columns if y == z else y.columns,
             work_arrays._old_z_values if y == z else y.values,
+            z.nrow,
             z.offsets,
             z.columns,
             mm_values,