warp-lang 1.4.1__py3-none-win_amd64.whl → 1.5.0__py3-none-win_amd64.whl

warp/fem/space/basis_space.py

@@ -13,10 +13,10 @@ from .topology import RegularDiscontinuousSpaceTopology, SpaceTopology
 
 
 class BasisSpace:
-    """Interface class for defining a scalar-valued basis over a geometry.
+    """Interface class for defining a shape function space over a geometry.
 
     A basis space makes it easy to define multiple function spaces sharing the same basis (and thus nodes) but with different valuation functions;
-    however, it is not a required ingredient of a function space.
+    however, it is not a required component of a function space.
 
     See also: :func:`make_polynomial_basis_space`, :func:`make_collocated_function_space`
     """
@@ -40,6 +40,11 @@ class BasisSpace:
         """Underlying geometry of the basis space"""
         return self._topology.geometry
 
+    @property
+    def value(self) -> ShapeFunction.Value:
+        """Value type for the underlying shape functions"""
+        raise NotImplementedError()
+
     def basis_arg_value(self, device) -> "BasisArg":
         """Value for the argument structure to be passed to device functions"""
         return BasisSpace.BasisArg()
@@ -122,6 +127,20 @@ class BasisSpace:
     def trace(self) -> "TraceBasisSpace":
         return TraceBasisSpace(self)
 
+    @property
+    def weight_type(self):
+        if self.value is ShapeFunction.Value.Scalar:
+            return float
+
+        return cache.cached_vec_type(length=self.geometry.cell_dimension, dtype=float)
+
+    @property
+    def weight_gradient_type(self):
+        if self.value is ShapeFunction.Value.Scalar:
+            return wp.vec(length=self.geometry.cell_dimension, dtype=float)
+
+        return cache.cached_mat_type(shape=(self.geometry.cell_dimension, self.geometry.cell_dimension), dtype=float)
+
 
 class ShapeBasisSpace(BasisSpace):
     """Base class for defining shape-function-based basis spaces."""
@@ -130,6 +149,10 @@ class ShapeBasisSpace(BasisSpace):
         super().__init__(topology)
         self._shape = shape
 
+        if self.value is not ShapeFunction.Value.Scalar:
+            self.BasisArg = self.topology.TopologyArg
+            self.basis_arg_value = self.topology.topo_arg_value
+
         self.ORDER = self._shape.ORDER
 
         if hasattr(shape, "element_node_triangulation"):
@@ -148,6 +171,10 @@ class ShapeBasisSpace(BasisSpace):
         """Shape functions used for defining individual element basis"""
         return self._shape
 
+    @property
+    def value(self) -> ShapeFunction.Value:
+        return self.shape.value
+
     @property
     def name(self):
         return f"{self.topology.name}_{self._shape.name}"
@@ -169,6 +196,9 @@ class ShapeBasisSpace(BasisSpace):
     def make_node_quadrature_weight(self):
         shape_node_quadrature_weight = self._shape.make_node_quadrature_weight()
 
+        if shape_node_quadrature_weight is None:
+            return None
+
         @cache.dynamic_func(suffix=self.name)
         def node_quadrature_weight(
             elt_arg: self.geometry.CellArg,
@@ -191,7 +221,11 @@ class ShapeBasisSpace(BasisSpace):
             coords: Coords,
             node_index_in_elt: int,
         ):
-            return shape_element_inner_weight(coords, node_index_in_elt)
+            if wp.static(self.value == ShapeFunction.Value.Scalar):
+                return shape_element_inner_weight(coords, node_index_in_elt)
+            else:
+                sign = self.topology.element_node_sign(elt_arg, basis_arg, element_index, node_index_in_elt)
+                return sign * shape_element_inner_weight(coords, node_index_in_elt)
 
         return element_inner_weight
 
@@ -206,13 +240,20 @@ class ShapeBasisSpace(BasisSpace):
             coords: Coords,
             node_index_in_elt: int,
         ):
-            return shape_element_inner_weight_gradient(coords, node_index_in_elt)
+            if wp.static(self.value == ShapeFunction.Value.Scalar):
+                return shape_element_inner_weight_gradient(coords, node_index_in_elt)
+            else:
+                sign = self.topology.element_node_sign(elt_arg, basis_arg, element_index, node_index_in_elt)
+                return sign * shape_element_inner_weight_gradient(coords, node_index_in_elt)
 
         return element_inner_weight_gradient
 
     def make_trace_node_quadrature_weight(self, trace_basis):
         shape_trace_node_quadrature_weight = self._shape.make_trace_node_quadrature_weight()
 
+        if shape_trace_node_quadrature_weight is None:
+            return None
+
         @cache.dynamic_func(suffix=self.name)
         def trace_node_quadrature_weight(
             geo_side_arg: trace_basis.geometry.SideArg,
@@ -275,6 +316,10 @@ class TraceBasisSpace(BasisSpace):
     def name(self):
         return f"{self._basis.name}_Trace"
 
+    @property
+    def value(self) -> ShapeFunction.Value:
+        return self._basis.value
+
     def make_node_coords_in_element(self):
         node_coords_in_cell = self._basis.make_node_coords_in_element()
 
@@ -316,7 +361,7 @@ class TraceBasisSpace(BasisSpace):
         ):
             cell_index, index_in_cell = self.topology.inner_cell_index(geo_side_arg, element_index, node_index_in_elt)
             if cell_index == NULL_ELEMENT_INDEX:
-                return 0.0
+                return self.weight_type(0.0)
 
             cell_coords = self.geometry.side_inner_cell_coords(geo_side_arg, element_index, coords)
 
@@ -344,7 +389,7 @@ class TraceBasisSpace(BasisSpace):
         ):
             cell_index, index_in_cell = self.topology.outer_cell_index(geo_side_arg, element_index, node_index_in_elt)
             if cell_index == NULL_ELEMENT_INDEX:
-                return 0.0
+                return self.weight_type(0.0)
 
             cell_coords = self.geometry.side_outer_cell_coords(geo_side_arg, element_index, coords)
 
@@ -361,7 +406,6 @@ class TraceBasisSpace(BasisSpace):
 
     def make_element_inner_weight_gradient(self):
         cell_inner_weight_gradient = self._basis.make_element_inner_weight_gradient()
-        grad_vec_type = wp.vec(length=self.geometry.dimension, dtype=float)
 
         @cache.dynamic_func(suffix=self._basis.name)
         def trace_element_inner_weight_gradient(
@@ -373,7 +417,7 @@ class TraceBasisSpace(BasisSpace):
         ):
             cell_index, index_in_cell = self.topology.inner_cell_index(geo_side_arg, element_index, node_index_in_elt)
             if cell_index == NULL_ELEMENT_INDEX:
-                return grad_vec_type(0.0)
+                return self.weight_gradient_type(0.0)
 
             cell_coords = self.geometry.side_inner_cell_coords(geo_side_arg, element_index, coords)
             geo_cell_arg = self.geometry.side_to_cell_arg(geo_side_arg)
@@ -383,7 +427,6 @@ class TraceBasisSpace(BasisSpace):
 
     def make_element_outer_weight_gradient(self):
         cell_outer_weight_gradient = self._basis.make_element_outer_weight_gradient()
-        grad_vec_type = wp.vec(length=self.geometry.dimension, dtype=float)
 
         @cache.dynamic_func(suffix=self._basis.name)
         def trace_element_outer_weight_gradient(
@@ -395,7 +438,7 @@ class TraceBasisSpace(BasisSpace):
         ):
             cell_index, index_in_cell = self.topology.outer_cell_index(geo_side_arg, element_index, node_index_in_elt)
             if cell_index == NULL_ELEMENT_INDEX:
-                return grad_vec_type(0.0)
+                return self.weight_gradient_type(0.0)
 
             cell_coords = self.geometry.side_outer_cell_coords(geo_side_arg, element_index, coords)
             geo_cell_arg = self.geometry.side_to_cell_arg(geo_side_arg)
@@ -528,6 +571,10 @@ class PointBasisSpace(BasisSpace):
     def name(self):
         return f"{self._quadrature.name}_Point"
 
+    @property
+    def value(self) -> ShapeFunction.Value:
+        return ShapeFunction.Value.Scalar
+
     def make_node_coords_in_element(self):
         @cache.dynamic_func(suffix=self.name)
         def node_coords_in_element(
@@ -571,7 +618,7 @@ class PointBasisSpace(BasisSpace):
         return element_inner_weight
 
     def make_element_inner_weight_gradient(self):
-        gradient_vec = cache.cached_vec_type(length=self.geometry.dimension, dtype=float)
+        gradient_vec = cache.cached_vec_type(length=self.geometry.cell_dimension, dtype=float)
 
         @cache.dynamic_func(suffix=self.name)
         def element_inner_weight_gradient(

warp/fem/space/function_space.py

@@ -1,6 +1,9 @@
+from typing import Any
+
 import warp as wp
+from warp.fem import cache
 from warp.fem.geometry import Geometry
-from warp.fem.types import Coords, DofIndex, ElementIndex, ElementKind
+from warp.fem.types import Coords, ElementIndex, ElementKind, Sample, make_free_sample
 
 from .topology import SpaceTopology
 
@@ -8,19 +11,39 @@ from .topology import SpaceTopology
 class FunctionSpace:
     """
     Interface class for function spaces, i.e. geometry + interpolation basis
