warp-lang 1.4.2__py3-none-manylinux2014_aarch64.whl → 1.5.1__py3-none-manylinux2014_aarch64.whl

warp/fem/space/nanogrid_function_space.py

@@ -3,14 +3,9 @@ from typing import Union
 import warp as wp
 from warp.fem import cache
 from warp.fem.geometry import AdaptiveNanogrid, Nanogrid
-from warp.fem.polynomial import is_closed
 from warp.fem.types import ElementIndex
 
-from .shape import (
-    CubeSerendipityShapeFunctions,
-    CubeTripolynomialShapeFunctions,
-    ShapeFunction,
-)
+from .shape import CubeShapeFunction
 from .topology import SpaceTopology, forward_base_topology
 
 
@@ -31,19 +26,17 @@ class NanogridSpaceTopology(SpaceTopology):
     def __init__(
         self,
         grid: Union[Nanogrid, AdaptiveNanogrid],
-        shape: ShapeFunction,
-        need_edge_indices: bool = True,
-        need_face_indices: bool = True,
+        shape: CubeShapeFunction,
     ):
-        if not is_closed(shape.family):
-            raise ValueError("A closed polynomial family is required to define a continuous function space")
-
+        self._shape = shape
         super().__init__(grid, shape.NODES_PER_ELEMENT)
         self._grid = grid
-        self._shape = shape
 
         self._vertex_grid = grid.vertex_grid.id
 
+        need_edge_indices = shape.EDGE_NODE_COUNT > 0
+        need_face_indices = shape.FACE_NODE_COUNT > 0
+
         if isinstance(grid, Nanogrid):
             self._edge_grid = grid.edge_grid.id if need_edge_indices else -1
             self._face_grid = grid.face_grid.id if need_face_indices else -1
@@ -55,6 +48,12 @@ class NanogridSpaceTopology(SpaceTopology):
             self._edge_count = grid.stacked_edge_count() if need_edge_indices else 0
             self._face_count = grid.stacked_face_count() if need_face_indices else 0
 
+        self.element_node_index = self._make_element_node_index()
+
+    @property
+    def name(self):
+        return f"{self.geometry.name}_{self._shape.name}"
+
     @cache.cached_arg_value
     def topo_arg_value(self, device):
         arg = NanogridTopologyArg()
@@ -97,51 +96,19 @@ class NanogridSpaceTopology(SpaceTopology):
 
         return element_node_index_adaptive
 
-
-@wp.func
-def _cell_vertex_coord(cell_ijk: wp.vec3i, cell_level: int, n: int):
-    return cell_ijk + AdaptiveNanogrid.fine_ijk(wp.vec3i((n & 4) >> 2, (n & 2) >> 1, n & 1), cell_level)
-
-
-@wp.func
-def _cell_edge_coord(cell_ijk: wp.vec3i, cell_level: int, axis: int, offset: int):
-    e_ijk = AdaptiveNanogrid.coarse_ijk(cell_ijk, cell_level)
-    e_ijk[(axis + 1) % 3] += offset >> 1
-    e_ijk[(axis + 2) % 3] += offset & 1
-    return AdaptiveNanogrid.encode_axis_and_level(e_ijk, axis, cell_level)
-
-
-@wp.func
-def _cell_face_coord(cell_ijk: wp.vec3i, cell_level: int, axis: int, offset: int):
-    f_ijk = AdaptiveNanogrid.coarse_ijk(cell_ijk, cell_level)
-    f_ijk[axis] += offset
-    return AdaptiveNanogrid.encode_axis_and_level(f_ijk, axis, cell_level)
-
-
-class NanogridTripolynomialSpaceTopology(NanogridSpaceTopology):
-    def __init__(self, grid: Union[Nanogrid, AdaptiveNanogrid], shape: CubeTripolynomialShapeFunctions):
-        super().__init__(grid, shape, need_edge_indices=shape.ORDER >= 2, need_face_indices=shape.ORDER >= 2)
-
-        self.element_node_index = self._make_element_node_index()
-
     def node_count(self) -> int:
-        ORDER = self._shape.ORDER
-        INTERIOR_NODES_PER_EDGE = max(0, ORDER - 1)
-        INTERIOR_NODES_PER_FACE = INTERIOR_NODES_PER_EDGE**2
-        INTERIOR_NODES_PER_CELL = INTERIOR_NODES_PER_EDGE**3
-
         return (
-            self._grid.vertex_count()
-            + self._edge_count * INTERIOR_NODES_PER_EDGE
-            + self._face_count * INTERIOR_NODES_PER_FACE
-            + self._grid.cell_count() * INTERIOR_NODES_PER_CELL
+            self._grid.vertex_count() * self._shape.VERTEX_NODE_COUNT
+            + self._edge_count * self._shape.EDGE_NODE_COUNT
+            + self._face_count * self._shape.FACE_NODE_COUNT
+            + self._grid.cell_count() * self._shape.INTERIOR_NODE_COUNT
         )
 
     def _make_element_node_index_generic(self):
-        ORDER = self._shape.ORDER
-        INTERIOR_NODES_PER_EDGE = wp.constant(max(0, ORDER - 1))
-        INTERIOR_NODES_PER_FACE = wp.constant(INTERIOR_NODES_PER_EDGE**2)
-        INTERIOR_NODES_PER_CELL = wp.constant(INTERIOR_NODES_PER_EDGE**3)
+        VERTEX_NODE_COUNT = self._shape.VERTEX_NODE_COUNT
+        EDGE_NODE_COUNT = self._shape.EDGE_NODE_COUNT
+        FACE_NODE_COUNT = self._shape.FACE_NODE_COUNT
+        INTERIOR_NODE_COUNT = self._shape.INTERIOR_NODE_COUNT
 
         @cache.dynamic_func(suffix=self.name)
         def element_node_index_generic(
@@ -153,82 +120,67 @@ class NanogridTripolynomialSpaceTopology(NanogridSpaceTopology):
         ):
             node_type, type_instance, type_index = self._shape.node_type_and_type_index(node_index_in_elt)
 
-            if node_type == CubeTripolynomialShapeFunctions.VERTEX:
-                n_ijk = _cell_vertex_coord(ijk, level, type_instance)
-                return wp.volume_lookup_index(topo_arg.vertex_grid, n_ijk[0], n_ijk[1], n_ijk[2])
+            if wp.static(VERTEX_NODE_COUNT > 0):
+                if node_type == CubeShapeFunction.VERTEX:
+                    n_ijk = _cell_vertex_coord(ijk, level, type_instance)
+                    return (
+                        wp.volume_lookup_index(topo_arg.vertex_grid, n_ijk[0], n_ijk[1], n_ijk[2]) * VERTEX_NODE_COUNT
+                        + type_index
+                    )
 
-            offset = topo_arg.vertex_count
+            offset = topo_arg.vertex_count * VERTEX_NODE_COUNT
 
-            if node_type == CubeTripolynomialShapeFunctions.EDGE:
-                axis = type_instance >> 2
-                node_offset = type_instance & 3
+            if wp.static(EDGE_NODE_COUNT > 0):
+                if node_type == CubeShapeFunction.EDGE:
+                    axis = type_instance >> 2
+                    node_offset = type_instance & 3
 
-                n_ijk = _cell_edge_coord(ijk, level, axis, node_offset)
+                    n_ijk = _cell_edge_coord(ijk, level, axis, node_offset)
 
-                edge_index = wp.volume_lookup_index(topo_arg.edge_grid, n_ijk[0], n_ijk[1], n_ijk[2])
-                return offset + INTERIOR_NODES_PER_EDGE * edge_index + type_index
+                    edge_index = wp.volume_lookup_index(topo_arg.edge_grid, n_ijk[0], n_ijk[1], n_ijk[2])
+                    return offset + EDGE_NODE_COUNT * edge_index + type_index
 
-            offset += INTERIOR_NODES_PER_EDGE * topo_arg.edge_count
+                offset += EDGE_NODE_COUNT * topo_arg.edge_count
 
-            if node_type == CubeTripolynomialShapeFunctions.FACE:
-                axis = type_instance >> 1
-                node_offset = type_instance & 1
+            if wp.static(FACE_NODE_COUNT > 0):
+                if node_type == CubeShapeFunction.FACE:
+                    axis = type_instance >> 1
+                    node_offset = type_instance & 1
 
-                n_ijk = _cell_face_coord(ijk, level, axis, node_offset)
+                    n_ijk = _cell_face_coord(ijk, level, axis, node_offset)
 
-                face_index = wp.volume_lookup_index(topo_arg.face_grid, n_ijk[0], n_ijk[1], n_ijk[2])
-                return offset + INTERIOR_NODES_PER_FACE * face_index + type_index
+                    face_index = wp.volume_lookup_index(topo_arg.face_grid, n_ijk[0], n_ijk[1], n_ijk[2])
+                    return offset + FACE_NODE_COUNT * face_index + type_index
 
-            offset += INTERIOR_NODES_PER_FACE * topo_arg.face_count
+                offset += FACE_NODE_COUNT * topo_arg.face_count
 
-            return offset + INTERIOR_NODES_PER_CELL * element_index + type_index
+            return offset + INTERIOR_NODE_COUNT * element_index + type_index
 
         return element_node_index_generic
 
 
-class NanogridSerendipitySpaceTopology(NanogridSpaceTopology):
-    def __init__(self, grid: Nanogrid, shape: CubeSerendipityShapeFunctions):
-        super().__init__(grid, shape, need_edge_indices=True, need_face_indices=False)
-
-        self.element_node_index = self._make_element_node_index()
-
-    def node_count(self) -> int:
-        return self.geometry.vertex_count() + (self._shape.ORDER - 1) * self._edge_count
-
-    def _make_element_node_index_generic(self):
-        ORDER = self._shape.ORDER
-
-        @cache.dynamic_func(suffix=self.name)
-        def element_node_index_generic(
-            topo_arg: NanogridTopologyArg,
-            element_index: ElementIndex,
-            node_index_in_elt: int,
-            ijk: wp.vec3i,
-            level: int,
-        ):
-            node_type, type_index = self._shape.node_type_and_type_index(node_index_in_elt)
-
-            if node_type == CubeSerendipityShapeFunctions.VERTEX:
-                n_ijk = _cell_vertex_coord(ijk, level, type_index)
-                return wp.volume_lookup_index(topo_arg.vertex_grid, n_ijk[0], n_ijk[1], n_ijk[2])
-
-            type_instance, index_in_edge = CubeSerendipityShapeFunctions._cube_edge_index(node_type, type_index)
-            axis = type_instance >> 2
-            node_offset = type_instance & 3
+@wp.func
+def _cell_vertex_coord(cell_ijk: wp.vec3i, cell_level: int, n: int):
+    return cell_ijk + AdaptiveNanogrid.fine_ijk(wp.vec3i((n & 4) >> 2, (n & 2) >> 1, n & 1), cell_level)
 
-            n_ijk = _cell_edge_coord(ijk, level, axis, node_offset)
 
-            edge_index = wp.volume_lookup_index(topo_arg.edge_grid, n_ijk[0], n_ijk[1], n_ijk[2])
-            return topo_arg.vertex_count + (ORDER - 1) * edge_index + index_in_edge
+@wp.func
+def _cell_edge_coord(cell_ijk: wp.vec3i, cell_level: int, axis: int, offset: int):
+    e_ijk = AdaptiveNanogrid.coarse_ijk(cell_ijk, cell_level)
+    e_ijk[(axis + 1) % 3] += offset >> 1
+    e_ijk[(axis + 2) % 3] += offset & 1
+    return AdaptiveNanogrid.encode_axis_and_level(e_ijk, axis, cell_level)
 
-        return element_node_index_generic
 
+@wp.func
+def _cell_face_coord(cell_ijk: wp.vec3i, cell_level: int, axis: int, offset: int):
+    f_ijk = AdaptiveNanogrid.coarse_ijk(cell_ijk, cell_level)
+    f_ijk[axis] += offset
+    return AdaptiveNanogrid.encode_axis_and_level(f_ijk, axis, cell_level)
 
