Repository: Genesis-Embodied-AI/Genesis
File: genesis/engine/solvers/rigid/rigid_solver_decomp.py (Python)

Code before:

    def get_weld_constraints(self, envs_idx=None):
        if envs_idx is None:
            envs_idx = np.arange(self.n_envs, dtype=np.int32)
        envs_idx = np.atleast_1d(envs_idx)

        env_col, obj_a_col, obj_b_col = [], [], []
        WELD = gs.EQUALITY_TYPE.WELD

        for env in envs_idx:
            n_eq = int(self.constraint_solver.ti_n_equalities[env])
            for i in range(n_eq):
                rec = self.equalities_info[i, env]
                if rec.eq_type == WELD:
                    env_col.append(env)
                    obj_a_col.append(int(rec.eq_obj1id))
                    obj_b_col.append(int(rec.eq_obj2id))

        obj_a = np.asarray(obj_a_col, dtype=np.int32)
        obj_b = np.asarray(obj_b_col, dtype=np.int32)

        if self.n_envs == 0:
            return {"obj_a": obj_a, "obj_b": obj_b}

        env_arr = np.asarray(env_col, dtype=np.int32)
        valid_mask = np.zeros((self.n_envs, obj_a.shape[0]), dtype=bool)
        for j, e in enumerate(env_arr):
            valid_mask[e, j] = True

        return {
            "env": env_arr,
            "obj_a": obj_a,
            "obj_b": obj_b,
            "valid_mask": valid_mask,
        }

    # ------------------------------------------------------------------------------------
    # ----------------------------------- properties -------------------------------------
    # ------------------------------------------------------------------------------------

    @property
    def links(self):
        if self.is_built:
            return self._links
        return gs.List(link for entity in self._entities for link in entity.links)

    @property
    def joints(self):
        if self.is_built:
            return self._joints
        return gs.List(joint for entity in self._entities for joint in entity.joints)

Review comment (bug):
You took inspiration from the wrong function. It was supposed to be `Collider.get_contacts`, not `RigidEntity.get_contacts`, sorry about that. Which mean that the API should be `def get_weld_constraints(self, as_tensor: bool = True, to_torch: bool = True):`

Then you could also add its counterpart in `RigidEntity` if you think it is worth it.

Code after:
        for i_l in range(entities_info.link_start[i_e], entities_info.link_end[i_e]):
            n_awake_links = ti.atomic_add(rigid_global_info.n_awake_links[i_b], 1)
            rigid_global_info.awake_links[n_awake_links, i_b] = i_l


@ti.func
def func_hibernate_entity(
    i_e,
    i_b,
    entities_state,
    entities_info,
    dofs_state,
    links_state,
    geoms_state,
):

    entities_state.hibernated[i_e, i_b] = True

    for i_d in range(entities_info.dof_start[i_e], entities_info.dof_end[i_e]):
        dofs_state.hibernated[i_d, i_b] = True
        dofs_state.vel[i_d, i_b] = 0.0
        dofs_state.acc[i_d, i_b] = 0.0

    for i_l in range(entities_info.link_start[i_e], entities_info.link_end[i_e]):
        links_state.hibernated[i_l, i_b] = True

    for i_g in range(entities_info.geom_start[i_e], entities_info.geom_end[i_e]):
        geoms_state.hibernated[i_g, i_b] = True


@ti.func
def func_wakeup_entity(
    i_e,
    i_b,
    entities_state: array_class.EntitiesState,
    entities_info: array_class.EntitiesInfo,
    dofs_state: array_class.DofsState,
    links_state: array_class.LinksState,
    rigid_global_info: array_class.RigidGlobalInfo,
):
    if entities_state.hibernated[i_e, i_b]:
        entities_state.hibernated[i_e, i_b] = False
        n_awake_entities = ti.atomic_add(rigid_global_info.n_awake_entities[i_b], 1)
        rigid_global_info.awake_entities[n_awake_entities, i_b] = i_e

        for i_d in range(entities_info.dof_start[i_e], entities_info.dof_end[i_e]):
            dofs_state.hibernated[i_d, i_b] = False
            n_awake_dofs = ti.atomic_add(rigid_global_info.n_awake_dofs[i_b], 1)
            rigid_global_info.awake_dofs[n_awake_dofs, i_b] = i_d

        for i_l in range(entities_info.link_start[i_e], entities_info.link_end[i_e]):