warp-lang 1.1.0__py3-none-manylinux2014_x86_64.whl → 1.2.0__py3-none-manylinux2014_x86_64.whl

warp/fem/space/hexmesh_function_space.py

@@ -6,23 +6,24 @@ from warp.fem.geometry.hexmesh import (
     FACE_ORIENTATION,
     FACE_TRANSLATION,
 )
-from warp.fem.polynomial import Polynomial, is_closed
-from warp.fem.types import Coords, ElementIndex
+from warp.fem.polynomial import is_closed
+from warp.fem.types import ElementIndex
 
-from .basis_space import ShapeBasisSpace, TraceBasisSpace
 from .shape import (
-    ConstantShapeFunction,
-    CubeNonConformingPolynomialShapeFunctions,
     CubeSerendipityShapeFunctions,
     CubeTripolynomialShapeFunctions,
     ShapeFunction,
 )
-from .topology import DiscontinuousSpaceTopologyMixin, SpaceTopology, forward_base_topology
+from .topology import SpaceTopology, forward_base_topology
 
 _FACE_ORIENTATION_I = wp.constant(wp.mat(shape=(16, 2), dtype=int)(FACE_ORIENTATION))
 _FACE_TRANSLATION_I = wp.constant(wp.mat(shape=(4, 2), dtype=int)(FACE_TRANSLATION))
 
-_CUBE_VERTEX_INDICES = wp.constant(wp.vec(length=8, dtype=int)([0, 4, 3, 7, 1, 5, 2, 6]))
+# map from shape function vertex indexing to hexmesh vertex indexing
+_CUBE_TO_HEX_VERTEX = wp.constant(wp.vec(length=8, dtype=int)([0, 4, 3, 7, 1, 5, 2, 6]))
+
+# map from shape function edge indexing to hexmesh edge indexing
+_CUBE_TO_HEX_EDGE = wp.constant(wp.vec(length=12, dtype=int)([0, 4, 2, 6, 3, 1, 7, 5, 8, 11, 9, 10]))
 
 
 @wp.struct
@@ -45,9 +46,12 @@ class HexmeshSpaceTopology(SpaceTopology):
         need_hex_edge_indices: bool = True,
         need_hex_face_indices: bool = True,
     ):
+        if not is_closed(shape.family):
+            raise ValueError("A closed polynomial family is required to define a continuous function space")
+
         super().__init__(mesh, shape.NODES_PER_ELEMENT)
         self._mesh = mesh
-        self._shape = shape
+        self.shape = shape
 
         if need_hex_edge_indices:
             self._hex_edge_indices = self._mesh.hex_edge_indices
@@ -111,44 +115,6 @@ class HexmeshSpaceTopology(SpaceTopology):
         hex_face_indices[hx1, local_face_1] = wp.vec2i(f, ori_1)
 
 
-class HexmeshDiscontinuousSpaceTopology(
-    DiscontinuousSpaceTopologyMixin,
-    SpaceTopology,
-):
-    def __init__(self, mesh: Hexmesh, shape: ShapeFunction):
-        super().__init__(mesh, shape.NODES_PER_ELEMENT)
-
-
-class HexmeshBasisSpace(ShapeBasisSpace):
-    def __init__(self, topology: HexmeshSpaceTopology, shape: ShapeFunction):
-        super().__init__(topology, shape)
-
-        self._mesh: Hexmesh = topology.geometry
-
-
-class HexmeshPiecewiseConstantBasis(HexmeshBasisSpace):
-    def __init__(self, mesh: Hexmesh):
-        shape = ConstantShapeFunction(mesh.reference_cell(), space_dimension=3)
-        topology = HexmeshDiscontinuousSpaceTopology(mesh, shape)
-        super().__init__(shape=shape, topology=topology)
-
-    class Trace(TraceBasisSpace):
-        @wp.func
-        def _node_coords_in_element(
-            side_arg: Hexmesh.SideArg,
-            basis_arg: HexmeshBasisSpace.BasisArg,
-            element_index: ElementIndex,
-            node_index_in_element: int,
-        ):
-            return Coords(0.5, 0.5, 0.0)
-
-        def make_node_coords_in_element(self):
-            return self._node_coords_in_element
-
-    def trace(self):
-        return HexmeshPiecewiseConstantBasis.Trace(self)
-
-
 class HexmeshTripolynomialSpaceTopology(HexmeshSpaceTopology):
     def __init__(self, mesh: Hexmesh, shape: CubeTripolynomialShapeFunctions):
         super().__init__(mesh, shape, need_hex_edge_indices=shape.ORDER >= 2, need_hex_face_indices=shape.ORDER >= 2)
