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

warp/fem/space/basis_space.py

@@ -1,13 +1,15 @@
 from typing import Optional
 
+import numpy as np
+
 import warp as wp
 from warp.fem import cache
 from warp.fem.geometry import Geometry
 from warp.fem.quadrature import Quadrature
-from warp.fem.types import Coords, ElementIndex, make_free_sample
+from warp.fem.types import NULL_ELEMENT_INDEX, Coords, ElementIndex, make_free_sample
 
 from .shape import ShapeFunction
-from .topology import DiscontinuousSpaceTopology, SpaceTopology
+from .topology import RegularDiscontinuousSpaceTopology, SpaceTopology
 
 
 class BasisSpace:
@@ -28,8 +30,6 @@ class BasisSpace:
     def __init__(self, topology: SpaceTopology):
         self._topology = topology
 
-        self.NODES_PER_ELEMENT = self._topology.NODES_PER_ELEMENT
-
     @property
     def topology(self) -> SpaceTopology:
         """Underlying topology of the basis space"""
@@ -49,8 +49,6 @@ class BasisSpace:
     def node_positions(self, out: Optional[wp.array] = None) -> wp.array:
         """Returns a temporary array containing the world position for each node"""
 
-        NODES_PER_ELEMENT = self.NODES_PER_ELEMENT
-
         pos_type = cache.cached_vec_type(length=self.geometry.dimension, dtype=float)
 
         node_coords_in_element = self.make_node_coords_in_element()
@@ -64,7 +62,8 @@ class BasisSpace:
         ):
             element_index = wp.tid()
 
-            for n in range(NODES_PER_ELEMENT):
+            element_node_count = self.topology.element_node_count(geo_cell_arg, topo_arg, element_index)
+            for n in range(element_node_count):
                 node_index = self.topology.element_node_index(geo_cell_arg, topo_arg, element_index, n)
                 coords = node_coords_in_element(geo_cell_arg, basis_arg, element_index, n)
 
@@ -139,6 +138,10 @@ class ShapeBasisSpace(BasisSpace):
             self.node_tets = self._node_tets
         if hasattr(shape, "element_node_hexes"):
             self.node_hexes = self._node_hexes
+        if hasattr(shape, "element_vtk_cells"):
+            self.vtk_cells = self._vtk_cells
+        if hasattr(topology, "node_grid"):
+            self.node_grid = topology.node_grid
 
     @property
     def shape(self) -> ShapeFunction:
@@ -245,6 +248,16 @@ class ShapeBasisSpace(BasisSpace):
         hex_indices = element_node_indices[:, element_hexes].reshape(-1, 8)
         return hex_indices
 
+    def _vtk_cells(self):
+        element_node_indices = self._topology.element_node_indices().numpy()
+        element_vtk_cells, element_vtk_cell_types = self._shape.element_vtk_cells()
+
+        idx_per_cell = element_vtk_cells.shape[1]
+        cell_indices = element_node_indices[:, element_vtk_cells].reshape(-1, idx_per_cell)
+        cells = np.hstack((np.full((cell_indices.shape[0], 1), idx_per_cell), cell_indices))
+
+        return cells.flatten(), np.tile(element_vtk_cell_types, element_node_indices.shape[0])
+
 
 class TraceBasisSpace(BasisSpace):
     """Auto-generated trace space evaluating the cell-defined basis on the geometry sides"""
@@ -302,7 +315,7 @@ class TraceBasisSpace(BasisSpace):
             node_index_in_elt: int,
         ):
             cell_index, index_in_cell = self.topology.inner_cell_index(geo_side_arg, element_index, node_index_in_elt)
-            if index_in_cell < 0:
+            if cell_index == NULL_ELEMENT_INDEX:
                 return 0.0
 
             cell_coords = self.geometry.side_inner_cell_coords(geo_side_arg, element_index, coords)
@@ -330,7 +343,7 @@ class TraceBasisSpace(BasisSpace):
             node_index_in_elt: int,
         ):
             cell_index, index_in_cell = self.topology.outer_cell_index(geo_side_arg, element_index, node_index_in_elt)
-            if index_in_cell < 0:
+            if cell_index == NULL_ELEMENT_INDEX:
                 return 0.0
 
             cell_coords = self.geometry.side_outer_cell_coords(geo_side_arg, element_index, coords)
