[FEATURE] Add linear_corotated elastic material for fem (#1304)

Author: Libero0809

genesis/engine/materials/FEM/base.py

@@ -57,6 +57,13 @@ def __init__(
         # will be set when added to solver
         self._idx = None
 
+    def build(self, fem_solver):
+        pass
+
+    @ti.func
+    def pre_compute(self, J, F, i_e, i_b):
+        pass
+
     @ti.func
     def update_stress(self, mu, lam, J, F, actu, m_dir):
         raise NotImplementedError
@@ -66,7 +73,7 @@ def compute_energy_gradient_hessian(self, mu, lam, J, F, actu, m_dir, i_e, i_b,
         raise NotImplementedError
 
     @ti.func
-    def compute_energy(self, mu, lam, J, F, actu, m_dir):
+    def compute_energy(self, mu, lam, J, F, actu, m_dir, i_e, i_b):
         raise NotImplementedError
 
     @property

genesis/engine/materials/FEM/elastic.py

@@ -35,6 +35,7 @@ class Elastic(Base):
         Constitutive model to use for stress computation. Options are:
         - 'linear': Linear elasticity model
         - 'stable_neohookean': A numerically stable Neo-Hookean model
+        - 'linear_corotated': Linear corotated elasticity model
         Default is 'linear'.
     """
 
@@ -63,11 +64,28 @@ def __init__(
             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)
@@ -87,6 +105,10 @@ def update_stress_stable_neohookean(self, mu, lam, J, F, actu, m_dir):
 
         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.")
+
     @ti.func
     def compute_energy_gradient_hessian_linear(self, mu, lam, J, F, actu, m_dir, i_e, i_b, hessian_field):
         """
@@ -197,7 +219,7 @@ def compute_energy_gradient_linear(self, mu, lam, J, F, actu, m_dir, i_e, i_b):
         return energy, gradient
 
     @ti.func
-    def compute_energy_linear(self, mu, lam, J, F, actu, m_dir):
+    def compute_energy_linear(self, mu, lam, J, F, actu, m_dir, i_e, i_b):
         """
         Compute the energy for linear elasticity.
 
@@ -316,7 +338,7 @@ def compute_energy_gradient_stable_neohookean(self, mu, lam, J, F, actu, m_dir,
         raise NotImplementedError("Gradient computation is not implemented for stable_neohookean model.")
 
     @ti.func
-    def compute_energy_stable_neohookean(self, mu, lam, J, F, actu, m_dir):
+    def compute_energy_stable_neohookean(self, mu, lam, J, F, actu, m_dir, i_e, i_b):
         """
         Compute the energy for the stable Neo-Hookean model.
 
@@ -354,6 +376,153 @@ def compute_energy_stable_neohookean(self, mu, lam, J, F, actu, m_dir):
 
         return energy
 
+    @ti.func
+    def compute_energy_gradient_hessian_linear_corotated(self, mu, lam, J, F, actu, m_dir, i_e, i_b, hessian_field):
+        """
+        Compute the energy, gradient, and Hessian for linear elasticity.
+
+        Parameters
+        ----------
+        mu: float
+            The first Lame parameter (shear modulus).
+        lam: float
+            The second Lame parameter (related to volume change).
+        J: float
+            The determinant of the deformation gradient F.
+        F: ti.Matrix
+            The deformation gradient matrix.
+        actu: ti.Matrix
+            The activation matrix (not used in linear elasticity).
+        m_dir: ti.Matrix
+            The material direction (not used in linear elasticity).
+        hessian_field: ti.Matrix
+            The Hessian of the energy with respect to the deformation gradient F.
+
+        Returns
+        -------
+        energy: float
+            The computed energy.
+        gradient: ti.Matrix
+            The gradient of the energy with respect to the deformation gradient F.
+
+        Notes
+        -------
+        This implementation assumes small deformations and linear stress-strain relationship.
+        It is adapted from the HOBAKv1 implementation for linear elasticity:
+        https://github.com/theodorekim/HOBAKv1/blob/main/src/Hyperelastic/Volume/LINEAR.cpp
+
+        """
+        R = self.R[i_b, i_e]
+        F_hat = R.transpose() @ F
+        # E = 1/2(F_hat + F_hat.transpose()) - I
+        eps = 0.5 * (F_hat + F_hat.transpose())
+        for i in ti.static(range(3)):
+            eps[i, i] -= 1.0
+        trEps = eps.trace()
+        energy = mu * eps.norm_sqr() + 0.5 * lam * trEps**2
+
+        gradient = 2.0 * mu * R @ eps + lam * trEps * R
+
+        # Zero out the matrix
+        for i in ti.static(ti.grouped(ti.ndrange(3, 3))):
+            hessian_field[i_b, i, i_e].fill(0.0)
+
+        # Identity part
+        for i, k in ti.static(ti.ndrange(3, 3)):
+            hessian_field[i_b, i, i, i_e][k, k] = mu
+
+        for i, j, alpha, beta in ti.static(ti.ndrange(3, 3, 3, 3)):
+            hessian_field[i_b, j, beta, i_e][i, alpha] += mu * R[i, beta] * R[alpha, j] + lam * R[alpha, beta] * R[i, j]
+
+        return energy, gradient
+
+    @ti.func
+    def compute_energy_gradient_linear_corotated(self, mu, lam, J, F, actu, m_dir, i_e, i_b):
+        """
+        Compute the energy, gradient for linear elasticity.
+
+        Parameters
+        ----------
+        mu: float
+            The first Lame parameter (shear modulus).
+        lam: float
+            The second Lame parameter (related to volume change).
+        J: float
+            The determinant of the deformation gradient F.
+        F: ti.Matrix
+            The deformation gradient matrix.
+        actu: ti.Matrix
+            The activation matrix (not used in linear elasticity).
+        m_dir: ti.Matrix
+            The material direction (not used in linear elasticity).
+
+        Returns
+        -------
+        energy: float
+            The computed energy.
+        gradient: ti.Matrix
+            The gradient of the energy with respect to the deformation gradient F.
+
+        Notes
+        -------
+        This implementation assumes small deformations and linear stress-strain relationship.
+        It is adapted from the HOBAKv1 implementation for linear elasticity:
+        https://github.com/theodorekim/HOBAKv1/blob/main/src/Hyperelastic/Volume/LINEAR.cpp
+
+        """
+        F_hat = self.R[i_b, i_e].transpose() @ F
+        # E = 1/2(F_hat + F_hat.transpose()) - I
+        eps = 0.5 * (F_hat + F_hat.transpose())
+        for i in ti.static(range(3)):
+            eps[i, i] -= 1.0
+        trEps = eps.trace()
+        energy = mu * eps.norm_sqr() + 0.5 * lam * trEps**2
+        gradient = 2.0 * mu * self.R[i_b, i_e] @ eps + lam * trEps * self.R[i_b, i_e]
+
+        return energy, gradient
+
+    @ti.func
+    def compute_energy_linear_corotated(self, mu, lam, J, F, actu, m_dir, i_e, i_b):
+        """
+        Compute the energy for linear elasticity.
+
+        Parameters
+        ----------
+        mu: float
+            The first Lame parameter (shear modulus).
+        lam: float
+            The second Lame parameter (related to volume change).
+        J: float
+            The determinant of the deformation gradient F.
+        F: ti.Matrix
+            The deformation gradient matrix.
+        actu: ti.Matrix
+            The activation matrix (not used in linear elasticity).
+        m_dir: ti.Matrix
+            The material direction (not used in linear elasticity).
+
+        Returns
+        -------
+        energy: float
+            The computed energy.
+
+        Notes
+        -------
+        This implementation assumes small deformations and linear stress-strain relationship.
+        It is adapted from the HOBAKv1 implementation for linear elasticity:
+        https://github.com/theodorekim/HOBAKv1/blob/main/src/Hyperelastic/Volume/LINEAR.cpp
+
+        """
+        F_hat = self.R[i_b, i_e].transpose() @ F
+        # E = 1/2(F_hat + F_hat.transpose()) - I
+        eps = 0.5 * (F_hat + F_hat.transpose())
+        for i in ti.static(range(3)):
+            eps[i, i] -= 1.0
+        trEps = eps.trace()
+        energy = mu * eps.norm_sqr() + 0.5 * lam * trEps**2
+
+        return energy
+
     @property
     def model(self):
         """The name of the constitutive model ('linear' or 'stable_neohookean')."""

genesis/engine/solvers/fem_solver.py

@@ -293,6 +293,9 @@ def build(self):
             for entity in self._entities:
                 entity._add_to_solver()
 
+        for mat in self._mats:
+            mat.build(self)
+
     def add_entity(self, idx, material, morph, surface):
         # add material's update methods if not matching any existing material
         exist = False
@@ -396,6 +399,18 @@ def compute_pos(self, f: ti.i32):
                 dt * self.elements_v[f + 1, i_v, i_b].vel + self.elements_v[f, i_v, i_b].pos
             )
 
+    @ti.kernel
+    def precompute_material_data(self, f: ti.i32):
+        for i_b, i_e in ti.ndrange(self._B, self.n_elements):
+            if not self.batch_active[i_b]:
+                continue
+
+            J, F = self._compute_ele_J_F(f, i_e, i_b)  # use last time step's pos to compute
+
+            for mat_idx in ti.static(self._mats_idx):
+                if self.elements_i[i_e].mat_idx == mat_idx:
+                    self._mats[mat_idx].pre_compute(J=J, F=F, i_e=i_e, i_b=i_b)
+
     @ti.kernel
     def init_pos_and_inertia(self, f: ti.i32):
         dt = self.substep_dt
@@ -411,10 +426,10 @@ def _compute_ele_J_F(self, f: ti.i32, i_e: ti.i32, i_b: ti.i32):
         Compute the determinant (J) and deformation gradient (F) for an element.
         """
         i_v0, i_v1, i_v2, i_v3 = self.elements_i[i_e].el2v
-        pos_v0 = self.elements_v[f + 1, i_v0, i_b].pos
-        pos_v1 = self.elements_v[f + 1, i_v1, i_b].pos
-        pos_v2 = self.elements_v[f + 1, i_v2, i_b].pos
-        pos_v3 = self.elements_v[f + 1, i_v3, i_b].pos
+        pos_v0 = self.elements_v[f, i_v0, i_b].pos
+        pos_v1 = self.elements_v[f, i_v1, i_b].pos
+        pos_v2 = self.elements_v[f, i_v2, i_b].pos
+        pos_v3 = self.elements_v[f, i_v3, i_b].pos
         D = ti.Matrix.cols([pos_v0 - pos_v3, pos_v1 - pos_v3, pos_v2 - pos_v3])
 
         B = self.elements_i[i_e].B
@@ -429,7 +444,7 @@ def compute_ele_hessian_gradient(self, f: ti.i32):
             if not self.batch_active[i_b]:
                 continue
 
-            J, F = self._compute_ele_J_F(f, i_e, i_b)
+            J, F = self._compute_ele_J_F(f + 1, i_e, i_b)
 
             for mat_idx in ti.static(self._mats_idx):
                 if self.elements_i[i_e].mat_idx == mat_idx:
@@ -472,7 +487,7 @@ def _func_compute_ele_energy(self, f: ti.i32):
             if not self.batch_linesearch_active[i_b]:
                 continue
 
-            J, F = self._compute_ele_J_F(f, i_e, i_b)
+            J, F = self._compute_ele_J_F(f + 1, i_e, i_b)
 
             for mat_idx in ti.static(self._mats_idx):
                 if self.elements_i[i_e].mat_idx == mat_idx:
@@ -483,6 +498,8 @@ def _func_compute_ele_energy(self, f: ti.i32):
                         F=F,
                         actu=self.elements_el[f, i_e, i_b].actu,
                         m_dir=self.elements_i[i_e].muscle_direction,
+                        i_e=i_e,
+                        i_b=i_b,
                     )
 
             # add linearized damping energy
@@ -851,6 +868,7 @@ def process_input_grad(self):
     def substep_pre_coupling(self, f):
         if self.is_active():
             if self._use_implicit_solver:
+                self.precompute_material_data(f)
                 self.init_pos_and_inertia(f)
                 self.batch_solve(f)
                 self.setup_pos_vel(f)