-def make_nanogrid_space_topology(grid: Union[Nanogrid, AdaptiveNanogrid], shape: ShapeFunction):
-    if isinstance(shape, CubeSerendipityShapeFunctions):
-        return forward_base_topology(NanogridSerendipitySpaceTopology, grid, shape)
 
-    if isinstance(shape, CubeTripolynomialShapeFunctions):
-        return forward_base_topology(NanogridTripolynomialSpaceTopology, grid, shape)
+def make_nanogrid_space_topology(grid: Union[Nanogrid, AdaptiveNanogrid], shape: CubeShapeFunction):
+    if isinstance(shape, CubeShapeFunction):
+        return forward_base_topology(NanogridSpaceTopology, grid, shape)
 
     raise ValueError(f"Unsupported shape function {shape.name}")

warp/fem/space/quadmesh_function_space.py

@@ -0,0 +1,208 @@
+import warp as wp
+from warp.fem import cache
+from warp.fem.geometry import Quadmesh2D
+from warp.fem.polynomial import is_closed
+from warp.fem.types import NULL_NODE_INDEX, ElementIndex
+
+from .shape import SquareShapeFunction
+from .topology import SpaceTopology, forward_base_topology
+
+
+@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}"
+
+    @cache.cached_arg_value
+    def topo_arg_value(self, device):
+        arg = Quadmesh2DTopologyArg()
+        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()
+        return arg
+
+    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.select(
+                    node_type == SquareShapeFunction.EDGE_X,
+                    wp.select(type_instance == 0, 1, 3),
+                    wp.select(type_instance == 0, 2, 0),
+                )
+
+                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.select(flipped, type_index, EDGE_NODE_COUNT - 1 - 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.select(
+                    node_type == SquareShapeFunction.EDGE_X,
+                    wp.select(type_instance == 0, 1, 3),
+                    wp.select(type_instance == 0, 2, 0),
+                )
+
+                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.select(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/fem/space/shape/__init__.py

@@ -5,29 +5,51 @@ from warp.fem.geometry import element as _element
 from warp.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,
 )
-from .tet_shape_function import TetrahedronNonConformingPolynomialShapeFunctions, TetrahedronPolynomialShapeFunctions
-from .triangle_shape_function import Triangle2DNonConformingPolynomialShapeFunctions, Triangle2DPolynomialShapeFunctions
 
 
 class ElementBasis(Enum):
     """Choice of basis function to equip individual elements"""
 
-    LAGRANGE = 0
+    LAGRANGE = "P"
     """Lagrange basis functions :math:`P_k` for simplices, tensor products :math:`Q_k` for squares and cubes"""
-    SERENDIPITY = 1
+    SERENDIPITY = "S"
     """Serendipity elements :math:`S_k`, corresponding to Lagrange nodes with interior points removed (for degree <= 3)"""
-    NONCONFORMING_POLYNOMIAL = 2
+    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."""
 
 
 def get_shape_function(
@@ -58,6 +80,10 @@ def get_shape_function(
         family = Polynomial.LOBATTO_GAUSS_LEGENDRE
 
     if isinstance(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:
@@ -65,14 +91,22 @@ def get_shape_function(
 
         return SquareBipolynomialShapeFunctions(degree=degree, family=family)
     if isinstance(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 Triangle2DNonConformingPolynomialShapeFunctions(degree=degree)
+            return TriangleNonConformingPolynomialShapeFunctions(degree=degree)
         if element_basis == ElementBasis.SERENDIPITY and degree > 2:
             raise NotImplementedError("Serendipity variant not implemented yet for Triangle elements")
 
-        return Triangle2DPolynomialShapeFunctions(degree=degree)
+        return TrianglePolynomialShapeFunctions(degree=degree)
 
     if isinstance(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:
@@ -80,6 +114,10 @@ def get_shape_function(
 
         return CubeTripolynomialShapeFunctions(degree=degree, family=family)
     if isinstance(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: