warp-lang 1.7.2__py3-none-macosx_10_13_universal2.whl → 1.8.0__py3-none-macosx_10_13_universal2.whl

warp/sim/model.py

@@ -19,7 +19,7 @@ from __future__ import annotations
 
 import copy
 import math
-from typing import List, Optional, Tuple
+from typing import List, Tuple
 
 import numpy as np
 
@@ -187,7 +187,7 @@ class SDF:
         return self.volume.id
 
     def __hash__(self):
-        return hash((self.volume.id))
+        return hash(self.volume.id)
 
 
 class Mesh:
@@ -219,7 +219,7 @@ class Mesh:
         com (Vec3): The center of mass of the body
     """
 
-    def __init__(self, vertices: List[Vec3], indices: List[int], compute_inertia=True, is_solid=True):
+    def __init__(self, vertices: list[Vec3], indices: list[int], compute_inertia=True, is_solid=True):
         """Construct a Mesh object from a triangle mesh
 
         The mesh center of mass and inertia tensor will automatically be
@@ -282,24 +282,24 @@ class State:
     """
 
     def __init__(self):
-        self.particle_q: Optional[wp.array] = None
+        self.particle_q: wp.array | None = None
         """Array of 3D particle positions with shape ``(particle_count,)`` and type :class:`vec3`."""
 
-        self.particle_qd: Optional[wp.array] = None
+        self.particle_qd: wp.array | None = None
         """Array of 3D particle velocities with shape ``(particle_count,)`` and type :class:`vec3`."""
 
-        self.particle_f: Optional[wp.array] = None
+        self.particle_f: wp.array | None = None
         """Array of 3D particle forces with shape ``(particle_count,)`` and type :class:`vec3`."""
 
-        self.body_q: Optional[wp.array] = None
+        self.body_q: wp.array | None = None
         """Array of body coordinates (7-dof transforms) in maximal coordinates with shape ``(body_count,)`` and type :class:`transform`."""
 
-        self.body_qd: Optional[wp.array] = None
+        self.body_qd: wp.array | None = None
         """Array of body velocities in maximal coordinates (first three entries represent angular velocity,
         last three entries represent linear velocity) with shape ``(body_count,)`` and type :class:`spatial_vector`.
         """
 
-        self.body_f: Optional[wp.array] = None
+        self.body_f: wp.array | None = None
         """Array of body forces in maximal coordinates (first three entries represent torque, last three
         entries represent linear force) with shape ``(body_count,)`` and type :class:`spatial_vector`.
 
@@ -309,10 +309,10 @@ class State:
             assumes the wrenches are measured w.r.t. world origin.
         """
 
-        self.joint_q: Optional[wp.array] = None
+        self.joint_q: wp.array | None = None
         """Array of generalized joint coordinates with shape ``(joint_coord_count,)`` and type ``float``."""
 
-        self.joint_qd: Optional[wp.array] = None
+        self.joint_qd: wp.array | None = None
         """Array of generalized joint velocities with shape ``(joint_dof_count,)`` and type ``float``."""
 
     def clear_forces(self) -> None:
@@ -372,7 +372,7 @@ class Control:
     should generally be created using the :func:`Model.control()` function.
     """
 
-    def __init__(self, model: Model = None):
+    def __init__(self, model: Model | None = None):
         if model:
             wp.utils.warn(
                 "Passing arguments to Control's __init__ is deprecated\n"
@@ -381,16 +381,16 @@ class Control:
                 stacklevel=2,
             )
 
-        self.joint_act: Optional[wp.array] = None
+        self.joint_act: wp.array | None = None
         """Array of joint control inputs with shape ``(joint_axis_count,)`` and type ``float``."""
 
-        self.tri_activations: Optional[wp.array] = None
+        self.tri_activations: wp.array | None = None
         """Array of triangle element activations with shape ``(tri_count,)`` and type ``float``."""
 
-        self.tet_activations: Optional[wp.array] = None
+        self.tet_activations: wp.array | None = None
         """Array of tetrahedral element activations with shape with shape ``(tet_count,) and type ``float``."""
 
-        self.muscle_activations: Optional[wp.array] = None
+        self.muscle_activations: wp.array | None = None
         """Array of muscle activations with shape ``(muscle_count,)`` and type ``float``."""
 
     def clear(self) -> None:
@@ -499,7 +499,7 @@ def compute_shape_mass(type, scale, src, density, is_solid, thickness):
             vertices = np.array(src.vertices) * np.array(scale)
             m, c, I, vol = compute_mesh_inertia(density, vertices, src.indices, is_solid, thickness)
             return m, c, I
