Added uvs to PBD solver for rendering

Author: alelievr

genesis/engine/entities/particle_entity.py

@@ -184,6 +184,7 @@ def _add_to_solver(self):
         self._add_particles_to_solver()
         if self._need_skinning:
             self._add_vverts_to_solver()
+            self._add_uvs_and_faces_to_solver()
 
     def _add_particles_to_solver(self):
         raise NotImplementedError
@@ -229,6 +230,45 @@ def _kernel_add_vverts_to_solver(
             for j in range(self.solver._n_vvert_supports):
                 self.solver.vverts_info.support_idxs[i_vv][j] = support_idxs_local[i_vv_, j] + self._particle_start
 
+    def _add_uvs_and_faces_to_solver(self):
+        """Add UV coordinates and face indices to the solver's global buffers."""
+        # Get UVs from vmesh (may be None if no texture)
+        uvs = self._vmesh.uvs if self._vmesh is not None else None
+        if uvs is not None and len(uvs) == self.n_vverts:
+            uvs_np = np.asarray(uvs, dtype=gs.np_float)
+        else:
+            # No UVs available, use zeros
+            uvs_np = np.zeros((self.n_vverts, 2), dtype=gs.np_float)
+
+        # Get face indices and offset them by vvert_start for global indexing
+        faces_np = np.asarray(self._vfaces, dtype=gs.np_int)
+
+        self._kernel_add_uvs_and_faces_to_solver(
+            uvs=uvs_np,
+            faces=faces_np,
+        )
+
+    @ti.kernel
+    def _kernel_add_uvs_and_faces_to_solver(
+        self,
+        uvs: ti.types.ndarray(element_dim=1),
+        faces: ti.types.ndarray(element_dim=1),
+    ):
+        # Copy UVs to solver's global UV buffer
+        for i_vv_ in range(self.n_vverts):
+            i_vv = i_vv_ + self._vvert_start
+            self.solver.vverts_uvs[i_vv] = uvs[i_vv_]
+
+        # Copy faces to solver's global face buffer (with global vertex indices)
+        for i_vf_ in range(self.n_vfaces):
+            i_vf = i_vf_ + self._vface_start
+            # Offset face indices by vvert_start for global vertex indexing
+            self.solver.vfaces_indices[i_vf] = ti.Vector([
+                ti.cast(faces[i_vf_][0], gs.ti_int) + self._vvert_start,
+                ti.cast(faces[i_vf_][1], gs.ti_int) + self._vvert_start,
+                ti.cast(faces[i_vf_][2], gs.ti_int) + self._vvert_start,
+            ])
+
     def sample(self):
         """
         Sample particles from the morph based on its type and the specified sampler.

genesis/engine/solvers/pbd_solver.py

@@ -102,6 +102,12 @@ def init_vvert_fields(self):
             shape=(max(self._n_vverts, 1), self._B), layout=ti.Layout.SOA
         )
 
+        # UV coordinates for visual vertices (static, same across all batch envs)
+        self.vverts_uvs = ti.field(dtype=gs.ti_vec2, shape=(max(self._n_vverts, 1),))
+
+        # Triangle face indices for visual mesh (static)
+        self.vfaces_indices = ti.field(dtype=gs.ti_ivec3, shape=(max(self._n_vfaces, 1),))
+
     def init_particle_fields(self):
         # particles information (static)
         struct_particle_info = ti.types.struct(
@@ -829,6 +835,24 @@ def get_state(self, f):
             state = None
         return state
 
+    def get_state_render(self):
+        """
+        Get visual vertex positions, UVs, and face indices for rendering.
+
+        Returns
+        -------
+        tuple
+            (vverts_pos, vverts_uvs, vfaces_indices) - vertex positions, UV coords, and triangle indices
+        """
+        if not self.is_active or self._n_vverts == 0:
+            return None, None, None
+
+        # Make sure render fields are up to date
+        self.update_render_fields()
+
+        # Return the Taichi fields directly for GPU access
+        return self.vverts_render.pos, self.vverts_uvs, self.vfaces_indices
+
     @ti.kernel
     def _kernel_get_state(
         self,