+
+    The value of a function `f` at a position `x` is generally computed as
+      ``f(x) = L(x)[sum_i f_i N_i(x)]``
+    with:
+        - ``f_i`` the value of the ith node's degrees-of-freedom (dof)
+        - ``N_i(x)`` the weight associated to the node at `x`
+        - ``L(x)`` local linear transformation from node-space to world-space
     """
 
     dtype: type
     """Value type of the interpolation functions"""
 
+    dof_dtype: type
+    """Data type of the degrees of freedom of each node"""
+
     SpaceArg: wp.codegen.Struct
     """Structure containing arguments to be passed to device function"""
 
+    LocalValueMap: type
+    """Type of the local map for transforming vector-valued functions from reference to world space"""
+
     VALUE_DOF_COUNT: int
+    """Number of degrees of freedom per value, as a Warp constant"""
+
+    NODE_DOF_COUNT: int
     """Number of degrees of freedom per node, as a Warp constant"""
 
+    ORDER: int
+    """Polynomial degree of the function space, used to determine integration order"""
+
     def __init__(self, topology: SpaceTopology):
         self._topology = topology
+        self.ElementArg = self.topology.ElementArg
 
         if self._topology.is_trace:
             self.element_inner_reference_gradient_transform = self.geometry.side_inner_inverse_deformation_gradient
@@ -31,7 +54,7 @@ class FunctionSpace:
 
     def node_count(self) -> int:
         """Number of nodes in the interpolation basis"""
-        raise NotImplementedError
+        return self.topology.node_count()
 
     def space_arg_value(self, device) -> wp.codegen.StructInstance:
         """Value of the arguments to be passed to device functions"""
@@ -60,7 +83,7 @@ class FunctionSpace:
     @property
     def degree(self) -> int:
         """Maximum polynomial degree of the underlying basis"""
-        raise NotImplementedError
+        return self.ORDER
 
     @property
     def name(self):
@@ -83,9 +106,50 @@ class FunctionSpace:
         """
         raise NotImplementedError
 
