warp-lang 1.2.1__py3-none-manylinux2014_aarch64.whl → 1.3.0__py3-none-manylinux2014_aarch64.whl

warp/fem/space/topology.py

@@ -1,9 +1,9 @@
-from typing import Optional, Type
+from typing import Optional, Tuple, Type
 
 import warp as wp
 from warp.fem import cache
 from warp.fem.geometry import DeformedGeometry, Geometry
-from warp.fem.types import ElementIndex
+from warp.fem.types import NULL_ELEMENT_INDEX, NULL_NODE_INDEX, ElementIndex
 
 
 class SpaceTopology:
@@ -18,8 +18,8 @@ class SpaceTopology:
     dimension: int
     """Embedding dimension of the function space"""
 
-    NODES_PER_ELEMENT: int
-    """Number of interpolation nodes per element of the geometry.
+    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
     """
@@ -30,12 +30,14 @@ class SpaceTopology:
 
         pass
 
-    def __init__(self, geometry: Geometry, nodes_per_element: int):
+    def __init__(self, geometry: Geometry, max_nodes_per_element: int):
         self._geometry = geometry
         self.dimension = geometry.dimension
-        self.NODES_PER_ELEMENT = wp.constant(nodes_per_element)
+        self.MAX_NODES_PER_ELEMENT = wp.constant(max_nodes_per_element)
         self.ElementArg = geometry.CellArg
 
+        self._make_constant_element_node_count()
+
     @property
     def geometry(self) -> Geometry:
         """Underlying geometry"""
@@ -51,25 +53,42 @@ class SpaceTopology:
 
     @property
     def name(self):
