code cleanup

Author: Libero0809

genesis/engine/couplers/sap_coupler.py

@@ -729,10 +729,10 @@ def compute_rigid_pcg_matrix_vector_product(self):
         )
 
     @ti.func
-    def compute_elastic_products(self, i_b, i_e, S, i_v, src):
+    def compute_elastic_products(self, i_b, i_e, S, i_vs, src):
         p9 = ti.Vector.zero(gs.ti_float, 9)
         for i, j in ti.static(ti.ndrange(3, 4)):
-            p9[i * 3 : i * 3 + 3] = p9[i * 3 : i * 3 + 3] + (S[j, i] * src[i_b, i_v[j]])
+            p9[i * 3 : i * 3 + 3] = p9[i * 3 : i * 3 + 3] + (S[j, i] * src[i_b, i_vs[j]])
 
         H9_p9 = ti.Vector.zero(gs.ti_float, 9)
 

genesis/engine/solvers/fem_solver.py

@@ -595,6 +595,21 @@ def compute_ele_hessian_gradient(self, f: ti.i32):
                             hessian_field=self.elements_el_hessian,
                         )
 
+    @ti.func
+    def _func_compute_element_mapping_matrix(self, i_vs, B, i_b):
+        """
+        Compute the element mapping matrix S for an element.
+        """
+        S = ti.Matrix.zero(gs.ti_float, 4, 3)
+        S[:3, :] = B
+        S[3, :] = -B[0, :] - B[1, :] - B[2, :]
+
+        if ti.static(self._enable_vertex_constraints):
+            for i in ti.static(range(4)):
+                if self.vertex_constraints.is_constrained[i_vs[i], i_b]:
+                    S[i, :] = ti.Vector.zero(gs.ti_float, 3)
+        return S
+
     @ti.func
     def _func_compute_ele_energy(self, f: ti.i32):
         """
@@ -622,21 +637,14 @@ def _func_compute_ele_energy(self, f: ti.i32):
             # add linearized damping energy
             if self._damping_beta > gs.EPS:
                 damping_beta_over_dt = self._damping_beta / self._substep_dt
-                i_v = self.elements_i[i_e].el2v
-                S = ti.Matrix.zero(gs.ti_float, 4, 3)
+                i_vs = self.elements_i[i_e].el2v
                 B = self.elements_i[i_e].B
-                S[:3, :] = B
-                S[3, :] = -B[0, :] - B[1, :] - B[2, :]
-
-                if ti.static(self._enable_vertex_constraints):
-                    for i in ti.static(range(4)):
-                        if self.vertex_constraints.is_constrained[i_v[i], i_b]:
-                            S[i, :] = ti.Vector.zero(gs.ti_float, 3)
+                S = self._func_compute_element_mapping_matrix(i_vs, B, i_b)
 
                 x_diff = ti.Vector.zero(gs.ti_float, 12)
                 for i in ti.static(range(4)):
                     x_diff[i * 3 : i * 3 + 3] = (
-                        self.elements_v[f + 1, i_v[i], i_b].pos - self.elements_v[f, i_v[i], i_b].pos
+                        self.elements_v[f + 1, i_vs[i], i_b].pos - self.elements_v[f, i_vs[i], i_b].pos
                     )
                 St_x_diff = ti.Vector.zero(gs.ti_float, 9)
                 for i, j in ti.static(ti.ndrange(3, 4)):
@@ -657,17 +665,17 @@ def accumulate_vertex_force_preconditioner(self, f: ti.i32):
         damping_beta_over_dt = self._damping_beta / self._substep_dt
         damping_beta_factor = damping_beta_over_dt + 1.0
         # inertia
-        for i_b, i_v in ti.ndrange(self._B, self.n_vertices):
+        for i_b, i_vs in ti.ndrange(self._B, self.n_vertices):
             if not self.batch_active[i_b]:
                 continue
-            self.elements_v_energy[i_b, i_v].force = -self.elements_v_info[i_v].mass_over_dt2 * (
-                (self.elements_v[f + 1, i_v, i_b].pos - self.elements_v_energy[i_b, i_v].inertia)
-                + (self.elements_v[f + 1, i_v, i_b].pos - self.elements_v[f, i_v, i_b].pos) * damping_alpha_dt
+            self.elements_v_energy[i_b, i_vs].force = -self.elements_v_info[i_vs].mass_over_dt2 * (
+                (self.elements_v[f + 1, i_vs, i_b].pos - self.elements_v_energy[i_b, i_vs].inertia)
+                + (self.elements_v[f + 1, i_vs, i_b].pos - self.elements_v[f, i_vs, i_b].pos) * damping_alpha_dt
             )
-            self.pcg_state_v[i_b, i_v].diag3x3 = ti.Matrix.zero(gs.ti_float, 3, 3)
+            self.pcg_state_v[i_b, i_vs].diag3x3 = ti.Matrix.zero(gs.ti_float, 3, 3)
             for i in ti.static(range(3)):
-                self.pcg_state_v[i_b, i_v].diag3x3[i, i] = (
-                    self.elements_v_info[i_v].mass_over_dt2 * damping_alpha_factor
+                self.pcg_state_v[i_b, i_vs].diag3x3[i, i] = (
+                    self.elements_v_info[i_vs].mass_over_dt2 * damping_alpha_factor
                 )
 
         # elastic
@@ -677,27 +685,19 @@ def accumulate_vertex_force_preconditioner(self, f: ti.i32):
             V = self.elements_i[i_e].V
             B = self.elements_i[i_e].B
             gradient = self.elements_el_energy[i_b, i_e].gradient
-            i_v = self.elements_i[i_e].el2v
-            S = ti.Matrix.zero(gs.ti_float, 4, 3)
-            S[:3, :] = B
-            S[3, :] = -B[0, :] - B[1, :] - B[2, :]
-
-            if ti.static(self._enable_vertex_constraints):
-                for i in ti.static(range(4)):
-                    if self.vertex_constraints.is_constrained[i_v[i], i_b]:
-                        S[i, :] = ti.Vector.zero(gs.ti_float, 3)
-
+            i_vs = self.elements_i[i_e].el2v
+            S = self._func_compute_element_mapping_matrix(i_vs, B, i_b)
             force = -V * gradient @ S.transpose()
 
             # atomic
             for i in ti.static(range(4)):
-                self.elements_v_energy[i_b, i_v[i]].force += force[:, i]
+                self.elements_v_energy[i_b, i_vs[i]].force += force[:, i]
 
             if self._damping_beta > gs.EPS:
                 x_diff = ti.Vector.zero(gs.ti_float, 12)
                 for i in ti.static(range(4)):
                     x_diff[i * 3 : i * 3 + 3] = (
-                        self.elements_v[f + 1, i_v[i], i_b].pos - self.elements_v[f, i_v[i], i_b].pos
+                        self.elements_v[f + 1, i_vs[i], i_b].pos - self.elements_v[f, i_vs[i], i_b].pos
                     )
                 St_x_diff = ti.Vector.zero(gs.ti_float, 9)
                 for i, j in ti.static(ti.ndrange(3, 4)):
@@ -712,22 +712,22 @@ def accumulate_vertex_force_preconditioner(self, f: ti.i32):
                 for i, j in ti.static(ti.ndrange(4, 3)):
                     S_H_St_x_diff[i * 3 : i * 3 + 3] += S[i, j] * H_St_x_diff[j * 3 : j * 3 + 3]
                 for i in ti.static(range(4)):
-                    self.elements_v_energy[i_b, i_v[i]].force += (
+                    self.elements_v_energy[i_b, i_vs[i]].force += (
                         -damping_beta_over_dt * V * S_H_St_x_diff[i * 3 : i * 3 + 3]
                     )
 
             # diagonal 3-by-3 block of hessian
             for k, i, j in ti.ndrange(4, 3, 3):
-                self.pcg_state_v[i_b, i_v[k]].diag3x3 += (
+                self.pcg_state_v[i_b, i_vs[k]].diag3x3 += (
                     V * damping_beta_factor * S[k, i] * S[k, j] * self.elements_el_hessian[i_b, i, j, i_e]
                 )
 
         # inverse
-        for i_b, i_v in ti.ndrange(self._B, self.n_vertices):
+        for i_b, i_vs in ti.ndrange(self._B, self.n_vertices):
             if not self.batch_active[i_b]:
                 continue
             # Use 3-by-3 diagonal block inverse for preconditioner
-            self.pcg_state_v[i_b, i_v].prec = self.pcg_state_v[i_b, i_v].diag3x3.inverse()
+            self.pcg_state_v[i_b, i_vs].prec = self.pcg_state_v[i_b, i_vs].diag3x3.inverse()
 
             # Other options for preconditioner:
             # Uncomment one of the following lines to test different preconditioners
@@ -745,32 +745,25 @@ def compute_Ap(self):
         damping_alpha_factor = damping_alpha_dt + 1.0
         damping_beta_over_dt = self._damping_beta / self._substep_dt
         damping_beta_factor = damping_beta_over_dt + 1.0
-        for i_b, i_v in ti.ndrange(self._B, self.n_vertices):
+        for i_b, i_vs in ti.ndrange(self._B, self.n_vertices):
             if not self.batch_pcg_active[i_b]:
                 continue
-            self.pcg_state_v[i_b, i_v].Ap = (
-                self.elements_v_info[i_v].mass_over_dt2 * damping_alpha_factor * self.pcg_state_v[i_b, i_v].p
+            self.pcg_state_v[i_b, i_vs].Ap = (
+                self.elements_v_info[i_vs].mass_over_dt2 * damping_alpha_factor * self.pcg_state_v[i_b, i_vs].p
             )
 
         for i_b, i_e in ti.ndrange(self._B, self.n_elements):
             if not self.batch_pcg_active[i_b]:
                 continue
             V = self.elements_i[i_e].V
             B = self.elements_i[i_e].B
-            i_v = self.elements_i[i_e].el2v
-            S = ti.Matrix.zero(gs.ti_float, 4, 3)
-            S[:3, :] = B
-            S[3, :] = -B[0, :] - B[1, :] - B[2, :]
-
-            if ti.static(self._enable_vertex_constraints):
-                for i in ti.static(range(4)):
-                    if self.vertex_constraints.is_constrained[i_v[i], i_b]:
-                        S[i, :] = ti.Vector.zero(gs.ti_float, 3)
+            i_vs = self.elements_i[i_e].el2v
+            S = self._func_compute_element_mapping_matrix(i_vs, B, i_b)
 
             p9 = ti.Vector([0.0] * 9, dt=gs.ti_float)
 
             for i, j in ti.static(ti.ndrange(3, 4)):
-                p9[i * 3 : i * 3 + 3] = p9[i * 3 : i * 3 + 3] + S[j, i] * self.pcg_state_v[i_b, i_v[j]].p
+                p9[i * 3 : i * 3 + 3] = p9[i * 3 : i * 3 + 3] + S[j, i] * self.pcg_state_v[i_b, i_vs[j]].p
 
             new_p9 = ti.Vector([0.0] * 9, dt=gs.ti_float)
 
@@ -781,7 +774,7 @@ def compute_Ap(self):
 
             # atomic
             for i in ti.static(range(4)):
-                self.pcg_state_v[i_b, i_v[i]].Ap += (
+                self.pcg_state_v[i_b, i_vs[i]].Ap += (
                     (S[i, 0] * new_p9[0:3] + S[i, 1] * new_p9[3:6] + S[i, 2] * new_p9[6:9]) * V * damping_beta_factor
                 )
 

genesis/options/solvers.py

@@ -114,6 +114,11 @@ class SAPCouplerOptions(BaseCouplerOptions):
     """
     Options configuring the inter-solver coupling for the Semi-Analytic Primal (SAP) contact solver used in Drake.
 
