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

warp/fem/space/restriction.py

@@ -1,160 +1,161 @@
-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
-
-wp.set_module_options({"enable_backward": False})
-
-
-class SpaceRestriction:
-    """Restriction of a space partition to a given GeometryDomain"""
-
-    def __init__(
-        self,
-        space_partition: SpacePartition,
-        domain: GeometryDomain,
-        device=None,
-        temporary_store: TemporaryStore = None,
-    ):
-        space_topology = space_partition.space_topology
-
-        if domain.dimension == space_topology.dimension - 1:
-            space_topology = space_topology.trace()
-
-        if domain.dimension != space_topology.dimension:
-            raise ValueError("Incompatible space and domain dimensions")
-
-        self.space_partition = space_partition
-        self.space_topology = space_topology
-        self.domain = domain
-
-        self._compute_node_element_indices(device=device, temporary_store=temporary_store)
-
-    def _compute_node_element_indices(self, device, temporary_store: TemporaryStore):
-        from warp.fem import cache
-
-        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}")
-        def fill_element_node_indices(
-            element_arg: self.domain.ElementArg,
-            domain_index_arg: self.domain.ElementIndexArg,
-            topo_arg: self.space_topology.TopologyArg,
-            partition_arg: self.space_partition.PartitionArg,
-            element_node_indices: wp.array2d(dtype=int),
-        ):
-            domain_element_index = wp.tid()
-            element_index = self.domain.element_index(domain_index_arg, domain_element_index)
-            for n in range(NODES_PER_ELEMENT):
-                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
-
-        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.array.shape[0],
-            kernel=fill_element_node_indices,
-            inputs=[
-                self.domain.element_arg_value(device),
-                self.domain.element_index_arg_value(device),
-                self.space_topology.topo_arg_value(device),
-                self.space_partition.partition_arg_value(device),
-                element_node_indices.array,
-            ],
-            device=device,
-        )
-
-        # Build compressed map from node to element indices
-        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, temporary_store=temporary_store
-        )
-
-        # Extract element index and index in element
-        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.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.array
-
-    def node_partition_indices(self):
-        return self._dof_partition_indices.array
-
-    def total_node_element_count(self):
-        return self._dof_element_indices.array.size
-
-    @wp.struct
-    class NodeArg:
-        dof_element_offsets: wp.array(dtype=int)
-        dof_element_indices: wp.array(dtype=int)
-        dof_partition_indices: wp.array(dtype=int)
-        dof_indices_in_element: wp.array(dtype=int)
-
-    @cached_arg_value
-    def node_arg(self, device):
-        arg = SpaceRestriction.NodeArg()
-        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
-    def node_partition_index(args: NodeArg, node_index: int):
-        return args.dof_partition_indices[node_index]
-
-    @wp.func
-    def node_element_count(args: NodeArg, node_index: int):
-        partition_node_index = SpaceRestriction.node_partition_index(args, node_index)
-        return args.dof_element_offsets[partition_node_index + 1] - args.dof_element_offsets[partition_node_index]
-
-    @wp.func
-    def node_element_index(args: NodeArg, node_index: int, element_index: int):
-        partition_node_index = SpaceRestriction.node_partition_index(args, node_index)
-        offset = args.dof_element_offsets[partition_node_index] + element_index
-        domain_element_index = args.dof_element_indices[offset]
-        index_in_element = args.dof_indices_in_element[offset]
-        return NodeElementIndex(domain_element_index, index_in_element)
-
-    @wp.kernel
-    def _split_vertex_element_index(
-        vertex_per_element: int,
-        sorted_indices: wp.array(dtype=int),
-        vertex_element_index: wp.array(dtype=int),
-        vertex_index_in_element: wp.array(dtype=int),
-    ):
-        idx = sorted_indices[wp.tid()]
-        element_index = idx // vertex_per_element
-        vertex_element_index[wp.tid()] = element_index
-        vertex_index_in_element[wp.tid()] = idx - vertex_per_element * element_index
+import warp as wp
+from warp.fem.cache import TemporaryStore, borrow_temporary, borrow_temporary_like, cached_arg_value
+from warp.fem.domain import GeometryDomain
+from warp.fem.types import NodeElementIndex
+from warp.fem.utils import compress_node_indices
+
+from .partition import SpacePartition
+
+wp.set_module_options({"enable_backward": False})
+
+
+class SpaceRestriction:
+    """Restriction of a space partition to a given GeometryDomain"""
+
+    def __init__(
+        self,
+        space_partition: SpacePartition,
+        domain: GeometryDomain,
+        device=None,
+        temporary_store: TemporaryStore = None,
+    ):
+        space_topology = space_partition.space_topology
+
+        if domain.dimension == space_topology.dimension - 1:
+            space_topology = space_topology.trace()
+
+        if domain.dimension != space_topology.dimension:
+            raise ValueError("Incompatible space and domain dimensions")
+
+        self.space_partition = space_partition
+        self.space_topology = space_topology
+        self.domain = domain
+
+        self._compute_node_element_indices(device=device, temporary_store=temporary_store)
+
+    def _compute_node_element_indices(self, device, temporary_store: TemporaryStore):
+        from warp.fem import cache
+
+        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}",
+            kernel_options={"max_unroll": 8},
+        )
+        def fill_element_node_indices(
+            element_arg: self.domain.ElementArg,
+            domain_index_arg: self.domain.ElementIndexArg,
+            topo_arg: self.space_topology.TopologyArg,
+            partition_arg: self.space_partition.PartitionArg,
+            element_node_indices: wp.array2d(dtype=int),
+        ):
+            domain_element_index = wp.tid()
+            element_index = self.domain.element_index(domain_index_arg, domain_element_index)
+            for n in range(NODES_PER_ELEMENT):
+                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
+
+        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.array.shape[0],
+            kernel=fill_element_node_indices,
+            inputs=[
+                self.domain.element_arg_value(device),
+                self.domain.element_index_arg_value(device),
+                self.space_topology.topo_arg_value(device),
+                self.space_partition.partition_arg_value(device),
+                element_node_indices.array,
+            ],
+            device=device,
+        )
+
+        # Build compressed map from node to element indices
+        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, temporary_store=temporary_store
+        )
+
+        # Extract element index and index in element
+        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.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.array
+
+    def node_partition_indices(self):
+        return self._dof_partition_indices.array
+
+    def total_node_element_count(self):
+        return self._dof_element_indices.array.size
+
+    @wp.struct
+    class NodeArg:
+        dof_element_offsets: wp.array(dtype=int)
+        dof_element_indices: wp.array(dtype=int)
+        dof_partition_indices: wp.array(dtype=int)
+        dof_indices_in_element: wp.array(dtype=int)
+
+    @cached_arg_value
+    def node_arg(self, device):
+        arg = SpaceRestriction.NodeArg()
+        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
+    def node_partition_index(args: NodeArg, node_index: int):
+        return args.dof_partition_indices[node_index]
+
+    @wp.func
+    def node_element_count(args: NodeArg, node_index: int):
+        partition_node_index = SpaceRestriction.node_partition_index(args, node_index)
+        return args.dof_element_offsets[partition_node_index + 1] - args.dof_element_offsets[partition_node_index]
+
+    @wp.func
+    def node_element_index(args: NodeArg, node_index: int, element_index: int):
+        partition_node_index = SpaceRestriction.node_partition_index(args, node_index)
+        offset = args.dof_element_offsets[partition_node_index] + element_index
+        domain_element_index = args.dof_element_indices[offset]
+        index_in_element = args.dof_indices_in_element[offset]
+        return NodeElementIndex(domain_element_index, index_in_element)
+
+    @wp.kernel
+    def _split_vertex_element_index(
+        vertex_per_element: int,
+        sorted_indices: wp.array(dtype=int),
+        vertex_element_index: wp.array(dtype=int),
+        vertex_index_in_element: wp.array(dtype=int),
+    ):
+        idx = sorted_indices[wp.tid()]
+        element_index = idx // vertex_per_element
+        vertex_element_index[wp.tid()] = element_index
+        vertex_index_in_element[wp.tid()] = idx - vertex_per_element * element_index

warp/fem/space/shape/__init__.py

@@ -1,15 +1,90 @@
-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,
-)
+from enum import Enum
+from typing import Optional
+
+from warp.fem.geometry import element as _element
+from warp.fem.polynomial import Polynomial
+
+from .cube_shape_function import (
+    CubeNonConformingPolynomialShapeFunctions,
+    CubeSerendipityShapeFunctions,
+    CubeTripolynomialShapeFunctions,
+)
+from .shape_function import ConstantShapeFunction, ShapeFunction
+from .square_shape_function import (
+    SquareBipolynomialShapeFunctions,
+    SquareNonConformingPolynomialShapeFunctions,
+    SquareSerendipityShapeFunctions,
+)
+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 basis functions :math:`P_k` for simplices, tensor products :math:`Q_k` for squares and cubes"""
+    SERENDIPITY = 1
+    """Serendipity elements :math:`S_k`, corresponding to Lagrange nodes with interior points removed (for degree <= 3)"""
+    NONCONFORMING_POLYNOMIAL = 2
+    """Simplex Lagrange basis functions :math:`P_{kd}` embedded into non conforming reference elements (e.g. squares or cubes). Discontinuous only."""
+
+
+def get_shape_function(
+    element: _element.Element,
+    space_dimension: int,
+    degree: int,
+    element_basis: ElementBasis,
+    family: Optional[Polynomial] = None,
+):
+    """
+    Equips a reference element with a shape function basis.
+
+    Args:
+        element: the reference element on which to build the shape function
+        space_dimension: the dimension of the embedding space
+        degree: polynomial degree of the per-element shape functions
+        element_basis: type of basis function for the individual elements
+        family: Polynomial family used to generate the shape function basis. If not provided, a reasonable basis is chosen.
+
+    Returns:
+        the corresponding shape function
+    """
+
+    if degree == 0:
+        return ConstantShapeFunction(element, space_dimension)
+
+    if family is None:
+        family = Polynomial.LOBATTO_GAUSS_LEGENDRE
+
+    if isinstance(element, _element.Square):
+        if element_basis == ElementBasis.NONCONFORMING_POLYNOMIAL:
+            return SquareNonConformingPolynomialShapeFunctions(degree=degree)
+        if element_basis == ElementBasis.SERENDIPITY and degree > 1:
+            return SquareSerendipityShapeFunctions(degree=degree, family=family)
+
+        return SquareBipolynomialShapeFunctions(degree=degree, family=family)
+    if isinstance(element, _element.Triangle):
+        if element_basis == ElementBasis.NONCONFORMING_POLYNOMIAL:
+            return Triangle2DNonConformingPolynomialShapeFunctions(degree=degree)
+        if element_basis == ElementBasis.SERENDIPITY and degree > 2:
+            raise NotImplementedError("Serendipity variant not implemented yet for Triangle elements")
+
+        return Triangle2DPolynomialShapeFunctions(degree=degree)
+
+    if isinstance(element, _element.Cube):
+        if element_basis == ElementBasis.NONCONFORMING_POLYNOMIAL:
+            return CubeNonConformingPolynomialShapeFunctions(degree=degree)
+        if element_basis == ElementBasis.SERENDIPITY and degree > 1:
+            return CubeSerendipityShapeFunctions(degree=degree, family=family)
+
+        return CubeTripolynomialShapeFunctions(degree=degree, family=family)
+    if isinstance(element, _element.Tetrahedron):
+        if element_basis == ElementBasis.NONCONFORMING_POLYNOMIAL:
+            return TetrahedronNonConformingPolynomialShapeFunctions(degree=degree)
+        if element_basis == ElementBasis.SERENDIPITY and degree > 2:
+            raise NotImplementedError("Serendipity variant not implemented yet for Tet elements")
+
+        return TetrahedronPolynomialShapeFunctions(degree=degree)
+
+    return NotImplementedError("Unrecognized element type")