Spaces:

qbhf2
/

GarmentCode

Paused

GarmentCode / NvidiaWarp-GarmentCode /warp /fem /space /shape /cube_shape_function.py

added NvidiaWarp and GarmentCode repos

66c9c8a over 1 year ago

26.8 kB

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

	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.52 + (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.52 + (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