Merge branch 'main' into feature/jacobian

Author: duburcqa

.github/workflows/generic-cpu.yml .github/workflows/generic.yml.github/workflows/generic-cpu.yml renamed to .github/workflows/generic.yml

@@ -1,4 +1,4 @@
-name: Generic - CPU only
+name: Generic
 
 on:
   pull_request:

.github/workflows/linux-gpu.yml .github/workflows/production.yml.github/workflows/linux-gpu.yml renamed to .github/workflows/production.yml

@@ -1,4 +1,4 @@
-name: Linux x86 - Nvidia GPU
+name: Production
 
 on:
   # Trigger the workflow on push on the master branch, or for any pull request

.pre-commit-config.yaml

@@ -1,8 +1,5 @@
 repos:
   - repo: https://github.com/psf/black
-    rev: 25.1.0  # Use the latest version or specify the version you prefer
+    rev: 25.1.0
     hooks:
       - id: black
-        args:
-          - --line-length=120
-          - .

examples/rigid/suction_cup.py

@@ -96,8 +96,8 @@ def main():
 
     # add suction / weld constraint
     rigid = scene.sim.rigid_solver
-    link_cube = np.array([cube.get_link("box_baselink").idx], dtype=gs.np_int)
-    link_franka = np.array([franka.get_link("hand").idx], dtype=gs.np_int)
+    link_cube = cube.get_link("box_baselink").idx
+    link_franka = franka.get_link("hand").idx
     rigid.add_weld_constraint(link_cube, link_franka)
 
     # lift

genesis/engine/entities/emitter.py

@@ -150,7 +150,7 @@ def emit(
             positions = gu.transform_by_trans_R(
                 positions,
                 pos,
-                gu.z_up_to_R(direction) @ gu.axis_angle_to_R(np.array([0, 0, 1]), theta),
+                gu.z_up_to_R(direction) @ gu.axis_angle_to_R(np.array([0.0, 0.0, 1.0], dtype=gs.np_float), theta),
             ).astype(gs.np_float)
 
             positions = np.tile(positions[np.newaxis], (self._sim._B, 1, 1))
@@ -206,17 +206,21 @@ def emit_omni(self, source_radius=0.1, pos=(0.5, 0.5, 1.0), speed=1.0, particle_
         Parameters:
         ----------
         source_radius: float, optional
-            The radius of the sphere source. Particles will be emitted from a shell with inner radius using 0.8 * source_radius and outer radius using source_radius.
+            The radius of the sphere source. Particles will be emitted from a shell with inner radius using
+            '0.8 * source_radius' and outer radius using source_radius.
         pos: array_like, shape=(3,)
             The center of the sphere source.
         speed: float
             The speed of the emitted particles.
         particle_size: float | None
-            The size (diameter) of the emitted particles. The actual number of particles emitted is determined by the volume of the sphere source and the size of the particles. If None, the solver's particle size is used. Note that this particle size only affects computation for number of particles emitted, not the actual size of the particles in simulation and rendering.
+            The size (diameter) of the emitted particles. The actual number of particles emitted is determined by the
+            volume of the sphere source and the size of the particles. If None, the solver's particle size is used.
+            Note that this particle size only affects computation for number of particles emitted, not the actual size
+            of the particles in simulation and rendering.
         """
         assert self._entity is not None
 
-        pos = np.array(pos)
+        pos = np.asarray(pos, dtype=gs.np_float)
 
         if particle_size is None:
             particle_size = self._solver.particle_size
@@ -227,20 +231,20 @@ def emit_omni(self, source_radius=0.1, pos=(0.5, 0.5, 1.0), speed=1.0, particle_
             inner_radius=source_radius * 0.4,
             sampler=self._entity.sampler,
         )
-
-        positions = gu.transform_by_T(positions_, gu.trans_to_T(pos)).astype(gs.np_float)
+        positions = pos + positions_
 
         if not self._solver.boundary.is_inside(positions):
             gs.raise_exception("Emitted particles are outside the boundary.")
 
-        n_particles = len(positions)
-        dists = np.linalg.norm(positions_, axis=1, keepdims=True)
-        positions[np.where(dists < gs.EPS)[0]] = np.array([gs.EPS, gs.EPS, gs.EPS])
-        vels = (positions_ / dists * speed).astype(gs.np_float)
+        dists = np.linalg.norm(positions_, axis=1)
+        positions[dists < gs.EPS] = gs.EPS
+        vels = (speed / (dists + gs.EPS)) * positions_
 
+        n_particles = len(positions)
         if n_particles > self._entity.n_particles:
             gs.logger.warning(
-                f"Number of particles to emit ({n_particles}) at the current step is larger than the maximum number of particles ({self._entity.n_particles})."
+                f"Number of particles to emit ({n_particles}) at the current step is larger than the maximum number "
+                f"of particles ({self._entity.n_particles})."
             )
 
         self._solver._kernel_set_particles_pos(

genesis/engine/entities/mpm_entity.py

@@ -89,14 +89,9 @@ def set_pos(self, f, pos):
         f : int
             The current substep index.
         pos : gs.Tensor
-            A tensor of shape (n_particles, 3) representing particle positions.
+            A tensor of shape (n_envs, n_particles, 3) representing particle positions.
         """
-        self.solver._kernel_set_particles_pos(
-            f,
-            self._particle_start,
-            self._n_particles,
-            pos,
-        )
+        self.solver._kernel_set_particles_pos(f, self._particle_start, self._n_particles, pos)
 
     def set_pos_grad(self, f, pos_grad):
         """

genesis/engine/entities/rigid_entity/rigid_link.py

@@ -55,24 +55,24 @@ def __init__(
         self._entity_idx_in_solver = entity.idx
 
         self._uid = gs.UID()
-        self._idx = idx
-        self._parent_idx = parent_idx
-        self._root_idx = root_idx
-        self._child_idxs = list()
-        self._invweight = invweight
-
-        self._joint_start = joint_start
-        self._n_joints = n_joints
-
-        self._geom_start = geom_start
-        self._cell_start = cell_start
-        self._vert_start = vert_start
-        self._face_start = face_start
-        self._edge_start = edge_start
-        self._verts_state_start = verts_state_start
-        self._vgeom_start = vgeom_start
-        self._vvert_start = vvert_start
-        self._vface_start = vface_start
+        self._idx: int = idx
+        self._parent_idx: int = parent_idx  # -1 if no parent
+        self._root_idx: int | None = root_idx  # None if no root
+        self._child_idxs: list[int] = list()
+        self._invweight: float | None = invweight
+
+        self._joint_start: int = joint_start
+        self._n_joints: int = n_joints
+
+        self._geom_start: int = geom_start
+        self._cell_start: int = cell_start
+        self._vert_start: int = vert_start
+        self._face_start: int = face_start
+        self._edge_start: int = edge_start
+        self._verts_state_start: int = verts_state_start
+        self._vgeom_start: int = vgeom_start
+        self._vvert_start: int = vvert_start
+        self._vface_start: int = vface_start
 
         # Link position & rotation at creation time:
         self._pos: ArrayLike = pos

genesis/engine/entities/sph_entity.py

@@ -92,12 +92,7 @@ def set_pos(self, f, pos):
         pos : ndarray
             Array of particle positions of shape (n_envs, n_particles, 3).
         """
-        self.solver._kernel_set_particles_pos(
-            f,
-            self._particle_start,
-            self._n_particles,
-            pos,
-        )
+        self.solver._kernel_set_particles_pos(f, self._particle_start, self._n_particles, pos)
 
     def set_pos_grad(self, f: ti.i32, pos_grad: ti.types.ndarray()):
         """

genesis/engine/solvers/mpm_solver.py

@@ -720,9 +720,9 @@ def _kernel_set_particles_pos(
                 self.particles[f, i_global, i_b].pos[k] = pos[i_b, i_p, k]
 
             # we restore these whenever directly setting positions
-            self.particles[f, i_global, i_b].vel = ti.Vector.zero(gs.ti_float, 3)
+            self.particles[f, i_global, i_b].vel.fill(0.0)
             self.particles[f, i_global, i_b].F = ti.Matrix.identity(gs.ti_float, 3)
-            self.particles[f, i_global, i_b].C = ti.Matrix.zero(gs.ti_float, 3, 3)
+            self.particles[f, i_global, i_b].C.fill(0.0)
             self.particles[f, i_global, i_b].Jp = self.particles_info[i_global].default_Jp
 
     @ti.kernel

genesis/engine/solvers/pbd_solver.py

@@ -840,8 +840,8 @@ def _kernel_set_particles_pos(
         for i_p, i_b in ti.ndrange(n_particles, self._B):
             i_global = i_p + particle_start
             for k in ti.static(range(3)):
-                self.particles[i_global, i_b].pos[k] = pos[i_p, k]
-            self.particles[i_global, i_b].vel = ti.Vector.zero(gs.ti_float, 3)
+                self.particles[i_global, i_b].pos[k] = pos[i_b, i_p, k]
+            self.particles[i_global, i_b].vel.fill(0.0)
 
     @ti.kernel
     def _kernel_set_particles_vel(