-    raise ValueError("Unsupported shape type: {}".format(type))
+    raise ValueError(f"Unsupported shape type: {type}")
 
 
 class Model:
@@ -678,7 +678,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.
+        particle_color_groups (list of array): The coloring of all the particles, used for VBD's Gauss-Seidel iteration. Each array contains indices of particles sharing the same color.
+        particle_colors (array): Contains the color assignment for every particle
 
         device (wp.Device): Device on which the Model was allocated
 
@@ -858,7 +859,10 @@ class Model:
         self.joint_dof_count = 0
         self.joint_coord_count = 0
 
-        self.particle_coloring = []
+        # indices of particles sharing the same color
+        self.particle_color_groups = []
+        # the color of each particles
+        self.particle_colors = None
 
         self.device = wp.get_device(device)
 
@@ -1183,7 +1187,7 @@ class ModelBuilder:
         self.particle_flags = []
         self.particle_max_velocity = 1e5
         # list of np.array
-        self.particle_coloring = []
+        self.particle_color_groups = []
 
         # shapes (each shape has an entry in these arrays)
         # transform from shape to body
@@ -1433,9 +1437,9 @@ 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
+        builder_coloring_translated = [group + start_particle_idx for group in builder.particle_color_groups]
+        self.particle_color_groups = combine_independent_particle_coloring(
+            self.particle_color_groups, builder_coloring_translated
         )
 
         start_body_idx = self.body_count
