warp-lang 1.0.0b2__py3-none-win_amd64.whl → 1.0.0b6__py3-none-win_amd64.whl

warp/fem/space/trimesh_2d_function_space.py

@@ -1,134 +1,66 @@
 import warp as wp
-import numpy as np
 
-
-from warp.fem.types import ElementIndex, Coords, OUTSIDE, vec2i, vec3i
+from warp.fem.types import ElementIndex, Coords
 from warp.fem.geometry import Trimesh2D
+from warp.fem import cache
 
-from .dof_mapper import DofMapper
-from .nodal_function_space import NodalFunctionSpace, NodalFunctionSpaceTrace
+from .topology import SpaceTopology, DiscontinuousSpaceTopologyMixin, forward_base_topology
+from .basis_space import ShapeBasisSpace, TraceBasisSpace
 
+from .shape import ShapeFunction, ConstantShapeFunction
+from .shape import Triangle2DPolynomialShapeFunctions, Triangle2DNonConformingPolynomialShapeFunctions
 
-class Trimesh2DFunctionSpace(NodalFunctionSpace):
-    DIMENSION = wp.constant(2)
 
-    @wp.struct
-    class SpaceArg:
-        geo_arg: Trimesh2D.SideArg
+@wp.struct
+class Trimesh2DTopologyArg:
+    edge_vertex_indices: wp.array(dtype=wp.vec2i)
+    tri_edge_indices: wp.array2d(dtype=int)
 
-        reference_transforms: wp.array(dtype=wp.mat22f)
-        tri_edge_indices: wp.array2d(dtype=int)
+    vertex_count: int
+    edge_count: int
 
-        vertex_count: int
-        edge_count: int
 
-    def __init__(self, mesh: Trimesh2D, dtype: type = float, dof_mapper: DofMapper = None):
-        super().__init__(dtype, dof_mapper)
-        self._mesh = mesh
+class Trimesh2DSpaceTopology(SpaceTopology):
+    TopologyArg = Trimesh2DTopologyArg
 
-        self._reference_transforms: wp.array = None
-        self._tri_edge_indices: wp.array = None
+    def __init__(self, mesh: Trimesh2D, shape: ShapeFunction):
+        super().__init__(mesh, shape.NODES_PER_ELEMENT)
+        self._mesh = mesh
+        self._shape = shape
 
-        self._compute_reference_transforms()
         self._compute_tri_edge_indices()
 
-    @property
-    def geometry(self) -> Trimesh2D:
-        return self._mesh
-
-    def space_arg_value(self, device):
-        arg = self.SpaceArg()
-        arg.geo_arg = self.geometry.side_arg_value(device)
-        arg.reference_transforms = self._reference_transforms.to(device)
+    @cache.cached_arg_value
+    def topo_arg_value(self, device):
+        arg = Trimesh2DTopologyArg()
         arg.tri_edge_indices = self._tri_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()
         return arg
 
-    class Trace(NodalFunctionSpaceTrace):
-        def __init__(self, space: NodalFunctionSpace):
-            super().__init__(space)
-            self.ORDER = space.ORDER
-
-    @wp.func
-    def _inner_cell_index(args: SpaceArg, side_index: ElementIndex):
-        return Trimesh2D.side_inner_cell_index(args.geo_arg, side_index)
-
-    @wp.func
-    def _outer_cell_index(args: SpaceArg, side_index: ElementIndex):
-        return Trimesh2D.side_outer_cell_index(args.geo_arg, side_index)
-
-    @wp.func
-    def _inner_cell_coords(args: SpaceArg, side_index: ElementIndex, side_coords: Coords):
-        tri_index = Trimesh2D.side_inner_cell_index(args.geo_arg, side_index)
-        return Trimesh2D.edge_to_tri_coords(args.geo_arg, side_index, tri_index, side_coords)
-
-    @wp.func
-    def _outer_cell_coords(args: SpaceArg, side_index: ElementIndex, side_coords: Coords):
-        tri_index = Trimesh2D.side_outer_cell_index(args.geo_arg, side_index)
-        return Trimesh2D.edge_to_tri_coords(args.geo_arg, side_index, tri_index, side_coords)
-
-    @wp.func
-    def _cell_to_side_coords(
-        args: SpaceArg,
-        side_index: ElementIndex,
-        element_index: ElementIndex,
-        element_coords: Coords,
-    ):
-        return Trimesh2D.tri_to_edge_coords(args.geo_arg, side_index, element_index, element_coords)
-
-    def _compute_reference_transforms(self):
-        self._reference_transforms = wp.empty(
-            dtype=wp.mat22f, device=self._mesh.positions.device, shape=(self._mesh.cell_count())
-        )
-
-        wp.launch(
-            kernel=Trimesh2DFunctionSpace._compute_reference_transforms_kernel,
-            dim=self._reference_transforms.shape,
-            device=self._reference_transforms.device,
-            inputs=[self._mesh.tri_vertex_indices, self._mesh.positions, self._reference_transforms],
-        )
-
     def _compute_tri_edge_indices(self):
         self._tri_edge_indices = wp.empty(
             dtype=int, device=self._mesh.tri_vertex_indices.device, shape=(self._mesh.cell_count(), 3)
         )
 
         wp.launch(
-            kernel=Trimesh2DFunctionSpace._compute_tri_edge_indices_kernel,
-            dim=self._mesh._edge_tri_indices.shape,
+            kernel=Trimesh2DSpaceTopology._compute_tri_edge_indices_kernel,
+            dim=self._mesh.edge_tri_indices.shape,
             device=self._mesh.tri_vertex_indices.device,
             inputs=[
-                self._mesh._edge_tri_indices,
-                self._mesh._edge_vertex_indices,
+                self._mesh.edge_tri_indices,
+                self._mesh.edge_vertex_indices,
                 self._mesh.tri_vertex_indices,
                 self._tri_edge_indices,
             ],
         )
 
-    @wp.kernel
-    def _compute_reference_transforms_kernel(
-        tri_vertex_indices: wp.array2d(dtype=int),
-        positions: wp.array(dtype=wp.vec2f),
-        transforms: wp.array(dtype=wp.mat22f),
-    ):
-        t = wp.tid()
-
-        p0 = positions[tri_vertex_indices[t, 0]]
-        p1 = positions[tri_vertex_indices[t, 1]]
-        p2 = positions[tri_vertex_indices[t, 2]]
-
-        e1 = p1 - p0
-        e2 = p2 - p0
-
-        mat = wp.mat22(e1, e2)
-        transforms[t] = wp.transpose(wp.inverse(mat))
-
     @wp.func
     def _find_edge_index_in_tri(
-        edge_vtx: vec2i,
-        tri_vtx: vec3i,
+        edge_vtx: wp.vec2i,
+        tri_vtx: wp.vec3i,
     ):
         for k in range(2):
             if (edge_vtx[0] == tri_vtx[k] and edge_vtx[1] == tri_vtx[k + 1]) or (
@@ -139,8 +71,8 @@ class Trimesh2DFunctionSpace(NodalFunctionSpace):
 
     @wp.kernel
     def _compute_tri_edge_indices_kernel(
-        edge_tri_indices: wp.array(dtype=vec2i),
-        edge_vertex_indices: wp.array(dtype=vec2i),
+        edge_tri_indices: wp.array(dtype=wp.vec2i),
+        edge_vertex_indices: wp.array(dtype=wp.vec2i),
         tri_vertex_indices: wp.array2d(dtype=int),
         tri_edge_indices: wp.array2d(dtype=int),
     ):
@@ -150,710 +82,140 @@ class Trimesh2DFunctionSpace(NodalFunctionSpace):
         edge_tris = edge_tri_indices[e]
 
         t0 = edge_tris[0]
-        t0_vtx = vec3i(tri_vertex_indices[t0, 0], tri_vertex_indices[t0, 1], tri_vertex_indices[t0, 2])
-        t0_edge = Trimesh2DFunctionSpace._find_edge_index_in_tri(edge_vtx, t0_vtx)
+        t0_vtx = wp.vec3i(tri_vertex_indices[t0, 0], tri_vertex_indices[t0, 1], tri_vertex_indices[t0, 2])
+        t0_edge = Trimesh2DSpaceTopology._find_edge_index_in_tri(edge_vtx, t0_vtx)
         tri_edge_indices[t0, t0_edge] = e
 
         t1 = edge_tris[1]
         if t1 != t0:
-            t1_vtx = vec3i(tri_vertex_indices[t1, 0], tri_vertex_indices[t1, 1], tri_vertex_indices[t1, 2])
-            t1_edge = Trimesh2DFunctionSpace._find_edge_index_in_tri(edge_vtx, t1_vtx)
+            t1_vtx = wp.vec3i(tri_vertex_indices[t1, 0], tri_vertex_indices[t1, 1], tri_vertex_indices[t1, 2])
+            t1_edge = Trimesh2DSpaceTopology._find_edge_index_in_tri(edge_vtx, t1_vtx)
             tri_edge_indices[t1, t1_edge] = e
 
 
-class Trimesh2DPiecewiseConstantSpace(Trimesh2DFunctionSpace):
-    ORDER = wp.constant(0)
-    NODES_PER_ELEMENT = wp.constant(1)
-
-    def __init__(self, grid: Trimesh2D, dtype: type = float, dof_mapper: DofMapper = None):
-        super().__init__(grid, dtype, dof_mapper)
-
-        self.element_outer_weight = self.element_inner_weight
-        self.element_outer_weight_gradient = self.element_inner_weight_gradient
-
-    def node_count(self) -> int:
-        return self._mesh.cell_count()
-
-    def node_positions(self):
-        vtx_pos = self._mesh.positions.numpy()
-        tri_vtx = self._mesh.tri_vertex_indices.numpy()
-
-        tri_pos = vtx_pos[tri_vtx]
-        centers = tri_pos.sum(axis=1) / 3.0
-
-        return centers[:,0], centers[:,1]
-
-    @wp.func
-    def element_node_index(
-        args: Trimesh2DFunctionSpace.SpaceArg,
-        element_index: ElementIndex,
-        node_index_in_elt: int,
-    ):
-        return element_index
-
-    @wp.func
-    def node_coords_in_element(
-        args: Trimesh2DFunctionSpace.SpaceArg,
-        element_index: ElementIndex,
-        node_index_in_elt: int,
-    ):
-        if node_index_in_elt == 0:
-            return Coords(1.0 / 3.0, 1.0 / 3.0, 1.0 / 3.0)
-
-        return Coords(OUTSIDE)
-
-    @wp.func
-    def node_quadrature_weight(
-        args: Trimesh2DFunctionSpace.SpaceArg,
-        element_index: ElementIndex,
-        node_index_in_elt: int,
-    ):
-        return 1.0
-
-    @wp.func
-    def element_inner_weight(
-        args: Trimesh2DFunctionSpace.SpaceArg,
-        element_index: ElementIndex,
-        coords: Coords,
-        node_index_in_elt: int,
-    ):
-        if node_index_in_elt == 0:
-            return 1.0
-        return 0.0
-
-    @wp.func
-    def element_inner_weight_gradient(
-        args: Trimesh2DFunctionSpace.SpaceArg,
-        element_index: ElementIndex,
-        coords: Coords,
-        node_index_in_elt: int,
-    ):
-        return wp.vec2(0.0)
+class Trimesh2DDiscontinuousSpaceTopology(
+    DiscontinuousSpaceTopologyMixin,
+    SpaceTopology,
+):
+    def __init__(self, mesh: Trimesh2D, shape: ShapeFunction):
+        super().__init__(mesh, shape.NODES_PER_ELEMENT)
 
-    class Trace(Trimesh2DFunctionSpace.Trace):
-        NODES_PER_ELEMENT = wp.constant(2)
-        ORDER = wp.constant(0)
 
-        def __init__(self, space: "Trimesh2DPiecewiseConstantSpace"):
-            super().__init__(space)
+class Trimesh2DBasisSpace(ShapeBasisSpace):
+    def __init__(self, topology: Trimesh2DSpaceTopology, shape: ShapeFunction):
+        super().__init__(topology, shape)
 
-            self.element_node_index = self._make_element_node_index(space)
+        self._mesh: Trimesh2D = topology.geometry
 
-            self.element_inner_weight = self._make_element_inner_weight(space)
-            self.element_inner_weight_gradient = self._make_element_inner_weight_gradient(space)
 
-            self.element_outer_weight = self._make_element_outer_weight(space)
-            self.element_outer_weight_gradient = self._make_element_outer_weight_gradient(space)
+class Trimesh2DPiecewiseConstantBasis(Trimesh2DBasisSpace):
+    def __init__(self, mesh: Trimesh2D):
+        shape = ConstantShapeFunction(mesh.reference_cell(), space_dimension=2)
+        topology = Trimesh2DDiscontinuousSpaceTopology(mesh, shape)
+        super().__init__(shape=shape, topology=topology)
 
+    class Trace(TraceBasisSpace):
         @wp.func
-        def node_coords_in_element(
-            args: Trimesh2DFunctionSpace.SpaceArg,
+        def _node_coords_in_element(
+            side_arg: Trimesh2D.SideArg,
+            basis_arg: Trimesh2DBasisSpace.BasisArg,
             element_index: ElementIndex,
             node_index_in_element: int,
         ):
-            if node_index_in_element == 0:
-                return Coords(0.5, 0.0, 0.0)
-            elif node_index_in_element == 1:
-                return Coords(0.5, 0.0, 0.0)
-
-            return Coords(OUTSIDE)
+            return Coords(0.5, 0.0, 0.0)
 
-        @wp.func
-        def node_quadrature_weight(
-            args: Trimesh2DFunctionSpace.SpaceArg,
-            element_index: ElementIndex,
-            node_index_in_elt: int,
-        ):
-            return 1.0
+        def make_node_coords_in_element(self):
+            return self._node_coords_in_element
 
     def trace(self):
-        return Trimesh2DPiecewiseConstantSpace.Trace(self)
-
-
-def _triangle_node_index(tx: int, ty: int, degree: int):
-    VERTEX_NODE_COUNT = 3
-    SIDE_INTERIOR_NODE_COUNT = degree - 1
-
-    # Index in similar order to e.g. VTK
-    # First vertices, then edge (counterclokwise) then interior points (recursively)
-
-    if tx == 0:
-        if ty == 0:
-            return 0
-        elif ty == degree:
-            return 2
-        else:
-            edge_index = 2
-            return VERTEX_NODE_COUNT + SIDE_INTERIOR_NODE_COUNT * edge_index + (SIDE_INTERIOR_NODE_COUNT - ty)
-    elif ty == 0:
-        if tx == degree:
-            return 1
-        else:
-            edge_index = 0
-            return VERTEX_NODE_COUNT + SIDE_INTERIOR_NODE_COUNT * edge_index + tx - 1
-    elif tx + ty == degree:
-        edge_index = 1
-        return VERTEX_NODE_COUNT + SIDE_INTERIOR_NODE_COUNT * edge_index + ty - 1
-
-    vertex_edge_node_count = 3 * degree
-    return vertex_edge_node_count + _triangle_node_index(tx - 1, ty - 1, degree - 2)
-
-
-class Trimesh2DPolynomialShapeFunctions:
-    INVALID = wp.constant(-1)
-    VERTEX = wp.constant(0)
-    EDGE = wp.constant(1)
-    INTERIOR = wp.constant(2)
-
-    def __init__(self, degree: int):
-        self.ORDER = wp.constant(degree)
+        return Trimesh2DPiecewiseConstantBasis.Trace(self)
 
-        self.NODES_PER_ELEMENT = wp.constant((degree + 1) * (degree + 2) // 2)
-        self.NODES_PER_SIDE = wp.constant(degree + 1)
-
-        triangle_coords = np.empty((self.NODES_PER_ELEMENT, 2), dtype=int)
-
-        for tx in range(degree + 1):
-            for ty in range(degree + 1 - tx):
-                index = _triangle_node_index(tx, ty, degree)
-                triangle_coords[index] = [tx, ty]
-
-        CoordTypeVec = wp.mat(dtype=int, shape=(self.NODES_PER_ELEMENT, 2))
-        self.NODE_TRIANGLE_COORDS = wp.constant(CoordTypeVec(triangle_coords))
-
-        self.node_type_and_type_index = self._get_node_type_and_type_index()
-        self._node_triangle_coordinates = self._get_node_triangle_coordinates()
-
-    @property
-    def name(self) -> str:
-        return f"{self.ORDER}"
-
-    def _get_node_triangle_coordinates(self):
-        NODE_TRIANGLE_COORDS = self.NODE_TRIANGLE_COORDS
-
-        def node_triangle_coordinates(
-            node_index_in_elt: int,
-        ):
-            return vec2i(NODE_TRIANGLE_COORDS[node_index_in_elt, 0], NODE_TRIANGLE_COORDS[node_index_in_elt, 1])
-
-        from warp.fem import cache
-
-        return cache.get_func(node_triangle_coordinates, self.name)
-
-    def _get_node_type_and_type_index(self):
-        ORDER = self.ORDER
-        NODES_PER_ELEMENT = self.NODES_PER_ELEMENT
-
-        def node_type_and_index(
-            node_index_in_elt: int,
-        ):
-            if node_index_in_elt < 0 or node_index_in_elt >= NODES_PER_ELEMENT:
-                return Trimesh2DPolynomialShapeFunctions.INVALID, Trimesh2DPolynomialShapeFunctions.INVALID
-
-            if node_index_in_elt < 3:
-                return Trimesh2DPolynomialShapeFunctions.VERTEX, node_index_in_elt
-
-            if node_index_in_elt < 3 * ORDER:
-                return Trimesh2DPolynomialShapeFunctions.EDGE, (node_index_in_elt - 3)
-
-            return Trimesh2DPolynomialShapeFunctions.INTERIOR, (node_index_in_elt - 3 * ORDER)
-
-        from warp.fem import cache
-
-        return cache.get_func(node_type_and_index, self.name)
-
-    def make_node_coords_in_element(self):
-        ORDER = self.ORDER
-
-        def node_coords_in_element(
-            args: Trimesh2DFunctionSpace.SpaceArg,
-            element_index: ElementIndex,
-            node_index_in_elt: int,
-        ):
-            tri_coords = self._node_triangle_coordinates(node_index_in_elt)
-            cx = float(tri_coords[0]) / float(ORDER)
-            cy = float(tri_coords[1]) / float(ORDER)
-            return Coords(1.0 - cx - cy, cx, cy)
-
-        from warp.fem import cache
-
-        return cache.get_func(node_coords_in_element, self.name)
-
-    def make_node_quadrature_weight(self):
-        ORDER = self.ORDER
-
-        def node_uniform_quadrature_weight(
-            args: Trimesh2DFunctionSpace.SpaceArg,
-            element_index: ElementIndex,
-            node_index_in_elt: int,
-        ):
-            node_type, type_index = self.node_type_and_type_index(node_index_in_elt)
-
-            base_weight = 1.0 / float(3 * ORDER * ORDER)
-            if node_type == Trimesh2DPolynomialShapeFunctions.VERTEX:
-                return base_weight
-            if node_type == Trimesh2DPolynomialShapeFunctions.EDGE:
-                return 2.0 * base_weight
-            return 4.0 * base_weight
-
-        def node_linear_quadrature_weight(
-            args: Trimesh2DFunctionSpace.SpaceArg,
-            element_index: ElementIndex,
-            node_index_in_elt: int,
-        ):
-            return 1.0 / 3.0
-
-        from warp.fem import cache
-
-        if ORDER == 1:
-            return cache.get_func(node_linear_quadrature_weight, self.name)
-        return cache.get_func(node_uniform_quadrature_weight, self.name)
-
-    def make_trace_node_quadrature_weight(self):
-        ORDER = self.ORDER
-        NODES_PER_ELEMENT = self.NODES_PER_ELEMENT
-
-        def trace_uniform_node_quadrature_weight(
-            args: Trimesh2DFunctionSpace.SpaceArg,
-            element_index: ElementIndex,
-            node_index_in_elt: int,
-        ):
-            if node_index_in_elt >= NODES_PER_ELEMENT:
-                node_index_in_cell = node_index_in_elt - NODES_PER_ELEMENT
-            else:
-                node_index_in_cell = node_index_in_elt
-
-            # We're either on a side interior or at a vertex
-            node_type, type_index = self.node_type_and_type_index(node_index_in_cell)
-
-            base_weight = 1.0 / float(ORDER)
-            return wp.select(node_type == Trimesh2DPolynomialShapeFunctions.VERTEX, base_weight, 0.5 * base_weight)
-
-        def trace_linear_node_quadrature_weight(
-            args: Trimesh2DFunctionSpace.SpaceArg,
-            element_index: ElementIndex,
-            node_index_in_elt: int,
-        ):
-            return 0.5
-
-        from warp.fem import cache
-
-        if ORDER == 1:
-            return cache.get_func(trace_linear_node_quadrature_weight, self.name)
-
-        return cache.get_func(trace_uniform_node_quadrature_weight, self.name)
-
-    def make_element_inner_weight(self):
-        ORDER = self.ORDER
-
-        def element_inner_weight_linear(
-            args: Trimesh2DFunctionSpace.SpaceArg,
-            element_index: ElementIndex,
-            coords: Coords,
-            node_index_in_elt: int,
-        ):
-            if node_index_in_elt < 0 or node_index_in_elt >= 3:
-                return 0.0
-
-            return coords[node_index_in_elt]
-
-        def element_inner_weight_quadratic(
-            args: Trimesh2DFunctionSpace.SpaceArg,
-            element_index: ElementIndex,
-            coords: Coords,
-            node_index_in_elt: int,
-        ):
-            node_type, type_index = self.node_type_and_type_index(node_index_in_elt)
 
-            if node_type == Trimesh2DPolynomialShapeFunctions.VERTEX:
-                # Vertex
-                return coords[type_index] * (2.0 * coords[type_index] - 1.0)
+class Trimesh2DPolynomialSpaceTopology(Trimesh2DSpaceTopology):
+    def __init__(self, mesh: Trimesh2D, shape: Triangle2DPolynomialShapeFunctions):
+        super().__init__(mesh, shape)
 
-            elif node_type == Trimesh2DPolynomialShapeFunctions.EDGE:
-                # Edge
-                c1 = type_index
-                c2 = (type_index + 1) % 3
-                return 4.0 * coords[c1] * coords[c2]
-
-            return 0.0
-
-        def element_inner_weight_cubic(
-            args: Trimesh2DFunctionSpace.SpaceArg,
-            element_index: ElementIndex,
-            coords: Coords,
-            node_index_in_elt: int,
-        ):
-            node_type, type_index = self.node_type_and_type_index(node_index_in_elt)
-
-            if node_type == Trimesh2DPolynomialShapeFunctions.VERTEX:
-                # Vertex
-                return 0.5 * coords[type_index] * (3.0 * coords[type_index] - 1.0) * (3.0 * coords[type_index] - 2.0)
-
-            elif node_type == Trimesh2DPolynomialShapeFunctions.EDGE:
-                # Edge
-                edge = type_index // 2
-                k = type_index - 2 * edge
-                c1 = (edge + k) % 3
-                c2 = (edge + 1 - k) % 3
-
-                return 4.5 * coords[c1] * coords[c2] * (3.0 * coords[c1] - 1.0)
-
-            elif node_type == Trimesh2DPolynomialShapeFunctions.INTERIOR:
-                # Interior
-                return 27.0 * coords[0] * coords[1] * coords[2]
-
-            return 0.0
-
-        from warp.fem import cache
-
-        if ORDER == 1:
-            return cache.get_func(element_inner_weight_linear, self.name)
-        elif ORDER == 2:
-            return cache.get_func(element_inner_weight_quadratic, self.name)
-        elif ORDER == 3:
-            return cache.get_func(element_inner_weight_cubic, self.name)
-
-        return None
-
-    def make_element_inner_weight_gradient(self):
-        ORDER = self.ORDER
-
-        def element_inner_weight_gradient_linear(
-            args: Trimesh2DFunctionSpace.SpaceArg,
-            element_index: ElementIndex,
-            coords: Coords,
-            node_index_in_elt: int,
-        ):
-            if node_index_in_elt < 0 or node_index_in_elt >= 3:
-                return wp.vec2(0.0)
-
-            dw_dc = wp.vec3(0.0)
-            dw_dc[node_index_in_elt] = 1.0
-
-            dw_du = wp.vec2(dw_dc[1] - dw_dc[0], dw_dc[2] - dw_dc[0])
-            return args.reference_transforms[element_index] * dw_du
-
-        def element_inner_weight_gradient_quadratic(
-            args: Trimesh2DFunctionSpace.SpaceArg,
-            element_index: ElementIndex,
-            coords: Coords,
-            node_index_in_elt: int,
-        ):
-            node_type, type_index = self.node_type_and_type_index(node_index_in_elt)
-
-            dw_dc = wp.vec3(0.0)
-
-            if node_type == Trimesh2DPolynomialShapeFunctions.VERTEX:
-                # Vertex
-                dw_dc[type_index] = 4.0 * coords[type_index] - 1.0
-
-            elif node_type == Trimesh2DPolynomialShapeFunctions.EDGE:
-                # Edge
-                c1 = type_index
-                c2 = (type_index + 1) % 3
-                dw_dc[c1] = 4.0 * coords[c2]
-                dw_dc[c2] = 4.0 * coords[c1]
-
-            dw_du = wp.vec2(dw_dc[1] - dw_dc[0], dw_dc[2] - dw_dc[0])
-            return args.reference_transforms[element_index] * dw_du
-
-        def element_inner_weight_gradient_cubic(
-            args: Trimesh2DFunctionSpace.SpaceArg,
-            element_index: ElementIndex,
-            coords: Coords,
-            node_index_in_elt: int,
-        ):
-            node_type, type_index = self.node_type_and_type_index(node_index_in_elt)
-
-            dw_dc = wp.vec3(0.0)
-
-            if node_type == Trimesh2DPolynomialShapeFunctions.VERTEX:
-                # Vertex
-                dw_dc[type_index] = (
-                    0.5 * 27.0 * coords[type_index] * coords[type_index] - 9.0 * coords[type_index] + 1.0
-                )
-
-            elif node_type == Trimesh2DPolynomialShapeFunctions.EDGE:
-                # Edge
-                edge = type_index // 2
-                k = type_index - 2 * edge
-                c1 = (edge + k) % 3
-                c2 = (edge + 1 - k) % 3
-
-                dw_dc[c1] = 4.5 * coords[c2] * (6.0 * coords[c1] - 1.0)
-                dw_dc[c2] = 4.5 * coords[c1] * (3.0 * coords[c1] - 1.0)
-
-            elif node_type == Trimesh2DPolynomialShapeFunctions.INTERIOR:
-                # Interior
-                dw_dc = wp.vec3(
-                    27.0 * coords[1] * coords[2], 27.0 * coords[2] * coords[0], 27.0 * coords[0] * coords[1]
-                )
-
-            dw_du = wp.vec2(dw_dc[1] - dw_dc[0], dw_dc[2] - dw_dc[0])
-            return args.reference_transforms[element_index] * dw_du
-
-        from warp.fem import cache
-
-        if ORDER == 1:
-            return cache.get_func(element_inner_weight_gradient_linear, self.name)
-        elif ORDER == 2:
-            return cache.get_func(element_inner_weight_gradient_quadratic, self.name)
-        elif ORDER == 3:
-            return cache.get_func(element_inner_weight_gradient_cubic, self.name)
-
-        return None
-
-    @staticmethod
-    def node_positions(space):
-        if space.ORDER == 1:
-            return np.transpose(space._mesh.positions.numpy())
-
-        NODES_PER_ELEMENT = space.NODES_PER_ELEMENT
-
-        def fill_node_positions_fn(
-            space_arg: space.SpaceArg,
-            node_positions: wp.array(dtype=wp.vec2),
-        ):
-            element_index = wp.tid()
-            tri_idx = space_arg.geo_arg.tri_vertex_indices[element_index]
-            p0 = space_arg.geo_arg.positions[tri_idx[0]]
-            p1 = space_arg.geo_arg.positions[tri_idx[1]]
-            p2 = space_arg.geo_arg.positions[tri_idx[2]]
-
-            for n in range(NODES_PER_ELEMENT):
-                node_index = space.element_node_index(space_arg, element_index, n)
-                coords = space.node_coords_in_element(space_arg, element_index, n)
-
-                pos = coords[0] * p0 + coords[1] * p1 + coords[2] * p2
-
-                node_positions[node_index] = pos
-
-        from warp.fem import cache
-
-        fill_node_positions = cache.get_kernel(
-            fill_node_positions_fn,
-            suffix=space.name,
-        )
-
-        device = space._mesh.tri_vertex_indices.device
-        node_positions = wp.empty(
-            shape=space.node_count(),
-            dtype=wp.vec2,
-            device=device,
-        )
-        wp.launch(
-            dim=space._mesh.cell_count(),
-            kernel=fill_node_positions,
-            inputs=[
-                space.space_arg_value(device),
-                node_positions,
-            ],
-            device=device,
-        )
-
-        return np.transpose(node_positions.numpy())
-
-    @staticmethod
-    def node_triangulation(space):
-        if space.ORDER == 1:
-            return space._mesh.tri_vertex_indices.numpy()
-
-        NODES_PER_ELEMENT = space.NODES_PER_ELEMENT
-
-        def fill_element_node_indices_fn(
-            space_arg: space.SpaceArg,
-            element_node_indices: wp.array2d(dtype=int),
-        ):
-            element_index = wp.tid()
-            for n in range(NODES_PER_ELEMENT):
-                element_node_indices[element_index, n] = space.element_node_index(space_arg, element_index, n)
-
-        from warp.fem import cache
+        self.element_node_index = self._make_element_node_index()
 
-        fill_element_node_indices = cache.get_kernel(
-            fill_element_node_indices_fn,
-            suffix=space.name,
-        )
+    def node_count(self) -> int:
+        INTERIOR_NODES_PER_SIDE = max(0, self._shape.ORDER - 1)
+        INTERIOR_NODES_PER_CELL = max(0, self._shape.ORDER - 2) * max(0, self._shape.ORDER - 1) // 2
 
-        device = space._mesh.tri_vertex_indices.device
-        element_node_indices = wp.empty(
-            shape=(space._mesh.cell_count(), NODES_PER_ELEMENT),
-            dtype=int,
-            device=device,
-        )
-        wp.launch(
-            dim=element_node_indices.shape[0],
-            kernel=fill_element_node_indices,
-            inputs=[
-                space.space_arg_value(device),
-                element_node_indices,
-            ],
-            device=device,
+        return (
+            self._mesh.vertex_count()
+            + self._mesh.side_count() * INTERIOR_NODES_PER_SIDE
+            + self._mesh.cell_count() * INTERIOR_NODES_PER_CELL
         )
 
-        element_node_indices = element_node_indices.numpy()
-        if space.ORDER == 2:
-            element_triangles = [[0, 3, 5], [3, 1, 4], [2, 5, 4], [3, 4, 5]]
-        elif space.ORDER == 3:
-            element_triangles = [
-                [0, 3, 8],
-                [3, 4, 9],
-                [4, 1, 5],
-                [8, 3, 9],
-                [4, 5, 9],
-                [8, 9, 7],
-                [9, 5, 6],
-                [6, 7, 9],
-                [7, 6, 2],
-            ]
-
-        tri_indices = element_node_indices[:, element_triangles].reshape(-1, 3)
-        return tri_indices
-
-
-class Trimesh2DPolynomialSpace(Trimesh2DFunctionSpace):
-
-    def __init__(self, grid: Trimesh2D, degree: int, dtype: type = float, dof_mapper: DofMapper = None):
-        super().__init__(grid, dtype, dof_mapper)
-
-        self._shape = Trimesh2DPolynomialShapeFunctions(degree)
-
-        self.ORDER = self._shape.ORDER
-        self.NODES_PER_ELEMENT = self._shape.NODES_PER_ELEMENT
-
-        self.element_node_index = self._make_element_node_index()
-        self.node_coords_in_element = self._shape.make_node_coords_in_element()
-        self.node_quadrature_weight = self._shape.make_node_quadrature_weight()
-        self.element_inner_weight = self._shape.make_element_inner_weight()
-        self.element_inner_weight_gradient = self._shape.make_element_inner_weight_gradient()
-
-        self.element_outer_weight = self.element_inner_weight
-        self.element_outer_weight_gradient = self.element_inner_weight_gradient
-
     def _make_element_node_index(self):
-        INTERIOR_NODES_PER_SIDE = wp.constant(max(0, self.ORDER - 1))
-        INTERIOR_NODES_PER_CELL = wp.constant(max(0, self.ORDER - 2) * max(0, self.ORDER - 1) // 2)
+        INTERIOR_NODES_PER_SIDE = wp.constant(max(0, self._shape.ORDER - 1))
+        INTERIOR_NODES_PER_CELL = wp.constant(max(0, self._shape.ORDER - 2) * max(0, self._shape.ORDER - 1) // 2)
 
+        @cache.dynamic_func(suffix=self.name)
         def element_node_index(
-            args: Trimesh2DFunctionSpace.SpaceArg,
+            geo_arg: Trimesh2D.CellArg,
+            topo_arg: Trimesh2DTopologyArg,
             element_index: ElementIndex,
             node_index_in_elt: int,
         ):
             node_type, type_index = self._shape.node_type_and_type_index(node_index_in_elt)
 
-            if node_type == Trimesh2DPolynomialShapeFunctions.VERTEX:
-                return args.geo_arg.tri_vertex_indices[element_index][type_index]
+            if node_type == Triangle2DPolynomialShapeFunctions.VERTEX:
+                return geo_arg.tri_vertex_indices[element_index][type_index]
 
-            global_offset = args.vertex_count
+            global_offset = topo_arg.vertex_count
 
-            if node_type == Trimesh2DPolynomialShapeFunctions.EDGE:
+            if node_type == Triangle2DPolynomialShapeFunctions.EDGE:
                 edge = type_index // INTERIOR_NODES_PER_SIDE
                 edge_node = type_index - INTERIOR_NODES_PER_SIDE * edge
 
-                global_edge_index = args.tri_edge_indices[element_index][edge]
+                global_edge_index = topo_arg.tri_edge_indices[element_index][edge]
 
                 if (
-                    args.geo_arg.edge_vertex_indices[global_edge_index][0]
-                    != args.geo_arg.tri_vertex_indices[element_index][edge]
+                    topo_arg.edge_vertex_indices[global_edge_index][0]
+                    != geo_arg.tri_vertex_indices[element_index][edge]
                 ):
                     edge_node = INTERIOR_NODES_PER_SIDE - 1 - edge_node
 
                 return global_offset + INTERIOR_NODES_PER_SIDE * global_edge_index + edge_node
 
-            global_offset += INTERIOR_NODES_PER_SIDE * args.edge_count
+            global_offset += INTERIOR_NODES_PER_SIDE * topo_arg.edge_count
             return global_offset + INTERIOR_NODES_PER_CELL * element_index + type_index
 
-        from warp.fem import cache
-
-        return cache.get_func(element_node_index, self.name)
-
-    def node_count(self) -> int:
-        INTERIOR_NODES_PER_SIDE = wp.constant(max(0, self.ORDER - 1))
-        INTERIOR_NODES_PER_CELL = wp.constant(max(0, self.ORDER - 2) * max(0, self.ORDER - 1) // 2)
-
-        return (
-            self._mesh.vertex_count()
-            + self._mesh.side_count() * INTERIOR_NODES_PER_SIDE
-            + self._mesh.cell_count() * INTERIOR_NODES_PER_CELL
-        )
-
-    def node_positions(self):
-        return Trimesh2DPolynomialShapeFunctions.node_positions(self)
-
-    def node_triangulation(self):
-        return Trimesh2DPolynomialShapeFunctions.node_triangulation(self)
+        return element_node_index
 
-    class Trace(Trimesh2DFunctionSpace.Trace):
-        NODES_PER_ELEMENT = wp.constant(2)
-        ORDER = wp.constant(0)
 
-        def __init__(self, space: "Trimesh2DPolynomialSpace"):
-            super().__init__(space)
-
-            self.element_node_index = self._make_element_node_index(space)
-            self.node_coords_in_element = self._make_node_coords_in_element(space)
-            self.node_quadrature_weight = space._shape.make_trace_node_quadrature_weight()
-
-            self.element_inner_weight = self._make_element_inner_weight(space)
-            self.element_inner_weight_gradient = self._make_element_inner_weight_gradient(space)
-
-            self.element_outer_weight = self._make_element_outer_weight(space)
-            self.element_outer_weight_gradient = self._make_element_outer_weight_gradient(space)
-
-    def trace(self):
-        return Trimesh2DPolynomialSpace.Trace(self)
-
-
-class Trimesh2DDGPolynomialSpace(Trimesh2DFunctionSpace):
+class Trimesh2DPolynomialBasisSpace(Trimesh2DBasisSpace):
     def __init__(
         self,
         mesh: Trimesh2D,
         degree: int,
-        dtype: type = float,
-        dof_mapper: DofMapper = None,
     ):
-        super().__init__(mesh, dtype, dof_mapper)
-
-        self._shape = Trimesh2DPolynomialShapeFunctions(degree)
+        shape = Triangle2DPolynomialShapeFunctions(degree)
+        topology = forward_base_topology(Trimesh2DPolynomialSpaceTopology, mesh, shape)
 
-        self.ORDER = self._shape.ORDER
-        self.NODES_PER_ELEMENT = self._shape.NODES_PER_ELEMENT
+        super().__init__(topology, shape)
 
-        self.element_node_index = self._make_element_node_index()
-        self.node_coords_in_element = self._shape.make_node_coords_in_element()
-        self.node_quadrature_weight = self._shape.make_node_quadrature_weight()
-        self.element_inner_weight = self._shape.make_element_inner_weight()
-        self.element_inner_weight_gradient = self._shape.make_element_inner_weight_gradient()
-
-        self.element_outer_weight = self.element_inner_weight
-        self.element_outer_weight_gradient = self.element_inner_weight_gradient
-
-    def node_count(self) -> int:
-        return self._mesh.cell_count() * self.NODES_PER_ELEMENT
-
-    def node_positions(self):
-        return Trimesh2DPolynomialShapeFunctions.node_positions(self)
-
-    def node_triangulation(self):
-        return Trimesh2DPolynomialShapeFunctions.node_triangulation(self)
-
-    def _make_element_node_index(self):
-        NODES_PER_ELEMENT = self.NODES_PER_ELEMENT
 
-        def element_node_index(
-            args: Trimesh2DFunctionSpace.SpaceArg,
-            element_index: ElementIndex,
-            node_index_in_elt: int,
-        ):
-            return element_index * NODES_PER_ELEMENT + node_index_in_elt
+class Trimesh2DDGPolynomialBasisSpace(Trimesh2DBasisSpace):
+    def __init__(
+        self,
+        mesh: Trimesh2D,
+        degree: int,
+    ):
+        shape = Triangle2DPolynomialShapeFunctions(degree)
+        topology = Trimesh2DDiscontinuousSpaceTopology(mesh, shape)
 
-        from warp.fem import cache
+        super().__init__(topology, shape)
 
-        return cache.get_func(element_node_index, f"{self.name}_{self.ORDER}")
 
-    class Trace(Trimesh2DPolynomialSpace.Trace):
-        def __init__(self, space: "Trimesh2DDGPolynomialSpace"):
-            super().__init__(space)
+class Trimesh2DNonConformingPolynomialBasisSpace(Trimesh2DBasisSpace):
+    def __init__(
+        self,
+        mesh: Trimesh2D,
+        degree: int,
+    ):
+        shape = Triangle2DNonConformingPolynomialShapeFunctions(degree)
+        topology = Trimesh2DDiscontinuousSpaceTopology(mesh, shape)
 
-    def trace(self):
-        return Trimesh2DDGPolynomialSpace.Trace(self)
+        super().__init__(topology, shape)