import math import warp as wp import numpy as np from warp.fem.types import Coords from warp.fem import cache def _tet_node_index(tx: int, ty: int, tz: int, degree: int): from .triangle_shape_function import _triangle_node_index VERTEX_NODE_COUNT = 4 EDGE_INTERIOR_NODE_COUNT = degree - 1 VERTEX_EDGE_NODE_COUNT = VERTEX_NODE_COUNT + 6 * EDGE_INTERIOR_NODE_COUNT FACE_INTERIOR_NODE_COUNT = (degree - 1) * (degree - 2) // 2 VERTEX_EDGE_FACE_NODE_COUNT = VERTEX_EDGE_NODE_COUNT + 4 * FACE_INTERIOR_NODE_COUNT # Index in similar order to e.g. VTK # First vertices, then edges (counterclokwise), then faces, then interior points (recursively) if tx == 0: if ty == 0: if tz == 0: return 0 elif tz == degree: return 3 else: # 0-3 edge edge_index = 3 return VERTEX_NODE_COUNT + EDGE_INTERIOR_NODE_COUNT * edge_index + (tz - 1) elif tz == 0: if ty == degree: return 2 else: # 2-0 edge edge_index = 2 return VERTEX_NODE_COUNT + EDGE_INTERIOR_NODE_COUNT * edge_index + (EDGE_INTERIOR_NODE_COUNT - ty) elif tz + ty == degree: # 2-3 edge edge_index = 5 return VERTEX_NODE_COUNT + EDGE_INTERIOR_NODE_COUNT * edge_index + (tz - 1) else: # 2-3-0 face face_index = 2 return ( VERTEX_EDGE_NODE_COUNT + FACE_INTERIOR_NODE_COUNT * face_index + _triangle_node_index(degree - 1 - ty - tz, tz - 1, degree - 3) ) elif ty == 0: if tz == 0: if tx == degree: return 1 else: # 0-1 edge edge_index = 0 return VERTEX_NODE_COUNT + EDGE_INTERIOR_NODE_COUNT * edge_index + (tx - 1) elif tz + tx == degree: # 1-3 edge edge_index = 4 return VERTEX_NODE_COUNT + EDGE_INTERIOR_NODE_COUNT * edge_index + (tz - 1) else: # 3-0-1 face face_index = 3 return ( VERTEX_EDGE_NODE_COUNT + FACE_INTERIOR_NODE_COUNT * face_index + _triangle_node_index(tx - 1, tz - 1, degree - 3) ) elif tz == 0: if tx + ty == degree: # 1-2 edge edge_index = 1 return VERTEX_NODE_COUNT + EDGE_INTERIOR_NODE_COUNT * edge_index + (ty - 1) else: # 0-1-2 face face_index = 0 return ( VERTEX_EDGE_NODE_COUNT + FACE_INTERIOR_NODE_COUNT * face_index + _triangle_node_index(tx - 1, ty - 1, degree - 3) ) elif tx + ty + tz == degree: # 1-2-3 face face_index = 1 return ( VERTEX_EDGE_NODE_COUNT + FACE_INTERIOR_NODE_COUNT * face_index + _triangle_node_index(tx - 1, tz - 1, degree - 3) ) return VERTEX_EDGE_FACE_NODE_COUNT + _tet_node_index(tx - 1, ty - 1, tz - 1, degree - 4) class TetrahedronPolynomialShapeFunctions: INVALID = wp.constant(-1) VERTEX = wp.constant(0) EDGE = wp.constant(1) FACE = wp.constant(2) INTERIOR = wp.constant(3) def __init__(self, degree: int): self.ORDER = wp.constant(degree) self.NODES_PER_ELEMENT = wp.constant((degree + 1) * (degree + 2) * (degree + 3) // 6) self.NODES_PER_SIDE = wp.constant((degree + 1) * (degree + 2) // 2) tet_coords = np.empty((self.NODES_PER_ELEMENT, 3), dtype=int) for tx in range(degree + 1): for ty in range(degree + 1 - tx): for tz in range(degree + 1 - tx - ty): index = _tet_node_index(tx, ty, tz, degree) tet_coords[index] = [tx, ty, tz] CoordTypeVec = wp.mat(dtype=int, shape=(self.NODES_PER_ELEMENT, 3)) self.NODE_TET_COORDS = wp.constant(CoordTypeVec(tet_coords)) self.node_type_and_type_index = self._get_node_type_and_type_index() self._node_tet_coordinates = self._get_node_tet_coordinates() @property def name(self) -> str: return f"Tet_P{self.ORDER}" def _get_node_tet_coordinates(self): NODE_TET_COORDS = self.NODE_TET_COORDS def node_tet_coordinates( node_index_in_elt: int, ): return wp.vec3i( NODE_TET_COORDS[node_index_in_elt, 0], NODE_TET_COORDS[node_index_in_elt, 1], NODE_TET_COORDS[node_index_in_elt, 2], ) return cache.get_func(node_tet_coordinates, self.name) def _get_node_type_and_type_index(self): ORDER = self.ORDER NODES_PER_ELEMENT = self.NODES_PER_ELEMENT def node_type_and_index( node_index_in_elt: int, ): if node_index_in_elt < 0 or node_index_in_elt >= NODES_PER_ELEMENT: return TetrahedronPolynomialShapeFunctions.INVALID, TetrahedronPolynomialShapeFunctions.INVALID if node_index_in_elt < 4: return TetrahedronPolynomialShapeFunctions.VERTEX, node_index_in_elt if node_index_in_elt < (6 * ORDER - 2): return TetrahedronPolynomialShapeFunctions.EDGE, (node_index_in_elt - 4) if node_index_in_elt < (2 * ORDER * ORDER + 2): return TetrahedronPolynomialShapeFunctions.FACE, (node_index_in_elt - (6 * ORDER - 2)) return TetrahedronPolynomialShapeFunctions.INTERIOR, (node_index_in_elt - (2 * ORDER * ORDER + 2)) return cache.get_func(node_type_and_index, self.name) def make_node_coords_in_element(self): ORDER = self.ORDER def node_coords_in_element( node_index_in_elt: int, ): tet_coords = self._node_tet_coordinates(node_index_in_elt) cx = float(tet_coords[0]) / float(ORDER) cy = float(tet_coords[1]) / float(ORDER) cz = float(tet_coords[2]) / float(ORDER) return Coords(cx, cy, cz) return cache.get_func(node_coords_in_element, self.name) def make_node_quadrature_weight(self): if self.ORDER == 3: # Order 1, but optimized quadrature weights for monomials of order <= 6 vertex_weight = 0.007348845656 edge_weight = 0.020688129855 face_weight = 0.180586764778 interior_weight = 0.0 else: vertex_weight = 1.0 / self.NODES_PER_ELEMENT edge_weight = 1.0 / self.NODES_PER_ELEMENT face_weight = 1.0 / self.NODES_PER_ELEMENT interior_weight = 1.0 / self.NODES_PER_ELEMENT VERTEX_WEIGHT = wp.constant(vertex_weight) EDGE_WEIGHT = wp.constant(edge_weight) FACE_WEIGHT = wp.constant(face_weight) INTERIOR_WEIGHT = wp.constant(interior_weight) @cache.dynamic_func(suffix=self.name) def node_quadrature_weight(node_index_in_element: int): node_type, type_index = self.node_type_and_type_index(node_index_in_element) if node_type == TetrahedronPolynomialShapeFunctions.VERTEX: return VERTEX_WEIGHT elif node_type == TetrahedronPolynomialShapeFunctions.EDGE: return EDGE_WEIGHT elif node_type == TetrahedronPolynomialShapeFunctions.FACE: return FACE_WEIGHT return INTERIOR_WEIGHT return node_quadrature_weight def make_trace_node_quadrature_weight(self): if self.ORDER == 3: # P3 intrisic quadrature vertex_weight = 1.0 / 30 edge_weight = 0.075 interior_weight = 0.45 elif self.ORDER == 2: # Order 1, but optimized quadrature weights for monomials of order <= 4 vertex_weight = 0.022335964126 edge_weight = 0.310997369207 interior_weight = 0.0 else: vertex_weight = 1.0 / self.NODES_PER_SIDE edge_weight = 1.0 / self.NODES_PER_SIDE interior_weight = 1.0 / self.NODES_PER_SIDE VERTEX_WEIGHT = wp.constant(vertex_weight) EDGE_WEIGHT = wp.constant(edge_weight) FACE_INTERIOR_WEIGHT = wp.constant(interior_weight) @cache.dynamic_func(suffix=self.name) def trace_node_quadrature_weight(node_index_in_element: int): node_type, type_index = self.node_type_and_type_index(node_index_in_element) if node_type == TetrahedronPolynomialShapeFunctions.VERTEX: return VERTEX_WEIGHT elif node_type == TetrahedronPolynomialShapeFunctions.EDGE: return EDGE_WEIGHT return FACE_INTERIOR_WEIGHT return trace_node_quadrature_weight def make_element_inner_weight(self): ORDER = self.ORDER def element_inner_weight_linear( coords: Coords, node_index_in_elt: int, ): if node_index_in_elt < 0 or node_index_in_elt >= 4: return 0.0 tet_coords = wp.vec4(1.0 - coords[0] - coords[1] - coords[2], coords[0], coords[1], coords[2]) return tet_coords[node_index_in_elt] def element_inner_weight_quadratic( coords: Coords, node_index_in_elt: int, ): node_type, type_index = self.node_type_and_type_index(node_index_in_elt) tet_coords = wp.vec4(1.0 - coords[0] - coords[1] - coords[2], coords[0], coords[1], coords[2]) if node_type == TetrahedronPolynomialShapeFunctions.VERTEX: # Vertex return tet_coords[type_index] * (2.0 * tet_coords[type_index] - 1.0) elif node_type == TetrahedronPolynomialShapeFunctions.EDGE: # Edge if type_index < 3: c1 = type_index c2 = (type_index + 1) % 3 else: c1 = type_index - 3 c2 = 3 return 4.0 * tet_coords[c1] * tet_coords[c2] return 0.0 def element_inner_weight_cubic( coords: Coords, node_index_in_elt: int, ): node_type, type_index = self.node_type_and_type_index(node_index_in_elt) tet_coords = wp.vec4(1.0 - coords[0] - coords[1] - coords[2], coords[0], coords[1], coords[2]) if node_type == TetrahedronPolynomialShapeFunctions.VERTEX: # Vertex return ( 0.5 * tet_coords[type_index] * (3.0 * tet_coords[type_index] - 1.0) * (3.0 * tet_coords[type_index] - 2.0) ) elif node_type == TetrahedronPolynomialShapeFunctions.EDGE: # Edge edge = type_index // 2 edge_node = type_index - 2 * edge if edge < 3: c1 = (edge + edge_node) % 3 c2 = (edge + 1 - edge_node) % 3 elif edge_node == 0: c1 = edge - 3 c2 = 3 else: c1 = 3 c2 = edge - 3 return 4.5 * tet_coords[c1] * tet_coords[c2] * (3.0 * tet_coords[c1] - 1.0) elif node_type == TetrahedronPolynomialShapeFunctions.FACE: # Interior c1 = type_index c2 = (c1 + 1) % 4 c3 = (c1 + 2) % 4 return 27.0 * tet_coords[c1] * tet_coords[c2] * tet_coords[c3] return 0.0 if ORDER == 1: return cache.get_func(element_inner_weight_linear, self.name) elif ORDER == 2: return cache.get_func(element_inner_weight_quadratic, self.name) elif ORDER == 3: return cache.get_func(element_inner_weight_cubic, self.name) return None def make_element_inner_weight_gradient(self): ORDER = self.ORDER def element_inner_weight_gradient_linear( coords: Coords, node_index_in_elt: int, ): if node_index_in_elt < 0 or node_index_in_elt >= 4: return wp.vec3(0.0) dw_dc = wp.vec4(0.0) dw_dc[node_index_in_elt] = 1.0 dw_du = wp.vec3(dw_dc[1] - dw_dc[0], dw_dc[2] - dw_dc[0], dw_dc[3] - dw_dc[0]) return dw_du def element_inner_weight_gradient_quadratic( coords: Coords, node_index_in_elt: int, ): node_type, type_index = self.node_type_and_type_index(node_index_in_elt) tet_coords = wp.vec4(1.0 - coords[0] - coords[1] - coords[2], coords[0], coords[1], coords[2]) dw_dc = wp.vec4(0.0) if node_type == TetrahedronPolynomialShapeFunctions.VERTEX: # Vertex dw_dc[type_index] = 4.0 * tet_coords[type_index] - 1.0 elif node_type == TetrahedronPolynomialShapeFunctions.EDGE: # Edge if type_index < 3: c1 = type_index c2 = (type_index + 1) % 3 else: c1 = type_index - 3 c2 = 3 dw_dc[c1] = 4.0 * tet_coords[c2] dw_dc[c2] = 4.0 * tet_coords[c1] dw_du = wp.vec3(dw_dc[1] - dw_dc[0], dw_dc[2] - dw_dc[0], dw_dc[3] - dw_dc[0]) return dw_du def element_inner_weight_gradient_cubic( coords: Coords, node_index_in_elt: int, ): node_type, type_index = self.node_type_and_type_index(node_index_in_elt) tet_coords = wp.vec4(1.0 - coords[0] - coords[1] - coords[2], coords[0], coords[1], coords[2]) dw_dc = wp.vec4(0.0) if node_type == TetrahedronPolynomialShapeFunctions.VERTEX: # Vertex dw_dc[type_index] = ( 0.5 * 27.0 * tet_coords[type_index] * tet_coords[type_index] - 9.0 * tet_coords[type_index] + 1.0 ) elif node_type == TetrahedronPolynomialShapeFunctions.EDGE: # Edge edge = type_index // 2 edge_node = type_index - 2 * edge if edge < 3: c1 = (edge + edge_node) % 3 c2 = (edge + 1 - edge_node) % 3 elif edge_node == 0: c1 = edge - 3 c2 = 3 else: c1 = 3 c2 = edge - 3 dw_dc[c1] = 4.5 * tet_coords[c2] * (6.0 * tet_coords[c1] - 1.0) dw_dc[c2] = 4.5 * tet_coords[c1] * (3.0 * tet_coords[c1] - 1.0) elif node_type == TetrahedronPolynomialShapeFunctions.FACE: # Interior c1 = type_index c2 = (c1 + 1) % 4 c3 = (c1 + 2) % 4 dw_dc[c1] = 27.0 * tet_coords[c2] * tet_coords[c3] dw_dc[c2] = 27.0 * tet_coords[c3] * tet_coords[c1] dw_dc[c3] = 27.0 * tet_coords[c1] * tet_coords[c2] dw_du = wp.vec3(dw_dc[1] - dw_dc[0], dw_dc[2] - dw_dc[0], dw_dc[3] - dw_dc[0]) return dw_du if ORDER == 1: return cache.get_func(element_inner_weight_gradient_linear, self.name) elif ORDER == 2: return cache.get_func(element_inner_weight_gradient_quadratic, self.name) elif ORDER == 3: return cache.get_func(element_inner_weight_gradient_cubic, self.name) return None def element_node_tets(self): if self.ORDER == 1: element_tets = [[0, 1, 2, 3]] if self.ORDER == 2: element_tets = [ [0, 4, 6, 7], [1, 5, 4, 8], [2, 6, 5, 9], [3, 7, 8, 9], [4, 5, 6, 8], [8, 7, 9, 6], [6, 5, 9, 8], [6, 8, 7, 4], ] elif self.ORDER == 3: raise NotImplementedError() return np.array(element_tets) class TetrahedronNonConformingPolynomialShapeFunctions: def __init__(self, degree: int): self._tet_shape = TetrahedronPolynomialShapeFunctions(degree=degree) self.ORDER = self._tet_shape.ORDER self.NODES_PER_ELEMENT = self._tet_shape.NODES_PER_ELEMENT self.element_node_tets = self._tet_shape.element_node_tets if self.ORDER == 1: self._TET_SCALE = 0.4472135955 # so v at 0.5854101966249680 (order 2) elif self.ORDER == 2: self._TET_SCALE = 0.6123779296874996 # optimized for low intrinsic quadrature error of deg 4 elif self.ORDER == 3: self._TET_SCALE = 0.7153564453124999 # optimized for low intrinsic quadrature error of deg 6 else: self._TET_SCALE = 1.0 self._TET_SCALE = wp.constant(self._TET_SCALE) self._TET_OFFSET = wp.constant((1.0 - self._TET_SCALE) * wp.vec3(0.25, 0.25, 0.25)) @property def name(self) -> str: return f"Tet_P{self.ORDER}d" def make_node_coords_in_element(self): node_coords_in_tet = self._tet_shape.make_node_coords_in_element() TET_SCALE = self._TET_SCALE TET_OFFSET = self._TET_OFFSET @cache.dynamic_func(suffix=self.name) def node_coords_in_element( node_index_in_elt: int, ): tet_coords = node_coords_in_tet(node_index_in_elt) return TET_SCALE * tet_coords + TET_OFFSET return node_coords_in_element def make_node_quadrature_weight(self): # Intrinsic quadrature -- precomputed integral of node shape functions # over element. Order euqla to self.ORDER if self.ORDER == 2: vertex_weight = 0.07499641 edge_weight = 0.11666908 face_interior_weight = 0.0 elif self.ORDER == 3: vertex_weight = 0.03345134 edge_weight = 0.04521887 face_interior_weight = 0.08089206 else: vertex_weight = 1.0 / self.NODES_PER_ELEMENT edge_weight = 1.0 / self.NODES_PER_ELEMENT face_interior_weight = 1.0 / self.NODES_PER_ELEMENT VERTEX_WEIGHT = wp.constant(vertex_weight) EDGE_WEIGHT = wp.constant(edge_weight) FACE_INTERIOR_WEIGHT = wp.constant(face_interior_weight) @cache.dynamic_func(suffix=self.name) def node_quadrature_weight(node_index_in_element: int): node_type, type_index = self._tet_shape.node_type_and_type_index(node_index_in_element) if node_type == TetrahedronPolynomialShapeFunctions.VERTEX: return VERTEX_WEIGHT elif node_type == TetrahedronPolynomialShapeFunctions.EDGE: return EDGE_WEIGHT return FACE_INTERIOR_WEIGHT return node_quadrature_weight def make_trace_node_quadrature_weight(self): # Non-conforming, zero measure on sides @wp.func def zero(node_index_in_elt: int): return 0.0 return zero def make_element_inner_weight(self): tet_inner_weight = self._tet_shape.make_element_inner_weight() TET_SCALE = self._TET_SCALE TET_OFFSET = self._TET_OFFSET @cache.dynamic_func(suffix=self.name) def element_inner_weight( coords: Coords, node_index_in_elt: int, ): tet_coords = (coords - TET_OFFSET) / TET_SCALE return tet_inner_weight(tet_coords, node_index_in_elt) return element_inner_weight def make_element_inner_weight_gradient(self): tet_inner_weight_gradient = self._tet_shape.make_element_inner_weight_gradient() TET_SCALE = self._TET_SCALE TET_OFFSET = self._TET_OFFSET @cache.dynamic_func(suffix=self.name) def element_inner_weight_gradient( coords: Coords, node_index_in_elt: int, ): tet_coords = (coords - TET_OFFSET) / TET_SCALE grad = tet_inner_weight_gradient(tet_coords, node_index_in_elt) return grad / TET_SCALE return element_inner_weight_gradient