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

warp/fem/space/quadmesh_2d_function_space.py

@@ -0,0 +1,369 @@
+import warp as wp
+
+from warp.fem.types import ElementIndex, Coords
+from warp.fem.polynomial import Polynomial, is_closed
+from warp.fem.geometry import Quadmesh2D
+from warp.fem import cache
+
+from .topology import SpaceTopology, DiscontinuousSpaceTopologyMixin, forward_base_topology
+from .basis_space import ShapeBasisSpace, TraceBasisSpace
+
+from .shape import ShapeFunction, ConstantShapeFunction
+from .shape import (
+    SquareBipolynomialShapeFunctions,
+    SquareSerendipityShapeFunctions,
+    SquareNonConformingPolynomialShapeFunctions,
+)
+
+
+@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
+
+
+class Quadmesh2DSpaceTopology(SpaceTopology):
+    TopologyArg = Quadmesh2DTopologyArg
+
+    def __init__(self, mesh: Quadmesh2D, shape: ShapeFunction):
+        super().__init__(mesh, shape.NODES_PER_ELEMENT)
+        self._mesh = mesh
+        self._shape = shape
+
+        self._compute_quad_edge_indices()
+
+    @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()
+        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=Quadmesh2DSpaceTopology._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 = Quadmesh2DSpaceTopology._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 = Quadmesh2DSpaceTopology._find_edge_index_in_quad(edge_vtx, t1_vtx)
+            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)
+
+        self.element_node_index = self._make_element_node_index()
+
+    def node_count(self) -> int:
+        ORDER = self._shape.ORDER
+        INTERIOR_NODES_PER_SIDE = max(0, ORDER - 1)
+        INTERIOR_NODES_PER_CELL = INTERIOR_NODES_PER_SIDE**2
+
+        return (
+            self._mesh.vertex_count()
+            + self._mesh.side_count() * INTERIOR_NODES_PER_SIDE
+            + self._mesh.cell_count() * INTERIOR_NODES_PER_CELL
+        )
+
+    def _make_element_node_index(self):
+        ORDER = self._shape.ORDER
+        INTERIOR_NODES_PER_SIDE = wp.constant(max(0, ORDER - 1))
+        INTERIOR_NODES_PER_CELL = wp.constant(INTERIOR_NODES_PER_SIDE**2)
+
+        @cache.dynamic_func(suffix=self.name)
+        def element_node_index(
+            geo_arg: Quadmesh2D.CellArg,
+            topo_arg: Quadmesh2DTopologyArg,
+            element_index: ElementIndex,
+            node_index_in_elt: int,
+        ):
+            node_i = node_index_in_elt // (ORDER + 1)
+            node_j = node_index_in_elt - (ORDER + 1) * node_i
+
+            # Vertices
+            if node_i == 0:
+                if node_j == 0:
+                    return geo_arg.quad_vertex_indices[element_index, 0]
+                elif node_j == ORDER:
+                    return geo_arg.quad_vertex_indices[element_index, 3]
+
+                # 3-0 edge
+                side_index = topo_arg.quad_edge_indices[element_index, 3]
+                local_vs = geo_arg.quad_vertex_indices[element_index, 3]
+                global_vs = topo_arg.edge_vertex_indices[side_index][0]
+                index_in_side = wp.select(local_vs == global_vs, ORDER - node_j, node_j) - 1
+
+                return topo_arg.vertex_count + (ORDER - 1) * side_index + index_in_side
+
+            elif node_i == ORDER:
+                if node_j == 0:
+                    return geo_arg.quad_vertex_indices[element_index, 1]
+                elif node_j == ORDER:
+                    return geo_arg.quad_vertex_indices[element_index, 2]
+
+                # 1-2 edge
+                side_index = topo_arg.quad_edge_indices[element_index, 1]
+                local_vs = geo_arg.quad_vertex_indices[element_index, 1]
+                global_vs = topo_arg.edge_vertex_indices[side_index][0]
+                index_in_side = wp.select(local_vs == global_vs, ORDER - node_j, node_j) - 1
+
+                return topo_arg.vertex_count + (ORDER - 1) * side_index + index_in_side
+
+            if node_j == 0:
+                # 0-1 edge
+                side_index = topo_arg.quad_edge_indices[element_index, 0]
+                local_vs = geo_arg.quad_vertex_indices[element_index, 0]
+                global_vs = topo_arg.edge_vertex_indices[side_index][0]
+                index_in_side = wp.select(local_vs == global_vs, node_i, ORDER - node_i) - 1
+
+                return topo_arg.vertex_count + (ORDER - 1) * side_index + index_in_side
+
+            elif node_j == ORDER:
+                # 2-3 edge
+                side_index = topo_arg.quad_edge_indices[element_index, 2]
+                local_vs = geo_arg.quad_vertex_indices[element_index, 2]
+                global_vs = topo_arg.edge_vertex_indices[side_index][0]
+                index_in_side = wp.select(local_vs == global_vs, node_i, ORDER - node_i) - 1
+
+                return topo_arg.vertex_count + (ORDER - 1) * side_index + index_in_side
+
+            return (
+                topo_arg.vertex_count
+                + topo_arg.edge_count * INTERIOR_NODES_PER_SIDE
+                + element_index * INTERIOR_NODES_PER_CELL
+                + (node_i - 1) * INTERIOR_NODES_PER_SIDE
+                + node_j
+                - 1
+            )
+
+        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)
+
+        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.side_count()
+
+    def _make_element_node_index(self):
+        ORDER = self._shape.ORDER
+
+        SHAPE_TO_QUAD_IDX = wp.constant(wp.vec4i([0, 3, 1, 2]))
+
+        @cache.dynamic_func(suffix=self.name)
+        def element_node_index(
+            cell_arg: Quadmesh2D.CellArg,
+            topo_arg: Quadmesh2DSpaceTopology.TopologyArg,
+            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 == SquareSerendipityShapeFunctions.VERTEX:
+                return cell_arg.quad_vertex_indices[element_index, SHAPE_TO_QUAD_IDX[type_index]]
+
+            side_offset, index_in_side = SquareSerendipityShapeFunctions.side_offset_and_index(type_index)
+
+            if node_type == SquareSerendipityShapeFunctions.EDGE_X:
+                if side_offset == 0:
+                    side_start = 0
+                else:
+                    side_start = 2
+                    index_in_side = ORDER - 2 - index_in_side
+            else:
+                if side_offset == 0:
+                    side_start = 3
+                    index_in_side = ORDER - 2 - index_in_side
+                else:
+                    side_start = 1
+
+            side_index = topo_arg.quad_edge_indices[element_index, side_start]
+            local_vs = cell_arg.quad_vertex_indices[element_index, side_start]
+            global_vs = topo_arg.edge_vertex_indices[side_index][0]
+            if local_vs != global_vs:
+                # Flip indexing direction
+                index_in_side = ORDER - 2 - index_in_side
+
+            return topo_arg.vertex_count + (ORDER - 1) * side_index + index_in_side
+
+        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
+
+        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)
+
+        super().__init__(topology, shape)

warp/fem/space/restriction.py

@@ -3,6 +3,7 @@ import warp as wp
 from warp.fem.domain import GeometryDomain
 from warp.fem.types import NodeElementIndex
 from warp.fem.utils import compress_node_indices
+from warp.fem.cache import cached_arg_value, borrow_temporary, borrow_temporary_like, TemporaryStore
 
 from .function_space import FunctionSpace
 from .partition import SpacePartition
@@ -11,97 +12,107 @@ wp.set_module_options({"enable_backward": False})
 
 
 class SpaceRestriction:
-    """Restriction of a space to a given GeometryDomain"""
+    """Restriction of a space partition to a given GeometryDomain"""
 
     def __init__(
         self,
-        space: FunctionSpace,
-        domain: GeometryDomain,
         space_partition: SpacePartition,
+        domain: GeometryDomain,
         device=None,
+        temporary_store: TemporaryStore = None,
     ):
-        if domain.dimension() == space.DIMENSION - 1:
-            space = space.trace()
+        space_topology = space_partition.space_topology
+
+        if domain.dimension == space_topology.dimension - 1:
+            space_topology = space_topology.trace()
 
-        if domain.dimension() != space.DIMENSION:
+        if domain.dimension != space_topology.dimension:
             raise ValueError("Incompatible space and domain dimensions")
 
-        self.space = space
         self.space_partition = space_partition
+        self.space_topology = space_topology
         self.domain = domain
 
-        self._compute_node_element_indices(device=device)
+        self._compute_node_element_indices(device=device, temporary_store=temporary_store)
 
-    def _compute_node_element_indices(self, device):
+    def _compute_node_element_indices(self, device, temporary_store: TemporaryStore):
         from warp.fem import cache
 
-        NODES_PER_ELEMENT = self.space.NODES_PER_ELEMENT
+        NODES_PER_ELEMENT = self.space_topology.NODES_PER_ELEMENT
 
-        def fill_element_node_indices_fn(
+        @cache.dynamic_kernel(suffix=f"{self.domain.name}_{self.space_topology.name}_{self.space_partition.name}")
+        def fill_element_node_indices(
+            element_arg: self.domain.ElementArg,
             domain_index_arg: self.domain.ElementIndexArg,
-            space_arg: self.space.SpaceArg,
+            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)
             for n in range(NODES_PER_ELEMENT):
-                space_nidx = self.space.element_node_index(space_arg, element_index, n)
+                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
 
-        fill_element_node_indices = cache.get_kernel(
-            fill_element_node_indices_fn,
-            suffix=f"{self.domain.name}_{self.space.name}_{self.space_partition.name}",
-        )
-
-        element_node_indices = wp.empty(
+        element_node_indices = borrow_temporary(
+            temporary_store,
             shape=(self.domain.element_count(), NODES_PER_ELEMENT),
             dtype=int,
             device=device,
         )
         wp.launch(
-            dim=element_node_indices.shape[0],
+            dim=element_node_indices.array.shape[0],
             kernel=fill_element_node_indices,
             inputs=[
+                self.domain.element_arg_value(device),
                 self.domain.element_index_arg_value(device),
-                self.space.space_arg_value(device),
+                self.space_topology.topo_arg_value(device),
                 self.space_partition.partition_arg_value(device),
-                element_node_indices,
+                element_node_indices.array,
             ],
             device=device,
         )
 
         # Build compressed map from node to element indices
-        flattened_node_indices = element_node_indices.reshape(element_node_indices.size)
+        flattened_node_indices = element_node_indices.array.flatten()
         (
             self._dof_partition_element_offsets,
             node_array_indices,
             self._node_count,
             self._dof_partition_indices,
-        ) = compress_node_indices(self.space_partition.node_count(), flattened_node_indices)
+        ) = compress_node_indices(
+            self.space_partition.node_count(), flattened_node_indices, temporary_store=temporary_store
+        )
 
         # Extract element index and index in element
-        self._dof_element_indices = wp.empty_like(flattened_node_indices)
-        self._dof_indices_in_element = wp.empty_like(flattened_node_indices)
+        self._dof_element_indices = borrow_temporary_like(flattened_node_indices, temporary_store)
+        self._dof_indices_in_element = borrow_temporary_like(flattened_node_indices, temporary_store)
         wp.launch(
             kernel=SpaceRestriction._split_vertex_element_index,
             dim=flattened_node_indices.shape,
-            inputs=[NODES_PER_ELEMENT, node_array_indices, self._dof_element_indices, self._dof_indices_in_element],
+            inputs=[
+                NODES_PER_ELEMENT,
+                node_array_indices.array,
+                self._dof_element_indices.array,
+                self._dof_indices_in_element.array,
+            ],
             device=flattened_node_indices.device,
         )
 
+        node_array_indices.release()
+
     def node_count(self):
         return self._node_count
 
     def partition_element_offsets(self):
-        return self._dof_partition_element_offsets
+        return self._dof_partition_element_offsets.array
 
     def node_partition_indices(self):
-        return self._dof_partition_indices
+        return self._dof_partition_indices.array
 
     def total_node_element_count(self):
-        return self._dof_element_indices.size
+        return self._dof_element_indices.array.size
 
     @wp.struct
     class NodeArg:
@@ -110,12 +121,13 @@ class SpaceRestriction:
         dof_partition_indices: wp.array(dtype=int)
         dof_indices_in_element: wp.array(dtype=int)
 
+    @cached_arg_value
     def node_arg(self, device):
         arg = SpaceRestriction.NodeArg()
-        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)
+        arg.dof_element_offsets = self._dof_partition_element_offsets.array.to(device)
+        arg.dof_element_indices = self._dof_element_indices.array.to(device)
+        arg.dof_partition_indices = self._dof_partition_indices.array.to(device)
+        arg.dof_indices_in_element = self._dof_indices_in_element.array.to(device)
         return arg
 
     @wp.func

warp/fem/space/shape/__init__.py

@@ -0,0 +1,15 @@
+from .shape_function import ShapeFunction, ConstantShapeFunction
+
+from .triangle_shape_function import Triangle2DPolynomialShapeFunctions, Triangle2DNonConformingPolynomialShapeFunctions
+from .tet_shape_function import TetrahedronPolynomialShapeFunctions, TetrahedronNonConformingPolynomialShapeFunctions
+
+from .square_shape_function import (
+    SquareBipolynomialShapeFunctions,
+    SquareSerendipityShapeFunctions,
+    SquareNonConformingPolynomialShapeFunctions,
+)
+from .cube_shape_function import (
+    CubeSerendipityShapeFunctions,
+    CubeTripolynomialShapeFunctions,
+    CubeNonConformingPolynomialShapeFunctions,
+)