warp-lang 1.2.2__py3-none-win_amd64.whl → 1.3.0__py3-none-win_amd64.whl

warp/fem/space/restriction.py

@@ -1,7 +1,7 @@
 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 NodeElementIndex
+from warp.fem.types import NULL_NODE_INDEX, NodeElementIndex
 from warp.fem.utils import compress_node_indices
 
 from .partition import SpacePartition
@@ -36,7 +36,7 @@ class SpaceRestriction:
     def _compute_node_element_indices(self, device, temporary_store: TemporaryStore):
         from warp.fem import cache
 
-        NODES_PER_ELEMENT = self.space_topology.NODES_PER_ELEMENT
+        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}",
@@ -51,14 +51,17 @@ class SpaceRestriction:
         ):
             domain_element_index = wp.tid()
             element_index = self.domain.element_index(domain_index_arg, domain_element_index)
-            for n in range(NODES_PER_ELEMENT):
+            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):
+                element_node_indices[domain_element_index, n] = NULL_NODE_INDEX
 
         element_node_indices = borrow_temporary(
             temporary_store,
-            shape=(self.domain.element_count(), NODES_PER_ELEMENT),
+            shape=(self.domain.element_count(), MAX_NODES_PER_ELEMENT),
             dtype=int,
             device=device,
         )
@@ -83,7 +86,10 @@ class SpaceRestriction:
             self._node_count,
             self._dof_partition_indices,
         ) = compress_node_indices(
-            self.space_partition.node_count(), flattened_node_indices, temporary_store=temporary_store
+            self.space_partition.node_count(),
+            flattened_node_indices,
+            return_unique_nodes=True,
+            temporary_store=temporary_store,
         )
 
         # Extract element index and index in element
@@ -93,7 +99,7 @@ class SpaceRestriction:
             kernel=SpaceRestriction._split_vertex_element_index,
             dim=flattened_node_indices.shape,
             inputs=[
-                NODES_PER_ELEMENT,
+                MAX_NODES_PER_ELEMENT,
                 node_array_indices.array,
                 self._dof_element_indices.array,
                 self._dof_indices_in_element.array,
@@ -132,20 +138,17 @@ class SpaceRestriction:
         return arg
 
     @wp.func
-    def node_partition_index(args: NodeArg, node_index: int):
-        return args.dof_partition_indices[node_index]
+    def node_partition_index(args: NodeArg, restriction_node_index: int):
+        return args.dof_partition_indices[restriction_node_index]
 
     @wp.func
-    def node_element_count(args: NodeArg, node_index: int):
-        partition_node_index = SpaceRestriction.node_partition_index(args, node_index)
-        return args.dof_element_offsets[partition_node_index + 1] - args.dof_element_offsets[partition_node_index]
+    def node_element_range(args: NodeArg, partition_node_index: int):
+        return args.dof_element_offsets[partition_node_index], args.dof_element_offsets[partition_node_index + 1]
 
     @wp.func
-    def node_element_index(args: NodeArg, node_index: int, element_index: int):
-        partition_node_index = SpaceRestriction.node_partition_index(args, node_index)
-        offset = args.dof_element_offsets[partition_node_index] + element_index
-        domain_element_index = args.dof_element_indices[offset]
-        index_in_element = args.dof_indices_in_element[offset]
+    def node_element_index(args: NodeArg, node_element_offset: int):
+        domain_element_index = args.dof_element_indices[node_element_offset]
+        index_in_element = args.dof_indices_in_element[node_element_offset]
         return NodeElementIndex(domain_element_index, index_in_element)
 
     @wp.kernel

warp/fem/space/shape/cube_shape_function.py

@@ -98,7 +98,7 @@ class CubeTripolynomialShapeFunctions:
 
                 # x face
                 type_instance = mi
-                type_index = wp.select(mi == 1, (j - 1) * (ORDER - 1) + k - 1, (k - 1) * (ORDER - 1) + j - 1)
+                type_index = (j - 1) * (ORDER - 1) + k - 1
                 return CubeTripolynomialShapeFunctions.FACE, type_instance, type_index
 
             if zj + mj == 1:
@@ -110,13 +110,13 @@ class CubeTripolynomialShapeFunctions:
 
                 # y face
                 type_instance = 2 + mj
-                type_index = wp.select(mj == 1, (i - 1) * (ORDER - 1) + k - 1, (k - 1) * (ORDER - 1) + i - 1)
+                type_index = (k - 1) * (ORDER - 1) + i - 1
                 return CubeTripolynomialShapeFunctions.FACE, type_instance, type_index
 
             if zk + mk == 1:
                 # z face
                 type_instance = 4 + mk
-                type_index = wp.select(mk == 1, (j - 1) * (ORDER - 1) + i - 1, (i - 1) * (ORDER - 1) + j - 1)
+                type_index = (i - 1) * (ORDER - 1) + j - 1
                 return CubeTripolynomialShapeFunctions.FACE, type_instance, type_index
 
             type_index = ((i - 1) * (ORDER - 1) + (j - 1)) * (ORDER - 1) + k - 1
@@ -322,6 +322,87 @@ class CubeTripolynomialShapeFunctions:
 
         return grid_to_tets(self.ORDER, self.ORDER, self.ORDER)
 
+    def element_vtk_cells(self):
+        n = self.ORDER + 1
+
+        # vertices
+        cells = [
+            [
+                [0, 0, 0],
+                [n - 1, 0, 0],
+                [n - 1, n - 1, 0],
+                [0, n - 1, 0],
+                [0, 0, n - 1],
+                [n - 1, 0, n - 1],
+                [n - 1, n - 1, n - 1],
+                [0, n - 1, n - 1],
+            ]
+        ]
+
+        if self.ORDER == 1:
+            cell_type = 12  # vtk_hexahedron
+        else:
+            middle = np.arange(1, n - 1)
+            front = np.zeros(n - 2, dtype=int)
+            back = np.full(n - 2, n - 1)
+
+            # edges
+            cells.append(np.column_stack((middle, front, front)))
+            cells.append(np.column_stack((back, middle, front)))
+            cells.append(np.column_stack((middle, back, front)))
+            cells.append(np.column_stack((front, middle, front)))
+
+            cells.append(np.column_stack((middle, front, back)))
+            cells.append(np.column_stack((back, middle, back)))
+            cells.append(np.column_stack((middle, back, back)))
+            cells.append(np.column_stack((front, middle, back)))
+
+            cells.append(np.column_stack((front, front, middle)))
+            cells.append(np.column_stack((back, front, middle)))
+            cells.append(np.column_stack((back, back, middle)))
+            cells.append(np.column_stack((front, back, middle)))
+
+            # faces
+
+            face = np.meshgrid(middle, middle)
+            front = np.zeros((n - 2) ** 2, dtype=int)
+            back = np.full((n - 2) ** 2, n - 1)
+
+            # YZ
+            cells.append(
+                np.column_stack((front, face[0].flatten(), face[1].flatten())),
+            )
+            cells.append(
+                np.column_stack((back, face[0].flatten(), face[1].flatten())),
+            )
+            # XZ
+            cells.append(
+                np.column_stack((face[0].flatten(), front, face[1].flatten())),
+            )
+            cells.append(
+                np.column_stack((face[0].flatten(), back, face[1].flatten())),
+            )
+            # XY
+            cells.append(
+                np.column_stack((face[0].flatten(), face[1].flatten(), front)),
+            )
+            cells.append(
+                np.column_stack((face[0].flatten(), face[1].flatten(), back)),
+            )
+
+            # interior
+            interior = np.meshgrid(middle, middle, middle)
+            cells.append(
+                np.column_stack((interior[0].flatten(), interior[1].flatten(), interior[2].flatten())),
+            )
+
+            cell_type = 72  # vtk_lagrange_hexahedron
+
+        cells = np.concatenate(cells)
+        cell_indices = cells[:, 0] * n * n + cells[:, 1] * n + cells[:, 2]
+
+        return cell_indices[np.newaxis, :], np.array([cell_type], dtype=np.int8)
+
 
 class CubeSerendipityShapeFunctions:
     """
@@ -632,6 +713,12 @@ class CubeSerendipityShapeFunctions:
 
         raise NotImplementedError()
 
+    def element_vtk_cells(self):
+        tets = np.array(self.element_node_tets())
+        cell_type = 10  # VTK_TETRA
+
+        return tets, np.full(tets.shape[0], cell_type, dtype=np.int8)
+
 
 class CubeNonConformingPolynomialShapeFunctions:
     # embeds the largest regular tet centered at (0.5, 0.5, 0.5) into the reference cube
@@ -656,6 +743,7 @@ class CubeNonConformingPolynomialShapeFunctions:
         self.NODES_PER_ELEMENT = self._tet_shape.NODES_PER_ELEMENT
 
         self.element_node_tets = self._tet_shape.element_node_tets
+        self.element_vtk_cells = self._tet_shape.element_vtk_cells
 
     @property
     def name(self) -> str:

warp/fem/space/shape/shape_function.py

@@ -1,3 +1,5 @@
+import numpy as np
+
 import warp as wp
 from warp.fem import cache
 from warp.fem.geometry import Element
@@ -100,3 +102,8 @@ class ConstantShapeFunction:
             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)

warp/fem/space/shape/square_shape_function.py

@@ -206,6 +206,31 @@ class SquareBipolynomialShapeFunctions:
 
         return grid_to_tris(self.ORDER, self.ORDER)
 
+    def element_vtk_cells(self):
+        n = self.ORDER + 1
+
+        # vertices
+        cells = [[0, (n - 1) * n, n * n - 1, n - 1]]
+
+        if self.ORDER == 1:
+            cell_type = 9  # VTK_QUAD
+        else:
+            middle = np.arange(1, n - 1)
+
+            # edges
+            cells.append(middle * n)
+            cells.append(middle + (n - 1) * n)
+            cells.append(middle * n + n - 1)
+            cells.append(middle)
+
+            # faces
+            interior = np.broadcast_to(middle, (n - 2, n - 2))
+            cells.append((interior * n + interior.transpose()).flatten())
+
+            cell_type = 70  # VTK_LAGRANGE_QUADRILATERAL
+
+        return np.concatenate(cells)[np.newaxis, :], np.array([cell_type], dtype=np.int8)
+
 
 class SquareSerendipityShapeFunctions:
     """
@@ -501,6 +526,12 @@ class SquareSerendipityShapeFunctions:
 
         return element_triangles
 
+    def element_vtk_cells(self):
+        tris = np.array(self.element_node_triangulation())
+        cell_type = 5  # VTK_TRIANGLE
+
+        return tris, np.full(tris.shape[0], cell_type, dtype=np.int8)
+
 
 class SquareNonConformingPolynomialShapeFunctions:
     # embeds the largest equilateral triangle centered at (0.5, 0.5) into the reference square
@@ -516,6 +547,7 @@ class SquareNonConformingPolynomialShapeFunctions:
         self.NODES_PER_ELEMENT = self._tri_shape.NODES_PER_ELEMENT
 
         self.element_node_triangulation = self._tri_shape.element_node_triangulation
+        self.element_vtk_cells = self._tri_shape.element_vtk_cells
 
     @property
     def name(self) -> str:

warp/fem/space/shape/tet_shape_function.py

@@ -92,7 +92,6 @@ def _tet_node_index(tx: int, ty: int, tz: int, degree: int):
 
 
 class TetrahedronPolynomialShapeFunctions:
-    INVALID = wp.constant(-1)
     VERTEX = wp.constant(0)
     EDGE = wp.constant(1)
     FACE = wp.constant(2)
@@ -138,14 +137,10 @@ class TetrahedronPolynomialShapeFunctions:
 
     def _get_node_type_and_type_index(self):
         ORDER = self.ORDER
-        NODES_PER_ELEMENT = self.NODES_PER_ELEMENT
 
         def node_type_and_index(
             node_index_in_elt: int,
         ):
-            if node_index_in_elt < 0 or node_index_in_elt >= NODES_PER_ELEMENT:
-                return TetrahedronPolynomialShapeFunctions.INVALID, TetrahedronPolynomialShapeFunctions.INVALID
-
             if node_index_in_elt < 4:
                 return TetrahedronPolynomialShapeFunctions.VERTEX, node_index_in_elt
 
@@ -208,7 +203,7 @@ class TetrahedronPolynomialShapeFunctions:
 
     def make_trace_node_quadrature_weight(self):
         if self.ORDER == 3:
-            # P3 intrisic quadrature
+            # P3 intrinsic quadrature
             vertex_weight = 1.0 / 30
             edge_weight = 0.075
             interior_weight = 0.45
@@ -447,6 +442,14 @@ class TetrahedronPolynomialShapeFunctions:
 
         return np.array(element_tets)
 
+    def element_vtk_cells(self):
+        cells = np.arange(self.NODES_PER_ELEMENT)
+        if self.ORDER == 1:
+            cell_type = 10  # VTK_TETRA
+        else:
+            cell_type = 71  # VTK_LAGRANGE_TETRAHEDRON
+        return cells[np.newaxis, :], np.array([cell_type], dtype=np.int8)
+
 
 class TetrahedronNonConformingPolynomialShapeFunctions:
     def __init__(self, degree: int):
@@ -455,6 +458,7 @@ class TetrahedronNonConformingPolynomialShapeFunctions:
         self.NODES_PER_ELEMENT = self._tet_shape.NODES_PER_ELEMENT
 
         self.element_node_tets = self._tet_shape.element_node_tets
+        self.element_vtk_cells = self._tet_shape.element_vtk_cells
 
         if self.ORDER == 1:
             self._TET_SCALE = 0.4472135955  # so v at 0.5854101966249680 (order 2)
@@ -489,7 +493,7 @@ class TetrahedronNonConformingPolynomialShapeFunctions:
 
     def make_node_quadrature_weight(self):
         # Intrinsic quadrature -- precomputed integral of node shape functions
-        # over element. Order euqla to self.ORDER
+        # over element. Order equal to self.ORDER
 
         if self.ORDER == 2:
             vertex_weight = 0.07499641

warp/fem/space/shape/triangle_shape_function.py

@@ -103,7 +103,7 @@ class Triangle2DPolynomialShapeFunctions:
 
     def make_node_quadrature_weight(self):
         if self.ORDER == 3:
-            # P3 intrisic quadrature
+            # P3 intrinsic quadrature
             vertex_weight = 1.0 / 30
             edge_weight = 0.075
             interior_weight = 0.45
@@ -309,6 +309,14 @@ class Triangle2DPolynomialShapeFunctions:
 
         return np.array(element_triangles)
 
+    def element_vtk_cells(self):
+        cells = np.arange(self.NODES_PER_ELEMENT)
+        if self.ORDER == 1:
+            cell_type = 5  # VTK_TRIANGLE
+        else:
+            cell_type = 69  # VTK_LAGRANGE_TRIANGLE
+        return cells[np.newaxis, :], np.array([cell_type], dtype=np.int8)
+
 
 class Triangle2DNonConformingPolynomialShapeFunctions:
     def __init__(self, degree: int):
@@ -317,6 +325,7 @@ class Triangle2DNonConformingPolynomialShapeFunctions:
         self.NODES_PER_ELEMENT = self._tri_shape.NODES_PER_ELEMENT
 
         self.element_node_triangulation = self._tri_shape.element_node_triangulation
+        self.element_vtk_cells = self._tri_shape.element_vtk_cells
 
         # Coordinates (a, b, b) of embedded triangle
         if self.ORDER == 1: