warp-lang 1.10.0__py3-none-macosx_11_0_arm64.whl

warp/_src/fem/space/quadmesh_function_space.py

@@ -0,0 +1,222 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import warp as wp
+from warp._src.fem import cache
+from warp._src.fem.geometry import Quadmesh2D
+from warp._src.fem.polynomial import is_closed
+from warp._src.fem.types import NULL_NODE_INDEX, ElementIndex
+
+from .shape import SquareShapeFunction
+from .topology import SpaceTopology, forward_base_topology
+
+_wp_module_name_ = "warp.fem.space.quadmesh_function_space"
+
+
+@wp.struct
+class Quadmesh2DTopologyArg:
+    edge_vertex_indices: wp.array(dtype=wp.vec2i)
+    quad_edge_indices: wp.array2d(dtype=int)
+
+    vertex_count: int
+    edge_count: int
+    cell_count: int
+
+
+class QuadmeshSpaceTopology(SpaceTopology):
+    TopologyArg = Quadmesh2DTopologyArg
+
+    def __init__(self, mesh: Quadmesh2D, shape: SquareShapeFunction):
+        if shape.value == SquareShapeFunction.Value.Scalar and not is_closed(shape.family):
+            raise ValueError("A closed polynomial family is required to define a continuous function space")
+
+        self._shape = shape
+        super().__init__(mesh, shape.NODES_PER_ELEMENT)
+        self._mesh = mesh
+
+        self._compute_quad_edge_indices()
+        self.element_node_index = self._make_element_node_index()
+        self.element_node_sign = self._make_element_node_sign()
+
+    @property
+    def name(self):
+        return f"{self.geometry.name}_{self._shape.name}"
+
+    def fill_topo_arg(self, arg: Quadmesh2DTopologyArg, device):
+        arg.quad_edge_indices = self._quad_edge_indices.to(device)
+        arg.edge_vertex_indices = self._mesh.edge_vertex_indices.to(device)
+
+        arg.vertex_count = self._mesh.vertex_count()
+        arg.edge_count = self._mesh.side_count()
+        arg.cell_count = self._mesh.cell_count()
+
+    def _compute_quad_edge_indices(self):
+        self._quad_edge_indices = wp.empty(
+            dtype=int, device=self._mesh.quad_vertex_indices.device, shape=(self._mesh.cell_count(), 4)
+        )
+
+        wp.launch(
+            kernel=QuadmeshSpaceTopology._compute_quad_edge_indices_kernel,
+            dim=self._mesh.edge_quad_indices.shape,
+            device=self._mesh.quad_vertex_indices.device,
+            inputs=[
+                self._mesh.edge_quad_indices,
+                self._mesh.edge_vertex_indices,
+                self._mesh.quad_vertex_indices,
+                self._quad_edge_indices,
+            ],
+        )
+
+    @wp.func
+    def _find_edge_index_in_quad(
+        edge_vtx: wp.vec2i,
+        quad_vtx: wp.vec4i,
+    ):
+        for k in range(3):
+            if (edge_vtx[0] == quad_vtx[k] and edge_vtx[1] == quad_vtx[k + 1]) or (
+                edge_vtx[1] == quad_vtx[k] and edge_vtx[0] == quad_vtx[k + 1]
+            ):
+                return k
+        return 3
+
+    @wp.kernel
+    def _compute_quad_edge_indices_kernel(
+        edge_quad_indices: wp.array(dtype=wp.vec2i),
+        edge_vertex_indices: wp.array(dtype=wp.vec2i),
+        quad_vertex_indices: wp.array2d(dtype=int),
+        quad_edge_indices: wp.array2d(dtype=int),
+    ):
+        e = wp.tid()
+
+        edge_vtx = edge_vertex_indices[e]
+        edge_quads = edge_quad_indices[e]
+
+        q0 = edge_quads[0]
+        q0_vtx = wp.vec4i(
+            quad_vertex_indices[q0, 0],
+            quad_vertex_indices[q0, 1],
+            quad_vertex_indices[q0, 2],
+            quad_vertex_indices[q0, 3],
+        )
+        q0_edge = QuadmeshSpaceTopology._find_edge_index_in_quad(edge_vtx, q0_vtx)
+        quad_edge_indices[q0, q0_edge] = e
+
+        q1 = edge_quads[1]
+        if q1 != q0:
+            t1_vtx = wp.vec4i(
+                quad_vertex_indices[q1, 0],
+                quad_vertex_indices[q1, 1],
+                quad_vertex_indices[q1, 2],
+                quad_vertex_indices[q1, 3],
+            )
+            t1_edge = QuadmeshSpaceTopology._find_edge_index_in_quad(edge_vtx, t1_vtx)
+            quad_edge_indices[q1, t1_edge] = e
+
+    def node_count(self) -> int:
+        return (
+            self.geometry.vertex_count() * self._shape.VERTEX_NODE_COUNT
+            + self.geometry.side_count() * self._shape.EDGE_NODE_COUNT
+            + self.geometry.cell_count() * self._shape.INTERIOR_NODE_COUNT
+        )
+
+    def _make_element_node_index(self):
+        VERTEX_NODE_COUNT = self._shape.VERTEX_NODE_COUNT
+        EDGE_NODE_COUNT = self._shape.EDGE_NODE_COUNT
+        INTERIOR_NODE_COUNT = self._shape.INTERIOR_NODE_COUNT
+
+        SHAPE_TO_QUAD_IDX = wp.constant(wp.vec4i([0, 3, 1, 2]))
+
+        @cache.dynamic_func(suffix=self.name)
+        def element_node_index(
+            cell_arg: self._mesh.CellArg,
+            topo_arg: QuadmeshSpaceTopology.TopologyArg,
+            element_index: ElementIndex,
+            node_index_in_elt: int,
+        ):
+            node_type, type_instance, type_index = self._shape.node_type_and_type_index(node_index_in_elt)
+
+            if wp.static(VERTEX_NODE_COUNT > 0):
+                if node_type == SquareShapeFunction.VERTEX:
+                    return (
+                        cell_arg.topology.quad_vertex_indices[element_index, SHAPE_TO_QUAD_IDX[type_instance]]
+                        * VERTEX_NODE_COUNT
+                        + type_index
+                    )
+
+            global_offset = topo_arg.vertex_count * VERTEX_NODE_COUNT
+
+            if wp.static(INTERIOR_NODE_COUNT > 0):
+                if node_type == SquareShapeFunction.INTERIOR:
+                    return global_offset + element_index * INTERIOR_NODE_COUNT + type_index
+
+                global_offset += INTERIOR_NODE_COUNT * topo_arg.cell_count
+
+            if wp.static(EDGE_NODE_COUNT > 0):
+                # EDGE_X, EDGE_Y
+                side_start = wp.where(
+                    node_type == SquareShapeFunction.EDGE_X,
+                    wp.where(type_instance == 0, 0, 2),
+                    wp.where(type_instance == 0, 3, 1),
+                )
+
+                side_index = topo_arg.quad_edge_indices[element_index, side_start]
+                local_vs = cell_arg.topology.quad_vertex_indices[element_index, side_start]
+                global_vs = topo_arg.edge_vertex_indices[side_index][0]
+
+                # Flip indexing direction
+                flipped = int(side_start >= 2) ^ int(local_vs != global_vs)
+                index_in_side = wp.where(flipped, EDGE_NODE_COUNT - 1 - type_index, type_index)
+
+                return global_offset + EDGE_NODE_COUNT * side_index + index_in_side
+
+            return NULL_NODE_INDEX  # should never happen
+
+        return element_node_index
+
+    def _make_element_node_sign(self):
+        @cache.dynamic_func(suffix=self.name)
+        def element_node_sign(
+            cell_arg: self._mesh.CellArg,
+            topo_arg: QuadmeshSpaceTopology.TopologyArg,
+            element_index: ElementIndex,
+            node_index_in_elt: int,
+        ):
+            node_type, type_instance, type_index = self._shape.node_type_and_type_index(node_index_in_elt)
+
+            if node_type == SquareShapeFunction.EDGE_X or node_type == SquareShapeFunction.EDGE_Y:
+                side_start = wp.where(
+                    node_type == SquareShapeFunction.EDGE_X,
+                    wp.where(type_instance == 0, 0, 2),
+                    wp.where(type_instance == 0, 3, 1),
+                )
+
+                side_index = topo_arg.quad_edge_indices[element_index, side_start]
+                local_vs = cell_arg.topology.quad_vertex_indices[element_index, side_start]
+                global_vs = topo_arg.edge_vertex_indices[side_index][0]
+
+                # Flip indexing direction
+                flipped = int(side_start >= 2) ^ int(local_vs != global_vs)
+                return wp.where(flipped, -1.0, 1.0)
+
+            return 1.0
+
+        return element_node_sign
+
+
+def make_quadmesh_space_topology(mesh: Quadmesh2D, shape: SquareShapeFunction):
+    if isinstance(shape, SquareShapeFunction):
+        return forward_base_topology(QuadmeshSpaceTopology, mesh, shape)
+
+    raise ValueError(f"Unsupported shape function {shape.name}")

warp/_src/fem/space/restriction.py

@@ -0,0 +1,221 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Optional
+
+import warp as wp
+from warp._src.fem import cache
+from warp._src.fem.domain import GeometryDomain
+from warp._src.fem.types import NULL_ELEMENT_INDEX, NULL_NODE_INDEX, NodeElementIndex
+from warp._src.fem.utils import compress_node_indices, host_read_at_index
+
+from .partition import SpacePartition
+
+_wp_module_name_ = "warp.fem.space.restriction"
+
+wp.set_module_options({"enable_backward": False})
+
+
+class SpaceRestriction:
+    """Restriction of a space partition to a given GeometryDomain"""
+
+    def __init__(
+        self,
+        space_partition: SpacePartition,
+        domain: GeometryDomain,
+        device=None,
+        temporary_store: cache.TemporaryStore = None,
+    ):
+        space_topology = space_partition.space_topology
+
+        if domain.dimension == space_topology.dimension - 1:
+            space_topology = space_topology.trace()
+
+        if domain.dimension != space_topology.dimension:
+            raise ValueError("Incompatible space and domain dimensions")
+
+        self.space_partition = space_partition
+        self.space_topology = space_topology
+        self.domain = domain
+
+        self._node_count_dev: wp.array = None
+        """Number of unique partition node indices"""
+        self._dof_partition_indices: wp.array = None
+        """Array of unique partition node indices"""
+
+        self._dof_partition_element_offsets: wp.array = None
+        """Mapping from partition node to offset in the per-node element indices array"""
+        self._dof_element_indices: wp.array = None
+        """Concatenation of neighboring elements indices for each partition node"""
+        self._dof_indices_in_element: wp.array = None
+        """Concatenation of node index in element for each partition node"""
+
+        self.rebuild(device=device, temporary_store=temporary_store)
+
+    def rebuild(self, device: Optional = None, temporary_store: Optional[cache.TemporaryStore] = None):
+        max_nodes_per_element = self.space_topology.MAX_NODES_PER_ELEMENT
+
+        @cache.dynamic_kernel(
+            suffix=(self.domain.name, self.space_topology.name, self.space_partition.name),
+            kernel_options={"max_unroll": 8},
+        )
+        def fill_element_node_indices(
+            element_arg: self.domain.ElementArg,
+            domain_index_arg: self.domain.ElementIndexArg,
+            topo_arg: self.space_topology.TopologyArg,
+            partition_arg: self.space_partition.PartitionArg,
+            element_node_indices: wp.array2d(dtype=int),
+        ):
+            domain_element_index = wp.tid()
+            element_index = self.domain.element_index(domain_index_arg, domain_element_index)
+
+            if element_index == NULL_ELEMENT_INDEX:
+                element_node_count = 0
+            else:
+                element_node_count = self.space_topology.element_node_count(element_arg, topo_arg, element_index)
+
+            for n in range(element_node_count):
+                space_nidx = self.space_topology.element_node_index(element_arg, topo_arg, element_index, n)
+                partition_nidx = self.space_partition.partition_node_index(partition_arg, space_nidx)
+                element_node_indices[domain_element_index, n] = partition_nidx
+            for n in range(element_node_count, element_node_indices.shape[1]):
+                element_node_indices[domain_element_index, n] = NULL_NODE_INDEX
+
+        element_node_indices = cache.borrow_temporary(
+            temporary_store,
+            shape=(self.domain.element_count(), max_nodes_per_element),
+            dtype=int,
+            device=device,
+        )
+        wp.launch(
+            dim=element_node_indices.shape[0],
+            kernel=fill_element_node_indices,
+            inputs=[
+                self.domain.element_arg_value(device),
+                self.domain.element_index_arg_value(device),
+                self.space_topology.topo_arg_value(device),
+                self.space_partition.partition_arg_value(device),
+                element_node_indices,
+            ],
+            device=device,
+        )
+
+        # Build compressed map from node to element indices
+        flattened_node_indices = element_node_indices.flatten()
+        (
+            self._dof_partition_element_offsets,
+            node_array_indices,
+            self._node_count_dev,
+            self._dof_partition_indices,
+        ) = compress_node_indices(
+            self.space_partition.node_count(),
+            flattened_node_indices,
+            node_offsets=self._dof_partition_element_offsets,
+            unique_node_count=self._node_count_dev,
+            unique_node_indices=self._dof_partition_indices,
+            return_unique_nodes=True,
+            temporary_store=temporary_store,
+        )
+
+        # Extract element index and index in element
+        if self._dof_element_indices is None or self._dof_element_indices.shape != flattened_node_indices.shape:
+            self._dof_element_indices = cache.borrow_temporary_like(flattened_node_indices, temporary_store)
+            self._dof_indices_in_element = cache.borrow_temporary_like(flattened_node_indices, temporary_store)
+
+        wp.launch(
+            kernel=SpaceRestriction._split_vertex_element_index,
+            dim=flattened_node_indices.shape,
+            inputs=[
+                max_nodes_per_element,
+                node_array_indices,
+                self._dof_element_indices,
+                self._dof_indices_in_element,
+            ],
+            device=flattened_node_indices.device,
+        )
+
+        node_array_indices.release()
+
+        # Upper bound on node count, use `node_count_sync` to get the actual value
+        self._node_count = min(self.space_partition.node_count(), self._dof_partition_indices.shape[0])
+
+    def node_count_sync(self) -> int:
+        """Ensures that the node count is synchronized with the device and returns it"""
+        if self._node_count_dev is not None:
+            self._node_count = int(host_read_at_index(self._node_count_dev, index=0))
+            self._node_count_dev = None
+        return self.node_count()
+
+    def node_count(self) -> int:
+        """Upper bound for the node count, use `node_count_sync` to get the actual value"""
+        return self._node_count
+
+    def partition_element_offsets(self):
+        return self._dof_partition_element_offsets
+
+    def node_partition_indices(self):
+        return self._dof_partition_indices
+
+    def total_node_element_count(self):
+        return self._dof_element_indices.size
+
+    @wp.struct
+    class NodeArg:
+        dof_element_offsets: wp.array(dtype=int)
+        dof_element_indices: wp.array(dtype=int)
+        dof_partition_indices: wp.array(dtype=int)
+        dof_indices_in_element: wp.array(dtype=int)
+
+    @cache.cached_arg_value
+    def node_arg_value(self, device):
+        arg = SpaceRestriction.NodeArg()
+        self.fill_node_arg(arg, device)
+        return arg
+
+    def fill_node_arg(self, arg: NodeArg, device):
+        arg.dof_element_offsets = self._dof_partition_element_offsets.to(device)
+        arg.dof_element_indices = self._dof_element_indices.to(device)
+        arg.dof_partition_indices = self._dof_partition_indices.to(device)
+        arg.dof_indices_in_element = self._dof_indices_in_element.to(device)
+
+    @wp.func
+    def node_partition_index(args: NodeArg, restriction_node_index: int):
+        return args.dof_partition_indices[restriction_node_index]
+
+    @wp.func
+    def node_partition_index_from_element_offset(args: NodeArg, element_offset: int):
+        return wp.lower_bound(args.dof_element_offsets, element_offset + 1) - 1
+
+    @wp.func
+    def node_element_range(args: NodeArg, partition_node_index: int):
+        return args.dof_element_offsets[partition_node_index], args.dof_element_offsets[partition_node_index + 1]
+
+    @wp.func
+    def node_element_index(args: NodeArg, node_element_offset: int):
+        domain_element_index = args.dof_element_indices[node_element_offset]
+        index_in_element = args.dof_indices_in_element[node_element_offset]
+        return NodeElementIndex(domain_element_index, index_in_element)
+
+    @wp.kernel
+    def _split_vertex_element_index(
+        vertex_per_element: int,
+        sorted_indices: wp.array(dtype=int),
+        vertex_element_index: wp.array(dtype=int),
+        vertex_index_in_element: wp.array(dtype=int),
+    ):
+        idx = sorted_indices[wp.tid()]
+        element_index = idx // vertex_per_element
+        vertex_element_index[wp.tid()] = element_index
+        vertex_index_in_element[wp.tid()] = idx - vertex_per_element * element_index

warp/_src/fem/space/shape/__init__.py

@@ -0,0 +1,152 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import functools
+from enum import Enum
+from typing import Optional
+
+from warp._src.fem.geometry import Element
+from warp._src.fem.polynomial import Polynomial
+
+from .cube_shape_function import (
+    CubeNedelecFirstKindShapeFunctions,
+    CubeNonConformingPolynomialShapeFunctions,
+    CubeRaviartThomasShapeFunctions,
+    CubeSerendipityShapeFunctions,
+    CubeShapeFunction,
+    CubeTripolynomialShapeFunctions,
+)
+from .shape_function import ConstantShapeFunction, ShapeFunction
+from .square_shape_function import (
+    SquareBipolynomialShapeFunctions,
+    SquareNedelecFirstKindShapeFunctions,
+    SquareNonConformingPolynomialShapeFunctions,
+    SquareRaviartThomasShapeFunctions,
+    SquareSerendipityShapeFunctions,
+    SquareShapeFunction,
+)
+from .tet_shape_function import (
+    TetrahedronNedelecFirstKindShapeFunctions,
+    TetrahedronNonConformingPolynomialShapeFunctions,
+    TetrahedronPolynomialShapeFunctions,
+    TetrahedronRaviartThomasShapeFunctions,
+    TetrahedronShapeFunction,
+)
+from .triangle_shape_function import (
+    TriangleNedelecFirstKindShapeFunctions,
+    TriangleNonConformingPolynomialShapeFunctions,
+    TrianglePolynomialShapeFunctions,
+    TriangleRaviartThomasShapeFunctions,
+    TriangleShapeFunction,
+)
+
+
+class ElementBasis(Enum):
+    """Choice of basis function to equip individual elements"""
+
+    LAGRANGE = "P"
+    """Lagrange basis functions :math:`P_k` for simplices, tensor products :math:`Q_k` for squares and cubes"""
+    SERENDIPITY = "S"
+    """Serendipity elements :math:`S_k`, corresponding to Lagrange nodes with interior points removed (for degree <= 3)"""
+    NONCONFORMING_POLYNOMIAL = "dP"
+    """Simplex Lagrange basis functions :math:`P_{kd}` embedded into non conforming reference elements (e.g. squares or cubes). Discontinuous only."""
+    NEDELEC_FIRST_KIND = "N1"
+    """Nédélec (first kind) H(curl) shape functions. Should be used with covariant function space."""
+    RAVIART_THOMAS = "RT"
+    """Raviart-Thomas H(div) shape functions. Should be used with contravariant function space."""
+
+
+@functools.lru_cache(maxsize=None)
+def make_element_shape_function(
+    element: Element,
+    degree: int,
+    element_basis: Optional[ElementBasis] = None,
+    family: Optional[Polynomial] = None,
+) -> ShapeFunction:
+    """
+    Equips a reference element with a shape function basis.
+
+    Args:
+        element: the type of reference element on which to build the shape function
+        degree: polynomial degree of the per-element shape functions
+        element_basis: type of basis function for the individual elements
+        family: Polynomial family used to generate the shape function basis. If not provided, a reasonable basis is chosen.
+
+    Returns:
+        the corresponding shape function
+
+    Raises:
+        NotImplementedError: If the shape function is not implemented for the given element type
+    """
+
+    if element_basis is None:
+        element_basis = ElementBasis.LAGRANGE
+    elif element_basis == ElementBasis.SERENDIPITY and degree == 1:
+        # Degree-1 serendipity is always equivalent to Lagrange
+        element_basis = ElementBasis.LAGRANGE
+
+    if degree == 0:
+        return ConstantShapeFunction(element)
+
+    if family is None:
+        family = Polynomial.LOBATTO_GAUSS_LEGENDRE
+
+    if element == Element.SQUARE:
+        if element_basis == ElementBasis.NEDELEC_FIRST_KIND:
+            return SquareNedelecFirstKindShapeFunctions(degree=degree)
+        if element_basis == ElementBasis.RAVIART_THOMAS:
+            return SquareRaviartThomasShapeFunctions(degree=degree)
+        if element_basis == ElementBasis.NONCONFORMING_POLYNOMIAL:
+            return SquareNonConformingPolynomialShapeFunctions(degree=degree)
+        if element_basis == ElementBasis.SERENDIPITY and degree > 1:
+            return SquareSerendipityShapeFunctions(degree=degree, family=family)
+
+        return SquareBipolynomialShapeFunctions(degree=degree, family=family)
+    if element == Element.TRIANGLE:
+        if element_basis == ElementBasis.NEDELEC_FIRST_KIND:
+            return TriangleNedelecFirstKindShapeFunctions(degree=degree)
+        if element_basis == ElementBasis.RAVIART_THOMAS:
+            return TriangleRaviartThomasShapeFunctions(degree=degree)
+        if element_basis == ElementBasis.NONCONFORMING_POLYNOMIAL:
+            return TriangleNonConformingPolynomialShapeFunctions(degree=degree)
+        if element_basis == ElementBasis.SERENDIPITY and degree > 2:
+            raise NotImplementedError("Serendipity variant not implemented yet for Triangle elements")
+
+        return TrianglePolynomialShapeFunctions(degree=degree)
+
+    if element == Element.CUBE:
+        if element_basis == ElementBasis.NEDELEC_FIRST_KIND:
+            return CubeNedelecFirstKindShapeFunctions(degree=degree)
+        if element_basis == ElementBasis.RAVIART_THOMAS:
+            return CubeRaviartThomasShapeFunctions(degree=degree)
+        if element_basis == ElementBasis.NONCONFORMING_POLYNOMIAL:
+            return CubeNonConformingPolynomialShapeFunctions(degree=degree)
+        if element_basis == ElementBasis.SERENDIPITY and degree > 1:
+            return CubeSerendipityShapeFunctions(degree=degree, family=family)
+
+        return CubeTripolynomialShapeFunctions(degree=degree, family=family)
+    if element == Element.TETRAHEDRON:
+        if element_basis == ElementBasis.NEDELEC_FIRST_KIND:
+            return TetrahedronNedelecFirstKindShapeFunctions(degree=degree)
+        if element_basis == ElementBasis.RAVIART_THOMAS:
+            return TetrahedronRaviartThomasShapeFunctions(degree=degree)
+        if element_basis == ElementBasis.NONCONFORMING_POLYNOMIAL:
+            return TetrahedronNonConformingPolynomialShapeFunctions(degree=degree)
+        if element_basis == ElementBasis.SERENDIPITY and degree > 2:
+            raise NotImplementedError("Serendipity variant not implemented yet for Tet elements")
+
+        return TetrahedronPolynomialShapeFunctions(degree=degree)
+
+    raise NotImplementedError(f"Unrecognized element type {element}")