@@ -359,7 +372,7 @@ class TraceBasisSpace(BasisSpace):
             node_index_in_elt: int,
         ):
             cell_index, index_in_cell = self.topology.inner_cell_index(geo_side_arg, element_index, node_index_in_elt)
-            if index_in_cell < 0:
+            if cell_index == NULL_ELEMENT_INDEX:
                 return grad_vec_type(0.0)
 
             cell_coords = self.geometry.side_inner_cell_coords(geo_side_arg, element_index, coords)
@@ -381,7 +394,7 @@ class TraceBasisSpace(BasisSpace):
             node_index_in_elt: int,
         ):
             cell_index, index_in_cell = self.topology.outer_cell_index(geo_side_arg, element_index, node_index_in_elt)
-            if index_in_cell < 0:
+            if cell_index == NULL_ELEMENT_INDEX:
                 return grad_vec_type(0.0)
 
             cell_coords = self.geometry.side_outer_cell_coords(geo_side_arg, element_index, coords)
@@ -419,7 +432,7 @@ class PiecewiseConstantBasisSpace(ShapeBasisSpace):
 
 
 def make_discontinuous_basis_space(geometry: Geometry, shape: ShapeFunction):
-    topology = DiscontinuousSpaceTopology(geometry, shape.NODES_PER_ELEMENT)
+    topology = RegularDiscontinuousSpaceTopology(geometry, shape.NODES_PER_ELEMENT)
 
     if shape.NODES_PER_ELEMENT == 1:
         # piecewise-constant space
@@ -428,6 +441,70 @@ def make_discontinuous_basis_space(geometry: Geometry, shape: ShapeFunction):
     return ShapeBasisSpace(topology=topology, shape=shape)
 
 
