warp-lang 1.10.0__py3-none-macosx_11_0_arm64.whl

warp/_src/fem/space/grid_2d_function_space.py

@@ -0,0 +1,179 @@
+# 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 numpy as np
+
+import warp as wp
+from warp._src.fem import cache
+from warp._src.fem.geometry import Grid2D
+from warp._src.fem.polynomial import is_closed
+from warp._src.fem.types import NULL_NODE_INDEX, ElementIndex
+
+from .shape import SquareBipolynomialShapeFunctions, SquareShapeFunction
+from .topology import SpaceTopology, forward_base_topology
+
+_wp_module_name_ = "warp.fem.space.grid_2d_function_space"
+
+
+class Grid2DSpaceTopology(SpaceTopology):
+    def __init__(self, grid: Grid2D, shape: SquareShapeFunction):
+        self._shape = shape
+        super().__init__(grid, shape.NODES_PER_ELEMENT)
+
+        self.element_node_index = self._make_element_node_index()
+
+    TopologyArg = Grid2D.SideArg
+
+    @property
+    def name(self):
+        return f"{self.geometry.name}_{self._shape.name}"
+
+    def fill_topo_arg(self, arg: Grid2D.SideArg, device):
+        self.geometry.fill_side_arg(arg, device)
+
+    def node_count(self) -> int:
+        return (
+            self.geometry.vertex_count() * self._shape.VERTEX_NODE_COUNT
+            + self.geometry.side_count() * self._shape.EDGE_NODE_COUNT
+            + self.geometry.cell_count() * self._shape.INTERIOR_NODE_COUNT
+        )
+
+    def _make_element_node_index(self):
+        VERTEX_NODE_COUNT = self._shape.VERTEX_NODE_COUNT
+        EDGE_NODE_COUNT = self._shape.EDGE_NODE_COUNT
+        INTERIOR_NODE_COUNT = self._shape.INTERIOR_NODE_COUNT
+
+        @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_instance, type_index = self._shape.node_type_and_type_index(node_index_in_elt)
+
+            if wp.static(VERTEX_NODE_COUNT > 0):
+                if node_type == SquareShapeFunction.VERTEX:
+                    return (
+                        Grid2DSpaceTopology._vertex_index(cell_arg, element_index, type_instance) * VERTEX_NODE_COUNT
+                        + type_index
+                    )
+
+            res = cell_arg.res
+            vertex_count = (res[0] + 1) * (res[1] + 1)
+            global_offset = vertex_count
+
+            if wp.static(INTERIOR_NODE_COUNT > 0):
+                if node_type == SquareShapeFunction.INTERIOR:
+                    return global_offset + element_index * INTERIOR_NODE_COUNT + type_index
+
+                cell_count = res[0] * res[1]
+                global_offset += INTERIOR_NODE_COUNT * cell_count
+
+            if wp.static(EDGE_NODE_COUNT > 0):
+                axis = 1 - (node_type - SquareShapeFunction.EDGE_X)
+
+                cell = Grid2D.get_cell(cell_arg.res, element_index)
+                origin = Grid2D.orient(axis, cell) + wp.vec2i(type_instance, 0)
+
+                side = Grid2D.Side(axis, origin)
+                side_index = Grid2D.side_index(topo_arg, side)
+
+                vertex_count = (res[0] + 1) * (res[1] + 1)
+
+                return global_offset + EDGE_NODE_COUNT * side_index + type_index
+
+            return NULL_NODE_INDEX  # unreachable
+
+        return element_node_index
+
+    @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 GridBipolynomialSpaceTopology(SpaceTopology):
+    def __init__(self, grid: Grid2D, shape: SquareBipolynomialShapeFunctions):
+        super().__init__(grid, shape.NODES_PER_ELEMENT)
+        self._shape = 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: self.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
+
+    def node_grid(self):
+        res = self.geometry.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.geometry.cell_size[0] + self.geometry.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.geometry.cell_size[1] + self.geometry.origin[1]
+
+        return np.meshgrid(X, Y, indexing="ij")
+
+
+def make_grid_2d_space_topology(grid: Grid2D, shape: SquareShapeFunction):
+    if isinstance(shape, SquareBipolynomialShapeFunctions) and is_closed(shape.family):
+        return forward_base_topology(GridBipolynomialSpaceTopology, grid, shape)
+
+    if isinstance(shape, SquareShapeFunction):
+        return forward_base_topology(Grid2DSpaceTopology, grid, shape)
+
+    raise ValueError(f"Unsupported shape function {shape.name}")

