Spaces:
Sleeping
Sleeping
| 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() | |
| def name(self) -> str: | |
| return f"Cube_Q{self.ORDER}_{self.family}" | |
| 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 | |
| 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 | |
| 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() | |
| def name(self) -> str: | |
| return f"Cube_S{self.ORDER}_{self.family}" | |
| def _get_node_type_and_type_index(self): | |
| 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 | |
| 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) | |
| 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),) | |
| def _edge_axis(node_type: int): | |
| return node_type - CubeSerendipityShapeFunctions.EDGE_X | |
| 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 | |
| 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 | |
| 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 | |
| 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 | |
| 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) | |
| 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) | |
| 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 | |
| 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)) | |
| 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 | |
| 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 | |
| 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))) | |
| 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))) | |
| 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 | |