warp-lang 0.10.1__py3-none-win_amd64.whl → 0.11.0__py3-none-win_amd64.whl

warp/fem/space/dof_mapper.py

@@ -0,0 +1,236 @@
+from typing import Any
+from enum import Enum
+import math
+
+import warp as wp
+import warp.types
+
+vec6 = wp.types.vector(length=6, dtype=wp.float32)
+
+_SQRT_2 = wp.constant(math.sqrt(2.0))
+_SQRT_3 = wp.constant(math.sqrt(3.0))
+_SQRT_1_2 = wp.constant(math.sqrt(1.0 / 2.0))
+_SQRT_1_3 = wp.constant(math.sqrt(1.0 / 3.0))
+
+
+class DofMapper:
+    """Base class from mapping node degrees of freedom to function values"""
+
+    value_dtype: type
+    dof_dtype: type
+    DOF_SIZE: int
+
+    @wp.func
+    def dof_to_value(dof: Any):
+        raise NotImplementedError
+
+    @wp.func
+    def value_to_dof(val: Any):
+        raise NotImplementedError
+
+    def __str__(self):
+        return f"{self.value_dtype.__name__}_{self.DOF_SIZE}"
+
+
+class IdentityMapper(DofMapper):
+    """Identity mapper"""
+
+    def __init__(self, dtype: type):
+        if dtype == float:
+            dtype = wp.float32
+
+        self.value_dtype = dtype
+        self.dof_dtype = dtype
+
+        size = warp.types.type_length(dtype)
+        self.DOF_SIZE = wp.constant(size)
+
+    @wp.func
+    def dof_to_value(dof: Any):
+        return dof
+
+    @wp.func
+    def value_to_dof(val: Any):
+        return val
+
+
+class SymmetricTensorMapper(DofMapper):
+    """Orthonormal isomorphism from R^{n (n+1)} to nxn symmetric tensors,
+    using usual L2 norm for vectors and half Frobenius norm, (tau : tau)/2 for tensors.
+    """
+
+    class Mapping(Enum):
+        VOIGT = 0
+        """Voigt ordering of vector coefficients:
+            first the three diagonal terms, then off-diagonal coefficients"""
+        DB16 = 1
+        """Ordering that also separates normal from tangential coefficients:
+           first trace, then other diagonal terms, then off-diagonal coefficients.
+           See [Daviet and Bertails-Descoubes 2016]"""
+
+    def __init__(self, dtype: type, mapping: Mapping = Mapping.VOIGT):
+        self.value_dtype = dtype
+        self.mapping = mapping
+
+        if dtype == wp.mat22:
+            self.dof_dtype = wp.vec3
+            self.DOF_SIZE = wp.constant(3)
+            if mapping == SymmetricTensorMapper.Mapping.VOIGT:
+                self.dof_to_value = SymmetricTensorMapper.dof_to_value_2d_voigt
+                self.value_to_dof = SymmetricTensorMapper.value_to_dof_2d_voigt
+            else:
+                self.dof_to_value = SymmetricTensorMapper.dof_to_value_2d
+                self.value_to_dof = SymmetricTensorMapper.value_to_dof_2d
+        elif dtype == wp.mat33:
+            self.dof_dtype = vec6
+            self.DOF_SIZE = wp.constant(6)
+            if mapping == SymmetricTensorMapper.Mapping.VOIGT:
+                self.dof_to_value = SymmetricTensorMapper.dof_to_value_3d_voigt
+                self.value_to_dof = SymmetricTensorMapper.value_to_dof_3d_voigt
+            else:
+                self.dof_to_value = SymmetricTensorMapper.dof_to_value_3d
+                self.value_to_dof = SymmetricTensorMapper.value_to_dof_3d
+        else:
+            raise ValueError("Unsupported value dtype: ", dtype)
+
+    def __str__(self):
+        return f"{self.mapping}_{self.DOF_SIZE}"
+
+    @wp.func
+    def dof_to_value_2d(dof: wp.vec3):
+        a = dof[0]
+        b = dof[1]
+        c = dof[2]
+        return wp.mat22(a + b, c, c, a - b)
+
+    @wp.func
+    def value_to_dof_2d(val: wp.mat22):
+        a = 0.5 * (val[0, 0] + val[1, 1])
+        b = 0.5 * (val[0, 0] - val[1, 1])
+        c = 0.5 * (val[0, 1] + val[1, 0])
+        return wp.vec3(a, b, c)
+
+    @wp.func
+    def dof_to_value_2d_voigt(dof: wp.vec3):
+        a = _SQRT_2 * dof[0]
+        b = _SQRT_2 * dof[1]
+        c = dof[2]
+        return wp.mat22(a, c, c, b)
+
+    @wp.func
+    def value_to_dof_2d_voigt(val: wp.mat22):
+        a = _SQRT_1_2 * val[0, 0]
+        b = _SQRT_1_2 * val[1, 1]
+        c = 0.5 * (val[0, 1] + val[1, 0])
+        return wp.vec3(a, b, c)
+
+    @wp.func
+    def dof_to_value_3d(dof: vec6):
+        a = dof[0] * _SQRT_2 * _SQRT_1_3
+        b = dof[1]
+        c = dof[2] * _SQRT_1_3
+        d = dof[3]
+        e = dof[4]
+        f = dof[5]
+        return wp.mat33(
+            a + b - c,
+            f,
+            e,
+            f,
+            a - b - c,
+            d,
+            e,
+            d,
+            a + 2.0 * c,
+        )
+
+    @wp.func
+    def value_to_dof_3d(val: wp.mat33):
+        a = (val[0, 0] + val[1, 1] + val[2, 2]) * _SQRT_1_3 * _SQRT_1_2
+        b = 0.5 * (val[0, 0] - val[1, 1])
+        c = 0.5 * (val[2, 2] - (val[0, 0] + val[1, 1] + val[2, 2]) / 3.0) * _SQRT_3
+
+        d = 0.5 * (val[2, 1] + val[1, 2])
+        e = 0.5 * (val[0, 2] + val[2, 0])
+        f = 0.5 * (val[1, 0] + val[0, 1])
+
+        return vec6(a, b, c, d, e, f)
+
+    @wp.func
+    def dof_to_value_3d_voigt(dof: vec6):
+        a = _SQRT_2 * dof[0]
+        b = _SQRT_2 * dof[1]
+        c = _SQRT_2 * dof[2]
+        d = dof[3]
+        e = dof[4]
+        f = dof[5]
+        return wp.mat33(
+            a,
+            f,
+            e,
+            f,
+            b,
+            d,
+            e,
+            d,
+            c,
+        )
+
+    @wp.func
+    def value_to_dof_3d_voigt(val: wp.mat33):
+        a = _SQRT_1_2 * val[0, 0]
+        b = _SQRT_1_2 * val[1, 1]
+        c = _SQRT_1_2 * val[2, 2]
+
+        d = 0.5 * (val[2, 1] + val[1, 2])
+        e = 0.5 * (val[0, 2] + val[2, 0])
+        f = 0.5 * (val[1, 0] + val[0, 1])
+
+        return vec6(a, b, c, d, e, f)
+
+
+class SkewSymmetricTensorMapper(DofMapper):
+    """Orthonormal isomorphism from R^{n (n-1)} to nxn skew-symmetric tensors,
+    using usual L2 norm for vectors and half Frobenius norm, (tau : tau)/2 for tensors.
+    """
+
+    def __init__(self, dtype: type):
+        self.value_dtype = dtype
+
+        if dtype == wp.mat22:
+            self.dof_dtype = float
+            self.DOF_SIZE = wp.constant(1)
+            self.dof_to_value = SkewSymmetricTensorMapper.dof_to_value_2d
+            self.value_to_dof = SkewSymmetricTensorMapper.value_to_dof_2d
+        elif dtype == wp.mat33:
+            self.dof_dtype = wp.vec3
+            self.DOF_SIZE = wp.constant(3)
+            self.dof_to_value = SkewSymmetricTensorMapper.dof_to_value_3d
+            self.value_to_dof = SkewSymmetricTensorMapper.value_to_dof_3d
+        else:
+            raise ValueError("Unsupported value dtype: ", dtype)
+
+    def __str__(self):
+        return f"{self.__class__.__name__}_{self.DOF_SIZE}"
+
+    @wp.func
+    def dof_to_value_2d(dof: float):
+        return wp.mat22(0.0, -dof, dof, 0.0)
+
+    @wp.func
+    def value_to_dof_2d(val: wp.mat22):
+        return 0.5 * (val[1, 0] - val[0, 1])
+
+    @wp.func
+    def dof_to_value_3d(dof: wp.vec3):
+        a = dof[0]
+        b = dof[1]
+        c = dof[2]
+        return wp.mat33(0.0, -c, b, c, 0.0, -a, -b, a, 0.0)
+
+    @wp.func
+    def value_to_dof_3d(val: wp.mat33):
+        a = 0.5 * (val[2, 1] - val[1, 2])
+        b = 0.5 * (val[0, 2] - val[2, 0])
+        c = 0.5 * (val[1, 0] - val[0, 1])
+        return wp.vec3(a, b, c)

warp/fem/space/function_space.py

@@ -0,0 +1,145 @@
+from typing import Any
+
+import warp as wp
+
+from warp.fem.types import DofIndex, ElementIndex, Coords
+from warp.fem.geometry import Geometry
+
+from .topology import SpaceTopology
+
+
+class FunctionSpace:
+    """
+    Interface class for function spaces, i.e. geometry + interpolation basis
+    """
+
+    dtype: type
+    """Value type of the interpolation functions"""
+
+    SpaceArg: wp.codegen.Struct
+    """Structure containing arguments to be passed to device function"""
+
+    VALUE_DOF_COUNT: int
+    """Number of degrees of freedom per node, as a Warp constant"""
+
+    def __init__(self, topology: SpaceTopology):
+        self._topology = topology
+
+        if self._topology.is_trace:
+            self.element_inner_reference_gradient_transform = self.geometry.side_inner_inverse_deformation_gradient
+            self.element_outer_reference_gradient_transform = self.geometry.side_outer_inverse_deformation_gradient
+        else:
+            self.element_inner_reference_gradient_transform = self.geometry.cell_inverse_deformation_gradient
+            self.element_outer_reference_gradient_transform = self.geometry.cell_inverse_deformation_gradient
+
+    def node_count(self) -> int:
+        """Number of nodes in the interpolation basis"""
+        raise NotImplementedError
+
+    def space_arg_value(self, device) -> wp.codegen.StructInstance:
+        """Value of the arguments to be passed to device functions"""
+        raise NotImplementedError
+
+    @property
+    def topology(self) -> SpaceTopology:
+        """Underlying geometry"""
+        return self._topology
+
+    @property
+    def geometry(self) -> Geometry:
+        """Underlying geometry"""
+        return self.topology.geometry
+
+    @property
+    def dimension(self) -> int:
+        """Function space embedding dimension"""
+        return self.topology.dimension
+
+    @property
+    def degree(self) -> int:
+        """Maximum polynomial degree of the underlying basis"""
+        raise NotImplementedError
+
+    @property
+    def name(self):
+        raise NotImplementedError
+
+    def __str__(self):
+        return self.name
+
+    def trace(self) -> "FunctionSpace":
+        """Trace of the function space over lower-dimensional elements of the geometry"""
+        raise NotImplementedError
+
+    def make_field(self, space_partition=None):
+        """Creates a zero-initialized discrete field over the function space holding values for all degrees of freedom of nodes in a space partition
+
+        space_arg:
+            space_partition: If provided, the subset of nodes to consider
+
+        See also: :func:`make_space_partition`
+        """
+        raise NotImplementedError
+
+    @staticmethod
+    def unit_dof_value(elt_arg: "SpaceTopology.ElementArg", space_arg: "SpaceArg", dof: DofIndex):
+        """Unit value for a given degree of freedom. Typically a rank-1 tensor"""
+        raise NotImplementedError
+
+    @staticmethod
+    def node_coords_in_element(
+        elt_arg: "SpaceTopology.ElementArg", space_arg: "SpaceArg", element_index: ElementIndex, node_index_in_elt: int
+    ):
+        """Coordinates inside element of a given node"""
+        raise NotImplementedError
+
+    @staticmethod
+    def node_quadrature_weight(
+        elt_arg: "SpaceTopology.ElementArg", space_arg: "SpaceArg", element_index: ElementIndex, node_index_in_elt: int
+    ):
+        """Weight of a given node when used as a quadrature point"""
+        raise NotImplementedError
+
+    @staticmethod
+    def element_inner_weight(
+        elt_arg: "SpaceTopology.ElementArg",
+        space_arg: "SpaceArg",
+        element_index: ElementIndex,
+        coords: Coords,
+        node_index_in_elt: int,
+    ):
+        """Inner weight for a node at given coordinates"""
+        raise NotImplementedError
+
+    @staticmethod
+    def element_inner_weight_gradient(
+        elt_arg: "SpaceTopology.ElementArg",
+        space_arg: "SpaceArg",
+        element_index: ElementIndex,
+        coords: Coords,
+        node_index_in_elt: int,
+    ):
+        """Inner weight gradient w.r.t. reference space for a node at given coordinates"""
+        raise NotImplementedError
+
+    @staticmethod
+    def element_outer_weight(
+        elt_arg: "SpaceTopology.ElementArg",
+        space_arg: "SpaceArg",
+        element_index: ElementIndex,
+        coords: Coords,
+        node_index_in_elt: int,
+    ):
+        """Outer weight for a node at given coordinates"""
+        raise NotImplementedError
+
+    @staticmethod
+    def element_outer_weight_gradient(
+        elt_arg: "SpaceTopology.ElementArg",
+        space_arg: "SpaceArg",
+        element_index: ElementIndex,
+        coords: Coords,
+        node_index_in_elt: int,
+    ):
+        """Outer weight gradient w.r.t reference space for a node at given coordinates"""
+        raise NotImplementedError