warp/_src/fem/space/grid_3d_function_space.py

@@ -0,0 +1,229 @@
+# 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 numpy as np
+
+import warp as wp
+from warp._src.fem import cache
+from warp._src.fem.geometry import Grid3D
+from warp._src.fem.polynomial import is_closed
+from warp._src.fem.types import ElementIndex
+
+from .shape import (
+    CubeShapeFunction,
+    CubeTripolynomialShapeFunctions,
+)
+from .topology import SpaceTopology, forward_base_topology
+
+_wp_module_name_ = "warp.fem.space.grid_3d_function_space"
+
+
+class Grid3DSpaceTopology(SpaceTopology):
+    def __init__(self, grid: Grid3D, shape: CubeShapeFunction):
+        self._shape = shape
+        super().__init__(grid, shape.NODES_PER_ELEMENT)
+        self.element_node_index = self._make_element_node_index()
+
+    @property
+    def name(self):
+        return f"{self.geometry.name}_{self._shape.name}"
+
+    @wp.func
+    def _vertex_coords(vidx_in_cell: int):
+        x = vidx_in_cell // 4
+        y = (vidx_in_cell - 4 * x) // 2
+        z = vidx_in_cell - 4 * x - 2 * y
+        return wp.vec3i(x, y, z)
+
+    @wp.func
+    def _vertex_index(cell_arg: Grid3D.CellArg, cell_index: ElementIndex, vidx_in_cell: int):
+        res = cell_arg.res
+        strides = wp.vec2i((res[1] + 1) * (res[2] + 1), res[2] + 1)
+
+        corner = Grid3D.get_cell(res, cell_index) + Grid3DSpaceTopology._vertex_coords(vidx_in_cell)
+        return Grid3D._from_3d_index(strides, corner)
+
+    def node_count(self) -> int:
+        return (
+            self.geometry.vertex_count() * self._shape.VERTEX_NODE_COUNT
+            + self.geometry.edge_count() * self._shape.EDGE_NODE_COUNT
+            + self.geometry.side_count() * self._shape.FACE_NODE_COUNT
+            + self.geometry.cell_count() * self._shape.INTERIOR_NODE_COUNT
+        )
+
+    def _make_element_node_index(self):
+        VERTEX_NODE_COUNT = self._shape.VERTEX_NODE_COUNT
+        EDGE_NODE_COUNT = self._shape.EDGE_NODE_COUNT
+        FACE_NODE_COUNT = self._shape.FACE_NODE_COUNT
+        INTERIOR_NODE_COUNT = self._shape.INTERIOR_NODE_COUNT
+
+        @cache.dynamic_func(suffix=self.name)
+        def element_node_index(
+            cell_arg: Grid3D.CellArg,
+            topo_arg: Grid3DSpaceTopology.TopologyArg,
+            element_index: ElementIndex,
+            node_index_in_elt: int,
+        ):
+            res = cell_arg.res
+            cell = Grid3D.get_cell(res, element_index)
+
+            node_type, type_instance, type_index = self._shape.node_type_and_type_index(node_index_in_elt)
+
+            if wp.static(VERTEX_NODE_COUNT > 0):
+                if node_type == CubeShapeFunction.VERTEX:
+                    return (
+                        Grid3DSpaceTopology._vertex_index(cell_arg, element_index, type_instance) * VERTEX_NODE_COUNT
+                        + type_index
+                    )
+
+            res = cell_arg.res
+            vertex_count = (res[0] + 1) * (res[1] + 1) * (res[2] + 1)
+            global_offset = vertex_count * VERTEX_NODE_COUNT
+
+            if wp.static(EDGE_NODE_COUNT > 0):
+                if node_type == CubeShapeFunction.EDGE:
+                    axis = CubeShapeFunction._edge_axis(type_instance)
+                    node_all = CubeShapeFunction._edge_coords(type_instance, type_index)
+
+                    res = cell_arg.res
+
+                    edge_index = 0
+                    if axis > 0:
+                        edge_index += (res[1] + 1) * (res[2] + 1) * res[0]
+                    if axis > 1:
+                        edge_index += (res[0] + 1) * (res[2] + 1) * res[1]
+
+                    res_loc = Grid3D._world_to_local(axis, res)
+                    cell_loc = Grid3D._world_to_local(axis, cell)
+
+                    edge_index += (res_loc[1] + 1) * (res_loc[2] + 1) * cell_loc[0]
+                    edge_index += (res_loc[2] + 1) * (cell_loc[1] + node_all[1])
+                    edge_index += cell_loc[2] + node_all[2]
+
+                    return global_offset + EDGE_NODE_COUNT * edge_index + type_index
+
+                edge_count = (
+                    (res[0] + 1) * (res[1] + 1) * (res[2])
+                    + (res[0]) * (res[1] + 1) * (res[2] + 1)
+                    + (res[0] + 1) * (res[1]) * (res[2] + 1)
+                )
+                global_offset += edge_count * EDGE_NODE_COUNT
+
+            if wp.static(FACE_NODE_COUNT > 0):
+                if node_type == CubeShapeFunction.FACE:
+                    axis = CubeShapeFunction._face_axis(type_instance)
+                    face_offset = CubeShapeFunction._face_offset(type_instance)
+
+                    face_index = 0
+                    if axis > 0:
+                        face_index += (res[0] + 1) * res[1] * res[2]
+                    if axis > 1:
+                        face_index += (res[1] + 1) * res[2] * res[0]
+
+                    res_loc = Grid3D._world_to_local(axis, res)
+                    cell_loc = Grid3D._world_to_local(axis, cell)
+
+                    face_index += res_loc[1] * res_loc[2] * (cell_loc[0] + face_offset)
+                    face_index += res_loc[2] * cell_loc[1]
+                    face_index += cell_loc[2]
+
+                    return global_offset + FACE_NODE_COUNT * face_index + type_index
+
+                face_count = (
+                    (res[0] + 1) * res[1] * res[2] + res[0] * (res[1] + 1) * res[2] + res[0] * res[1] * (res[2] + 1)
+                )
+                global_offset += face_count * FACE_NODE_COUNT
+
+            # interior
+            return global_offset + element_index * INTERIOR_NODE_COUNT + type_index
+
+        return element_node_index
+
+
+class GridTripolynomialSpaceTopology(SpaceTopology):
+    def __init__(self, grid: Grid3D, shape: CubeTripolynomialShapeFunctions):
+        super().__init__(grid, shape.NODES_PER_ELEMENT)
+        self._shape = 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)
+            * (self.geometry.res[2] * 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: Grid3D.CellArg,
+            topo_arg: self.TopologyArg,
+            element_index: ElementIndex,
+            node_index_in_elt: int,
+        ):
+            res = cell_arg.res
+            cell = Grid3D.get_cell(res, element_index)
+
+            node_i, node_j, node_k = self._shape._node_ijk(node_index_in_elt)
+
+            node_x = ORDER * cell[0] + node_i
+            node_y = ORDER * cell[1] + node_j
+            node_z = ORDER * cell[2] + node_k
+
+            node_pitch_y = (res[2] * ORDER) + 1
+            node_pitch_x = node_pitch_y * ((res[1] * ORDER) + 1)
+            node_index = node_pitch_x * node_x + node_pitch_y * node_y + node_z
+
+            return node_index
+
+        return element_node_index
+
+    def node_grid(self):
+        res = self.geometry.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.geometry.cell_size[0] + self.geometry.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.geometry.cell_size[1] + self.geometry.origin[1]
+
+        grid_coords_z = np.repeat(np.arange(0, res[2], dtype=float), len(cell_coords)) + np.tile(
+            cell_coords, reps=res[2]
+        )
+        grid_coords_z = np.append(grid_coords_z, res[2])
+        Z = grid_coords_z * self.geometry.cell_size[2] + self.geometry.origin[2]
+
+        return np.meshgrid(X, Y, Z, indexing="ij")
+
+
+def make_grid_3d_space_topology(grid: Grid3D, shape: CubeShapeFunction):
+    if isinstance(shape, CubeTripolynomialShapeFunctions) and is_closed(shape.family):
+        return forward_base_topology(GridTripolynomialSpaceTopology, grid, shape)
+
+    if isinstance(shape, CubeShapeFunction):
+        return forward_base_topology(Grid3DSpaceTopology, grid, shape)
+
+    raise ValueError(f"Unsupported shape function {shape.name}")

