Repository: Genesis-Embodied-AI/Genesis
File: genesis/engine/entities/pbd_entity.py (Python)

Code before:

        self._kernel_add_uvs_and_faces_to_solver(
            uvs=uvs,
            faces=self._vfaces,
        )

    @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 (skip if no UVs provided)
        n_uvs = uvs.shape[0]
        for i_vv_ in range(n_uvs):
            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
            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,
                ]
            )

    @gs.assert_built
    def set_particles_pos(self, poss, particles_idx_local=None, envs_idx=None):
        envs_idx = self._scene._sanitize_envs_idx(envs_idx)
        particles_idx_local = self._sanitize_particles_idx_local(particles_idx_local, envs_idx)
        particles_idx = particles_idx_local + self._particle_start
        poss = self._sanitize_particles_tensor(poss, gs.tc_float, particles_idx, envs_idx, (3,))
        self.solver._kernel_set_particles_pos(particles_idx, envs_idx, poss)

    def get_particles_pos(self, envs_idx=None):
        envs_idx = self._scene._sanitize_envs_idx(envs_idx)
        poss = self._sanitize_particles_tensor(None, gs.tc_float, None, envs_idx, (3,))
        self.solver._kernel_get_particles_pos(self._particle_start, self.n_particles, envs_idx, poss)
        if self._scene.n_envs == 0:
            poss = poss[0]
        return poss

    @gs.assert_built
    def set_particles_vel(self, vels, particles_idx_local=None, envs_idx=None):
        envs_idx = self._scene._sanitize_envs_idx(envs_idx)
        particles_idx_local = self._sanitize_particles_idx_local(particles_idx_local, envs_idx)
        particles_idx = particles_idx_local + self._particle_start
        vels = self._sanitize_particles_tensor(vels, gs.tc_float, particles_idx, envs_idx, (3,))

Review comment (suggestion):
Remove cast, it should not be necessary. You should rather cast the input argument before calling this kernel.

Code after:
    @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 (skip if no UVs provided)
        n_uvs = uvs.shape[0]
        for i_vv_ in range(n_uvs):
            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
            self.solver.vfaces_indices[i_vf] = ti.Vector(
                [
                    faces[i_vf_][0] + self._vvert_start,
                    faces[i_vf_][1] + self._vvert_start,
                    faces[i_vf_][2] + self._vvert_start,
                ]
            )

    @gs.assert_built
    def set_particles_pos(self, poss, particles_idx_local=None, envs_idx=None):
        envs_idx = self._scene._sanitize_envs_idx(envs_idx)
        particles_idx_local = self._sanitize_particles_idx_local(particles_idx_local, envs_idx)
        particles_idx = particles_idx_local + self._particle_start
        poss = self._sanitize_particles_tensor(poss, gs.tc_float, particles_idx, envs_idx, (3,))
        self.solver._kernel_set_particles_pos(particles_idx, envs_idx, poss)

    def get_particles_pos(self, envs_idx=None):
        envs_idx = self._scene._sanitize_envs_idx(envs_idx)
        poss = self._sanitize_particles_tensor(None, gs.tc_float, None, envs_idx, (3,))
        self.solver._kernel_get_particles_pos(self._particle_start, self.n_particles, envs_idx, poss)
        if self._scene.n_envs == 0:
            poss = poss[0]
        return poss

    @gs.assert_built
    def set_particles_vel(self, vels, particles_idx_local=None, envs_idx=None):
        envs_idx = self._scene._sanitize_envs_idx(envs_idx)
        particles_idx_local = self._sanitize_particles_idx_local(particles_idx_local, envs_idx)
        particles_idx = particles_idx_local + self._particle_start
        vels = self._sanitize_particles_tensor(vels, gs.tc_float, particles_idx, envs_idx, (3,))
        self.solver._kernel_set_particles_vel(particles_idx, envs_idx, vels)

    def get_particles_vel(self, envs_idx=None):
        envs_idx = self._scene._sanitize_envs_idx(envs_idx)
        vels = self._sanitize_particles_tensor(None, gs.tc_float, None, envs_idx, (3,))
        self.solver._kernel_get_particles_vel(self._particle_start, self.n_particles, envs_idx, vels)