Spaces:
Sleeping
Sleeping
| from typing import Optional | |
| import warp as wp | |
| from warp.fem.cache import ( | |
| TemporaryStore, | |
| borrow_temporary, | |
| borrow_temporary_like, | |
| cached_arg_value, | |
| ) | |
| from warp.fem.types import OUTSIDE, Coords, ElementIndex, Sample, make_free_sample | |
| from .element import LinearEdge, Square | |
| from .geometry import Geometry | |
| # from .closest_point import project_on_tet_at_origin | |
| class Quadmesh2DCellArg: | |
| quad_vertex_indices: wp.array2d(dtype=int) | |
| positions: wp.array(dtype=wp.vec2) | |
| # for neighbor cell lookup | |
| vertex_quad_offsets: wp.array(dtype=int) | |
| vertex_quad_indices: wp.array(dtype=int) | |
| class Quadmesh2DSideArg: | |
| cell_arg: Quadmesh2DCellArg | |
| edge_vertex_indices: wp.array(dtype=wp.vec2i) | |
| edge_quad_indices: wp.array(dtype=wp.vec2i) | |
| class Quadmesh2D(Geometry): | |
| """Two-dimensional quadrilateral mesh geometry""" | |
| dimension = 2 | |
| def __init__( | |
| self, quad_vertex_indices: wp.array, positions: wp.array, temporary_store: Optional[TemporaryStore] = None | |
| ): | |
| """ | |
| Constructs a two-dimensional quadrilateral mesh. | |
| Args: | |
| quad_vertex_indices: warp array of shape (num_tris, 4) containing vertex indices for each quad, in counter-clockwise order | |
| positions: warp array of shape (num_vertices, 2) containing 2d position for each vertex | |
| temporary_store: shared pool from which to allocate temporary arrays | |
| """ | |
| self.quad_vertex_indices = quad_vertex_indices | |
| self.positions = positions | |
| self._edge_vertex_indices: wp.array = None | |
| self._edge_quad_indices: wp.array = None | |
| self._vertex_quad_offsets: wp.array = None | |
| self._vertex_quad_indices: wp.array = None | |
| self._build_topology(temporary_store) | |
| def cell_count(self): | |
| return self.quad_vertex_indices.shape[0] | |
| def vertex_count(self): | |
| return self.positions.shape[0] | |
| def side_count(self): | |
| return self._edge_vertex_indices.shape[0] | |
| def boundary_side_count(self): | |
| return self._boundary_edge_indices.shape[0] | |
| def reference_cell(self) -> Square: | |
| return Square() | |
| def reference_side(self) -> LinearEdge: | |
| return LinearEdge() | |
| def edge_quad_indices(self) -> wp.array: | |
| return self._edge_quad_indices | |
| def edge_vertex_indices(self) -> wp.array: | |
| return self._edge_vertex_indices | |
| CellArg = Quadmesh2DCellArg | |
| SideArg = Quadmesh2DSideArg | |
| class SideIndexArg: | |
| boundary_edge_indices: wp.array(dtype=int) | |
| # Geometry device interface | |
| def cell_arg_value(self, device) -> CellArg: | |
| args = self.CellArg() | |
| args.quad_vertex_indices = self.quad_vertex_indices.to(device) | |
| args.positions = self.positions.to(device) | |
| args.vertex_quad_offsets = self._vertex_quad_offsets.to(device) | |
| args.vertex_quad_indices = self._vertex_quad_indices.to(device) | |
| return args | |
| def cell_position(args: CellArg, s: Sample): | |
| quad_idx = args.quad_vertex_indices[s.element_index] | |
| w_p = s.element_coords | |
| w_m = Coords(1.0) - s.element_coords | |
| # 0 : m m | |
| # 1 : p m | |
| # 2 : p p | |
| # 3 : m p | |
| return ( | |
| w_m[0] * w_m[1] * args.positions[quad_idx[0]] | |
| + w_p[0] * w_m[1] * args.positions[quad_idx[1]] | |
| + w_p[0] * w_p[1] * args.positions[quad_idx[2]] | |
| + w_m[0] * w_p[1] * args.positions[quad_idx[3]] | |
| ) | |
| def cell_deformation_gradient(cell_arg: CellArg, s: Sample): | |
| """Deformation gradient at `coords`""" | |
| quad_idx = cell_arg.quad_vertex_indices[s.element_index] | |
| w_p = s.element_coords | |
| w_m = Coords(1.0) - s.element_coords | |
| return ( | |
| wp.outer(cell_arg.positions[quad_idx[0]], wp.vec2(-w_m[1], -w_m[0])) | |
| + wp.outer(cell_arg.positions[quad_idx[1]], wp.vec2(w_m[1], -w_p[0])) | |
| + wp.outer(cell_arg.positions[quad_idx[2]], wp.vec2(w_p[1], w_p[0])) | |
| + wp.outer(cell_arg.positions[quad_idx[3]], wp.vec2(-w_p[1], w_m[0])) | |
| ) | |
| def cell_inverse_deformation_gradient(cell_arg: CellArg, s: Sample): | |
| return wp.inverse(Quadmesh2D.cell_deformation_gradient(cell_arg, s)) | |
| def cell_measure(args: CellArg, s: Sample): | |
| return wp.abs(wp.determinant(Quadmesh2D.cell_deformation_gradient(args, s))) | |
| def cell_normal(args: CellArg, s: Sample): | |
| return wp.vec2(0.0) | |
| def side_index_arg_value(self, device) -> SideIndexArg: | |
| args = self.SideIndexArg() | |
| args.boundary_edge_indices = self._boundary_edge_indices.to(device) | |
| return args | |
| def boundary_side_index(args: SideIndexArg, boundary_side_index: int): | |
| """Boundary side to side index""" | |
| return args.boundary_edge_indices[boundary_side_index] | |
| def side_arg_value(self, device) -> CellArg: | |
| args = self.SideArg() | |
| args.cell_arg = self.cell_arg_value(device) | |
| args.edge_vertex_indices = self._edge_vertex_indices.to(device) | |
| args.edge_quad_indices = self._edge_quad_indices.to(device) | |
| return args | |
| def side_position(args: SideArg, s: Sample): | |
| edge_idx = args.edge_vertex_indices[s.element_index] | |
| return (1.0 - s.element_coords[0]) * args.cell_arg.positions[edge_idx[0]] + s.element_coords[ | |
| 0 | |
| ] * args.cell_arg.positions[edge_idx[1]] | |
| def side_deformation_gradient(args: SideArg, s: Sample): | |
| edge_idx = args.edge_vertex_indices[s.element_index] | |
| v0 = args.cell_arg.positions[edge_idx[0]] | |
| v1 = args.cell_arg.positions[edge_idx[1]] | |
| return v1 - v0 | |
| def side_inner_inverse_deformation_gradient(args: SideArg, s: Sample): | |
| cell_index = Quadmesh2D.side_inner_cell_index(args, s.element_index) | |
| cell_coords = Quadmesh2D.side_inner_cell_coords(args, s.element_index, s.element_coords) | |
| return Quadmesh2D.cell_inverse_deformation_gradient(args.cell_arg, make_free_sample(cell_index, cell_coords)) | |
| def side_outer_inverse_deformation_gradient(args: SideArg, s: Sample): | |
| cell_index = Quadmesh2D.side_outer_cell_index(args, s.element_index) | |
| cell_coords = Quadmesh2D.side_outer_cell_coords(args, s.element_index, s.element_coords) | |
| return Quadmesh2D.cell_inverse_deformation_gradient(args.cell_arg, make_free_sample(cell_index, cell_coords)) | |
| def side_measure(args: SideArg, s: Sample): | |
| edge_idx = args.edge_vertex_indices[s.element_index] | |
| v0 = args.cell_arg.positions[edge_idx[0]] | |
| v1 = args.cell_arg.positions[edge_idx[1]] | |
| return wp.length(v1 - v0) | |
| def side_measure_ratio(args: SideArg, s: Sample): | |
| inner = Quadmesh2D.side_inner_cell_index(args, s.element_index) | |
| outer = Quadmesh2D.side_outer_cell_index(args, s.element_index) | |
| inner_coords = Quadmesh2D.side_inner_cell_coords(args, s.element_index, s.element_coords) | |
| outer_coords = Quadmesh2D.side_outer_cell_coords(args, s.element_index, s.element_coords) | |
| return Quadmesh2D.side_measure(args, s) / wp.min( | |
| Quadmesh2D.cell_measure(args.cell_arg, make_free_sample(inner, inner_coords)), | |
| Quadmesh2D.cell_measure(args.cell_arg, make_free_sample(outer, outer_coords)), | |
| ) | |
| def side_normal(args: SideArg, s: Sample): | |
| edge_idx = args.edge_vertex_indices[s.element_index] | |
| v0 = args.cell_arg.positions[edge_idx[0]] | |
| v1 = args.cell_arg.positions[edge_idx[1]] | |
| e = v1 - v0 | |
| return wp.normalize(wp.vec2(-e[1], e[0])) | |
| def side_inner_cell_index(arg: SideArg, side_index: ElementIndex): | |
| return arg.edge_quad_indices[side_index][0] | |
| def side_outer_cell_index(arg: SideArg, side_index: ElementIndex): | |
| return arg.edge_quad_indices[side_index][1] | |
| def edge_to_quad_coords(args: SideArg, side_index: ElementIndex, quad_index: ElementIndex, side_coords: Coords): | |
| edge_vidx = args.edge_vertex_indices[side_index] | |
| quad_vidx = args.cell_arg.quad_vertex_indices[quad_index] | |
| vs = edge_vidx[0] | |
| ve = edge_vidx[1] | |
| s = side_coords[0] | |
| if vs == quad_vidx[0]: | |
| return wp.select(ve == quad_vidx[1], Coords(0.0, s, 0.0), Coords(s, 0.0, 0.0)) | |
| elif vs == quad_vidx[1]: | |
| return wp.select(ve == quad_vidx[2], Coords(1.0 - s, 0.0, 0.0), Coords(1.0, s, 0.0)) | |
| elif vs == quad_vidx[2]: | |
| return wp.select(ve == quad_vidx[3], Coords(1.0, 1.0 - s, 0.0), Coords(1.0 - s, 1.0, 0.0)) | |
| return wp.select(ve == quad_vidx[0], Coords(s, 1.0, 0.0), Coords(0.0, 1.0 - s, 0.0)) | |
| def side_inner_cell_coords(args: SideArg, side_index: ElementIndex, side_coords: Coords): | |
| inner_cell_index = Quadmesh2D.side_inner_cell_index(args, side_index) | |
| return Quadmesh2D.edge_to_quad_coords(args, side_index, inner_cell_index, side_coords) | |
| def side_outer_cell_coords(args: SideArg, side_index: ElementIndex, side_coords: Coords): | |
| outer_cell_index = Quadmesh2D.side_outer_cell_index(args, side_index) | |
| return Quadmesh2D.edge_to_quad_coords(args, side_index, outer_cell_index, side_coords) | |
| def side_from_cell_coords( | |
| args: SideArg, | |
| side_index: ElementIndex, | |
| quad_index: ElementIndex, | |
| quad_coords: Coords, | |
| ): | |
| edge_vidx = args.edge_vertex_indices[side_index] | |
| quad_vidx = args.cell_arg.quad_vertex_indices[quad_index] | |
| vs = edge_vidx[0] | |
| ve = edge_vidx[1] | |
| cx = quad_coords[0] | |
| cy = quad_coords[1] | |
| if vs == quad_vidx[0]: | |
| oc = wp.select(ve == quad_vidx[1], cx, cy) | |
| ec = wp.select(ve == quad_vidx[1], cy, cx) | |
| elif vs == quad_vidx[1]: | |
| oc = wp.select(ve == quad_vidx[2], cy, 1.0 - cx) | |
| ec = wp.select(ve == quad_vidx[2], 1.0 - cx, cy) | |
| elif vs == quad_vidx[2]: | |
| oc = wp.select(ve == quad_vidx[3], 1.0 - cx, 1.0 - cy) | |
| ec = wp.select(ve == quad_vidx[3], 1.0 - cy, 1.0 - cx) | |
| else: | |
| oc = wp.select(ve == quad_vidx[0], 1.0 - cy, cx) | |
| ec = wp.select(ve == quad_vidx[0], cx, 1.0 - cy) | |
| return wp.select(oc == 0.0, Coords(OUTSIDE), Coords(ec, 0.0, 0.0)) | |
| def side_to_cell_arg(side_arg: SideArg): | |
| return side_arg.cell_arg | |
| def _build_topology(self, temporary_store: TemporaryStore): | |
| from warp.fem.utils import compress_node_indices, masked_indices | |
| from warp.utils import array_scan | |
| device = self.quad_vertex_indices.device | |
| vertex_quad_offsets, vertex_quad_indices, _, __ = compress_node_indices( | |
| self.vertex_count(), self.quad_vertex_indices, temporary_store=temporary_store | |
| ) | |
| self._vertex_quad_offsets = vertex_quad_offsets.detach() | |
| self._vertex_quad_indices = vertex_quad_indices.detach() | |
| vertex_start_edge_count = borrow_temporary(temporary_store, dtype=int, device=device, shape=self.vertex_count()) | |
| vertex_start_edge_count.array.zero_() | |
| vertex_start_edge_offsets = borrow_temporary_like(vertex_start_edge_count, temporary_store=temporary_store) | |
| vertex_edge_ends = borrow_temporary(temporary_store, dtype=int, device=device, shape=(4 * self.cell_count())) | |
| vertex_edge_quads = borrow_temporary( | |
| temporary_store, dtype=int, device=device, shape=(4 * self.cell_count(), 2) | |
| ) | |
| # Count face edges starting at each vertex | |
| wp.launch( | |
| kernel=Quadmesh2D._count_starting_edges_kernel, | |
| device=device, | |
| dim=self.cell_count(), | |
| inputs=[self.quad_vertex_indices, vertex_start_edge_count.array], | |
| ) | |
| array_scan(in_array=vertex_start_edge_count.array, out_array=vertex_start_edge_offsets.array, inclusive=False) | |
| # Count number of unique edges (deduplicate across faces) | |
| vertex_unique_edge_count = vertex_start_edge_count | |
| wp.launch( | |
| kernel=Quadmesh2D._count_unique_starting_edges_kernel, | |
| device=device, | |
| dim=self.vertex_count(), | |
| inputs=[ | |
| self._vertex_quad_offsets, | |
| self._vertex_quad_indices, | |
| self.quad_vertex_indices, | |
| vertex_start_edge_offsets.array, | |
| vertex_unique_edge_count.array, | |
| vertex_edge_ends.array, | |
| vertex_edge_quads.array, | |
| ], | |
| ) | |
| vertex_unique_edge_offsets = borrow_temporary_like(vertex_start_edge_offsets, temporary_store=temporary_store) | |
| array_scan(in_array=vertex_start_edge_count.array, out_array=vertex_unique_edge_offsets.array, inclusive=False) | |
| # Get back edge count to host | |
| if device.is_cuda: | |
| edge_count = borrow_temporary(temporary_store, shape=(1,), dtype=int, device="cpu", pinned=True) | |
| # Last vertex will not own any edge, so its count will be zero; just fetching last prefix count is ok | |
| wp.copy( | |
| dest=edge_count.array, src=vertex_unique_edge_offsets.array, src_offset=self.vertex_count() - 1, count=1 | |
| ) | |
| wp.synchronize_stream(wp.get_stream(device)) | |
| edge_count = int(edge_count.array.numpy()[0]) | |
| else: | |
| edge_count = int(vertex_unique_edge_offsets.array.numpy()[self.vertex_count() - 1]) | |
| self._edge_vertex_indices = wp.empty(shape=(edge_count,), dtype=wp.vec2i, device=device) | |
| self._edge_quad_indices = wp.empty(shape=(edge_count,), dtype=wp.vec2i, device=device) | |
| boundary_mask = borrow_temporary(temporary_store=temporary_store, shape=(edge_count,), dtype=int, device=device) | |
| # Compress edge data | |
| wp.launch( | |
| kernel=Quadmesh2D._compress_edges_kernel, | |
| device=device, | |
| dim=self.vertex_count(), | |
| inputs=[ | |
| vertex_start_edge_offsets.array, | |
| vertex_unique_edge_offsets.array, | |
| vertex_unique_edge_count.array, | |
| vertex_edge_ends.array, | |
| vertex_edge_quads.array, | |
| self._edge_vertex_indices, | |
| self._edge_quad_indices, | |
| boundary_mask.array, | |
| ], | |
| ) | |
| vertex_start_edge_offsets.release() | |
| vertex_unique_edge_offsets.release() | |
| vertex_unique_edge_count.release() | |
| vertex_edge_ends.release() | |
| vertex_edge_quads.release() | |
| # Flip normals if necessary | |
| wp.launch( | |
| kernel=Quadmesh2D._flip_edge_normals, | |
| device=device, | |
| dim=self.side_count(), | |
| inputs=[self._edge_vertex_indices, self._edge_quad_indices, self.quad_vertex_indices, self.positions], | |
| ) | |
| boundary_edge_indices, _ = masked_indices(boundary_mask.array, temporary_store=temporary_store) | |
| self._boundary_edge_indices = boundary_edge_indices.detach() | |
| boundary_mask.release() | |
| def _count_starting_edges_kernel( | |
| quad_vertex_indices: wp.array2d(dtype=int), vertex_start_edge_count: wp.array(dtype=int) | |
| ): | |
| t = wp.tid() | |
| for k in range(4): | |
| v0 = quad_vertex_indices[t, k] | |
| v1 = quad_vertex_indices[t, (k + 1) % 4] | |
| if v0 < v1: | |
| wp.atomic_add(vertex_start_edge_count, v0, 1) | |
| else: | |
| wp.atomic_add(vertex_start_edge_count, v1, 1) | |
| def _find( | |
| needle: int, | |
| values: wp.array(dtype=int), | |
| beg: int, | |
| end: int, | |
| ): | |
| for i in range(beg, end): | |
| if values[i] == needle: | |
| return i | |
| return -1 | |
| def _count_unique_starting_edges_kernel( | |
| vertex_quad_offsets: wp.array(dtype=int), | |
| vertex_quad_indices: wp.array(dtype=int), | |
| quad_vertex_indices: wp.array2d(dtype=int), | |
| vertex_start_edge_offsets: wp.array(dtype=int), | |
| vertex_start_edge_count: wp.array(dtype=int), | |
| edge_ends: wp.array(dtype=int), | |
| edge_quads: wp.array2d(dtype=int), | |
| ): | |
| v = wp.tid() | |
| edge_beg = vertex_start_edge_offsets[v] | |
| quad_beg = vertex_quad_offsets[v] | |
| quad_end = vertex_quad_offsets[v + 1] | |
| edge_cur = edge_beg | |
| for quad in range(quad_beg, quad_end): | |
| q = vertex_quad_indices[quad] | |
| for k in range(4): | |
| v0 = quad_vertex_indices[q, k] | |
| v1 = quad_vertex_indices[q, (k + 1) % 4] | |
| if v == wp.min(v0, v1): | |
| other_v = wp.max(v0, v1) | |
| # Check if other_v has been seen | |
| seen_idx = Quadmesh2D._find(other_v, edge_ends, edge_beg, edge_cur) | |
| if seen_idx == -1: | |
| edge_ends[edge_cur] = other_v | |
| edge_quads[edge_cur, 0] = q | |
| edge_quads[edge_cur, 1] = q | |
| edge_cur += 1 | |
| else: | |
| edge_quads[seen_idx, 1] = q | |
| vertex_start_edge_count[v] = edge_cur - edge_beg | |
| def _compress_edges_kernel( | |
| vertex_start_edge_offsets: wp.array(dtype=int), | |
| vertex_unique_edge_offsets: wp.array(dtype=int), | |
| vertex_unique_edge_count: wp.array(dtype=int), | |
| uncompressed_edge_ends: wp.array(dtype=int), | |
| uncompressed_edge_quads: wp.array2d(dtype=int), | |
| edge_vertex_indices: wp.array(dtype=wp.vec2i), | |
| edge_quad_indices: wp.array(dtype=wp.vec2i), | |
| boundary_mask: wp.array(dtype=int), | |
| ): | |
| v = wp.tid() | |
| start_beg = vertex_start_edge_offsets[v] | |
| unique_beg = vertex_unique_edge_offsets[v] | |
| unique_count = vertex_unique_edge_count[v] | |
| for e in range(unique_count): | |
| src_index = start_beg + e | |
| edge_index = unique_beg + e | |
| edge_vertex_indices[edge_index] = wp.vec2i(v, uncompressed_edge_ends[src_index]) | |
| q0 = uncompressed_edge_quads[src_index, 0] | |
| q1 = uncompressed_edge_quads[src_index, 1] | |
| edge_quad_indices[edge_index] = wp.vec2i(q0, q1) | |
| if q0 == q1: | |
| boundary_mask[edge_index] = 1 | |
| else: | |
| boundary_mask[edge_index] = 0 | |
| def _flip_edge_normals( | |
| edge_vertex_indices: wp.array(dtype=wp.vec2i), | |
| edge_quad_indices: wp.array(dtype=wp.vec2i), | |
| quad_vertex_indices: wp.array2d(dtype=int), | |
| positions: wp.array(dtype=wp.vec2), | |
| ): | |
| e = wp.tid() | |
| tri = edge_quad_indices[e][0] | |
| quad_vidx = quad_vertex_indices[tri] | |
| edge_vidx = edge_vertex_indices[e] | |
| quad_centroid = ( | |
| positions[quad_vidx[0]] + positions[quad_vidx[1]] + positions[quad_vidx[2]] + positions[quad_vidx[3]] | |
| ) / 4.0 | |
| v0 = positions[edge_vidx[0]] | |
| v1 = positions[edge_vidx[1]] | |
| edge_center = 0.5 * (v1 + v0) | |
| edge_vec = v1 - v0 | |
| edge_normal = wp.vec2(-edge_vec[1], edge_vec[0]) | |
| # if edge normal points toward first triangle centroid, flip indices | |
| if wp.dot(quad_centroid - edge_center, edge_normal) > 0.0: | |
| edge_vertex_indices[e] = wp.vec2i(edge_vidx[1], edge_vidx[0]) | |