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

Code before:
                        area = 0.5 * e1.cross(e2).norm()
                        total_area += area
                        total_area_weighted_centroid += (
                            area * (polygon_vertices[:, 0] + polygon_vertices[:, i - 1] + polygon_vertices[:, i]) / 3.0
                        )

            centroid = total_area_weighted_centroid / total_area

            # Compute barycentric coordinates
            barycentric = tet_barycentric(centroid, tet_vertices)
            pressure = (
                barycentric[0] * tet_pressures[0]
                + barycentric[1] * tet_pressures[1]
                + barycentric[2] * tet_pressures[2]
                + barycentric[3] * tet_pressures[3]
            )
            self.fem_floor_contact_pairs[i_c].barycentric = barycentric

            C = ti.static(1.0e8)
            deformable_g = C
            rigid_g = self.fem_pressure_gradient[i_b, i_e].z
            # TODO A better way to handle corner cases where pressure and pressure gradient are ill defined
            if total_area < gs.EPS or rigid_g < gs.EPS:
                self.fem_floor_contact_pairs[i_c].active = 0
                continue
            g = 1.0 / (1.0 / deformable_g + 1.0 / rigid_g)  # harmonic average
            deformable_k = total_area * C
            rigid_k = total_area * g
            rigid_phi0 = -pressure / g
            rigid_fn0 = total_area * pressure
            # TODO custom dissipation
            self.fem_floor_contact_pairs[i_c].k = rigid_k  # contact stiffness
            self.fem_floor_contact_pairs[i_c].phi0 = rigid_phi0
            self.fem_floor_contact_pairs[i_c].fn0 = rigid_fn0
            self.fem_floor_contact_pairs[i_c].taud = self.sim._substep_dt * 1.0 / ti.math.pi

    def sap_solve(self, f):
        self.init_sap_solve(f)
        for i in range(self._n_sap_iterations):
            # init gradient and preconditioner
            self.compute_non_contact_gradient_diag(f, i)

            # compute contact hessian and gradient
            self.compute_contact_gradient_hessian_diag_prec(f)

            # solve for the vertex velocity
            self.pcg_solve()

            # line search
            self.linesearch(f)
            # TODO add convergence check

Review comment (suggestion):
`rigid_g < gs.EPS` would work just fine if you do `g = deformable_g * rigid_g * ti.rsqrt(deformable_g + rigid_g)` Anyway, it would be more robust numerically and faster to compute (even if it is not really a concern here...).

Code after:
            for i_b, i_e in ti.ndrange(fem_solver._B, fem_solver.n_elements):
                V_dt2 = fem_solver.elements_i[i_e].V * dt2
                B = fem_solver.elements_i[i_e].B
                s = -B[0, :] - B[1, :] - B[2, :]  # s is the negative sum of B rows
                p9 = ti.Vector([0.0] * 9, dt=gs.ti_float)
                i_v0, i_v1, i_v2, i_v3 = fem_solver.elements_i[i_e].el2v

                for i in ti.static(range(3)):
                    p9[i * 3 : i * 3 + 3] = (
                        B[0, i] * self.v_diff[i_b, i_v0]
                        + B[1, i] * self.v_diff[i_b, i_v1]
                        + B[2, i] * self.v_diff[i_b, i_v2]
                        + s[i] * self.v_diff[i_b, i_v3]
                    )

                new_p9 = ti.Vector([0.0] * 9, dt=gs.ti_float)

                for i in ti.static(range(3)):
                    new_p9[i * 3 : i * 3 + 3] = (
                        fem_solver.elements_el_hessian[i_b, i, 0, i_e] @ p9[0:3]
                        + fem_solver.elements_el_hessian[i_b, i, 1, i_e] @ p9[3:6]
                        + fem_solver.elements_el_hessian[i_b, i, 2, i_e] @ p9[6:9]
                    )

                # atomic
                self.gradient[i_b, i_v0] += (
                    (B[0, 0] * new_p9[0:3] + B[0, 1] * new_p9[3:6] + B[0, 2] * new_p9[6:9])
                    * V_dt2
                    * damping_beta_factor
                )
                self.gradient[i_b, i_v1] += (
                    (B[1, 0] * new_p9[0:3] + B[1, 1] * new_p9[3:6] + B[1, 2] * new_p9[6:9])
                    * V_dt2
                    * damping_beta_factor
                )
                self.gradient[i_b, i_v2] += (
                    (B[2, 0] * new_p9[0:3] + B[2, 1] * new_p9[3:6] + B[2, 2] * new_p9[6:9])
                    * V_dt2
                    * damping_beta_factor
                )
                self.gradient[i_b, i_v3] += (
                    (s[0] * new_p9[0:3] + s[1] * new_p9[3:6] + s[2] * new_p9[6:9]) * V_dt2 * damping_beta_factor
                )

    @ti.kernel
    def compute_contact_gradient_hessian_diag_prec(self, f: ti.i32):
        pairs = ti.static(self.fem_floor_contact_pairs)
        fem_solver = self.fem_solver

        for i_c in range(self.n_fem_floor_contact_pairs[None]):
            if pairs[i_c].active == 0: