warp-lang 1.4.1__py3-none-manylinux2014_aarch64.whl → 1.5.0__py3-none-manylinux2014_aarch64.whl

warp/sim/model.py

@@ -15,6 +15,7 @@ import numpy as np
 
 import warp as wp
 
+from .graph_coloring import ColoringAlgorithm, color_trimesh, combine_independent_particle_coloring
 from .inertia import (
     compute_box_inertia,
     compute_capsule_inertia,
@@ -891,7 +892,7 @@ class Model:
             target.soft_contact_body_pos = wp.zeros(count, dtype=wp.vec3, requires_grad=requires_grad)
             target.soft_contact_body_vel = wp.zeros(count, dtype=wp.vec3, requires_grad=requires_grad)
             target.soft_contact_normal = wp.zeros(count, dtype=wp.vec3, requires_grad=requires_grad)
-            target.soft_contact_tids = wp.zeros(count, dtype=int)
+            target.soft_contact_tids = wp.zeros(self.particle_count * (self.shape_count - 1), dtype=int)
 
     def allocate_soft_contacts(self, count, requires_grad=False):
         self._allocate_soft_contacts(self, count, requires_grad)
@@ -1134,6 +1135,8 @@ class ModelBuilder:
         self.particle_radius = []
         self.particle_flags = []
         self.particle_max_velocity = 1e5
+        # list of np.array
+        self.particle_coloring = []
 
         # shapes (each shape has an entry in these arrays)
         # transform from shape to body
@@ -1372,6 +1375,11 @@ class ModelBuilder:
         if builder.tet_count:
             self.tet_indices.extend((np.array(builder.tet_indices, dtype=np.int32) + start_particle_idx).tolist())
 
+        builder_coloring_translated = [group + start_particle_idx for group in builder.particle_coloring]
+        self.particle_coloring = combine_independent_particle_coloring(
+            self.particle_coloring, builder_coloring_translated
+        )
+
         start_body_idx = self.body_count
         start_shape_idx = self.shape_count
         for s, b in enumerate(builder.shape_body):
@@ -3421,7 +3429,12 @@ class ModelBuilder:
 
     # particles
     def add_particle(
-        self, pos: Vec3, vel: Vec3, mass: float, radius: float = None, flags: wp.uint32 = PARTICLE_FLAG_ACTIVE
+        self,
+        pos: Vec3,
+        vel: Vec3,
+        mass: float,
+        radius: float = None,
+        flags: wp.uint32 = PARTICLE_FLAG_ACTIVE,
     ) -> int:
         """Adds a single particle to the model
 
@@ -3446,7 +3459,9 @@ class ModelBuilder:
         self.particle_radius.append(radius)
         self.particle_flags.append(flags)
 
-        return len(self.particle_q) - 1
+        particle_id = self.particle_count - 1
+
+        return particle_id
 
     def add_spring(self, i: int, j, ke: float, kd: float, control: float):
         """Adds a spring between two particles in the system
@@ -3826,6 +3841,7 @@ class ModelBuilder:
         add_springs: bool = False,
         spring_ke: float = default_spring_ke,
         spring_kd: float = default_spring_kd,
+        particle_radius: float = default_particle_radius,
     ):
         """Helper to create a regular planar cloth grid
 
@@ -3846,7 +3862,6 @@ class ModelBuilder:
             fix_right: Make the right-most edge of particles kinematic
             fix_top: Make the top-most edge of particles kinematic
             fix_bottom: Make the bottom-most edge of particles kinematic
-
         """
 
         def grid_index(x, y, dim_x):
@@ -3876,7 +3891,7 @@ class ModelBuilder:
                     m = 0.0
                     particle_flag = wp.uint32(int(particle_flag) & ~int(PARTICLE_FLAG_ACTIVE))
 
-                self.add_particle(p, vel, m, flags=particle_flag)
+                self.add_particle(p, vel, m, flags=particle_flag, radius=particle_radius)
 
                 if x > 0 and y > 0:
                     if reverse_winding:
@@ -3960,6 +3975,7 @@ class ModelBuilder:
         add_springs: bool = False,
         spring_ke: float = default_spring_ke,
         spring_kd: float = default_spring_kd,
