warp-lang 1.10.0__py3-none-macosx_11_0_arm64.whl

warp/_src/fem/space/topology.py

@@ -0,0 +1,461 @@
+# 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 functools import cached_property
+from typing import ClassVar, Optional, Tuple, Type
+
+import warp as wp
+from warp._src.fem import cache
+from warp._src.fem.geometry import DeformedGeometry, Geometry
+from warp._src.fem.types import NULL_ELEMENT_INDEX, NULL_NODE_INDEX, ElementIndex
+
+_wp_module_name_ = "warp.fem.space.topology"
+
+
+class SpaceTopology:
+    """
+    Interface class for defining the topology of a function space.
+
+    The topology only considers the indices of the nodes in each element, and as such,
+    the connectivity pattern of the function space.
+    It does not specify the actual location of the nodes within the elements, or the valuation function.
+    """
+
+    dimension: int
+    """Embedding dimension of the function space"""
+
+    MAX_NODES_PER_ELEMENT: int
+    """maximum number of interpolation nodes per element of the geometry.
+
+    .. note:: This will change to be defined per-element in future versions
+    """
+
+    _dynamic_attribute_constructors: ClassVar = {
+        "element_node_count": lambda obj: obj._make_constant_element_node_count(),
+        "element_node_sign": lambda obj: obj._make_constant_element_node_sign(),
+        "side_neighbor_node_counts": lambda obj: obj._make_constant_side_neighbor_node_counts(),
+    }
+
+    @wp.struct
+    class TopologyArg:
+        """Structure containing arguments to be passed to device functions"""
+
+        pass
+
+    def __init__(self, geometry: Geometry, max_nodes_per_element: int):
+        self._geometry = geometry
+        self.dimension = geometry.dimension
+        self.MAX_NODES_PER_ELEMENT = wp.constant(max_nodes_per_element)
+        self.ElementArg = geometry.CellArg
+
+        cache.setup_dynamic_attributes(self, cls=__class__)
+
+    @property
+    def geometry(self) -> Geometry:
+        """Underlying geometry"""
+        return self._geometry
+
+    def node_count(self) -> int:
+        """Number of nodes in the interpolation basis"""
+        raise NotImplementedError
+
+    @cache.cached_arg_value
+    def topo_arg_value(self, device) -> "TopologyArg":
+        """Value of the topology argument structure to be passed to device functions"""
+        arg = self.TopologyArg()
+        self.fill_topo_arg(arg, device)
+        return arg
+
+    def fill_topo_arg(self, arg, device):
+        pass
+
+    @cached_property
+    def name(self):
+        return f"{self.__class__.__name__}_{self.MAX_NODES_PER_ELEMENT}"
+
+    def __str__(self):
+        return self.name
+
+    @staticmethod
+    def element_node_count(
+        geo_arg: "ElementArg",  # noqa: F821
+        topo_arg: "TopologyArg",
+        element_index: ElementIndex,
+    ) -> int:
+        """Returns the actual number of nodes in a given element"""
+        raise NotImplementedError
+
+    @staticmethod
+    def element_node_index(
+        geo_arg: "ElementArg",  # noqa: F821
+        topo_arg: "TopologyArg",
+        element_index: ElementIndex,
+        node_index_in_elt: int,
+    ) -> int:
+        """Global node index for a given node in a given element"""
+        raise NotImplementedError
+
+    @staticmethod
+    def side_neighbor_node_counts(
+        side_arg: "ElementArg",  # noqa: F821
+        side_index: ElementIndex,
+    ) -> Tuple[int, int]:
+        """Returns the number of nodes for both the inner and outer cells of a given sides"""
+        raise NotImplementedError
+
+    def element_node_indices(self, out: Optional[wp.array] = None) -> wp.array:
+        """Returns a temporary array containing the global index for each node of each element"""
+
+        MAX_NODES_PER_ELEMENT = self.MAX_NODES_PER_ELEMENT
+
+        @cache.dynamic_kernel(suffix=self.name)
+        def fill_element_node_indices(
+            geo_cell_arg: self.geometry.CellArg,
+            topo_arg: self.TopologyArg,
+            element_node_indices: wp.array2d(dtype=int),
+        ):
+            element_index = wp.tid()
+            element_node_count = self.element_node_count(geo_cell_arg, topo_arg, element_index)
+            for n in range(element_node_count):
+                element_node_indices[element_index, n] = self.element_node_index(
+                    geo_cell_arg, topo_arg, element_index, n
+                )
+
+        shape = (self.geometry.cell_count(), MAX_NODES_PER_ELEMENT)
+        if out is None:
+            element_node_indices = wp.empty(
+                shape=shape,
+                dtype=int,
+            )
+        else:
+            if out.shape != shape or out.dtype != wp.int32:
+                raise ValueError(f"Out element node indices array must have shape {shape} and data type 'int32'")
+            element_node_indices = out
+
+        wp.launch(
+            dim=element_node_indices.shape[0],
+            kernel=fill_element_node_indices,
+            inputs=[
+                self.geometry.cell_arg_value(device=element_node_indices.device),
+                self.topo_arg_value(device=element_node_indices.device),
+                element_node_indices,
+            ],
+            device=element_node_indices.device,
+        )
+
+        return element_node_indices
+
+    # Interface generating trace space topology
+
+    def trace(self) -> "TraceSpaceTopology":
+        """Trace of the function space over lower-dimensional elements of the geometry"""
+
+        return TraceSpaceTopology(self)
+
+    @property
+    def is_trace(self) -> bool:
+        """Whether this topology is defined on the trace of the geometry"""
+        return self.dimension == self.geometry.dimension - 1
+
+    def full_space_topology(self) -> "SpaceTopology":
+        """Returns the full space topology from which this topology is derived"""
+        return self
+
+    def __eq__(self, other: "SpaceTopology") -> bool:
+        """Checks whether two topologies are compatible"""
+        return self.geometry == other.geometry and self.name == other.name
+
+    def is_derived_from(self, other: "SpaceTopology") -> bool:
+        """Checks whether two topologies are equal, or `self` is the trace of `other`"""
+        if self.dimension == other.dimension:
+            return self == other
+        if self.dimension + 1 == other.dimension:
+            return self.full_space_topology() == other
+        return False
+
+    def _make_constant_element_node_count(self):
+        NODES_PER_ELEMENT = wp.constant(self.MAX_NODES_PER_ELEMENT)
+
+        @cache.dynamic_func(suffix=self.name)
+        def constant_element_node_count(
+            geo_arg: self.geometry.CellArg,
+            topo_arg: self.TopologyArg,
+            element_index: ElementIndex,
+        ):
+            return NODES_PER_ELEMENT
+
+        return constant_element_node_count
+
+    def _make_constant_side_neighbor_node_counts(self):
+        NODES_PER_ELEMENT = wp.constant(self.MAX_NODES_PER_ELEMENT)
+
+        @cache.dynamic_func(suffix=self.name)
+        def constant_side_neighbor_node_counts(
+            side_arg: self.geometry.SideArg,
+            element_index: ElementIndex,
+        ):
+            return NODES_PER_ELEMENT, NODES_PER_ELEMENT
+
+        return constant_side_neighbor_node_counts
+
+    def _make_constant_element_node_sign(self):
+        @cache.dynamic_func(suffix=self.name)
+        def constant_element_node_sign(
+            geo_arg: self.geometry.CellArg,
+            topo_arg: self.TopologyArg,
+            element_index: ElementIndex,
+            node_index_in_element: int,
+        ):
+            return 1.0
+
+        return constant_element_node_sign
+
+
+class TraceSpaceTopology(SpaceTopology):
+    """Auto-generated trace topology defining the node indices associated to the geometry sides"""
+
+    _dynamic_attribute_constructors: ClassVar = {
+        "inner_cell_index": lambda obj: obj._make_inner_cell_index(),
+        "outer_cell_index": lambda obj: obj._make_outer_cell_index(),
+        "neighbor_cell_index": lambda obj: obj._make_neighbor_cell_index(),
+        "element_node_index": lambda obj: obj._make_element_node_index(),
+        "element_node_count": lambda obj: obj._make_element_node_count(),
+        "element_node_sign": lambda obj: obj._make_element_node_sign(),
+    }
+
+    def __init__(self, topo: SpaceTopology):
+        self._topo = topo
+
+        super().__init__(topo.geometry, 2 * topo.MAX_NODES_PER_ELEMENT)
+
+        self.dimension = topo.dimension - 1
+        self.ElementArg = topo.geometry.SideArg
+
+        self.TopologyArg = topo.TopologyArg
+        self.topo_arg_value = topo.topo_arg_value
+        self.fill_topo_arg = topo.fill_topo_arg
+
+        self.side_neighbor_node_counts = None
+        cache.setup_dynamic_attributes(self, cls=__class__)
+
+    def node_count(self) -> int:
+        return self._topo.node_count()
+
+    @cached_property
+    def name(self):
+        return f"{self._topo.name}_Trace"
+
+    def _make_inner_cell_index(self):
+        @cache.dynamic_func(suffix=self.name)
+        def inner_cell_index(side_arg: self.geometry.SideArg, element_index: ElementIndex, node_index_in_elt: int):
+            inner_count, outer_count = self._topo.side_neighbor_node_counts(side_arg, element_index)
+            if node_index_in_elt >= inner_count:
+                return NULL_ELEMENT_INDEX, NULL_NODE_INDEX
+            return self.geometry.side_inner_cell_index(side_arg, element_index), node_index_in_elt
+
+        return inner_cell_index
+
+    def _make_outer_cell_index(self):
+        @cache.dynamic_func(suffix=self.name)
+        def outer_cell_index(side_arg: self.geometry.SideArg, element_index: ElementIndex, node_index_in_elt: int):
+            inner_count, outer_count = self._topo.side_neighbor_node_counts(side_arg, element_index)
+            if node_index_in_elt < inner_count:
+                return NULL_ELEMENT_INDEX, NULL_NODE_INDEX
+            return self.geometry.side_outer_cell_index(side_arg, element_index), node_index_in_elt - inner_count
+
+        return outer_cell_index
+
+    def _make_neighbor_cell_index(self):
+        @cache.dynamic_func(suffix=self.name)
+        def neighbor_cell_index(side_arg: self.geometry.SideArg, element_index: ElementIndex, node_index_in_elt: int):
+            inner_count, outer_count = self._topo.side_neighbor_node_counts(side_arg, element_index)
+            if node_index_in_elt < inner_count:
+                return self.geometry.side_inner_cell_index(side_arg, element_index), node_index_in_elt
+
+            return (
+                self.geometry.side_outer_cell_index(side_arg, element_index),
+                node_index_in_elt - inner_count,
+            )
+
+        return neighbor_cell_index
+
+    def _make_element_node_count(self):
+        @cache.dynamic_func(suffix=self.name)
+        def trace_element_node_count(
+            geo_side_arg: self.geometry.SideArg,
+            topo_arg: self._topo.TopologyArg,
+            element_index: ElementIndex,
+        ):
+            inner_count, outer_count = self._topo.side_neighbor_node_counts(geo_side_arg, element_index)
+            return inner_count + outer_count
+
+        return trace_element_node_count
+
+    def _make_element_node_index(self):
+        @cache.dynamic_func(suffix=self.name)
+        def trace_element_node_index(
+            geo_side_arg: self.geometry.SideArg,
+            topo_arg: self._topo.TopologyArg,
+            element_index: ElementIndex,
+            node_index_in_elt: int,
+        ):
+            cell_index, index_in_cell = self.neighbor_cell_index(geo_side_arg, element_index, node_index_in_elt)
+
+            geo_cell_arg = self.geometry.side_to_cell_arg(geo_side_arg)
+            return self._topo.element_node_index(geo_cell_arg, topo_arg, cell_index, index_in_cell)
+
+        return trace_element_node_index
+
+    def _make_element_node_sign(self):
+        @cache.dynamic_func(suffix=self.name)
+        def trace_element_node_sign(
+            geo_side_arg: self.geometry.SideArg,
+            topo_arg: self._topo.TopologyArg,
+            element_index: ElementIndex,
+            node_index_in_elt: int,
+        ):
+            cell_index, index_in_cell = self.neighbor_cell_index(geo_side_arg, element_index, node_index_in_elt)
+
+            geo_cell_arg = self.geometry.side_to_cell_arg(geo_side_arg)
+            return self._topo.element_node_sign(geo_cell_arg, topo_arg, cell_index, index_in_cell)
+
+        return trace_element_node_sign
+
+    def full_space_topology(self) -> SpaceTopology:
+        """Returns the full space topology from which this topology is derived"""
+        return self._topo
+
+    def __eq__(self, other: "TraceSpaceTopology") -> bool:
+        return self._topo == other._topo
+
+
+class RegularDiscontinuousSpaceTopologyMixin:
+    """Helper for defining discontinuous topologies (per-element nodes)"""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.element_node_index = self._make_element_node_index()
+
+    def node_count(self):
+        return self.geometry.cell_count() * self.MAX_NODES_PER_ELEMENT
+
+    @cached_property
+    def name(self):
+        return f"{self.geometry.name}_D{self.MAX_NODES_PER_ELEMENT}"
+
+    def _make_element_node_index(self):
+        NODES_PER_ELEMENT = self.MAX_NODES_PER_ELEMENT
+
+        @cache.dynamic_func(suffix=self.name)
+        def element_node_index(
+            elt_arg: self.geometry.CellArg,
+            topo_arg: self.TopologyArg,
+            element_index: ElementIndex,
+            node_index_in_elt: int,
+        ):
+            return NODES_PER_ELEMENT * element_index + node_index_in_elt
+
+        return element_node_index
+
+
+class RegularDiscontinuousSpaceTopology(RegularDiscontinuousSpaceTopologyMixin, SpaceTopology):
+    """Topology for generic discontinuous spaces"""
+
+    pass
+
+
+class DeformedGeometrySpaceTopology(SpaceTopology):
+    _dynamic_attribute_constructors: ClassVar = {
+        "element_node_index": lambda obj: obj._make_element_node_index(),
+        "element_node_count": lambda obj: obj._make_element_node_count(),
+        "element_node_sign": lambda obj: obj._make_element_node_sign(),
+        "side_neighbor_node_counts": lambda obj: obj._make_side_neighbor_node_counts(),
+    }
+
+    def __init__(self, geometry: DeformedGeometry, base_topology: SpaceTopology):
+        self.base = base_topology
+        super().__init__(geometry, base_topology.MAX_NODES_PER_ELEMENT)
+
+        self.node_count = self.base.node_count
+        self.topo_arg_value = self.base.topo_arg_value
+        self.fill_topo_arg = self.base.fill_topo_arg
+        self.TopologyArg = self.base.TopologyArg
+
+        cache.setup_dynamic_attributes(self, cls=__class__)
+
+    @cached_property
+    def name(self):
+        return f"{self.base.name}_{self.geometry.field.name}"
+
+    def _make_element_node_index(self):
+        @cache.dynamic_func(suffix=self.name)
+        def element_node_index(
+            elt_arg: self.geometry.CellArg,
+            topo_arg: self.TopologyArg,
+            element_index: ElementIndex,
+            node_index_in_elt: int,
+        ):
+            return self.base.element_node_index(elt_arg.base_arg, topo_arg, element_index, node_index_in_elt)
+
+        return element_node_index
+
+    def _make_element_node_count(self):
+        @cache.dynamic_func(suffix=self.name)
+        def element_node_count(
+            elt_arg: self.geometry.CellArg,
+            topo_arg: self.TopologyArg,
+            element_count: ElementIndex,
+        ):
+            return self.base.element_node_count(elt_arg.base_arg, topo_arg, element_count)
+
+        return element_node_count
+
+    def _make_side_neighbor_node_counts(self):
+        @cache.dynamic_func(suffix=self.name)
+        def side_neighbor_node_counts(
+            side_arg: self.geometry.SideArg,
+            element_index: ElementIndex,
+        ):
+            inner_count, outer_count = self.base.side_neighbor_node_counts(side_arg.base_arg, element_index)
+            return inner_count, outer_count
+
+        return side_neighbor_node_counts
+
+    def _make_element_node_sign(self):
+        @cache.dynamic_func(suffix=self.name)
+        def element_node_sign(
+            elt_arg: self.geometry.CellArg,
+            topo_arg: self.TopologyArg,
+            element_index: ElementIndex,
+            node_index_in_elt: int,
+        ):
+            return self.base.element_node_sign(elt_arg.base_arg, topo_arg, element_index, node_index_in_elt)
+
+        return element_node_sign
+
+
+def forward_base_topology(topology_class: Type[SpaceTopology], geometry: Geometry, *args, **kwargs) -> SpaceTopology:
+    """
+    If `geometry` is *not* a :class:`DeformedGeometry`, constructs a normal instance of `topology_class` over `geometry`, forwarding additional arguments.
+
+    If `geometry` *is* a :class:`DeformedGeometry`, constructs an instance of `topology_class` over the base (undeformed) geometry of `geometry`, then warp it
+    in a :class:`DeformedGeometrySpaceTopology` forwarding the calls to the underlying topology.
+    """
+
+    if isinstance(geometry, DeformedGeometry):
+        base_topo = topology_class(geometry.base, *args, **kwargs)
+        return DeformedGeometrySpaceTopology(geometry, base_topo)
+
+    return topology_class(geometry, *args, **kwargs)

