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

Code before:
        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

        if not self._solver._enable_particle_constraints:
            gs.raise_exception(

Review comment (suggestion):
Do not pass link but rather link_idx. We do not want to pass complex objects as input arguments. It is just asking for trouble in terms of memory management, and it is harder to keep track of the simulation state.

Code after:
        """
        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,
            particles_mask,
            self._particle_start,
            stiffness,