+    def gradient_valid(self) -> bool:
+        """Whether gradient operator can be computed. Only for scalar and vector fields as higher-order tensors are not support yet"""
+        return not wp.types.type_is_matrix(self.dtype)
+
+    def divergence_valid(self) -> bool:
+        """Whether divergence of this field can be computed. Only for vector and tensor fields with same dimension as embedding geometry"""
+        if wp.types.type_is_vector(self.dtype):
+            return wp.types.type_length(self.dtype) == self.geometry.dimension
+        if wp.types.type_is_matrix(self.dtype):
+            return self.dtype._shape_[0] == self.geometry.dimension
+        return False
+
     @staticmethod
-    def unit_dof_value(elt_arg: "SpaceTopology.ElementArg", space_arg: "SpaceArg", dof: DofIndex):  # noqa: F821
-        """Unit value for a given degree of freedom. Typically a rank-1 tensor"""
+    def node_basis_element(dof_coord: int):
+        """Basis element for node degrees of freedom.
+
+        Assumes 0 <= dof_coord < NODE_DOF_COUNT
+        """
+        raise NotImplementedError
+
+    @staticmethod
+    def value_basis_element(dof_coord: int):
+        """Basis element for the function space values
+
+        Assumes 0 <= dof_coord < VALUE_DOF_COUNT
+        """
+        raise NotImplementedError
+
+    @staticmethod
+    def local_value_map_inner(
+        elt_arg: "SpaceTopology.ElementArg",
+        element_index: ElementIndex,
+        coords: Coords,
+    ):
+        """Builds the local value map transforming from node to world space"""
+        raise NotImplementedError
+
+    @staticmethod
+    def local_value_map_outer(
+        elt_arg: "SpaceTopology.ElementArg",
+        element_index: ElementIndex,
+        coords: Coords,
+    ):
+        """Builds the local value map transforming vector-valued from node to world space"""
         raise NotImplementedError
 
     @staticmethod
@@ -151,3 +215,80 @@ class FunctionSpace:
     ):
         """Outer weight gradient w.r.t reference space for a node at given coordinates"""
         raise NotImplementedError
+
+    def space_value(
+        dof_value: "FunctionSpace.dof_dtype",
+        node_weight: Any,
+        local_value_map: "FunctionSpace.LocalValueMap",
+    ):
+        """
+        Assembles the world-space value of the function space
+        Args:
+         - dof_value: node value in the degrees-of-freedom basis
+         - node_weight: weight associated to the node, as given per `element_(inn|out)er_weight`
+         - local_value_map: local transformation from node space to world space, as given per `local_map_value_(inn|out)er`
+        """
+        raise NotADirectoryError
+
+    def space_gradient(
+        dof_value: "FunctionSpace.dof_dtype",
+        node_weight: Any,
+        local_value_map: "FunctionSpace.LocalValueMap",
+        grad_transform: Any,
+    ):
+        """
+        Assembles the world-space gradient of the function space
+        Args:
+         - dof_value: node value in the degrees-of-freedom basis
+         - node_weight_gradient: gradient of the weight associated to the node, as given per `element_(inn|out)er_weight_gradient`
+         - local_value_map: local transformation from node space to world space, as given per `local_map_value_(inn|out)er`
+         - grad_transform: transform mapping the reference space gradient to worls-space gradient (inverse deformation gradient)
+        """
+        raise NotImplementedError
+
+    def space_divergence(
+        dof_value: "FunctionSpace.dof_dtype",
+        node_weight: Any,
+        local_value_map: "FunctionSpace.LocalValueMap",
+        grad_transform: Any,
+    ):
+        """ "
+        Assembles the world-space divergence of the function space
+        Args:
+         - dof_value: node value in the degrees-of-freedom basis
+         - node_weight_gradient: gradient of the weight associated to the node, as given per `element_(inn|out)er_weight_gradient`
+         - local_value_map: local transformation from node space to world space, as given per `local_map_value_(inn|out)er`
+         - grad_transform: transform mapping the reference space gradient to worls-space gradient (inverse deformation gradient)
+        """
+        raise NotImplementedError
+
+    @staticmethod
+    def node_dof_value(
+        elt_arg: "FunctionSpace.ElementArg",
+        space_arg: "FunctionSpace.SpaceArg",
+        element_index: ElementIndex,
+        node_index_in_elt: int,
+        space_value: "FunctionSpace.dtype",
+    ):
+        """Converts space value to node degrees of freedom"""
+        raise NotImplementedError
+
+    def _make_side_inner_inverse_deformation_gradient(self):
+        @cache.dynamic_func(suffix=self.name)
+        def side_inner_inverse_deformation_gradient(args: self.ElementArg, s: Sample):
+            cell_index = self.side_inner_cell_index(args, s.element_index)
+            cell_coords = self.side_inner_cell_coords(args, s.element_index, s.element_coords)
+            cell_arg = self.side_to_cell_arg(args)
+            return self.geometry.cell_inverse_deformation_gradient(cell_arg, make_free_sample(cell_index, cell_coords))
+
+        return side_inner_inverse_deformation_gradient
+
+    def _make_side_outer_inverse_deformation_gradient(self):
+        @cache.dynamic_func(suffix=self.name)
+        def side_outer_inverse_deformation_gradient(args: self.ElementArg, s: Sample):
+            cell_index = self.side_outer_cell_index(args, s.element_index)
+            cell_coords = self.side_outer_cell_coords(args, s.element_index, s.element_coords)
+            cell_arg = self.side_to_cell_arg(args)
+            return self.geometry.cell_inverse_deformation_gradient(cell_arg, make_free_sample(cell_index, cell_coords))
+
+        return side_outer_inverse_deformation_gradient