@@ -1593,12 +1597,12 @@ class ModelBuilder:
     # register a rigid body and return its index.
     def add_body(
         self,
-        origin: Optional[Transform] = None,
+        origin: Transform | None = None,
         armature: float = 0.0,
-        com: Optional[Vec3] = None,
-        I_m: Optional[Mat33] = None,
+        com: Vec3 | None = None,
+        I_m: Mat33 | None = None,
         m: float = 0.0,
-        name: str = None,
+        name: str | None = None,
     ) -> int:
         """Adds a rigid body to the model.
 
@@ -1657,11 +1661,11 @@ class ModelBuilder:
         joint_type: wp.constant,
         parent: int,
         child: int,
-        linear_axes: Optional[List[JointAxis]] = None,
-        angular_axes: Optional[List[JointAxis]] = None,
-        name: str = None,
-        parent_xform: Optional[wp.transform] = None,
-        child_xform: Optional[wp.transform] = None,
+        linear_axes: list[JointAxis] | None = None,
+        angular_axes: list[JointAxis] | None = None,
+        name: str | None = None,
+        parent_xform: wp.transform | None = None,
+        child_xform: wp.transform | None = None,
         linear_compliance: float = 0.0,
         angular_compliance: float = 0.0,
         armature: float = 1e-2,
@@ -1797,21 +1801,21 @@ class ModelBuilder:
         self,
         parent: int,
         child: int,
-        parent_xform: Optional[wp.transform] = None,
-        child_xform: Optional[wp.transform] = None,
+        parent_xform: wp.transform | None = None,
+        child_xform: wp.transform | None = None,
         axis: Vec3 = (1.0, 0.0, 0.0),
-        target: float = None,
+        target: float | None = None,
         target_ke: float = 0.0,
         target_kd: float = 0.0,
         mode: int = JOINT_MODE_FORCE,
         limit_lower: float = -2 * math.pi,
         limit_upper: float = 2 * math.pi,
-        limit_ke: float = None,
-        limit_kd: float = None,
+        limit_ke: float | None = None,
+        limit_kd: float | None = None,
         linear_compliance: float = 0.0,
         angular_compliance: float = 0.0,
         armature: float = 1e-2,
-        name: str = None,
+        name: str | None = None,
         collision_filter_parent: bool = True,
         enabled: bool = True,
     ) -> int:
@@ -1887,21 +1891,21 @@ class ModelBuilder:
         self,
         parent: int,
         child: int,
-        parent_xform: Optional[wp.transform] = None,
-        child_xform: Optional[wp.transform] = None,
+        parent_xform: wp.transform | None = None,
+        child_xform: wp.transform | None = None,
         axis: Vec3 = (1.0, 0.0, 0.0),
-        target: float = None,
+        target: float | None = None,
         target_ke: float = 0.0,
         target_kd: float = 0.0,
         mode: int = JOINT_MODE_FORCE,
         limit_lower: float = -1e4,
         limit_upper: float = 1e4,
-        limit_ke: float = None,
-        limit_kd: float = None,
+        limit_ke: float | None = None,
+        limit_kd: float | None = None,
         linear_compliance: float = 0.0,
         angular_compliance: float = 0.0,
         armature: float = 1e-2,
-        name: str = None,
+        name: str | None = None,
         collision_filter_parent: bool = True,
         enabled: bool = True,
     ) -> int:
@@ -1977,12 +1981,12 @@ class ModelBuilder:
         self,
         parent: int,
         child: int,
-        parent_xform: Optional[wp.transform] = None,
-        child_xform: Optional[wp.transform] = None,
+        parent_xform: wp.transform | None = None,
+        child_xform: wp.transform | None = None,
         linear_compliance: float = 0.0,
         angular_compliance: float = 0.0,
         armature: float = 1e-2,
-        name: str = None,
+        name: str | None = None,
         collision_filter_parent: bool = True,
         enabled: bool = True,
     ) -> int:
@@ -2028,12 +2032,12 @@ class ModelBuilder:
         self,
         parent: int,
         child: int,
-        parent_xform: Optional[wp.transform] = None,
-        child_xform: Optional[wp.transform] = None,
+        parent_xform: wp.transform | None = None,
+        child_xform: wp.transform | None = None,
         linear_compliance: float = 0.0,
         angular_compliance: float = 0.0,
         armature: float = 1e-2,
-        name: str = None,
+        name: str | None = None,
         collision_filter_parent: bool = True,
         enabled: bool = True,
     ) -> int:
@@ -2079,11 +2083,11 @@ class ModelBuilder:
     def add_joint_free(
         self,
         child: int,
-        parent_xform: Optional[wp.transform] = None,
-        child_xform: Optional[wp.transform] = None,
+        parent_xform: wp.transform | None = None,
+        child_xform: wp.transform | None = None,
         armature: float = 0.0,
         parent: int = -1,
-        name: str = None,
+        name: str | None = None,
         collision_filter_parent: bool = True,
         enabled: bool = True,
     ) -> int:
@@ -2126,8 +2130,8 @@ class ModelBuilder:
         self,
         parent: int,
         child: int,
-        parent_xform: Optional[wp.transform] = None,
-        child_xform: Optional[wp.transform] = None,
+        parent_xform: wp.transform | None = None,
+        child_xform: wp.transform | None = None,
         min_distance: float = -1.0,
         max_distance: float = 1.0,
         compliance: float = 0.0,
@@ -2183,12 +2187,12 @@ class ModelBuilder:
         child: int,
         axis_0: JointAxis,
         axis_1: JointAxis,
-        parent_xform: Optional[wp.transform] = None,
-        child_xform: Optional[wp.transform] = None,
+        parent_xform: wp.transform | None = None,
+        child_xform: wp.transform | None = None,
         linear_compliance: float = 0.0,
         angular_compliance: float = 0.0,
         armature: float = 1e-2,
-        name: str = None,
+        name: str | None = None,
         collision_filter_parent: bool = True,
         enabled: bool = True,
     ) -> int:
@@ -2240,12 +2244,12 @@ class ModelBuilder:
         axis_0: JointAxis,
         axis_1: JointAxis,
         axis_2: JointAxis,
-        parent_xform: Optional[wp.transform] = None,
-        child_xform: Optional[wp.transform] = None,
+        parent_xform: wp.transform | None = None,
+        child_xform: wp.transform | None = None,
         linear_compliance: float = 0.0,
         angular_compliance: float = 0.0,
         armature: float = 1e-2,
-        name: str = None,
+        name: str | None = None,
         collision_filter_parent: bool = True,
         enabled: bool = True,
     ) -> int:
@@ -2298,11 +2302,11 @@ class ModelBuilder:
         self,
         parent: int,
         child: int,
-        linear_axes: Optional[List[JointAxis]] = None,
-        angular_axes: Optional[List[JointAxis]] = None,
-        name: str = None,
-        parent_xform: Optional[wp.transform] = None,
-        child_xform: Optional[wp.transform] = None,
+        linear_axes: list[JointAxis] | None = None,
+        angular_axes: list[JointAxis] | None = None,
+        name: str | None = None,
+        parent_xform: wp.transform | None = None,
+        child_xform: wp.transform | None = None,
         linear_compliance: float = 0.0,
         angular_compliance: float = 0.0,
         armature: float = 1e-2,
@@ -2425,7 +2429,7 @@ class ModelBuilder:
             return "unknown"
 
         if show_body_names:
-            vertices = ["world"] + self.body_name
+            vertices = ["world", *self.body_name]
         else:
             vertices = ["-1"] + [str(i) for i in range(self.body_count)]
         if plot_shapes:
@@ -2749,7 +2753,7 @@ class ModelBuilder:
 
     # muscles
     def add_muscle(
-        self, bodies: List[int], positions: List[Vec3], f0: float, lm: float, lt: float, lmax: float, pen: float
+        self, bodies: list[int], positions: list[Vec3], f0: float, lm: float, lt: float, lmax: float, pen: float
     ) -> float:
         """Adds a muscle-tendon activation unit.
 
@@ -2784,28 +2788,28 @@ class ModelBuilder:
     # shapes
     def add_shape_plane(
         self,
-        plane: Vec4 = (0.0, 1.0, 0.0, 0.0),
-        pos: Vec3 = None,
-        rot: Quat = None,
+        plane: Vec4 | tuple[float, float, float, float] = (0.0, 1.0, 0.0, 0.0),
+        pos: Vec3 | None = None,
+        rot: Quat | None = None,
         width: float = 10.0,
         length: float = 10.0,
         body: int = -1,
-        ke: float = None,
-        kd: float = None,
-        kf: float = None,
-        ka: float = None,
-        mu: float = None,
-        restitution: float = None,
-        thickness: float = None,
+        ke: float | None = None,
+        kd: float | None = None,
+        kf: float | None = None,
+        ka: float | None = None,
+        mu: float | None = None,
+        restitution: float | None = None,
+        thickness: float | None = None,
         has_ground_collision: bool = False,
         has_shape_collision: bool = True,
         is_visible: bool = True,
         collision_group: int = -1,
-    ):
-        """
-        Adds a plane collision shape.
-        If pos and rot are defined, the plane is assumed to have its normal as (0, 1, 0).
-        Otherwise, the plane equation defined through the `plane` argument is used.
+    ) -> int:
+        """Add a plane collision shape.
+
+        If ``pos`` and ``rot`` are defined, the plane is assumed to have its normal as (0, 1, 0).
+        Otherwise, the plane equation defined through the ``plane`` argument is used.
 
         Args:
             plane: The plane equation in form a*x + b*y + c*z + d = 0
@@ -2869,24 +2873,24 @@ class ModelBuilder:
     def add_shape_sphere(
         self,
         body,
-        pos: Vec3 = (0.0, 0.0, 0.0),
-        rot: Quat = (0.0, 0.0, 0.0, 1.0),
+        pos: Vec3 | tuple[float, float, float] = (0.0, 0.0, 0.0),
+        rot: Quat | tuple[float, float, float, float] = (0.0, 0.0, 0.0, 1.0),
         radius: float = 1.0,
-        density: float = None,
-        ke: float = None,
-        kd: float = None,
-        kf: float = None,
-        ka: float = None,
-        mu: float = None,
-        restitution: float = None,
+        density: float | None = None,
+        ke: float | None = None,
+        kd: float | None = None,
+        kf: float | None = None,
+        ka: float | None = None,
+        mu: float | None = None,
+        restitution: float | None = None,
         is_solid: bool = True,
-        thickness: float = None,
+        thickness: float | None = None,
         has_ground_collision: bool = True,
         has_shape_collision: bool = True,
         collision_group: int = -1,
         is_visible: bool = True,
-    ):
-        """Adds a sphere collision shape to a body.
+    ) -> int:
+        """Add a sphere collision shape to a body.
 
         Args:
             body: The index of the parent body this shape belongs to (use -1 for static shapes)
@@ -2938,26 +2942,26 @@ class ModelBuilder:
     def add_shape_box(
         self,
         body: int,
-        pos: Vec3 = (0.0, 0.0, 0.0),
-        rot: Quat = (0.0, 0.0, 0.0, 1.0),
+        pos: Vec3 | tuple[float, float, float] = (0.0, 0.0, 0.0),
+        rot: Quat | tuple[float, float, float, float] = (0.0, 0.0, 0.0, 1.0),
         hx: float = 0.5,
         hy: float = 0.5,
         hz: float = 0.5,
-        density: float = None,
-        ke: float = None,
-        kd: float = None,
-        kf: float = None,
-        ka: float = None,
-        mu: float = None,
-        restitution: float = None,
+        density: float | None = None,
+        ke: float | None = None,
+        kd: float | None = None,
+        kf: float | None = None,
+        ka: float | None = None,
+        mu: float | None = None,
+        restitution: float | None = None,
         is_solid: bool = True,
-        thickness: float = None,
+        thickness: float | None = None,
         has_ground_collision: bool = True,
         has_shape_collision: bool = True,
         collision_group: int = -1,
         is_visible: bool = True,
-    ):
-        """Adds a box collision shape to a body.
+    ) -> int:
+        """Add a box collision shape to a body.
 
         Args:
             body: The index of the parent body this shape belongs to (use -1 for static shapes)
