warp-lang 1.9.0__py3-none-manylinux_2_34_aarch64.whl → 1.10.0__py3-none-manylinux_2_34_aarch64.whl

warp/{fem → _src/fem}/space/restriction.py

@@ -13,14 +13,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import Optional
+
 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 NULL_NODE_INDEX, NodeElementIndex
-from warp.fem.utils import compress_node_indices
+from warp._src.fem import cache
+from warp._src.fem.domain import GeometryDomain
+from warp._src.fem.types import NULL_ELEMENT_INDEX, NULL_NODE_INDEX, NodeElementIndex
+from warp._src.fem.utils import compress_node_indices, host_read_at_index
 
 from .partition import SpacePartition
 
+_wp_module_name_ = "warp.fem.space.restriction"
+
 wp.set_module_options({"enable_backward": False})
 
 
@@ -32,7 +36,7 @@ class SpaceRestriction:
         space_partition: SpacePartition,
         domain: GeometryDomain,
         device=None,
-        temporary_store: TemporaryStore = None,
+        temporary_store: cache.TemporaryStore = None,
     ):
         space_topology = space_partition.space_topology
 
@@ -46,15 +50,25 @@ class SpaceRestriction:
         self.space_topology = space_topology
         self.domain = domain
 
-        self._compute_node_element_indices(device=device, temporary_store=temporary_store)
+        self._node_count_dev: wp.array = None
+        """Number of unique partition node indices"""
+        self._dof_partition_indices: wp.array = None
+        """Array of unique partition node indices"""
 
-    def _compute_node_element_indices(self, device, temporary_store: TemporaryStore):
-        from warp.fem import cache
+        self._dof_partition_element_offsets: wp.array = None
+        """Mapping from partition node to offset in the per-node element indices array"""
+        self._dof_element_indices: wp.array = None
+        """Concatenation of neighboring elements indices for each partition node"""
+        self._dof_indices_in_element: wp.array = None
+        """Concatenation of node index in element for each partition node"""
 
-        MAX_NODES_PER_ELEMENT = self.space_topology.MAX_NODES_PER_ELEMENT
+        self.rebuild(device=device, temporary_store=temporary_store)
+
+    def rebuild(self, device: Optional = None, temporary_store: Optional[cache.TemporaryStore] = None):
+        max_nodes_per_element = self.space_topology.MAX_NODES_PER_ELEMENT
 
         @cache.dynamic_kernel(
-            suffix=f"{self.domain.name}_{self.space_topology.name}_{self.space_partition.name}",
+            suffix=(self.domain.name, self.space_topology.name, self.space_partition.name),
             kernel_options={"max_unroll": 8},
         )
         def fill_element_node_indices(
@@ -66,75 +80,96 @@ class SpaceRestriction:
         ):
             domain_element_index = wp.tid()
             element_index = self.domain.element_index(domain_index_arg, domain_element_index)
-            element_node_count = self.space_topology.element_node_count(element_arg, topo_arg, element_index)
+
+            if element_index == NULL_ELEMENT_INDEX:
+                element_node_count = 0
+            else:
+                element_node_count = self.space_topology.element_node_count(element_arg, topo_arg, element_index)
+
             for n in range(element_node_count):
                 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
-            for n in range(element_node_count, MAX_NODES_PER_ELEMENT):
+            for n in range(element_node_count, element_node_indices.shape[1]):
                 element_node_indices[domain_element_index, n] = NULL_NODE_INDEX
 
-        element_node_indices = borrow_temporary(
+        element_node_indices = cache.borrow_temporary(
             temporary_store,
-            shape=(self.domain.element_count(), MAX_NODES_PER_ELEMENT),
+            shape=(self.domain.element_count(), max_nodes_per_element),
             dtype=int,
             device=device,
         )
         wp.launch(
-            dim=element_node_indices.array.shape[0],
+            dim=element_node_indices.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,
+                element_node_indices,
             ],
             device=device,
         )
 
         # Build compressed map from node to element indices
-        flattened_node_indices = element_node_indices.array.flatten()
+        flattened_node_indices = element_node_indices.flatten()
         (
             self._dof_partition_element_offsets,
             node_array_indices,
-            self._node_count,
+            self._node_count_dev,
             self._dof_partition_indices,
         ) = compress_node_indices(
             self.space_partition.node_count(),
             flattened_node_indices,
+            node_offsets=self._dof_partition_element_offsets,
+            unique_node_count=self._node_count_dev,
+            unique_node_indices=self._dof_partition_indices,
             return_unique_nodes=True,
             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)
+        if self._dof_element_indices is None or self._dof_element_indices.shape != flattened_node_indices.shape:
+            self._dof_element_indices = cache.borrow_temporary_like(flattened_node_indices, temporary_store)
+            self._dof_indices_in_element = cache.borrow_temporary_like(flattened_node_indices, temporary_store)
+
         wp.launch(
             kernel=SpaceRestriction._split_vertex_element_index,
             dim=flattened_node_indices.shape,
             inputs=[
-                MAX_NODES_PER_ELEMENT,
-                node_array_indices.array,
-                self._dof_element_indices.array,
-                self._dof_indices_in_element.array,
+                max_nodes_per_element,
+                node_array_indices,
+                self._dof_element_indices,
+                self._dof_indices_in_element,
             ],
             device=flattened_node_indices.device,
         )
 
         node_array_indices.release()
 
-    def node_count(self):
+        # Upper bound on node count, use `node_count_sync` to get the actual value
+        self._node_count = min(self.space_partition.node_count(), self._dof_partition_indices.shape[0])
+
+    def node_count_sync(self) -> int:
+        """Ensures that the node count is synchronized with the device and returns it"""
+        if self._node_count_dev is not None:
+            self._node_count = int(host_read_at_index(self._node_count_dev, index=0))
+            self._node_count_dev = None
+        return self.node_count()
+
+    def node_count(self) -> int:
+        """Upper bound for the node count, use `node_count_sync` to get the actual value"""
         return self._node_count
 
     def partition_element_offsets(self):
-        return self._dof_partition_element_offsets.array
+        return self._dof_partition_element_offsets
 
     def node_partition_indices(self):
-        return self._dof_partition_indices.array
+        return self._dof_partition_indices
 
     def total_node_element_count(self):
-        return self._dof_element_indices.array.size
+        return self._dof_element_indices.size
 
     @wp.struct
     class NodeArg:
@@ -143,17 +178,17 @@ class SpaceRestriction:
         dof_partition_indices: wp.array(dtype=int)
         dof_indices_in_element: wp.array(dtype=int)
 
-    @cached_arg_value
+    @cache.cached_arg_value
     def node_arg_value(self, device):
         arg = SpaceRestriction.NodeArg()
         self.fill_node_arg(arg, device)
         return arg
 
     def fill_node_arg(self, arg: NodeArg, device):
-        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)
+        arg.dof_element_offsets = self._dof_partition_element_offsets.to(device)
+        arg.dof_element_indices = self._dof_element_indices.to(device)
+        arg.dof_partition_indices = self._dof_partition_indices.to(device)
+        arg.dof_indices_in_element = self._dof_indices_in_element.to(device)
 
     @wp.func
     def node_partition_index(args: NodeArg, restriction_node_index: int):

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

@@ -0,0 +1,152 @@
+# 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 functools
+from enum import Enum
+from typing import Optional
+
+from warp._src.fem.geometry import Element
+from warp._src.fem.polynomial import Polynomial
+
+from .cube_shape_function import (
+    CubeNedelecFirstKindShapeFunctions,
+    CubeNonConformingPolynomialShapeFunctions,
+    CubeRaviartThomasShapeFunctions,
+    CubeSerendipityShapeFunctions,
+    CubeShapeFunction,
+    CubeTripolynomialShapeFunctions,
+)
+from .shape_function import ConstantShapeFunction, ShapeFunction
+from .square_shape_function import (
+    SquareBipolynomialShapeFunctions,
+    SquareNedelecFirstKindShapeFunctions,
+    SquareNonConformingPolynomialShapeFunctions,
+    SquareRaviartThomasShapeFunctions,
+    SquareSerendipityShapeFunctions,
+    SquareShapeFunction,
+)
+from .tet_shape_function import (
+    TetrahedronNedelecFirstKindShapeFunctions,
+    TetrahedronNonConformingPolynomialShapeFunctions,
+    TetrahedronPolynomialShapeFunctions,
+    TetrahedronRaviartThomasShapeFunctions,
+    TetrahedronShapeFunction,
+)
+from .triangle_shape_function import (
+    TriangleNedelecFirstKindShapeFunctions,
+    TriangleNonConformingPolynomialShapeFunctions,
+    TrianglePolynomialShapeFunctions,
+    TriangleRaviartThomasShapeFunctions,
+    TriangleShapeFunction,
+)
+
+
+class ElementBasis(Enum):
+    """Choice of basis function to equip individual elements"""
+
+    LAGRANGE = "P"
+    """Lagrange basis functions :math:`P_k` for simplices, tensor products :math:`Q_k` for squares and cubes"""
+    SERENDIPITY = "S"
+    """Serendipity elements :math:`S_k`, corresponding to Lagrange nodes with interior points removed (for degree <= 3)"""
+    NONCONFORMING_POLYNOMIAL = "dP"
+    """Simplex Lagrange basis functions :math:`P_{kd}` embedded into non conforming reference elements (e.g. squares or cubes). Discontinuous only."""
+    NEDELEC_FIRST_KIND = "N1"
+    """Nédélec (first kind) H(curl) shape functions. Should be used with covariant function space."""
+    RAVIART_THOMAS = "RT"
+    """Raviart-Thomas H(div) shape functions. Should be used with contravariant function space."""
+
+
+@functools.lru_cache(maxsize=None)
+def make_element_shape_function(
+    element: Element,
+    degree: int,
+    element_basis: Optional[ElementBasis] = None,
+    family: Optional[Polynomial] = None,
+) -> ShapeFunction:
+    """
+    Equips a reference element with a shape function basis.
+
+    Args:
+        element: the type of reference element on which to build the shape function
+        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
+
+    Raises:
+        NotImplementedError: If the shape function is not implemented for the given element type
+    """
+
+    if element_basis is None:
+        element_basis = ElementBasis.LAGRANGE
+    elif element_basis == ElementBasis.SERENDIPITY and degree == 1:
+        # Degree-1 serendipity is always equivalent to Lagrange
+        element_basis = ElementBasis.LAGRANGE
+
+    if degree == 0:
+        return ConstantShapeFunction(element)
+
+    if family is None:
+        family = Polynomial.LOBATTO_GAUSS_LEGENDRE
+
+    if element == Element.SQUARE:
+        if element_basis == ElementBasis.NEDELEC_FIRST_KIND:
+            return SquareNedelecFirstKindShapeFunctions(degree=degree)
+        if element_basis == ElementBasis.RAVIART_THOMAS:
+            return SquareRaviartThomasShapeFunctions(degree=degree)
+        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 element == Element.TRIANGLE:
+        if element_basis == ElementBasis.NEDELEC_FIRST_KIND:
+            return TriangleNedelecFirstKindShapeFunctions(degree=degree)
+        if element_basis == ElementBasis.RAVIART_THOMAS:
+            return TriangleRaviartThomasShapeFunctions(degree=degree)
+        if element_basis == ElementBasis.NONCONFORMING_POLYNOMIAL:
+            return TriangleNonConformingPolynomialShapeFunctions(degree=degree)
+        if element_basis == ElementBasis.SERENDIPITY and degree > 2:
+            raise NotImplementedError("Serendipity variant not implemented yet for Triangle elements")
+
+        return TrianglePolynomialShapeFunctions(degree=degree)
+
+    if element == Element.CUBE:
+        if element_basis == ElementBasis.NEDELEC_FIRST_KIND:
+            return CubeNedelecFirstKindShapeFunctions(degree=degree)
+        if element_basis == ElementBasis.RAVIART_THOMAS:
+            return CubeRaviartThomasShapeFunctions(degree=degree)
+        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 element == Element.TETRAHEDRON:
+        if element_basis == ElementBasis.NEDELEC_FIRST_KIND:
+            return TetrahedronNedelecFirstKindShapeFunctions(degree=degree)
+        if element_basis == ElementBasis.RAVIART_THOMAS:
+            return TetrahedronRaviartThomasShapeFunctions(degree=degree)
+        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)
+
+    raise NotImplementedError(f"Unrecognized element type {element}")

warp/{fem → _src/fem}/space/shape/cube_shape_function.py

@@ -18,14 +18,16 @@ import math
 import numpy as np
 
 import warp as wp
-from warp.fem import cache
-from warp.fem.geometry import Grid3D
-from warp.fem.polynomial import Polynomial, is_closed, lagrange_scales, quadrature_1d
-from warp.fem.types import Coords
+from warp._src.fem import cache
+from warp._src.fem.geometry import Grid3D
+from warp._src.fem.polynomial import Polynomial, is_closed, lagrange_scales, quadrature_1d
+from warp._src.fem.types import Coords
 
 from .shape_function import ShapeFunction
 from .tet_shape_function import TetrahedronPolynomialShapeFunctions
 
+_wp_module_name_ = "warp.fem.space.shape.cube_shape_function"
+
 
 class CubeShapeFunction(ShapeFunction):
     VERTEX = 0
@@ -79,7 +81,7 @@ class CubeTripolynomialShapeFunctions(CubeShapeFunction):
         lobatto_coords, lobatto_weight = quadrature_1d(point_count=degree + 1, family=family)
         lagrange_scale = lagrange_scales(lobatto_coords)
 
-        NodeVec = wp.types.vector(length=degree + 1, dtype=wp.float32)
+        NodeVec = cache.cached_vec_type(length=degree + 1, dtype=wp.float32)
         self.LOBATTO_COORDS = wp.constant(NodeVec(lobatto_coords))
         self.LOBATTO_WEIGHT = wp.constant(NodeVec(lobatto_weight))
         self.LAGRANGE_SCALE = wp.constant(NodeVec(lagrange_scale))
@@ -395,12 +397,12 @@ class CubeTripolynomialShapeFunctions(CubeShapeFunction):
         return cache.get_func(element_inner_weight_gradient, self.name)
 
     def element_node_hexes(self):
-        from warp.fem.utils import grid_to_hexes
+        from warp._src.fem.utils import grid_to_hexes
 
         return grid_to_hexes(self.ORDER, self.ORDER, self.ORDER)
 
     def element_node_tets(self):
-        from warp.fem.utils import grid_to_tets
+        from warp._src.fem.utils import grid_to_tets
 
         return grid_to_tets(self.ORDER, self.ORDER, self.ORDER)
 
@@ -514,7 +516,7 @@ class CubeSerendipityShapeFunctions(CubeShapeFunction):
         lobatto_coords, lobatto_weight = quadrature_1d(point_count=degree + 1, family=family)
         lagrange_scale = lagrange_scales(lobatto_coords)
 
-        NodeVec = wp.types.vector(length=degree + 1, dtype=wp.float32)
+        NodeVec = cache.cached_vec_type(length=degree + 1, dtype=wp.float32)
         self.LOBATTO_COORDS = wp.constant(NodeVec(lobatto_coords))
         self.LOBATTO_WEIGHT = wp.constant(NodeVec(lobatto_weight))
         self.LAGRANGE_SCALE = wp.constant(NodeVec(lagrange_scale))
@@ -751,7 +753,7 @@ class CubeSerendipityShapeFunctions(CubeShapeFunction):
         return element_inner_weight_gradient
 
     def element_node_tets(self):
