warp-lang 1.0.1__py3-none-manylinux2014_aarch64.whl → 1.1.0__py3-none-manylinux2014_aarch64.whl

warp/fem/geometry/partition.py

@@ -1,376 +1,374 @@
-from typing import Any
-
-import warp as wp
-
-from warp.fem.types import ElementIndex, NULL_ELEMENT_INDEX
-from warp.fem.utils import masked_indices
-from warp.fem.cache import cached_arg_value, TemporaryStore, borrow_temporary
-
-from .geometry import Geometry
-
-
-wp.set_module_options({"enable_backward": False})
-
-
-class GeometryPartition:
-    """Base class for geometry partitions, i.e. subset of cells and sides"""
-
-    class CellArg:
-        pass
-
-    class SideArg:
-        pass
-
-    def __init__(self, geometry: Geometry):
-        self.geometry = geometry
-
-    def cell_count(self) -> int:
-        """Number of cells that are 'owned' by this partition"""
-        raise NotImplementedError()
-
-    def side_count(self) -> int:
-        """Number of sides that are 'owned' by this partition"""
-        raise NotImplementedError()
-
-    def boundary_side_count(self) -> int:
-        """Number of geo-boundary sides that are 'owned' by this partition"""
-        raise NotImplementedError()
-
-    def frontier_side_count(self) -> int:
-        """Number of sides with neighbors owned by this and another partition"""
-        raise NotImplementedError()
-
-    @property
-    def name(self) -> str:
-        return f"{self.geometry.name}_{self.__class__.__name__}"
-
-    def __str__(self) -> str:
-        return self.name
-
-    def cell_arg_value(self, device):
-        raise NotImplementedError()
-
-    def side_arg_value(self, device):
-        raise NotImplementedError()
-
-    @staticmethod
-    def cell_index(args: CellArg, partition_cell_index: int):
-        """Index in the geometry of a partition cell"""
-        raise NotImplementedError()
-
-    @staticmethod
-    def partition_cell_index(args: CellArg, cell_index: int):
-        """Index of a geometry cell in the partition (or ``NULL_ELEMENT_INDEX``)"""
-        raise NotImplementedError()
-
-    @staticmethod
-    def side_index(args: SideArg, partition_side_index: int):
-        """Partition side to side index"""
-        raise NotImplementedError()
-
-    @staticmethod
-    def boundary_side_index(args: SideArg, boundary_side_index: int):
-        """Boundary side to side index"""
-        raise NotImplementedError()
-
-    @staticmethod
-    def frontier_side_index(args: SideArg, frontier_side_index: int):
-        """Frontier side to side index"""
-        raise NotImplementedError()
-
-
-class WholeGeometryPartition(GeometryPartition):
-    """Trivial (NOP) partition"""
-
-    def __init__(
-        self,
-        geometry: Geometry,
-    ):
-        super().__init__(geometry)
-
-        self.SideArg = geometry.SideIndexArg
-        self.side_arg_value = geometry.side_index_arg_value
-
-        self.cell_index = WholeGeometryPartition._identity_element_index
-        self.partition_cell_index = WholeGeometryPartition._identity_element_index
-
-        self.side_index = WholeGeometryPartition._identity_element_index
-        self.boundary_side_index = geometry.boundary_side_index
-        self.frontier_side_index = WholeGeometryPartition._identity_element_index
-
-    def __eq__(self, other: GeometryPartition) -> bool:
-        # Ensures that two whole partition instances of the same geometry are considered equal
-        return isinstance(other, WholeGeometryPartition) and self.geometry == other.geometry
-
-    def cell_count(self) -> int:
-        return self.geometry.cell_count()
-
-    def side_count(self) -> int:
-        return self.geometry.side_count()
-
-    def boundary_side_count(self) -> int:
-        return self.geometry.boundary_side_count()
-
-    def frontier_side_count(self) -> int:
-        return 0
-
-    @wp.struct
-    class CellArg:
-        pass
-
-    def cell_arg_value(self, device):
-        arg = WholeGeometryPartition.CellArg()
-        return arg
-
-    @wp.func
-    def _identity_element_index(args: Any, idx: ElementIndex):
-        return idx
-
-    @property
-    def name(self) -> str:
-        return self.geometry.name
-
-
-class CellBasedGeometryPartition(GeometryPartition):
-    """Geometry partition based on a subset of cells. Interior, boundary and frontier sides are automatically categorized."""
-
-    def __init__(
-        self,
-        geometry: Geometry,
-        device=None,
-    ):
-        super().__init__(geometry)
-
-    @wp.struct
-    class SideArg:
-        partition_side_indices: wp.array(dtype=int)
-        boundary_side_indices: wp.array(dtype=int)
-        frontier_side_indices: wp.array(dtype=int)
-
-    def side_count(self) -> int:
-        return self._partition_side_indices.array.shape[0]
-
-    def boundary_side_count(self) -> int:
-        return self._boundary_side_indices.array.shape[0]
-
-    def frontier_side_count(self) -> int:
-        return self._frontier_side_indices.array.shape[0]
-
-    @cached_arg_value
-    def side_arg_value(self, device):
-        arg = LinearGeometryPartition.SideArg()
-        arg.partition_side_indices = self._partition_side_indices.array.to(device)
-        arg.boundary_side_indices = self._boundary_side_indices.array.to(device)
-        arg.frontier_side_indices = self._frontier_side_indices.array.to(device)
-        return arg
-
-    @wp.func
-    def side_index(args: SideArg, partition_side_index: int):
-        """partition side to side index"""
-        return args.partition_side_indices[partition_side_index]
-
-    @wp.func
-    def boundary_side_index(args: SideArg, boundary_side_index: int):
-        """Boundary side to side index"""
-        return args.boundary_side_indices[boundary_side_index]
-
-    @wp.func
-    def frontier_side_index(args: SideArg, frontier_side_index: int):
-        """Frontier side to side index"""
-        return args.frontier_side_indices[frontier_side_index]
-
-    def compute_side_indices_from_cells(
-        self, cell_arg_value: Any, cell_inclusion_test_func: wp.Function, device, temporary_store: TemporaryStore = None
-    ):
-        from warp.fem import cache
-
-        cell_arg_type = next(iter(cell_inclusion_test_func.input_types.values()))
-
-        @cache.dynamic_kernel(suffix=f"{self.geometry.name}_{cell_inclusion_test_func.key}")
-        def count_sides(
-            geo_arg: self.geometry.SideArg,
-            cell_arg_value: cell_arg_type,
-            partition_side_mask: wp.array(dtype=int),
-            boundary_side_mask: wp.array(dtype=int),
-            frontier_side_mask: wp.array(dtype=int),
-        ):
-            side_index = wp.tid()
-            inner_cell_index = self.geometry.side_inner_cell_index(geo_arg, side_index)
-            outer_cell_index = self.geometry.side_outer_cell_index(geo_arg, side_index)
-
-            inner_in = cell_inclusion_test_func(cell_arg_value, inner_cell_index)
-            outer_in = cell_inclusion_test_func(cell_arg_value, outer_cell_index)
-
-            if inner_in:
-                # Inner neighbor in partition; count as partition side
-                partition_side_mask[side_index] = 1
-
-                # Inner and outer element as the same -- this is a boundary side
-                if inner_cell_index == outer_cell_index:
-                    boundary_side_mask[side_index] = 1
-
-            if inner_in != outer_in:
-                # Exactly one neighbor in partition; count as frontier side
-                frontier_side_mask[side_index] = 1
-
-        partition_side_mask = borrow_temporary(
-            temporary_store,
-            shape=(self.geometry.side_count(),),
-            dtype=int,
-            device=device,
-        )
-        boundary_side_mask = borrow_temporary(
-            temporary_store,
-            shape=(self.geometry.side_count(),),
-            dtype=int,
-            device=device,
-        )
-        frontier_side_mask = borrow_temporary(
-            temporary_store,
-            shape=(self.geometry.side_count(),),
-            dtype=int,
-            device=device,
-        )
-
-        partition_side_mask.array.zero_()
-        boundary_side_mask.array.zero_()
-        frontier_side_mask.array.zero_()
-
-        wp.launch(
-            dim=partition_side_mask.array.shape[0],
-            kernel=count_sides,
-            inputs=[
-                self.geometry.side_arg_value(device),
-                cell_arg_value,
-                partition_side_mask.array,
-                boundary_side_mask.array,
-                frontier_side_mask.array,
-            ],
-            device=device,
-        )
-
-        # Convert counts to indices
-        self._partition_side_indices, _ = masked_indices(partition_side_mask.array, temporary_store=temporary_store)
-        self._boundary_side_indices, _ = masked_indices(boundary_side_mask.array, temporary_store=temporary_store)
-        self._frontier_side_indices, _ = masked_indices(frontier_side_mask.array, temporary_store=temporary_store)
-
-        partition_side_mask.release()
-        boundary_side_mask.release()
-        frontier_side_mask.release()
-
-
-class LinearGeometryPartition(CellBasedGeometryPartition):
-    def __init__(
-        self,
-        geometry: Geometry,
-        partition_rank: int,
-        partition_count: int,
-        device=None,
-        temporary_store: TemporaryStore = None,
-    ):
-        """Creates a geometry partition by uniformly partionning cell indices
-
-        Args:
-            geometry: the geometry to partition
-            partition_rank: the index of the partition being created
-            partition_count: the number of partitions that will be created over the geometry
-            device: Warp device on which to perform and store computations
-        """
-        super().__init__(geometry)
-
-        total_cell_count = geometry.cell_count()
-
-        cells_per_partition = (total_cell_count + partition_count - 1) // partition_count
-        self.cell_begin = cells_per_partition * partition_rank
-        self.cell_end = min(self.cell_begin + cells_per_partition, total_cell_count)
-
-        super().compute_side_indices_from_cells(
-            self.cell_arg_value(device),
-            LinearGeometryPartition._cell_inclusion_test,
-            device,
-            temporary_store=temporary_store,
-        )
-
-    def cell_count(self) -> int:
-        return self.cell_end - self.cell_begin
-
-    @wp.struct
-    class CellArg:
-        cell_begin: int
-        cell_end: int
-
-    def cell_arg_value(self, device):
-        arg = LinearGeometryPartition.CellArg()
-        arg.cell_begin = self.cell_begin
-        arg.cell_end = self.cell_end
-        return arg
-
-    @wp.func
-    def cell_index(args: CellArg, partition_cell_index: int):
-        """Partition cell to cell index"""
-        return args.cell_begin + partition_cell_index
-
-    @wp.func
-    def partition_cell_index(args: CellArg, cell_index: int):
-        """Partition cell to cell index"""
-        if cell_index > args.cell_end:
-            return NULL_ELEMENT_INDEX
-
-        partition_cell_index = cell_index - args.cell_begin
-        if partition_cell_index < 0:
-            return NULL_ELEMENT_INDEX
-
-        return partition_cell_index
-
-    @wp.func
-    def _cell_inclusion_test(arg: CellArg, cell_index: int):
-        return cell_index >= arg.cell_begin and cell_index < arg.cell_end
-
-
-class ExplicitGeometryPartition(CellBasedGeometryPartition):
-    def __init__(self, geometry: Geometry, cell_mask: "wp.array(dtype=int)", temporary_store: TemporaryStore = None):
-        """Creates a geometry partition by uniformly partionning cell indices
-
-        Args:
-            geometry: the geometry to partition
-            cell_mask: warp array of length ``geometry.cell_count()`` indicating which cells are selected. Array values must be either ``1`` (selected) or ``0`` (not selected).
-        """
-
-        super().__init__(geometry)
-
-        self._cell_mask = cell_mask
-        self._cells, self._partition_cells = masked_indices(self._cell_mask, temporary_store=temporary_store)
-
-        super().compute_side_indices_from_cells(
-            self._cell_mask,
-            ExplicitGeometryPartition._cell_inclusion_test,
-            self._cell_mask.device,
-            temporary_store=temporary_store,
-        )
-
-    def cell_count(self) -> int:
-        return self._cells.array.shape[0]
-
-    @wp.struct
-    class CellArg:
-        cell_index: wp.array(dtype=int)
-        partition_cell_index: wp.array(dtype=int)
-
-    @cached_arg_value
-    def cell_arg_value(self, device):
-        arg = ExplicitGeometryPartition.CellArg()
-        arg.cell_index = self._cells.array.to(device)
-        arg.partition_cell_index = self._partition_cells.array.to(device)
-        return arg
-
-    @wp.func
-    def cell_index(args: CellArg, partition_cell_index: int):
-        return args.cell_index[partition_cell_index]
-
-    @wp.func
-    def partition_cell_index(args: CellArg, cell_index: int):
-        return args.partition_cell_index[cell_index]
-
-    @wp.func
-    def _cell_inclusion_test(mask: wp.array(dtype=int), cell_index: int):
-        return mask[cell_index] > 0
+from typing import Any
+
+import warp as wp
+from warp.fem.cache import TemporaryStore, borrow_temporary, cached_arg_value
+from warp.fem.types import NULL_ELEMENT_INDEX, ElementIndex
+from warp.fem.utils import masked_indices
+
+from .geometry import Geometry
+
+wp.set_module_options({"enable_backward": False})
+
+
+class GeometryPartition:
+    """Base class for geometry partitions, i.e. subset of cells and sides"""
+
+    class CellArg:
+        pass
+
+    class SideArg:
+        pass
+
+    def __init__(self, geometry: Geometry):
+        self.geometry = geometry
+
+    def cell_count(self) -> int:
+        """Number of cells that are 'owned' by this partition"""
+        raise NotImplementedError()
+
+    def side_count(self) -> int:
+        """Number of sides that are 'owned' by this partition"""
+        raise NotImplementedError()
+
+    def boundary_side_count(self) -> int:
+        """Number of geo-boundary sides that are 'owned' by this partition"""
+        raise NotImplementedError()
+
+    def frontier_side_count(self) -> int:
+        """Number of sides with neighbors owned by this and another partition"""
+        raise NotImplementedError()
+
+    @property
+    def name(self) -> str:
+        return f"{self.geometry.name}_{self.__class__.__name__}"
+
+    def __str__(self) -> str:
+        return self.name
+
+    def cell_arg_value(self, device):
+        raise NotImplementedError()
+
+    def side_arg_value(self, device):
+        raise NotImplementedError()
+
+    @staticmethod
+    def cell_index(args: CellArg, partition_cell_index: int):
+        """Index in the geometry of a partition cell"""
+        raise NotImplementedError()
+
+    @staticmethod
+    def partition_cell_index(args: CellArg, cell_index: int):
+        """Index of a geometry cell in the partition (or ``NULL_ELEMENT_INDEX``)"""
+        raise NotImplementedError()
+
+    @staticmethod
+    def side_index(args: SideArg, partition_side_index: int):
+        """Partition side to side index"""
+        raise NotImplementedError()
+
+    @staticmethod
+    def boundary_side_index(args: SideArg, boundary_side_index: int):
+        """Boundary side to side index"""
+        raise NotImplementedError()
+
+    @staticmethod
+    def frontier_side_index(args: SideArg, frontier_side_index: int):
+        """Frontier side to side index"""
+        raise NotImplementedError()
+
+
+class WholeGeometryPartition(GeometryPartition):
+    """Trivial (NOP) partition"""
+
+    def __init__(
+        self,
+        geometry: Geometry,
+    ):
+        super().__init__(geometry)
+
+        self.SideArg = geometry.SideIndexArg
+        self.side_arg_value = geometry.side_index_arg_value
+
+        self.cell_index = WholeGeometryPartition._identity_element_index
+        self.partition_cell_index = WholeGeometryPartition._identity_element_index
+
+        self.side_index = WholeGeometryPartition._identity_element_index
+        self.boundary_side_index = geometry.boundary_side_index
+        self.frontier_side_index = WholeGeometryPartition._identity_element_index
+
+    def __eq__(self, other: GeometryPartition) -> bool:
+        # Ensures that two whole partition instances of the same geometry are considered equal
+        return isinstance(other, WholeGeometryPartition) and self.geometry == other.geometry
+
+    def cell_count(self) -> int:
+        return self.geometry.cell_count()
+
+    def side_count(self) -> int:
+        return self.geometry.side_count()
+
+    def boundary_side_count(self) -> int:
+        return self.geometry.boundary_side_count()
+
+    def frontier_side_count(self) -> int:
+        return 0
+
+    @wp.struct
+    class CellArg:
+        pass
+
+    def cell_arg_value(self, device):
+        arg = WholeGeometryPartition.CellArg()
+        return arg
+
+    @wp.func
+    def _identity_element_index(args: Any, idx: ElementIndex):
+        return idx
+
+    @property
+    def name(self) -> str:
+        return self.geometry.name
+
+
+class CellBasedGeometryPartition(GeometryPartition):
+    """Geometry partition based on a subset of cells. Interior, boundary and frontier sides are automatically categorized."""
+
+    def __init__(
+        self,
+        geometry: Geometry,
+        device=None,
+    ):
+        super().__init__(geometry)
+
+    @wp.struct
+    class SideArg:
+        partition_side_indices: wp.array(dtype=int)
+        boundary_side_indices: wp.array(dtype=int)
+        frontier_side_indices: wp.array(dtype=int)
+
+    def side_count(self) -> int:
+        return self._partition_side_indices.array.shape[0]
+
+    def boundary_side_count(self) -> int:
+        return self._boundary_side_indices.array.shape[0]
+
+    def frontier_side_count(self) -> int:
+        return self._frontier_side_indices.array.shape[0]
+
+    @cached_arg_value
+    def side_arg_value(self, device):
+        arg = LinearGeometryPartition.SideArg()
+        arg.partition_side_indices = self._partition_side_indices.array.to(device)
+        arg.boundary_side_indices = self._boundary_side_indices.array.to(device)
+        arg.frontier_side_indices = self._frontier_side_indices.array.to(device)
+        return arg
+
+    @wp.func
+    def side_index(args: SideArg, partition_side_index: int):
+        """partition side to side index"""
+        return args.partition_side_indices[partition_side_index]
+
+    @wp.func
+    def boundary_side_index(args: SideArg, boundary_side_index: int):
+        """Boundary side to side index"""
+        return args.boundary_side_indices[boundary_side_index]
+
+    @wp.func
+    def frontier_side_index(args: SideArg, frontier_side_index: int):
+        """Frontier side to side index"""
+        return args.frontier_side_indices[frontier_side_index]
+
+    def compute_side_indices_from_cells(
+        self, cell_arg_value: Any, cell_inclusion_test_func: wp.Function, device, temporary_store: TemporaryStore = None
+    ):
+        from warp.fem import cache
+
+        cell_arg_type = next(iter(cell_inclusion_test_func.input_types.values()))
+
+        @cache.dynamic_kernel(suffix=f"{self.geometry.name}_{cell_inclusion_test_func.key}")
+        def count_sides(
+            geo_arg: self.geometry.SideArg,
+            cell_arg_value: cell_arg_type,
+            partition_side_mask: wp.array(dtype=int),
+            boundary_side_mask: wp.array(dtype=int),
+            frontier_side_mask: wp.array(dtype=int),
+        ):
+            side_index = wp.tid()
+            inner_cell_index = self.geometry.side_inner_cell_index(geo_arg, side_index)
+            outer_cell_index = self.geometry.side_outer_cell_index(geo_arg, side_index)
+
+            inner_in = cell_inclusion_test_func(cell_arg_value, inner_cell_index)
+            outer_in = cell_inclusion_test_func(cell_arg_value, outer_cell_index)
+
+            if inner_in:
+                # Inner neighbor in partition; count as partition side
+                partition_side_mask[side_index] = 1
+
+                # Inner and outer element as the same -- this is a boundary side
+                if inner_cell_index == outer_cell_index:
+                    boundary_side_mask[side_index] = 1
+
+            if inner_in != outer_in:
+                # Exactly one neighbor in partition; count as frontier side
+                frontier_side_mask[side_index] = 1
+
+        partition_side_mask = borrow_temporary(
+            temporary_store,
+            shape=(self.geometry.side_count(),),
+            dtype=int,
+            device=device,
+        )
+        boundary_side_mask = borrow_temporary(
+            temporary_store,
+            shape=(self.geometry.side_count(),),
+            dtype=int,
+            device=device,
+        )
+        frontier_side_mask = borrow_temporary(
+            temporary_store,
+            shape=(self.geometry.side_count(),),
+            dtype=int,
+            device=device,
+        )
+
+        partition_side_mask.array.zero_()
+        boundary_side_mask.array.zero_()
+        frontier_side_mask.array.zero_()
+
+        wp.launch(
+            dim=partition_side_mask.array.shape[0],
+            kernel=count_sides,
+            inputs=[
+                self.geometry.side_arg_value(device),
+                cell_arg_value,
+                partition_side_mask.array,
+                boundary_side_mask.array,
+                frontier_side_mask.array,
+            ],
+            device=device,
+        )
+
+        # Convert counts to indices
+        self._partition_side_indices, _ = masked_indices(partition_side_mask.array, temporary_store=temporary_store)
+        self._boundary_side_indices, _ = masked_indices(boundary_side_mask.array, temporary_store=temporary_store)
+        self._frontier_side_indices, _ = masked_indices(frontier_side_mask.array, temporary_store=temporary_store)
+
+        partition_side_mask.release()
+        boundary_side_mask.release()
+        frontier_side_mask.release()
+
+
+class LinearGeometryPartition(CellBasedGeometryPartition):
+    def __init__(
+        self,
+        geometry: Geometry,
+        partition_rank: int,
+        partition_count: int,
+        device=None,
+        temporary_store: TemporaryStore = None,
+    ):
+        """Creates a geometry partition by uniformly partionning cell indices
+
+        Args:
+            geometry: the geometry to partition
+            partition_rank: the index of the partition being created
+            partition_count: the number of partitions that will be created over the geometry
+            device: Warp device on which to perform and store computations
+        """
+        super().__init__(geometry)
+
+        total_cell_count = geometry.cell_count()
+
+        cells_per_partition = (total_cell_count + partition_count - 1) // partition_count
+        self.cell_begin = cells_per_partition * partition_rank
+        self.cell_end = min(self.cell_begin + cells_per_partition, total_cell_count)
+
+        super().compute_side_indices_from_cells(
+            self.cell_arg_value(device),
+            LinearGeometryPartition._cell_inclusion_test,
+            device,
+            temporary_store=temporary_store,
+        )
+
+    def cell_count(self) -> int:
+        return self.cell_end - self.cell_begin
+
+    @wp.struct
+    class CellArg:
+        cell_begin: int
+        cell_end: int
+
+    def cell_arg_value(self, device):
+        arg = LinearGeometryPartition.CellArg()
+        arg.cell_begin = self.cell_begin
+        arg.cell_end = self.cell_end
+        return arg
+
+    @wp.func
+    def cell_index(args: CellArg, partition_cell_index: int):
+        """Partition cell to cell index"""
+        return args.cell_begin + partition_cell_index
+
+    @wp.func
+    def partition_cell_index(args: CellArg, cell_index: int):
+        """Partition cell to cell index"""
+        if cell_index > args.cell_end:
+            return NULL_ELEMENT_INDEX
+
+        partition_cell_index = cell_index - args.cell_begin
+        if partition_cell_index < 0:
+            return NULL_ELEMENT_INDEX
+
+        return partition_cell_index
+
+    @wp.func
+    def _cell_inclusion_test(arg: CellArg, cell_index: int):
+        return cell_index >= arg.cell_begin and cell_index < arg.cell_end
+
+
+class ExplicitGeometryPartition(CellBasedGeometryPartition):
+    def __init__(self, geometry: Geometry, cell_mask: "wp.array(dtype=int)", temporary_store: TemporaryStore = None):
+        """Creates a geometry partition by uniformly partionning cell indices
+
+        Args:
+            geometry: the geometry to partition
+            cell_mask: warp array of length ``geometry.cell_count()`` indicating which cells are selected. Array values must be either ``1`` (selected) or ``0`` (not selected).
+        """
+
+        super().__init__(geometry)
+
+        self._cell_mask = cell_mask
+        self._cells, self._partition_cells = masked_indices(self._cell_mask, temporary_store=temporary_store)
+
+        super().compute_side_indices_from_cells(
+            self._cell_mask,
+            ExplicitGeometryPartition._cell_inclusion_test,
+            self._cell_mask.device,
+            temporary_store=temporary_store,
+        )
+
+    def cell_count(self) -> int:
+        return self._cells.array.shape[0]
+
+    @wp.struct
+    class CellArg:
+        cell_index: wp.array(dtype=int)
+        partition_cell_index: wp.array(dtype=int)
+
+    @cached_arg_value
+    def cell_arg_value(self, device):
+        arg = ExplicitGeometryPartition.CellArg()
+        arg.cell_index = self._cells.array.to(device)
+        arg.partition_cell_index = self._partition_cells.array.to(device)
+        return arg
+
+    @wp.func
+    def cell_index(args: CellArg, partition_cell_index: int):
+        return args.cell_index[partition_cell_index]
+
+    @wp.func
+    def partition_cell_index(args: CellArg, cell_index: int):
+        return args.partition_cell_index[cell_index]
+
+    @wp.func
+    def _cell_inclusion_test(mask: wp.array(dtype=int), cell_index: int):
+        return mask[cell_index] > 0