warp-lang 1.9.0__py3-none-manylinux_2_34_aarch64.whl → 1.10.0__py3-none-manylinux_2_34_aarch64.whl

warp/{fem → _src/fem}/geometry/nanogrid.py

@@ -18,104 +18,74 @@ from typing import Any, Optional
 import numpy as np
 
 import warp as wp
-from warp.fem import cache, utils
-from warp.fem.types import NULL_ELEMENT_INDEX, OUTSIDE, Coords, ElementIndex, Sample, make_free_sample
+from warp._src.fem import cache, utils
+from warp._src.fem.types import NULL_ELEMENT_INDEX, OUTSIDE, Coords, ElementIndex, Sample, make_free_sample
 
-from .element import Cube, Square
+from .element import Element
 from .geometry import Geometry
 
-# Flag used for building edge/face grids to disambiguiate axis within the grid
-# Morton indexing allows for
-GRID_AXIS_FLAG = wp.constant(wp.int32(1 << 20))
+_wp_module_name_ = "warp.fem.geometry.nanogrid"
 
-FACE_AXIS_MASK = wp.constant(wp.uint8((1 << 2) - 1))
-FACE_INNER_OFFSET_BIT = wp.constant(wp.uint8(2))
-FACE_OUTER_OFFSET_BIT = wp.constant(wp.uint8(3))
 
-
-@wp.func
-def _add_axis_flag(ijk: wp.vec3i, axis: int):
-    coord = ijk[axis]
-    ijk[axis] = wp.where(coord < 0, coord & (~GRID_AXIS_FLAG), coord | GRID_AXIS_FLAG)
-    return ijk
-
-
-@wp.func
-def _extract_axis_flag(ijk: wp.vec3i):
-    for ax in range(3):
-        coord = ijk[ax]
-        if coord < 0:
-            if (ijk[ax] & GRID_AXIS_FLAG) == 0:
-                ijk[ax] = ijk[ax] | GRID_AXIS_FLAG
-                return ax, ijk
-        else:
-            if (ijk[ax] & GRID_AXIS_FLAG) != 0:
-                ijk[ax] = ijk[ax] & (~GRID_AXIS_FLAG)
-                return ax, ijk
-
-    return -1, ijk
-
-
-@wp.struct
-class NanogridCellArg:
-    # Utility device functions
-    cell_grid: wp.uint64
-    cell_ijk: wp.array(dtype=wp.vec3i)
-    inverse_transform: wp.mat33
-    cell_volume: float
-
-
-@wp.struct
-class NanogridSideArg:
-    # Utility device functions
-    cell_arg: NanogridCellArg
-    face_ijk: wp.array(dtype=wp.vec3i)
-    face_flags: wp.array(dtype=wp.uint8)
-    face_areas: wp.vec3
-
-
-class Nanogrid(Geometry):
-    """Sparse grid geometry"""
+class NanogridBase(Geometry):
+    """Base class for regular and adaptive Nanogrid"""
 
     dimension = 3
 
-    def __init__(self, grid: wp.Volume, temporary_store: Optional[cache.TemporaryStore] = None):
-        """
-        Constructs a sparse grid geometry from an in-memory NanoVDB volume.
+    # Flag used for building edge/face grids to disambiguiate axis within the grid
+    # Morton indexing allows for
+    GRID_AXIS_FLAG = wp.constant(wp.int32(1 << 20))
 
-        Args:
-            grid: The NanoVDB volume. Any type is accepted, but for indexing efficiency an index grid is recommended.
-                If `grid` is an 'on' index grid, cells will be created for active voxels only, otherwise cells will
-                be created for all leaf voxels.
-            temporary_store: shared pool from which to allocate temporary arrays
-        """
+    FACE_AXIS_MASK = wp.constant(wp.uint8((1 << 2) - 1))
+    FACE_INNER_OFFSET_BIT = wp.constant(wp.uint8(2))
+    FACE_OUTER_OFFSET_BIT = wp.constant(wp.uint8(3))
 
-        self._cell_grid = grid
-        self._cell_grid_info = grid.get_grid_info()
+    def __init__(
+        self,
+        cell_grid: wp.Volume,
+        cell_ijk: wp.array(dtype=wp.vec3i),
+        node_grid: wp.Volume,
+        node_ijk: wp.array(dtype=wp.vec3i),
+    ):
+        self._cell_grid = cell_grid
+        self._cell_ijk = cell_ijk
+        self._node_grid = node_grid
+        self._node_ijk = node_ijk
 
-        device = grid.device
+        self._face_grid = None
+        self._face_ijk = None
+        self._boundary_face_indices = None
 
-        cell_count = grid.get_voxel_count()
-        self._cell_ijk = wp.array(shape=(cell_count,), dtype=wp.vec3i, device=device)
-        grid.get_voxels(out=self._cell_ijk)
+        self._cell_grid_info = cell_grid.get_grid_info()
+        self._init_transform()
 
-        self._node_grid = _build_node_grid(self._cell_ijk, grid, temporary_store)
-        node_count = self._node_grid.get_voxel_count()
-        self._node_ijk = wp.array(shape=(node_count,), dtype=wp.vec3i, device=device)
-        self._node_grid.get_voxels(out=self._node_ijk)
+    def reference_cell(self) -> Element:
+        return Element.CUBE
 
-        self._face_grid = None
-        self._face_ijk = None
+    def reference_side(self) -> Element:
+        return Element.SQUARE
 
-        self._edge_grid = None
-        self._edge_count = 0
+    def _init_transform(self):
+        transform = np.array(self.transform).reshape(3, 3)
 
-        transform = np.array(self._cell_grid_info.transform_matrix).reshape(3, 3)
-        self._inverse_transform = wp.mat33f(np.linalg.inv(transform))
-        self._cell_volume = abs(np.linalg.det(transform))
-        self._face_areas = wp.vec3(
-            tuple(np.linalg.norm(np.cross(transform[:, k - 2], transform[:, k - 1])) for k in range(3))
-        )
+        diag = np.diag(transform)
+        if np.max(np.abs(transform - np.diag(diag))) < 1.0e-6 * np.max(diag):
+            # Rectangular voxels
+            self._inverse_transform = wp.mat33f(
+                1.0 / diag[0], 0.0, 0.0, 0.0, 1.0 / diag[1], 0.0, 0.0, 0.0, 1.0 / diag[2]
+            )
+            self._cell_volume = wp.float32(diag[0] * diag[1] * diag[2])
+            self._face_areas = wp.vec3(diag[1] * diag[2], diag[2] * diag[0], diag[0] * diag[1])
+        else:
+            self._inverse_transform = wp.mat33f(np.linalg.inv(transform))
+            self._cell_volume = wp.float32(abs(np.linalg.det(transform)))
+            self._face_areas = wp.vec3(
+                tuple(np.linalg.norm(np.cross(transform[:, k - 2], transform[:, k - 1])) for k in range(3))
+            )
+
+    @property
+    def transform(self):
+        return self._cell_grid_info.transform_matrix
 
     @property
     def cell_grid(self) -> wp.Volume:
@@ -130,11 +100,6 @@ class Nanogrid(Geometry):
         self._ensure_face_grid()
         return self._face_grid
 
-    @property
-    def edge_grid(self) -> wp.Volume:
-        self._ensure_edge_grid()
-        return self._edge_grid
-
     def cell_count(self):
         return self._cell_ijk.shape[0]
 
@@ -149,77 +114,19 @@ class Nanogrid(Geometry):
         self._ensure_face_grid()
         return self._boundary_face_indices.shape[0]
 
-    def edge_count(self):
-        self._ensure_edge_grid()
-        return self._edge_count
-
-    def reference_cell(self) -> Cube:
-        return Cube()
-
-    def reference_side(self) -> Square:
-        return Square()
-
-    CellArg = NanogridCellArg
-
-    @cache.cached_arg_value
-    def cell_arg_value(self, device) -> CellArg:
-        args = self.CellArg()
-        self.fill_cell_arg(args, device)
-        return args
-
-    def fill_cell_arg(self, arg, device):
-        arg.cell_grid = self._cell_grid.id
-        arg.cell_ijk = self._cell_ijk
-
-        arg.inverse_transform = self._inverse_transform
-        arg.cell_volume = self._cell_volume
-
-    @wp.func
-    def cell_position(args: CellArg, s: Sample):
-        uvw = wp.vec3(args.cell_ijk[s.element_index]) + s.element_coords
-        return wp.volume_index_to_world(args.cell_grid, uvw - wp.vec3(0.5))
-
-    @wp.func
-    def cell_deformation_gradient(args: CellArg, s: Sample):
-        return wp.inverse(args.inverse_transform)
-
-    @wp.func
-    def cell_inverse_deformation_gradient(args: CellArg, s: Sample):
-        return args.inverse_transform
-
-    def supports_cell_lookup(self, device):
-        return True
-
-    @wp.func
-    def _lookup_cell_index(args: NanogridCellArg, i: int, j: int, k: int):
-        return wp.volume_lookup_index(args.cell_grid, i, j, k)
-
-    @wp.func
-    def _cell_coordinates_local(args: NanogridCellArg, cell_index: int, uvw: wp.vec3):
-        ijk = wp.vec3(args.cell_ijk[cell_index])
-        rel_pos = uvw - ijk
-        return rel_pos
-
-    @wp.func
-    def _cell_closest_point_local(args: NanogridCellArg, cell_index: int, uvw: wp.vec3):
-        ijk = wp.vec3(args.cell_ijk[cell_index])
-        rel_pos = uvw - ijk
-        coords = wp.min(wp.max(rel_pos, wp.vec3(0.0)), wp.vec3(1.0))
-        return wp.length_sq(wp.volume_index_to_world_dir(args.cell_grid, coords - rel_pos)), coords
+    @wp.struct
+    class SideIndexArg:
+        boundary_face_indices: wp.array(dtype=int)
 
-    @wp.func
-    def cell_coordinates(args: NanogridCellArg, cell_index: int, pos: wp.vec3):
-        uvw = wp.volume_world_to_index(args.cell_grid, pos) + wp.vec3(0.5)
-        return Nanogrid._cell_coordinates_local(args, cell_index, uvw)
+    def fill_side_index_arg(self, arg: SideIndexArg, device):
+        self._ensure_face_grid()
+        arg.boundary_face_indices = self._boundary_face_indices.to(device)
 
     @wp.func
-    def cell_closest_point(args: NanogridCellArg, cell_index: int, pos: wp.vec3):
-        uvw = wp.volume_world_to_index(args.cell_grid, pos) + wp.vec3(0.5)
-        dist, coords = Nanogrid._cell_closest_point_local(args, cell_index, uvw)
-        return coords, dist
+    def boundary_side_index(args: SideIndexArg, boundary_side_index: int):
+        return args.boundary_face_indices[boundary_side_index]
 
-    @staticmethod
-    def _make_filtered_cell_lookup(grid_geo, filter_func: wp.Function = None):
+    def make_filtered_cell_lookup(grid_geo, filter_func: wp.Function = None):
         suffix = f"{grid_geo.name}{filter_func.key if filter_func is not None else ''}"
 
         @cache.dynamic_func(suffix=suffix)
@@ -285,49 +192,53 @@ class Nanogrid(Geometry):
 
         return cell_lookup
 
-    def make_filtered_cell_lookup(self, filter_func):
-        return Nanogrid._make_filtered_cell_lookup(self, filter_func)
+    def _ensure_face_grid(self):
+        if self._face_ijk is None:
+            self._build_face_grid()
 
     @wp.func
-    def cell_measure(args: CellArg, s: Sample):
-        return args.cell_volume
+    def _add_axis_flag(ijk: wp.vec3i, axis: int):
+        coord = ijk[axis]
+        ijk[axis] = wp.where(coord < 0, coord & (~NanogridBase.GRID_AXIS_FLAG), coord | NanogridBase.GRID_AXIS_FLAG)
+        return ijk
 
     @wp.func
-    def cell_normal(args: CellArg, s: Sample):
-        return wp.vec3(0.0)
+    def _extract_axis_flag(ijk: wp.vec3i):
+        for ax in range(3):
+            coord = ijk[ax]
+            if coord < 0:
+                if (ijk[ax] & NanogridBase.GRID_AXIS_FLAG) == 0:
+                    ijk[ax] = ijk[ax] | NanogridBase.GRID_AXIS_FLAG
+                    return ax, ijk
+            else:
+                if (ijk[ax] & NanogridBase.GRID_AXIS_FLAG) != 0:
+                    ijk[ax] = ijk[ax] & (~NanogridBase.GRID_AXIS_FLAG)
+                    return ax, ijk
 
-    SideArg = NanogridSideArg
+        return -1, ijk
 
-    @cache.cached_arg_value
-    def side_arg_value(self, device) -> SideArg:
-        args = self.SideArg()
-        self.fill_side_arg(args, device)
-        return args
+    @wp.func
+    def _make_face_flags(axis: int, plus_cell_index: int, minus_cell_index: int):
+        plus_boundary = wp.uint8(wp.where(plus_cell_index == -1, 1, 0)) << NanogridBase.FACE_OUTER_OFFSET_BIT
+        minus_boundary = wp.uint8(wp.where(minus_cell_index == -1, 1, 0)) << NanogridBase.FACE_INNER_OFFSET_BIT
 
-    def fill_side_arg(self, arg: SideArg, device):
-        self._ensure_face_grid()
-        self.fill_cell_arg(arg.cell_arg, device)
-        arg.face_ijk = self._face_ijk.to(device)
-        arg.face_flags = self._face_flags.to(device)
-        arg.face_areas = self._face_areas
+        return wp.uint8(axis) | plus_boundary | minus_boundary
 
-    @wp.struct
-    class SideIndexArg:
-        boundary_face_indices: wp.array(dtype=int)
+    @wp.func
+    def _get_boundary_mask(flags: wp.uint8):
+        return int((flags >> NanogridBase.FACE_OUTER_OFFSET_BIT) | (flags >> NanogridBase.FACE_INNER_OFFSET_BIT)) & 1
 
-    @cache.cached_arg_value
-    def side_index_arg_value(self, device) -> SideIndexArg:
-        args = self.SideIndexArg()
-        self.fill_side_index_arg(args, device)
-        return args
+    @wp.func
+    def _get_face_axis(flags: wp.uint8):
+        return wp.int32(flags & NanogridBase.FACE_AXIS_MASK)
 
-    def fill_side_index_arg(self, arg: SideIndexArg, device):
-        self._ensure_face_grid()
-        arg.boundary_face_indices = self._boundary_face_indices.to(device)
+    @wp.func
+    def _get_face_inner_offset(flags: wp.uint8):
+        return wp.int32(flags >> NanogridBase.FACE_INNER_OFFSET_BIT) & 1
 
     @wp.func
-    def boundary_side_index(args: SideIndexArg, boundary_side_index: int):
-        return args.boundary_face_indices[boundary_side_index]
+    def _get_face_outer_offset(flags: wp.uint8):
+        return wp.int32(flags >> NanogridBase.FACE_OUTER_OFFSET_BIT) & 1
 
     @wp.func
     def _side_to_cell_coords(axis: int, flip: int, inner: float, side_coords: Coords):
@@ -342,16 +253,141 @@ class Nanogrid(Geometry):
         return Coords(cell_coords[(axis + 1 + flip) % 3], cell_coords[(axis + 2 - flip) % 3], 0.0)
 
     @wp.func
