warp-lang 1.10.0__py3-none-macosx_11_0_arm64.whl

warp/_src/fem/space/shape/shape_function.py

@@ -0,0 +1,134 @@
+# 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 enum import Enum
+
+import numpy as np
+
+import warp as wp
+from warp._src.fem import cache
+from warp._src.fem.geometry import Element
+from warp._src.fem.types import Coords
+
+_wp_module_name_ = "warp.fem.space.shape.shape_function"
+
+
+class ShapeFunction:
+    """Interface class for defining scalar-valued shape functions over a single element"""
+
+    ORDER: int
+    """Maximum degree of the polynomials used to define the shape function"""
+
+    NODES_PER_ELEMENT: int
+    """Number of shape function nodes"""
+
+    class Value(Enum):
+        Scalar = 0
+        CovariantVector = 1
+        ContravariantVector = 2
+
+    value: Value = Value.Scalar
+
+    @property
+    def name(self) -> str:
+        """Unique name encoding all parameters defining the shape function"""
+        raise NotImplementedError()
+
+    def make_node_coords_in_element(self):
+        """Creates a device function returning the coordinates of each node"""
+        raise NotImplementedError()
+
+    def make_node_quadrature_weight(self):
+        """Creates a device function returning the weight of each node when use as a quadrature point over the element"""
+        raise NotImplementedError()
+
+    def make_trace_node_quadrature_weight(self):
+        """Creates a device function returning the weight of each node when use as a quadrature point over the element boundary"""
+        raise NotImplementedError()
+
+    def make_element_inner_weight(self):
+        """Creates a device function returning the value of the shape function associated to a given node at given coordinates"""
+        raise NotImplementedError()
+
+    def make_element_inner_weight_gradient(self):
+        """Creates a device function returning the gradient of the shape function associated to a given node at given coordinates"""
+        raise NotImplementedError()
+
+
+class ConstantShapeFunction(ShapeFunction):
+    """Shape function that is constant over the element"""
+
+    def __init__(self, element: Element):
+        self._element_prototype = element.prototype
+
+        self.ORDER = wp.constant(0)
+        self.NODES_PER_ELEMENT = wp.constant(1)
+
+        coords, _ = self._element_prototype.instantiate_quadrature(order=0, family=None)
+        self.COORDS = wp.constant(coords[0])
+
+    @property
+    def name(self) -> str:
+        return f"{self._element_prototype.__name__}"
+
+    def make_node_coords_in_element(self):
+        COORDS = self.COORDS
+
+        @cache.dynamic_func(suffix=self.name)
+        def node_coords_in_element(
+            node_index_in_elt: int,
+        ):
+            return COORDS
+
+        return node_coords_in_element
+
+    @wp.func
+    def _node_quadrature_weight(
+        node_index_in_elt: int,
+    ):
+        return 1.0
+
+    def make_node_quadrature_weight(self):
+        return ConstantShapeFunction._node_quadrature_weight
+
+    def make_trace_node_quadrature_weight(self):
+        return ConstantShapeFunction._node_quadrature_weight
+
+    @wp.func
+    def _element_inner_weight(
+        coords: Coords,
+        node_index_in_elt: int,
+    ):
+        return 1.0
+
+    def make_element_inner_weight(self):
+        return ConstantShapeFunction._element_inner_weight
+
+    def make_element_inner_weight_gradient(self):
+        grad_type = cache.cached_vec_type(length=self._element_prototype.dimension, dtype=float)
+
+        @cache.dynamic_func(suffix=self.name)
+        def element_inner_weight_gradient(
+            coords: Coords,
+            node_index_in_elt: int,
+        ):
+            return grad_type(0.0)
+
+        return element_inner_weight_gradient
+
+    def element_vtk_cells(self):
+        cell_type = 1  # VTK_VERTEX
+
+        return np.zeros((1, 1), dtype=int), np.full(1, cell_type, dtype=np.int8)