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

Code before:
        tmp_quad_total0 = constraint_state.quad_gauss[_i0 + 0, i_b]
        tmp_quad_total1 = constraint_state.quad_gauss[_i0 + 1, i_b]
        tmp_quad_total2 = constraint_state.quad_gauss[_i0 + 2, i_b]
        for i_c in range(constraint_state.n_constraints[i_b]):
            x = constraint_state.Jaref[i_c, i_b] + alpha * constraint_state.jv[i_c, i_b]
            active = 1
            ne = constraint_state.n_constraints_equality[i_b]
            nef = ne + constraint_state.n_constraints_frictionloss[i_b]
            if i_c >= ne and i_c < nef:
                f = constraint_state.efc_frictionloss[i_c, i_b]
                r = 1.0 / ti.max(constraint_state.efc_D[i_c, i_b], gs.EPS)
                rf = r * f
                linear_neg = x <= -rf
                linear_pos = x >= rf
                # active = (~linear_neg) & (~linear_pos)

                if linear_neg or linear_pos:
                    qf = linear_neg * ti.Vector(
                        [f * (-0.5 * rf - constraint_state.Jaref[i_c, i_b]), -f * constraint_state.jv[i_c, i_b], 0.0]
                    )
                    qf += linear_pos * ti.Vector(
                        [f * (-0.5 * rf + constraint_state.Jaref[i_c, i_b]), f * constraint_state.jv[i_c, i_b], 0.0]
                    )
                    constraint_state.quad[i_c, _i0 + 0, i_b] = qf[0]
                    constraint_state.quad[i_c, _i0 + 1, i_b] = qf[1]
                    constraint_state.quad[i_c, _i0 + 2, i_b] = qf[2]

            elif i_c >= nef:
                active = x < 0

            tmp_quad_total0 += constraint_state.quad[i_c, _i0 + 0, i_b] * active
            tmp_quad_total1 += constraint_state.quad[i_c, _i0 + 1, i_b] * active
            tmp_quad_total2 += constraint_state.quad[i_c, _i0 + 2, i_b] * active

    cost = alpha * alpha * tmp_quad_total2 + alpha * tmp_quad_total1 + tmp_quad_total0

    deriv_0 = 2 * alpha * tmp_quad_total2 + tmp_quad_total1
    deriv_1 = 2 * tmp_quad_total2 + gs.EPS * (ti.abs(tmp_quad_total2) < gs.EPS)

    constraint_state.ls_its[i_b] = constraint_state.ls_its[i_b] + 1

    return alpha, cost, deriv_0, deriv_1


@ti.func
def func_no_linesearch(i_b, constraint_state: array_class.ConstraintState):
    func_ls_init(i_b)
    n_dofs = constraint_state.search.shape[0]

    constraint_state.improved[i_b] = 1
    for i_d in range(n_dofs):

Review comment (performance):
You are setting value equal to zero. Just do this directly. Moreover, I don't this it is worth making `qf` a vector. Just define `qf_1` and `qf_2`. That would be more efficient.

Code after:
    tmp_quad_total1 = constraint_state.quad_gauss[1, i_b]
    tmp_quad_total2 = constraint_state.quad_gauss[2, i_b]
    for i_c in range(constraint_state.n_constraints[i_b]):
        x = constraint_state.Jaref[i_c, i_b] + alpha * constraint_state.jv[i_c, i_b]
        active = 1
        ne = constraint_state.n_constraints_equality[i_b]
        nef = ne + constraint_state.n_constraints_frictionloss[i_b]
        qf_0 = constraint_state.quad[i_c, 0, i_b]
        qf_1 = constraint_state.quad[i_c, 1, i_b]
        qf_2 = constraint_state.quad[i_c, 2, i_b]
        if i_c >= ne and i_c < nef:
            f = constraint_state.efc_frictionloss[i_c, i_b]
            r = 1.0 / ti.max(constraint_state.efc_D[i_c, i_b], gs.EPS)
            rf = r * f
            linear_neg = x <= -rf
            linear_pos = x >= rf

            if linear_neg or linear_pos:
                qf_0 = linear_neg * f * (-0.5 * rf - constraint_state.Jaref[i_c, i_b]) + linear_pos * f * (
                    -0.5 * rf + constraint_state.Jaref[i_c, i_b]
                )
                qf_1 = linear_neg * (-f * constraint_state.jv[i_c, i_b]) + linear_pos * (
                    f * constraint_state.jv[i_c, i_b]
                )
                qf_2 = 0.0
        elif i_c >= nef:
            active = x < 0

        tmp_quad_total0 += qf_0 * active
        tmp_quad_total1 += qf_1 * active
        tmp_quad_total2 += qf_2 * active

    cost = alpha * alpha * tmp_quad_total2 + alpha * tmp_quad_total1 + tmp_quad_total0

    deriv_0 = 2 * alpha * tmp_quad_total2 + tmp_quad_total1
    deriv_1 = 2 * tmp_quad_total2 + gs.EPS * (ti.abs(tmp_quad_total2) < gs.EPS)

    constraint_state.ls_its[i_b] = constraint_state.ls_its[i_b] + 1

    return alpha, cost, deriv_0, deriv_1


@ti.func
def func_no_linesearch(i_b, constraint_state: array_class.ConstraintState):
    func_ls_init(i_b)
    n_dofs = constraint_state.search.shape[0]

    constraint_state.improved[i_b] = 1
    for i_d in range(n_dofs):
        constraint_state.qacc[i_d, i_b] = constraint_state.qacc[i_d, i_b] + constraint_state.search[i_d, i_b]
        constraint_state.Ma[i_d, i_b] = constraint_state.Ma[i_d, i_b] + constraint_state.mv[i_d, i_b]