warp/_src/fem/space/hexmesh_function_space.py

@@ -0,0 +1,255 @@
+# 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 warp as wp
+from warp._src.fem import cache
+from warp._src.fem.geometry import Hexmesh
+from warp._src.fem.geometry.hexmesh import (
+    EDGE_VERTEX_INDICES,
+    FACE_ORIENTATION,
+    FACE_TRANSLATION,
+)
+from warp._src.fem.types import ElementIndex
+
+from .shape import CubeShapeFunction
+from .topology import SpaceTopology, forward_base_topology
+
+_wp_module_name_ = "warp.fem.space.hexmesh_function_space"
+
+_FACE_ORIENTATION_I = wp.constant(wp.mat(shape=(16, 2), dtype=int)(FACE_ORIENTATION))
+_FACE_TRANSLATION_I = wp.constant(wp.mat(shape=(4, 2), dtype=int)(FACE_TRANSLATION))
+
+# map from shape function vertex indexing to hexmesh vertex indexing
+_CUBE_TO_HEX_VERTEX = wp.constant(wp.vec(length=8, dtype=int)([0, 4, 3, 7, 1, 5, 2, 6]))
+
+# map from shape function edge indexing to hexmesh edge indexing
+_CUBE_TO_HEX_EDGE = wp.constant(wp.vec(length=12, dtype=int)([0, 4, 2, 6, 3, 1, 7, 5, 8, 11, 9, 10]))
+
+
+@wp.struct
+class HexmeshTopologyArg:
+    hex_edge_indices: wp.array2d(dtype=int)
+    hex_face_indices: wp.array2d(dtype=wp.vec2i)
+
+    vertex_count: int
+    edge_count: int
+    face_count: int
+
+
+class HexmeshSpaceTopology(SpaceTopology):
+    TopologyArg = HexmeshTopologyArg
+
+    def __init__(
+        self,
+        mesh: Hexmesh,
+        shape: CubeShapeFunction,
+    ):
+        self._shape = shape
+        super().__init__(mesh, shape.NODES_PER_ELEMENT)
+        self._mesh = mesh
+
+        need_edge_indices = shape.EDGE_NODE_COUNT > 0
+        need_face_indices = shape.FACE_NODE_COUNT > 0
+
+        if need_edge_indices:
+            self._hex_edge_indices = self._mesh.hex_edge_indices
+            self._edge_count = self._mesh.edge_count()
+        else:
+            self._hex_edge_indices = wp.empty(shape=(0, 0), dtype=int)
+            self._edge_count = 0
+
+        if need_face_indices:
+            self._compute_hex_face_indices()
+        else:
+            self._hex_face_indices = wp.empty(shape=(0, 0), dtype=wp.vec2i)
+
+        self._compute_hex_face_indices()
+
+        self.element_node_index = self._make_element_node_index()
+        self.element_node_sign = self._make_element_node_sign()
+
+    @property
+    def name(self):
+        return f"{self.geometry.name}_{self._shape.name}"
+
+    def fill_topo_arg(self, arg: HexmeshTopologyArg, device):
+        arg.hex_edge_indices = self._hex_edge_indices.to(device)
+        arg.hex_face_indices = self._hex_face_indices.to(device)
+        arg.vertex_count = self._mesh.vertex_count()
+        arg.face_count = self._mesh.side_count()
+        arg.edge_count = self._edge_count
+
+    def _compute_hex_face_indices(self):
+        self._hex_face_indices = wp.empty(
+            dtype=wp.vec2i, device=self._mesh.hex_vertex_indices.device, shape=(self._mesh.cell_count(), 6)
+        )
+
+        wp.launch(
+            kernel=HexmeshSpaceTopology._compute_hex_face_indices_kernel,
+            dim=self._mesh.side_count(),
+            device=self._mesh.hex_vertex_indices.device,
+            inputs=[
+                self._mesh.face_hex_indices,
+                self._mesh._face_hex_face_orientation,
+                self._hex_face_indices,
+            ],
+        )
+
+    @wp.kernel
+    def _compute_hex_face_indices_kernel(
+        face_hex_indices: wp.array(dtype=wp.vec2i),
+        face_hex_face_ori: wp.array(dtype=wp.vec4i),
+        hex_face_indices: wp.array2d(dtype=wp.vec2i),
+    ):
+        f = wp.tid()
+
+        # face indices from CubeShapeFunction always have positive orientation,
+        # while Hexmesh faces are oriented to point "outside"
+        # We need to flip orientation for faces at offset 0
+
+        hx0 = face_hex_indices[f][0]
+        local_face_0 = face_hex_face_ori[f][0]
+        ori_0 = face_hex_face_ori[f][1]
+
+        local_face_offset_0 = CubeShapeFunction._face_offset(local_face_0)
+        flip_0 = ori_0 ^ (1 - local_face_offset_0)
+
+        hex_face_indices[hx0, local_face_0] = wp.vec2i(f, flip_0)
+
+        hx1 = face_hex_indices[f][1]
+        local_face_1 = face_hex_face_ori[f][2]
+        ori_1 = face_hex_face_ori[f][3]
+
+        local_face_offset_1 = CubeShapeFunction._face_offset(local_face_1)
+        flip_1 = ori_1 ^ (1 - local_face_offset_1)
+
+        hex_face_indices[hx1, local_face_1] = wp.vec2i(f, flip_1)
+
+    def node_count(self) -> int:
+        return (
+            self._mesh.vertex_count() * self._shape.VERTEX_NODE_COUNT
+            + self._mesh.edge_count() * self._shape.EDGE_NODE_COUNT
+            + self._mesh.side_count() * self._shape.FACE_NODE_COUNT
+            + self._mesh.cell_count() * self._shape.INTERIOR_NODE_COUNT
+        )
+
+    @wp.func
+    def _rotate_face_coordinates(ori: int, offset: int, coords: wp.vec2i):
+        fv = ori // 2
+
+        rot_i = wp.dot(_FACE_ORIENTATION_I[2 * ori], coords)
+        rot_j = wp.dot(_FACE_ORIENTATION_I[2 * ori + 1], coords)
+
+        return wp.vec2i(rot_i, rot_j) + _FACE_TRANSLATION_I[fv]
+
+    def _make_element_node_index(self):
+        VERTEX_NODE_COUNT = self._shape.VERTEX_NODE_COUNT
+        EDGE_NODE_COUNT = self._shape.EDGE_NODE_COUNT
+        FACE_NODE_COUNT = self._shape.FACE_NODE_COUNT
+        INTERIOR_NODE_COUNT = self._shape.INTERIOR_NODE_COUNT
+
+        @cache.dynamic_func(suffix=self.name)
+        def element_node_index(
+            geo_arg: Hexmesh.CellArg,
+            topo_arg: HexmeshTopologyArg,
+            element_index: ElementIndex,
+            node_index_in_elt: int,
+        ):
+            node_type, type_instance, type_index = self._shape.node_type_and_type_index(node_index_in_elt)
+
+            if wp.static(VERTEX_NODE_COUNT > 0):
+                if node_type == CubeShapeFunction.VERTEX:
+                    return (
+                        geo_arg.hex_vertex_indices[element_index, _CUBE_TO_HEX_VERTEX[type_instance]]
+                        * VERTEX_NODE_COUNT
+                        + type_index
+                    )
+
+            offset = topo_arg.vertex_count * VERTEX_NODE_COUNT
+
+            if wp.static(EDGE_NODE_COUNT > 0):
+                if node_type == CubeShapeFunction.EDGE:
+                    hex_edge = _CUBE_TO_HEX_EDGE[type_instance]
+                    edge_index = topo_arg.hex_edge_indices[element_index, hex_edge]
+
+                    v0 = geo_arg.hex_vertex_indices[element_index, EDGE_VERTEX_INDICES[hex_edge, 0]]
+                    v1 = geo_arg.hex_vertex_indices[element_index, EDGE_VERTEX_INDICES[hex_edge, 1]]
+
+                    if v0 > v1:
+                        type_index = EDGE_NODE_COUNT - 1 - type_index
+
+                    return offset + EDGE_NODE_COUNT * edge_index + type_index
+
+                offset += EDGE_NODE_COUNT * topo_arg.edge_count
+
+            if wp.static(FACE_NODE_COUNT > 0):
+                if node_type == CubeShapeFunction.FACE:
+                    face_index_and_ori = topo_arg.hex_face_indices[element_index, type_instance]
+                    face_index = face_index_and_ori[0]
+                    face_orientation = face_index_and_ori[1]
+
+                    face_offset = CubeShapeFunction._face_offset(type_instance)
+
+                    if wp.static(FACE_NODE_COUNT > 1):
+                        face_coords = self._shape._face_node_ij(type_index)
+                        rot_face_coords = HexmeshSpaceTopology._rotate_face_coordinates(
+                            face_orientation, face_offset, face_coords
+                        )
+                        type_index = self._shape._linear_face_node_index(type_index, rot_face_coords)
+
+                    return offset + FACE_NODE_COUNT * face_index + type_index
+
+                offset += FACE_NODE_COUNT * topo_arg.face_count
+
+            return offset + INTERIOR_NODE_COUNT * element_index + type_index
+
+        return element_node_index
+
+    def _make_element_node_sign(self):
+        EDGE_NODE_COUNT = self._shape.EDGE_NODE_COUNT
+        FACE_NODE_COUNT = self._shape.FACE_NODE_COUNT
+
+        @cache.dynamic_func(suffix=self.name)
+        def element_node_sign(
+            geo_arg: self.geometry.CellArg,
+            topo_arg: HexmeshTopologyArg,
+            element_index: ElementIndex,
+            node_index_in_elt: int,
+        ):
+            node_type, type_instance, type_index = self._shape.node_type_and_type_index(node_index_in_elt)
+
+            if wp.static(EDGE_NODE_COUNT > 0):
+                if node_type == CubeShapeFunction.EDGE:
+                    hex_edge = _CUBE_TO_HEX_EDGE[type_instance]
+                    v0 = geo_arg.hex_vertex_indices[element_index, EDGE_VERTEX_INDICES[hex_edge, 0]]
+                    v1 = geo_arg.hex_vertex_indices[element_index, EDGE_VERTEX_INDICES[hex_edge, 1]]
+                    return wp.where(v0 > v1, -1.0, 1.0)
+
+            if wp.static(FACE_NODE_COUNT > 0):
+                if node_type == CubeShapeFunction.FACE:
+                    face_index_and_ori = topo_arg.hex_face_indices[element_index, type_instance]
+                    flip = face_index_and_ori[1] & 1
+                    return wp.where(flip == 0, 1.0, -1.0)
+
+            return 1.0
+
+        return element_node_sign
+
+
+def make_hexmesh_space_topology(mesh: Hexmesh, shape: CubeShapeFunction):
+    if isinstance(shape, CubeShapeFunction):
+        return forward_base_topology(HexmeshSpaceTopology, mesh, shape)
+
+    raise ValueError(f"Unsupported shape function {shape.name}")