-    def _get_face_axis(flags: wp.uint8):
-        return wp.int32(flags & FACE_AXIS_MASK)
+    def _face_tangent_vecs(cell_grid: wp.uint64, axis: int, flip: int):
+        u_axis = wp.vec3()
+        v_axis = wp.vec3()
+        u_axis[(axis + 1 + flip) % 3] = 1.0
+        v_axis[(axis + 2 - flip) % 3] = 1.0
+        return wp.volume_index_to_world_dir(cell_grid, u_axis), wp.volume_index_to_world_dir(cell_grid, v_axis)
+
+
+@wp.struct
+class NanogridCellArg:
+    # Utility device functions
+    cell_grid: wp.uint64
+    cell_ijk: wp.array(dtype=wp.vec3i)
+    inverse_transform: wp.mat33
+    cell_volume: float
+
+
+@wp.struct
+class NanogridSideArg:
+    # Utility device functions
+    cell_arg: NanogridCellArg
+    face_ijk: wp.array(dtype=wp.vec3i)
+    face_flags: wp.array(dtype=wp.uint8)
+    face_areas: wp.vec3
+
+
+class Nanogrid(NanogridBase):
+    """Sparse grid geometry"""
+
+    def __init__(self, grid: wp.Volume, temporary_store: Optional[cache.TemporaryStore] = None):
+        """
+        Constructs a sparse grid geometry from an in-memory NanoVDB volume.
+
+        Args:
+            grid: The NanoVDB volume. Any type is accepted, but for indexing efficiency an index grid is recommended.
+                If `grid` is an 'on' index grid, cells will be created for active voxels only, otherwise cells will
+                be created for all leaf voxels.
+            temporary_store: shared pool from which to allocate temporary arrays
+        """
+
+        self._cell_grid = grid
+        self._cell_grid_info = grid.get_grid_info()
+
+        device = self._cell_grid.device
+        cell_ijk = wp.array(dtype=wp.vec3i, shape=(grid.get_voxel_count(),), device=device)
+        grid.get_voxels(out=cell_ijk)
+
+        node_grid = _build_node_grid(cell_ijk, grid, temporary_store)
+        node_count = node_grid.get_voxel_count()
+        node_ijk = wp.array(shape=(node_count,), dtype=wp.vec3i, device=device)
+        node_grid.get_voxels(out=node_ijk)
+
+        super().__init__(grid, cell_ijk, node_grid, node_ijk)
+
+        self._edge_count = 0
+        self._edge_grid = None
+
+    @property
+    def edge_grid(self) -> wp.Volume:
+        self._ensure_edge_grid()
+        return self._edge_grid
+
+    def edge_count(self):
+        self._ensure_edge_grid()
+        return self._edge_count
+
+    CellArg = NanogridCellArg
+
+    def fill_cell_arg(self, arg, device):
+        arg.cell_grid = self._cell_grid.id
+        arg.cell_ijk = self._cell_ijk
+
+        arg.inverse_transform = self._inverse_transform
+        arg.cell_volume = self._cell_volume
 
     @wp.func
-    def _get_face_inner_offset(flags: wp.uint8):
-        return wp.int32(flags >> FACE_INNER_OFFSET_BIT) & 1
+    def cell_position(args: CellArg, s: Sample):
+        uvw = wp.vec3(args.cell_ijk[s.element_index]) + s.element_coords
+        return wp.volume_index_to_world(args.cell_grid, uvw - wp.vec3(0.5))
 
     @wp.func
-    def _get_face_outer_offset(flags: wp.uint8):
-        return wp.int32(flags >> FACE_OUTER_OFFSET_BIT) & 1
+    def cell_deformation_gradient(args: CellArg, s: Sample):
+        return wp.inverse(args.inverse_transform)
+
+    @wp.func
+    def cell_inverse_deformation_gradient(args: CellArg, s: Sample):
+        return args.inverse_transform
+
+    def supports_cell_lookup(self, device):
+        return True
+
+    @wp.func
+    def _lookup_cell_index(args: NanogridCellArg, i: int, j: int, k: int):
+        return wp.volume_lookup_index(args.cell_grid, i, j, k)
+
+    @wp.func
+    def _cell_coordinates_local(args: NanogridCellArg, cell_index: int, uvw: wp.vec3):
+        ijk = wp.vec3(args.cell_ijk[cell_index])
+        rel_pos = uvw - ijk
+        return rel_pos
+
+    @wp.func
+    def _cell_closest_point_local(args: NanogridCellArg, cell_index: int, uvw: wp.vec3):
+        ijk = wp.vec3(args.cell_ijk[cell_index])
+        rel_pos = uvw - ijk
+        coords = wp.min(wp.max(rel_pos, wp.vec3(0.0)), wp.vec3(1.0))
+        return wp.length_sq(wp.volume_index_to_world_dir(args.cell_grid, coords - rel_pos)), coords
+
+    @wp.func
+    def cell_coordinates(args: NanogridCellArg, cell_index: int, pos: wp.vec3):
+        uvw = wp.volume_world_to_index(args.cell_grid, pos) + wp.vec3(0.5)
+        return Nanogrid._cell_coordinates_local(args, cell_index, uvw)
+
+    @wp.func
+    def cell_closest_point(args: NanogridCellArg, cell_index: int, pos: wp.vec3):
+        uvw = wp.volume_world_to_index(args.cell_grid, pos) + wp.vec3(0.5)
+        dist, coords = Nanogrid._cell_closest_point_local(args, cell_index, uvw)
+        return coords, dist
+
+    @wp.func
+    def cell_measure(args: CellArg, s: Sample):
+        return args.cell_volume
+
+    @wp.func
+    def cell_normal(args: CellArg, s: Sample):
+        return wp.vec3(0.0)
+
+    SideArg = NanogridSideArg
+
+    def fill_side_arg(self, arg: SideArg, device):
+        self._ensure_face_grid()
+        self.fill_cell_arg(arg.cell_arg, device)
+        arg.face_ijk = self._face_ijk.to(device)
+        arg.face_flags = self._face_flags.to(device)
+        arg.face_areas = self._face_areas
 
     @wp.func
     def side_position(args: SideArg, s: Sample):
@@ -365,14 +401,6 @@ class Nanogrid(Geometry):
         cell_grid = args.cell_arg.cell_grid
         return wp.volume_index_to_world(cell_grid, uvw - wp.vec3(0.5))
 
-    @wp.func
-    def _face_tangent_vecs(cell_grid: wp.uint64, axis: int, flip: int):
-        u_axis = wp.vec3()
-        v_axis = wp.vec3()
-        u_axis[(axis + 1 + flip) % 3] = 1.0
-        v_axis[(axis + 2 - flip) % 3] = 1.0
-        return wp.volume_index_to_world_dir(cell_grid, u_axis), wp.volume_index_to_world_dir(cell_grid, v_axis)
-
     @wp.func
     def side_deformation_gradient(args: SideArg, s: Sample):
         flags = args.face_flags[s.element_index]
@@ -509,19 +537,15 @@ class Nanogrid(Geometry):
             _build_face_flags,
             dim=face_count,
             device=device,
-            inputs=[self._cell_grid.id, self._face_ijk, self._face_flags, boundary_face_mask.array],
+            inputs=[self._cell_grid.id, self._face_ijk, self._face_flags, boundary_face_mask],
         )
-        boundary_face_indices, _ = utils.masked_indices(boundary_face_mask.array)
+        boundary_face_indices, _ = utils.masked_indices(boundary_face_mask)
         self._boundary_face_indices = boundary_face_indices.detach()
 
     def _build_edge_grid(self, temporary_store: Optional[cache.TemporaryStore] = None):
         self._edge_grid = _build_edge_grid(self._cell_ijk, self._cell_grid, temporary_store)
         self._edge_count = self._edge_grid.get_voxel_count()
 
-    def _ensure_face_grid(self):
-        if self._face_ijk is None:
-            self._build_face_grid()
-
     def _ensure_edge_grid(self):
         if self._edge_grid is None:
             self._build_edge_grid()
@@ -543,13 +567,13 @@ def _cell_face_indices(
 ):
     cell = wp.tid()
     ijk = cell_ijk[cell]
-    node_ijk[cell, 0] = _add_axis_flag(ijk, 0)
-    node_ijk[cell, 1] = _add_axis_flag(ijk, 1)
-    node_ijk[cell, 2] = _add_axis_flag(ijk, 2)
+    node_ijk[cell, 0] = NanogridBase._add_axis_flag(ijk, 0)
+    node_ijk[cell, 1] = NanogridBase._add_axis_flag(ijk, 1)
+    node_ijk[cell, 2] = NanogridBase._add_axis_flag(ijk, 2)
 
-    node_ijk[cell, 3] = _add_axis_flag(ijk + wp.vec3i(1, 0, 0), 0)
-    node_ijk[cell, 4] = _add_axis_flag(ijk + wp.vec3i(0, 1, 0), 1)
-    node_ijk[cell, 5] = _add_axis_flag(ijk + wp.vec3i(0, 0, 1), 2)
+    node_ijk[cell, 3] = NanogridBase._add_axis_flag(ijk + wp.vec3i(1, 0, 0), 0)
+    node_ijk[cell, 4] = NanogridBase._add_axis_flag(ijk + wp.vec3i(0, 1, 0), 1)
+    node_ijk[cell, 5] = NanogridBase._add_axis_flag(ijk + wp.vec3i(0, 0, 1), 2)
 
 
 @wp.kernel
@@ -559,32 +583,30 @@ def _cell_edge_indices(
 ):
     cell = wp.tid()
     ijk = cell_ijk[cell]
-    edge_ijk[cell, 0] = _add_axis_flag(ijk, 0)
-    edge_ijk[cell, 1] = _add_axis_flag(ijk, 1)
-    edge_ijk[cell, 2] = _add_axis_flag(ijk, 2)
+    edge_ijk[cell, 0] = NanogridBase._add_axis_flag(ijk, 0)
+    edge_ijk[cell, 1] = NanogridBase._add_axis_flag(ijk, 1)
+    edge_ijk[cell, 2] = NanogridBase._add_axis_flag(ijk, 2)
 
-    edge_ijk[cell, 3] = _add_axis_flag(ijk + wp.vec3i(0, 1, 0), 0)
-    edge_ijk[cell, 4] = _add_axis_flag(ijk + wp.vec3i(0, 0, 1), 1)
-    edge_ijk[cell, 5] = _add_axis_flag(ijk + wp.vec3i(1, 0, 0), 2)
+    edge_ijk[cell, 3] = NanogridBase._add_axis_flag(ijk + wp.vec3i(0, 1, 0), 0)
+    edge_ijk[cell, 4] = NanogridBase._add_axis_flag(ijk + wp.vec3i(0, 0, 1), 1)
+    edge_ijk[cell, 5] = NanogridBase._add_axis_flag(ijk + wp.vec3i(1, 0, 0), 2)
 
-    edge_ijk[cell, 6] = _add_axis_flag(ijk + wp.vec3i(0, 1, 1), 0)
-    edge_ijk[cell, 7] = _add_axis_flag(ijk + wp.vec3i(1, 0, 1), 1)
-    edge_ijk[cell, 8] = _add_axis_flag(ijk + wp.vec3i(1, 1, 0), 2)
+    edge_ijk[cell, 6] = NanogridBase._add_axis_flag(ijk + wp.vec3i(0, 1, 1), 0)
+    edge_ijk[cell, 7] = NanogridBase._add_axis_flag(ijk + wp.vec3i(1, 0, 1), 1)
+    edge_ijk[cell, 8] = NanogridBase._add_axis_flag(ijk + wp.vec3i(1, 1, 0), 2)
 
-    edge_ijk[cell, 9] = _add_axis_flag(ijk + wp.vec3i(0, 0, 1), 0)
-    edge_ijk[cell, 10] = _add_axis_flag(ijk + wp.vec3i(1, 0, 0), 1)
-    edge_ijk[cell, 11] = _add_axis_flag(ijk + wp.vec3i(0, 1, 0), 2)
+    edge_ijk[cell, 9] = NanogridBase._add_axis_flag(ijk + wp.vec3i(0, 0, 1), 0)
+    edge_ijk[cell, 10] = NanogridBase._add_axis_flag(ijk + wp.vec3i(1, 0, 0), 1)
+    edge_ijk[cell, 11] = NanogridBase._add_axis_flag(ijk + wp.vec3i(0, 1, 0), 2)
 
 
 def _build_node_grid(cell_ijk, grid: wp.Volume, temporary_store: cache.TemporaryStore):
     cell_count = cell_ijk.shape[0]
 
     cell_nodes = cache.borrow_temporary(temporary_store, shape=(cell_count, 8), dtype=wp.vec3i, device=cell_ijk.device)
