warp-lang 1.10.0__py3-none-macosx_11_0_arm64.whl

warp/_src/fem/space/dof_mapper.py

@@ -0,0 +1,253 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+from enum import Enum
+from typing import Any
+
+import warp as wp
+from warp._src.types import type_size
+
+_wp_module_name_ = "warp.fem.space.dof_mapper"
+
+vec6 = wp.vec(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 = type_size(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/_src/fem/space/function_space.py

@@ -0,0 +1,312 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any
+
+from warp._src.codegen import Struct, StructInstance
+from warp._src.fem import cache
+from warp._src.fem.geometry import Geometry
+from warp._src.fem.types import Coords, ElementIndex, ElementKind, Sample, make_free_sample
+from warp._src.types import type_is_matrix, type_is_vector, type_size
+
+from .topology import SpaceTopology
+
+_wp_module_name_ = "warp.fem.space.function_space"
+
+
+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: 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
+            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"""
+        return self.topology.node_count()
+
+    def space_arg_value(self, device) -> 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 element_kind(self) -> ElementKind:
+        """Kind of element the function space is expressed over"""
+        return ElementKind.CELL if self.dimension == self.geometry.dimension else ElementKind.SIDE
+
+    @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"""
+        return self.ORDER
+
+    @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
+
+        Args:
+            space_partition: If provided, the subset of nodes to consider
+
+        See also: :func:`make_space_partition`
+        """
+        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 supported yet"""
+        return not 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 type_is_vector(self.dtype):
+            return type_size(self.dtype) == self.geometry.dimension
+        if type_is_matrix(self.dtype):
+            return self.dtype._shape_[0] == self.geometry.dimension
+        return False
+
+    @staticmethod
+    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
+    def node_coords_in_element(
+        elt_arg: "SpaceTopology.ElementArg",
+        space_arg: "SpaceArg",  # noqa: F821
+        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",  # noqa: F821
+        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",  # noqa: F821
+        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",  # noqa: F821
+        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",  # noqa: F821
+        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",  # noqa: F821
+        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
+
+    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 NotImplementedError
+
+    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 world-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 world-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