diff --git a/genesis/engine/bvh.py b/genesis/engine/bvh.py index 454218e265..a5f266f7a7 100644 --- a/genesis/engine/bvh.py +++ b/genesis/engine/bvh.py @@ -112,7 +112,7 @@ class LBVH(RBC): https://research.nvidia.com/sites/default/files/pubs/2012-06_Maximizing-Parallelism-in/karras2012hpg_paper.pdf """ - def __init__(self, aabb: AABB, max_n_query_result_per_aabb: int = 8): + def __init__(self, aabb: AABB, max_n_query_result_per_aabb: int = 8, n_radix_sort_groups: int = 256): self.aabbs = aabb.aabbs self.n_aabbs = aabb.n_aabbs self.n_batches = aabb.n_batches @@ -130,12 +130,18 @@ def __init__(self, aabb: AABB, max_n_query_result_per_aabb: int = 8): # Histogram for radix sort self.hist = ti.field(ti.u32, shape=(self.n_batches, 256)) # Prefix sum for histogram - self.prefix_sum = ti.field(ti.u32, shape=(self.n_batches, 256)) + self.prefix_sum = ti.field(ti.u32, shape=(self.n_batches, 256 + 1)) # Offset for radix sort self.offset = ti.field(ti.u32, shape=(self.n_batches, self.n_aabbs)) # Temporary storage for radix sort self.tmp_morton_codes = ti.field(ti.types.vector(2, ti.u32), shape=(self.n_batches, self.n_aabbs)) + self.n_radix_sort_groups = n_radix_sort_groups + self.hist_group = ti.field(ti.u32, shape=(self.n_batches, self.n_radix_sort_groups, 256 + 1)) + self.prefix_sum_group = ti.field(ti.u32, shape=(self.n_batches, self.n_radix_sort_groups + 1, 256)) + self.group_size = self.n_aabbs // self.n_radix_sort_groups + self.visited = ti.field(ti.u8, shape=(self.n_aabbs,)) + @ti.dataclass class Node: """ @@ -176,6 +182,19 @@ def build(self): self.build_radix_tree() self.compute_bounds() + @ti.func + def filter(self, i_a, i_q): + """ + Filter function that always returns False. + + This function does not filter out any AABB by default. + It can be overridden in subclasses to implement custom filtering logic. + + i_a: index of the found AABB + i_q: index of the query AABB + """ + return False + @ti.kernel def compute_aabb_centers_and_scales(self): for i_b, i_a in ti.ndrange(self.n_batches, self.n_aabbs): @@ -230,8 +249,12 @@ def radix_sort_morton_codes(self): """ Radix sort the morton codes, using 8 bits at a time. """ - for i in range(8): - self._kernel_radix_sort_morton_codes_one_round(i) + # The last 32 bits are the index of the AABB which are already sorted, no need to sort + for i in range(4, 8): + if self.n_radix_sort_groups == 1: + self._kernel_radix_sort_morton_codes_one_round(i) + else: + self._kernel_radix_sort_morton_codes_one_round_group(i) @ti.kernel def _kernel_radix_sort_morton_codes_one_round(self, i: int): @@ -262,6 +285,50 @@ def _kernel_radix_sort_morton_codes_one_round(self, i: int): for i_b, i_a in ti.ndrange(self.n_batches, self.n_aabbs): self.morton_codes[i_b, i_a] = self.tmp_morton_codes[i_b, i_a] + @ti.kernel + def _kernel_radix_sort_morton_codes_one_round_group(self, i: int): + # Clear histogram + self.hist_group.fill(0) + + # Fill histogram + for i_b, i_g in ti.ndrange(self.n_batches, self.n_radix_sort_groups): + start = i_g * self.group_size + end = ti.select( + i_g == self.n_radix_sort_groups - 1, + self.n_aabbs, + (i_g + 1) * self.group_size, + ) + for i_a in range(start, end): + code = (self.morton_codes[i_b, i_a][1 - (i // 4)] >> ((i % 4) * 8)) & 0xFF + self.offset[i_b, i_a] = self.hist_group[i_b, i_g, code] + self.hist_group[i_b, i_g, code] = self.hist_group[i_b, i_g, code] + 1 + + # Compute prefix sum + for i_b, i_c in ti.ndrange(self.n_batches, 256): + self.prefix_sum_group[i_b, 0, i_c] = 0 + for i_g in range(1, self.n_radix_sort_groups + 1): # sequential prefix sum + self.prefix_sum_group[i_b, i_g, i_c] = ( + self.prefix_sum_group[i_b, i_g - 1, i_c] + self.hist_group[i_b, i_g - 1, i_c] + ) + for i_b in range(self.n_batches): + self.prefix_sum[i_b, 0] = 0 + for i_c in range(1, 256 + 1): # sequential prefix sum + self.prefix_sum[i_b, i_c] = ( + self.prefix_sum[i_b, i_c - 1] + self.prefix_sum_group[i_b, self.n_radix_sort_groups, i_c - 1] + ) + + # Reorder morton codes + for i_b, i_a in ti.ndrange(self.n_batches, self.n_aabbs): + code = (self.morton_codes[i_b, i_a][1 - (i // 4)] >> ((i % 4) * 8)) & 0xFF + i_g = ti.min(i_a // self.group_size, self.n_radix_sort_groups - 1) + idx = ti.i32(self.prefix_sum[i_b, code] + self.prefix_sum_group[i_b, i_g, code] + self.offset[i_b, i_a]) + # Use the group prefix sum to find the correct index + self.tmp_morton_codes[i_b, idx] = self.morton_codes[i_b, i_a] + + # Swap the temporary and original morton codes + for i_b, i_a in ti.ndrange(self.n_batches, self.n_aabbs): + self.morton_codes[i_b, i_a] = self.tmp_morton_codes[i_b, i_a] + @ti.kernel def build_radix_tree(self): """ @@ -396,10 +463,13 @@ def query(self, aabbs: ti.template()): if aabbs[i_b, i_q].intersects(node.bound): # If it's a leaf node, add the AABB index to the query results if node.left == -1 and node.right == -1: + i_a = ti.i32(self.morton_codes[i_b, node_idx - (self.n_aabbs - 1)][1]) + # Check if the filter condition is met + if self.filter(i_a, i_q): + continue idx = ti.atomic_add(self.query_result_count[None], 1) if idx < self.max_n_query_results: - code = self.morton_codes[i_b, node_idx - (self.n_aabbs - 1)][1] - self.query_result[idx] = gs.ti_ivec3(i_b, ti.i32(code), i_q) # Store the AABB index + self.query_result[idx] = gs.ti_ivec3(i_b, i_a, i_q) # Store the AABB index else: # Push children onto the stack if node.right != -1: @@ -408,3 +478,35 @@ def query(self, aabbs: ti.template()): if node.left != -1: query_stack[stack_depth] = node.left stack_depth += 1 + + +@ti.data_oriented +class FEMSurfaceTetLBVH(LBVH): + """ + FEMSurfaceTetLBVH is a specialized Linear BVH for FEM surface tetrahedrals. + + It extends the LBVH class to support filtering based on FEM surface tetrahedral elements. + """ + + def __init__(self, fem_solver, aabb: AABB, max_n_query_result_per_aabb: int = 8, n_radix_sort_groups: int = 256): + super().__init__(aabb, max_n_query_result_per_aabb, n_radix_sort_groups) + self.fem_solver = fem_solver + + @ti.func + def filter(self, i_a, i_q): + """ + Filter function for FEM surface tets. Filter out tet that share vertices. + + This is used to avoid self-collisions in FEM surface tets. + + i_a: index of the found AABB + i_q: index of the query AABB + """ + + result = i_a >= i_q + i_av = self.fem_solver.elements_i[self.fem_solver.surface_elements[i_a]].el2v + i_qv = self.fem_solver.elements_i[self.fem_solver.surface_elements[i_q]].el2v + for i, j in ti.static(ti.ndrange(4, 4)): + if i_av[i] == i_qv[j]: + result = True + return result diff --git a/genesis/engine/coupler.py b/genesis/engine/coupler.py deleted file mode 100644 index 564b283acf..0000000000 --- a/genesis/engine/coupler.py +++ /dev/null @@ -1,1511 +0,0 @@ -from typing import TYPE_CHECKING - -import numpy as np -import taichi as ti - -import genesis as gs -from genesis.options.solvers import CouplerOptions, SAPCouplerOptions -from genesis.repr_base import RBC - -if TYPE_CHECKING: - from genesis.engine.simulator import Simulator - - -@ti.data_oriented -class Coupler(RBC): - """ - This class handles all the coupling between different solvers. - """ - - # ------------------------------------------------------------------------------------ - # --------------------------------- Initialization ----------------------------------- - # ------------------------------------------------------------------------------------ - - def __init__( - self, - simulator: "Simulator", - options: "CouplerOptions", - ) -> None: - self.sim = simulator - self.options = options - - self.tool_solver = self.sim.tool_solver - self.rigid_solver = self.sim.rigid_solver - self.avatar_solver = self.sim.avatar_solver - self.mpm_solver = self.sim.mpm_solver - self.sph_solver = self.sim.sph_solver - self.pbd_solver = self.sim.pbd_solver - self.fem_solver = self.sim.fem_solver - self.sf_solver = self.sim.sf_solver - - def build(self) -> None: - self._rigid_mpm = self.rigid_solver.is_active() and self.mpm_solver.is_active() and self.options.rigid_mpm - self._rigid_sph = self.rigid_solver.is_active() and self.sph_solver.is_active() and self.options.rigid_sph - self._rigid_pbd = self.rigid_solver.is_active() and self.pbd_solver.is_active() and self.options.rigid_pbd - self._rigid_fem = self.rigid_solver.is_active() and self.fem_solver.is_active() and self.options.rigid_fem - self._mpm_sph = self.mpm_solver.is_active() and self.sph_solver.is_active() and self.options.mpm_sph - self._mpm_pbd = self.mpm_solver.is_active() and self.pbd_solver.is_active() and self.options.mpm_pbd - self._fem_mpm = self.fem_solver.is_active() and self.mpm_solver.is_active() and self.options.fem_mpm - self._fem_sph = self.fem_solver.is_active() and self.sph_solver.is_active() and self.options.fem_sph - - if self._rigid_mpm and self.mpm_solver.enable_CPIC: - # this field stores the geom index of the thin shell rigid object (if any) that separates particle and its surrounding grid cell - self.cpic_flag = ti.field(gs.ti_int, shape=(self.mpm_solver.n_particles, 3, 3, 3, self.mpm_solver._B)) - self.mpm_rigid_normal = ti.Vector.field( - 3, - dtype=gs.ti_float, - shape=(self.mpm_solver.n_particles, self.rigid_solver.n_geoms_, self.mpm_solver._B), - ) - - if self._rigid_sph: - self.sph_rigid_normal = ti.Vector.field( - 3, - dtype=gs.ti_float, - shape=(self.sph_solver.n_particles, self.rigid_solver.n_geoms_, self.sph_solver._B), - ) - self.sph_rigid_normal_reordered = ti.Vector.field( - 3, - dtype=gs.ti_float, - shape=(self.sph_solver.n_particles, self.rigid_solver.n_geoms_, self.sph_solver._B), - ) - - if self._rigid_pbd: - self.pbd_rigid_normal_reordered = ti.Vector.field( - 3, dtype=gs.ti_float, shape=(self.pbd_solver.n_particles, self.pbd_solver._B, self.rigid_solver.n_geoms) - ) - - if self._mpm_sph: - self.mpm_sph_stencil_size = int(np.floor(self.mpm_solver.dx / self.sph_solver.hash_grid_cell_size) + 2) - - if self._mpm_pbd: - self.mpm_pbd_stencil_size = int(np.floor(self.mpm_solver.dx / self.pbd_solver.hash_grid_cell_size) + 2) - - ## DEBUG - self._dx = 1 / 1024 - self._stencil_size = int(np.floor(self._dx / self.sph_solver.hash_grid_cell_size) + 2) - - self.reset(envs_idx=self.sim.scene._envs_idx) - - def reset(self, envs_idx=None) -> None: - if self._rigid_mpm and self.mpm_solver.enable_CPIC: - if envs_idx is None: - self.mpm_rigid_normal.fill(0) - else: - self._kernel_reset_mpm(envs_idx) - - if self._rigid_sph: - if envs_idx is None: - self.sph_rigid_normal.fill(0) - else: - self._kernel_reset_sph(envs_idx) - - @ti.kernel - def _kernel_reset_mpm(self, envs_idx: ti.types.ndarray()): - for i_p, i_g, i_b_ in ti.ndrange(self.mpm_solver.n_particles, self.rigid_solver.n_geoms, envs_idx.shape[0]): - self.mpm_rigid_normal[i_p, i_g, envs_idx[i_b_]] = 0.0 - - @ti.kernel - def _kernel_reset_sph(self, envs_idx: ti.types.ndarray()): - for i_p, i_g, i_b_ in ti.ndrange(self.sph_solver.n_particles, self.rigid_solver.n_geoms, envs_idx.shape[0]): - self.sph_rigid_normal[i_p, i_g, envs_idx[i_b_]] = 0.0 - - @ti.func - def _func_collide_with_rigid(self, f, pos_world, vel, mass, i_b): - for i_g in range(self.rigid_solver.n_geoms): - if self.rigid_solver.geoms_info[i_g].needs_coup: - vel = self._func_collide_with_rigid_geom(pos_world, vel, mass, i_g, i_b) - return vel - - @ti.func - def _func_collide_with_rigid_geom(self, pos_world, vel, mass, geom_idx, batch_idx): - g_info = self.rigid_solver.geoms_info[geom_idx] - signed_dist = self.rigid_solver.sdf.sdf_world(pos_world, geom_idx, batch_idx) - - # bigger coup_softness implies that the coupling influence extends further away from the object. - influence = ti.min(ti.exp(-signed_dist / max(1e-10, g_info.coup_softness)), 1) - - if influence > 0.1: - normal_rigid = self.rigid_solver.sdf.sdf_normal_world(pos_world, geom_idx, batch_idx) - vel = self._func_collide_in_rigid_geom(pos_world, vel, mass, normal_rigid, influence, geom_idx, batch_idx) - - return vel - - @ti.func - def _func_collide_with_rigid_geom_robust(self, pos_world, vel, mass, normal_prev, geom_idx, batch_idx): - """ - Similar to _func_collide_with_rigid_geom, but additionally handles potential side flip due to penetration. - """ - g_info = self.rigid_solver.geoms_info[geom_idx] - signed_dist = self.rigid_solver.sdf.sdf_world(pos_world, geom_idx, batch_idx) - normal_rigid = self.rigid_solver.sdf.sdf_normal_world(pos_world, geom_idx, batch_idx) - - # bigger coup_softness implies that the coupling influence extends further away from the object. - influence = ti.min(ti.exp(-signed_dist / max(1e-10, g_info.coup_softness)), 1) - - # if normal_rigid.dot(normal_prev) < 0: # side flip due to penetration - # influence = 1.0 - # normal_rigid = normal_prev - if influence > 0.1: - vel = self._func_collide_in_rigid_geom(pos_world, vel, mass, normal_rigid, influence, geom_idx, batch_idx) - - # attraction force - # if 0.001 < signed_dist < 0.01: - # vel = vel - normal_rigid * 0.1 * signed_dist - - return vel, normal_rigid - - @ti.func - def _func_collide_in_rigid_geom(self, pos_world, vel, mass, normal_rigid, influence, geom_idx, batch_idx): - """ - Resolves collision when a particle is already in collision with a rigid object. - This function assumes known normal_rigid and influence. - """ - g_info = self.rigid_solver.geoms_info[geom_idx] - vel_rigid = self.rigid_solver._func_vel_at_point(pos_world, g_info.link_idx, batch_idx) - - # v w.r.t rigid - rvel = vel - vel_rigid - rvel_normal_magnitude = rvel.dot(normal_rigid) # negative if inward - - if rvel_normal_magnitude < 0: # colliding - #################### rigid -> particle #################### - # tangential component - rvel_tan = rvel - rvel_normal_magnitude * normal_rigid - rvel_tan_norm = rvel_tan.norm(gs.EPS) - - # tangential component after friction - rvel_tan = ( - rvel_tan / rvel_tan_norm * ti.max(0, rvel_tan_norm + rvel_normal_magnitude * g_info.coup_friction) - ) - - # normal component after collision - rvel_normal = -normal_rigid * rvel_normal_magnitude * g_info.coup_restitution - - # normal + tangential component - rvel_new = rvel_tan + rvel_normal - - # apply influence - vel_old = vel - vel = vel_rigid + rvel_new * influence + rvel * (1 - influence) - - #################### particle -> rigid #################### - # Compute delta momentum and apply to rigid body. - delta_mv = mass * (vel - vel_old) - force = -delta_mv / self.rigid_solver.substep_dt - self.rigid_solver._func_apply_external_force(pos_world, force, g_info.link_idx, batch_idx) - - return vel - - @ti.func - def _func_mpm_tool(self, f, pos_world, vel, i_b): - for entity in ti.static(self.tool_solver.entities): - if ti.static(entity.material.collision): - vel = entity.collide(f, pos_world, vel, i_b) - return vel - - @ti.kernel - def mpm_grid_op(self, f: ti.i32, t: ti.f32): - """ - This combines mpm's grid_op with coupling operations. - If we decouple grid_op with coupling with different solvers, we need to run grid-level operations for each coupling pair, which is inefficient. - """ - for ii, jj, kk, i_b in ti.ndrange(*self.mpm_solver.grid_res, self.mpm_solver._B): - I = (ii, jj, kk) - if self.mpm_solver.grid[f, I, i_b].mass > gs.EPS: - #################### MPM grid op #################### - # Momentum to velocity - vel_mpm = (1 / self.mpm_solver.grid[f, I, i_b].mass) * self.mpm_solver.grid[f, I, i_b].vel_in - - # gravity - vel_mpm += self.mpm_solver.substep_dt * self.mpm_solver._gravity[i_b] - - pos = (I + self.mpm_solver.grid_offset) * self.mpm_solver.dx - mass_mpm = self.mpm_solver.grid[f, I, i_b].mass / self.mpm_solver._p_vol_scale - - # external force fields - for i_ff in ti.static(range(len(self.mpm_solver._ffs))): - vel_mpm += self.mpm_solver._ffs[i_ff].get_acc(pos, vel_mpm, t, -1) * self.mpm_solver.substep_dt - - #################### MPM <-> Tool #################### - if ti.static(self.tool_solver.is_active()): - vel_mpm = self._func_mpm_tool(f, pos, vel_mpm, i_b) - - #################### MPM <-> Rigid #################### - if ti.static(self._rigid_mpm): - vel_mpm = self._func_collide_with_rigid(f, pos, vel_mpm, mass_mpm, i_b) - - #################### MPM <-> SPH #################### - if ti.static(self._mpm_sph): - # using the lower corner of MPM cell to find the corresponding SPH base cell - base = self.sph_solver.sh.pos_to_grid(pos - 0.5 * self.mpm_solver.dx) - - # ---------- SPH -> MPM ---------- - sph_vel = ti.Vector([0.0, 0.0, 0.0]) - colliding_particles = 0 - for offset in ti.grouped( - ti.ndrange(self.mpm_sph_stencil_size, self.mpm_sph_stencil_size, self.mpm_sph_stencil_size) - ): - slot_idx = self.sph_solver.sh.grid_to_slot(base + offset) - for i in range( - self.sph_solver.sh.slot_start[slot_idx, i_b], - self.sph_solver.sh.slot_start[slot_idx, i_b] + self.sph_solver.sh.slot_size[slot_idx, i_b], - ): - if ( - ti.abs(pos - self.sph_solver.particles_reordered.pos[i, i_b]).max() - < self.mpm_solver.dx * 0.5 - ): - sph_vel += self.sph_solver.particles_reordered.vel[i, i_b] - colliding_particles += 1 - if colliding_particles > 0: - vel_old = vel_mpm - vel_mpm = sph_vel / colliding_particles - - # ---------- MPM -> SPH ---------- - delta_mv = mass_mpm * (vel_mpm - vel_old) - - for offset in ti.grouped( - ti.ndrange(self.mpm_sph_stencil_size, self.mpm_sph_stencil_size, self.mpm_sph_stencil_size) - ): - slot_idx = self.sph_solver.sh.grid_to_slot(base + offset) - for i in range( - self.sph_solver.sh.slot_start[slot_idx, i_b], - self.sph_solver.sh.slot_start[slot_idx, i_b] - + self.sph_solver.sh.slot_size[slot_idx, i_b], - ): - if ( - ti.abs(pos - self.sph_solver.particles_reordered.pos[i, i_b]).max() - < self.mpm_solver.dx * 0.5 - ): - self.sph_solver.particles_reordered[i, i_b].vel = ( - self.sph_solver.particles_reordered[i, i_b].vel - - delta_mv / self.sph_solver.particles_info_reordered[i, i_b].mass - ) - - #################### MPM <-> PBD #################### - if ti.static(self._mpm_pbd): - # using the lower corner of MPM cell to find the corresponding PBD base cell - base = self.pbd_solver.sh.pos_to_grid(pos - 0.5 * self.mpm_solver.dx) - - # ---------- PBD -> MPM ---------- - pbd_vel = ti.Vector([0.0, 0.0, 0.0]) - colliding_particles = 0 - for offset in ti.grouped( - ti.ndrange(self.mpm_pbd_stencil_size, self.mpm_pbd_stencil_size, self.mpm_pbd_stencil_size) - ): - slot_idx = self.pbd_solver.sh.grid_to_slot(base + offset) - for i in range( - self.pbd_solver.sh.slot_start[slot_idx, i_b], - self.pbd_solver.sh.slot_start[slot_idx, i_b] + self.pbd_solver.sh.slot_size[slot_idx, i_b], - ): - if ( - ti.abs(pos - self.pbd_solver.particles_reordered.pos[i, i_b]).max() - < self.mpm_solver.dx * 0.5 - ): - pbd_vel += self.pbd_solver.particles_reordered.vel[i, i_b] - colliding_particles += 1 - if colliding_particles > 0: - vel_old = vel_mpm - vel_mpm = pbd_vel / colliding_particles - - # ---------- MPM -> PBD ---------- - delta_mv = mass_mpm * (vel_mpm - vel_old) - - for offset in ti.grouped( - ti.ndrange(self.mpm_pbd_stencil_size, self.mpm_pbd_stencil_size, self.mpm_pbd_stencil_size) - ): - slot_idx = self.pbd_solver.sh.grid_to_slot(base + offset) - for i in range( - self.pbd_solver.sh.slot_start[slot_idx, i_b], - self.pbd_solver.sh.slot_start[slot_idx, i_b] - + self.pbd_solver.sh.slot_size[slot_idx, i_b], - ): - if ( - ti.abs(pos - self.pbd_solver.particles_reordered.pos[i, i_b]).max() - < self.mpm_solver.dx * 0.5 - ): - if self.pbd_solver.particles_reordered[i, i_b].free: - self.pbd_solver.particles_reordered[i, i_b].vel = ( - self.pbd_solver.particles_reordered[i, i_b].vel - - delta_mv / self.pbd_solver.particles_info_reordered[i, i_b].mass - ) - - #################### MPM boundary #################### - _, self.mpm_solver.grid[f, I, i_b].vel_out = self.mpm_solver.boundary.impose_pos_vel(pos, vel_mpm) - - @ti.kernel - def mpm_surface_to_particle(self, f: ti.i32): - for i_p, i_b in ti.ndrange(self.mpm_solver.n_particles, self.mpm_solver._B): - if self.mpm_solver.particles_ng[f, i_p, i_b].active: - for i_g in range(self.rigid_solver.n_geoms): - if self.rigid_solver.geoms_info[i_g].needs_coup: - sdf_normal = self.rigid_solver.sdf.sdf_normal_world( - self.mpm_solver.particles[f, i_p, i_b].pos, i_g, i_b - ) - # we only update the normal if the particle does not the object - if sdf_normal.dot(self.mpm_rigid_normal[i_p, i_g, i_b]) >= 0: - self.mpm_rigid_normal[i_p, i_g, i_b] = sdf_normal - - @ti.kernel - def fem_surface_force(self, f: ti.i32): - # TODO: all collisions are on vertices instead of surface and edge - for i_s, i_b in ti.ndrange(self.fem_solver.n_surfaces, self.fem_solver._B): - if self.fem_solver.surface[i_s].active: - dt = self.fem_solver.substep_dt - iel = self.fem_solver.surface[i_s].tri2el - mass = self.fem_solver.elements_i[iel].mass_scaled / self.fem_solver.vol_scale - - p1 = self.fem_solver.elements_v[f, self.fem_solver.surface[i_s].tri2v[0], i_b].pos - p2 = self.fem_solver.elements_v[f, self.fem_solver.surface[i_s].tri2v[1], i_b].pos - p3 = self.fem_solver.elements_v[f, self.fem_solver.surface[i_s].tri2v[2], i_b].pos - u = p2 - p1 - v = p3 - p1 - surface_normal = ti.math.cross(u, v) - surface_normal = surface_normal / surface_normal.norm(gs.EPS) - - # FEM <-> Rigid - if ti.static(self._rigid_fem): - # NOTE: collision only on surface vertices - for j in ti.static(range(3)): - iv = self.fem_solver.surface[i_s].tri2v[j] - vel_fem_sv = self._func_collide_with_rigid( - f, - self.fem_solver.elements_v[f, iv, i_b].pos, - self.fem_solver.elements_v[f + 1, iv, i_b].vel, - mass / 3.0, # assume element mass uniformly distributed to vertices - i_b, - ) - self.fem_solver.elements_v[f + 1, iv, i_b].vel = vel_fem_sv - - # FEM <-> MPM (interact with MPM grid instead of particles) - # NOTE: not doing this in mpm_grid_op otherwise we need to search for fem surface for each particles - # however, this function is called after mpm boundary conditions. - if ti.static(self._fem_mpm): - for j in ti.static(range(3)): - iv = self.fem_solver.surface[i_s].tri2v[j] - pos = self.fem_solver.elements_v[f, iv, i_b].pos - vel_fem_sv = self.fem_solver.elements_v[f + 1, iv, i_b].vel - mass_fem_sv = mass / 4.0 # assume element mass uniformly distributed - - # follow MPM p2g scheme - vel_mpm = ti.Vector([0.0, 0.0, 0.0]) - mass_mpm = 0.0 - mpm_base = ti.floor(pos * self.mpm_solver.inv_dx - 0.5).cast(gs.ti_int) - mpm_fx = pos * self.mpm_solver.inv_dx - mpm_base.cast(gs.ti_float) - mpm_w = [0.5 * (1.5 - mpm_fx) ** 2, 0.75 - (mpm_fx - 1.0) ** 2, 0.5 * (mpm_fx - 0.5) ** 2] - new_vel_fem_sv = vel_fem_sv - for mpm_offset in ti.static(ti.grouped(self.mpm_solver.stencil_range())): - mpm_grid_I = mpm_base - self.mpm_solver.grid_offset + mpm_offset - mpm_grid_mass = self.mpm_solver.grid[f, mpm_grid_I, i_b].mass / self.mpm_solver.p_vol_scale - - mpm_weight = ti.cast(1.0, gs.ti_float) - for d in ti.static(range(3)): - mpm_weight *= mpm_w[mpm_offset[d]][d] - - # FEM -> MPM - mpm_grid_pos = (mpm_grid_I + self.mpm_solver.grid_offset) * self.mpm_solver.dx - signed_dist = (mpm_grid_pos - pos).dot(surface_normal) - if signed_dist <= self.mpm_solver.dx: # NOTE: use dx as minimal unit for collision - vel_mpm_at_cell = mpm_weight * self.mpm_solver.grid[f, mpm_grid_I, i_b].vel_out - mass_mpm_at_cell = mpm_weight * mpm_grid_mass - - vel_mpm += vel_mpm_at_cell - mass_mpm += mass_mpm_at_cell - - if mass_mpm_at_cell > gs.EPS: - delta_mpm_vel_at_cell_unmul = ( - vel_fem_sv * mpm_weight - self.mpm_solver.grid[f, mpm_grid_I, i_b].vel_out - ) - mass_mul_at_cell = ( - mpm_grid_mass / mass_fem_sv - ) # NOTE: use un-reweighted mass instead of mass_mpm_at_cell - delta_mpm_vel_at_cell = delta_mpm_vel_at_cell_unmul * mass_mul_at_cell - self.mpm_solver.grid[f, mpm_grid_I, i_b].vel_out += delta_mpm_vel_at_cell - - new_vel_fem_sv -= delta_mpm_vel_at_cell * mass_mpm_at_cell / mass_fem_sv - - # MPM -> FEM - if mass_mpm > gs.EPS: - # delta_mv = (vel_mpm - vel_fem_sv) * mass_mpm - # delta_vel_fem_sv = delta_mv / mass_fem_sv - # self.fem_solver.elements_v[f + 1, iv].vel += delta_vel_fem_sv - self.fem_solver.elements_v[f + 1, iv, i_b].vel = new_vel_fem_sv - - # FEM <-> SPH TODO: this doesn't work well - if ti.static(self._fem_sph): - for j in ti.static(range(3)): - iv = self.fem_solver.surface[i_s].tri2v[j] - pos = self.fem_solver.elements_v[f, iv, i_b].pos - vel_fem_sv = self.fem_solver.elements_v[f + 1, iv, i_b].vel - mass_fem_sv = mass / 4.0 - - dx = self.sph_solver.hash_grid_cell_size # self._dx - stencil_size = 2 # self._stencil_size - - base = self.sph_solver.sh.pos_to_grid(pos - 0.5 * dx) - - # ---------- SPH -> FEM ---------- - sph_vel = ti.Vector([0.0, 0.0, 0.0]) - colliding_particles = 0 - for offset in ti.grouped(ti.ndrange(stencil_size, stencil_size, stencil_size)): - slot_idx = self.sph_solver.sh.grid_to_slot(base + offset) - for k in range( - self.sph_solver.sh.slot_start[slot_idx, i_b], - self.sph_solver.sh.slot_start[slot_idx, i_b] - + self.sph_solver.sh.slot_size[slot_idx, i_b], - ): - if ti.abs(pos - self.sph_solver.particles_reordered.pos[k, i_b]).max() < dx * 0.5: - sph_vel += self.sph_solver.particles_reordered.vel[k, i_b] - colliding_particles += 1 - - if colliding_particles > 0: - vel_old = vel_fem_sv - vel_fem_sv_unprojected = sph_vel / colliding_particles - vel_fem_sv = ( - vel_fem_sv_unprojected.dot(surface_normal) * surface_normal - ) # exclude tangential velocity - - # ---------- FEM -> SPH ---------- - delta_mv = mass_fem_sv * (vel_fem_sv - vel_old) - - for offset in ti.grouped(ti.ndrange(stencil_size, stencil_size, stencil_size)): - slot_idx = self.sph_solver.sh.grid_to_slot(base + offset) - for k in range( - self.sph_solver.sh.slot_start[slot_idx, i_b], - self.sph_solver.sh.slot_start[slot_idx, i_b] - + self.sph_solver.sh.slot_size[slot_idx, i_b], - ): - if ti.abs(pos - self.sph_solver.particles_reordered.pos[k, i_b]).max() < dx * 0.5: - self.sph_solver.particles_reordered[k, i_b].vel = ( - self.sph_solver.particles_reordered[k, i_b].vel - - delta_mv / self.sph_solver.particles_info_reordered[k, i_b].mass - ) - - self.fem_solver.elements_v[f + 1, iv, i_b].vel = vel_fem_sv - - # boundary condition - for j in ti.static(range(3)): - iv = self.fem_solver.surface[i_s].tri2v[j] - _, self.fem_solver.elements_v[f + 1, iv, i_b].vel = self.fem_solver.boundary.impose_pos_vel( - self.fem_solver.elements_v[f, iv, i_b].pos, self.fem_solver.elements_v[f + 1, iv, i_b].vel - ) - - def fem_hydroelastic(self, f: ti.i32): - # Floor contact - - # collision detection - self.fem_solver.floor_hydroelastic_detection(f) - - @ti.kernel - def sph_rigid(self, f: ti.i32): - for i_p, i_b in ti.ndrange(self.sph_solver._n_particles, self.sph_solver._B): - if self.sph_solver.particles_ng_reordered[i_p, i_b].active: - for i_g in range(self.rigid_solver.n_geoms): - if self.rigid_solver.geoms_info[i_g].needs_coup: - ( - self.sph_solver.particles_reordered[i_p, i_b].vel, - self.sph_rigid_normal_reordered[i_p, i_g, i_b], - ) = self._func_collide_with_rigid_geom_robust( - self.sph_solver.particles_reordered[i_p, i_b].pos, - self.sph_solver.particles_reordered[i_p, i_b].vel, - self.sph_solver.particles_info_reordered[i_p, i_b].mass, - self.sph_rigid_normal_reordered[i_p, i_g, i_b], - i_g, - i_b, - ) - - @ti.kernel - def pbd_rigid(self, f: ti.i32): - for i_p, i_b in ti.ndrange(self.pbd_solver._n_particles, self.sph_solver._B): - if self.pbd_solver.particles_ng_reordered[i_p, i_b].active: - # NOTE: Couldn't figure out a good way to handle collision with non-free particle. Such collision is not phsically plausible anyway. - for i_g in range(self.rigid_solver.n_geoms): - if self.rigid_solver.geoms_info[i_g].needs_coup: - ( - self.pbd_solver.particles_reordered[i_p, i_b].pos, - self.pbd_solver.particles_reordered[i_p, i_b].vel, - self.pbd_rigid_normal_reordered[i_p, i_b, i_g], - ) = self._func_pbd_collide_with_rigid_geom( - i_p, - self.pbd_solver.particles_reordered[i_p, i_b].pos, - self.pbd_solver.particles_reordered[i_p, i_b].vel, - self.pbd_solver.particles_info_reordered[i_p, i_b].mass, - self.pbd_rigid_normal_reordered[i_p, i_b, i_g], - i_g, - i_b, - ) - - @ti.func - def _func_pbd_collide_with_rigid_geom(self, i, pos_world, vel, mass, normal_prev, geom_idx, batch_idx): - """ - Resolves collision when a particle is already in collision with a rigid object. - This function assumes known normal_rigid and influence. - """ - g_info = self.rigid_solver.geoms_info[geom_idx] - signed_dist = self.rigid_solver.sdf.sdf_world(pos_world, geom_idx, batch_idx) - vel_rigid = self.rigid_solver._func_vel_at_point(pos_world, g_info.link_idx, batch_idx) - normal_rigid = self.rigid_solver.sdf.sdf_normal_world(pos_world, geom_idx, batch_idx) - new_pos = pos_world - if signed_dist < self.pbd_solver.particle_size / 2: # skip non-penetration particles - - rvel = vel - vel_rigid - rvel_normal_magnitude = rvel.dot(normal_rigid) # negative if inward - rvel_tan = rvel - rvel_normal_magnitude * normal_rigid - rvel_tan_norm = rvel_tan.norm(gs.EPS) - - #################### rigid -> particle #################### - stiffness = 1.0 # value in [0, 1] - friction = 0.15 - energy_loss = 0.0 # value in [0, 1] - new_pos = pos_world + stiffness * normal_rigid * (self.pbd_solver.particle_size / 2 - signed_dist) - v_norm = (new_pos - self.pbd_solver.particles_reordered[i, batch_idx].ipos) / self.pbd_solver._substep_dt - - delta_normal_magnitude = (v_norm - vel).dot(normal_rigid) - - delta_v_norm = delta_normal_magnitude * normal_rigid - vel = v_norm - - #################### particle -> rigid #################### - delta_mv = mass * delta_v_norm - force = (-delta_mv / self.rigid_solver._substep_dt) * (1 - energy_loss) - - self.rigid_solver._func_apply_external_force(pos_world, force, g_info.link_idx, batch_idx) - - return new_pos, vel, normal_rigid - - def preprocess(self, f): - # preprocess for MPM CPIC - if self.mpm_solver.is_active() and self.mpm_solver.enable_CPIC: - self.mpm_surface_to_particle(f) - - def couple(self, f): - # MPM <-> all others - if self.mpm_solver.is_active(): - self.mpm_grid_op(f, self.sim.cur_t) - - # SPH <-> Rigid - if self._rigid_sph: - self.sph_rigid(f) - - # PBD <-> Rigid - if self._rigid_pbd: - self.pbd_rigid(f) - - if self.fem_solver.is_active(): - self.fem_surface_force(f) - - def couple_grad(self, f): - if self.mpm_solver.is_active(): - self.mpm_grid_op.grad(f, self.sim.cur_t) - - if self.fem_solver.is_active(): - self.fem_surface_force.grad(f) - - @property - def active_solvers(self): - """All the active solvers managed by the scene's simulator.""" - return self.sim.active_solvers - - -@ti.func -def tet_barycentric(p, tet_vertices): - """ - Compute the barycentric coordinates of point p with respect to the tetrahedron defined by tet_vertices. - tet_vertices is a matrix of shape (3, 4) where each column is a vertex of the tetrahedron. - """ - v0 = tet_vertices[:, 0] - v1 = tet_vertices[:, 1] - v2 = tet_vertices[:, 2] - v3 = tet_vertices[:, 3] - - # Compute the vectors from the vertices to the point p - v1_p = p - v1 - v2_p = p - v2 - v3_p = p - v3 - - # Compute the volumes of the tetrahedra formed by the point and the vertices - vol_tet = ti.math.dot(v1 - v0, ti.math.cross(v2 - v0, v3 - v0)) - - # Compute the barycentric coordinates - b0 = -ti.math.dot(v1_p, ti.math.cross(v2 - v1, v3 - v1)) / vol_tet - b1 = ti.math.dot(v2_p, ti.math.cross(v3 - v2, v0 - v2)) / vol_tet - b2 = -ti.math.dot(v3_p, ti.math.cross(v0 - v3, v1 - v3)) / vol_tet - b3 = 1.0 - b0 - b1 - b2 - - return ti.Vector([b0, b1, b2, b3]) - - -@ti.data_oriented -class SAPCoupler(RBC): - """ - This class handles all the coupling between different solvers using 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 - """ - - # ------------------------------------------------------------------------------------ - # --------------------------------- Initialization ----------------------------------- - # ------------------------------------------------------------------------------------ - - def __init__( - self, - simulator: "Simulator", - options: "SAPCouplerOptions", - ) -> None: - self.sim = simulator - self.options = options - self.rigid_solver = self.sim.rigid_solver - self.fem_solver = self.sim.fem_solver - self._n_sap_iterations = options.n_sap_iterations - self._sap_threshold = options.sap_threshold - self._n_pcg_iterations = options.n_pcg_iterations - self._pcg_threshold = options.pcg_threshold - self._n_linesearch_iterations = options.n_linesearch_iterations - self._linesearch_c = options.linesearch_c - self._linesearch_tau = options.linesearch_tau - self.default_deformable_g = 1.0e8 # default deformable geometry size - - def build(self) -> None: - self._B = self.sim._B - self._rigid_fem = self.rigid_solver.is_active() and self.fem_solver.is_active() and self.options.rigid_fem - - if self.fem_solver.is_active(): - if self.fem_solver._use_implicit_solver is False: - gs.raise_exception( - "SAPCoupler requires FEM to use implicit solver. " - "Please set `use_implicit_solver=True` in FEM options." - ) - self.init_fem_fields() - - self.init_sap_fields() - self.init_pcg_fields() - self.init_linesearch_fields() - - def reset(self, envs_idx=None) -> None: - pass - - def init_fem_fields(self): - fem_solver = self.fem_solver - self.fem_pressure = ti.field(gs.ti_float, shape=(fem_solver.n_vertices)) - fem_pressure_np = np.concatenate([fem_entity.pressure_field_np for fem_entity in fem_solver.entities]) - self.fem_pressure.from_numpy(fem_pressure_np) - self.fem_pressure_gradient = ti.field(gs.ti_vec3, shape=(fem_solver._B, fem_solver.n_elements)) - self.fem_floor_contact_pair_type = ti.types.struct( - active=gs.ti_bool, # whether the contact pair is active - batch_idx=gs.ti_int, # batch index - geom_idx=gs.ti_int, # index of the FEM element - intersection_code=gs.ti_int, # intersection code for the element - distance=gs.ti_vec4, # distance vector for the element - k=gs.ti_float, # contact stiffness - phi0=gs.ti_float, # initial signed distance - fn0=gs.ti_float, # initial normal force magnitude - taud=gs.ti_float, # dissipation time scale - barycentric=gs.ti_vec4, # barycentric coordinates of the contact point - Rn=gs.ti_float, # Regularitaion for normal - Rt=gs.ti_float, # Regularitaion for tangential - vn_hat=gs.ti_float, # Stablization for normal velocity - mu=gs.ti_float, # friction coefficient - mu_hat=gs.ti_float, # friction coefficient regularized - mu_factor=gs.ti_float, # friction coefficient factor, 1/(1+mu_tilde**2) - energy=gs.ti_float, # energy - G=gs.ti_mat3, # Hessian matrix - ) - self.max_fem_floor_contact_pairs = fem_solver.n_surfaces * fem_solver._B - self.n_fem_floor_contact_pairs = ti.field(gs.ti_int, shape=()) - self.fem_floor_contact_pairs = self.fem_floor_contact_pair_type.field(shape=(self.max_fem_floor_contact_pairs,)) - - # Lookup table for marching tetrahedra edges - kMarchingTetsEdgeTable_np = np.array( - [ - [-1, -1, -1, -1], - [0, 3, 2, -1], - [0, 1, 4, -1], - [4, 3, 2, 1], - [1, 2, 5, -1], - [0, 3, 5, 1], - [0, 2, 5, 4], - [3, 5, 4, -1], - [3, 4, 5, -1], - [4, 5, 2, 0], - [1, 5, 3, 0], - [1, 5, 2, -1], - [1, 2, 3, 4], - [0, 4, 1, -1], - [0, 2, 3, -1], - [-1, -1, -1, -1], - ] - ) - self.kMarchingTetsEdgeTable = ti.field(gs.ti_ivec4, shape=kMarchingTetsEdgeTable_np.shape[0]) - self.kMarchingTetsEdgeTable.from_numpy(kMarchingTetsEdgeTable_np) - - kTetEdges_np = np.array([[0, 1], [1, 2], [2, 0], [0, 3], [1, 3], [2, 3]]) - self.kTetEdges = ti.field(gs.ti_ivec2, shape=kTetEdges_np.shape[0]) - self.kTetEdges.from_numpy(kTetEdges_np) - - def init_sap_fields(self): - self.batch_active = ti.field( - dtype=gs.ti_bool, - shape=self.sim._B, - needs_grad=False, - ) - self.v = ti.field(gs.ti_vec3, shape=(self.fem_solver._B, self.fem_solver.n_vertices)) - self.v_diff = ti.field(gs.ti_vec3, shape=(self.fem_solver._B, self.fem_solver.n_vertices)) - self.gradient = ti.field(gs.ti_vec3, shape=(self.fem_solver._B, self.fem_solver.n_vertices)) - - def init_pcg_fields(self): - self.batch_pcg_active = ti.field( - dtype=gs.ti_bool, - shape=self.sim._B, - needs_grad=False, - ) - - pcg_state = ti.types.struct( - rTr=gs.ti_float, - rTz=gs.ti_float, - rTr_new=gs.ti_float, - rTz_new=gs.ti_float, - pTAp=gs.ti_float, - alpha=gs.ti_float, - beta=gs.ti_float, - ) - - self.pcg_state = pcg_state.field( - shape=self.sim._B, - needs_grad=False, - layout=ti.Layout.SOA, - ) - - pcg_state_v = ti.types.struct( - diag3x3=gs.ti_mat3, # diagonal 3-by-3 block of the hessian - prec=gs.ti_mat3, # preconditioner - x=gs.ti_vec3, # solution vector - r=gs.ti_vec3, # residual vector - z=gs.ti_vec3, # preconditioned residual vector - p=gs.ti_vec3, # search direction vector - Ap=gs.ti_vec3, # matrix-vector product - ) - - self.pcg_state_v = pcg_state_v.field( - shape=(self.sim._B, self.fem_solver.n_vertices), - needs_grad=False, - layout=ti.Layout.SOA, - ) - - def init_linesearch_fields(self): - self.batch_linesearch_active = ti.field( - dtype=gs.ti_bool, - shape=self.sim._B, - needs_grad=False, - ) - - linesearch_state = ti.types.struct( - prev_energy=gs.ti_float, - energy=gs.ti_float, - step_size=gs.ti_float, - m=gs.ti_float, - ) - - self.linesearch_state = linesearch_state.field( - shape=self.sim._B, - needs_grad=False, - layout=ti.Layout.SOA, - ) - - linesearch_state_v = ti.types.struct( - x_prev=gs.ti_vec3, # solution vector - ) - - self.linesearch_state_v = linesearch_state_v.field( - shape=(self.sim._B, self.fem_solver.n_vertices), - needs_grad=False, - layout=ti.Layout.SOA, - ) - - def preprocess(self, f): - pass - - def couple(self, f): - self.has_contact = False - if self.fem_solver.is_active(): - self.fem_compute_pressure_gradient(f) - self.fem_floor_detection(f) - self.has_fem_floor_contact = self.n_fem_floor_contact_pairs[None] > 0 - self.has_contact = self.has_contact or self.has_fem_floor_contact - - if self.has_contact: - self.sap_solve(f) - self.update_vel(f) - - def couple_grad(self, f): - gs.raise_exception("couple_grad is not available for HydroelasticCoupler. Please use Coupler instead.") - - @ti.kernel - def update_vel(self, f: ti.i32): - fem_solver = self.fem_solver - for i_b, i_v in ti.ndrange(fem_solver._B, fem_solver.n_vertices): - self.fem_solver.elements_v[f + 1, i_v, i_b].vel = self.v[i_b, i_v] - - @ti.kernel - def fem_compute_pressure_gradient(self, f: ti.i32): - fem_solver = self.fem_solver - for i_b, i_e in ti.ndrange(fem_solver._B, fem_solver.n_elements): - grad = ti.static(self.fem_pressure_gradient) - grad[i_b, i_e].fill(0.0) - - for i in ti.static(range(4)): - i_v0 = fem_solver.elements_i[i_e].el2v[i] - i_v1 = fem_solver.elements_i[i_e].el2v[(i + 1) % 4] - i_v2 = fem_solver.elements_i[i_e].el2v[(i + 2) % 4] - i_V3 = fem_solver.elements_i[i_e].el2v[(i + 3) % 4] - pos_v0 = fem_solver.elements_v[f + 1, i_v0, i_b].pos - pos_v1 = fem_solver.elements_v[f + 1, i_v1, i_b].pos - pos_v2 = fem_solver.elements_v[f + 1, i_v2, i_b].pos - pos_v3 = fem_solver.elements_v[f + 1, i_V3, i_b].pos - - e10 = pos_v0 - pos_v1 - e12 = pos_v2 - pos_v1 - e13 = pos_v3 - pos_v1 - - area_vector = e12.cross(e13) # area vector of the triangle formed by v1, v2, v3 - signed_volume = area_vector.dot(e10) # signed volume of the tetrahedron formed by v0, v1, v2, v3 - if ti.abs(signed_volume) > gs.EPS: - grad_i = area_vector / signed_volume - grad[i_b, i_e] += grad_i * self.fem_pressure[i_v0] - - @ti.kernel - def fem_floor_detection(self, f: ti.i32): - fem_solver = self.fem_solver - - # Compute contact pairs - self.n_fem_floor_contact_pairs[None] = 0 - # TODO Check surface element only instead of all elements - for i_b, i_e in ti.ndrange(fem_solver._B, fem_solver.n_elements): - intersection_code = ti.int32(0) - distance = ti.Vector([0.0, 0.0, 0.0, 0.0]) - for i in ti.static(range(4)): - i_v = fem_solver.elements_i[i_e].el2v[i] - pos_v = fem_solver.elements_v[f + 1, i_v, i_b].pos - distance[i] = pos_v.z - fem_solver.floor_height - if distance[i] > 0: - intersection_code |= 1 << i - - # check if the element intersect with the floor - if intersection_code != 0 and intersection_code != 15: - pair_idx = ti.atomic_add(self.n_fem_floor_contact_pairs[None], 1) - if pair_idx < self.max_fem_floor_contact_pairs: - self.fem_floor_contact_pairs[pair_idx].batch_idx = i_b - self.fem_floor_contact_pairs[pair_idx].geom_idx = i_e - self.fem_floor_contact_pairs[pair_idx].intersection_code = intersection_code - self.fem_floor_contact_pairs[pair_idx].distance = distance - - # Compute data for each contact pair - for i_c in range(self.n_fem_floor_contact_pairs[None]): - pair = self.fem_floor_contact_pairs[i_c] - self.fem_floor_contact_pairs[i_c].active = True # mark the contact pair as active - i_b = pair.batch_idx - i_e = pair.geom_idx - intersection_code = pair.intersection_code - distance = pair.distance - intersected_edges = ti.static(self.kMarchingTetsEdgeTable)[intersection_code] - tet_vertices = ti.Matrix.zero(gs.ti_float, 3, 4) # 4 vertices - tet_pressures = ti.Vector.zero(gs.ti_float, 4) # pressures at the vertices - - for i in ti.static(range(4)): - i_v = fem_solver.elements_i[i_e].el2v[i] - tet_vertices[:, i] = fem_solver.elements_v[f + 1, i_v, i_b].pos - tet_pressures[i] = self.fem_pressure[i_v] - - polygon_vertices = ti.Matrix.zero(gs.ti_float, 3, 4) # 3 or 4 vertices - total_area = gs.EPS # avoid division by zero - total_area_weighted_centroid = ti.Vector([0.0, 0.0, 0.0]) - for i in range(4): - if intersected_edges[i] >= 0: - edge = ti.static(self.kTetEdges)[intersected_edges[i]] - pos_v0 = tet_vertices[:, edge[0]] - pos_v1 = tet_vertices[:, edge[1]] - d_v0 = distance[edge[0]] - d_v1 = distance[edge[1]] - t = d_v0 / (d_v0 - d_v1) - polygon_vertices[:, i] = pos_v0 + t * (pos_v1 - pos_v0) - - # Compute tirangle area and centroid - if i >= 2: - e1 = polygon_vertices[:, i - 1] - polygon_vertices[:, 0] - e2 = polygon_vertices[:, i] - polygon_vertices[:, 0] - 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 - - 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 = False - continue - g = self.default_deformable_g * rigid_g / (self.default_deformable_g + rigid_g) # harmonic average - 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 - - def init_sap_solve(self, f: ti.i32): - self.init_v(f) - self.batch_active.fill(1) - if self.has_fem_floor_contact: - self.compute_fem_floor_regularization(f) - - @ti.kernel - def init_v(self, f: ti.i32): - fem_solver = self.fem_solver - for i_b, i_v in ti.ndrange(fem_solver._B, fem_solver.n_vertices): - self.v[i_b, i_v] = fem_solver.elements_v[f + 1, i_v, i_b].vel - - @ti.kernel - def compute_fem_floor_regularization(self, f: ti.i32): - pairs = ti.static(self.fem_floor_contact_pairs) - time_step = self.sim._substep_dt - dt2_inv = 1.0 / (time_step**2) - fem_solver = self.fem_solver - - for i_c in range(self.n_fem_floor_contact_pairs[None]): - if not pairs[i_c].active: - continue - i_b = pairs[i_c].batch_idx - i_e = pairs[i_c].geom_idx - W = ti.Matrix.zero(gs.ti_float, 3, 3) - # W = sum (JA^-1J^T) - # With floor, J is Identity times the barycentric coordinates - for i in ti.static(range(4)): - i_v = fem_solver.elements_i[i_e].el2v[i] - W += pairs[i_c].barycentric[i] ** 2 * dt2_inv * fem_solver.pcg_state_v[i_b, i_v].prec - w_rms = W.norm() / 3.0 - beta = ti.static(1.0) - beta_factor = ti.static(beta**2 / (4.0 * ti.math.pi**2)) - k = pairs[i_c].k - taud = pairs[i_c].taud - Rn = max(beta_factor * w_rms, 1.0 / (time_step * k * (time_step + taud))) - sigma = ti.static(1.0e-3) - Rt = sigma * w_rms - vn_hat = -pairs[i_c].phi0 / (time_step + taud) - - pairs[i_c].Rn = Rn - pairs[i_c].Rt = Rt - pairs[i_c].vn_hat = vn_hat - pairs[i_c].mu = fem_solver.elements_i[i_e].friction_mu - pairs[i_c].mu_hat = pairs[i_c].mu * Rt / Rn - pairs[i_c].mu_factor = 1.0 / (1.0 + pairs[i_c].mu * pairs[i_c].mu_hat) - - @ti.kernel - def compute_non_contact_gradient_diag(self, f: ti.i32, iter: int): - fem_solver = self.fem_solver - dt2 = fem_solver._substep_dt**2 - damping_alpha_dt = fem_solver._damping_alpha * fem_solver._substep_dt - damping_alpha_factor = damping_alpha_dt + 1.0 - damping_beta_over_dt = fem_solver._damping_beta / fem_solver._substep_dt - damping_beta_factor = damping_beta_over_dt + 1.0 - - for i_b, i_v in ti.ndrange(fem_solver._B, fem_solver.n_vertices): - self.gradient[i_b, i_v].fill(0.0) - # was using position now using velocity, need to multiply dt^2 - self.pcg_state_v[i_b, i_v].diag3x3 = fem_solver.pcg_state_v[i_b, i_v].diag3x3 * dt2 - self.v_diff[i_b, i_v] = self.v[i_b, i_v] - fem_solver.elements_v[f + 1, i_v, i_b].vel - - # No need to do this for iter=0 because v=v* and A(v-v*) = 0 - if iter > 0: - for i_b, i_v in ti.ndrange(fem_solver._B, fem_solver.n_vertices): - self.gradient[i_b, i_v] = ( - fem_solver.elements_v_info[i_v].mass_over_dt2 * self.v_diff[i_b, i_v] * dt2 * damping_alpha_factor - ) - - 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 not pairs[i_c].active: - continue - i_b = pairs[i_c].batch_idx - i_e = pairs[i_c].geom_idx - vc = ti.Vector([0.0, 0.0, 0.0]) - # With floor, the contact frame is the same as the world frame - for i in ti.static(range(4)): - i_v = fem_solver.elements_i[i_e].el2v[i] - vc += pairs[i_c].barycentric[i] * self.v[i_b, i_v] - y = ti.Vector([0.0, 0.0, pairs[i_c].vn_hat]) - vc - y[0] /= pairs[i_c].Rt - y[1] /= pairs[i_c].Rt - y[2] /= pairs[i_c].Rn - yr = y[:2].norm(gs.EPS) - yn = y[2] - - t_hat = y[:2] / yr - contact_mode = self.compute_contact_mode(pairs[i_c].mu, pairs[i_c].mu_hat, yr, yn) - gamma = ti.Vector.zero(gs.ti_float, 3) - pairs[i_c].G.fill(0.0) - if contact_mode == 0: # Sticking - gamma = y - pairs[i_c].G[0, 0] = 1.0 / pairs[i_c].Rt - pairs[i_c].G[1, 1] = 1.0 / pairs[i_c].Rt - pairs[i_c].G[2, 2] = 1.0 / pairs[i_c].Rn - elif contact_mode == 1: # Sliding - gn = (yn + pairs[i_c].mu_hat * yr) * pairs[i_c].mu_factor - gt = pairs[i_c].mu * gn * t_hat - gamma = ti.Vector([gt[0], gt[1], gn]) - P = t_hat.outer_product(t_hat) - Pperp = ti.Matrix.identity(gs.ti_float, 2) - P - dgt_dyt = pairs[i_c].mu * (gn / yr * Pperp + pairs[i_c].mu_hat * pairs[i_c].mu_factor * P) - dgt_dyn = pairs[i_c].mu * pairs[i_c].mu_factor * t_hat - dgn_dyt = pairs[i_c].mu_hat * pairs[i_c].mu_factor * t_hat - dgn_dyn = pairs[i_c].mu_factor - - pairs[i_c].G[:2, :2] = dgt_dyt - pairs[i_c].G[:2, 2] = dgt_dyn - pairs[i_c].G[2, :2] = dgn_dyt - pairs[i_c].G[2, 2] = dgn_dyn - - pairs[i_c].G[:, :2] *= 1.0 / pairs[i_c].Rt - pairs[i_c].G[:, 2] *= 1.0 / pairs[i_c].Rn - - else: # No contact - pass - - R_gamma = gamma - R_gamma[0] *= pairs[i_c].Rt - R_gamma[1] *= pairs[i_c].Rt - R_gamma[2] *= pairs[i_c].Rn - pairs[i_c].energy = 0.5 * gamma.dot(R_gamma) - for i in ti.static(range(4)): - i_v = fem_solver.elements_i[i_e].el2v[i] - self.gradient[i_b, i_v] -= pairs[i_c].barycentric[i] * gamma - self.pcg_state_v[i_b, i_v].diag3x3 += pairs[i_c].barycentric[i] ** 2 * pairs[i_c].G - - for i_b, i_v in ti.ndrange(fem_solver._B, fem_solver.n_vertices): - if not self.batch_active[i_b]: - continue - self.pcg_state_v[i_b, i_v].prec = self.pcg_state_v[i_b, i_v].diag3x3.inverse() - - @ti.func - def compute_contact_energy(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 not pairs[i_c].active: - continue - i_b = pairs[i_c].batch_idx - if not self.batch_linesearch_active[i_b]: - continue - i_e = pairs[i_c].geom_idx - vc = ti.Vector([0.0, 0.0, 0.0]) - # With floor, the contact frame is the same as the world frame - for i in ti.static(range(4)): - i_v = fem_solver.elements_i[i_e].el2v[i] - vc += pairs[i_c].barycentric[i] * self.v[i_b, i_v] - y = ti.Vector([0.0, 0.0, pairs[i_c].vn_hat]) - vc - y[0] /= pairs[i_c].Rt - y[1] /= pairs[i_c].Rt - y[2] /= pairs[i_c].Rn - yr = y[:2].norm(gs.EPS) - yn = y[2] - - t_hat = y[:2] / yr - contact_mode = self.compute_contact_mode(pairs[i_c].mu, pairs[i_c].mu_hat, yr, yn) - gamma = ti.Vector.zero(gs.ti_float, 3) - if contact_mode == 0: # Sticking - gamma = y - elif contact_mode == 1: # Sliding - gn = (yn + pairs[i_c].mu_hat * yr) * pairs[i_c].mu_factor - gt = pairs[i_c].mu * gn * t_hat - gamma = ti.Vector([gt[0], gt[1], gn]) - else: # No contact - pass - - R_gamma = gamma - R_gamma[0] *= pairs[i_c].Rt - R_gamma[1] *= pairs[i_c].Rt - R_gamma[2] *= pairs[i_c].Rn - pairs[i_c].energy = 0.5 * gamma.dot(R_gamma) - - @ti.func - def compute_contact_mode(self, mu, mu_hat, yr, yn): - """ - Compute the contact mode based on the friction coefficients and the relative velocities. - Returns: - 0: Sticking - 1: Sliding - 2: No contact - """ - result = 2 # No contact - if yr <= mu * yn: - result = 0 # Sticking - elif -mu_hat * yr < yn and yn < yr / mu: - result = 1 # Sliding - return result - - @ti.func - def compute_Ap(self): - fem_solver = self.fem_solver - dt2 = fem_solver._substep_dt**2 - damping_alpha_dt = fem_solver._damping_alpha * fem_solver._substep_dt - damping_alpha_factor = damping_alpha_dt + 1.0 - damping_beta_over_dt = fem_solver._damping_beta / fem_solver._substep_dt - damping_beta_factor = damping_beta_over_dt + 1.0 - - for i_b, i_v in ti.ndrange(fem_solver._B, fem_solver.n_vertices): - if not self.batch_pcg_active[i_b]: - continue - self.pcg_state_v[i_b, i_v].Ap = ( - fem_solver.elements_v_info[i_v].mass_over_dt2 - * self.pcg_state_v[i_b, i_v].p - * dt2 - * damping_alpha_factor - ) - - for i_b, i_e in ti.ndrange(fem_solver._B, fem_solver.n_elements): - if not self.batch_pcg_active[i_b]: - continue - 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.pcg_state_v[i_b, i_v0].p - + B[1, i] * self.pcg_state_v[i_b, i_v1].p - + B[2, i] * self.pcg_state_v[i_b, i_v2].p - + s[i] * self.pcg_state_v[i_b, i_v3].p - ) - - 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.pcg_state_v[i_b, i_v0].Ap += ( - (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.pcg_state_v[i_b, i_v1].Ap += ( - (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.pcg_state_v[i_b, i_v2].Ap += ( - (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.pcg_state_v[i_b, i_v3].Ap += ( - (s[0] * new_p9[0:3] + s[1] * new_p9[3:6] + s[2] * new_p9[6:9]) * V_dt2 * damping_beta_factor - ) - - pairs = ti.static(self.fem_floor_contact_pairs) - for i_c in range(self.n_fem_floor_contact_pairs[None]): - if not pairs[i_c].active: - continue - i_b = pairs[i_c].batch_idx - i_e = pairs[i_c].geom_idx - - x = ti.Vector.zero(gs.ti_float, 3) - for i in ti.static(range(4)): - i_v = fem_solver.elements_i[i_e].el2v[i] - x += pairs[i_c].barycentric[i] * self.pcg_state_v[i_b, i_v].p - - x = pairs[i_c].G @ x - - for i in ti.static(range(4)): - i_v = fem_solver.elements_i[i_e].el2v[i] - self.pcg_state_v[i_b, i_v].Ap += pairs[i_c].barycentric[i] * x - - @ti.kernel - def init_pcg_solve(self): - fem_solver = self.fem_solver - for i_b in ti.ndrange(self._B): - self.batch_pcg_active[i_b] = self.batch_active[i_b] - if not self.batch_pcg_active[i_b]: - continue - self.pcg_state[i_b].rTr = 0.0 - self.pcg_state[i_b].rTz = 0.0 - for i_b, i_v in ti.ndrange(self._B, fem_solver.n_vertices): - if not self.batch_pcg_active[i_b]: - continue - self.pcg_state_v[i_b, i_v].x = 0 - self.pcg_state_v[i_b, i_v].r = -self.gradient[i_b, i_v] - self.pcg_state_v[i_b, i_v].z = self.pcg_state_v[i_b, i_v].prec @ self.pcg_state_v[i_b, i_v].r - self.pcg_state_v[i_b, i_v].p = self.pcg_state_v[i_b, i_v].z - ti.atomic_add(self.pcg_state[i_b].rTr, self.pcg_state_v[i_b, i_v].r.dot(self.pcg_state_v[i_b, i_v].r)) - ti.atomic_add(self.pcg_state[i_b].rTz, self.pcg_state_v[i_b, i_v].r.dot(self.pcg_state_v[i_b, i_v].z)) - for i_b in ti.ndrange(self._B): - if not self.batch_pcg_active[i_b]: - continue - self.batch_pcg_active[i_b] = self.pcg_state[i_b].rTr > self._pcg_threshold - - @ti.kernel - def one_pcg_iter(self): - self.compute_Ap() - - fem_solver = self.fem_solver - - # compute pTAp - for i_b in ti.ndrange(self._B): - if not self.batch_pcg_active[i_b]: - continue - self.pcg_state[i_b].pTAp = 0.0 - for i_b, i_v in ti.ndrange(self._B, fem_solver.n_vertices): - if not self.batch_pcg_active[i_b]: - continue - ti.atomic_add(self.pcg_state[i_b].pTAp, self.pcg_state_v[i_b, i_v].p.dot(self.pcg_state_v[i_b, i_v].Ap)) - - # compute alpha and update x, r, z, rTr, rTz - for i_b in ti.ndrange(self._B): - if not self.batch_pcg_active[i_b]: - continue - self.pcg_state[i_b].alpha = self.pcg_state[i_b].rTz / self.pcg_state[i_b].pTAp - self.pcg_state[i_b].rTr_new = 0.0 - self.pcg_state[i_b].rTz_new = 0.0 - for i_b, i_v in ti.ndrange(self._B, fem_solver.n_vertices): - if not self.batch_pcg_active[i_b]: - continue - self.pcg_state_v[i_b, i_v].x += self.pcg_state[i_b].alpha * self.pcg_state_v[i_b, i_v].p - self.pcg_state_v[i_b, i_v].r -= self.pcg_state[i_b].alpha * self.pcg_state_v[i_b, i_v].Ap - self.pcg_state_v[i_b, i_v].z = self.pcg_state_v[i_b, i_v].prec @ self.pcg_state_v[i_b, i_v].r - ti.atomic_add(self.pcg_state[i_b].rTr_new, self.pcg_state_v[i_b, i_v].r.dot(self.pcg_state_v[i_b, i_v].r)) - ti.atomic_add(self.pcg_state[i_b].rTz_new, self.pcg_state_v[i_b, i_v].r.dot(self.pcg_state_v[i_b, i_v].z)) - - # check convergence - for i_b in ti.ndrange(self._B): - if not self.batch_pcg_active[i_b]: - continue - self.batch_pcg_active[i_b] = self.pcg_state[i_b].rTr_new > self._pcg_threshold - - # update beta, rTr, rTz - for i_b in ti.ndrange(self._B): - if not self.batch_pcg_active[i_b]: - continue - self.pcg_state[i_b].beta = self.pcg_state[i_b].rTr_new / self.pcg_state[i_b].rTr - self.pcg_state[i_b].rTr = self.pcg_state[i_b].rTr_new - self.pcg_state[i_b].rTz = self.pcg_state[i_b].rTz_new - - # update p - for i_b, i_v in ti.ndrange(self._B, fem_solver.n_vertices): - if not self.batch_pcg_active[i_b]: - continue - self.pcg_state_v[i_b, i_v].p = ( - self.pcg_state_v[i_b, i_v].z + self.pcg_state[i_b].beta * self.pcg_state_v[i_b, i_v].p - ) - - def pcg_solve(self): - self.init_pcg_solve() - for i in range(self._n_pcg_iterations): - self.one_pcg_iter() - - @ti.func - def compute_total_energy(self, f, energy): - fem_solver = self.fem_solver - dt2 = fem_solver._substep_dt**2 - damping_alpha_dt = fem_solver._damping_alpha * fem_solver._substep_dt - damping_alpha_factor = damping_alpha_dt + 1.0 - damping_beta_over_dt = fem_solver._damping_beta / fem_solver._substep_dt - damping_beta_factor = damping_beta_over_dt + 1.0 - - for i_b in ti.ndrange(self._B): - if not self.batch_linesearch_active[i_b]: - continue - energy[i_b] = 0.0 - - # Inertia - for i_b, i_v in ti.ndrange(self._B, fem_solver.n_vertices): - if not self.batch_linesearch_active[i_b]: - continue - self.v_diff[i_b, i_v] = self.v[i_b, i_v] - fem_solver.elements_v[f + 1, i_v, i_b].vel - energy[i_b] += ( - 0.5 - * fem_solver.elements_v_info[i_v].mass_over_dt2 - * self.v_diff[i_b, i_v].dot(self.v_diff[i_b, i_v]) - * dt2 - * damping_alpha_factor - ) - - # Elastic - for i_b, i_e in ti.ndrange(self._B, fem_solver.n_elements): - if not self.batch_linesearch_active[i_b]: - continue - - 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.zero(gs.ti_float, 9) - 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] - ) - - H9_p9 = ti.Vector.zero(gs.ti_float, 9) - - for i in ti.static(range(3)): - H9_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] - ) - - energy[i_b] += 0.5 * V_dt2 * p9.dot(H9_p9) * damping_beta_factor - - # Contact - self.compute_contact_energy(f) - for i_c in range(self.n_fem_floor_contact_pairs[None]): - pair = self.fem_floor_contact_pairs[i_c] - i_b = pair.batch_idx - if not self.batch_linesearch_active[i_b] or not pair.active: - continue - energy[i_b] += pair.energy - - @ti.kernel - def init_linesearch(self, f: ti.i32): - fem_solver = self.fem_solver - for i_b in ti.ndrange(self._B): - self.batch_linesearch_active[i_b] = self.batch_active[i_b] - if not self.batch_linesearch_active[i_b]: - continue - self.linesearch_state[i_b].step_size = 1.0 / self._linesearch_tau - self.linesearch_state[i_b].m = 0.0 - - # x_prev, m - for i_b, i_v in ti.ndrange(self._B, fem_solver.n_vertices): - if not self.batch_linesearch_active[i_b]: - continue - self.linesearch_state_v[i_b, i_v].x_prev = self.v[i_b, i_v] - self.linesearch_state[i_b].m += self.pcg_state_v[i_b, i_v].x.dot(self.gradient[i_b, i_v]) - - self.compute_total_energy(f, self.linesearch_state.prev_energy) - - @ti.kernel - def one_linesearch_iter(self, f: ti.i32): - fem_solver = self.fem_solver - - # update vel - for i_b, i_v in ti.ndrange(self._B, fem_solver.n_vertices): - if not self.batch_linesearch_active[i_b]: - continue - self.v[i_b, i_v] = ( - self.linesearch_state_v[i_b, i_v].x_prev - + self.linesearch_state[i_b].step_size * self.pcg_state_v[i_b, i_v].x - ) - - self.compute_total_energy(f, self.linesearch_state.energy) - - # check condition - for i_b in ti.ndrange(self._B): - if not self.batch_linesearch_active[i_b]: - continue - self.batch_linesearch_active[i_b] = ( - self.linesearch_state[i_b].energy - > self.linesearch_state[i_b].prev_energy - + self._linesearch_c * self.linesearch_state[i_b].step_size * self.linesearch_state[i_b].m - ) - if not self.batch_linesearch_active[i_b]: - continue - self.linesearch_state[i_b].step_size *= self._linesearch_tau - - def linesearch(self, f: ti.i32): - """ - Note - ------ - https://en.wikipedia.org/wiki/Backtracking_line_search#Algorithm - """ - self.init_linesearch(f) - for i in range(self._n_linesearch_iterations): - self.one_linesearch_iter(f) - - @property - def active_solvers(self): - """All the active solvers managed by the scene's simulator.""" - return self.sim.active_solvers diff --git a/genesis/engine/couplers/__init__.py b/genesis/engine/couplers/__init__.py new file mode 100644 index 0000000000..56a748265e --- /dev/null +++ b/genesis/engine/couplers/__init__.py @@ -0,0 +1,2 @@ +from .legacy_coupler import LegacyCoupler +from .sap_coupler import SAPCoupler diff --git a/genesis/engine/couplers/legacy_coupler.py b/genesis/engine/couplers/legacy_coupler.py new file mode 100644 index 0000000000..81be696d3b --- /dev/null +++ b/genesis/engine/couplers/legacy_coupler.py @@ -0,0 +1,606 @@ +from typing import TYPE_CHECKING + +import numpy as np +import taichi as ti + +import genesis as gs +from genesis.options.solvers import LegacyCouplerOptions +from genesis.repr_base import RBC + +if TYPE_CHECKING: + from genesis.engine.simulator import Simulator + + +@ti.data_oriented +class LegacyCoupler(RBC): + """ + This class handles all the coupling between different solvers. LegacyCoupler will be deprecated in the future. + """ + + # ------------------------------------------------------------------------------------ + # --------------------------------- Initialization ----------------------------------- + # ------------------------------------------------------------------------------------ + + def __init__( + self, + simulator: "Simulator", + options: "LegacyCouplerOptions", + ) -> None: + self.sim = simulator + self.options = options + + self.tool_solver = self.sim.tool_solver + self.rigid_solver = self.sim.rigid_solver + self.avatar_solver = self.sim.avatar_solver + self.mpm_solver = self.sim.mpm_solver + self.sph_solver = self.sim.sph_solver + self.pbd_solver = self.sim.pbd_solver + self.fem_solver = self.sim.fem_solver + self.sf_solver = self.sim.sf_solver + + def build(self) -> None: + self._rigid_mpm = self.rigid_solver.is_active() and self.mpm_solver.is_active() and self.options.rigid_mpm + self._rigid_sph = self.rigid_solver.is_active() and self.sph_solver.is_active() and self.options.rigid_sph + self._rigid_pbd = self.rigid_solver.is_active() and self.pbd_solver.is_active() and self.options.rigid_pbd + self._rigid_fem = self.rigid_solver.is_active() and self.fem_solver.is_active() and self.options.rigid_fem + self._mpm_sph = self.mpm_solver.is_active() and self.sph_solver.is_active() and self.options.mpm_sph + self._mpm_pbd = self.mpm_solver.is_active() and self.pbd_solver.is_active() and self.options.mpm_pbd + self._fem_mpm = self.fem_solver.is_active() and self.mpm_solver.is_active() and self.options.fem_mpm + self._fem_sph = self.fem_solver.is_active() and self.sph_solver.is_active() and self.options.fem_sph + + if self._rigid_mpm and self.mpm_solver.enable_CPIC: + # this field stores the geom index of the thin shell rigid object (if any) that separates particle and its surrounding grid cell + self.cpic_flag = ti.field(gs.ti_int, shape=(self.mpm_solver.n_particles, 3, 3, 3, self.mpm_solver._B)) + self.mpm_rigid_normal = ti.Vector.field( + 3, + dtype=gs.ti_float, + shape=(self.mpm_solver.n_particles, self.rigid_solver.n_geoms_, self.mpm_solver._B), + ) + + if self._rigid_sph: + self.sph_rigid_normal = ti.Vector.field( + 3, + dtype=gs.ti_float, + shape=(self.sph_solver.n_particles, self.rigid_solver.n_geoms_, self.sph_solver._B), + ) + self.sph_rigid_normal_reordered = ti.Vector.field( + 3, + dtype=gs.ti_float, + shape=(self.sph_solver.n_particles, self.rigid_solver.n_geoms_, self.sph_solver._B), + ) + + if self._rigid_pbd: + self.pbd_rigid_normal_reordered = ti.Vector.field( + 3, dtype=gs.ti_float, shape=(self.pbd_solver.n_particles, self.pbd_solver._B, self.rigid_solver.n_geoms) + ) + + if self._mpm_sph: + self.mpm_sph_stencil_size = int(np.floor(self.mpm_solver.dx / self.sph_solver.hash_grid_cell_size) + 2) + + if self._mpm_pbd: + self.mpm_pbd_stencil_size = int(np.floor(self.mpm_solver.dx / self.pbd_solver.hash_grid_cell_size) + 2) + + ## DEBUG + self._dx = 1 / 1024 + self._stencil_size = int(np.floor(self._dx / self.sph_solver.hash_grid_cell_size) + 2) + + self.reset(envs_idx=self.sim.scene._envs_idx) + + def reset(self, envs_idx=None) -> None: + if self._rigid_mpm and self.mpm_solver.enable_CPIC: + if envs_idx is None: + self.mpm_rigid_normal.fill(0) + else: + self._kernel_reset_mpm(envs_idx) + + if self._rigid_sph: + if envs_idx is None: + self.sph_rigid_normal.fill(0) + else: + self._kernel_reset_sph(envs_idx) + + @ti.kernel + def _kernel_reset_mpm(self, envs_idx: ti.types.ndarray()): + for i_p, i_g, i_b_ in ti.ndrange(self.mpm_solver.n_particles, self.rigid_solver.n_geoms, envs_idx.shape[0]): + self.mpm_rigid_normal[i_p, i_g, envs_idx[i_b_]] = 0.0 + + @ti.kernel + def _kernel_reset_sph(self, envs_idx: ti.types.ndarray()): + for i_p, i_g, i_b_ in ti.ndrange(self.sph_solver.n_particles, self.rigid_solver.n_geoms, envs_idx.shape[0]): + self.sph_rigid_normal[i_p, i_g, envs_idx[i_b_]] = 0.0 + + @ti.func + def _func_collide_with_rigid(self, f, pos_world, vel, mass, i_b): + for i_g in range(self.rigid_solver.n_geoms): + if self.rigid_solver.geoms_info[i_g].needs_coup: + vel = self._func_collide_with_rigid_geom(pos_world, vel, mass, i_g, i_b) + return vel + + @ti.func + def _func_collide_with_rigid_geom(self, pos_world, vel, mass, geom_idx, batch_idx): + g_info = self.rigid_solver.geoms_info[geom_idx] + signed_dist = self.rigid_solver.sdf.sdf_world(pos_world, geom_idx, batch_idx) + + # bigger coup_softness implies that the coupling influence extends further away from the object. + influence = ti.min(ti.exp(-signed_dist / max(1e-10, g_info.coup_softness)), 1) + + if influence > 0.1: + normal_rigid = self.rigid_solver.sdf.sdf_normal_world(pos_world, geom_idx, batch_idx) + vel = self._func_collide_in_rigid_geom(pos_world, vel, mass, normal_rigid, influence, geom_idx, batch_idx) + + return vel + + @ti.func + def _func_collide_with_rigid_geom_robust(self, pos_world, vel, mass, normal_prev, geom_idx, batch_idx): + """ + Similar to _func_collide_with_rigid_geom, but additionally handles potential side flip due to penetration. + """ + g_info = self.rigid_solver.geoms_info[geom_idx] + signed_dist = self.rigid_solver.sdf.sdf_world(pos_world, geom_idx, batch_idx) + normal_rigid = self.rigid_solver.sdf.sdf_normal_world(pos_world, geom_idx, batch_idx) + + # bigger coup_softness implies that the coupling influence extends further away from the object. + influence = ti.min(ti.exp(-signed_dist / max(1e-10, g_info.coup_softness)), 1) + + # if normal_rigid.dot(normal_prev) < 0: # side flip due to penetration + # influence = 1.0 + # normal_rigid = normal_prev + if influence > 0.1: + vel = self._func_collide_in_rigid_geom(pos_world, vel, mass, normal_rigid, influence, geom_idx, batch_idx) + + # attraction force + # if 0.001 < signed_dist < 0.01: + # vel = vel - normal_rigid * 0.1 * signed_dist + + return vel, normal_rigid + + @ti.func + def _func_collide_in_rigid_geom(self, pos_world, vel, mass, normal_rigid, influence, geom_idx, batch_idx): + """ + Resolves collision when a particle is already in collision with a rigid object. + This function assumes known normal_rigid and influence. + """ + g_info = self.rigid_solver.geoms_info[geom_idx] + vel_rigid = self.rigid_solver._func_vel_at_point(pos_world, g_info.link_idx, batch_idx) + + # v w.r.t rigid + rvel = vel - vel_rigid + rvel_normal_magnitude = rvel.dot(normal_rigid) # negative if inward + + if rvel_normal_magnitude < 0: # colliding + #################### rigid -> particle #################### + # tangential component + rvel_tan = rvel - rvel_normal_magnitude * normal_rigid + rvel_tan_norm = rvel_tan.norm(gs.EPS) + + # tangential component after friction + rvel_tan = ( + rvel_tan / rvel_tan_norm * ti.max(0, rvel_tan_norm + rvel_normal_magnitude * g_info.coup_friction) + ) + + # normal component after collision + rvel_normal = -normal_rigid * rvel_normal_magnitude * g_info.coup_restitution + + # normal + tangential component + rvel_new = rvel_tan + rvel_normal + + # apply influence + vel_old = vel + vel = vel_rigid + rvel_new * influence + rvel * (1 - influence) + + #################### particle -> rigid #################### + # Compute delta momentum and apply to rigid body. + delta_mv = mass * (vel - vel_old) + force = -delta_mv / self.rigid_solver.substep_dt + self.rigid_solver._func_apply_external_force(pos_world, force, g_info.link_idx, batch_idx) + + return vel + + @ti.func + def _func_mpm_tool(self, f, pos_world, vel, i_b): + for entity in ti.static(self.tool_solver.entities): + if ti.static(entity.material.collision): + vel = entity.collide(f, pos_world, vel, i_b) + return vel + + @ti.kernel + def mpm_grid_op(self, f: ti.i32, t: ti.f32): + """ + This combines mpm's grid_op with coupling operations. + If we decouple grid_op with coupling with different solvers, we need to run grid-level operations for each coupling pair, which is inefficient. + """ + for ii, jj, kk, i_b in ti.ndrange(*self.mpm_solver.grid_res, self.mpm_solver._B): + I = (ii, jj, kk) + if self.mpm_solver.grid[f, I, i_b].mass > gs.EPS: + #################### MPM grid op #################### + # Momentum to velocity + vel_mpm = (1 / self.mpm_solver.grid[f, I, i_b].mass) * self.mpm_solver.grid[f, I, i_b].vel_in + + # gravity + vel_mpm += self.mpm_solver.substep_dt * self.mpm_solver._gravity[i_b] + + pos = (I + self.mpm_solver.grid_offset) * self.mpm_solver.dx + mass_mpm = self.mpm_solver.grid[f, I, i_b].mass / self.mpm_solver._p_vol_scale + + # external force fields + for i_ff in ti.static(range(len(self.mpm_solver._ffs))): + vel_mpm += self.mpm_solver._ffs[i_ff].get_acc(pos, vel_mpm, t, -1) * self.mpm_solver.substep_dt + + #################### MPM <-> Tool #################### + if ti.static(self.tool_solver.is_active()): + vel_mpm = self._func_mpm_tool(f, pos, vel_mpm, i_b) + + #################### MPM <-> Rigid #################### + if ti.static(self._rigid_mpm): + vel_mpm = self._func_collide_with_rigid(f, pos, vel_mpm, mass_mpm, i_b) + + #################### MPM <-> SPH #################### + if ti.static(self._mpm_sph): + # using the lower corner of MPM cell to find the corresponding SPH base cell + base = self.sph_solver.sh.pos_to_grid(pos - 0.5 * self.mpm_solver.dx) + + # ---------- SPH -> MPM ---------- + sph_vel = ti.Vector([0.0, 0.0, 0.0]) + colliding_particles = 0 + for offset in ti.grouped( + ti.ndrange(self.mpm_sph_stencil_size, self.mpm_sph_stencil_size, self.mpm_sph_stencil_size) + ): + slot_idx = self.sph_solver.sh.grid_to_slot(base + offset) + for i in range( + self.sph_solver.sh.slot_start[slot_idx, i_b], + self.sph_solver.sh.slot_start[slot_idx, i_b] + self.sph_solver.sh.slot_size[slot_idx, i_b], + ): + if ( + ti.abs(pos - self.sph_solver.particles_reordered.pos[i, i_b]).max() + < self.mpm_solver.dx * 0.5 + ): + sph_vel += self.sph_solver.particles_reordered.vel[i, i_b] + colliding_particles += 1 + if colliding_particles > 0: + vel_old = vel_mpm + vel_mpm = sph_vel / colliding_particles + + # ---------- MPM -> SPH ---------- + delta_mv = mass_mpm * (vel_mpm - vel_old) + + for offset in ti.grouped( + ti.ndrange(self.mpm_sph_stencil_size, self.mpm_sph_stencil_size, self.mpm_sph_stencil_size) + ): + slot_idx = self.sph_solver.sh.grid_to_slot(base + offset) + for i in range( + self.sph_solver.sh.slot_start[slot_idx, i_b], + self.sph_solver.sh.slot_start[slot_idx, i_b] + + self.sph_solver.sh.slot_size[slot_idx, i_b], + ): + if ( + ti.abs(pos - self.sph_solver.particles_reordered.pos[i, i_b]).max() + < self.mpm_solver.dx * 0.5 + ): + self.sph_solver.particles_reordered[i, i_b].vel = ( + self.sph_solver.particles_reordered[i, i_b].vel + - delta_mv / self.sph_solver.particles_info_reordered[i, i_b].mass + ) + + #################### MPM <-> PBD #################### + if ti.static(self._mpm_pbd): + # using the lower corner of MPM cell to find the corresponding PBD base cell + base = self.pbd_solver.sh.pos_to_grid(pos - 0.5 * self.mpm_solver.dx) + + # ---------- PBD -> MPM ---------- + pbd_vel = ti.Vector([0.0, 0.0, 0.0]) + colliding_particles = 0 + for offset in ti.grouped( + ti.ndrange(self.mpm_pbd_stencil_size, self.mpm_pbd_stencil_size, self.mpm_pbd_stencil_size) + ): + slot_idx = self.pbd_solver.sh.grid_to_slot(base + offset) + for i in range( + self.pbd_solver.sh.slot_start[slot_idx, i_b], + self.pbd_solver.sh.slot_start[slot_idx, i_b] + self.pbd_solver.sh.slot_size[slot_idx, i_b], + ): + if ( + ti.abs(pos - self.pbd_solver.particles_reordered.pos[i, i_b]).max() + < self.mpm_solver.dx * 0.5 + ): + pbd_vel += self.pbd_solver.particles_reordered.vel[i, i_b] + colliding_particles += 1 + if colliding_particles > 0: + vel_old = vel_mpm + vel_mpm = pbd_vel / colliding_particles + + # ---------- MPM -> PBD ---------- + delta_mv = mass_mpm * (vel_mpm - vel_old) + + for offset in ti.grouped( + ti.ndrange(self.mpm_pbd_stencil_size, self.mpm_pbd_stencil_size, self.mpm_pbd_stencil_size) + ): + slot_idx = self.pbd_solver.sh.grid_to_slot(base + offset) + for i in range( + self.pbd_solver.sh.slot_start[slot_idx, i_b], + self.pbd_solver.sh.slot_start[slot_idx, i_b] + + self.pbd_solver.sh.slot_size[slot_idx, i_b], + ): + if ( + ti.abs(pos - self.pbd_solver.particles_reordered.pos[i, i_b]).max() + < self.mpm_solver.dx * 0.5 + ): + if self.pbd_solver.particles_reordered[i, i_b].free: + self.pbd_solver.particles_reordered[i, i_b].vel = ( + self.pbd_solver.particles_reordered[i, i_b].vel + - delta_mv / self.pbd_solver.particles_info_reordered[i, i_b].mass + ) + + #################### MPM boundary #################### + _, self.mpm_solver.grid[f, I, i_b].vel_out = self.mpm_solver.boundary.impose_pos_vel(pos, vel_mpm) + + @ti.kernel + def mpm_surface_to_particle(self, f: ti.i32): + for i_p, i_b in ti.ndrange(self.mpm_solver.n_particles, self.mpm_solver._B): + if self.mpm_solver.particles_ng[f, i_p, i_b].active: + for i_g in range(self.rigid_solver.n_geoms): + if self.rigid_solver.geoms_info[i_g].needs_coup: + sdf_normal = self.rigid_solver.sdf.sdf_normal_world( + self.mpm_solver.particles[f, i_p, i_b].pos, i_g, i_b + ) + # we only update the normal if the particle does not the object + if sdf_normal.dot(self.mpm_rigid_normal[i_p, i_g, i_b]) >= 0: + self.mpm_rigid_normal[i_p, i_g, i_b] = sdf_normal + + @ti.kernel + def fem_surface_force(self, f: ti.i32): + # TODO: all collisions are on vertices instead of surface and edge + for i_s, i_b in ti.ndrange(self.fem_solver.n_surfaces, self.fem_solver._B): + if self.fem_solver.surface[i_s].active: + dt = self.fem_solver.substep_dt + iel = self.fem_solver.surface[i_s].tri2el + mass = self.fem_solver.elements_i[iel].mass_scaled / self.fem_solver.vol_scale + + p1 = self.fem_solver.elements_v[f, self.fem_solver.surface[i_s].tri2v[0], i_b].pos + p2 = self.fem_solver.elements_v[f, self.fem_solver.surface[i_s].tri2v[1], i_b].pos + p3 = self.fem_solver.elements_v[f, self.fem_solver.surface[i_s].tri2v[2], i_b].pos + u = p2 - p1 + v = p3 - p1 + surface_normal = ti.math.cross(u, v) + surface_normal = surface_normal / surface_normal.norm(gs.EPS) + + # FEM <-> Rigid + if ti.static(self._rigid_fem): + # NOTE: collision only on surface vertices + for j in ti.static(range(3)): + iv = self.fem_solver.surface[i_s].tri2v[j] + vel_fem_sv = self._func_collide_with_rigid( + f, + self.fem_solver.elements_v[f, iv, i_b].pos, + self.fem_solver.elements_v[f + 1, iv, i_b].vel, + mass / 3.0, # assume element mass uniformly distributed to vertices + i_b, + ) + self.fem_solver.elements_v[f + 1, iv, i_b].vel = vel_fem_sv + + # FEM <-> MPM (interact with MPM grid instead of particles) + # NOTE: not doing this in mpm_grid_op otherwise we need to search for fem surface for each particles + # however, this function is called after mpm boundary conditions. + if ti.static(self._fem_mpm): + for j in ti.static(range(3)): + iv = self.fem_solver.surface[i_s].tri2v[j] + pos = self.fem_solver.elements_v[f, iv, i_b].pos + vel_fem_sv = self.fem_solver.elements_v[f + 1, iv, i_b].vel + mass_fem_sv = mass / 4.0 # assume element mass uniformly distributed + + # follow MPM p2g scheme + vel_mpm = ti.Vector([0.0, 0.0, 0.0]) + mass_mpm = 0.0 + mpm_base = ti.floor(pos * self.mpm_solver.inv_dx - 0.5).cast(gs.ti_int) + mpm_fx = pos * self.mpm_solver.inv_dx - mpm_base.cast(gs.ti_float) + mpm_w = [0.5 * (1.5 - mpm_fx) ** 2, 0.75 - (mpm_fx - 1.0) ** 2, 0.5 * (mpm_fx - 0.5) ** 2] + new_vel_fem_sv = vel_fem_sv + for mpm_offset in ti.static(ti.grouped(self.mpm_solver.stencil_range())): + mpm_grid_I = mpm_base - self.mpm_solver.grid_offset + mpm_offset + mpm_grid_mass = self.mpm_solver.grid[f, mpm_grid_I, i_b].mass / self.mpm_solver.p_vol_scale + + mpm_weight = ti.cast(1.0, gs.ti_float) + for d in ti.static(range(3)): + mpm_weight *= mpm_w[mpm_offset[d]][d] + + # FEM -> MPM + mpm_grid_pos = (mpm_grid_I + self.mpm_solver.grid_offset) * self.mpm_solver.dx + signed_dist = (mpm_grid_pos - pos).dot(surface_normal) + if signed_dist <= self.mpm_solver.dx: # NOTE: use dx as minimal unit for collision + vel_mpm_at_cell = mpm_weight * self.mpm_solver.grid[f, mpm_grid_I, i_b].vel_out + mass_mpm_at_cell = mpm_weight * mpm_grid_mass + + vel_mpm += vel_mpm_at_cell + mass_mpm += mass_mpm_at_cell + + if mass_mpm_at_cell > gs.EPS: + delta_mpm_vel_at_cell_unmul = ( + vel_fem_sv * mpm_weight - self.mpm_solver.grid[f, mpm_grid_I, i_b].vel_out + ) + mass_mul_at_cell = ( + mpm_grid_mass / mass_fem_sv + ) # NOTE: use un-reweighted mass instead of mass_mpm_at_cell + delta_mpm_vel_at_cell = delta_mpm_vel_at_cell_unmul * mass_mul_at_cell + self.mpm_solver.grid[f, mpm_grid_I, i_b].vel_out += delta_mpm_vel_at_cell + + new_vel_fem_sv -= delta_mpm_vel_at_cell * mass_mpm_at_cell / mass_fem_sv + + # MPM -> FEM + if mass_mpm > gs.EPS: + # delta_mv = (vel_mpm - vel_fem_sv) * mass_mpm + # delta_vel_fem_sv = delta_mv / mass_fem_sv + # self.fem_solver.elements_v[f + 1, iv].vel += delta_vel_fem_sv + self.fem_solver.elements_v[f + 1, iv, i_b].vel = new_vel_fem_sv + + # FEM <-> SPH TODO: this doesn't work well + if ti.static(self._fem_sph): + for j in ti.static(range(3)): + iv = self.fem_solver.surface[i_s].tri2v[j] + pos = self.fem_solver.elements_v[f, iv, i_b].pos + vel_fem_sv = self.fem_solver.elements_v[f + 1, iv, i_b].vel + mass_fem_sv = mass / 4.0 + + dx = self.sph_solver.hash_grid_cell_size # self._dx + stencil_size = 2 # self._stencil_size + + base = self.sph_solver.sh.pos_to_grid(pos - 0.5 * dx) + + # ---------- SPH -> FEM ---------- + sph_vel = ti.Vector([0.0, 0.0, 0.0]) + colliding_particles = 0 + for offset in ti.grouped(ti.ndrange(stencil_size, stencil_size, stencil_size)): + slot_idx = self.sph_solver.sh.grid_to_slot(base + offset) + for k in range( + self.sph_solver.sh.slot_start[slot_idx, i_b], + self.sph_solver.sh.slot_start[slot_idx, i_b] + + self.sph_solver.sh.slot_size[slot_idx, i_b], + ): + if ti.abs(pos - self.sph_solver.particles_reordered.pos[k, i_b]).max() < dx * 0.5: + sph_vel += self.sph_solver.particles_reordered.vel[k, i_b] + colliding_particles += 1 + + if colliding_particles > 0: + vel_old = vel_fem_sv + vel_fem_sv_unprojected = sph_vel / colliding_particles + vel_fem_sv = ( + vel_fem_sv_unprojected.dot(surface_normal) * surface_normal + ) # exclude tangential velocity + + # ---------- FEM -> SPH ---------- + delta_mv = mass_fem_sv * (vel_fem_sv - vel_old) + + for offset in ti.grouped(ti.ndrange(stencil_size, stencil_size, stencil_size)): + slot_idx = self.sph_solver.sh.grid_to_slot(base + offset) + for k in range( + self.sph_solver.sh.slot_start[slot_idx, i_b], + self.sph_solver.sh.slot_start[slot_idx, i_b] + + self.sph_solver.sh.slot_size[slot_idx, i_b], + ): + if ti.abs(pos - self.sph_solver.particles_reordered.pos[k, i_b]).max() < dx * 0.5: + self.sph_solver.particles_reordered[k, i_b].vel = ( + self.sph_solver.particles_reordered[k, i_b].vel + - delta_mv / self.sph_solver.particles_info_reordered[k, i_b].mass + ) + + self.fem_solver.elements_v[f + 1, iv, i_b].vel = vel_fem_sv + + # boundary condition + for j in ti.static(range(3)): + iv = self.fem_solver.surface[i_s].tri2v[j] + _, self.fem_solver.elements_v[f + 1, iv, i_b].vel = self.fem_solver.boundary.impose_pos_vel( + self.fem_solver.elements_v[f, iv, i_b].pos, self.fem_solver.elements_v[f + 1, iv, i_b].vel + ) + + def fem_hydroelastic(self, f: ti.i32): + # Floor contact + + # collision detection + self.fem_solver.floor_hydroelastic_detection(f) + + @ti.kernel + def sph_rigid(self, f: ti.i32): + for i_p, i_b in ti.ndrange(self.sph_solver._n_particles, self.sph_solver._B): + if self.sph_solver.particles_ng_reordered[i_p, i_b].active: + for i_g in range(self.rigid_solver.n_geoms): + if self.rigid_solver.geoms_info[i_g].needs_coup: + ( + self.sph_solver.particles_reordered[i_p, i_b].vel, + self.sph_rigid_normal_reordered[i_p, i_g, i_b], + ) = self._func_collide_with_rigid_geom_robust( + self.sph_solver.particles_reordered[i_p, i_b].pos, + self.sph_solver.particles_reordered[i_p, i_b].vel, + self.sph_solver.particles_info_reordered[i_p, i_b].mass, + self.sph_rigid_normal_reordered[i_p, i_g, i_b], + i_g, + i_b, + ) + + @ti.kernel + def pbd_rigid(self, f: ti.i32): + for i_p, i_b in ti.ndrange(self.pbd_solver._n_particles, self.sph_solver._B): + if self.pbd_solver.particles_ng_reordered[i_p, i_b].active: + # NOTE: Couldn't figure out a good way to handle collision with non-free particle. Such collision is not phsically plausible anyway. + for i_g in range(self.rigid_solver.n_geoms): + if self.rigid_solver.geoms_info[i_g].needs_coup: + ( + self.pbd_solver.particles_reordered[i_p, i_b].pos, + self.pbd_solver.particles_reordered[i_p, i_b].vel, + self.pbd_rigid_normal_reordered[i_p, i_b, i_g], + ) = self._func_pbd_collide_with_rigid_geom( + i_p, + self.pbd_solver.particles_reordered[i_p, i_b].pos, + self.pbd_solver.particles_reordered[i_p, i_b].vel, + self.pbd_solver.particles_info_reordered[i_p, i_b].mass, + self.pbd_rigid_normal_reordered[i_p, i_b, i_g], + i_g, + i_b, + ) + + @ti.func + def _func_pbd_collide_with_rigid_geom(self, i, pos_world, vel, mass, normal_prev, geom_idx, batch_idx): + """ + Resolves collision when a particle is already in collision with a rigid object. + This function assumes known normal_rigid and influence. + """ + g_info = self.rigid_solver.geoms_info[geom_idx] + signed_dist = self.rigid_solver.sdf.sdf_world(pos_world, geom_idx, batch_idx) + vel_rigid = self.rigid_solver._func_vel_at_point(pos_world, g_info.link_idx, batch_idx) + normal_rigid = self.rigid_solver.sdf.sdf_normal_world(pos_world, geom_idx, batch_idx) + new_pos = pos_world + if signed_dist < self.pbd_solver.particle_size / 2: # skip non-penetration particles + + rvel = vel - vel_rigid + rvel_normal_magnitude = rvel.dot(normal_rigid) # negative if inward + rvel_tan = rvel - rvel_normal_magnitude * normal_rigid + rvel_tan_norm = rvel_tan.norm(gs.EPS) + + #################### rigid -> particle #################### + stiffness = 1.0 # value in [0, 1] + friction = 0.15 + energy_loss = 0.0 # value in [0, 1] + new_pos = pos_world + stiffness * normal_rigid * (self.pbd_solver.particle_size / 2 - signed_dist) + v_norm = (new_pos - self.pbd_solver.particles_reordered[i, batch_idx].ipos) / self.pbd_solver._substep_dt + + delta_normal_magnitude = (v_norm - vel).dot(normal_rigid) + + delta_v_norm = delta_normal_magnitude * normal_rigid + vel = v_norm + + #################### particle -> rigid #################### + delta_mv = mass * delta_v_norm + force = (-delta_mv / self.rigid_solver._substep_dt) * (1 - energy_loss) + + self.rigid_solver._func_apply_external_force(pos_world, force, g_info.link_idx, batch_idx) + + return new_pos, vel, normal_rigid + + def preprocess(self, f): + # preprocess for MPM CPIC + if self.mpm_solver.is_active() and self.mpm_solver.enable_CPIC: + self.mpm_surface_to_particle(f) + + def couple(self, f): + # MPM <-> all others + if self.mpm_solver.is_active(): + self.mpm_grid_op(f, self.sim.cur_t) + + # SPH <-> Rigid + if self._rigid_sph: + self.sph_rigid(f) + + # PBD <-> Rigid + if self._rigid_pbd: + self.pbd_rigid(f) + + if self.fem_solver.is_active(): + self.fem_surface_force(f) + + def couple_grad(self, f): + if self.mpm_solver.is_active(): + self.mpm_grid_op.grad(f, self.sim.cur_t) + + if self.fem_solver.is_active(): + self.fem_surface_force.grad(f) + + @property + def active_solvers(self): + """All the active solvers managed by the scene's simulator.""" + return self.sim.active_solvers diff --git a/genesis/engine/couplers/sap_coupler.py b/genesis/engine/couplers/sap_coupler.py new file mode 100644 index 0000000000..d968299e1c --- /dev/null +++ b/genesis/engine/couplers/sap_coupler.py @@ -0,0 +1,1746 @@ +from typing import TYPE_CHECKING +import math + +import numpy as np +import taichi as ti + +import genesis as gs +from genesis.options.solvers import SAPCouplerOptions +from genesis.repr_base import RBC +from genesis.engine.bvh import AABB, LBVH, FEMSurfaceTetLBVH +from genesis.constants import IntEnum + +if TYPE_CHECKING: + from genesis.engine.simulator import Simulator + +MARCHING_TETS_EDGE_TABLE = ( + (-1, -1, -1, -1), + (0, 3, 2, -1), + (0, 1, 4, -1), + (4, 3, 2, 1), + (1, 2, 5, -1), + (0, 3, 5, 1), + (0, 2, 5, 4), + (3, 5, 4, -1), + (3, 4, 5, -1), + (4, 5, 2, 0), + (1, 5, 3, 0), + (1, 5, 2, -1), + (1, 2, 3, 4), + (0, 4, 1, -1), + (0, 2, 3, -1), + (-1, -1, -1, -1), +) + +TET_EDGES = ( + (0, 1), + (1, 2), + (2, 0), + (0, 3), + (1, 3), + (2, 3), +) + +# Cosine threshold for whether two vectors are considered to be in the same direction. Set to zero for strictly positive. +COS_ANGLE_THRESHOLD = math.cos(math.pi * 5.0 / 8.0) + + +@ti.func +def tet_barycentric(p, tet_vertices): + """ + Compute the barycentric coordinates of point p with respect to the tetrahedron defined by tet_vertices. + tet_vertices is a matrix of shape (3, 4) where each column is a vertex of the tetrahedron. + """ + v0 = tet_vertices[:, 0] + v1 = tet_vertices[:, 1] + v2 = tet_vertices[:, 2] + v3 = tet_vertices[:, 3] + + # Compute the volumes of the tetrahedra formed by the point and the vertices + vol_tet_inv = 1.0 / ((v1 - v0).dot((v2 - v0).cross(v3 - v0))) + + # Compute the barycentric coordinates + b0 = (p - v1).dot((v3 - v1).cross(v2 - v1)) * vol_tet_inv + b1 = (p - v2).dot((v3 - v2).cross(v0 - v2)) * vol_tet_inv + b2 = (p - v3).dot((v1 - v3).cross(v0 - v3)) * vol_tet_inv + b3 = 1.0 - b0 - b1 - b2 + + return ti.Vector([b0, b1, b2, b3], dt=gs.ti_float) + + +@ti.data_oriented +class SAPCoupler(RBC): + """ + This class handles all the coupling between different solvers using 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 + Code reference: https://github.com/RobotLocomotion/drake/blob/d7a5096c6d0f131705c374390202ad95d0607fd4/multibody/plant/sap_driver.cc + """ + + # ------------------------------------------------------------------------------------ + # --------------------------------- Initialization ----------------------------------- + # ------------------------------------------------------------------------------------ + + def __init__( + self, + simulator: "Simulator", + options: "SAPCouplerOptions", + ) -> None: + self.sim = simulator + self.options = options + self.rigid_solver = self.sim.rigid_solver + self.fem_solver = self.sim.fem_solver + self._n_sap_iterations = options.n_sap_iterations + self._n_pcg_iterations = options.n_pcg_iterations + self._n_linesearch_iterations = options.n_linesearch_iterations + self._sap_convergence_atol = options.sap_convergence_atol + self._sap_convergence_rtol = options.sap_convergence_rtol + self._sap_taud = options.sap_taud + self._sap_beta = options.sap_beta + self._sap_sigma = options.sap_sigma + self._pcg_threshold = options.pcg_threshold + self._linesearch_ftol = options.linesearch_ftol + self._linesearch_max_step_size = options.linesearch_max_step_size + self._hydroelastic_stiffness = options.hydroelastic_stiffness + self._point_contact_stiffness = options.point_contact_stiffness + self._fem_floor_type = options.fem_floor_type + self._fem_self_tet = options.fem_self_tet + + # ------------------------------------------------------------------------------------ + # --------------------------------- Initialization ----------------------------------- + # ------------------------------------------------------------------------------------ + + def build(self) -> None: + self._B = self.sim._B + self._rigid_fem = self.rigid_solver.is_active() and self.fem_solver.is_active() and self.options.rigid_fem + self.contacts = [] + self._init_bvh() + + if self.fem_solver.is_active(): + if self.fem_solver._use_implicit_solver is False: + gs.raise_exception( + "SAPCoupler requires FEM to use implicit solver. " + "Please set `use_implicit_solver=True` in FEM options." + ) + if self._fem_floor_type == "tet" or self._fem_self_tet: + # Hydroelastic + self._init_hydroelastic_fem_fields_and_info() + + if self._fem_floor_type == "tet": + self.fem_floor_tet_contact = FEMFloorTetContact(self.sim) + self.contacts.append(self.fem_floor_tet_contact) + + if self._fem_floor_type == "vert": + self.fem_floor_vert_contact = FEMFloorVertContact(self.sim) + self.contacts.append(self.fem_floor_vert_contact) + + if self._fem_self_tet: + self.fem_self_tet_contact = FEMSelfTetContact(self.sim) + self.contacts.append(self.fem_self_tet_contact) + + self._init_sap_fields() + self._init_pcg_fields() + self._init_linesearch_fields() + + def reset(self, envs_idx=None): + pass + + def _init_hydroelastic_fem_fields_and_info(self): + self.fem_pressure = ti.field(gs.ti_float, shape=(self.fem_solver.n_vertices)) + fem_pressure_np = np.concatenate([fem_entity.pressure_field_np for fem_entity in self.fem_solver.entities]) + self.fem_pressure.from_numpy(fem_pressure_np) + self.fem_pressure_gradient = ti.field(gs.ti_vec3, shape=(self.fem_solver._B, self.fem_solver.n_elements)) + + # Lookup table for marching tetrahedra edges + self.MarchingTetsEdgeTable = ti.field(gs.ti_ivec4, shape=len(MARCHING_TETS_EDGE_TABLE)) + self.MarchingTetsEdgeTable.from_numpy(np.array(MARCHING_TETS_EDGE_TABLE, dtype=np.int32)) + + self.TetEdges = ti.field(gs.ti_ivec2, shape=len(TET_EDGES)) + self.TetEdges.from_numpy(np.array(TET_EDGES, dtype=np.int32)) + + def _init_bvh(self): + if self.fem_solver.is_active() and self._fem_self_tet: + self.fem_surface_tet_aabb = AABB(self.fem_solver._B, self.fem_solver.n_surface_elements) + self.fem_surface_tet_bvh = FEMSurfaceTetLBVH( + self.fem_solver, self.fem_surface_tet_aabb, max_n_query_result_per_aabb=32 + ) + + def _init_sap_fields(self): + self.batch_active = ti.field(dtype=ti.u1, shape=self.sim._B, needs_grad=False) + self.v = ti.field(gs.ti_vec3, shape=(self.fem_solver._B, self.fem_solver.n_vertices)) + self.v_diff = ti.field(gs.ti_vec3, shape=(self.fem_solver._B, self.fem_solver.n_vertices)) + self.gradient = ti.field(gs.ti_vec3, shape=(self.fem_solver._B, self.fem_solver.n_vertices)) + + sap_state = ti.types.struct( + gradient_norm=gs.ti_float, # norm of the gradient + momentum_norm=gs.ti_float, # norm of the momentum + impulse_norm=gs.ti_float, # norm of the impulse + ) + + self.sap_state = sap_state.field(shape=self.sim._B, needs_grad=False, layout=ti.Layout.SOA) + + sap_state_v = ti.types.struct( + impulse=gs.ti_vec3, # impulse vector + ) + + self.sap_state_v = sap_state_v.field( + shape=(self.sim._B, self.fem_solver.n_vertices), + needs_grad=False, + layout=ti.Layout.SOA, + ) + + def _init_pcg_fields(self): + self.batch_pcg_active = ti.field(dtype=ti.u1, shape=self.sim._B, needs_grad=False) + + pcg_state = ti.types.struct( + rTr=gs.ti_float, + rTz=gs.ti_float, + rTr_new=gs.ti_float, + rTz_new=gs.ti_float, + pTAp=gs.ti_float, + alpha=gs.ti_float, + beta=gs.ti_float, + ) + + self.pcg_state = pcg_state.field(shape=self.sim._B, needs_grad=False, layout=ti.Layout.SOA) + + pcg_state_v = ti.types.struct( + diag3x3=gs.ti_mat3, # diagonal 3-by-3 block of the hessian + prec=gs.ti_mat3, # preconditioner + x=gs.ti_vec3, # solution vector + r=gs.ti_vec3, # residual vector + z=gs.ti_vec3, # preconditioned residual vector + p=gs.ti_vec3, # search direction vector + Ap=gs.ti_vec3, # matrix-vector product + ) + + self.pcg_state_v = pcg_state_v.field( + shape=(self.sim._B, self.fem_solver.n_vertices), needs_grad=False, layout=ti.Layout.SOA + ) + + def _init_linesearch_fields(self): + self.batch_linesearch_active = ti.field( + dtype=ti.u1, + shape=self.sim._B, + needs_grad=False, + ) + + linesearch_state = ti.types.struct( + prev_energy=gs.ti_float, + energy=gs.ti_float, + step_size=gs.ti_float, + m=gs.ti_float, + dell_dalpha=gs.ti_float, # first derivative of the total energy w.r.t. alpha + d2ellA_dalpha2=gs.ti_float, # second derivative of the dynamic energy w.r.t. alpha + d2ell_dalpha2=gs.ti_float, # second derivative of the total energy w.r.t. alpha + dell_scale=gs.ti_float, # scale factor for the first derivative + alpha_min=gs.ti_float, # minimum stepsize value + alpha_max=gs.ti_float, # maximum stepsize value + alpha_tol=gs.ti_float, # stepsize tolerance for convergence + f_lower=gs.ti_float, # minimum f value + f_upper=gs.ti_float, # maximum f value + f=gs.ti_float, # f value + df=gs.ti_float, # f gradient + minus_dalpha=gs.ti_float, # negative stepsize + minus_dalpha_prev=gs.ti_float, # previous negative stepsize + ) + + self.linesearch_state = linesearch_state.field(shape=self.sim._B, needs_grad=False, layout=ti.Layout.SOA) + + linesearch_state_v = ti.types.struct( + x_prev=gs.ti_vec3, # solution vector + dp=gs.ti_vec3, # A @ dv + ) + + self.linesearch_state_v = linesearch_state_v.field( + shape=(self.sim._B, self.fem_solver.n_vertices), + needs_grad=False, + layout=ti.Layout.SOA, + ) + + # ------------------------------------------------------------------------------------ + # -------------------------------------- Main ---------------------------------------- + # ------------------------------------------------------------------------------------ + + def preprocess(self, f): + pass + + def couple(self, i_step): + self.has_contact = False + if self.fem_solver.is_active(): + if self._fem_floor_type == "tet" or self._fem_self_tet: + self.fem_compute_pressure_gradient(i_step) + + for contact in self.contacts: + contact.detection(i_step) + contact.update_has_contact() + self.has_contact = self.has_contact or contact.has_contact + + if self.has_contact: + self.sap_solve(i_step) + self.update_vel(i_step) + + def couple_grad(self, i_step): + gs.raise_exception("couple_grad is not available for SAPCoupler. Please use LegacyCoupler instead.") + + @ti.kernel + def update_vel(self, i_step: ti.i32): + for i_b, i_v in ti.ndrange(self.fem_solver._B, self.fem_solver.n_vertices): + self.fem_solver.elements_v[i_step + 1, i_v, i_b].vel = self.v[i_b, i_v] + + @ti.kernel + def fem_compute_pressure_gradient(self, i_step: ti.i32): + for i_b, i_e in ti.ndrange(self.fem_solver._B, self.fem_solver.n_elements): + grad = ti.static(self.fem_pressure_gradient) + grad[i_b, i_e].fill(0.0) + + for i in ti.static(range(4)): + i_v0 = self.fem_solver.elements_i[i_e].el2v[i] + i_v1 = self.fem_solver.elements_i[i_e].el2v[(i + 1) % 4] + i_v2 = self.fem_solver.elements_i[i_e].el2v[(i + 2) % 4] + i_v3 = self.fem_solver.elements_i[i_e].el2v[(i + 3) % 4] + pos_v0 = self.fem_solver.elements_v[i_step, i_v0, i_b].pos + pos_v1 = self.fem_solver.elements_v[i_step, i_v1, i_b].pos + pos_v2 = self.fem_solver.elements_v[i_step, i_v2, i_b].pos + pos_v3 = self.fem_solver.elements_v[i_step, i_v3, i_b].pos + + e10 = pos_v0 - pos_v1 + e12 = pos_v2 - pos_v1 + e13 = pos_v3 - pos_v1 + + area_vector = e12.cross(e13) + signed_volume = area_vector.dot(e10) + if ti.abs(signed_volume) > gs.EPS: + grad_i = area_vector / signed_volume + grad[i_b, i_e] += grad_i * self.fem_pressure[i_v0] + + # ------------------------------------------------------------------------------------ + # ------------------------------------- Solve ---------------------------------------- + # ------------------------------------------------------------------------------------ + + def sap_solve(self, i_step): + self._init_sap_solve(i_step) + for iter in range(self._n_sap_iterations): + # init gradient and preconditioner + self.compute_non_contact_gradient_diag(i_step, iter) + + # compute contact hessian and gradient + self.compute_contact_gradient_hessian_diag_prec() + self.check_sap_convergence() + # solve for the vertex velocity + self.pcg_solve() + + # line search + self.exact_linesearch(i_step) + + @ti.kernel + def check_sap_convergence(self): + a_tol = 1e-6 + r_tol = 1e-5 + for i_b in range(self.fem_solver._B): + if not self.batch_active[i_b]: + continue + self.sap_state[i_b].gradient_norm = 0.0 + self.sap_state[i_b].momentum_norm = 0.0 + self.sap_state[i_b].impulse_norm = 0.0 + + for i_b, i_v in ti.ndrange(self.fem_solver._B, self.fem_solver.n_vertices): + if not self.batch_active[i_b]: + continue + self.sap_state[i_b].gradient_norm += ( + self.gradient[i_b, i_v].norm_sqr() * self.fem_solver.elements_v_info[i_v].mass_inv + ) + self.sap_state[i_b].momentum_norm += self.v[i_b, i_v].norm_sqr() * self.fem_solver.elements_v_info[i_v].mass + self.sap_state[i_b].impulse_norm += ( + self.sap_state_v.impulse[i_b, i_v].norm_sqr() * self.fem_solver.elements_v_info[i_v].mass_inv + ) + for i_b in range(self.fem_solver._B): + if not self.batch_active[i_b]: + continue + self.batch_active[i_b] = self.sap_state[i_b].gradient_norm >= a_tol + r_tol * ti.max( + self.sap_state[i_b].momentum_norm, self.sap_state[i_b].impulse_norm + ) + + def _init_sap_solve(self, i_step: ti.i32): + self._init_v(i_step) + self.batch_active.fill(True) + for contact in self.contacts: + if contact.has_contact: + contact.compute_regularization() + + @ti.kernel + def _init_v(self, i_step: ti.i32): + for i_b, i_v in ti.ndrange(self.fem_solver._B, self.fem_solver.n_vertices): + self.v[i_b, i_v] = self.fem_solver.elements_v[i_step + 1, i_v, i_b].vel + + def compute_non_contact_gradient_diag(self, i_step: ti.i32, iter: int): + self.init_non_contact_gradient_diag(i_step) + # No need to do this for iter=0 because v=v* and A(v-v*) = 0 + if iter > 0: + self.compute_inertia_elastic_gradient() + + @ti.kernel + def init_non_contact_gradient_diag(self, i_step: ti.i32): + dt2 = self.fem_solver._substep_dt**2 + for i_b, i_v in ti.ndrange(self.fem_solver._B, self.fem_solver.n_vertices): + self.gradient[i_b, i_v].fill(0.0) + # was using position now using velocity, need to multiply dt^2 + self.pcg_state_v[i_b, i_v].diag3x3 = self.fem_solver.pcg_state_v[i_b, i_v].diag3x3 * dt2 + self.v_diff[i_b, i_v] = self.v[i_b, i_v] - self.fem_solver.elements_v[i_step + 1, i_v, i_b].vel + + @ti.kernel + def compute_inertia_elastic_gradient(self): + self._func_compute_inertia_elastic_Ap(self.v_diff, self.gradient, self.batch_active) + + def compute_contact_gradient_hessian_diag_prec(self): + self.clear_impulses() + for contact in self.contacts: + if contact.has_contact: + contact.compute_gradient_hessian_diag() + self.compute_preconditioner() + + @ti.kernel + def clear_impulses(self): + for i_b, i_v in ti.ndrange(self.fem_solver._B, self.fem_solver.n_vertices): + if not self.batch_active[i_b]: + continue + self.sap_state_v[i_b, i_v].impulse.fill(0.0) + + @ti.kernel + def compute_preconditioner(self): + for i_b, i_v in ti.ndrange(self.fem_solver._B, self.fem_solver.n_vertices): + if not self.batch_active[i_b]: + continue + self.pcg_state_v[i_b, i_v].prec = self.pcg_state_v[i_b, i_v].diag3x3.inverse() + + def compute_Ap(self): + self.compute_inertia_elastic_Ap() + # Contact + for contact in self.contacts: + if contact.has_contact: + contact.compute_Ap() + + @ti.kernel + def compute_inertia_elastic_Ap(self): + self._func_compute_inertia_elastic_Ap(self.pcg_state_v.p, self.pcg_state_v.Ap, self.batch_pcg_active) + + @ti.func + def compute_elastic_products(self, i_b, i_e, B, s, i_v0, i_v1, i_v2, i_v3, src): + p9 = ti.Vector.zero(gs.ti_float, 9) + for i in ti.static(range(3)): + p9[i * 3 : i * 3 + 3] = ( + B[0, i] * src[i_b, i_v0] + B[1, i] * src[i_b, i_v1] + B[2, i] * src[i_b, i_v2] + s[i] * src[i_b, i_v3] + ) + H9_p9 = ti.Vector.zero(gs.ti_float, 9) + for i in ti.static(range(3)): + H9_p9[i * 3 : i * 3 + 3] = ( + self.fem_solver.elements_el_hessian[i_b, i, 0, i_e] @ p9[0:3] + + self.fem_solver.elements_el_hessian[i_b, i, 1, i_e] @ p9[3:6] + + self.fem_solver.elements_el_hessian[i_b, i, 2, i_e] @ p9[6:9] + ) + return p9, H9_p9 + + @ti.func + def _func_compute_inertia_elastic_Ap(self, src, dst, active): + dt2 = self.fem_solver._substep_dt**2 + damping_alpha_factor = self.fem_solver._damping_alpha * self.fem_solver._substep_dt + 1.0 + damping_beta_factor = self.fem_solver._damping_beta / self.fem_solver._substep_dt + 1.0 + + # Inerita + for i_b, i_v in ti.ndrange(self.fem_solver._B, self.fem_solver.n_vertices): + if not active[i_b]: + continue + dst[i_b, i_v] = ( + self.fem_solver.elements_v_info[i_v].mass_over_dt2 * src[i_b, i_v] * dt2 * damping_alpha_factor + ) + + # Elasticity + for i_b, i_e in ti.ndrange(self.fem_solver._B, self.fem_solver.n_elements): + if not active[i_b]: + continue + V_dt2 = self.fem_solver.elements_i[i_e].V * dt2 + B = self.fem_solver.elements_i[i_e].B + s = -B[0, :] - B[1, :] - B[2, :] # s is the negative sum of B rows + i_v0, i_v1, i_v2, i_v3 = self.fem_solver.elements_i[i_e].el2v + + _, new_p9 = self.compute_elastic_products(i_b, i_e, B, s, i_v0, i_v1, i_v2, i_v3, src) + # atomic + scale = V_dt2 * damping_beta_factor + dst[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]) * scale + dst[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]) * scale + dst[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]) * scale + dst[i_b, i_v3] += (s[0] * new_p9[0:3] + s[1] * new_p9[3:6] + s[2] * new_p9[6:9]) * scale + + @ti.kernel + def init_pcg_solve(self): + for i_b in ti.ndrange(self._B): + self.batch_pcg_active[i_b] = self.batch_active[i_b] + if not self.batch_pcg_active[i_b]: + continue + self.pcg_state[i_b].rTr = 0.0 + self.pcg_state[i_b].rTz = 0.0 + for i_b, i_v in ti.ndrange(self._B, self.fem_solver.n_vertices): + if not self.batch_pcg_active[i_b]: + continue + self.pcg_state_v[i_b, i_v].x = 0.0 + self.pcg_state_v[i_b, i_v].r = -self.gradient[i_b, i_v] + self.pcg_state_v[i_b, i_v].z = self.pcg_state_v[i_b, i_v].prec @ self.pcg_state_v[i_b, i_v].r + self.pcg_state_v[i_b, i_v].p = self.pcg_state_v[i_b, i_v].z + self.pcg_state[i_b].rTr += self.pcg_state_v[i_b, i_v].r.dot(self.pcg_state_v[i_b, i_v].r) + self.pcg_state[i_b].rTz += self.pcg_state_v[i_b, i_v].r.dot(self.pcg_state_v[i_b, i_v].z) + for i_b in ti.ndrange(self._B): + if not self.batch_pcg_active[i_b]: + continue + self.batch_pcg_active[i_b] = self.pcg_state[i_b].rTr > self._pcg_threshold + + def one_pcg_iter(self): + self.compute_Ap() + self._kernel_one_pcg_iter() + + @ti.kernel + def _kernel_one_pcg_iter(self): + # compute pTAp + for i_b in ti.ndrange(self._B): + if not self.batch_pcg_active[i_b]: + continue + self.pcg_state[i_b].pTAp = 0.0 + for i_b, i_v in ti.ndrange(self._B, self.fem_solver.n_vertices): + if not self.batch_pcg_active[i_b]: + continue + ti.atomic_add(self.pcg_state[i_b].pTAp, self.pcg_state_v[i_b, i_v].p.dot(self.pcg_state_v[i_b, i_v].Ap)) + + # compute alpha and update x, r, z, rTr, rTz + for i_b in ti.ndrange(self._B): + if not self.batch_pcg_active[i_b]: + continue + self.pcg_state[i_b].alpha = self.pcg_state[i_b].rTz / self.pcg_state[i_b].pTAp + self.pcg_state[i_b].rTr_new = 0.0 + self.pcg_state[i_b].rTz_new = 0.0 + for i_b, i_v in ti.ndrange(self._B, self.fem_solver.n_vertices): + if not self.batch_pcg_active[i_b]: + continue + self.pcg_state_v[i_b, i_v].x = ( + self.pcg_state_v[i_b, i_v].x + self.pcg_state[i_b].alpha * self.pcg_state_v[i_b, i_v].p + ) + self.pcg_state_v[i_b, i_v].r = ( + self.pcg_state_v[i_b, i_v].r - self.pcg_state[i_b].alpha * self.pcg_state_v[i_b, i_v].Ap + ) + self.pcg_state_v[i_b, i_v].z = self.pcg_state_v[i_b, i_v].prec @ self.pcg_state_v[i_b, i_v].r + self.pcg_state[i_b].rTr_new += self.pcg_state_v[i_b, i_v].r.dot(self.pcg_state_v[i_b, i_v].r) + self.pcg_state[i_b].rTz_new += self.pcg_state_v[i_b, i_v].r.dot(self.pcg_state_v[i_b, i_v].z) + + # check convergence + for i_b in ti.ndrange(self._B): + if not self.batch_pcg_active[i_b]: + continue + self.batch_pcg_active[i_b] = self.pcg_state[i_b].rTr_new > self._pcg_threshold + # update beta, rTr, rTz + for i_b in ti.ndrange(self._B): + if not self.batch_pcg_active[i_b]: + continue + self.pcg_state[i_b].beta = self.pcg_state[i_b].rTz_new / self.pcg_state[i_b].rTz + self.pcg_state[i_b].rTr = self.pcg_state[i_b].rTr_new + self.pcg_state[i_b].rTz = self.pcg_state[i_b].rTz_new + + # update p + for i_b, i_v in ti.ndrange(self._B, self.fem_solver.n_vertices): + if not self.batch_pcg_active[i_b]: + continue + self.pcg_state_v[i_b, i_v].p = ( + self.pcg_state_v[i_b, i_v].z + self.pcg_state[i_b].beta * self.pcg_state_v[i_b, i_v].p + ) + + def pcg_solve(self): + self.init_pcg_solve() + for i in range(self._n_pcg_iterations): + self.one_pcg_iter() + + def compute_total_energy(self, i_step: ti.i32, energy): + self.compute_inertia_elastic_energy(i_step, energy) + # Contact + for contact in self.contacts: + if contact.has_contact: + contact.compute_energy(energy) + + @ti.kernel + def compute_inertia_elastic_energy(self, i_step: ti.i32, energy: ti.template()): + dt2 = self.fem_solver._substep_dt**2 + damping_alpha_factor = self.fem_solver._damping_alpha * self.fem_solver._substep_dt + 1.0 + damping_beta_factor = self.fem_solver._damping_beta / self.fem_solver._substep_dt + 1.0 + + for i_b in ti.ndrange(self._B): + energy[i_b] = 0.0 + if not self.batch_linesearch_active[i_b]: + continue + + # Inertia + for i_b, i_v in ti.ndrange(self._B, self.fem_solver.n_vertices): + if not self.batch_linesearch_active[i_b]: + continue + self.v_diff[i_b, i_v] = self.v[i_b, i_v] - self.fem_solver.elements_v[i_step + 1, i_v, i_b].vel + energy[i_b] += ( + 0.5 + * self.fem_solver.elements_v_info[i_v].mass_over_dt2 + * self.v_diff[i_b, i_v].dot(self.v_diff[i_b, i_v]) + * dt2 + * damping_alpha_factor + ) + + # Elastic + for i_b, i_e in ti.ndrange(self._B, self.fem_solver.n_elements): + if not self.batch_linesearch_active[i_b]: + continue + + V_dt2 = self.fem_solver.elements_i[i_e].V * dt2 + B = self.fem_solver.elements_i[i_e].B + s = -B[0, :] - B[1, :] - B[2, :] # s is the negative sum of B rows + i_v0, i_v1, i_v2, i_v3 = self.fem_solver.elements_i[i_e].el2v + + p9, H9_p9 = self.compute_elastic_products(i_b, i_e, B, s, i_v0, i_v1, i_v2, i_v3, self.v_diff) + energy[i_b] += 0.5 * p9.dot(H9_p9) * damping_beta_factor * V_dt2 + + def init_exact_linesearch(self, i_step: ti.i32): + self._kernel_init_exact_linesearch() + self.prepare_search_direction_data() + self.compute_inertia_elastic_energy(i_step, self.linesearch_state.prev_energy) + self.update_velocity_linesearch() + self.compute_line_energy_gradient_hessian(i_step) + self.check_initial_exact_linesearch_convergence() + self.init_newton_linesearch() + + @ti.kernel + def init_newton_linesearch(self): + for i_b in ti.ndrange(self._B): + if not self.batch_linesearch_active[i_b]: + continue + self.linesearch_state[i_b].dell_scale = -self.linesearch_state[i_b].m + self.linesearch_state[i_b].step_size = ti.min( + -self.linesearch_state[i_b].m / self.linesearch_state[i_b].d2ell_dalpha2, self._linesearch_max_step_size + ) + self.linesearch_state[i_b].alpha_min = 0.0 + self.linesearch_state[i_b].alpha_max = self._linesearch_max_step_size + self.linesearch_state[i_b].f_lower = -1.0 + self.linesearch_state[i_b].f_upper = ( + self.linesearch_state[i_b].dell_dalpha / self.linesearch_state[i_b].dell_scale + ) + self.linesearch_state[i_b].alpha_tol = self._linesearch_ftol * self.linesearch_state[i_b].step_size + self.linesearch_state[i_b].minus_dalpha = ( + self.linesearch_state[i_b].alpha_min - self.linesearch_state[i_b].alpha_max + ) + self.linesearch_state[i_b].minus_dalpha_prev = self.linesearch_state[i_b].minus_dalpha + if ti.abs(self.linesearch_state[i_b].f_lower) < self._linesearch_ftol: + self.batch_linesearch_active[i_b] = False + self.linesearch_state[i_b].step_size = self.linesearch_state[i_b].alpha_min + if ti.abs(self.linesearch_state[i_b].f_upper) < self._linesearch_ftol: + self.batch_linesearch_active[i_b] = False + self.linesearch_state[i_b].step_size = self.linesearch_state[i_b].alpha_max + + def compute_line_energy_gradient_hessian(self, i_step: ti.i32): + for contact in self.contacts: + if contact.has_contact: + contact.compute_energy_gamma_G() + self.compute_inertia_elastic_energy_alpha(i_step, self.linesearch_state.energy) + self.compute_inertia_elastic_gradient_alpha(i_step) + self.compute_inertia_elastic_hessian_alpha() + for contact in self.contacts: + if contact.has_contact: + contact.compute_gradient_hessian_alpha() + + @ti.kernel + def compute_inertia_elastic_gradient_alpha(self, i_step: ti.i32): + self.linesearch_state.dell_dalpha.fill(0.0) + dp = ti.static(self.linesearch_state_v.dp) + v = ti.static(self.v) + v_star = ti.static(self.fem_solver.elements_v.vel) + for i_b, i_v in ti.ndrange(self._B, self.fem_solver.n_vertices): + if not self.batch_linesearch_active[i_b]: + continue + self.linesearch_state.dell_dalpha[i_b] += dp[i_b, i_v].dot(v[i_b, i_v] - v_star[i_step + 1, i_b, i_v]) + + @ti.kernel + def compute_inertia_elastic_hessian_alpha(self): + for i_b in ti.ndrange(self._B): + self.linesearch_state.d2ell_dalpha2[i_b] = self.linesearch_state.d2ellA_dalpha2[i_b] + + @ti.kernel + def compute_inertia_elastic_energy_alpha(self, i_step: ti.i32, energy: ti.template()): + alpha = ti.static(self.linesearch_state.step_size) + dp = ti.static(self.linesearch_state_v.dp) + v = ti.static(self.v) + v_star = ti.static(self.fem_solver.elements_v.vel) + for i_b in ti.ndrange(self._B): + if not self.batch_linesearch_active[i_b]: + continue + energy[i_b] = ( + self.linesearch_state.prev_energy[i_b] + + 0.5 * alpha[i_b] ** 2 * self.linesearch_state[i_b].d2ellA_dalpha2 + ) + + for i_b, i_v in ti.ndrange(self._B, self.fem_solver.n_vertices): + if not self.batch_linesearch_active[i_b]: + continue + energy[i_b] += alpha[i_b] * dp[i_b, i_v].dot(v[i_b, i_v] - v_star[i_step + 1, i_b, i_v]) + + def prepare_search_direction_data(self): + self.prepare_inertia_elastic_search_direction_data() + for contact in self.contacts: + if contact.has_contact: + contact.prepare_search_direction_data() + self.compute_d2ellA_dalpha2() + + @ti.kernel + def compute_d2ellA_dalpha2(self): + for i_b in ti.ndrange(self._B): + self.linesearch_state[i_b].d2ellA_dalpha2 = 0.0 + + for i_b, i_v in ti.ndrange(self._B, self.fem_solver.n_vertices): + if not self.batch_linesearch_active[i_b]: + continue + self.linesearch_state[i_b].d2ellA_dalpha2 += self.pcg_state_v[i_b, i_v].x.dot( + self.linesearch_state_v[i_b, i_v].dp + ) + + @ti.kernel + def prepare_inertia_elastic_search_direction_data(self): + self._func_compute_inertia_elastic_Ap( + self.pcg_state_v.x, self.linesearch_state_v.dp, self.batch_linesearch_active + ) + + @ti.kernel + def _kernel_init_exact_linesearch(self): + for i_b in ti.ndrange(self._B): + self.batch_linesearch_active[i_b] = self.batch_active[i_b] + if not self.batch_linesearch_active[i_b]: + continue + self.linesearch_state[i_b].m = 0.0 + self.linesearch_state[i_b].step_size = self._linesearch_max_step_size + + # x_prev, m + for i_b, i_v in ti.ndrange(self._B, self.fem_solver.n_vertices): + if not self.batch_linesearch_active[i_b]: + continue + self.linesearch_state[i_b].m += self.pcg_state_v[i_b, i_v].x.dot(self.gradient[i_b, i_v]) + self.linesearch_state_v[i_b, i_v].x_prev = self.v[i_b, i_v] + + @ti.kernel + def check_initial_exact_linesearch_convergence(self): + for i_b in ti.ndrange(self._B): + if not self.batch_linesearch_active[i_b]: + continue + self.batch_linesearch_active[i_b] = self.linesearch_state[i_b].dell_dalpha > 0.0 + # When tolerance is small but gradient norm is small, take step 1.0 and end + for i_b, i_v in ti.ndrange(self._B, self.fem_solver.n_vertices): + if not self.batch_linesearch_active[i_b]: + continue + if ( + -self.linesearch_state[i_b].m + < self._sap_convergence_atol + self._sap_convergence_rtol * self.linesearch_state[i_b].prev_energy + ): + self.v[i_b, i_v] = self.linesearch_state_v[i_b, i_v].x_prev + self.pcg_state_v[i_b, i_v].x + for i_b, i_v in ti.ndrange(self._B, self.fem_solver.n_vertices): + if not self.batch_linesearch_active[i_b]: + continue + if ( + -self.linesearch_state[i_b].m + < self._sap_convergence_atol + self._sap_convergence_rtol * self.linesearch_state[i_b].prev_energy + ): + self.batch_linesearch_active[i_b] = False + self.linesearch_state[i_b].step_size = 1.0 + + @ti.kernel + def update_velocity_linesearch(self): + for i_b, i_v in ti.ndrange(self._B, self.fem_solver.n_vertices): + if not self.batch_linesearch_active[i_b]: + continue + self.v[i_b, i_v] = ( + self.linesearch_state_v[i_b, i_v].x_prev + + self.linesearch_state[i_b].step_size * self.pcg_state_v[i_b, i_v].x + ) + + def exact_linesearch(self, i_step: ti.i32): + """ + Exact line search using rtsafe (Numerical Recipes book). + + This is a hybrid of Newton's method and bisection to find root of df/dalpha = 0. + + Note + ------ + Code Reference: + https://github.com/RobotLocomotion/drake/blob/5fbb89e6e380c418b3f651ebde22a8f9203b6b1e/multibody/contact_solvers/sap/sap_solver.h#L393 + """ + self.init_exact_linesearch(i_step) + for i in range(self._n_linesearch_iterations): + self.one_exact_linesearch_iter(i_step) + + def one_exact_linesearch_iter(self, i_step: ti.i32): + self.update_velocity_linesearch() + self.compute_line_energy_gradient_hessian(i_step) + self.compute_f_df_bracket() + self.find_next_step_size() + + @ti.kernel + def compute_f_df_bracket(self): + """ + Compute the function (derivative of total energy) value and its derivative to alpha. + Update the bracket for the next step size. + + The bracket is defined by [alpha_min, alpha_max] which is the range that contains the root of df/dalpha = 0. + """ + for i_b in ti.ndrange(self._B): + if not self.batch_linesearch_active[i_b]: + continue + self.linesearch_state[i_b].f = ( + self.linesearch_state[i_b].dell_dalpha / self.linesearch_state[i_b].dell_scale + ) + self.linesearch_state[i_b].df = ( + self.linesearch_state[i_b].d2ell_dalpha2 / self.linesearch_state[i_b].dell_scale + ) + if ti.math.sign(self.linesearch_state[i_b].f) != ti.math.sign(self.linesearch_state[i_b].f_upper): + self.linesearch_state[i_b].alpha_min = self.linesearch_state[i_b].step_size + self.linesearch_state[i_b].f_lower = self.linesearch_state[i_b].f + else: + self.linesearch_state[i_b].alpha_max = self.linesearch_state[i_b].step_size + self.linesearch_state[i_b].f_upper = self.linesearch_state[i_b].f + if ti.abs(self.linesearch_state[i_b].f) < self._linesearch_ftol: + self.batch_linesearch_active[i_b] = False + + @ti.kernel + def find_next_step_size(self): + for i_b in ti.ndrange(self._B): + if not self.batch_linesearch_active[i_b]: + continue + newton_is_slow = 2.0 * ti.abs(self.linesearch_state[i_b].f) > ti.abs( + self.linesearch_state[i_b].minus_dalpha_prev * self.linesearch_state[i_b].df + ) + self.linesearch_state[i_b].minus_dalpha_prev = self.linesearch_state[i_b].minus_dalpha + if newton_is_slow: + # bisect + self.linesearch_state[i_b].minus_dalpha = 0.5 * ( + self.linesearch_state[i_b].alpha_min - self.linesearch_state[i_b].alpha_max + ) + self.linesearch_state[i_b].step_size = ( + self.linesearch_state[i_b].alpha_min - self.linesearch_state[i_b].minus_dalpha + ) + else: + # newton + self.linesearch_state[i_b].minus_dalpha = self.linesearch_state[i_b].f / self.linesearch_state[i_b].df + self.linesearch_state[i_b].step_size = ( + self.linesearch_state[i_b].step_size - self.linesearch_state[i_b].minus_dalpha + ) + if ( + self.linesearch_state[i_b].step_size <= self.linesearch_state[i_b].alpha_min + or self.linesearch_state[i_b].step_size >= self.linesearch_state[i_b].alpha_max + ): + # bisect + self.linesearch_state[i_b].minus_dalpha = 0.5 * ( + self.linesearch_state[i_b].alpha_min - self.linesearch_state[i_b].alpha_max + ) + self.linesearch_state[i_b].step_size = ( + self.linesearch_state[i_b].alpha_min - self.linesearch_state[i_b].minus_dalpha + ) + if ti.abs(self.linesearch_state[i_b].minus_dalpha) < self.linesearch_state[i_b].alpha_tol: + self.batch_linesearch_active[i_b] = False + + # ------------------------------------------------------------------------------------ + # ----------------------------------- Properties ------------------------------------- + # ------------------------------------------------------------------------------------ + @property + def active_solvers(self): + """All the active solvers managed by the scene's simulator.""" + return self.sim.active_solvers + + +class ContactMode(IntEnum): + STICK = 0 + SLIDE = 1 + NO_CONTACT = 2 + + +@ti.data_oriented +class BaseContact(RBC): + """ + Base class for contact handling in SAPCoupler. + + This class provides a framework for managing contact pairs, computing gradients, + and handling contact-related computations. + """ + + def __init__( + self, + simulator: "Simulator", + ) -> None: + self.sim = simulator + self.coupler = simulator.coupler + self.n_contact_pairs = ti.field(gs.ti_int, shape=()) + self._has_contact = True + self.sap_contact_info_type = ti.types.struct( + k=gs.ti_float, # contact stiffness + phi0=gs.ti_float, # initial signed distance + Rn=gs.ti_float, # Regularization for normal + Rt=gs.ti_float, # Regularization for tangential + Rn_inv=gs.ti_float, # Inverse of Rn + Rt_inv=gs.ti_float, # Inverse of Rt + vn_hat=gs.ti_float, # Stablization for normal velocity + mu=gs.ti_float, # friction coefficient + mu_hat=gs.ti_float, # friction coefficient regularized + mu_factor=gs.ti_float, # friction coefficient factor, 1/(1+mu_tilde**2) + energy=gs.ti_float, # energy + gamma=gs.ti_vec3, # contact impulse + G=gs.ti_mat3, # Hessian matrix + dvc=gs.ti_vec3, # velocity change at contact point, for exact line search + ) + + @property + def has_contact(self): + return self._has_contact + + def update_has_contact(self): + self._has_contact = self.n_contact_pairs[None] > 0 + + @ti.kernel + def compute_gradient_hessian_diag(self): + pairs = ti.static(self.contact_pairs) + sap_info = ti.static(pairs.sap_info) + for i_p in range(self.n_contact_pairs[None]): + vc = self.compute_Jx(i_p, self.coupler.v) + # With floor, the contact frame is the same as the world frame + self.compute_contact_gamma_G(sap_info, i_p, vc) + self.add_Jt_x(self.coupler.gradient, i_p, -sap_info[i_p].gamma) + self.add_Jt_x(self.coupler.sap_state_v.impulse, i_p, sap_info[i_p].gamma) + self.add_Jt_A_J_diag3x3(self.coupler.pcg_state_v.diag3x3, i_p, sap_info[i_p].G) + + @ti.kernel + def compute_gradient_hessian_alpha(self): + dvc = ti.static(self.contact_pairs.sap_info.dvc) + gamma = ti.static(self.contact_pairs.sap_info.gamma) + G = ti.static(self.contact_pairs.sap_info.G) + for i_p in ti.ndrange(self.n_contact_pairs[None]): + i_b = self.contact_pairs[i_p].batch_idx + if not self.coupler.batch_linesearch_active[i_b]: + continue + self.coupler.linesearch_state.dell_dalpha[i_b] -= dvc[i_p].dot(gamma[i_p]) + self.coupler.linesearch_state.d2ell_dalpha2[i_b] += dvc[i_p].dot(G[i_p] @ dvc[i_p]) + + @ti.kernel + def compute_regularization(self): + pairs = ti.static(self.contact_pairs) + sap_info = ti.static(pairs.sap_info) + dt2_inv = 1.0 / (self.sim._substep_dt**2) + + for i_p in range(self.n_contact_pairs[None]): + W = self.compute_delassus(i_p) + w_rms = W.norm() / 3.0 * dt2_inv + self.compute_contact_regularization(sap_info, i_p, w_rms, self.sim._substep_dt) + + @ti.kernel + def compute_energy_gamma_G(self): + pairs = ti.static(self.contact_pairs) + sap_info = ti.static(pairs.sap_info) + for i_p in range(self.n_contact_pairs[None]): + vc = self.compute_Jx(i_p, self.coupler.v) + self.compute_contact_energy_gamma_G(sap_info, i_p, vc) + + @ti.kernel + def compute_energy(self, energy: ti.template()): + pairs = ti.static(self.contact_pairs) + sap_info = ti.static(pairs.sap_info) + for i_p in range(self.n_contact_pairs[None]): + i_b = pairs[i_p].batch_idx + if not self.batch_linesearch_active[i_b]: + continue + vc = self.compute_Jx(i_p, self.coupler.v) + self.compute_contact_energy(sap_info, i_p, vc) + energy[i_b] += sap_info[i_p].energy + + @ti.kernel + def prepare_search_direction_data(self): + pairs = ti.static(self.contact_pairs) + sap_info = ti.static(pairs.sap_info) + for i_p in ti.ndrange(self.n_contact_pairs[None]): + i_b = pairs[i_p].batch_idx + if not self.coupler.batch_linesearch_active[i_b]: + continue + sap_info[i_p].dvc = self.compute_Jx(i_p, self.coupler.pcg_state_v.x) + + @ti.kernel + def compute_Ap(self): + pairs = ti.static(self.contact_pairs) + sap_info = ti.static(pairs.sap_info) + for i_p in range(self.n_contact_pairs[None]): + # Jt @ G @ J @ p + x = self.compute_Jx(i_p, self.coupler.pcg_state_v.p) + x = sap_info[i_p].G @ x + self.add_Jt_x(self.coupler.pcg_state_v.Ap, i_p, x) + + @ti.kernel + def compute_contact_pos(self, i_step: ti.i32): + pairs = ti.static(self.contact_pairs) + for i_p in range(self.n_contact_pairs[None]): + pairs[i_p].contact_pos = self.compute_contact_point(i_p, self.fem_solver.elements_v.pos, i_step) + + @ti.func + def compute_contact_gamma_G(self, sap_info, i_p, vc): + y = ti.Vector([0.0, 0.0, sap_info[i_p].vn_hat]) - vc + y[0] *= sap_info[i_p].Rt_inv + y[1] *= sap_info[i_p].Rt_inv + y[2] *= sap_info[i_p].Rn_inv + yr = y[:2].norm(gs.EPS) + yn = y[2] + + t_hat = y[:2] / yr + contact_mode = self.compute_contact_mode(sap_info[i_p].mu, sap_info[i_p].mu_hat, yr, yn) + sap_info[i_p].gamma.fill(0.0) + sap_info[i_p].G.fill(0.0) + if contact_mode == ContactMode.STICK: + sap_info[i_p].gamma = y + sap_info[i_p].G[0, 0] = sap_info[i_p].Rt_inv + sap_info[i_p].G[1, 1] = sap_info[i_p].Rt_inv + sap_info[i_p].G[2, 2] = sap_info[i_p].Rn_inv + elif contact_mode == ContactMode.SLIDE: + gn = (yn + sap_info[i_p].mu_hat * yr) * sap_info[i_p].mu_factor + gt = sap_info[i_p].mu * gn * t_hat + sap_info[i_p].gamma = ti.Vector([gt[0], gt[1], gn]) + P = t_hat.outer_product(t_hat) + Pperp = ti.Matrix.identity(gs.ti_float, 2) - P + dgt_dyt = sap_info[i_p].mu * (gn / yr * Pperp + sap_info[i_p].mu_hat * sap_info[i_p].mu_factor * P) + dgt_dyn = sap_info[i_p].mu * sap_info[i_p].mu_factor * t_hat + dgn_dyt = sap_info[i_p].mu_hat * sap_info[i_p].mu_factor * t_hat + dgn_dyn = sap_info[i_p].mu_factor + + sap_info[i_p].G[:2, :2] = dgt_dyt * sap_info[i_p].Rt_inv + sap_info[i_p].G[:2, 2] = dgt_dyn * sap_info[i_p].Rn_inv + sap_info[i_p].G[2, :2] = dgn_dyt * sap_info[i_p].Rt_inv + sap_info[i_p].G[2, 2] = dgn_dyn * sap_info[i_p].Rn_inv + else: # No contact + pass + + @ti.func + def compute_contact_energy_gamma_G(self, sap_info, i_p, vc): + self.compute_contact_gamma_G(sap_info, i_p, vc) + R_gamma = sap_info[i_p].gamma + R_gamma[0] *= sap_info[i_p].Rt + R_gamma[1] *= sap_info[i_p].Rt + R_gamma[2] *= sap_info[i_p].Rn + sap_info[i_p].energy = 0.5 * sap_info[i_p].gamma.dot(R_gamma) + + @ti.func + def compute_contact_energy(self, sap_info, i_p, vc): + y = ti.Vector([0.0, 0.0, sap_info[i_p].vn_hat]) - vc + y[0] *= sap_info[i_p].Rt_inv + y[1] *= sap_info[i_p].Rt_inv + y[2] *= sap_info[i_p].Rn_inv + yr = y[:2].norm(gs.EPS) + yn = y[2] + + t_hat = y[:2] / yr + contact_mode = self.compute_contact_mode(sap_info[i_p].mu, sap_info[i_p].mu_hat, yr, yn) + sap_info[i_p].gamma.fill(0.0) + if contact_mode == ContactMode.STICK: + sap_info[i_p].gamma = y + elif contact_mode == ContactMode.SLIDE: + gn = (yn + sap_info[i_p].mu_hat * yr) * sap_info[i_p].mu_factor + gt = sap_info[i_p].mu * gn * t_hat + sap_info[i_p].gamma = ti.Vector([gt[0], gt[1], gn]) + else: # No contact + pass + + R_gamma = sap_info[i_p].gamma + R_gamma[0] *= sap_info[i_p].Rt + R_gamma[1] *= sap_info[i_p].Rt + R_gamma[2] *= sap_info[i_p].Rn + sap_info[i_p].energy = 0.5 * sap_info[i_p].gamma.dot(R_gamma) + + @ti.func + def compute_contact_mode(self, mu, mu_hat, yr, yn): + """ + Compute the contact mode based on the friction coefficients and the relative velocities. + """ + result = ContactMode.NO_CONTACT + if yr <= mu * yn: + result = ContactMode.STICK + elif -mu_hat * yr < yn and yn < yr / mu: + result = ContactMode.SLIDE + return result + + @ti.func + def compute_contact_regularization(self, sap_info, i_p, w_rms, time_step): + beta_factor = self.coupler._sap_beta**2 / (4.0 * ti.math.pi**2) + k = sap_info[i_p].k + Rn = max(beta_factor * w_rms, 1.0 / (time_step * k * (time_step + self.coupler._sap_taud))) + Rt = self.coupler._sap_sigma * w_rms + vn_hat = -sap_info[i_p].phi0 / (time_step + self.coupler._sap_taud) + sap_info[i_p].Rn = Rn + sap_info[i_p].Rt = Rt + sap_info[i_p].Rn_inv = 1.0 / Rn + sap_info[i_p].Rt_inv = 1.0 / Rt + sap_info[i_p].vn_hat = vn_hat + sap_info[i_p].mu_hat = sap_info[i_p].mu * Rt * sap_info[i_p].Rn_inv + sap_info[i_p].mu_factor = 1.0 / (1.0 + sap_info[i_p].mu * sap_info[i_p].mu_hat) + + +@ti.data_oriented +class FEMFloorTetContact(BaseContact): + """ + Class for handling contact between a tetrahedral mesh and a floor in a simulation using hydroelastic model. + + This class extends the BaseContact class and provides methods for detecting contact + between the tetrahedral elements and the floor, computing contact pairs, and managing + contact-related computations. + """ + + def __init__( + self, + simulator: "Simulator", + ) -> None: + super().__init__(simulator) + self.name = "FEMFloorTetContact" + self.fem_solver = self.sim.fem_solver + self.contact_candidate_type = ti.types.struct( + batch_idx=gs.ti_int, # batch index + geom_idx=gs.ti_int, # index of the FEM element + intersection_code=gs.ti_int, # intersection code for the element + distance=gs.ti_vec4, # distance vector for the element + ) + self.n_contact_candidates = ti.field(gs.ti_int, shape=()) + self.max_contact_candidates = self.fem_solver.n_surface_elements * self.fem_solver._B + self.contact_candidates = self.contact_candidate_type.field(shape=(self.max_contact_candidates,)) + + self.contact_pair_type = ti.types.struct( + batch_idx=gs.ti_int, # batch index + geom_idx=gs.ti_int, # index of the FEM element + barycentric=gs.ti_vec4, # barycentric coordinates of the contact point + contact_pos=gs.ti_vec3, # contact position + sap_info=self.sap_contact_info_type, # contact info + ) + self.max_contact_pairs = self.fem_solver.n_surface_elements * self.fem_solver._B + self.contact_pairs = self.contact_pair_type.field(shape=(self.max_contact_pairs,)) + + @ti.kernel + def detection(self, i_step: ti.i32): + candidates = ti.static(self.contact_candidates) + # Compute contact pairs + self.n_contact_candidates[None] = 0 + # TODO Check surface element only instead of all elements + for i_b, i_e in ti.ndrange(self.coupler._B, self.fem_solver.n_elements): + intersection_code = ti.int32(0) + distance = ti.Vector.zero(gs.ti_float, 4) + for i in ti.static(range(4)): + i_v = self.fem_solver.elements_i[i_e].el2v[i] + pos_v = self.fem_solver.elements_v[i_step, i_v, i_b].pos + distance[i] = pos_v.z - self.fem_solver.floor_height + if distance[i] > 0.0: + intersection_code |= 1 << i + + # check if the element intersect with the floor + if intersection_code != 0 and intersection_code != 15: + i_c = ti.atomic_add(self.n_contact_candidates[None], 1) + if i_c < self.max_contact_candidates: + candidates[i_c].batch_idx = i_b + candidates[i_c].geom_idx = i_e + candidates[i_c].intersection_code = intersection_code + candidates[i_c].distance = distance + + pairs = ti.static(self.contact_pairs) + sap_info = ti.static(pairs.sap_info) + self.n_contact_pairs[None] = 0 + # Compute pair from candidates + for i_c in range(self.n_contact_candidates[None]): + candidate = candidates[i_c] + i_b = candidate.batch_idx + i_e = candidate.geom_idx + intersection_code = candidate.intersection_code + distance = candidate.distance + intersected_edges = self.coupler.MarchingTetsEdgeTable[intersection_code] + tet_vertices = ti.Matrix.zero(gs.ti_float, 3, 4) # 4 vertices + tet_pressures = ti.Vector.zero(gs.ti_float, 4) # pressures at the vertices + + for i in ti.static(range(4)): + i_v = self.fem_solver.elements_i[i_e].el2v[i] + tet_vertices[:, i] = self.fem_solver.elements_v[i_step, i_v, i_b].pos + tet_pressures[i] = self.coupler.fem_pressure[i_v] + + polygon_vertices = ti.Matrix.zero(gs.ti_float, 3, 4) # 3 or 4 vertices + total_area = gs.EPS # avoid division by zero + total_area_weighted_centroid = ti.Vector([0.0, 0.0, 0.0]) + for i in range(4): + if intersected_edges[i] >= 0: + edge = self.coupler.TetEdges[intersected_edges[i]] + pos_v0 = tet_vertices[:, edge[0]] + pos_v1 = tet_vertices[:, edge[1]] + d_v0 = distance[edge[0]] + d_v1 = distance[edge[1]] + t = d_v0 / (d_v0 - d_v1) + polygon_vertices[:, i] = pos_v0 + t * (pos_v1 - pos_v0) + + # Compute tirangle area and centroid + if i >= 2: + e1 = polygon_vertices[:, i - 1] - polygon_vertices[:, 0] + e2 = polygon_vertices[:, i] - polygon_vertices[:, 0] + 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] + ) + + deformable_g = self.coupler._hydroelastic_stiffness + rigid_g = self.coupler.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: + continue + g = 1.0 / (1.0 / deformable_g + 1.0 / rigid_g) # harmonic average + rigid_k = total_area * g + rigid_phi0 = -pressure / g + i_p = ti.atomic_add(self.n_contact_pairs[None], 1) + if i_p < self.max_contact_pairs: + pairs[i_p].batch_idx = i_b + pairs[i_p].geom_idx = i_e + pairs[i_p].barycentric = barycentric + # TODO custom dissipation + sap_info[i_p].k = rigid_k # contact stiffness + sap_info[i_p].phi0 = rigid_phi0 + sap_info[i_p].mu = self.fem_solver.elements_i[i_e].friction_mu # friction coefficient + + @ti.func + def compute_Jx(self, i_p, x): + """ + Compute the contact Jacobian J times a vector x. + """ + pairs = ti.static(self.contact_pairs) + i_b = pairs[i_p].batch_idx + i_g = pairs[i_p].geom_idx + Jx = ti.Vector.zero(gs.ti_float, 3) + for i in ti.static(range(4)): + i_v = self.fem_solver.elements_i[i_g].el2v[i] + Jx += pairs[i_p].barycentric[i] * x[i_b, i_v] + return Jx + + @ti.func + def compute_contact_point(self, i_p, x, f): + """ + Compute the contact point for a given contact pair. + """ + pairs = ti.static(self.contact_pairs) + i_b = pairs[i_p].batch_idx + i_g = pairs[i_p].geom_idx + Jx = ti.Vector.zero(gs.ti_float, 3) + for i in ti.static(range(4)): + i_v = self.fem_solver.elements_i[i_g].el2v[i] + Jx += pairs[i_p].barycentric[i] * x[f, i_v, i_b] + return Jx + + @ti.func + def add_Jt_x(self, y, i_p, x): + pairs = ti.static(self.contact_pairs) + i_b = pairs[i_p].batch_idx + i_g = pairs[i_p].geom_idx + for i in ti.static(range(4)): + i_v = self.fem_solver.elements_i[i_g].el2v[i] + y[i_b, i_v] += pairs[i_p].barycentric[i] * x + + @ti.func + def add_Jt_A_J_diag3x3(self, y, i_p, A): + pairs = ti.static(self.contact_pairs) + i_b = pairs[i_p].batch_idx + i_g = pairs[i_p].geom_idx + for i in ti.static(range(4)): + i_v = self.fem_solver.elements_i[i_g].el2v[i] + y[i_b, i_v] += pairs[i_p].barycentric[i] ** 2 * A + + @ti.func + def compute_delassus(self, i_p): + pairs = ti.static(self.contact_pairs) + i_b = pairs[i_p].batch_idx + i_g = pairs[i_p].geom_idx + W = ti.Matrix.zero(gs.ti_float, 3, 3) + # W = sum (JA^-1J^T) + # With floor, J is Identity times the barycentric coordinates + for i in ti.static(range(4)): + i_v = self.fem_solver.elements_i[i_g].el2v[i] + W += pairs[i_p].barycentric[i] ** 2 * self.fem_solver.pcg_state_v[i_b, i_v].prec + return W + + +@ti.data_oriented +class FEMSelfTetContact(BaseContact): + """ + Class for handling self-contact between tetrahedral elements in a simulation using hydroelastic model. + + This class extends the BaseContact class and provides methods for detecting self-contact + between tetrahedral elements, computing contact pairs, and managing contact-related computations. + """ + + def __init__( + self, + simulator: "Simulator", + ) -> None: + super().__init__(simulator) + self.name = "FEMSelfTetContact" + self.fem_solver = self.sim.fem_solver + self.contact_candidate_type = ti.types.struct( + batch_idx=gs.ti_int, # batch index + geom_idx0=gs.ti_int, # index of the FEM element0 + intersection_code0=gs.ti_int, # intersection code for element0 + geom_idx1=gs.ti_int, # index of the FEM element1 + normal=gs.ti_vec3, # contact plane normal + x=gs.ti_vec3, # a point on the contact plane + distance0=gs.ti_vec4, # distance vector for element0 + ) + self.n_contact_candidates = ti.field(gs.ti_int, shape=()) + self.max_contact_candidates = self.fem_solver.n_surface_elements * self.fem_solver._B * 8 + self.contact_candidates = self.contact_candidate_type.field(shape=(self.max_contact_candidates,)) + + self.contact_pair_type = ti.types.struct( + batch_idx=gs.ti_int, # batch index + normal=gs.ti_vec3, # contact plane normal + tangent0=gs.ti_vec3, # contact plane tangent0 + tangent1=gs.ti_vec3, # contact plane tangent1 + geom_idx0=gs.ti_int, # index of the FEM element0 + geom_idx1=gs.ti_int, # index of the FEM element1 + barycentric0=gs.ti_vec4, # barycentric coordinates of the contact point in tet 0 + barycentric1=gs.ti_vec4, # barycentric coordinates of the contact point in tet 1 + contact_pos=gs.ti_vec3, # contact position + sap_info=self.sap_contact_info_type, # contact info + ) + self.max_contact_pairs = self.fem_solver.n_surface_elements * self.fem_solver._B + self.contact_pairs = self.contact_pair_type.field(shape=(self.max_contact_pairs,)) + + @ti.kernel + def compute_aabb(self, i_step: ti.i32): + aabbs = ti.static(self.coupler.fem_surface_tet_aabb.aabbs) + for i_b, i_se in ti.ndrange(self.fem_solver._B, self.fem_solver.n_surface_elements): + aabbs[i_b, i_se].min.fill(np.inf) + aabbs[i_b, i_se].max.fill(-np.inf) + i_e = self.fem_solver.surface_elements[i_se] + i_v = self.fem_solver.elements_i[i_e].el2v + + for i in ti.static(range(4)): + pos_v = self.fem_solver.elements_v[i_step, i_v[i], i_b].pos + aabbs[i_b, i_se].min = ti.min(aabbs[i_b, i_se].min, pos_v) + aabbs[i_b, i_se].max = ti.max(aabbs[i_b, i_se].max, pos_v) + + @ti.kernel + def compute_candidates(self, i_step: ti.i32): + candidates = ti.static(self.contact_candidates) + self.n_contact_candidates[None] = 0 + for i_r in ti.ndrange(self.coupler.fem_surface_tet_bvh.query_result_count[None]): + i_b, i_sa, i_sq = self.coupler.fem_surface_tet_bvh.query_result[i_r] + i_a = self.fem_solver.surface_elements[i_sa] + i_q = self.fem_solver.surface_elements[i_sq] + i_v0 = self.fem_solver.elements_i[i_a].el2v[0] + i_v1 = self.fem_solver.elements_i[i_q].el2v[0] + x0 = self.fem_solver.elements_v[i_step, i_v0, i_b].pos + x1 = self.fem_solver.elements_v[i_step, i_v1, i_b].pos + p0 = self.coupler.fem_pressure[i_v0] + p1 = self.coupler.fem_pressure[i_v1] + g0 = self.coupler.fem_pressure_gradient[i_b, i_a] + g1 = self.coupler.fem_pressure_gradient[i_b, i_q] + g0_norm = g0.norm() + g1_norm = g1.norm() + if g0_norm < gs.EPS or g1_norm < gs.EPS: + continue + # Calculate the isosurface, i.e. equal pressure plane defined by x and normal + # Solve for p0 + g0.dot(x - x0) = p1 + g1.dot(x - x1) + normal = g0 - g1 + magnitude = normal.norm() + if magnitude < gs.EPS: + continue + normal /= magnitude + b = p1 - p0 - g1.dot(x1) + g0.dot(x0) + x = b / magnitude * normal + # Check that the normal is pointing along g0 and against g1, some allowance as used in Drake + if normal.dot(g0) < COS_ANGLE_THRESHOLD * g0_norm or normal.dot(g1) > -COS_ANGLE_THRESHOLD * g1_norm: + continue + + intersection_code0 = ti.int32(0) + distance0 = ti.Vector.zero(gs.ti_float, 4) + intersection_code1 = ti.int32(0) + distance1 = ti.Vector.zero(gs.ti_float, 4) + for i in ti.static(range(4)): + i_v = self.fem_solver.elements_i[i_a].el2v[i] + pos_v = self.fem_solver.elements_v[i_step, i_v, i_b].pos + distance0[i] = (pos_v - x).dot(normal) # signed distance + if distance0[i] > 0.0: + intersection_code0 |= 1 << i + for i in ti.static(range(4)): + i_v = self.fem_solver.elements_i[i_q].el2v[i] + pos_v = self.fem_solver.elements_v[i_step, i_v, i_b].pos + distance1[i] = (pos_v - x).dot(normal) + if distance1[i] > 0.0: + intersection_code1 |= 1 << i + # Fast check for whether both tets intersect with the plane + if ( + intersection_code0 == 0 + or intersection_code1 == 0 + or intersection_code0 == 15 + or intersection_code1 == 15 + ): + continue + i_c = ti.atomic_add(self.n_contact_candidates[None], 1) + if i_c < self.max_contact_candidates: + candidates[i_c].batch_idx = i_b + candidates[i_c].normal = normal + candidates[i_c].x = x + candidates[i_c].geom_idx0 = i_a + candidates[i_c].intersection_code0 = intersection_code0 + candidates[i_c].distance0 = distance0 + candidates[i_c].geom_idx1 = i_q + + @ti.kernel + def compute_pairs(self, i_step: ti.i32): + """ + Computes the FEM self contact pairs and their properties. + + Intersection code reference: + https://github.com/RobotLocomotion/drake/blob/8c3a249184ed09f0faab3c678536d66d732809ce/geometry/proximity/field_intersection.cc#L87 + """ + candidates = ti.static(self.contact_candidates) + pairs = ti.static(self.contact_pairs) + sap_info = ti.static(pairs.sap_info) + normal_signs = ti.Vector([1.0, -1.0, 1.0, -1.0], dt=gs.ti_float) # make normal point outward + self.n_contact_pairs[None] = 0 + for i_c in range(self.n_contact_candidates[None]): + i_b = candidates[i_c].batch_idx + i_e0 = candidates[i_c].geom_idx0 + i_e1 = candidates[i_c].geom_idx1 + intersection_code0 = candidates[i_c].intersection_code0 + distance0 = candidates[i_c].distance0 + intersected_edges0 = self.coupler.MarchingTetsEdgeTable[intersection_code0] + tet_vertices0 = ti.Matrix.zero(gs.ti_float, 3, 4) # 4 vertices of tet 0 + tet_pressures0 = ti.Vector.zero(gs.ti_float, 4) # pressures at the vertices of tet 0 + tet_vertices1 = ti.Matrix.zero(gs.ti_float, 3, 4) # 4 vertices of tet 1 + + for i in ti.static(range(4)): + i_v = self.fem_solver.elements_i[i_e0].el2v[i] + tet_vertices0[:, i] = self.fem_solver.elements_v[i_step, i_v, i_b].pos + tet_pressures0[i] = self.coupler.fem_pressure[i_v] + + for i in ti.static(range(4)): + i_v = self.fem_solver.elements_i[i_e1].el2v[i] + tet_vertices1[:, i] = self.fem_solver.elements_v[i_step, i_v, i_b].pos + + polygon_vertices = ti.Matrix.zero(gs.ti_float, 3, 8) # maximum 8 vertices + polygon_n_vertices = gs.ti_int(0) + clipped_vertices = ti.Matrix.zero(gs.ti_float, 3, 8) # maximum 8 vertices + clipped_n_vertices = gs.ti_int(0) + for i in range(4): + if intersected_edges0[i] >= 0: + edge = self.coupler.TetEdges[intersected_edges0[i]] + pos_v0 = tet_vertices0[:, edge[0]] + pos_v1 = tet_vertices0[:, edge[1]] + d_v0 = distance0[edge[0]] + d_v1 = distance0[edge[1]] + t = d_v0 / (d_v0 - d_v1) + polygon_vertices[:, polygon_n_vertices] = pos_v0 + t * (pos_v1 - pos_v0) + polygon_n_vertices += 1 + # Intersects the polygon with the four halfspaces of the four triangles + # of the tetrahedral element1. + for face in range(4): + clipped_n_vertices = 0 + x = tet_vertices1[:, (face + 1) % 4] + normal = (tet_vertices1[:, (face + 2) % 4] - x).cross( + tet_vertices1[:, (face + 3) % 4] - x + ) * normal_signs[face] + normal /= normal.norm() + + distances = ti.Vector.zero(gs.ti_float, 8) + for i in range(polygon_n_vertices): + distances[i] = (polygon_vertices[:, i] - x).dot(normal) + + for i in range(polygon_n_vertices): + j = (i + 1) % polygon_n_vertices + if distances[i] <= 0.0: + clipped_vertices[:, clipped_n_vertices] = polygon_vertices[:, i] + clipped_n_vertices += 1 + if distances[j] > 0.0: + wa = distances[j] / (distances[j] - distances[i]) + wb = 1.0 - wa + clipped_vertices[:, clipped_n_vertices] = ( + wa * polygon_vertices[:, i] + wb * polygon_vertices[:, j] + ) + clipped_n_vertices += 1 + elif distances[j] <= 0.0: + wa = distances[j] / (distances[j] - distances[i]) + wb = 1.0 - wa + clipped_vertices[:, clipped_n_vertices] = ( + wa * polygon_vertices[:, i] + wb * polygon_vertices[:, j] + ) + clipped_n_vertices += 1 + polygon_n_vertices = clipped_n_vertices + polygon_vertices = clipped_vertices + + if polygon_n_vertices < 3: + # If the polygon has less than 3 vertices, it is not a valid contact + break + + if polygon_n_vertices < 3: + continue + + # compute centroid and area of the polygon + total_area = gs.EPS # avoid division by zero + total_area_weighted_centroid = ti.Vector.zero(gs.ti_float, 3) + for i in range(2, polygon_n_vertices): + e1 = polygon_vertices[:, i - 1] - polygon_vertices[:, 0] + e2 = polygon_vertices[:, i] - polygon_vertices[:, 0] + 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 + ) + + if total_area < gs.EPS: + continue + + centroid = total_area_weighted_centroid / total_area + barycentric0 = tet_barycentric(centroid, tet_vertices0) + barycentric1 = tet_barycentric(centroid, tet_vertices1) + tangent0 = polygon_vertices[:, 0] - centroid + tangent0 /= tangent0.norm() + tangent1 = candidates[i_c].normal.cross(tangent0) + i_p = ti.atomic_add(self.n_contact_pairs[None], 1) + if i_p < self.max_contact_pairs: + pairs[i_p].batch_idx = i_b + pairs[i_p].normal = candidates[i_c].normal + pairs[i_p].tangent0 = tangent0 + pairs[i_p].tangent1 = tangent1 + pairs[i_p].geom_idx0 = i_e0 + pairs[i_p].geom_idx1 = i_e1 + pairs[i_p].barycentric0 = barycentric0 + pairs[i_p].barycentric1 = barycentric1 + pressure = ( + barycentric0[0] * tet_pressures0[0] + + barycentric0[1] * tet_pressures0[1] + + barycentric0[2] * tet_pressures0[2] + + barycentric0[3] * tet_pressures0[3] + ) + + deformable_g = self.coupler._hydroelastic_stiffness + deformable_k = total_area * deformable_g + # FIXME This is an approximated value, different from Drake, which actually calculates the distance + deformable_phi0 = -pressure / deformable_g * 2 + sap_info[i_p].k = deformable_k + sap_info[i_p].phi0 = deformable_phi0 + sap_info[i_p].mu = ti.sqrt( + self.fem_solver.elements_i[i_e0].friction_mu * self.fem_solver.elements_i[i_e1].friction_mu + ) + + def detection(self, i_step: ti.i32): + self.compute_aabb(i_step) + self.coupler.fem_surface_tet_bvh.build() + self.coupler.fem_surface_tet_bvh.query(self.coupler.fem_surface_tet_aabb.aabbs) + if ( + self.coupler.fem_surface_tet_bvh.query_result_count[None] + > self.coupler.fem_surface_tet_bvh.max_n_query_results + ): + raise ValueError( + f"Query result count {self.coupler.fem_surface_tet_bvh.query_result_count[None]} " + f"exceeds max_n_query_results {self.coupler.fem_surface_tet_bvh.max_n_query_results}" + ) + self.compute_candidates(i_step) + if self.n_contact_candidates[None] > self.max_contact_candidates: + raise ValueError( + f"{self.name} number of contact candidates {self.n_contact_candidates[None]} " + f"exceeds max_contact_candidates {self.max_contact_candidates}" + ) + self.compute_pairs(i_step) + if self.n_contact_pairs[None] > self.max_contact_pairs: + raise ValueError( + f"{self.name} number of contact pairs {self.n_contact_pairs[None]} " + f"exceeds max_contact_pairs {self.max_contact_pairs}" + ) + + @ti.func + def compute_Jx(self, i_p, x): + """ + Compute the contact Jacobian J times a vector x. + """ + pairs = ti.static(self.contact_pairs) + i_b = pairs[i_p].batch_idx + i_g0 = pairs[i_p].geom_idx0 + i_g1 = pairs[i_p].geom_idx1 + Jx = ti.Vector.zero(gs.ti_float, 3) + for i in ti.static(range(4)): + i_v = self.fem_solver.elements_i[i_g0].el2v[i] + Jx += pairs[i_p].barycentric0[i] * x[i_b, i_v] + for i in ti.static(range(4)): + i_v = self.fem_solver.elements_i[i_g1].el2v[i] + Jx -= pairs[i_p].barycentric1[i] * x[i_b, i_v] + Jx = ti.Vector( + [Jx.dot(pairs[i_p].tangent0), Jx.dot(pairs[i_p].tangent1), Jx.dot(pairs[i_p].normal)], dt=gs.ti_float + ) + return Jx + + @ti.func + def compute_contact_point(self, i_p, x, f): + """ + Compute the contact point for a given contact pair. + """ + pairs = ti.static(self.contact_pairs) + i_b = pairs[i_p].batch_idx + i_g0 = pairs[i_p].geom_idx0 + i_g1 = pairs[i_p].geom_idx1 + Jx = ti.Vector.zero(gs.ti_float, 3) + for i in ti.static(range(4)): + i_v = self.fem_solver.elements_i[i_g0].el2v[i] + Jx += pairs[i_p].barycentric0[i] * x[f, i_v, i_b] + for i in ti.static(range(4)): + i_v = self.fem_solver.elements_i[i_g1].el2v[i] + Jx += pairs[i_p].barycentric1[i] * x[f, i_v, i_b] + return Jx * 0.5 + + @ti.func + def add_Jt_x(self, y, i_p, x): + pairs = ti.static(self.contact_pairs) + i_b = pairs[i_p].batch_idx + i_g0 = pairs[i_p].geom_idx0 + i_g1 = pairs[i_p].geom_idx1 + world = ti.Matrix.cols([pairs[i_p].tangent0, pairs[i_p].tangent1, pairs[i_p].normal]) + x_ = world @ x + for i in ti.static(range(4)): + i_v = self.fem_solver.elements_i[i_g0].el2v[i] + y[i_b, i_v] += pairs[i_p].barycentric0[i] * x_ + for i in ti.static(range(4)): + i_v = self.fem_solver.elements_i[i_g1].el2v[i] + y[i_b, i_v] -= pairs[i_p].barycentric1[i] * x_ + + @ti.func + def add_Jt_A_J_diag3x3(self, y, i_p, A): + pairs = ti.static(self.contact_pairs) + i_b = pairs[i_p].batch_idx + i_g0 = pairs[i_p].geom_idx0 + i_g1 = pairs[i_p].geom_idx1 + world = ti.Matrix.cols([pairs[i_p].tangent0, pairs[i_p].tangent1, pairs[i_p].normal]) + B_ = world @ A @ world.transpose() + for i in ti.static(range(4)): + i_v = self.fem_solver.elements_i[i_g0].el2v[i] + y[i_b, i_v] += pairs[i_p].barycentric0[i] ** 2 * B_ + for i in ti.static(range(4)): + i_v = self.fem_solver.elements_i[i_g1].el2v[i] + y[i_b, i_v] += pairs[i_p].barycentric1[i] ** 2 * B_ + + @ti.func + def compute_delassus(self, i_p): + pairs = ti.static(self.contact_pairs) + i_b = pairs[i_p].batch_idx + i_g0 = pairs[i_p].geom_idx0 + i_g1 = pairs[i_p].geom_idx1 + world = ti.Matrix.cols([pairs[i_p].tangent0, pairs[i_p].tangent1, pairs[i_p].normal]) + W = ti.Matrix.zero(gs.ti_float, 3, 3) + # W = sum (JA^-1J^T) + # With floor, J is Identity + for i in ti.static(range(4)): + i_v = self.fem_solver.elements_i[i_g0].el2v[i] + W += pairs[i_p].barycentric0[i] ** 2 * self.fem_solver.pcg_state_v[i_b, i_v].prec + for i in ti.static(range(4)): + i_v = self.fem_solver.elements_i[i_g1].el2v[i] + W += pairs[i_p].barycentric1[i] ** 2 * self.fem_solver.pcg_state_v[i_b, i_v].prec + W = world.transpose() @ W @ world + return W + + +@ti.data_oriented +class FEMFloorVertContact(BaseContact): + """ + Class for handling contact between tetrahedral elements and a floor in a simulation using point contact model. + + This class extends the BaseContact class and provides methods for detecting contact + between the tetrahedral elements and the floor, computing contact pairs, and managing + contact-related computations. + """ + + def __init__( + self, + simulator: "Simulator", + ) -> None: + super().__init__(simulator) + self.name = "FEMFloorVertContact" + self.fem_solver = self.sim.fem_solver + + self.contact_pair_type = ti.types.struct( + batch_idx=gs.ti_int, # batch index + geom_idx=gs.ti_int, # index of the vertex + contact_pos=gs.ti_vec3, # contact position + sap_info=self.sap_contact_info_type, # contact info + ) + self.max_contact_pairs = self.fem_solver.n_surface_elements * self.fem_solver._B + self.contact_pairs = self.contact_pair_type.field(shape=(self.max_contact_pairs,)) + + @ti.kernel + def detection(self, i_step: ti.i32): + pairs = ti.static(self.contact_pairs) + sap_info = ti.static(pairs.sap_info) + # Compute contact pairs + self.n_contact_pairs[None] = 0 + for i_b, i_sv in ti.ndrange(self.coupler._B, self.fem_solver.n_surface_vertices): + i_v = self.fem_solver.surface_vertices[i_sv] + pos_v = self.fem_solver.elements_v[i_step, i_v, i_b].pos + distance = pos_v.z - self.fem_solver.floor_height + if distance > 0.0: + continue + i_p = ti.atomic_add(self.n_contact_pairs[None], 1) + if i_p < self.max_contact_pairs: + pairs[i_p].batch_idx = i_b + pairs[i_p].geom_idx = i_v + sap_info[i_p].k = self.coupler._point_contact_stiffness * self.fem_solver.surface_vert_mass[i_v] + sap_info[i_p].phi0 = distance + sap_info[i_p].mu = self.fem_solver.elements_v_info[i_v].friction_mu + + @ti.func + def compute_Jx(self, i_p, x): + """ + Compute the contact Jacobian J times a vector x. + """ + pairs = ti.static(self.contact_pairs) + i_b = pairs[i_p].batch_idx + i_g = pairs[i_p].geom_idx + Jx = x[i_b, i_g] + return Jx + + @ti.func + def compute_contact_point(self, i_p, x, i_step): + """ + Compute the contact point for a given contact pair. + """ + pairs = ti.static(self.contact_pairs) + i_b = pairs[i_p].batch_idx + i_g = pairs[i_p].geom_idx + Jx = x[i_step, i_g, i_b] + return Jx + + @ti.func + def add_Jt_x(self, y, i_p, x): + pairs = ti.static(self.contact_pairs) + i_b = pairs[i_p].batch_idx + i_g = pairs[i_p].geom_idx + y[i_b, i_g] += x + + @ti.func + def add_Jt_A_J_diag3x3(self, y, i_p, A): + pairs = ti.static(self.contact_pairs) + i_b = pairs[i_p].batch_idx + i_g = pairs[i_p].geom_idx + y[i_b, i_g] += A + + @ti.func + def compute_delassus(self, i_p): + pairs = ti.static(self.contact_pairs) + i_b = pairs[i_p].batch_idx + i_g = pairs[i_p].geom_idx + # W = sum (JA^-1J^T) + # With floor, J is Identity + W = self.fem_solver.pcg_state_v[i_b, i_g].prec + return W diff --git a/genesis/engine/entities/fem_entity.py b/genesis/engine/entities/fem_entity.py index f5f2aa4b2c..5a38a7e269 100644 --- a/genesis/engine/entities/fem_entity.py +++ b/genesis/engine/entities/fem_entity.py @@ -8,7 +8,7 @@ import genesis.utils.element as eu import genesis.utils.geom as gu import genesis.utils.mesh as mu -from genesis.engine.coupler import SAPCoupler +from genesis.engine.couplers import SAPCoupler from genesis.engine.states.cache import QueriedStates from genesis.engine.states.entities import FEMEntityState from genesis.utils.misc import to_gs_tensor, tensor_to_array @@ -356,7 +356,7 @@ def sample(self): else: gs.raise_exception(f"Unsupported morph: {self.morph}.") - self.instantiate(verts, elems) + self.instantiate(*eu.split_all_surface_tets(verts, elems)) def _add_to_solver(self, in_backward=False): if not in_backward: diff --git a/genesis/engine/scene.py b/genesis/engine/scene.py index 7f67385fc7..7ce549a233 100644 --- a/genesis/engine/scene.py +++ b/genesis/engine/scene.py @@ -16,7 +16,8 @@ from genesis.engine.simulator import Simulator from genesis.options import ( AvatarOptions, - CouplerOptions, + BaseCouplerOptions, + LegacyCouplerOptions, FEMOptions, MPMOptions, PBDOptions, @@ -82,7 +83,7 @@ class Scene(RBC): def __init__( self, sim_options: SimOptions | None = None, - coupler_options: CouplerOptions | None = None, + coupler_options: BaseCouplerOptions | None = None, tool_options: ToolOptions | None = None, rigid_options: RigidOptions | None = None, avatar_options: AvatarOptions | None = None, @@ -100,7 +101,7 @@ def __init__( ): # Handling of default arguments sim_options = sim_options or SimOptions() - coupler_options = coupler_options or CouplerOptions() + coupler_options = coupler_options or LegacyCouplerOptions() tool_options = tool_options or ToolOptions() rigid_options = rigid_options or RigidOptions() avatar_options = avatar_options or AvatarOptions() @@ -200,7 +201,7 @@ def __init__( def _validate_options( self, sim_options: SimOptions, - coupler_options: CouplerOptions, + coupler_options: BaseCouplerOptions, tool_options: ToolOptions, rigid_options: RigidOptions, avatar_options: AvatarOptions, @@ -217,8 +218,8 @@ def _validate_options( if not isinstance(sim_options, SimOptions): gs.raise_exception("`sim_options` should be an instance of `SimOptions`.") - if not isinstance(coupler_options, CouplerOptions): - gs.raise_exception("`coupler_options` should be an instance of `CouplerOptions`.") + if not isinstance(coupler_options, BaseCouplerOptions): + gs.raise_exception("`coupler_options` should be an instance of `BaseCouplerOptions`.") if not isinstance(tool_options, ToolOptions): gs.raise_exception("`tool_options` should be an instance of `ToolOptions`.") diff --git a/genesis/engine/simulator.py b/genesis/engine/simulator.py index fb09ffef2c..2ab5db2b09 100644 --- a/genesis/engine/simulator.py +++ b/genesis/engine/simulator.py @@ -7,7 +7,8 @@ from genesis.options.morphs import Morph from genesis.options.solvers import ( AvatarOptions, - CouplerOptions, + BaseCouplerOptions, + LegacyCouplerOptions, SAPCouplerOptions, FEMOptions, MPMOptions, @@ -20,7 +21,7 @@ ) from genesis.repr_base import RBC -from .coupler import Coupler, SAPCoupler +from .couplers import LegacyCoupler, SAPCoupler from .entities import HybridEntity from .solvers.base_solver import Solver from .solvers import ( @@ -75,7 +76,7 @@ def __init__( self, scene: "Scene", options: SimOptions, - coupler_options: CouplerOptions, + coupler_options: BaseCouplerOptions, tool_options: ToolOptions, rigid_options: RigidOptions, avatar_options: AvatarOptions, @@ -137,8 +138,12 @@ def __init__( # coupler if isinstance(self.coupler_options, SAPCouplerOptions): self._coupler = SAPCoupler(self, self.coupler_options) + elif isinstance(self.coupler_options, LegacyCouplerOptions): + self._coupler = LegacyCoupler(self, self.coupler_options) else: - self._coupler = Coupler(self, self.coupler_options) + gs.raise_exception( + f"Coupler options {self.coupler_options} not supported. Please use SAPCouplerOptions or LegacyCouplerOptions." + ) # states self._queried_states = QueriedStates() diff --git a/genesis/engine/solvers/fem_solver.py b/genesis/engine/solvers/fem_solver.py index 9b590eca71..434d94a9f4 100644 --- a/genesis/engine/solvers/fem_solver.py +++ b/genesis/engine/solvers/fem_solver.py @@ -8,6 +8,7 @@ from .base_solver import Solver import numpy as np +import igl @ti.data_oriented @@ -154,7 +155,9 @@ def init_element_fields(self): element_v_info = ti.types.struct( mass=gs.ti_float, # mass of the vertex + mass_inv=gs.ti_float, # inverse mass of the vertex mass_over_dt2=gs.ti_float, # scaled mass of the vertex over dt^2 + friction_mu=gs.ti_float, # friction coefficient for contact ) pcg_state_v = ti.types.struct( @@ -265,6 +268,62 @@ def init_surface_fields(self): layout=ti.Layout.SOA, ) + def _init_surface_info(self): + self.vertices_on_surface = ti.field(dtype=ti.u1, shape=(self.n_vertices)) + self.elements_on_surface = ti.field(dtype=ti.u1, shape=(self.n_elements)) + self.compute_surface_vertices() + self.compute_surface_elements() + vertices_on_surface_np = self.vertices_on_surface.to_numpy() + elements_on_surface_np = self.elements_on_surface.to_numpy() + surface_vertices_np = np.where(vertices_on_surface_np)[0].flatten() + self.surface_vertices = ti.field( + dtype=ti.i32, + shape=(len(surface_vertices_np),), + needs_grad=False, + ) + self.surface_vertices.from_numpy(surface_vertices_np) + surface_elements_np = elements_on_surface_np.nonzero()[0].reshape((-1,)) + self.surface_elements = ti.field( + dtype=ti.i32, + shape=(len(surface_elements_np),), + needs_grad=False, + ) + self.surface_elements.from_numpy(surface_elements_np) + + surface_triangles_np = self.surface.tri2v.to_numpy().reshape(-1, 3) + pos_np = self.elements_v.pos.to_numpy()[0, :, 0, :].reshape(-1, 3)[surface_vertices_np] + surface_vertices_mapping = np.full(self.n_vertices, -1, dtype=np.int32) + surface_vertices_mapping[surface_vertices_np] = np.arange(len(surface_vertices_np)) + mass = igl.massmatrix(pos_np, surface_vertices_mapping[surface_triangles_np]) + surface_vert_mass_np = mass.diagonal() + self.surface_vert_mass = ti.field( + dtype=gs.ti_float, + shape=(len(surface_vertices_np),), + needs_grad=False, + ) + self.surface_vert_mass.from_numpy(surface_vert_mass_np) + + @ti.kernel + def compute_surface_vertices(self): + for i_v in range(self.n_vertices): + self.vertices_on_surface[i_v] = 0 + + for i_s in range(self.n_surfaces): + tri2v = self.surface[i_s].tri2v + for i in ti.static(range(3)): + self.vertices_on_surface[tri2v[i]] = 1 + + @ti.kernel + def compute_surface_elements(self): + for i_e in range(self.n_elements): + i_v = self.elements_i[i_e].el2v + self.elements_on_surface[i_e] = ( + self.vertices_on_surface[i_v[0]] + or self.vertices_on_surface[i_v[1]] + or self.vertices_on_surface[i_v[2]] + or self.vertices_on_surface[i_v[3]] + ) + def init_ckpt(self): self._ckpt = dict() @@ -279,6 +338,7 @@ def build(self): super().build() self.n_envs = self.sim.n_envs self._B = self.sim._B + self.tet_wrong_order = ti.field(dtype=ti.u1, shape=(), needs_grad=False) # batch fields self.init_batch_fields() @@ -297,6 +357,14 @@ def build(self): for mat in self._mats: mat.build(self) + if self.n_elements_max > 0: + self._init_surface_info() + if self.tet_wrong_order[None] == 1: + raise RuntimeError( + "The order of vertices in the tetrahedral elements is not correct. " + "Please check the input mesh or the FEM solver implementation." + ) + def add_entity(self, idx, material, morph, surface): # add material's update methods if not matching any existing material exist = False @@ -403,11 +471,7 @@ def compute_pos(self, f: ti.i32): @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) @@ -584,6 +648,7 @@ def accumulate_vertex_force_preconditioner(self, f: ti.i32): gradient = self.elements_el_energy[i_b, i_e].gradient force = -V * gradient @ B.transpose() i_v = self.elements_i[i_e].el2v + # atomic self.elements_v_energy[i_b, i_v[0]].force += force[:, 0] self.elements_v_energy[i_b, i_v[1]].force += force[:, 1] @@ -633,12 +698,12 @@ def accumulate_vertex_force_preconditioner(self, f: ti.i32): # Other options for preconditioner: # Uncomment one of the following lines to test different preconditioners # Use identity for preconditioner - # self.pcg_state[i_v, i_b].prec = ti.Matrix.identity(gs.ti_float, 3) + # self.pcg_state_v[i_b, i_v].prec = ti.Matrix.identity(gs.ti_float, 3) # Use diagonal for preconditioner - # self.pcg_state[i_v, i_b].prec = ti.Matrix([[1 / self.pcg_state[i_v, i_b].diag3x3[0, 0], 0, 0], - # [0, 1 / self.pcg_state[i_v, i_b].diag3x3[1, 1], 0], - # [0, 0, 1 / self.pcg_state[i_v, i_b].diag3x3[2, 2]]]) + # self.pcg_state_v[i_b, i_v].prec = ti.Matrix([[1 / self.pcg_state_v[i_b, i_v].diag3x3[0, 0], 0, 0], + # [0, 1 / self.pcg_state_v[i_b, i_v].diag3x3[1, 1], 0], + # [0, 0, 1 / self.pcg_state_v[i_b, i_v].diag3x3[2, 2]]]) @ti.func def compute_Ap(self): @@ -755,7 +820,7 @@ def one_pcg_iter(self): for i_b in range(self._B): if not self.batch_pcg_active[i_b]: continue - self.pcg_state[i_b].beta = self.pcg_state[i_b].rTr_new / self.pcg_state[i_b].rTr + self.pcg_state[i_b].beta = self.pcg_state[i_b].rTz_new / self.pcg_state[i_b].rTz self.pcg_state[i_b].rTr = self.pcg_state[i_b].rTr_new self.pcg_state[i_b].rTz = self.pcg_state[i_b].rTz_new @@ -841,12 +906,23 @@ def one_linesearch_iter(self, f: ti.i32): continue self.linesearch_state[i_b].step_size *= self._linesearch_tau + @ti.kernel + def skip_linesearch(self, f: ti.i32): + # Inertia, x_prev, m + for i_b, i_v in ti.ndrange(self._B, self.n_vertices): + if not self.batch_active[i_b]: + continue + self.elements_v[f + 1, i_v, i_b].pos = self.elements_v[f + 1, i_v, i_b].pos + self.pcg_state_v[i_b, i_v].x + def linesearch(self, f: ti.i32): """ Note ------ https://en.wikipedia.org/wiki/Backtracking_line_search#Algorithm """ + if self._n_linesearch_iterations <= 0: + self.skip_linesearch(f) + return self.init_linesearch(f) for i in range(self._n_linesearch_iterations): self.one_linesearch_iter(f) @@ -1058,8 +1134,12 @@ def _kernel_add_elements( for j in ti.static(range(3)): self.elements_v[f, i_global, i_b].pos[j] = verts[i_v, j] self.elements_v[f, i_global, i_b].vel = ti.Vector.zero(gs.ti_float, 3) + + for i_v in range(n_verts_local): + i_global = i_v + v_start self.elements_v_info[i_global].mass = 0.0 self.elements_v_info[i_global].mass_over_dt2 = 0.0 + self.elements_v_info[i_global].friction_mu = mat_friction_mu one_over_dt2 = 1.0 / (self.substep_dt**2) n_elems_local = elems.shape[0] @@ -1072,7 +1152,11 @@ def _kernel_add_elements( d = self.elements_v[f, elems[i_e, 3] + v_start, 0].pos B_inv = ti.Matrix.cols([a - d, b - d, c - d]) self.elements_i[i_global].B = B_inv.inverse() - V = ti.abs(B_inv.determinant()) / 6 + det = B_inv.determinant() + # Determinant should be consistently smaller than 0 + if det >= 0.0: + self.tet_wrong_order[None] = True + V = ti.abs(det) / 6.0 self.elements_i[i_global].V = V V_scaled = V * self._vol_scale self.elements_i[i_global].V_scaled = V_scaled @@ -1091,6 +1175,10 @@ def _kernel_add_elements( self.elements_i[i_global].muscle_group = 0 self.elements_i[i_global].muscle_direction = ti.Vector([0.0, 0.0, 1.0], dt=gs.ti_float) + for i_v in range(n_verts_local): + i_global = i_v + v_start + self.elements_v_info[i_global].mass_inv = 1.0 / self.elements_v_info[i_global].mass + for i_e, i_b in ti.ndrange(n_elems_local, self._B): i_global = i_e + el_start self.elements_el[f, i_global, i_b].actu = 0.0 @@ -1317,3 +1405,11 @@ def n_elements_max(self): @property def vol_scale(self): return self._vol_scale + + @property + def n_surface_vertices(self): + return self.surface_vertices.shape[0] + + @property + def n_surface_elements(self): + return self.surface_elements.shape[0] diff --git a/genesis/options/solvers.py b/genesis/options/solvers.py index ff367722b0..3a9a801dc3 100644 --- a/genesis/options/solvers.py +++ b/genesis/options/solvers.py @@ -68,7 +68,15 @@ def __init__(self, **data): self._steps_local = None -class CouplerOptions(Options): +class BaseCouplerOptions(Options): + """ + Base class for all coupler options. + """ + + pass + + +class LegacyCouplerOptions(BaseCouplerOptions): """ Options configuring the inter-solver coupling. @@ -90,8 +98,6 @@ class CouplerOptions(Options): Whether to enable coupling between FEM and MPM solvers. Defaults to True. fem_sph : bool, optional Whether to enable coupling between FEM and SPH solvers. Defaults to True. - hydroelastic_contact : bool, optional - Whether to enable hydroelastic contact. Defaults to False. Experimental """ rigid_mpm: bool = True @@ -104,7 +110,7 @@ class CouplerOptions(Options): fem_sph: bool = True -class SAPCouplerOptions(CouplerOptions): +class SAPCouplerOptions(BaseCouplerOptions): """ Options configuring the inter-solver coupling for the Semi-Analytic Primal (SAP) contact solver used in Drake. @@ -115,16 +121,33 @@ class SAPCouplerOptions(CouplerOptions): n_pcg_iterations : int, optional Number of iterations for the Preconditioned Conjugate Gradient solver. Defaults to 100. n_linesearch_iterations : int, optional - Number of iterations for the line search solver. Defaults to 10. - sap_threshold : float, optional - Threshold for the SAP solver. Defaults to 1e-6. + Max number of iterations for the line search solver. Defaults to 10. + sap_convergence_atol : float, optional + Absolute tolerance for SAP convergence. Defaults to 1e-6. + sap_convergence_rtol : float, optional + Relative tolerance for SAP convergence. Defaults to 1e-5. + sap_taud : float, optional + Dissipation time scale for SAP. Defaults to 0.1. + sap_beta : float, optional + Normal regularization parameter for SAP. Defaults to 1.0. + sap_sigma : float, optional + Friction regularization parameter for SAP. Defaults to 1e-3. pcg_threshold : float, optional Threshold for the Preconditioned Conjugate Gradient solver. Defaults to 1e-6. - linesearch_c : float, optional - Line search sufficient decrease parameter. Defaults to 1e-4. - linesearch_tau : float, optional - Line search step size reduction factor. Defaults to 0.8. - + linesearch_ftol : float, optional + Line search sufficient value close to zero for exact linesearch. Defaults to 1e-6. + linesearch_max_step_size : float, optional + Maximum step size for exact linesearch. Defaults to 1.5. + hydroelastic_stiffness : float, optional + Stiffness for hydroelastic contact. Defaults to 1e8. + point_contact_stiffness : float, optional + Stiffness for point contact. Defaults to 1e8. + fem_floor_type : str, optional + Type of contact against the floor. Defaults to "tet". Can be "tet", "vert", or "none". + Tet would be the default choice for most cases. + Vert would be preferable when the mesh is very coarse, such as a single cube or a tetrahedron. + fem_self_tet : bool, optional + Whether to use tetrahedral based self-contact. Defaults to True. Note ---- Paper reference: https://arxiv.org/abs/2110.10107 @@ -134,10 +157,18 @@ class SAPCouplerOptions(CouplerOptions): n_sap_iterations: int = 5 n_pcg_iterations: int = 100 n_linesearch_iterations: int = 10 - sap_threshold: float = 1e-6 + sap_convergence_atol: float = 1e-6 + sap_convergence_rtol: float = 1e-5 + sap_taud: float = 0.1 + sap_beta: float = 1.0 + sap_sigma: float = 1e-3 pcg_threshold: float = 1e-6 - linesearch_c: float = 1e-4 - linesearch_tau: float = 0.8 + linesearch_ftol: float = 1e-6 + linesearch_max_step_size: float = 1.5 + hydroelastic_stiffness: float = 1e8 + point_contact_stiffness: float = 1e8 + fem_floor_type: str = "tet" + fem_self_tet: bool = True ############################ Solvers inside simulator ############################ @@ -560,11 +591,11 @@ class FEMOptions(Options): Whether to use the implicit solver. Defaults to False. Implicit solver is a more stable solver for FEM. It can be used with a large time step. n_newton_iterations : int, optional - Maximum number of Newton iterations. Defaults to 5. Only used when `use_implicit_solver` is True. + Maximum number of Newton iterations. Defaults to 1. Only used when `use_implicit_solver` is True. n_pcg_iterations : int, optional - Maximum number of PCG iterations. Defaults to 100. Only used when `use_implicit_solver` is True. + Maximum number of PCG iterations. Defaults to 500. Only used when `use_implicit_solver` is True. n_linesearch_iterations : int, optional - Maximum number of line search iterations. Defaults to 10. Only used when `use_implicit_solver` is True. + Maximum number of line search iterations. Defaults to 0. Only used when `use_implicit_solver` is True. newton_dx_threshold : float, optional Threshold for the Newton solver. Defaults to 1e-6. Only used when `use_implicit_solver` is True. pcg_threshold : float, optional @@ -576,7 +607,7 @@ class FEMOptions(Options): damping_alpha : float, optional Rayleigh Damping factor for the implicit solver. Defaults to 0.5. Only used when `use_implicit_solver` is True. damping_beta : float, optional - Rayleigh Damping factor for the implicit solver. Defaults to 1e-4. Only used when `use_implicit_solver` is True. + Rayleigh Damping factor for the implicit solver. Defaults to 5e-4. Only used when `use_implicit_solver` is True. Note ---- @@ -591,15 +622,15 @@ class FEMOptions(Options): damping: Optional[float] = 0.0 floor_height: float = None use_implicit_solver: bool = False - n_newton_iterations: int = 5 - n_pcg_iterations: int = 100 - n_linesearch_iterations: int = 10 + n_newton_iterations: int = 1 + n_pcg_iterations: int = 500 + n_linesearch_iterations: int = 0 newton_dx_threshold: float = 1e-6 pcg_threshold: float = 1e-6 linesearch_c: float = 1e-4 linesearch_tau: float = 0.5 damping_alpha: float = 0.5 - damping_beta: float = 1e-4 + damping_beta: float = 5e-4 class SFOptions(Options): diff --git a/genesis/utils/element.py b/genesis/utils/element.py index e413898b6c..abc4d80c38 100644 --- a/genesis/utils/element.py +++ b/genesis/utils/element.py @@ -4,6 +4,8 @@ import numpy as np import trimesh +import igl + import genesis as gs from . import geom as gu @@ -60,3 +62,34 @@ def mesh_to_elements(file, pos=(0, 0, 0), scale=1.0, tet_cfg=dict()): verts += np.array(pos) return verts, elems + + +def split_all_surface_tets(verts, elems): + """ + Splits tetrahedras that have 4 vertices on the surface into 4 smaller tetrahedras. + + This is useful for the hydroelastic contact model. + """ + F, *_ = igl.boundary_facets(elems) + on_surface = np.zeros(verts.shape[0], dtype=bool) + on_surface[ + F.reshape( + -1, + ) + ] = True + all_on_surface = np.all(on_surface[elems], axis=1) + if not all_on_surface.any(): + return verts, elems + bad_elems = elems[all_on_surface] + new_verts = np.mean(verts[bad_elems], axis=1, dtype=np.float32) + new_elems = [] + for idx, (v0, v1, v2, v3) in enumerate(bad_elems, len(verts)): + new_elems.append([v0, v1, v2, idx]) + new_elems.append([v0, v1, idx, v3]) + new_elems.append([v0, idx, v2, v3]) + new_elems.append([idx, v1, v2, v3]) + new_elems = np.array(new_elems, dtype=np.int32) + verts = np.concatenate([verts, new_verts], axis=0) + # remove the bad elements from the original elements + elems = np.concatenate([elems[~all_on_surface], new_elems], axis=0) + return verts, elems diff --git a/tests/test_fem.py b/tests/test_fem.py index f3ebcde627..5b67ef2619 100644 --- a/tests/test_fem.py +++ b/tests/test_fem.py @@ -5,7 +5,7 @@ import genesis as gs from genesis.utils.misc import tensor_to_array -from .utils import assert_allclose +from .utils import assert_allclose, get_hf_assets @pytest.fixture(scope="session") @@ -206,7 +206,6 @@ def fem_material_linear(): def test_sphere_box_fall_implicit_fem_coupler(fem_material_linear, show_viewer): - """Test adding multiple FEM entities to the scene""" scene = gs.Scene( sim_options=gs.options.SimOptions( dt=1.0 / 60.0, @@ -246,15 +245,14 @@ def test_sphere_box_fall_implicit_fem_coupler(fem_material_linear, show_viewer): for entity in scene.entities: state = entity.get_state() - pos = tensor_to_array(state.pos) - min_pos_z = np.min(pos[..., 2]) + min_pos_z = state.pos[..., 2].min() + # The contact requires some penetration to generate enough contact force to cancel out gravity assert_allclose( min_pos_z, 0.0, atol=5e-2 ), f"Entity {entity.uid} minimum Z position {min_pos_z} is not close to 0.0." def test_sphere_fall_implicit_fem_sap_coupler(fem_material_linear, show_viewer): - """Test adding multiple FEM entities to the scene""" scene = gs.Scene( sim_options=gs.options.SimOptions( dt=1.0 / 60.0, @@ -285,11 +283,11 @@ def test_sphere_fall_implicit_fem_sap_coupler(fem_material_linear, show_viewer): for entity in scene.entities: state = entity.get_state() - pos = tensor_to_array(state.pos) - min_pos_z = np.min(pos[..., 2]) + min_pos_z = state.pos[..., 2].min() + # The contact requires some penetration to generate enough contact force to cancel out gravity assert_allclose( - min_pos_z, 0.0, atol=1e-3 - ), f"Entity {entity.uid} minimum Z position {min_pos_z} is not close to 0.0." + min_pos_z, -1e-3, atol=1e-4 + ), f"Entity {entity.uid} minimum Z position {min_pos_z} is not close to -1e-3." @pytest.fixture(scope="session") @@ -299,12 +297,12 @@ def fem_material_linear_corotated(): def test_linear_corotated_sphere_fall_implicit_fem_sap_coupler(fem_material_linear_corotated, show_viewer): - """Test adding multiple FEM entities to the scene""" scene = gs.Scene( sim_options=gs.options.SimOptions( dt=1.0 / 60.0, substeps=2, ), + # Not using default fem_options to make it faster, linear material only need one iteration without linesearch fem_options=gs.options.FEMOptions( use_implicit_solver=True, ), @@ -332,9 +330,10 @@ def test_linear_corotated_sphere_fall_implicit_fem_sap_coupler(fem_material_line state = entity.get_state() pos = tensor_to_array(state.pos.reshape(-1, 3)) min_pos_z = np.min(pos[..., 2]) + # The contact requires some penetration to generate enough contact force to cancel out gravity assert_allclose( - min_pos_z, 0.0, atol=1.5e-3 - ), f"Entity {entity.uid} minimum Z position {min_pos_z} is not close to 0.0." + min_pos_z, -1e-3, atol=1e-4 + ), f"Entity {entity.uid} minimum Z position {min_pos_z} is not close to -1e-3." BV, BF = igl.bounding_box(pos) x_scale = BV[0, 0] - BV[-1, 0] y_scale = BV[0, 1] - BV[-1, 1] @@ -343,3 +342,60 @@ def test_linear_corotated_sphere_fall_implicit_fem_sap_coupler(fem_material_line assert_allclose(y_scale, 0.2, atol=1e-3), f"Entity {entity.uid} Y scale {y_scale} is not close to 0.2." # The Z scale is expected to be more squashed due to gravity assert_allclose(z_scale, 0.2, atol=2e-3), f"Entity {entity.uid} Z scale {z_scale} is not close to 0.2." + + +@pytest.fixture(scope="session") +def fem_material_linear_corotated_soft(): + """Fixture for common FEM linear material properties""" + return gs.materials.FEM.Elastic(model="linear_corotated", E=1.0e5, nu=0.4) + + +def test_fem_sphere_box_self(fem_material_linear_corotated, fem_material_linear_corotated_soft, show_viewer): + scene = gs.Scene( + sim_options=gs.options.SimOptions( + dt=1 / 60, + substeps=2, + ), + fem_options=gs.options.FEMOptions( + use_implicit_solver=True, + ), + coupler_options=gs.options.SAPCouplerOptions(), + show_viewer=show_viewer, + ) + + # Add first FEM entity + scene.add_entity( + morph=gs.morphs.Sphere( + pos=(0.0, 0.0, 0.1), + radius=0.1, + ), + material=fem_material_linear_corotated, + ) + + # Add second FEM entity + scale = 0.1 + asset_path = get_hf_assets(pattern="meshes/cube8.obj") + scene.add_entity( + morph=gs.morphs.Mesh( + file=f"{asset_path}/meshes/cube8.obj", + scale=scale, + pos=(0.0, 0.0, scale * 4.0), + ), + material=fem_material_linear_corotated, + ) + + # Build the scene + scene.build() + # Run simulation + for _ in range(200): + scene.step() + + depths = [-1e-3, -2e-5] + atols = [2e-4, 4e-6] + for i, entity in enumerate(scene.entities): + state = entity.get_state() + min_pos_z = state.pos[..., 2].min() + # The contact requires some penetration to generate enough contact force to cancel out gravity + assert_allclose( + min_pos_z, depths[i], atol=atols[i] + ), f"Entity {entity.uid} minimum Z position {min_pos_z} is not close to {depths[i]}."