-        from warp.fem.utils import grid_to_tets
+        from warp._src.fem.utils import grid_to_tets
 
         if self.ORDER == 2:
             element_tets = np.array(

warp/{fem → _src/fem}/space/shape/shape_function.py

@@ -18,9 +18,11 @@ from enum import Enum
 import numpy as np
 
 import warp as wp
-from warp.fem import cache
-from warp.fem.geometry import Element
-from warp.fem.types import Coords
+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:
@@ -68,19 +70,18 @@ class ShapeFunction:
 class ConstantShapeFunction(ShapeFunction):
     """Shape function that is constant over the element"""
 
-    def __init__(self, element: Element, space_dimension: int):
-        self._element = element
-        self._dimension = space_dimension
+    def __init__(self, element: Element):
+        self._element_prototype = element.prototype
 
         self.ORDER = wp.constant(0)
         self.NODES_PER_ELEMENT = wp.constant(1)
 
-        coords, _ = element.instantiate_quadrature(order=0, family=None)
+        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.__class__.__name__}{self._dimension}"
+        return f"{self._element_prototype.__name__}"
 
     def make_node_coords_in_element(self):
         COORDS = self.COORDS
@@ -116,7 +117,7 @@ class ConstantShapeFunction(ShapeFunction):
         return ConstantShapeFunction._element_inner_weight
 
     def make_element_inner_weight_gradient(self):
