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

Code before:

        Returns
        -------
        particles_idx_local : torch.Tensor, shape (n_in_bbox,)
            Local particle indices within the bounding box.
        """
        bbox_min = torch.tensor(bbox_min, dtype=gs.tc_float, device=gs.device)
        bbox_max = torch.tensor(bbox_max, dtype=gs.tc_float, device=gs.device)

        # Get particle positions (use first env if batched)
        poss = self.get_particles_pos(envs_idx=envs_idx)
        if poss.ndim == 3:
            poss = poss[0]  # Use first environment

        # Find particles in bounding box
        in_bbox = (
            (poss[:, 0] >= bbox_min[0])
            & (poss[:, 0] <= bbox_max[0])
            & (poss[:, 1] >= bbox_min[1])
            & (poss[:, 1] <= bbox_max[1])
            & (poss[:, 2] >= bbox_min[2])
            & (poss[:, 2] <= bbox_max[2])
        )

        particles_idx_local = torch.nonzero(in_bbox, as_tuple=False).squeeze(-1)
        return particles_idx_local

    @gs.assert_built
    def set_particle_constraints(self, particles_idx_local, link, stiffness=1e4, envs_idx=None):
        """
        Attach MPM particles to a rigid link using soft constraints.

        The particles will be pulled toward their relative position on the link
        using spring forces with critical damping.

        Parameters
        ----------
        particles_idx_local : array_like
            Local particle indices to constrain.
        link : RigidLink
            The rigid link to attach particles to.
        stiffness : float, optional
            Spring stiffness for the constraint. Defaults to 1e4.
        envs_idx : array_like, optional
            Environment indices. If None, applies to all environments.

        Note
        ----
        Requires `enable_particle_constraints=True` in MPMOptions.
        """
        from genesis.engine.entities.rigid_entity import RigidLink

Review comment (suggestion):
`torch.nonzero` should be avoided because it forces GPU sync. Try supporting boolean mask all the way for efficiency.

Code after:
        -------
        mask : torch.Tensor, shape (n_envs, n_particles)
            Boolean mask where True indicates particle is within the bounding box.
        """
        bbox_min = torch.as_tensor(bbox_min, dtype=gs.tc_float, device=gs.device)
        bbox_max = torch.as_tensor(bbox_max, dtype=gs.tc_float, device=gs.device)

        # Get particle positions: shape (n_envs, n_particles, 3)
        poss = self.get_particles_pos()
        if poss.ndim == 2:
            poss = poss.unsqueeze(0)  # (1, n_particles, 3)

        # Vectorized bbox check: (n_envs, n_particles)
        mask = ((bbox_min <= poss) & (poss <= bbox_max)).all(dim=-1)
        return mask

    @gs.assert_built
    def set_particle_constraints(self, particles_mask, link_idx, stiffness):
        """
        Attach MPM particles to a rigid link using soft constraints.

        The particles will be pulled toward their relative position on the link
        using spring forces with critical damping.

        Parameters
        ----------
        particles_mask : torch.Tensor, shape (n_envs, n_particles)
            Boolean mask indicating which particles to constrain.
        link_idx : int
            Index of the rigid link to attach particles to.
        stiffness : float
            Spring stiffness for the constraint.
        """
        if not isinstance(link_idx, int):
            gs.raise_exception("link_idx must be an integer.")

        if not self._solver._constraints_initialized:
            self._solver.init_constraints()

        # Get link position and quaternion for all envs
        rigid_solver = self._sim.coupler.rigid_solver
        link_pos = rigid_solver.get_links_pos(links_idx=[link_idx])  # (n_envs, 1, 3)
        link_quat = rigid_solver.get_links_quat(links_idx=[link_idx])  # (n_envs, 1, 4)
        if link_pos.ndim == 2:
            link_pos = link_pos.unsqueeze(0)
            link_quat = link_quat.unsqueeze(0)
        link_pos = link_pos[:, 0, :]  # (n_envs, 3)
        link_quat = link_quat[:, 0, :]  # (n_envs, 4)

        self._solver._kernel_set_particle_constraints(
            self._sim.cur_substep_local,