-    wp.launch(
-        _cell_node_indices, dim=cell_nodes.array.shape, inputs=[cell_ijk, cell_nodes.array], device=cell_ijk.device
-    )
+    wp.launch(_cell_node_indices, dim=cell_nodes.shape, inputs=[cell_ijk, cell_nodes], device=cell_ijk.device)
     node_grid = wp.Volume.allocate_by_voxels(
-        cell_nodes.array.flatten(), voxel_size=grid.get_voxel_size(), device=cell_ijk.device
+        cell_nodes.flatten(), voxel_size=grid.get_voxel_size(), device=cell_ijk.device
     )
 
     return node_grid
@@ -594,9 +616,9 @@ def _build_face_grid(cell_ijk, grid: wp.Volume, temporary_store: cache.Temporary
     cell_count = cell_ijk.shape[0]
 
     cell_faces = cache.borrow_temporary(temporary_store, shape=(cell_count, 6), dtype=wp.vec3i, device=cell_ijk.device)
-    wp.launch(_cell_face_indices, dim=cell_count, inputs=[cell_ijk, cell_faces.array], device=cell_ijk.device)
+    wp.launch(_cell_face_indices, dim=cell_count, inputs=[cell_ijk, cell_faces], device=cell_ijk.device)
     face_grid = wp.Volume.allocate_by_voxels(
-        cell_faces.array.flatten(), voxel_size=grid.get_voxel_size(), device=cell_ijk.device
+        cell_faces.flatten(), voxel_size=grid.get_voxel_size(), device=cell_ijk.device
     )
 
     return face_grid
@@ -606,27 +628,14 @@ def _build_edge_grid(cell_ijk, grid: wp.Volume, temporary_store: cache.Temporary
     cell_count = cell_ijk.shape[0]
 
     cell_edges = cache.borrow_temporary(temporary_store, shape=(cell_count, 12), dtype=wp.vec3i, device=cell_ijk.device)
-    wp.launch(_cell_edge_indices, dim=cell_count, inputs=[cell_ijk, cell_edges.array], device=cell_ijk.device)
+    wp.launch(_cell_edge_indices, dim=cell_count, inputs=[cell_ijk, cell_edges], device=cell_ijk.device)
     edge_grid = wp.Volume.allocate_by_voxels(
-        cell_edges.array.flatten(), voxel_size=grid.get_voxel_size(), device=cell_ijk.device
+        cell_edges.flatten(), voxel_size=grid.get_voxel_size(), device=cell_ijk.device
     )
 
     return edge_grid
 
 
-@wp.func
-def _make_face_flags(axis: int, plus_cell_index: int, minus_cell_index: int):
-    plus_boundary = wp.uint8(wp.where(plus_cell_index == -1, 1, 0)) << FACE_OUTER_OFFSET_BIT
-    minus_boundary = wp.uint8(wp.where(minus_cell_index == -1, 1, 0)) << FACE_INNER_OFFSET_BIT
-
-    return wp.uint8(axis) | plus_boundary | minus_boundary
-
-
-@wp.func
-def _get_boundary_mask(flags: wp.uint8):
-    return int((flags >> FACE_OUTER_OFFSET_BIT) | (flags >> FACE_INNER_OFFSET_BIT)) & 1
-
-
 @wp.kernel
 def _build_face_flags(
     cell_grid: wp.uint64,
@@ -636,7 +645,7 @@ def _build_face_flags(
 ):
     face = wp.tid()
 
-    axis, ijk = _extract_axis_flag(face_ijk[face])
+    axis, ijk = NanogridBase._extract_axis_flag(face_ijk[face])
 
     ijk_minus = ijk
     ijk_minus[axis] -= 1
@@ -646,6 +655,6 @@ def _build_face_flags(
 
     face_ijk[face] = ijk
 
-    flags = _make_face_flags(axis, plus_cell_index, minus_cell_index)
+    flags = NanogridBase._make_face_flags(axis, plus_cell_index, minus_cell_index)
     face_flags[face] = flags
-    boundary_face_mask[face] = _get_boundary_mask(flags)
+    boundary_face_mask[face] = NanogridBase._get_boundary_mask(flags)