-        grad_type = wp.vec(length=self._dimension, dtype=float)
+        grad_type = cache.cached_vec_type(length=self._element_prototype.dimension, dtype=float)
 
         @cache.dynamic_func(suffix=self.name)
         def element_inner_weight_gradient(

warp/{fem → _src/fem}/space/shape/square_shape_function.py

@@ -18,13 +18,15 @@ import math
 import numpy as np
 
 import warp as wp
-from warp.fem import cache
-from warp.fem.polynomial import Polynomial, is_closed, lagrange_scales, quadrature_1d
-from warp.fem.types import Coords
+from warp._src.fem import cache
+from warp._src.fem.polynomial import Polynomial, is_closed, lagrange_scales, quadrature_1d
+from warp._src.fem.types import Coords
 
 from .shape_function import ShapeFunction
 from .triangle_shape_function import TrianglePolynomialShapeFunctions
 
+_wp_module_name_ = "warp.fem.space.shape.square_shape_function"
+
 
 class SquareShapeFunction(ShapeFunction):
     VERTEX = 0
@@ -68,7 +70,7 @@ class SquareBipolynomialShapeFunctions(SquareShapeFunction):
         lobatto_coords, lobatto_weight = quadrature_1d(point_count=degree + 1, family=family)
         lagrange_scale = lagrange_scales(lobatto_coords)
 
-        NodeVec = wp.types.vector(length=degree + 1, dtype=wp.float32)
+        NodeVec = cache.cached_vec_type(length=degree + 1, dtype=wp.float32)
         self.LOBATTO_COORDS = wp.constant(NodeVec(lobatto_coords))
         self.LOBATTO_WEIGHT = wp.constant(NodeVec(lobatto_weight))
         self.LAGRANGE_SCALE = wp.constant(NodeVec(lagrange_scale))
@@ -291,7 +293,7 @@ class SquareBipolynomialShapeFunctions(SquareShapeFunction):
         return cache.get_func(element_inner_weight_gradient, self.name)
 
     def element_node_triangulation(self):
-        from warp.fem.utils import grid_to_tris
+        from warp._src.fem.utils import grid_to_tris
 
         return grid_to_tris(self.ORDER, self.ORDER)
 
@@ -348,7 +350,7 @@ class SquareSerendipityShapeFunctions(SquareShapeFunction):
         lobatto_coords, lobatto_weight = quadrature_1d(point_count=degree + 1, family=family)
         lagrange_scale = lagrange_scales(lobatto_coords)
 
-        NodeVec = wp.types.vector(length=degree + 1, dtype=wp.float32)
+        NodeVec = cache.cached_vec_type(length=degree + 1, dtype=wp.float32)
         self.LOBATTO_COORDS = wp.constant(NodeVec(lobatto_coords))
         self.LOBATTO_WEIGHT = wp.constant(NodeVec(lobatto_weight))
         self.LAGRANGE_SCALE = wp.constant(NodeVec(lagrange_scale))

warp/{fem → _src/fem}/space/shape/tet_shape_function.py

@@ -16,11 +16,13 @@
 import numpy as np
 
 import warp as wp
-from warp.fem import cache
-from warp.fem.types import Coords
+from warp._src.fem import cache
+from warp._src.fem.types import Coords
 
 from .shape_function import ShapeFunction
 
+_wp_module_name_ = "warp.fem.space.shape.tet_shape_function"
+
 
 def _tet_node_index(tx: int, ty: int, tz: int, degree: int):
     from .triangle_shape_function import _triangle_node_index
@@ -182,7 +184,7 @@ class TetrahedronPolynomialShapeFunctions(TetrahedronShapeFunction):
                     index = _tet_node_index(tx, ty, tz, degree)
                     tet_coords[index] = [tx, ty, tz]
 
-        CoordTypeVec = wp.mat(dtype=int, shape=(self.NODES_PER_ELEMENT, 3))
+        CoordTypeVec = cache.cached_mat_type(dtype=int, shape=(self.NODES_PER_ELEMENT, 3))
         self.NODE_TET_COORDS = wp.constant(CoordTypeVec(tet_coords))
 
         self.node_type_and_type_index = self._get_node_type_and_type_index()

warp/{fem → _src/fem}/space/shape/triangle_shape_function.py

@@ -16,11 +16,13 @@
 import numpy as np
 
 import warp as wp
-from warp.fem import cache
-from warp.fem.types import Coords
+from warp._src.fem import cache
+from warp._src.fem.types import Coords
 
 from .shape_function import ShapeFunction
 
+_wp_module_name_ = "warp.fem.space.shape.triangle_shape_function"
+
 
 def _triangle_node_index(tx: int, ty: int, degree: int):
     VERTEX_NODE_COUNT = 3
@@ -95,7 +97,7 @@ class TrianglePolynomialShapeFunctions(TriangleShapeFunction):
                 index = _triangle_node_index(tx, ty, degree)
                 triangle_coords[index] = [tx, ty]
 
-        CoordTypeVec = wp.mat(dtype=int, shape=(self.NODES_PER_ELEMENT, 2))
+        CoordTypeVec = cache.cached_mat_type(dtype=int, shape=(self.NODES_PER_ELEMENT, 2))
         self.NODE_TRIANGLE_COORDS = wp.constant(CoordTypeVec(triangle_coords))
 
         self.node_type_and_type_index = self._get_node_type_and_type_index()

warp/{fem → _src/fem}/space/tetmesh_function_space.py

@@ -14,9 +14,9 @@
 # limitations under the License.
 
 import warp as wp
-from warp.fem import cache
-from warp.fem.geometry import Tetmesh
-from warp.fem.types import ElementIndex
+from warp._src.fem import cache
+from warp._src.fem.geometry import Tetmesh
+from warp._src.fem.types import ElementIndex
 
 from .shape import (
     ShapeFunction,
@@ -25,6 +25,8 @@ from .shape import (
 )
 from .topology import SpaceTopology, forward_base_topology
 
+_wp_module_name_ = "warp.fem.space.tetmesh_function_space"
+
 
 @wp.struct
 class TetmeshTopologyArg:
@@ -72,12 +74,6 @@ class TetmeshSpaceTopology(SpaceTopology):
     def name(self):
         return f"{self.geometry.name}_{self._shape.name}"
 
-    @cache.cached_arg_value
-    def topo_arg_value(self, device):
-        arg = TetmeshTopologyArg()
-        self.fill_topo_arg(arg, device)
-        return arg
-
     def fill_topo_arg(self, arg: TetmeshTopologyArg, device):
         arg.tet_face_indices = self._tet_face_indices.to(device)
         arg.tet_edge_indices = self._tet_edge_indices.to(device)