warp/fem/space/grid_2d_function_space.py

@@ -4,23 +4,83 @@ import warp as wp
 from warp.fem import cache
 from warp.fem.geometry import Grid2D
 from warp.fem.polynomial import is_closed
-from warp.fem.types import ElementIndex
+from warp.fem.types import NULL_NODE_INDEX, ElementIndex
 
-from .shape import (
-    ShapeFunction,
-    SquareBipolynomialShapeFunctions,
-    SquareSerendipityShapeFunctions,
-)
+from .shape import SquareBipolynomialShapeFunctions, SquareShapeFunction
 from .topology import SpaceTopology, forward_base_topology
 
 
 class Grid2DSpaceTopology(SpaceTopology):
-    def __init__(self, grid: Grid2D, shape: ShapeFunction):
-        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)
+    def __init__(self, grid: Grid2D, shape: SquareShapeFunction):
         self._shape = shape
+        super().__init__(grid, shape.NODES_PER_ELEMENT)
+
+        self.element_node_index = self._make_element_node_index()
+
+    TopologyArg = Grid2D.SideArg
+
+    @property
+    def name(self):
+        return f"{self.geometry.name}_{self._shape.name}"
+
+    def topo_arg_value(self, device):
+        return self.geometry.side_arg_value(device)
+
+    def node_count(self) -> int:
+        return (
+            self.geometry.vertex_count() * self._shape.VERTEX_NODE_COUNT
+            + self.geometry.side_count() * self._shape.EDGE_NODE_COUNT
+            + self.geometry.cell_count() * self._shape.INTERIOR_NODE_COUNT
+        )
+
+    def _make_element_node_index(self):
+        VERTEX_NODE_COUNT = self._shape.VERTEX_NODE_COUNT
+        EDGE_NODE_COUNT = self._shape.EDGE_NODE_COUNT
+        INTERIOR_NODE_COUNT = self._shape.INTERIOR_NODE_COUNT
+
+        @cache.dynamic_func(suffix=self.name)
+        def element_node_index(
+            cell_arg: Grid2D.CellArg,
+            topo_arg: Grid2D.SideArg,
+            element_index: ElementIndex,
+            node_index_in_elt: int,
+        ):
+            node_type, type_instance, type_index = self._shape.node_type_and_type_index(node_index_in_elt)
+
+            if wp.static(VERTEX_NODE_COUNT > 0):
+                if node_type == SquareShapeFunction.VERTEX:
+                    return (
+                        Grid2DSpaceTopology._vertex_index(cell_arg, element_index, type_instance) * VERTEX_NODE_COUNT
+                        + type_index
+                    )
+
+            res = cell_arg.res
+            vertex_count = (res[0] + 1) * (res[1] + 1)
+            global_offset = vertex_count
+
+            if wp.static(INTERIOR_NODE_COUNT > 0):
+                if node_type == SquareShapeFunction.INTERIOR:
+                    return global_offset + element_index * INTERIOR_NODE_COUNT + type_index
+
+                cell_count = res[0] * res[1]
+                global_offset += INTERIOR_NODE_COUNT * cell_count
+
+            if wp.static(EDGE_NODE_COUNT > 0):
+                axis = 1 - (node_type - SquareShapeFunction.EDGE_X)
+
+                cell = Grid2D.get_cell(cell_arg.res, element_index)
+                origin = wp.vec2i(cell[Grid2D.ROTATION[axis, 0]] + type_instance, cell[Grid2D.ROTATION[axis, 1]])
+
+                side = Grid2D.Side(axis, origin)
+                side_index = Grid2D.side_index(topo_arg, side)
+
+                vertex_count = (res[0] + 1) * (res[1] + 1)
+
+                return global_offset + EDGE_NODE_COUNT * side_index + type_index
+
+            return NULL_NODE_INDEX  # unreachable
+
+        return element_node_index
 
     @wp.func
     def _vertex_coords(vidx_in_cell: int):