-        return f"{self.__class__.__name__}_{self.NODES_PER_ELEMENT}"
+        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"""
 
-        NODES_PER_ELEMENT = self.NODES_PER_ELEMENT
+        MAX_NODES_PER_ELEMENT = self.MAX_NODES_PER_ELEMENT
 
         @cache.dynamic_kernel(suffix=self.name)
         def fill_element_node_indices(
@@ -78,12 +97,13 @@ class SpaceTopology:
             element_node_indices: wp.array2d(dtype=int),
         ):
             element_index = wp.tid()
-            for n in range(NODES_PER_ELEMENT):
+            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(), NODES_PER_ELEMENT)
+        shape = (self.geometry.cell_count(), MAX_NODES_PER_ELEMENT)
         if out is None:
             element_node_indices = wp.empty(
                 shape=shape,
@@ -135,14 +155,36 @@ class SpaceTopology:
             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
+
+        @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
+
+        self.element_node_count = constant_element_node_count
+        self.side_neighbor_node_counts = constant_side_neighbor_node_counts
+
 
 class TraceSpaceTopology(SpaceTopology):
     """Auto-generated trace topology defining the node indices associated to the geometry sides"""
 
     def __init__(self, topo: SpaceTopology):
-        super().__init__(topo.geometry, 2 * topo.NODES_PER_ELEMENT)
-
         self._topo = topo
+
+        super().__init__(topo.geometry, 2 * topo.MAX_NODES_PER_ELEMENT)
+
         self.dimension = topo.dimension - 1
         self.ElementArg = topo.geometry.SideArg
 
@@ -154,6 +196,8 @@ class TraceSpaceTopology(SpaceTopology):
         self.neighbor_cell_index = self._make_neighbor_cell_index()
 
         self.element_node_index = self._make_element_node_index()
+        self.element_node_count = self._make_element_node_count()
+        self.side_neighbor_node_counts = None
 
     def node_count(self) -> int:
         return self._topo.node_count()
@@ -163,39 +207,51 @@ class TraceSpaceTopology(SpaceTopology):
         return f"{self._topo.name}_Trace"
 
     def _make_inner_cell_index(self):
-        NODES_PER_ELEMENT = self._topo.NODES_PER_ELEMENT
-
         @cache.dynamic_func(suffix=self.name)
-        def inner_cell_index(args: self.geometry.SideArg, element_index: ElementIndex, node_index_in_elt: int):
-            index_in_inner_cell = wp.select(node_index_in_elt < NODES_PER_ELEMENT, -1, node_index_in_elt)
-            return self.geometry.side_inner_cell_index(args, element_index), index_in_inner_cell
+        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):
-        NODES_PER_ELEMENT = self._topo.NODES_PER_ELEMENT
-
         @cache.dynamic_func(suffix=self.name)
-        def outer_cell_index(args: self.geometry.SideArg, element_index: ElementIndex, node_index_in_elt: int):
-            return self.geometry.side_outer_cell_index(args, element_index), node_index_in_elt - NODES_PER_ELEMENT
+        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):
-        NODES_PER_ELEMENT = self._topo.NODES_PER_ELEMENT
-
         @cache.dynamic_func(suffix=self.name)
-        def neighbor_cell_index(args: self.geometry.SideArg, element_index: ElementIndex, node_index_in_elt: int):
-            if node_index_in_elt < NODES_PER_ELEMENT:
-                return self.geometry.side_inner_cell_index(args, element_index), node_index_in_elt
-            else:
-                return (
-                    self.geometry.side_outer_cell_index(args, element_index),
-                    node_index_in_elt - NODES_PER_ELEMENT,
-                )
+        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(
@@ -219,7 +275,7 @@ class TraceSpaceTopology(SpaceTopology):
         return self._topo == other._topo
 
 
-class DiscontinuousSpaceTopologyMixin:
+class RegularDiscontinuousSpaceTopologyMixin:
     """Helper for defining discontinuous topologies (per-element nodes)"""
 
     def __init__(self, *args, **kwargs):
@@ -227,14 +283,14 @@ class DiscontinuousSpaceTopologyMixin:
         self.element_node_index = self._make_element_node_index()
 
     def node_count(self):
-        return self.geometry.cell_count() * self.NODES_PER_ELEMENT
+        return self.geometry.cell_count() * self.MAX_NODES_PER_ELEMENT
 
     @property
     def name(self):
-        return f"{self.geometry.name}_D{self.NODES_PER_ELEMENT}"
+        return f"{self.geometry.name}_D{self.MAX_NODES_PER_ELEMENT}"
 
     def _make_element_node_index(self):
-        NODES_PER_ELEMENT = self.NODES_PER_ELEMENT
+        NODES_PER_ELEMENT = self.MAX_NODES_PER_ELEMENT
 
         @cache.dynamic_func(suffix=self.name)
         def element_node_index(
@@ -248,7 +304,7 @@ class DiscontinuousSpaceTopologyMixin:
         return element_node_index
 
 
-class DiscontinuousSpaceTopology(DiscontinuousSpaceTopologyMixin, SpaceTopology):
+class RegularDiscontinuousSpaceTopology(RegularDiscontinuousSpaceTopologyMixin, SpaceTopology):
     """Topology for generic discontinuous spaces"""
 
     pass
@@ -256,20 +312,20 @@ class DiscontinuousSpaceTopology(DiscontinuousSpaceTopologyMixin, SpaceTopology)
 
 class DeformedGeometrySpaceTopology(SpaceTopology):
     def __init__(self, geometry: DeformedGeometry, base_topology: SpaceTopology):
-        super().__init__(geometry, base_topology.NODES_PER_ELEMENT)
-
         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.TopologyArg = self.base.TopologyArg
 
-        self.element_node_index = self._make_element_node_index()
+        self._make_passthrough_functions()
 
     @property
     def name(self):
         return f"{self.base.name}_{self.geometry.field.name}"
 
-    def _make_element_node_index(self):
+    def _make_passthrough_functions(self):
         @cache.dynamic_func(suffix=self.name)
         def element_node_index(
             elt_arg: self.geometry.CellArg,
@@ -279,7 +335,25 @@ class DeformedGeometrySpaceTopology(SpaceTopology):
         ):
             return self.base.element_node_index(elt_arg.elt_arg, topo_arg, element_index, node_index_in_elt)
 
-        return element_node_index
+        @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.elt_arg, topo_arg, element_count)
+
+        @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
+
+        self.element_node_index = element_node_index
+        self.element_node_count = element_node_count
+        self.side_neighbor_node_counts = side_neighbor_node_counts
 
 
 def forward_base_topology(topology_class: Type[SpaceTopology], geometry: Geometry, *args, **kwargs) -> SpaceTopology:

warp/fem/types.py

@@ -1,3 +1,5 @@
+from enum import Enum
+
 import warp as wp
 
 # kept to avoid breaking existing example code, no longer used internally
@@ -14,7 +16,7 @@ NodeIndex = int
 
 NULL_ELEMENT_INDEX = wp.constant(-1)
 NULL_QP_INDEX = wp.constant(-1)
-NULL_NODE_INDEX = wp.constant(-1)
+NULL_NODE_INDEX = wp.constant((1 << 31) - 1)  # this should be larger than normal nodes when sorting
 
 DofIndex = wp.vec2i
 """Opaque descriptor for indexing degrees of freedom within elements"""
@@ -31,6 +33,11 @@ def get_node_coord(dof_idx: DofIndex):
     return dof_idx[1]
 
 
+class ElementKind(Enum):
+    CELL = 0
+    SIDE = 1
+
+
 @wp.struct
 class NodeElementIndex:
     domain_element_index: ElementIndex