@@ -156,7 +122,7 @@ class HexmeshTripolynomialSpaceTopology(HexmeshSpaceTopology):
         self.element_node_index = self._make_element_node_index()
 
     def node_count(self) -> int:
-        ORDER = self._shape.ORDER
+        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
@@ -182,7 +148,7 @@ class HexmeshTripolynomialSpaceTopology(HexmeshSpaceTopology):
         return rot_i * size + rot_j
 
     def _make_element_node_index(self):
-        ORDER = self._shape.ORDER
+        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)
@@ -194,18 +160,19 @@ class HexmeshTripolynomialSpaceTopology(HexmeshSpaceTopology):
             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)
+            node_type, type_instance, type_index = self.shape.node_type_and_type_index(node_index_in_elt)
 
             if node_type == CubeTripolynomialShapeFunctions.VERTEX:
-                return geo_arg.hex_vertex_indices[element_index, _CUBE_VERTEX_INDICES[type_instance]]
+                return geo_arg.hex_vertex_indices[element_index, _CUBE_TO_HEX_VERTEX[type_instance]]
 
             offset = topo_arg.vertex_count
 
             if node_type == CubeTripolynomialShapeFunctions.EDGE:
-                edge_index = topo_arg.hex_edge_indices[element_index, type_instance]
+                hex_edge = _CUBE_TO_HEX_EDGE[type_instance]
+                edge_index = topo_arg.hex_edge_indices[element_index, hex_edge]
 
-                v0 = geo_arg.hex_vertex_indices[element_index, EDGE_VERTEX_INDICES[type_instance, 0]]
-                v1 = geo_arg.hex_vertex_indices[element_index, EDGE_VERTEX_INDICES[type_instance, 1]]
+                v0 = geo_arg.hex_vertex_indices[element_index, EDGE_VERTEX_INDICES[hex_edge, 0]]
+                v1 = geo_arg.hex_vertex_indices[element_index, EDGE_VERTEX_INDICES[hex_edge, 1]]
 
                 if v0 > v1:
                     type_index = ORDER - 1 - type_index
@@ -232,52 +199,21 @@ class HexmeshTripolynomialSpaceTopology(HexmeshSpaceTopology):
         return element_node_index
 
 
-class HexmeshTripolynomialBasisSpace(HexmeshBasisSpace):
-    def __init__(
-        self,
-        mesh: Hexmesh,
-        degree: int,
-        family: Polynomial,
-    ):
-        if family is None:
-            family = Polynomial.LOBATTO_GAUSS_LEGENDRE
-
-        if not is_closed(family):
-            raise ValueError("A closed polynomial family is required to define a continuous function space")
-
-        shape = CubeTripolynomialShapeFunctions(degree, family=family)
-        topology = forward_base_topology(HexmeshTripolynomialSpaceTopology, mesh, shape)
-
-        super().__init__(topology, shape)
-
-
-class HexmeshDGTripolynomialBasisSpace(HexmeshBasisSpace):
+class HexmeshSerendipitySpaceTopology(HexmeshSpaceTopology):
     def __init__(
         self,
-        mesh: Hexmesh,
-        degree: int,
-        family: Polynomial,
+        grid: Hexmesh,
+        shape: CubeSerendipityShapeFunctions,
     ):
-        if family is None:
-            family = Polynomial.LOBATTO_GAUSS_LEGENDRE
-
-        shape = CubeTripolynomialShapeFunctions(degree, family=family)
-        topology = HexmeshDiscontinuousSpaceTopology(mesh, shape)
-
-        super().__init__(topology, shape)
-
-
-class HexmeshSerendipitySpaceTopology(HexmeshSpaceTopology):
-    def __init__(self, grid: Hexmesh, shape: CubeSerendipityShapeFunctions):
         super().__init__(grid, shape, need_hex_edge_indices=True, need_hex_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.geometry.edge_count()
+        return self.geometry.vertex_count() + (self.shape.ORDER - 1) * self.geometry.edge_count()
 
     def _make_element_node_index(self):
-        ORDER = self._shape.ORDER
+        ORDER = self.shape.ORDER
 
         @cache.dynamic_func(suffix=self.name)
         def element_node_index(
@@ -286,17 +222,18 @@ class HexmeshSerendipitySpaceTopology(HexmeshSpaceTopology):
             element_index: ElementIndex,
             node_index_in_elt: int,
         ):
-            node_type, type_index = self._shape.node_type_and_type_index(node_index_in_elt)
+            node_type, type_index = self.shape.node_type_and_type_index(node_index_in_elt)
 
             if node_type == CubeSerendipityShapeFunctions.VERTEX:
-                return cell_arg.hex_vertex_indices[element_index, _CUBE_VERTEX_INDICES[type_index]]
+                return cell_arg.hex_vertex_indices[element_index, _CUBE_TO_HEX_VERTEX[type_index]]
 
             type_instance, index_in_edge = CubeSerendipityShapeFunctions._cube_edge_index(node_type, type_index)
+            hex_edge = _CUBE_TO_HEX_EDGE[type_instance]
 
-            edge_index = topo_arg.hex_edge_indices[element_index, type_instance]
+            edge_index = topo_arg.hex_edge_indices[element_index, hex_edge]
 
-            v0 = cell_arg.hex_vertex_indices[element_index, EDGE_VERTEX_INDICES[type_instance, 0]]
-            v1 = cell_arg.hex_vertex_indices[element_index, EDGE_VERTEX_INDICES[type_instance, 1]]
+            v0 = cell_arg.hex_vertex_indices[element_index, EDGE_VERTEX_INDICES[hex_edge, 0]]
+            v1 = cell_arg.hex_vertex_indices[element_index, EDGE_VERTEX_INDICES[hex_edge, 1]]
 
             if v0 > v1:
                 index_in_edge = ORDER - 1 - index_in_edge
@@ -306,45 +243,11 @@ class HexmeshSerendipitySpaceTopology(HexmeshSpaceTopology):
         return element_node_index
 
 
-class HexmeshSerendipityBasisSpace(HexmeshBasisSpace):
-    def __init__(
-        self,
-        mesh: Hexmesh,
-        degree: int,
-        family: Polynomial,
-    ):
-        if family is None:
-            family = Polynomial.LOBATTO_GAUSS_LEGENDRE
-
-        shape = CubeSerendipityShapeFunctions(degree, family=family)
-        topology = forward_base_topology(HexmeshSerendipitySpaceTopology, mesh, shape=shape)
-
-        super().__init__(topology=topology, shape=shape)
+def make_hexmesh_space_topology(mesh: Hexmesh, shape: ShapeFunction):
+    if isinstance(shape, CubeSerendipityShapeFunctions):
+        return forward_base_topology(HexmeshSerendipitySpaceTopology, mesh, shape)
 
+    if isinstance(shape, CubeTripolynomialShapeFunctions):
+        return forward_base_topology(HexmeshTripolynomialSpaceTopology, mesh, shape)
 
-class HexmeshDGSerendipityBasisSpace(HexmeshBasisSpace):
-    def __init__(
-        self,
-        mesh: Hexmesh,
-        degree: int,
-        family: Polynomial,
-    ):
-        if family is None:
-            family = Polynomial.LOBATTO_GAUSS_LEGENDRE
-
-        shape = CubeSerendipityShapeFunctions(degree, family=family)
-        topology = HexmeshDiscontinuousSpaceTopology(mesh, shape=shape)
-
-        super().__init__(topology=topology, shape=shape)
-
-
-class HexmeshPolynomialBasisSpace(HexmeshBasisSpace):
-    def __init__(
-        self,
-        mesh: Hexmesh,
-        degree: int,
-    ):
-        shape = CubeNonConformingPolynomialShapeFunctions(degree)
-        topology = HexmeshDiscontinuousSpaceTopology(mesh, shape)
-
-        super().__init__(topology, shape)
+    raise ValueError(f"Unsupported shape function {shape.name}")

warp/fem/space/nanogrid_function_space.py

@@ -0,0 +1,202 @@
+import warp as wp
+from warp.fem import cache
+from warp.fem.geometry import Nanogrid
+from warp.fem.geometry.nanogrid import _add_axis_flag
+from warp.fem.polynomial import is_closed
+from warp.fem.types import ElementIndex
+
+from .shape import (
+    CubeSerendipityShapeFunctions,
+    CubeTripolynomialShapeFunctions,
+    ShapeFunction,
+)
+from .topology import SpaceTopology, forward_base_topology
+
+
+@wp.struct
+class NanogridTopologyArg:
+    vertex_grid: wp.uint64
+    face_grid: wp.uint64
+    edge_grid: wp.uint64
+
+    vertex_count: int
+    edge_count: int
+    face_count: int
+
+
+class NanogridSpaceTopology(SpaceTopology):
+    TopologyArg = NanogridTopologyArg
+
+    def __init__(
+        self,
+        grid: Nanogrid,
+        shape: ShapeFunction,
+        need_edge_indices: bool = True,
+        need_face_indices: bool = True,
+    ):
+        if not is_closed(shape.family):
+            raise ValueError("A closed polynomial family is required to define a continuous function space")
+
+        super().__init__(grid, shape.NODES_PER_ELEMENT)
+        self._grid = grid
+        self._shape = shape
+
+        if need_edge_indices:
+            self._edge_count = self._grid.edge_count()
+        else:
+            self._edge_count = 0
+
+        self._vertex_grid = grid._node_grid
+        self._face_grid = grid._face_grid
+        self._edge_grid = grid._edge_grid
+
+    @cache.cached_arg_value
+    def topo_arg_value(self, device):
+        arg = NanogridTopologyArg()
+
+        arg.vertex_grid = self._vertex_grid.id
+        arg.face_grid = self._face_grid.id
+        arg.edge_grid = -1 if self._edge_grid is None else self._edge_grid.id
+
+        arg.vertex_count = self._grid.vertex_count()
+        arg.face_count = self._grid.side_count()
+        arg.edge_count = self._edge_count
+        return arg
+
+
+@wp.func
+def _cell_vertex_coord(cell_ijk: wp.vec3i, n: int):
+    return cell_ijk + wp.vec3i((n & 4) >> 2, (n & 2) >> 1, n & 1)
+
+
+@wp.func
+def _cell_edge_coord(cell_ijk: wp.vec3i, axis: int, offset: int):
+    e_ijk = cell_ijk
+    e_ijk[(axis + 1) % 3] += offset >> 1
+    e_ijk[(axis + 2) % 3] += offset & 1
+    return _add_axis_flag(e_ijk, axis)
+
+
+@wp.func
+def _cell_face_coord(cell_ijk: wp.vec3i, axis: int, offset: int):
+    f_ijk = cell_ijk
+    f_ijk[axis] += offset
+    return _add_axis_flag(f_ijk, axis)
+
+
+class NanogridTripolynomialSpaceTopology(NanogridSpaceTopology):
+    def __init__(self, grid: Nanogrid, 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._grid.edge_count() * INTERIOR_NODES_PER_EDGE
+            + self._grid.side_count() * INTERIOR_NODES_PER_FACE
+            + self._grid.cell_count() * INTERIOR_NODES_PER_CELL
+        )
+
+    def _make_element_node_index(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)
+
+        @cache.dynamic_func(suffix=self.name)
+        def element_node_index(
+            geo_arg: Nanogrid.CellArg,
+            topo_arg: NanogridTopologyArg,
+            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)
+
+            ijk = geo_arg.cell_ijk[element_index]
+
+            if node_type == CubeTripolynomialShapeFunctions.VERTEX:
+                n_ijk = _cell_vertex_coord(ijk, type_instance)
+                return wp.volume_lookup_index(topo_arg.vertex_grid, n_ijk[0], n_ijk[1], n_ijk[2])
+
+            offset = topo_arg.vertex_count
+
+            if node_type == CubeTripolynomialShapeFunctions.EDGE:
+                axis = type_instance >> 2
+                node_offset = type_instance & 3
+
+                n_ijk = _cell_edge_coord(ijk, 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
+
+            offset += INTERIOR_NODES_PER_EDGE * topo_arg.edge_count
+
+            if node_type == CubeTripolynomialShapeFunctions.FACE:
+                axis = type_instance >> 1
+                node_offset = type_instance & 1
+
+                n_ijk = _cell_face_coord(ijk, 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
+
+            offset += INTERIOR_NODES_PER_FACE * topo_arg.face_count
+
+            return offset + INTERIOR_NODES_PER_CELL * element_index + type_index
+
+        return element_node_index
+
+
+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.geometry.edge_count()
+
+    def _make_element_node_index(self):
+        ORDER = self._shape.ORDER
+
+        @cache.dynamic_func(suffix=self.name)
+        def element_node_index(
+            cell_arg: Nanogrid.CellArg,
+            topo_arg: NanogridSpaceTopology.TopologyArg,
+            element_index: ElementIndex,
+            node_index_in_elt: int,
+        ):
+            node_type, type_index = self._shape.node_type_and_type_index(node_index_in_elt)
+
+            ijk = cell_arg.cell_ijk[element_index]
+
+            if node_type == CubeSerendipityShapeFunctions.VERTEX:
+                n_ijk = _cell_vertex_coord(ijk, 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
+
+            n_ijk = _cell_edge_coord(ijk, 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
+
+        return element_node_index
+
+
+def make_nanogrid_space_topology(grid: Nanogrid, shape: ShapeFunction):
+    if isinstance(shape, CubeSerendipityShapeFunctions):
+        return forward_base_topology(NanogridSerendipitySpaceTopology, grid, shape)
+
+    if isinstance(shape, CubeTripolynomialShapeFunctions):
+        return forward_base_topology(NanogridTripolynomialSpaceTopology, grid, shape)
+
+    raise ValueError(f"Unsupported shape function {shape.name}")

warp/fem/space/quadmesh_2d_function_space.py

@@ -1,18 +1,15 @@
 import warp as wp
 from warp.fem import cache
 from warp.fem.geometry import Quadmesh2D
-from warp.fem.polynomial import Polynomial, is_closed
-from warp.fem.types import Coords, ElementIndex
+from warp.fem.polynomial import is_closed
+from warp.fem.types import ElementIndex
 
-from .basis_space import ShapeBasisSpace, TraceBasisSpace
 from .shape import (
-    ConstantShapeFunction,
     ShapeFunction,
     SquareBipolynomialShapeFunctions,
-    SquareNonConformingPolynomialShapeFunctions,
     SquareSerendipityShapeFunctions,
 )
-from .topology import DiscontinuousSpaceTopologyMixin, SpaceTopology, forward_base_topology
+from .topology import SpaceTopology, forward_base_topology
 
 
 @wp.struct
@@ -28,6 +25,9 @@ class Quadmesh2DSpaceTopology(SpaceTopology):
     TopologyArg = Quadmesh2DTopologyArg
 
     def __init__(self, mesh: Quadmesh2D, shape: ShapeFunction):
+        if not is_closed(shape.family):
+            raise ValueError("A closed polynomial family is required to define a continuous function space")
+
         super().__init__(mesh, shape.NODES_PER_ELEMENT)
         self._mesh = mesh
         self._shape = shape
@@ -107,44 +107,6 @@ class Quadmesh2DSpaceTopology(SpaceTopology):
             quad_edge_indices[q1, t1_edge] = e
 
 
-class Quadmesh2DDiscontinuousSpaceTopology(
-    DiscontinuousSpaceTopologyMixin,
-    SpaceTopology,
-):
-    def __init__(self, mesh: Quadmesh2D, shape: ShapeFunction):
-        super().__init__(mesh, shape.NODES_PER_ELEMENT)
-
-
-class Quadmesh2DBasisSpace(ShapeBasisSpace):
-    def __init__(self, topology: Quadmesh2DSpaceTopology, shape: ShapeFunction):
-        super().__init__(topology, shape)
-
-        self._mesh: Quadmesh2D = topology.geometry
-
-
-class Quadmesh2DPiecewiseConstantBasis(Quadmesh2DBasisSpace):
-    def __init__(self, mesh: Quadmesh2D):
-        shape = ConstantShapeFunction(mesh.reference_cell(), space_dimension=2)
-        topology = Quadmesh2DDiscontinuousSpaceTopology(mesh, shape)
-        super().__init__(shape=shape, topology=topology)
-
-    class Trace(TraceBasisSpace):
-        @wp.func
-        def _node_coords_in_element(
-            side_arg: Quadmesh2D.SideArg,
-            basis_arg: Quadmesh2DBasisSpace.BasisArg,
-            element_index: ElementIndex,
-            node_index_in_element: int,
-        ):
-            return Coords(0.5, 0.0, 0.0)
-
-        def make_node_coords_in_element(self):
-            return self._node_coords_in_element
-
-    def trace(self):
-        return Quadmesh2DPiecewiseConstantBasis.Trace(self)
-
-
 class Quadmesh2DBipolynomialSpaceTopology(Quadmesh2DSpaceTopology):
     def __init__(self, mesh: Quadmesh2D, shape: SquareBipolynomialShapeFunctions):
         super().__init__(mesh, shape)
@@ -236,41 +198,6 @@ class Quadmesh2DBipolynomialSpaceTopology(Quadmesh2DSpaceTopology):
         return element_node_index
 
 
-class Quadmesh2DBipolynomialBasisSpace(Quadmesh2DBasisSpace):
-    def __init__(
-        self,
-        mesh: Quadmesh2D,
-        degree: int,
-        family: Polynomial,
-    ):
-        if family is None:
-            family = Polynomial.LOBATTO_GAUSS_LEGENDRE
-
-        if not is_closed(family):
-            raise ValueError("A closed polynomial family is required to define a continuous function space")
-
-        shape = SquareBipolynomialShapeFunctions(degree, family=family)
-        topology = forward_base_topology(Quadmesh2DBipolynomialSpaceTopology, mesh, shape)
-
-        super().__init__(topology, shape)
-
-
-class Quadmesh2DDGBipolynomialBasisSpace(Quadmesh2DBasisSpace):
-    def __init__(
-        self,
-        mesh: Quadmesh2D,
-        degree: int,
-        family: Polynomial,
-    ):
-        if family is None:
-            family = Polynomial.LOBATTO_GAUSS_LEGENDRE
-
-        shape = SquareBipolynomialShapeFunctions(degree, family=family)
-        topology = Quadmesh2DDiscontinuousSpaceTopology(mesh, shape)
-
-        super().__init__(topology, shape)
-
-
 class Quadmesh2DSerendipitySpaceTopology(Quadmesh2DSpaceTopology):
     def __init__(self, grid: Quadmesh2D, shape: SquareSerendipityShapeFunctions):
         super().__init__(grid, shape)
@@ -324,45 +251,11 @@ class Quadmesh2DSerendipitySpaceTopology(Quadmesh2DSpaceTopology):
         return element_node_index
 
 
-class Quadmesh2DSerendipityBasisSpace(Quadmesh2DBasisSpace):
-    def __init__(
-        self,
-        mesh: Quadmesh2D,
-        degree: int,
-        family: Polynomial,
-    ):
-        if family is None:
-            family = Polynomial.LOBATTO_GAUSS_LEGENDRE
+def make_quadmesh_2d_space_topology(mesh: Quadmesh2D, shape: ShapeFunction):
+    if isinstance(shape, SquareSerendipityShapeFunctions):
+        return forward_base_topology(Quadmesh2DSerendipitySpaceTopology, mesh, shape)
 
-        shape = SquareSerendipityShapeFunctions(degree, family=family)
-        topology = forward_base_topology(Quadmesh2DSerendipitySpaceTopology, mesh, shape=shape)
-
-        super().__init__(topology=topology, shape=shape)
-
-
-class Quadmesh2DDGSerendipityBasisSpace(Quadmesh2DBasisSpace):
-    def __init__(
-        self,
-        mesh: Quadmesh2D,
-        degree: int,
-        family: Polynomial,
-    ):
-        if family is None:
-            family = Polynomial.LOBATTO_GAUSS_LEGENDRE
-
-        shape = SquareSerendipityShapeFunctions(degree, family=family)
-        topology = Quadmesh2DDiscontinuousSpaceTopology(mesh, shape=shape)
-
-        super().__init__(topology=topology, shape=shape)
-
-
-class Quadmesh2DPolynomialBasisSpace(Quadmesh2DBasisSpace):
-    def __init__(
-        self,
-        mesh: Quadmesh2D,
-        degree: int,
-    ):
-        shape = SquareNonConformingPolynomialShapeFunctions(degree)
-        topology = Quadmesh2DDiscontinuousSpaceTopology(mesh, shape)
+    if isinstance(shape, SquareBipolynomialShapeFunctions):
+        return forward_base_topology(Quadmesh2DBipolynomialSpaceTopology, mesh, shape)
 
-        super().__init__(topology, shape)
+    raise ValueError(f"Unsupported shape function {shape.name}")

warp/fem/space/restriction.py

@@ -38,7 +38,10 @@ class SpaceRestriction:
 
         NODES_PER_ELEMENT = self.space_topology.NODES_PER_ELEMENT
 
-        @cache.dynamic_kernel(suffix=f"{self.domain.name}_{self.space_topology.name}_{self.space_partition.name}")
+        @cache.dynamic_kernel(
+            suffix=f"{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,