@@ -2982,7 +2986,6 @@ class ModelBuilder:
 
         Returns:
             The index of the added shape
-
         """
 
         return self._add_shape(
@@ -3010,26 +3013,26 @@ class ModelBuilder:
     def add_shape_capsule(
         self,
         body: int,
-        pos: Vec3 = (0.0, 0.0, 0.0),
-        rot: Quat = (0.0, 0.0, 0.0, 1.0),
+        pos: Vec3 | tuple[float, float, float] = (0.0, 0.0, 0.0),
+        rot: Quat | tuple[float, float, float, float] = (0.0, 0.0, 0.0, 1.0),
         radius: float = 1.0,
         half_height: float = 0.5,
         up_axis: int = 1,
-        density: float = None,
-        ke: float = None,
-        kd: float = None,
-        kf: float = None,
-        ka: float = None,
-        mu: float = None,
-        restitution: float = None,
+        density: float | None = None,
+        ke: float | None = None,
+        kd: float | None = None,
+        kf: float | None = None,
+        ka: float | None = None,
+        mu: float | None = None,
+        restitution: float | None = None,
         is_solid: bool = True,
-        thickness: float = None,
+        thickness: float | None = None,
         has_ground_collision: bool = True,
         has_shape_collision: bool = True,
         collision_group: int = -1,
         is_visible: bool = True,
-    ):
-        """Adds a capsule collision shape to a body.
+    ) -> int:
+        """Add a capsule collision shape to a body.
 
         Args:
             body: The index of the parent body this shape belongs to (use -1 for static shapes)
@@ -3090,26 +3093,26 @@ class ModelBuilder:
     def add_shape_cylinder(
         self,
         body: int,
-        pos: Vec3 = (0.0, 0.0, 0.0),
-        rot: Quat = (0.0, 0.0, 0.0, 1.0),
+        pos: Vec3 | tuple[float, float, float] = (0.0, 0.0, 0.0),
+        rot: Quat | tuple[float, float, float, float] = (0.0, 0.0, 0.0, 1.0),
         radius: float = 1.0,
         half_height: float = 0.5,
         up_axis: int = 1,
-        density: float = None,
-        ke: float = None,
-        kd: float = None,
-        kf: float = None,
-        ka: float = None,
-        mu: float = None,
-        restitution: float = None,
+        density: float | None = None,
+        ke: float | None = None,
+        kd: float | None = None,
+        kf: float | None = None,
+        ka: float | None = None,
+        mu: float | None = None,
+        restitution: float | None = None,
         is_solid: bool = True,
-        thickness: float = None,
+        thickness: float | None = None,
         has_ground_collision: bool = True,
         has_shape_collision: bool = True,
         collision_group: int = -1,
         is_visible: bool = True,
-    ):
-        """Adds a cylinder collision shape to a body.
+    ) -> int:
+        """Add a cylinder collision shape to a body.
 
         Args:
             body: The index of the parent body this shape belongs to (use -1 for static shapes)
@@ -3172,26 +3175,26 @@ class ModelBuilder:
     def add_shape_cone(
         self,
         body: int,
-        pos: Vec3 = (0.0, 0.0, 0.0),
-        rot: Quat = (0.0, 0.0, 0.0, 1.0),
+        pos: Vec3 | tuple[float, float, float] = (0.0, 0.0, 0.0),
+        rot: Quat | tuple[float, float, float, float] = (0.0, 0.0, 0.0, 1.0),
         radius: float = 1.0,
         half_height: float = 0.5,
         up_axis: int = 1,
-        density: float = None,
-        ke: float = None,
-        kd: float = None,
-        kf: float = None,
-        ka: float = None,
-        mu: float = None,
-        restitution: float = None,
+        density: float | None = None,
+        ke: float | None = None,
+        kd: float | None = None,
+        kf: float | None = None,
+        ka: float | None = None,
+        mu: float | None = None,
+        restitution: float | None = None,
         is_solid: bool = True,
-        thickness: float = None,
+        thickness: float | None = None,
         has_ground_collision: bool = True,
         has_shape_collision: bool = True,
         collision_group: int = -1,
         is_visible: bool = True,
-    ):
-        """Adds a cone collision shape to a body.
+    ) -> int:
+        """Add a cone collision shape to a body.
 
         Args:
             body: The index of the parent body this shape belongs to (use -1 for static shapes)
@@ -3254,25 +3257,25 @@ class ModelBuilder:
     def add_shape_mesh(
         self,
         body: int,
-        pos: Optional[Vec3] = None,
-        rot: Optional[Quat] = None,
-        mesh: Optional[Mesh] = None,
-        scale: Optional[Vec3] = None,
-        density: float = None,
-        ke: float = None,
-        kd: float = None,
-        kf: float = None,
-        ka: float = None,
-        mu: float = None,
-        restitution: float = None,
+        pos: Vec3 | None = None,
+        rot: Quat | None = None,
+        mesh: Mesh | None = None,
+        scale: Vec3 | None = None,
+        density: float | None = None,
+        ke: float | None = None,
+        kd: float | None = None,
+        kf: float | None = None,
+        ka: float | None = None,
+        mu: float | None = None,
+        restitution: float | None = None,
         is_solid: bool = True,
-        thickness: float = None,
+        thickness: float | None = None,
         has_ground_collision: bool = True,
         has_shape_collision: bool = True,
         collision_group: int = -1,
         is_visible: bool = True,
-    ):
-        """Adds a triangle mesh collision shape to a body.
+    ) -> int:
+        """Add a triangle mesh collision shape to a body.
 
         Args:
             body: The index of the parent body this shape belongs to (use -1 for static shapes)
@@ -3335,25 +3338,25 @@ class ModelBuilder:
     def add_shape_sdf(
         self,
         body: int,
-        pos: Vec3 = (0.0, 0.0, 0.0),
-        rot: Quat = (0.0, 0.0, 0.0, 1.0),
-        sdf: SDF = None,
-        scale: Vec3 = (1.0, 1.0, 1.0),
-        density: float = None,
-        ke: float = None,
-        kd: float = None,
-        kf: float = None,
-        ka: float = None,
-        mu: float = None,
-        restitution: float = None,
+        pos: Vec3 | tuple[float, float, float] = (0.0, 0.0, 0.0),
+        rot: Quat | tuple[float, float, float, float] = (0.0, 0.0, 0.0, 1.0),
+        sdf: SDF | None = None,
+        scale: Vec3 | tuple[float, float, float] = (1.0, 1.0, 1.0),
+        density: float | None = None,
+        ke: float | None = None,
+        kd: float | None = None,
+        kf: float | None = None,
+        ka: float | None = None,
+        mu: float | None = None,
+        restitution: float | None = None,
         is_solid: bool = True,
-        thickness: float = None,
+        thickness: float | None = None,
         has_ground_collision: bool = True,
         has_shape_collision: bool = True,
         collision_group: int = -1,
         is_visible: bool = True,
-    ):
-        """Adds SDF collision shape to a body.
+    ) -> int:
+        """Add a SDF collision shape to a body.
 
         Args:
             body: The index of the parent body this shape belongs to (use -1 for static shapes)
@@ -3444,8 +3447,8 @@ class ModelBuilder:
         collision_filter_parent=True,
         has_ground_collision=True,
         has_shape_collision=True,
-        is_visible=True,
-    ):
+        is_visible: bool = True,
+    ) -> int:
         self.shape_body.append(body)
         shape = self.shape_count
         if body in self.body_shapes:
@@ -3504,7 +3507,7 @@ class ModelBuilder:
         pos: Vec3,
         vel: Vec3,
         mass: float,
-        radius: float = None,
+        radius: float | None = None,
         flags: wp.uint32 = PARTICLE_FLAG_ACTIVE,
     ) -> int:
         """Adds a single particle to the model
@@ -3569,11 +3572,11 @@ class ModelBuilder:
         i: int,
         j: int,
         k: int,
-        tri_ke: float = None,
-        tri_ka: float = None,
-        tri_kd: float = None,
-        tri_drag: float = None,
-        tri_lift: float = None,
+        tri_ke: float | None = None,
+        tri_ka: float | None = None,
+        tri_kd: float | None = None,
+        tri_drag: float | None = None,
+        tri_lift: float | None = None,
     ) -> float:
         """Adds a triangular FEM element between three particles in the system.
 
@@ -3634,15 +3637,15 @@ class ModelBuilder:
 
     def add_triangles(
         self,
-        i: List[int],
-        j: List[int],
-        k: List[int],
-        tri_ke: Optional[List[float]] = None,
-        tri_ka: Optional[List[float]] = None,
-        tri_kd: Optional[List[float]] = None,
-        tri_drag: Optional[List[float]] = None,
-        tri_lift: Optional[List[float]] = None,
-    ) -> List[float]:
+        i: list[int],
+        j: list[int],
+        k: list[int],
+        tri_ke: list[float] | None = None,
+        tri_ka: list[float] | None = None,
+        tri_kd: list[float] | None = None,
+        tri_drag: list[float] | None = None,
+        tri_lift: list[float] | None = None,
+    ) -> list[float]:
         """Adds triangular FEM elements between groups of three particles in the system.
 
         Triangles are modeled as viscoelastic elements with elastic stiffness and damping
@@ -3777,10 +3780,10 @@ class ModelBuilder:
         j: int,
         k: int,
         l: int,
-        rest: float = None,
-        edge_ke: float = None,
-        edge_kd: float = None,
-    ):
+        rest: float | None = None,
+        edge_ke: float | None = None,
+        edge_kd: float | None = None,
+    ) -> None:
         """Adds a bending edge element between four particles in the system.
 
         Bending elements are designed to be between two connected triangles. Then
@@ -3832,10 +3835,10 @@ class ModelBuilder:
         j,
         k,
         l,
-        rest: Optional[List[float]] = None,
-        edge_ke: Optional[List[float]] = None,
-        edge_kd: Optional[List[float]] = None,
-    ):
+        rest: list[float] | None = None,
+        edge_ke: list[float] | None = None,
+        edge_kd: list[float] | None = None,
+    ) -> None:
         """Adds bending edge elements between groups of four particles in the system.
 
         Bending elements are designed to be between two connected triangles. Then
@@ -3914,18 +3917,18 @@ class ModelBuilder:
         fix_right: bool = False,
         fix_top: bool = False,
         fix_bottom: bool = False,
-        tri_ke: float = None,
-        tri_ka: float = None,
-        tri_kd: float = None,
-        tri_drag: float = None,
-        tri_lift: float = None,
-        edge_ke: float = None,
-        edge_kd: float = None,
+        tri_ke: float | None = None,
+        tri_ka: float | None = None,
+        tri_kd: float | None = None,
+        tri_drag: float | None = None,
+        tri_lift: float | None = None,
+        edge_ke: float | None = None,
+        edge_kd: float | None = None,
         add_springs: bool = False,
-        spring_ke: float = None,
-        spring_kd: float = None,
-        particle_radius: float = None,
-    ):
+        spring_ke: float | None = None,
+        spring_kd: float | None = None,
+        particle_radius: float | None = None,
+    ) -> None:
         """Helper to create a regular planar cloth grid
 
         Creates a rectangular grid of particles with FEM triangles and bending elements
@@ -4054,23 +4057,23 @@ class ModelBuilder:
         rot: Quat,
         scale: float,
         vel: Vec3,
-        vertices: List[Vec3],
-        indices: List[int],
+        vertices: list[Vec3],
+        indices: list[int],
         density: float,
         edge_callback=None,
         face_callback=None,
-        tri_ke: float = None,
-        tri_ka: float = None,
-        tri_kd: float = None,
-        tri_drag: float = None,
-        tri_lift: float = None,
-        edge_ke: float = None,
-        edge_kd: float = None,
+        tri_ke: float | None = None,
+        tri_ka: float | None = None,
+        tri_kd: float | None = None,
+        tri_drag: float | None = None,
+        tri_lift: float | None = None,
+        edge_ke: float | None = None,
+        edge_kd: float | None = None,
         add_springs: bool = False,
-        spring_ke: float = None,
-        spring_kd: float = None,
-        particle_radius: float = None,
-    ):
+        spring_ke: float | None = None,
+        spring_kd: float | None = None,
+        particle_radius: float | None = None,
+    ) -> None:
         """Helper to create a cloth model from a regular triangle mesh
 
         Creates one FEM triangle element and one bending element for every face
@@ -4179,9 +4182,9 @@ class ModelBuilder:
         cell_z: float,
         mass: float,
         jitter: float,
-        radius_mean: float = None,
+        radius_mean: float | None = None,
         radius_std: float = 0.0,
-    ):
+    ) -> None:
         radius_mean = radius_mean if radius_mean is not None else self.default_particle_radius
 
         rng = np.random.default_rng(42)
@@ -4218,12 +4221,12 @@ class ModelBuilder:
         fix_right: bool = False,
         fix_top: bool = False,
         fix_bottom: bool = False,
-        tri_ke: float = None,
-        tri_ka: float = None,
-        tri_kd: float = None,
-        tri_drag: float = None,
-        tri_lift: float = None,
-    ):
+        tri_ke: float | None = None,
+        tri_ka: float | None = None,
+        tri_kd: float | None = None,
+        tri_drag: float | None = None,
+        tri_lift: float | None = None,
+    ) -> None:
         """Helper to create a rectangular tetrahedral FEM grid
 
         Creates a regular grid of FEM tetrahedra and surface triangles. Useful for example
@@ -4339,18 +4342,18 @@ class ModelBuilder:
         rot: Quat,
         scale: float,
         vel: Vec3,
-        vertices: List[Vec3],
-        indices: List[int],
+        vertices: list[Vec3],
+        indices: list[int],
         density: float,
         k_mu: float,
         k_lambda: float,
         k_damp: float,
-        tri_ke: float = None,
-        tri_ka: float = None,
-        tri_kd: float = None,
-        tri_drag: float = None,
-        tri_lift: float = None,
-    ):
+        tri_ke: float | None = None,
+        tri_ka: float | None = None,
+        tri_kd: float | None = None,
+        tri_drag: float | None = None,
+        tri_lift: float | None = None,
+    ) -> None:
         """Helper to create a tetrahedral model from an input tetrahedral mesh
 
         Args:
@@ -4460,15 +4463,15 @@ class ModelBuilder:
         self,
         normal=None,
         offset=0.0,
-        ke: float = None,
-        kd: float = None,
-        kf: float = None,
-        mu: float = None,
-        restitution: float = None,
-    ):
-        """
-        Creates a ground plane for the world. If the normal is not specified,
-        the up_vector of the ModelBuilder is used.
+        ke: float | None = None,
+        kd: float | None = None,
+        kf: float | None = None,
+        mu: float | None = None,
+        restitution: float | None = None,
+    ) -> None:
+        """Create a ground plane for the world.
+
+        If the normal is not specified, the ``up_vector`` of the :class:`ModelBuilder` is used.
         """
         ke = ke if ke is not None else self.default_shape_ke
         kd = kd if kd is not None else self.default_shape_kd
@@ -4489,26 +4492,25 @@ class ModelBuilder:
             "restitution": restitution,
         }
 
-    def _create_ground_plane(self):
+    def _create_ground_plane(self) -> None:
         ground_id = self.add_shape_plane(**self._ground_params)
         self._ground_created = True
         # disable ground collisions as they will be treated separately
         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.
+    def set_coloring(self, particle_color_groups):
+        """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_color_groups: 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 = [
+        particle_color_groups = [
             color_group if isinstance(color_group, np.ndarray) else np.array(color_group)
-            for color_group in particle_coloring
+            for color_group in particle_color_groups
         ]
-        self.particle_coloring = particle_coloring
+        self.particle_color_groups = particle_color_groups
 
     def color(
         self,
@@ -4516,9 +4518,8 @@ class ModelBuilder:
         balance_colors=True,
         target_max_min_color_ratio=1.1,
         coloring_algorithm=ColoringAlgorithm.MCS,
-    ):
-        """
-        Run coloring algorithm to generate coloring information.
+    ) -> None:
+        """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
@@ -4544,7 +4545,7 @@ class ModelBuilder:
         # ignore bending energy if it is too small
         edge_indices = np.array(self.edge_indices)
 
-        self.particle_coloring = color_trimesh(
+        self.particle_color_groups = color_trimesh(
             len(self.particle_q),
             edge_indices,
             include_bending,
@@ -4604,7 +4605,11 @@ 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]
+            particle_colors = np.empty(self.particle_count, dtype=int)
+            for color in range(len(self.particle_color_groups)):
+                particle_colors[self.particle_color_groups[color]] = color
+            m.particle_colors = wp.array(particle_colors, dtype=int)
+            m.particle_color_groups = [wp.array(group, dtype=int) for group in self.particle_color_groups]
 
             # hash-grid for particle interactions
             m.particle_grid = wp.HashGrid(128, 128, 128)