@@ -37,10 +97,10 @@ class Grid2DSpaceTopology(SpaceTopology):
         return Grid2D._from_2d_index(x_stride, corner)
 
 
-class GridBipolynomialSpaceTopology(Grid2DSpaceTopology):
+class GridBipolynomialSpaceTopology(SpaceTopology):
     def __init__(self, grid: Grid2D, shape: SquareBipolynomialShapeFunctions):
-        super().__init__(grid, shape)
-
+        super().__init__(grid, shape.NODES_PER_ELEMENT)
+        self._shape = shape
         self.element_node_index = self._make_element_node_index()
 
     def node_count(self) -> int:
@@ -52,7 +112,7 @@ class GridBipolynomialSpaceTopology(Grid2DSpaceTopology):
         @cache.dynamic_func(suffix=self.name)
         def element_node_index(
             cell_arg: Grid2D.CellArg,
-            topo_arg: Grid2DSpaceTopology.TopologyArg,
+            topo_arg: self.TopologyArg,
             element_index: ElementIndex,
             node_index_in_elt: int,
         ):
@@ -92,57 +152,11 @@ class GridBipolynomialSpaceTopology(Grid2DSpaceTopology):
         return np.meshgrid(X, Y, indexing="ij")
 
 
-class GridSerendipitySpaceTopology(Grid2DSpaceTopology):
-    def __init__(self, grid: Grid2D, shape: SquareSerendipityShapeFunctions):
-        super().__init__(grid, shape)
-
-        self.element_node_index = self._make_element_node_index()
-
-    TopologyArg = Grid2D.SideArg
-
-    def topo_arg_value(self, device):
-        return self.geometry.side_arg_value(device)
-
-    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
-
-        @cache.dynamic_func(suffix=self.name)
-        def element_node_index(
-            cell_arg: Grid2D.CellArg,
-            topo_arg: Grid2D.SideArg,
-            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 Grid2DSpaceTopology._vertex_index(cell_arg, element_index, type_index)
-
-            side_offset, index_in_side = SquareSerendipityShapeFunctions.side_offset_and_index(type_index)
-            axis = 1 - (node_type - SquareSerendipityShapeFunctions.EDGE_X)
-
-            cell = Grid2D.get_cell(cell_arg.res, element_index)
-            origin = wp.vec2i(cell[Grid2D.ROTATION[axis, 0]] + side_offset, cell[Grid2D.ROTATION[axis, 1]])
-
-            side = Grid2D.Side(axis, origin)
-            side_index = Grid2D.side_index(topo_arg, side)
-
-            res = cell_arg.res
-            vertex_count = (res[0] + 1) * (res[1] + 1)
-
-            return vertex_count + (ORDER - 1) * side_index + index_in_side
-
-        return element_node_index
-
-
-def make_grid_2d_space_topology(grid: Grid2D, shape: ShapeFunction):
-    if isinstance(shape, SquareSerendipityShapeFunctions):
-        return forward_base_topology(GridSerendipitySpaceTopology, grid, shape)
-
-    if isinstance(shape, SquareBipolynomialShapeFunctions):
+def make_grid_2d_space_topology(grid: Grid2D, shape: SquareShapeFunction):
+    if isinstance(shape, SquareBipolynomialShapeFunctions) and is_closed(shape.family):
         return forward_base_topology(GridBipolynomialSpaceTopology, grid, shape)
 
+    if isinstance(shape, SquareShapeFunction):
+        return forward_base_topology(Grid2DSpaceTopology, grid, shape)
+
     raise ValueError(f"Unsupported shape function {shape.name}")