warp-lang 1.0.1__py3-none-manylinux2014_aarch64.whl → 1.1.0__py3-none-manylinux2014_aarch64.whl

warp/fem/space/shape/tet_shape_function.py

@@ -1,567 +1,565 @@
-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
+import numpy as np
+
+import warp as wp
+from warp.fem import cache
+from warp.fem.types import Coords
+
+
+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 (counterclockwise), 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