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

Code before:

    ti.loop_config(serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.PARTIAL)
    for i_b in range(_B):
        nefc = constraint_state.n_constraints[i_b]
        # zero AR
        for r in range(nefc):
            for c in range(nefc):
                constraint_state.efc_AR[r, c, i_b] = gs.ti_float(0.0)

        # build AR = J * inv(M) * J^T
        # do it row-by-row: for each row r, tmp = inv(M) * J[r]^T, then AR[r,:] = J * tmp
        for r in range(nefc):
            # tmp = M^{-1} * Jr^T
            for i_d in range(n_dofs):
                constraint_state.Mgrad[i_d, i_b] = constraint_state.jac[r, i_d, i_b]

            rigid_solver.func_solve_mass_batched(
                constraint_state.Mgrad,
                constraint_state.Mgrad,
                i_b,
                entities_info=entities_info,
                rigid_global_info=rigid_global_info,
                static_rigid_sim_config=static_rigid_sim_config,
            )

            # AR[r, c] = J[c, :] dot tmp
            for c in range(nefc):
                s = gs.ti_float(0.0)
                for i_d in range(n_dofs):
                    s += constraint_state.jac[c, i_d, i_b] * constraint_state.Mgrad[i_d, i_b]
                constraint_state.efc_AR[r, c, i_b] = s

        # we don't need this term if we only use PGS for noslip
        # for r in range(nefc):
        #     constraint_state.efc_AR[r, r, i_b] += 1.0 / constraint_state.efc_D[r, i_b]

        for i_c in range(constraint_state.n_constraints[i_b]):
            v = -constraint_state.aref[i_c, i_b]
            for i_d in range(n_dofs):
                v += constraint_state.jac[i_c, i_d, i_b] * dofs_state.acc_smooth[i_d, i_b]
            constraint_state.efc_b[i_c, i_b] = v


@gs.maybe_pure
@ti.kernel
def kernel_noslip(
    collider_state: array_class.ColliderState,
    constraint_state: array_class.ConstraintState,
    rigid_global_info: array_class.RigidGlobalInfo,
    static_rigid_sim_config: ti.template(),
    static_rigid_sim_cache_key: array_class.StaticRigidSimCacheKey,

Review comment (suggestion):
replace `dot` by `*` for consistency with the above formula (even though it should probably be `@` everywhere instead)

Code after:

    ti.loop_config(serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.PARTIAL)
    for i_b in range(_B):
        nefc = constraint_state.n_constraints[i_b]
        # zero AR
        for i_row in range(nefc):
            for i_col in range(nefc):
                constraint_state.efc_AR[i_row, i_col, i_b] = gs.ti_float(0.0)

        # build AR = J * inv(M) * J^T
        # do it row-by-row: for each row r, tmp = inv(M) * J[r]^T, then AR[r,:] = J * tmp
        for i_row in range(nefc):
            # tmp = M^{-1} * Jr^T
            for i_d in range(n_dofs):
                constraint_state.Mgrad[i_d, i_b] = constraint_state.jac[i_row, i_d, i_b]

            rigid_solver.func_solve_mass_batched(
                constraint_state.Mgrad,
                constraint_state.Mgrad,
                i_b,
                entities_info=entities_info,
                rigid_global_info=rigid_global_info,
                static_rigid_sim_config=static_rigid_sim_config,
            )

            # AR[r, c] = J[c, :] * tmp
            for i_col in range(nefc):
                s = gs.ti_float(0.0)
                for i_d in range(n_dofs):
                    s += constraint_state.jac[i_col, i_d, i_b] * constraint_state.Mgrad[i_d, i_b]
                constraint_state.efc_AR[i_row, i_col, i_b] = s

        for i_c in range(constraint_state.n_constraints[i_b]):
            v = -constraint_state.aref[i_c, i_b]
            for i_d in range(n_dofs):
                v += constraint_state.jac[i_c, i_d, i_b] * dofs_state.acc_smooth[i_d, i_b]
            constraint_state.efc_b[i_c, i_b] = v


@gs.maybe_pure
@ti.kernel
def kernel_noslip(
    collider_state: array_class.ColliderState,
    constraint_state: array_class.ConstraintState,
    rigid_global_info: array_class.RigidGlobalInfo,
    static_rigid_sim_config: ti.template(),
    static_rigid_sim_cache_key: array_class.StaticRigidSimCacheKey,
):
    _B = constraint_state.jac.shape[2]
    n_dofs = constraint_state.jac.shape[1]