warp/fem/space/grid_2d_function_space.py

@@ -0,0 +1,267 @@
+import warp as wp
+import numpy as np
+
+from warp.fem.types import ElementIndex, Coords
+from warp.fem.polynomial import Polynomial, is_closed
+from warp.fem.geometry import Grid2D
+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,
+)
+
+
+class Grid2DSpaceTopology(SpaceTopology):
+    def __init__(self, grid: Grid2D, shape: ShapeFunction):
+        super().__init__(grid, shape.NODES_PER_ELEMENT)
+        self._shape = shape
+
+    @wp.func
+    def _vertex_coords(vidx_in_cell: int):
+        x = vidx_in_cell // 2
+        y = vidx_in_cell - 2 * x
+        return wp.vec2i(x, y)
+
+    @wp.func
+    def _vertex_index(cell_arg: Grid2D.CellArg, cell_index: ElementIndex, vidx_in_cell: int):
+        res = cell_arg.res
+        x_stride = res[1] + 1
+
+        corner = Grid2D.get_cell(res, cell_index) + Grid2DSpaceTopology._vertex_coords(vidx_in_cell)
+        return Grid2D._from_2d_index(x_stride, corner)
+
+
+class Grid2DDiscontinuousSpaceTopology(
+    DiscontinuousSpaceTopologyMixin,
+    Grid2DSpaceTopology,
+):
+    pass
+
+
+class Grid2DBasisSpace(ShapeBasisSpace):
+    def __init__(self, topology: Grid2DSpaceTopology, shape: ShapeFunction):
+        super().__init__(topology, shape)
+
+        self._grid: Grid2D = topology.geometry
+
+
+class GridPiecewiseConstantBasis(Grid2DBasisSpace):
+    def __init__(self, grid: Grid2D):
+        shape = ConstantShapeFunction(grid.reference_cell(), space_dimension=2)
+        topology = Grid2DDiscontinuousSpaceTopology(grid, shape)
+        super().__init__(shape=shape, topology=topology)
+
+        if isinstance(grid, Grid2D):
+            self.node_grid = self._node_grid
+
+    def _node_grid(self):
+        res = self._grid.res
+
+        X = (np.arange(0, res[0], dtype=float) + 0.5) * self._grid.cell_size[0] + self._grid.origin[0]
+        Y = (np.arange(0, res[1], dtype=float) + 0.5) * self._grid.cell_size[1] + self._grid.origin[1]
+        return np.meshgrid(X, Y, indexing="ij")
+
+    class Trace(TraceBasisSpace):
+        @wp.func
+        def _node_coords_in_element(
+            side_arg: Grid2D.SideArg,
+            basis_arg: Grid2DBasisSpace.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 GridPiecewiseConstantBasis.Trace(self)
+
+
+class GridBipolynomialSpaceTopology(Grid2DSpaceTopology):
+    def __init__(self, grid: Grid2D, shape: SquareBipolynomialShapeFunctions):
+        super().__init__(grid, shape)
+
+        self.element_node_index = self._make_element_node_index()
+
+    def node_count(self) -> int:
+        return (self.geometry.res[0] * self._shape.ORDER + 1) * (self.geometry.res[1] * self._shape.ORDER + 1)
+
+    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: Grid2DSpaceTopology.TopologyArg,
+            element_index: ElementIndex,
+            node_index_in_elt: int,
+        ):
+            res = cell_arg.res
+            cell = Grid2D.get_cell(res, element_index)
+
+            node_i = node_index_in_elt // (ORDER + 1)
+            node_j = node_index_in_elt - (ORDER + 1) * node_i
+
+            node_x = ORDER * cell[0] + node_i
+            node_y = ORDER * cell[1] + node_j
+
+            node_pitch = (res[1] * ORDER) + 1
+            node_index = node_pitch * node_x + node_y
+
+            return node_index
+
+        return element_node_index
+
+
+class GridBipolynomialBasisSpace(Grid2DBasisSpace):
+    def __init__(
+        self,
+        grid: Grid2D,
+        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(GridBipolynomialSpaceTopology, grid, shape)
+
+        super().__init__(topology, shape)
+
+        if isinstance(grid, Grid2D):
+            self.node_grid = self._node_grid
+
+    def _node_grid(self):
+        res = self._grid.res
+
+        cell_coords = np.array(self._shape.LOBATTO_COORDS)[:-1]
+
+        grid_coords_x = np.repeat(np.arange(0, res[0], dtype=float), len(cell_coords)) + np.tile(
+            cell_coords, reps=res[0]
+        )
+        grid_coords_x = np.append(grid_coords_x, res[0])
+        X = grid_coords_x * self._grid.cell_size[0] + self._grid.origin[0]
+
+        grid_coords_y = np.repeat(np.arange(0, res[1], dtype=float), len(cell_coords)) + np.tile(
+            cell_coords, reps=res[1]
+        )
+        grid_coords_y = np.append(grid_coords_y, res[1])
+        Y = grid_coords_y * self._grid.cell_size[1] + self._grid.origin[1]
+
+        return np.meshgrid(X, Y, indexing="ij")
+
+
+class GridDGBipolynomialBasisSpace(Grid2DBasisSpace):
+    def __init__(
+        self,
+        grid: Grid2D,
+        degree: int,
+        family: Polynomial,
+    ):
+        if family is None:
+            family = Polynomial.LOBATTO_GAUSS_LEGENDRE
+
+        shape = SquareBipolynomialShapeFunctions(degree, family=family)
+        topology = Grid2DDiscontinuousSpaceTopology(grid, shape)
+
+        super().__init__(shape=shape, topology=topology)
+
+
+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
+
+
+class GridSerendipityBasisSpace(Grid2DBasisSpace):
+    def __init__(
+        self,
+        grid: Grid2D,
+        degree: int,
+        family: Polynomial,
+    ):
+        if family is None:
+            family = Polynomial.LOBATTO_GAUSS_LEGENDRE
+
+        shape = SquareSerendipityShapeFunctions(degree, family=family)
+        topology = forward_base_topology(GridSerendipitySpaceTopology, grid, shape=shape)
+
+        super().__init__(topology=topology, shape=shape)
+
+
+class GridDGSerendipityBasisSpace(Grid2DBasisSpace):
+    def __init__(
+        self,
+        grid: Grid2D,
+        degree: int,
+        family: Polynomial,
+    ):
+        if family is None:
+            family = Polynomial.LOBATTO_GAUSS_LEGENDRE
+
+        shape = SquareSerendipityShapeFunctions(degree, family=family)
+        topology = Grid2DDiscontinuousSpaceTopology(grid, shape=shape)
+
+        super().__init__(topology=topology, shape=shape)
+
+
+class GridDGPolynomialBasisSpace(Grid2DBasisSpace):
+    def __init__(
+        self,
+        grid: Grid2D,
+        degree: int,
+    ):
+        shape = SquareNonConformingPolynomialShapeFunctions(degree)
+        topology = Grid2DDiscontinuousSpaceTopology(grid, shape=shape)
+
+        super().__init__(topology=topology, shape=shape)