import math import warp as wp import numpy as np from warp.fem.geometry import Grid3D from warp.fem.polynomial import Polynomial, quadrature_1d, lagrange_scales, is_closed from warp.fem.types import Coords from warp.fem import cache from .tet_shape_function import TetrahedronPolynomialShapeFunctions _CUBE_EDGE_INDICES = wp.constant( wp.mat(shape=(3, 4), dtype=int)( [ [0, 4, 2, 6], [3, 1, 7, 5], [8, 11, 9, 10], ] ) ) class CubeTripolynomialShapeFunctions: VERTEX = 0 EDGE = 1 FACE = 2 INTERIOR = 3 def __init__(self, degree: int, family: Polynomial): self.family = family self.ORDER = wp.constant(degree) self.NODES_PER_ELEMENT = wp.constant((degree + 1) ** 3) self.NODES_PER_EDGE = wp.constant(degree + 1) lobatto_coords, lobatto_weight = quadrature_1d(point_count=degree + 1, family=family) lagrange_scale = lagrange_scales(lobatto_coords) NodeVec = wp.types.vector(length=degree + 1, dtype=wp.float32) self.LOBATTO_COORDS = wp.constant(NodeVec(lobatto_coords)) self.LOBATTO_WEIGHT = wp.constant(NodeVec(lobatto_weight)) self.LAGRANGE_SCALE = wp.constant(NodeVec(lagrange_scale)) self.ORDER_PLUS_ONE = wp.constant(self.ORDER + 1) self._node_ijk = self._make_node_ijk() self.node_type_and_type_index = self._make_node_type_and_type_index() @property def name(self) -> str: return f"Cube_Q{self.ORDER}_{self.family}" @wp.func def _vertex_coords_f(vidx_in_cell: int): x = vidx_in_cell // 4 y = (vidx_in_cell - 4 * x) // 2 z = vidx_in_cell - 4 * x - 2 * y return wp.vec3(float(x), float(y), float(z)) def _make_node_ijk(self): ORDER_PLUS_ONE = self.ORDER_PLUS_ONE def node_ijk( node_index_in_elt: int, ): node_i = node_index_in_elt // (ORDER_PLUS_ONE * ORDER_PLUS_ONE) node_jk = node_index_in_elt - ORDER_PLUS_ONE * ORDER_PLUS_ONE * node_i node_j = node_jk // ORDER_PLUS_ONE node_k = node_jk - ORDER_PLUS_ONE * node_j return node_i, node_j, node_k return cache.get_func(node_ijk, self.name) def _make_node_type_and_type_index(self): ORDER = self.ORDER @cache.dynamic_func(suffix=self.name) def node_type_and_type_index( node_index_in_elt: int, ): i, j, k = self._node_ijk(node_index_in_elt) zi = wp.select(i == 0, 0, 1) zj = wp.select(j == 0, 0, 1) zk = wp.select(k == 0, 0, 1) mi = wp.select(i == ORDER, 0, 1) mj = wp.select(j == ORDER, 0, 1) mk = wp.select(k == ORDER, 0, 1) if zi + mi == 1: if zj + mj == 1: if zk + mk == 1: # vertex type_instance = mi * 4 + mj * 2 + mk return CubeTripolynomialShapeFunctions.VERTEX, type_instance, 0 # z edge type_instance = _CUBE_EDGE_INDICES[2, mi * 2 + mj] type_index = k - 1 return CubeTripolynomialShapeFunctions.EDGE, type_instance, type_index if zk + mk == 1: # y edge type_instance = _CUBE_EDGE_INDICES[1, mk * 2 + mi] type_index = j - 1 return CubeTripolynomialShapeFunctions.EDGE, type_instance, type_index # x face type_instance = mi type_index = wp.select(mi == 1, (j - 1) * (ORDER - 1) + k - 1, (k - 1) * (ORDER - 1) + j - 1) return CubeTripolynomialShapeFunctions.FACE, type_instance, type_index if zj + mj == 1: if zk + mk == 1: # x edge type_instance = _CUBE_EDGE_INDICES[0, mj * 2 + mk] type_index = i - 1 return CubeTripolynomialShapeFunctions.EDGE, type_instance, type_index # y face type_instance = 2 + mj type_index = wp.select(mj == 1, (i - 1) * (ORDER - 1) + k - 1, (k - 1) * (ORDER - 1) + i - 1) return CubeTripolynomialShapeFunctions.FACE, type_instance, type_index if zk + mk == 1: # z face type_instance = 4 + mk type_index = wp.select(mk == 1, (j - 1) * (ORDER - 1) + i - 1, (i - 1) * (ORDER - 1) + j - 1) return CubeTripolynomialShapeFunctions.FACE, type_instance, type_index type_index = ((i - 1) * (ORDER - 1) + (j - 1)) * (ORDER - 1) + k - 1 return CubeTripolynomialShapeFunctions.INTERIOR, 0, type_index return node_type_and_type_index def make_node_coords_in_element(self): LOBATTO_COORDS = self.LOBATTO_COORDS @cache.dynamic_func(suffix=self.name) def node_coords_in_element( node_index_in_elt: int, ): node_i, node_j, node_k = self._node_ijk(node_index_in_elt) return Coords(LOBATTO_COORDS[node_i], LOBATTO_COORDS[node_j], LOBATTO_COORDS[node_k]) return node_coords_in_element def make_node_quadrature_weight(self): ORDER = self.ORDER LOBATTO_WEIGHT = self.LOBATTO_WEIGHT def node_quadrature_weight( node_index_in_elt: int, ): node_i, node_j, node_k = self._node_ijk(node_index_in_elt) return LOBATTO_WEIGHT[node_i] * LOBATTO_WEIGHT[node_j] * LOBATTO_WEIGHT[node_k] def node_quadrature_weight_linear( node_index_in_elt: int, ): return 0.125 if ORDER == 1: return cache.get_func(node_quadrature_weight_linear, self.name) return cache.get_func(node_quadrature_weight, self.name) def make_trace_node_quadrature_weight(self): ORDER = self.ORDER LOBATTO_WEIGHT = self.LOBATTO_WEIGHT def trace_node_quadrature_weight( node_index_in_elt: int, ): # We're either on a side interior or at a vertex # If we find one index at extremum, pick the two other node_i, node_j, node_k = self._node_ijk(node_index_in_elt) if node_i == 0 or node_i == ORDER: return LOBATTO_WEIGHT[node_j] * LOBATTO_WEIGHT[node_k] if node_j == 0 or node_j == ORDER: return LOBATTO_WEIGHT[node_i] * LOBATTO_WEIGHT[node_k] return LOBATTO_WEIGHT[node_i] * LOBATTO_WEIGHT[node_j] def trace_node_quadrature_weight_linear( node_index_in_elt: int, ): return 0.25 def trace_node_quadrature_weight_open( node_index_in_elt: int, ): return 0.0 if not is_closed(self.family): return cache.get_func(trace_node_quadrature_weight_open, self.name) if ORDER == 1: return cache.get_func(trace_node_quadrature_weight_linear, self.name) return cache.get_func(trace_node_quadrature_weight, self.name) def make_element_inner_weight(self): ORDER_PLUS_ONE = self.ORDER_PLUS_ONE LOBATTO_COORDS = self.LOBATTO_COORDS LAGRANGE_SCALE = self.LAGRANGE_SCALE def element_inner_weight( coords: Coords, node_index_in_elt: int, ): node_i, node_j, node_k = self._node_ijk(node_index_in_elt) w = float(1.0) for k in range(ORDER_PLUS_ONE): if k != node_i: w *= coords[0] - LOBATTO_COORDS[k] if k != node_j: w *= coords[1] - LOBATTO_COORDS[k] if k != node_k: w *= coords[2] - LOBATTO_COORDS[k] w *= LAGRANGE_SCALE[node_i] * LAGRANGE_SCALE[node_j] * LAGRANGE_SCALE[node_k] return w def element_inner_weight_linear( coords: Coords, node_index_in_elt: int, ): v = CubeTripolynomialShapeFunctions._vertex_coords_f(node_index_in_elt) wx = (1.0 - coords[0]) * (1.0 - v[0]) + v[0] * coords[0] wy = (1.0 - coords[1]) * (1.0 - v[1]) + v[1] * coords[1] wz = (1.0 - coords[2]) * (1.0 - v[2]) + v[2] * coords[2] return wx * wy * wz if self.ORDER == 1 and is_closed(self.family): return cache.get_func(element_inner_weight_linear, self.name) return cache.get_func(element_inner_weight, self.name) def make_element_inner_weight_gradient(self): ORDER_PLUS_ONE = self.ORDER_PLUS_ONE LOBATTO_COORDS = self.LOBATTO_COORDS LAGRANGE_SCALE = self.LAGRANGE_SCALE def element_inner_weight_gradient( coords: Coords, node_index_in_elt: int, ): node_i, node_j, node_k = self._node_ijk(node_index_in_elt) prefix_xy = float(1.0) prefix_yz = float(1.0) prefix_zx = float(1.0) for k in range(ORDER_PLUS_ONE): if k != node_i: prefix_yz *= coords[0] - LOBATTO_COORDS[k] if k != node_j: prefix_zx *= coords[1] - LOBATTO_COORDS[k] if k != node_k: prefix_xy *= coords[2] - LOBATTO_COORDS[k] prefix_x = prefix_zx * prefix_xy prefix_y = prefix_yz * prefix_xy prefix_z = prefix_zx * prefix_yz grad_x = float(0.0) grad_y = float(0.0) grad_z = float(0.0) for k in range(ORDER_PLUS_ONE): if k != node_i: delta_x = coords[0] - LOBATTO_COORDS[k] grad_x = grad_x * delta_x + prefix_x prefix_x *= delta_x if k != node_j: delta_y = coords[1] - LOBATTO_COORDS[k] grad_y = grad_y * delta_y + prefix_y prefix_y *= delta_y if k != node_k: delta_z = coords[2] - LOBATTO_COORDS[k] grad_z = grad_z * delta_z + prefix_z prefix_z *= delta_z grad = ( LAGRANGE_SCALE[node_i] * LAGRANGE_SCALE[node_j] * LAGRANGE_SCALE[node_k] * wp.vec3( grad_x, grad_y, grad_z, ) ) return grad def element_inner_weight_gradient_linear( coords: Coords, node_index_in_elt: int, ): v = CubeTripolynomialShapeFunctions._vertex_coords_f(node_index_in_elt) wx = (1.0 - coords[0]) * (1.0 - v[0]) + v[0] * coords[0] wy = (1.0 - coords[1]) * (1.0 - v[1]) + v[1] * coords[1] wz = (1.0 - coords[2]) * (1.0 - v[2]) + v[2] * coords[2] dx = 2.0 * v[0] - 1.0 dy = 2.0 * v[1] - 1.0 dz = 2.0 * v[2] - 1.0 return wp.vec3(dx * wy * wz, dy * wz * wx, dz * wx * wy) if self.ORDER == 1 and is_closed(self.family): return cache.get_func(element_inner_weight_gradient_linear, self.name) return cache.get_func(element_inner_weight_gradient, self.name) def element_node_hexes(self): from warp.fem.utils import grid_to_hexes return grid_to_hexes(self.ORDER, self.ORDER, self.ORDER) def element_node_tets(self): from warp.fem.utils import grid_to_tets return grid_to_tets(self.ORDER, self.ORDER, self.ORDER) class CubeSerendipityShapeFunctions: """ Serendipity element ~ tensor product space without interior nodes Edge shape functions are usual Lagrange shape functions times a bilinear function in the normal directions Corner shape functions are trilinear shape functions times a function of (x^{d-1} + y^{d-1}) """ # Node categories VERTEX = wp.constant(0) EDGE_X = wp.constant(1) EDGE_Y = wp.constant(2) def __init__(self, degree: int, family: Polynomial): if not is_closed(family): raise ValueError("A closed polynomial family is required to define serendipity elements") if degree not in [2, 3]: raise NotImplementedError("Serendipity element only implemented for order 2 or 3") self.family = family self.ORDER = wp.constant(degree) self.NODES_PER_ELEMENT = wp.constant(8 + 12 * (degree - 1)) self.NODES_PER_EDGE = wp.constant(degree + 1) lobatto_coords, lobatto_weight = quadrature_1d(point_count=degree + 1, family=family) lagrange_scale = lagrange_scales(lobatto_coords) NodeVec = wp.types.vector(length=degree + 1, dtype=wp.float32) self.LOBATTO_COORDS = wp.constant(NodeVec(lobatto_coords)) self.LOBATTO_WEIGHT = wp.constant(NodeVec(lobatto_weight)) self.LAGRANGE_SCALE = wp.constant(NodeVec(lagrange_scale)) self.ORDER_PLUS_ONE = wp.constant(self.ORDER + 1) self.node_type_and_type_index = self._get_node_type_and_type_index() self._node_lobatto_indices = self._get_node_lobatto_indices() @property def name(self) -> str: return f"Cube_S{self.ORDER}_{self.family}" def _get_node_type_and_type_index(self): @cache.dynamic_func(suffix=self.name) def node_type_and_index( node_index_in_elt: int, ): if node_index_in_elt < 8: return CubeSerendipityShapeFunctions.VERTEX, node_index_in_elt type_index = (node_index_in_elt - 8) // 3 side = node_index_in_elt - 8 - 3 * type_index return CubeSerendipityShapeFunctions.EDGE_X + side, type_index return node_type_and_index @wp.func def _vertex_coords(vidx_in_cell: int): x = vidx_in_cell // 4 y = (vidx_in_cell - 4 * x) // 2 z = vidx_in_cell - 4 * x - 2 * y return wp.vec3i(x, y, z) @wp.func def _edge_coords(type_index: int): index_in_side = type_index // 4 side_offset = type_index - 4 * index_in_side return (wp.vec3i(index_in_side + 1, side_offset // 2, side_offset % 2),) @wp.func def _edge_axis(node_type: int): return node_type - CubeSerendipityShapeFunctions.EDGE_X @wp.func def _cube_edge_index(node_type: int, type_index: int): index_in_side = type_index // 4 side_offset = type_index - 4 * index_in_side return _CUBE_EDGE_INDICES[node_type - CubeSerendipityShapeFunctions.EDGE_X, side_offset], index_in_side def _get_node_lobatto_indices(self): ORDER = self.ORDER @cache.dynamic_func(suffix=self.name) def node_lobatto_indices(node_type: int, type_index: int): if node_type == CubeSerendipityShapeFunctions.VERTEX: return CubeSerendipityShapeFunctions._vertex_coords(type_index) * ORDER axis = CubeSerendipityShapeFunctions._edge_axis(node_type) local_coords = CubeSerendipityShapeFunctions._edge_coords(type_index) local_indices = wp.vec3i(local_coords[0], local_coords[1] * ORDER, local_coords[2] * ORDER) return Grid3D._local_to_world(axis, local_indices) return node_lobatto_indices def make_node_coords_in_element(self): LOBATTO_COORDS = self.LOBATTO_COORDS @cache.dynamic_func(suffix=self.name) def node_coords_in_element( node_index_in_elt: int, ): node_type, type_index = self.node_type_and_type_index(node_index_in_elt) node_coords = self._node_lobatto_indices(node_type, type_index) return Coords( LOBATTO_COORDS[node_coords[0]], LOBATTO_COORDS[node_coords[1]], LOBATTO_COORDS[node_coords[2]] ) return node_coords_in_element def make_node_quadrature_weight(self): ORDER = self.ORDER @cache.dynamic_func(suffix=self.name) def node_quadrature_weight( node_index_in_elt: int, ): node_type, type_index = self.node_type_and_type_index(node_index_in_elt) if node_type == CubeSerendipityShapeFunctions.VERTEX: return 1.0 / float(8 * ORDER * ORDER * ORDER) return (1.0 - 1.0 / float(ORDER * ORDER * ORDER)) / float(12 * (ORDER - 1)) return node_quadrature_weight def make_trace_node_quadrature_weight(self): ORDER = self.ORDER @cache.dynamic_func(suffix=self.name) def trace_node_quadrature_weight( node_index_in_elt: int, ): node_type, type_index = self.node_type_and_type_index(node_index_in_elt) if node_type == CubeSerendipityShapeFunctions.VERTEX: return 0.25 / float(ORDER * ORDER) return (0.25 - 0.25 / float(ORDER * ORDER)) / float(ORDER - 1) return trace_node_quadrature_weight def make_element_inner_weight(self): ORDER = self.ORDER ORDER_PLUS_ONE = self.ORDER_PLUS_ONE LOBATTO_COORDS = self.LOBATTO_COORDS LAGRANGE_SCALE = self.LAGRANGE_SCALE DEGREE_3_SPHERE_RAD = wp.constant(2 * 0.5**2 + (0.5 - LOBATTO_COORDS[1]) ** 2) DEGREE_3_SPHERE_SCALE = 1.0 / (0.75 - DEGREE_3_SPHERE_RAD) @cache.dynamic_func(suffix=self.name) def element_inner_weight( coords: Coords, node_index_in_elt: int, ): node_type, type_index = self.node_type_and_type_index(node_index_in_elt) if node_type == CubeSerendipityShapeFunctions.VERTEX: node_ijk = CubeSerendipityShapeFunctions._vertex_coords(type_index) cx = wp.select(node_ijk[0] == 0, coords[0], 1.0 - coords[0]) cy = wp.select(node_ijk[1] == 0, coords[1], 1.0 - coords[1]) cz = wp.select(node_ijk[2] == 0, coords[2], 1.0 - coords[2]) w = cx * cy * cz if ORDER == 2: w *= cx + cy + cz - 3.0 + LOBATTO_COORDS[1] return w * LAGRANGE_SCALE[0] if ORDER == 3: w *= ( (cx - 0.5) * (cx - 0.5) + (cy - 0.5) * (cy - 0.5) + (cz - 0.5) * (cz - 0.5) - DEGREE_3_SPHERE_RAD ) return w * DEGREE_3_SPHERE_SCALE axis = CubeSerendipityShapeFunctions._edge_axis(node_type) node_all = CubeSerendipityShapeFunctions._edge_coords(type_index) local_coords = Grid3D._world_to_local(axis, coords) w = float(1.0) w *= wp.select(node_all[1] == 0, local_coords[1], 1.0 - local_coords[1]) w *= wp.select(node_all[2] == 0, local_coords[2], 1.0 - local_coords[2]) for k in range(ORDER_PLUS_ONE): if k != node_all[0]: w *= local_coords[0] - LOBATTO_COORDS[k] w *= LAGRANGE_SCALE[node_all[0]] return w return element_inner_weight def make_element_inner_weight_gradient(self): ORDER = self.ORDER ORDER_PLUS_ONE = self.ORDER_PLUS_ONE LOBATTO_COORDS = self.LOBATTO_COORDS LAGRANGE_SCALE = self.LAGRANGE_SCALE DEGREE_3_SPHERE_RAD = wp.constant(2 * 0.5**2 + (0.5 - LOBATTO_COORDS[1]) ** 2) DEGREE_3_SPHERE_SCALE = 1.0 / (0.75 - DEGREE_3_SPHERE_RAD) @cache.dynamic_func(suffix=self.name) def element_inner_weight_gradient( coords: Coords, node_index_in_elt: int, ): node_type, type_index = self.node_type_and_type_index(node_index_in_elt) if node_type == CubeSerendipityShapeFunctions.VERTEX: node_ijk = CubeSerendipityShapeFunctions._vertex_coords(type_index) cx = wp.select(node_ijk[0] == 0, coords[0], 1.0 - coords[0]) cy = wp.select(node_ijk[1] == 0, coords[1], 1.0 - coords[1]) cz = wp.select(node_ijk[2] == 0, coords[2], 1.0 - coords[2]) gx = wp.select(node_ijk[0] == 0, 1.0, -1.0) gy = wp.select(node_ijk[1] == 0, 1.0, -1.0) gz = wp.select(node_ijk[2] == 0, 1.0, -1.0) if ORDER == 2: w = cx + cy + cz - 3.0 + LOBATTO_COORDS[1] grad_x = cy * cz * gx * (w + cx) grad_y = cz * cx * gy * (w + cy) grad_z = cx * cy * gz * (w + cz) return wp.vec3(grad_x, grad_y, grad_z) * LAGRANGE_SCALE[0] if ORDER == 3: w = ( (cx - 0.5) * (cx - 0.5) + (cy - 0.5) * (cy - 0.5) + (cz - 0.5) * (cz - 0.5) - DEGREE_3_SPHERE_RAD ) dw_dcx = 2.0 * cx - 1.0 dw_dcy = 2.0 * cy - 1.0 dw_dcz = 2.0 * cz - 1.0 grad_x = cy * cz * gx * (w + dw_dcx * cx) grad_y = cz * cx * gy * (w + dw_dcy * cy) grad_z = cx * cy * gz * (w + dw_dcz * cz) return wp.vec3(grad_x, grad_y, grad_z) * DEGREE_3_SPHERE_SCALE axis = CubeSerendipityShapeFunctions._edge_axis(node_type) node_all = CubeSerendipityShapeFunctions._edge_coords(type_index) local_coords = Grid3D._world_to_local(axis, coords) w_long = wp.select(node_all[1] == 0, local_coords[1], 1.0 - local_coords[1]) w_lat = wp.select(node_all[2] == 0, local_coords[2], 1.0 - local_coords[2]) g_long = wp.select(node_all[1] == 0, 1.0, -1.0) g_lat = wp.select(node_all[2] == 0, 1.0, -1.0) w_alt = LAGRANGE_SCALE[node_all[0]] g_alt = float(0.0) prefix_alt = LAGRANGE_SCALE[node_all[0]] for k in range(ORDER_PLUS_ONE): if k != node_all[0]: delta_alt = local_coords[0] - LOBATTO_COORDS[k] w_alt *= delta_alt g_alt = g_alt * delta_alt + prefix_alt prefix_alt *= delta_alt local_grad = wp.vec3(g_alt * w_long * w_lat, w_alt * g_long * w_lat, w_alt * w_long * g_lat) return Grid3D._local_to_world(axis, local_grad) return element_inner_weight_gradient def element_node_tets(self): from warp.fem.utils import grid_to_tets if self.ORDER == 2: element_tets = np.array( [ [0, 8, 9, 10], [1, 11, 10, 15], [2, 9, 14, 13], [3, 15, 13, 17], [4, 12, 8, 16], [5, 18, 16, 11], [6, 14, 12, 19], [7, 19, 18, 17], [16, 12, 18, 11], [8, 16, 12, 11], [12, 19, 18, 14], [14, 19, 17, 18], [10, 9, 15, 8], [10, 8, 11, 15], [9, 13, 15, 14], [13, 14, 17, 15], ] ) middle_hex = np.array([8, 11, 9, 15, 12, 18, 14, 17]) middle_tets = middle_hex[grid_to_tets(1, 1, 1)] return np.concatenate((element_tets, middle_tets)) raise NotImplementedError() class CubeNonConformingPolynomialShapeFunctions: # embeds the largest regular tet centered at (0.5, 0.5, 0.5) into the reference cube _tet_height = 2.0 / 3.0 _tet_side = math.sqrt(3.0 / 2.0) * _tet_height _tet_face_height = math.sqrt(3.0) / 2.0 * _tet_side _tet_to_cube = np.array( [ [_tet_side, _tet_side / 2.0, _tet_side / 2.0], [0.0, _tet_face_height, _tet_face_height / 3.0], [0.0, 0.0, _tet_height], ] ) _TET_OFFSET = wp.constant(wp.vec3(0.5 - 0.5 * _tet_side, 0.5 - _tet_face_height / 3.0, 0.5 - 0.25 * _tet_height)) 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 @property def name(self) -> str: return f"Cube_P{self.ORDER}d" def make_node_coords_in_element(self): node_coords_in_tet = self._tet_shape.make_node_coords_in_element() TET_TO_CUBE = wp.constant(wp.mat33(self._tet_to_cube)) @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_TO_CUBE * tet_coords + CubeNonConformingPolynomialShapeFunctions._TET_OFFSET return node_coords_in_element def make_node_quadrature_weight(self): NODES_PER_ELEMENT = self.NODES_PER_ELEMENT @cache.dynamic_func(suffix=self.name) def node_uniform_quadrature_weight( node_index_in_elt: int, ): return 1.0 / float(NODES_PER_ELEMENT) return node_uniform_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() CUBE_TO_TET = wp.constant(wp.mat33(np.linalg.inv(self._tet_to_cube))) @cache.dynamic_func(suffix=self.name) def element_inner_weight( coords: Coords, node_index_in_elt: int, ): tet_coords = CUBE_TO_TET * (coords - CubeNonConformingPolynomialShapeFunctions._TET_OFFSET) 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() CUBE_TO_TET = wp.constant(wp.mat33(np.linalg.inv(self._tet_to_cube))) @cache.dynamic_func(suffix=self.name) def element_inner_weight_gradient( coords: Coords, node_index_in_elt: int, ): tet_coords = CUBE_TO_TET * (coords - CubeNonConformingPolynomialShapeFunctions._TET_OFFSET) grad = tet_inner_weight_gradient(tet_coords, node_index_in_elt) return wp.transpose(CUBE_TO_TET) * grad return element_inner_weight_gradient