+        particle_radius: float = default_particle_radius,
     ):
         """Helper to create a cloth model from a regular triangle mesh
 
@@ -3975,7 +3991,7 @@ class ModelBuilder:
             density: The density per-area of the mesh
             edge_callback: A user callback when an edge is created
             face_callback: A user callback when a face is created
-
+            particle_radius: The particle_radius which controls particle based collisions.
         Note:
 
             The mesh should be two manifold.
@@ -3989,7 +4005,7 @@ class ModelBuilder:
         for v in vertices:
             p = wp.quat_rotate(rot, v * scale) + pos
 
-            self.add_particle(p, vel, 0.0)
+            self.add_particle(p, vel, 0.0, radius=particle_radius)
 
         # triangles
         inds = start_vertex + np.array(indices)
@@ -4016,22 +4032,22 @@ class ModelBuilder:
 
         adj = wp.utils.MeshAdjacency(self.tri_indices[start_tri:end_tri], end_tri - start_tri)
 
-        edgeinds = np.fromiter(
+        edge_indices = np.fromiter(
             (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[:, 3],
-            edge_ke=[edge_ke] * len(edgeinds),
-            edge_kd=[edge_kd] * len(edgeinds),
+            edge_indices[:, 0],
+            edge_indices[:, 1],
+            edge_indices[:, 2],
+            edge_indices[:, 3],
+            edge_ke=[edge_ke] * len(edge_indices),
+            edge_kd=[edge_kd] * len(edge_indices),
         )
 
         if add_springs:
             spring_indices = set()
-            for i, j, k, l in edgeinds:
+            for i, j, k, l in edge_indices:
                 spring_indices.add((min(i, j), max(i, j)))
                 spring_indices.add((min(i, k), max(i, k)))
                 spring_indices.add((min(i, l), max(i, l)))
@@ -4060,7 +4076,7 @@ class ModelBuilder:
         radius_mean: float = default_particle_radius,
         radius_std: float = 0.0,
     ):
-        rng = np.random.default_rng()
+        rng = np.random.default_rng(42)
         for z in range(dim_z):
             for y in range(dim_y):
                 for x in range(dim_x):
@@ -4070,7 +4086,7 @@ class ModelBuilder:
                     p = wp.quat_rotate(rot, v) + pos + wp.vec3(rng.random(3) * jitter)
 
                     if radius_std > 0.0:
-                        r = radius_mean + np.random.randn() * radius_std
+                        r = radius_mean + rng.standard_normal() * radius_std
                     else:
                         r = radius_mean
                     self.add_particle(p, vel, m, r)
@@ -4356,6 +4372,61 @@ class ModelBuilder:
         for i in range(self.shape_count - 1):
             self.shape_collision_filter_pairs.add((i, ground_id))
 
+    def set_coloring(self, particle_coloring):
+        """
+        Set coloring information with user-provided coloring.
+
+        Args:
+            particle_coloring: A list of list or `np.array` with `dtype`=`int`. The length of the list is the number of colors
+             and each list or `np.array` contains the indices of vertices with this color.
+        """
+        particle_coloring = [
+            color_group if isinstance(color_group, np.ndarray) else np.array(color_group)
+            for color_group in particle_coloring
+        ]
+        self.particle_coloring = particle_coloring
+
+    def color(
+        self,
+        include_bending=False,
+        balance_colors=True,
+        target_max_min_color_ratio=1.1,
+        coloring_algorithm=ColoringAlgorithm.MCS,
+    ):
+        """
+        Run coloring algorithm to generate coloring information.
+
+        Args:
+            include_bending_energy: Whether to consider bending energy for trimeshes in the coloring process. If set to `True`, the generated
+                graph will contain all the edges connecting o1 and o2; otherwise, the graph will be equivalent to the trimesh.
+            balance_colors: Whether to apply the color balancing algorithm to balance the size of each color
+            target_max_min_color_ratio: the color balancing algorithm will stop when the ratio between the largest color and
+                the smallest color reaches this value
+            algorithm: Value should an enum type of ColoringAlgorithm, otherwise it will raise an error. ColoringAlgorithm.mcs means using the MCS coloring algorithm,
+                while ColoringAlgorithm.ordered_greedy means using the degree-ordered greedy algorithm. The MCS algorithm typically generates 30% to 50% fewer colors
+                compared to the ordered greedy algorithm, while maintaining the same linear complexity. Although MCS has a constant overhead that makes it about twice
+                as slow as the greedy algorithm, it produces significantly better coloring results. We recommend using MCS, especially if coloring is only part of the
+                preprocessing stage.e.
+
+        Note:
+
+            References to the coloring algorithm:
+            MCS: Pereira, F. M. Q., & Palsberg, J. (2005, November). Register allocation via coloring of chordal graphs. In Asian Symposium on Programming Languages and Systems (pp. 315-329). Berlin, Heidelberg: Springer Berlin Heidelberg.
+            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.
+
+        """
+        # ignore bending energy if it is too small
+        edge_indices = np.array(self.edge_indices)
+
+        self.particle_coloring = color_trimesh(
+            len(self.particle_q),
+            edge_indices,
+            include_bending,
+            algorithm=coloring_algorithm,
+            balance_colors=balance_colors,
+            target_max_min_color_ratio=target_max_min_color_ratio,
+        )
+
     def finalize(self, device=None, requires_grad=False) -> Model:
         """Convert this builder object to a concrete model for simulation.
 
@@ -4407,6 +4478,8 @@ class ModelBuilder:
             m.particle_max_radius = np.max(self.particle_radius) if len(self.particle_radius) > 0 else 0.0
             m.particle_max_velocity = self.particle_max_velocity
 
+            m.particle_coloring = [wp.array(group, dtype=int) for group in self.particle_coloring]
+
             # hash-grid for particle interactions
             m.particle_grid = wp.HashGrid(128, 128, 128)