warp/_src/fem/space/trimesh_function_space.py

@@ -0,0 +1,193 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 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 Trimesh
+from warp._src.fem.types import ElementIndex
+
+from .shape import TriangleShapeFunction
+from .topology import SpaceTopology, forward_base_topology
+
+_wp_module_name_ = "warp.fem.space.trimesh_function_space"
+
+
+@wp.struct
+class TrimeshTopologyArg:
+    edge_vertex_indices: wp.array(dtype=wp.vec2i)
+    tri_edge_indices: wp.array2d(dtype=int)
+
+    vertex_count: int
+    edge_count: int
+
+
+class TrimeshSpaceTopology(SpaceTopology):
+    TopologyArg = TrimeshTopologyArg
+
+    def __init__(self, mesh: Trimesh, shape: TriangleShapeFunction):
+        self._shape = shape
+        super().__init__(mesh, shape.NODES_PER_ELEMENT)
+        self._mesh = mesh
+
+        self._compute_tri_edge_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: TrimeshTopologyArg, device):
+        arg.tri_edge_indices = self._tri_edge_indices.to(device)
+        arg.edge_vertex_indices = self._mesh.edge_vertex_indices.to(device)
+
+        arg.vertex_count = self._mesh.vertex_count()
+        arg.edge_count = self._mesh.side_count()
+
+    def _compute_tri_edge_indices(self):
+        self._tri_edge_indices = wp.empty(
+            dtype=int, device=self._mesh.tri_vertex_indices.device, shape=(self._mesh.cell_count(), 3)
+        )
+
+        wp.launch(
+            kernel=TrimeshSpaceTopology._compute_tri_edge_indices_kernel,
+            dim=self._mesh.edge_tri_indices.shape,
+            device=self._mesh.tri_vertex_indices.device,
+            inputs=[
+                self._mesh.edge_tri_indices,
+                self._mesh.edge_vertex_indices,
+                self._mesh.tri_vertex_indices,
+                self._tri_edge_indices,
+            ],
+        )
+
+    @wp.func
+    def _find_edge_index_in_tri(
+        edge_vtx: wp.vec2i,
+        tri_vtx: wp.vec3i,
+    ):
+        for k in range(2):
+            if (edge_vtx[0] == tri_vtx[k] and edge_vtx[1] == tri_vtx[k + 1]) or (
+                edge_vtx[1] == tri_vtx[k] and edge_vtx[0] == tri_vtx[k + 1]
+            ):
+                return k
+        return 2
+
+    @wp.kernel
+    def _compute_tri_edge_indices_kernel(
+        edge_tri_indices: wp.array(dtype=wp.vec2i),
+        edge_vertex_indices: wp.array(dtype=wp.vec2i),
+        tri_vertex_indices: wp.array2d(dtype=int),
+        tri_edge_indices: wp.array2d(dtype=int),
+    ):
+        e = wp.tid()
+
+        edge_vtx = edge_vertex_indices[e]
+        edge_tris = edge_tri_indices[e]
+
+        t0 = edge_tris[0]
+        t0_vtx = wp.vec3i(tri_vertex_indices[t0, 0], tri_vertex_indices[t0, 1], tri_vertex_indices[t0, 2])
+        t0_edge = TrimeshSpaceTopology._find_edge_index_in_tri(edge_vtx, t0_vtx)
+        tri_edge_indices[t0, t0_edge] = e
+
+        t1 = edge_tris[1]
+        if t1 != t0:
+            t1_vtx = wp.vec3i(tri_vertex_indices[t1, 0], tri_vertex_indices[t1, 1], tri_vertex_indices[t1, 2])
+            t1_edge = TrimeshSpaceTopology._find_edge_index_in_tri(edge_vtx, t1_vtx)
+            tri_edge_indices[t1, t1_edge] = e
+
+    def node_count(self) -> int:
+        return (
+            self._mesh.vertex_count() * self._shape.VERTEX_NODE_COUNT
+            + self._mesh.side_count() * self._shape.EDGE_NODE_COUNT
+            + self._mesh.cell_count() * self._shape.INTERIOR_NODE_COUNT
+        )
+
+    def _make_element_node_index(self):
+        VERTEX_NODE_COUNT = self._shape.VERTEX_NODE_COUNT
+        INTERIOR_NODES_PER_SIDE = self._shape.EDGE_NODE_COUNT
+        INTERIOR_NODES_PER_CELL = self._shape.INTERIOR_NODE_COUNT
+
+        @cache.dynamic_func(suffix=self.name)
+        def element_node_index(
+            geo_arg: self.geometry.CellArg,
+            topo_arg: TrimeshTopologyArg,
+            element_index: ElementIndex,
+            node_index_in_elt: int,
+        ):
+            node_type, type_index = self._shape.node_type_and_type_index(node_index_in_elt)
+
+            if wp.static(VERTEX_NODE_COUNT > 0):
+                if node_type == TriangleShapeFunction.VERTEX:
+                    vertex = type_index // VERTEX_NODE_COUNT
+                    vertex_node = type_index - VERTEX_NODE_COUNT * vertex
+                    return geo_arg.topology.tri_vertex_indices[element_index][vertex] * VERTEX_NODE_COUNT + vertex_node
+
+            global_offset = topo_arg.vertex_count * VERTEX_NODE_COUNT
+
+            if wp.static(INTERIOR_NODES_PER_SIDE > 0):
+                if node_type == TriangleShapeFunction.EDGE:
+                    edge = type_index // INTERIOR_NODES_PER_SIDE
+                    edge_node = type_index - INTERIOR_NODES_PER_SIDE * edge
+
+                    global_edge_index = topo_arg.tri_edge_indices[element_index][edge]
+
+                    if (
+                        topo_arg.edge_vertex_indices[global_edge_index][0]
+                        != geo_arg.topology.tri_vertex_indices[element_index][edge]
+                    ):
+                        edge_node = INTERIOR_NODES_PER_SIDE - 1 - edge_node
+
+                    return global_offset + INTERIOR_NODES_PER_SIDE * global_edge_index + edge_node
+
+                global_offset += INTERIOR_NODES_PER_SIDE * topo_arg.edge_count
+
+            return global_offset + INTERIOR_NODES_PER_CELL * element_index + type_index
+
+        return element_node_index
+
+    def _make_element_node_sign(self):
+        INTERIOR_NODES_PER_SIDE = self._shape.EDGE_NODE_COUNT
+
+        @cache.dynamic_func(suffix=self.name)
+        def element_node_sign(
+            geo_arg: self.geometry.CellArg,
+            topo_arg: TrimeshTopologyArg,
+            element_index: ElementIndex,
+            node_index_in_elt: int,
+        ):
+            node_type, type_index = self._shape.node_type_and_type_index(node_index_in_elt)
+
+            if node_type == TriangleShapeFunction.EDGE:
+                edge = type_index // INTERIOR_NODES_PER_SIDE
+
+                global_edge_index = topo_arg.tri_edge_indices[element_index][edge]
+                return wp.where(
+                    topo_arg.edge_vertex_indices[global_edge_index][0]
+                    == geo_arg.topology.tri_vertex_indices[element_index][edge],
+                    1.0,
+                    -1.0,
+                )
+
+            return 1.0
+
+        return element_node_sign
+
+
+def make_trimesh_space_topology(mesh: Trimesh, shape: TriangleShapeFunction):
+    if isinstance(shape, TriangleShapeFunction):
+        return forward_base_topology(TrimeshSpaceTopology, mesh, shape)
+
+    raise ValueError(f"Unsupported shape function {shape.name}")