Repository: Genesis-Embodied-AI/Genesis
File: genesis/engine/materials/FEM/elastic.py (Python)

Code before:
            self.compute_energy_gradient_hessian = self.compute_energy_gradient_hessian_stable_neohookean
            self.compute_energy_gradient = self.compute_energy_gradient_stable_neohookean
            self.compute_energy = self.compute_energy_stable_neohookean
            self.hessian_invariant = False
            if model == "stable_neohooken":
                gs.logger.warning("The 'stable_neohooken' model is deprecated. Use 'stable_neohookean' instead.")
        elif model == "linear_corotated":
            self.build = self.build_linear_corotated
            self.pre_compute = self.pre_compute_linear_corotated
            self.update_stress = self.update_stress_linear_corotated
            self.compute_energy_gradient_hessian = self.compute_energy_gradient_hessian_linear_corotated
            self.compute_energy_gradient = self.compute_energy_gradient_linear_corotated
            self.compute_energy = self.compute_energy_linear_corotated
            self.hessian_static = False
        else:
            gs.raise_exception(f"Unrecognized constitutive model: {model}")

        self._model = model

    def build_linear_corotated(self, fem_solver):
        self.R = ti.field(dtype=gs.ti_mat3, shape=(fem_solver._B, fem_solver.n_elements))

    @ti.func
    def pre_compute_linear_corotated(self, J, F, i_e, i_b):
        U, S, V = ti.svd(F)
        R = U @ V.transpose()
        self.R[i_b, i_e] = R

    @ti.func
    def update_stress_linear(self, mu, lam, J, F, actu, m_dir):
        I = ti.Matrix.identity(dt=gs.ti_float, n=3)
        stress = mu * (F + F.transpose() - 2 * I) + lam * (F - I).trace() * I

        return stress

    @ti.func
    def update_stress_stable_neohookean(self, mu, lam, J, F, actu, m_dir):
        IC = F.norm_sqr()
        dJdF0 = F[:, 1].cross(F[:, 2])
        dJdF1 = F[:, 2].cross(F[:, 0])
        dJdF2 = F[:, 0].cross(F[:, 1])
        dJdF = ti.Matrix.cols([dJdF0, dJdF1, dJdF2])
        alpha = 1 + 0.75 * mu / lam
        stress = mu * (1 - 1 / (IC + 1)) * F + lam * (J - alpha) * dJdF

        return stress

    @ti.func
    def update_stress_linear_corotated(self, mu, lam, J, F, actu, m_dir):
        raise NotImplementedError("Linear corotated stress update is not implemented yet.")


Review comment (refactor):
Add comment for clarity:
```python
# Computing Polar Decomposition instead of calling `R, P = ti.polar_decompose(F)` since `P` is not needed here
```

Code after:
            self.compute_energy_gradient_hessian = self.compute_energy_gradient_hessian_stable_neohookean
            self.compute_energy_gradient = self.compute_energy_gradient_stable_neohookean
            self.compute_energy = self.compute_energy_stable_neohookean
            self.hessian_invariant = False
            if model == "stable_neohooken":
                gs.logger.warning("The 'stable_neohooken' model is deprecated. Use 'stable_neohookean' instead.")
        elif model == "linear_corotated":
            self.build = self.build_linear_corotated
            self.pre_compute = self.pre_compute_linear_corotated
            self.update_stress = self.update_stress_linear_corotated
            self.compute_energy_gradient_hessian = self.compute_energy_gradient_hessian_linear_corotated
            self.compute_energy_gradient = self.compute_energy_gradient_linear_corotated
            self.compute_energy = self.compute_energy_linear_corotated
            self.hessian_static = False
        else:
            gs.raise_exception(f"Unrecognized constitutive model: {model}")

        self._model = model

    def build_linear_corotated(self, fem_solver):
        self.R = ti.field(dtype=gs.ti_mat3, shape=(fem_solver._B, fem_solver.n_elements))

    @ti.func
    def pre_compute_linear_corotated(self, J, F, i_e, i_b):
        # Computing Polar Decomposition instead of calling `R, P = ti.polar_decompose(F)` since `P` is not needed here
        U, S, V = ti.svd(F)
        R = U @ V.transpose()
        self.R[i_b, i_e] = R

    @ti.func
    def update_stress_linear(self, mu, lam, J, F, actu, m_dir):
        I = ti.Matrix.identity(dt=gs.ti_float, n=3)
        stress = mu * (F + F.transpose() - 2 * I) + lam * (F - I).trace() * I

        return stress

    @ti.func
    def update_stress_stable_neohookean(self, mu, lam, J, F, actu, m_dir):
        IC = F.norm_sqr()
        dJdF0 = F[:, 1].cross(F[:, 2])
        dJdF1 = F[:, 2].cross(F[:, 0])
        dJdF2 = F[:, 0].cross(F[:, 1])
        dJdF = ti.Matrix.cols([dJdF0, dJdF1, dJdF2])
        alpha = 1 + 0.75 * mu / lam
        stress = mu * (1 - 1 / (IC + 1)) * F + lam * (J - alpha) * dJdF

        return stress

    @ti.func
    def update_stress_linear_corotated(self, mu, lam, J, F, actu, m_dir):
        raise NotImplementedError("Linear corotated stress update is not implemented yet.")