+    Note
+    ----
+    Paper reference: https://arxiv.org/abs/2110.10107
+    Drake reference: https://drake.mit.edu/release_notes/v1.5.0.html
+
     Parameters
     ----------
     n_sap_iterations : int, optional
@@ -152,10 +157,6 @@ class SAPCouplerOptions(BaseCouplerOptions):
         Type of contact against the floor for rigid bodies. Defaults to "vert". Can be "vert" or "none".
     enable_rigid_fem_contact : bool, optional
         Whether to enable coupling between rigid and FEM solvers. Defaults to True.
-    Note
-    ----
-    Paper reference: https://arxiv.org/abs/2110.10107
-    Drake reference: https://drake.mit.edu/release_notes/v1.5.0.html
     """
 
     n_sap_iterations: int = 5
@@ -581,6 +582,13 @@ class FEMOptions(Options):
     """
     Options configuring the FEMSolver.
 
+    Note
+    ----
+    - Damping coefficients are used to control the damping effect in the simulation.
+    They are used in the Rayleigh Damping model, which is a common damping model in FEM simulations.
+    Reference: https://doc.comsol.com/5.5/doc/com.comsol.help.sme/sme_ug_modeling.05.083.html
+    - TODO Move it to material parameters in the future instead of solver options.
+
     Parameters
     ----------
     dt : float, optional
@@ -614,12 +622,6 @@ class FEMOptions(Options):
         Rayleigh Damping factor for the implicit solver. Defaults to 5e-4. Only used when `use_implicit_solver` is True.
     enable_vertex_constraints : bool, optional
         Whether to enable vertex constraints. Defaults to False.
-    Note
-    ----
-    - Damping coefficients are used to control the damping effect in the simulation.
-    They are used in the Rayleigh Damping model, which is a common damping model in FEM simulations.
-    Reference: https://doc.comsol.com/5.5/doc/com.comsol.help.sme/sme_ug_modeling.05.083.html
-    - TODO Move it to material parameters in the future instead of solver options.
     """
 
     dt: Optional[float] = None

tests/test_fem.py

@@ -586,8 +586,8 @@ def test_fem_articulated(fem_material_linear_corotated_soft, show_viewer):
 
 def test_implicit_hard_vertex_constraint(fem_material_linear_corotated, show_viewer):
     """
-    Test if a box with hard vertex constraints has those vertices fixed,
-    and that updating and removing constraints works correctly.
+    Test if a box with hard vertex constraints has those vertices fixed, and that updating and removing constraints
+    works correctly.
     """
     scene = gs.Scene(
         sim_options=gs.options.SimOptions(
@@ -628,7 +628,7 @@ def test_implicit_hard_vertex_constraint(fem_material_linear_corotated, show_vie
 
     positions = cube.get_state().pos[0][verts_idx]
     assert_allclose(
-        positions, initial_target_poss, tol=0.0
+        positions, initial_target_poss, tol=gs.EPS
     ), "Vertices should stay at initial target positions with hard constraints"
     new_target_poss = initial_target_poss + gs.tensor(
         [[0.1, 0.1, 0.1]],
@@ -642,7 +642,7 @@ def test_implicit_hard_vertex_constraint(fem_material_linear_corotated, show_vie
 
     positions_after_update = cube.get_state().pos[0][verts_idx]
     assert_allclose(
-        positions_after_update, new_target_poss, tol=0.0
+        positions_after_update, new_target_poss, tol=gs.EPS
     ), "Vertices should be at new target positions after updating constraints"
 
     cube.remove_vertex_constraints()
@@ -653,15 +653,30 @@ def test_implicit_hard_vertex_constraint(fem_material_linear_corotated, show_vie
         scene.step()
 
     state = cube.get_state()
-    min_pos_z = state.pos[..., 2].min()
+    center = state.pos.mean(axis=(0, 1))
+    assert_allclose(
+        center,
+        np.array([0.2, 0.13, 0.1], dtype=np.float32),
+        atol=0.2,
+        err_msg=f"Cube center {center} moves too far from [0.2, 0.13, 0.1] after removing constraints.",
+    )
+
+    velocity = state.vel.mean(axis=(0, 1))
+    assert_allclose(
+        velocity,
+        np.array([0.0, 0.0, 0.0], dtype=np.float32),
+        atol=1e-5,
+        err_msg=f"Cube velocity {velocity} should be close to zero after settling.",
+    )
+
     # The contact requires some penetration to generate enough contact force to cancel out gravity
+    min_pos_z = state.pos[..., 2].min()
     assert_allclose(min_pos_z, -2.0e-5, atol=5e-6), f"Cube minimum Z position {min_pos_z} is not close to -2.0e-5."
 
 
 def test_sphere_box_vertex_constraint(fem_material_linear_corotated, show_viewer):
     """
-    Test if a box with hard vertex constraints has those vertices fixed,
-    and that updating and removing constraints works correctly.
+    Test if a box with hard vertex constraints has those vertices fixed, and collisiong with a sphere works correctly.
     """
     scene = gs.Scene(
         sim_options=gs.options.SimOptions(
@@ -699,7 +714,6 @@ def test_sphere_box_vertex_constraint(fem_material_linear_corotated, show_viewer
     )
 
     scene.build()
-    # return
     if show_viewer:
         sphere_debug = scene.draw_debug_spheres(poss=initial_target_poss, radius=0.02, color=(1, 0, 1, 0.8))
 
@@ -711,10 +725,28 @@ def test_sphere_box_vertex_constraint(fem_material_linear_corotated, show_viewer
     pos = cube.get_state().pos
     fixed_pos = pos[0][verts_idx]
     assert_allclose(
-        fixed_pos, initial_target_poss, tol=0.0
+        fixed_pos, initial_target_poss, tol=gs.EPS
     ), "Vertices should stay at initial target positions with hard constraints"
 
-    min_sphere_pos_z = sphere.get_state().pos[..., 2].min()
+    state = sphere.get_state()
+    center = state.pos.mean(axis=(0, 1))
+    assert_allclose(
+        center,
+        np.array([0.4, 0.4, 0.1], dtype=np.float32),
+        atol=0.2,
+        err_msg=f"Sphere center {center} moved too much from initial position [0.4, 0.4, 0.1].",
+    )
+
+    # Using a larger tolerance here since the sphere is rolling, rolling friction is not accurately modeled.
+    velocity = state.vel.mean(axis=(0, 1))
+    assert_allclose(
+        velocity,
+        np.array([0.0, 0.0, 0.0], dtype=np.float32),
+        atol=0.03,
+        err_msg=f"Sphere velocity {velocity} should be close to zero after settling.",
+    )
+
+    min_sphere_pos_z = state.pos[..., 2].min()
     assert_allclose(
         min_sphere_pos_z, -1e-3, atol=2e-4
     ), f"Sphere minimum Z position {min_sphere_pos_z} is not close to cube bottom surface."