+class UnstructuredPointTopology(SpaceTopology):
+    """Topology for unstructured points defined from quadrature formula. See :class:`PointBasisSpace`"""
+
+    def __init__(self, quadrature: Quadrature):
+        if quadrature.max_points_per_element() is None:
+            raise ValueError("Quadrature must define a maximum number of points per element")
+
+        if quadrature.domain.element_count() != quadrature.domain.geometry_element_count():
+            raise ValueError("Point topology may only be defined on quadrature domains than span the whole geometry")
+
+        self._quadrature = quadrature
+        self.TopologyArg = quadrature.Arg
+
+        super().__init__(quadrature.domain.geometry, max_nodes_per_element=quadrature.max_points_per_element())
+
+        self.element_node_index = self._make_element_node_index()
+        self.element_node_count = self._make_element_node_count()
+        self.side_neighbor_node_counts = self._make_side_neighbor_node_counts()
+
+    def node_count(self):
+        return self._quadrature.total_point_count()
+
+    @property
+    def name(self):
+        return f"PointTopology_{self._quadrature}"
+
+    def topo_arg_value(self, device) -> SpaceTopology.TopologyArg:
+        """Value of the topology argument structure to be passed to device functions"""
+        return self._quadrature.arg_value(device)
+
+    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._quadrature.point_index(elt_arg, topo_arg, element_index, 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_index: ElementIndex,
+        ):
+            return self._quadrature.point_count(elt_arg, topo_arg, element_index, element_index)
+
+        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,
+        ):
+            return 0, 0
+
+        return side_neighbor_node_counts
+
+
 class PointBasisSpace(BasisSpace):
     """An unstructured :class:`BasisSpace` that is non-zero at a finite set of points only.
 
@@ -437,12 +514,7 @@ class PointBasisSpace(BasisSpace):
     def __init__(self, quadrature: Quadrature):
         self._quadrature = quadrature
 
-        if quadrature.points_per_element() is None:
-            raise NotImplementedError("Varying number of points per element is not supported yet")
-
-        topology = DiscontinuousSpaceTopology(
-            geometry=quadrature.domain.geometry, nodes_per_element=quadrature.points_per_element()
-        )
+        topology = UnstructuredPointTopology(quadrature)
         super().__init__(topology)
 
         self.BasisArg = quadrature.Arg
@@ -464,7 +536,7 @@ class PointBasisSpace(BasisSpace):
             element_index: ElementIndex,
             node_index_in_elt: int,
         ):
-            return self._quadrature.point_coords(elt_arg, basis_arg, element_index, node_index_in_elt)
+            return self._quadrature.point_coords(elt_arg, basis_arg, element_index, element_index, node_index_in_elt)
 
         return node_coords_in_element
 
@@ -476,11 +548,13 @@ class PointBasisSpace(BasisSpace):
             element_index: ElementIndex,
             node_index_in_elt: int,
         ):
-            return self._quadrature.point_weight(elt_arg, basis_arg, element_index, node_index_in_elt)
+            return self._quadrature.point_weight(elt_arg, basis_arg, element_index, element_index, node_index_in_elt)
 
         return node_quadrature_weight
 
     def make_element_inner_weight(self):
+        _DIRAC_INTEGRATION_RADIUS = wp.constant(1.0e-6)
+
         @cache.dynamic_func(suffix=self.name)
         def element_inner_weight(
             elt_arg: self._quadrature.domain.ElementArg,
@@ -489,8 +563,10 @@ class PointBasisSpace(BasisSpace):
             coords: Coords,
             node_index_in_elt: int,
         ):
-            qp_coord = self._quadrature.point_coords(elt_arg, basis_arg, element_index, node_index_in_elt)
-            return wp.select(wp.length_sq(coords - qp_coord) < 0.001, 0.0, 1.0)
+            qp_coord = self._quadrature.point_coords(
+                elt_arg, basis_arg, element_index, element_index, node_index_in_elt
+            )
+            return wp.select(wp.length_sq(coords - qp_coord) < _DIRAC_INTEGRATION_RADIUS, 0.0, 1.0)
 
         return element_inner_weight
 

warp/fem/space/collocated_function_space.py

@@ -36,7 +36,7 @@ class CollocatedFunctionSpace(FunctionSpace):
         self.element_outer_weight_gradient = self._basis.make_element_outer_weight_gradient()
 
         # For backward compatibility
-        if hasattr(basis.topology, "node_grid"):
+        if hasattr(basis, "node_grid"):
             self.node_grid = basis.node_grid
         if hasattr(basis, "node_triangulation"):
             self.node_triangulation = basis.node_triangulation
@@ -44,6 +44,8 @@ class CollocatedFunctionSpace(FunctionSpace):
             self.node_tets = basis.node_tets
         if hasattr(basis, "node_hexes"):
             self.node_hexes = basis.node_hexes
+        if hasattr(basis, "vtk_cells"):
+            self.vtk_cells = basis.vtk_cells
 
     def space_arg_value(self, device):
         return self._basis.basis_arg_value(device)

warp/fem/space/function_space.py

@@ -1,6 +1,6 @@
 import warp as wp
 from warp.fem.geometry import Geometry
-from warp.fem.types import Coords, DofIndex, ElementIndex
+from warp.fem.types import Coords, DofIndex, ElementIndex, ElementKind
 
 from .topology import SpaceTopology
 
@@ -47,6 +47,11 @@ class FunctionSpace:
         """Underlying geometry"""
         return self.topology.geometry
 
+    @property
+    def element_kind(self) -> ElementKind:
+        """Kind of element the function space is expressed over"""
+        return ElementKind.CELL if self.dimension == self.geometry.dimension else ElementKind.SIDE
+
     @property
     def dimension(self) -> int:
         """Function space embedding dimension"""
@@ -71,7 +76,7 @@ class FunctionSpace:
     def make_field(self, space_partition=None):
         """Creates a zero-initialized discrete field over the function space holding values for all degrees of freedom of nodes in a space partition
 
-        space_arg:
+        Args:
             space_partition: If provided, the subset of nodes to consider
 
         See also: :func:`make_space_partition`

warp/fem/space/grid_2d_function_space.py

@@ -72,7 +72,7 @@ class GridBipolynomialSpaceTopology(Grid2DSpaceTopology):
 
         return element_node_index
 
-    def _node_grid(self):
+    def node_grid(self):
         res = self.geometry.res
 
         cell_coords = np.array(self._shape.LOBATTO_COORDS)[:-1]
@@ -81,13 +81,13 @@ class GridBipolynomialSpaceTopology(Grid2DSpaceTopology):
             cell_coords, reps=res[0]
         )
         grid_coords_x = np.append(grid_coords_x, res[0])
-        X = grid_coords_x * self._grid.cell_size[0] + self._grid.origin[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._grid.cell_size[1] + self._grid.origin[1]
+        Y = grid_coords_y * self.geometry.cell_size[1] + self.geometry.origin[1]
 
         return np.meshgrid(X, Y, indexing="ij")
 

warp/fem/space/grid_3d_function_space.py

@@ -79,7 +79,7 @@ class GridTripolynomialSpaceTopology(Grid3DSpaceTopology):
 
         return element_node_index
 
-    def _node_grid(self):
+    def node_grid(self):
         res = self.geometry.res
 
         cell_coords = np.array(self._shape.LOBATTO_COORDS)[:-1]
@@ -88,19 +88,19 @@ class GridTripolynomialSpaceTopology(Grid3DSpaceTopology):
             cell_coords, reps=res[0]
         )
         grid_coords_x = np.append(grid_coords_x, res[0])
-        X = grid_coords_x * self._grid.cell_size[0] + self._grid.origin[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._grid.cell_size[1] + self._grid.origin[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._grid.cell_size[2] + self._grid.origin[2]
+        Z = grid_coords_z * self.geometry.cell_size[2] + self.geometry.origin[2]
 
         return np.meshgrid(X, Y, Z, indexing="ij")
 

warp/fem/space/hexmesh_function_space.py

@@ -142,8 +142,9 @@ class HexmeshTripolynomialSpaceTopology(HexmeshSpaceTopology):
 
         fv = ori // 2
 
-        rot_i = wp.dot(_FACE_ORIENTATION_I[2 * ori], coords) + _FACE_TRANSLATION_I[fv, 0]
-        rot_j = wp.dot(_FACE_ORIENTATION_I[2 * ori + 1], coords) + _FACE_TRANSLATION_I[fv, 1]
+        # face indices from shape function always have positive orientation, drop `ori % 2`
+        rot_i = wp.dot(_FACE_ORIENTATION_I[4 * fv], coords) + _FACE_TRANSLATION_I[fv, 0]
+        rot_j = wp.dot(_FACE_ORIENTATION_I[4 * fv + 1], coords) + _FACE_TRANSLATION_I[fv, 1]
 
         return rot_i * size + rot_j
 

warp/fem/space/nanogrid_function_space.py

@@ -41,25 +41,23 @@ class NanogridSpaceTopology(SpaceTopology):
         self._grid = grid
         self._shape = shape
 
-        if need_edge_indices:
-            self._edge_count = self._grid.edge_count()
-        else:
-            self._edge_count = 0
+        self._vertex_grid = grid.vertex_grid.id
 
-        self._vertex_grid = grid._node_grid
-        self._face_grid = grid._face_grid
-        self._edge_grid = grid._edge_grid
+        self._edge_grid = grid.edge_grid.id if need_edge_indices else -1
+        self._face_grid = grid.face_grid.id if need_face_indices else -1
+        self._edge_count = grid.edge_count() if need_edge_indices else 0
+        self._face_count = grid.side_count() if need_face_indices else 0
 
     @cache.cached_arg_value
     def topo_arg_value(self, device):
         arg = NanogridTopologyArg()
 
-        arg.vertex_grid = self._vertex_grid.id
-        arg.face_grid = self._face_grid.id
-        arg.edge_grid = -1 if self._edge_grid is None else self._edge_grid.id
+        arg.vertex_grid = self._vertex_grid
+        arg.face_grid = self._face_grid
+        arg.edge_grid = self._edge_grid
 
         arg.vertex_count = self._grid.vertex_count()
-        arg.face_count = self._grid.side_count()
+        arg.face_count = self._face_count
         arg.edge_count = self._edge_count
         return arg
 
@@ -98,8 +96,8 @@ class NanogridTripolynomialSpaceTopology(NanogridSpaceTopology):
 
         return (
             self._grid.vertex_count()
-            + self._grid.edge_count() * INTERIOR_NODES_PER_EDGE
-            + self._grid.side_count() * INTERIOR_NODES_PER_FACE
+            + self._edge_count * INTERIOR_NODES_PER_EDGE
+            + self._face_count * INTERIOR_NODES_PER_FACE
             + self._grid.cell_count() * INTERIOR_NODES_PER_CELL
         )
 
@@ -160,7 +158,7 @@ class NanogridSerendipitySpaceTopology(NanogridSpaceTopology):
         self.element_node_index = self._make_element_node_index()
 
     def node_count(self) -> int:
-        return self.geometry.vertex_count() + (self._shape.ORDER - 1) * self.geometry.edge_count()
+        return self.geometry.vertex_count() + (self._shape.ORDER - 1) * self._edge_count
 
     def _make_element_node_index(self):
         ORDER = self._shape.ORDER

warp/fem/space/partition.py

@@ -1,12 +1,7 @@
 from typing import Any, Optional
 
 import warp as wp
-from warp.fem.cache import (
-    TemporaryStore,
-    borrow_temporary,
-    borrow_temporary_like,
-    cached_arg_value,
-)
+import warp.fem.cache as cache
 from warp.fem.geometry import GeometryPartition, WholeGeometryPartition
 from warp.fem.types import NULL_NODE_INDEX
 from warp.fem.utils import _iota_kernel, compress_node_indices
@@ -42,7 +37,7 @@ class SpacePartition:
 
     @staticmethod
     def partition_node_index(args: "PartitionArg", space_node_index: int):
-        """Returns the index in the partition of a function space node, or -1 if it does not exist"""
+        """Returns the index in the partition of a function space node, or ``NULL_NODE_INDEX`` if it does not exist"""
 
     def __str__(self) -> str:
         return self.name
@@ -76,7 +71,7 @@ class WholeSpacePartition(SpacePartition):
     def space_node_indices(self):
         """Return the global function space indices for nodes in this partition"""
         if self._node_indices is None:
-            self._node_indices = borrow_temporary(temporary_store=None, shape=(self.node_count(),), dtype=int)
+            self._node_indices = cache.borrow_temporary(temporary_store=None, shape=(self.node_count(),), dtype=int)
             wp.launch(kernel=_iota_kernel, dim=self.node_count(), inputs=[self._node_indices.array, 1])
         return self._node_indices.array
 
@@ -121,7 +116,7 @@ class NodePartition(SpacePartition):
         geo_partition: GeometryPartition,
         with_halo: bool = True,
         device=None,
-        temporary_store: TemporaryStore = None,
+        temporary_store: cache.TemporaryStore = None,
     ):
         super().__init__(space_topology=space_topology, geo_partition=geo_partition)
 
@@ -143,7 +138,7 @@ class NodePartition(SpacePartition):
         """Return the global function space indices for nodes in this partition"""
         return self._node_indices.array
 
-    @cached_arg_value
+    @cache.cached_arg_value
     def partition_arg_value(self, device):
         arg = NodePartition.PartitionArg()
         arg.space_to_partition = self._space_to_partition.array.to(device)
@@ -153,12 +148,10 @@ class NodePartition(SpacePartition):
     def partition_node_index(args: PartitionArg, space_node_index: int):
         return args.space_to_partition[space_node_index]
 
-    def _compute_node_indices_from_sides(self, device, with_halo: bool, temporary_store: TemporaryStore):
+    def _compute_node_indices_from_sides(self, device, with_halo: bool, temporary_store: cache.TemporaryStore):
         from warp.fem import cache
 
         trace_topology = self.space_topology.trace()
-        NODES_PER_CELL = self.space_topology.NODES_PER_ELEMENT
-        NODES_PER_SIDE = trace_topology.NODES_PER_ELEMENT
 
         @cache.dynamic_kernel(suffix=f"{self.geo_partition.name}_{self.space_topology.name}")
         def node_category_from_cells_kernel(
@@ -171,7 +164,8 @@ class NodePartition(SpacePartition):
 
             cell_index = self.geo_partition.cell_index(geo_partition_arg, partition_cell_index)
 
-            for n in range(NODES_PER_CELL):
+            cell_node_count = self.space_topology.element_node_count(geo_arg, space_arg, cell_index)
+            for n in range(cell_node_count):
                 space_nidx = self.space_topology.element_node_index(geo_arg, space_arg, cell_index, n)
                 node_mask[space_nidx] = NodeCategory.OWNED_INTERIOR
 
@@ -186,7 +180,8 @@ class NodePartition(SpacePartition):
 
             side_index = self.geo_partition.side_index(geo_partition_arg, partition_side_index)
 
-            for n in range(NODES_PER_SIDE):
+            side_node_count = trace_topology.element_node_count(geo_arg, space_arg, side_index)
+            for n in range(side_node_count):
                 space_nidx = trace_topology.element_node_index(geo_arg, space_arg, side_index, n)
 
                 if node_mask[space_nidx] == NodeCategory.EXTERIOR:
@@ -203,14 +198,15 @@ class NodePartition(SpacePartition):
 
             side_index = self.geo_partition.frontier_side_index(geo_partition_arg, frontier_side_index)
 
-            for n in range(NODES_PER_SIDE):
+            side_node_count = trace_topology.element_node_count(geo_arg, space_arg, side_index)
+            for n in range(side_node_count):
                 space_nidx = trace_topology.element_node_index(geo_arg, space_arg, side_index, n)
                 if node_mask[space_nidx] == NodeCategory.EXTERIOR:
                     node_mask[space_nidx] = NodeCategory.HALO_OTHER_SIDE
                 elif node_mask[space_nidx] == NodeCategory.OWNED_INTERIOR:
                     node_mask[space_nidx] = NodeCategory.OWNED_FRONTIER
 
-        node_category = borrow_temporary(
+        node_category = cache.borrow_temporary(
             temporary_store,
             shape=(self.space_topology.node_count(),),
             dtype=int,
@@ -259,50 +255,52 @@ class NodePartition(SpacePartition):
 
         node_category.release()
 
-    def _finalize_node_indices(self, node_category: wp.array(dtype=int), temporary_store: TemporaryStore):
-        category_offsets, node_indices, _, __ = compress_node_indices(NodeCategory.COUNT, node_category)
+    def _finalize_node_indices(self, node_category: wp.array(dtype=int), temporary_store: cache.TemporaryStore):
+        category_offsets, node_indices = compress_node_indices(
+            NodeCategory.COUNT, node_category, temporary_store=temporary_store
+        )
 
         # Copy offsets to cpu
         device = node_category.device
-        self._category_offsets = borrow_temporary(
-            temporary_store,
-            shape=category_offsets.array.shape,
-            dtype=category_offsets.array.dtype,
-            pinned=device.is_cuda,
-            device="cpu",
-        )
-        wp.copy(src=category_offsets.array, dest=self._category_offsets.array)
-
-        if device.is_cuda:
-            # TODO switch to synchronize_event once available
-            wp.synchronize_stream(wp.get_stream(device))
-
-        category_offsets.release()
+        with wp.ScopedDevice(device):
+            self._category_offsets = cache.borrow_temporary(
+                temporary_store,
+                shape=category_offsets.array.shape,
+                dtype=category_offsets.array.dtype,
+                pinned=device.is_cuda,
+                device="cpu",
+            )
+            wp.copy(src=category_offsets.array, dest=self._category_offsets.array)
+            copy_event = cache.capture_event()
 
-        # Compute global to local indices
-        self._space_to_partition = borrow_temporary_like(node_indices, temporary_store)
-        wp.launch(
-            kernel=NodePartition._scatter_partition_indices,
-            dim=self.space_topology.node_count(),
-            device=device,
-            inputs=[self.node_count(), node_indices.array, self._space_to_partition.array],
-        )
+            # Compute global to local indices
+            self._space_to_partition = cache.borrow_temporary_like(node_indices, temporary_store)
+            wp.launch(
+                kernel=NodePartition._scatter_partition_indices,
+                dim=self.space_topology.node_count(),
+                device=device,
+                inputs=[category_offsets.array, node_indices.array, self._space_to_partition.array],
+            )
 
-        # Copy to shrinked-to-fit array
-        self._node_indices = borrow_temporary(temporary_store, shape=(self.node_count()), dtype=int, device=device)
-        wp.copy(dest=self._node_indices.array, src=node_indices.array, count=self.node_count())
+            # Copy to shrinked-to-fit array
+            cache.synchronize_event(copy_event)  # Transfer to host must be finished to access node_count()
+            self._node_indices = cache.borrow_temporary(
+                temporary_store, shape=(self.node_count()), dtype=int, device=device
+            )
+            wp.copy(dest=self._node_indices.array, src=node_indices.array, count=self.node_count())
 
-        node_indices.release()
+            node_indices.release()
 
     @wp.kernel
     def _scatter_partition_indices(
-        local_node_count: int,
+        category_offsets: wp.array(dtype=int),
         node_indices: wp.array(dtype=int),
         space_to_partition_indices: wp.array(dtype=int),
     ):
         local_idx = wp.tid()
         space_idx = node_indices[local_idx]
 
+        local_node_count = category_offsets[NodeCategory.EXTERIOR]  # all but exterior nodes
         if local_idx < local_node_count:
             space_to_partition_indices[space_idx] = local_idx
         else:
@@ -315,7 +313,7 @@ def make_space_partition(
     space_topology: Optional[SpaceTopology] = None,
     with_halo: bool = True,
     device=None,
-    temporary_store: TemporaryStore = None,
+    temporary_store: cache.TemporaryStore = None,
 ) -> SpacePartition:
     """Computes the subset of nodes from a function space topology that touch a geometry partition