warp-lang 1.7.2rc1__py3-none-manylinux_2_34_aarch64.whl → 1.8.0__py3-none-manylinux_2_34_aarch64.whl

warp/fem/geometry/hexmesh.py

@@ -33,9 +33,8 @@ class HexmeshCellArg:
     hex_vertex_indices: wp.array2d(dtype=int)
     positions: wp.array(dtype=wp.vec3)
 
-    # for neighbor cell lookup
-    vertex_hex_offsets: wp.array(dtype=int)
-    vertex_hex_indices: wp.array(dtype=int)
+    # for global cell lookup
+    hex_bvh: wp.uint64
 
 
 @wp.struct
@@ -131,20 +130,28 @@ class Hexmesh(Geometry):
     dimension = 3
 
     def __init__(
-        self, hex_vertex_indices: wp.array, positions: wp.array, temporary_store: Optional[TemporaryStore] = None
+        self,
+        hex_vertex_indices: wp.array,
+        positions: wp.array,
+        assume_parallelepiped_cells=False,
+        build_bvh: bool = False,
+        temporary_store: Optional[TemporaryStore] = None,
     ):
         """
-        Constructs a tetrahedral mesh.
+        Constructs a hexahedral mesh.
 
         Args:
             hex_vertex_indices: warp array of shape (num_hexes, 8) containing vertex indices for each hex
                 following standard ordering (bottom face vertices in counter-clockwise order, then similarly for upper face)
             positions: warp array of shape (num_vertices, 3) containing 3d position for each vertex
+            assume_parallelepiped: If true, assume that all cells are parallelepipeds (cheaper position/gradient evaluations)
+            build_bvh: Whether to also build the hex BVH, which is necessary for the global `fem.lookup` operator
             temporary_store: shared pool from which to allocate temporary arrays
         """
 
         self.hex_vertex_indices = hex_vertex_indices
         self.positions = positions
+        self.parallelepiped_cells = assume_parallelepiped_cells
 
         self._face_vertex_indices: wp.array = None
         self._face_hex_indices: wp.array = None
@@ -155,7 +162,25 @@ class Hexmesh(Geometry):
         self._edge_count = 0
         self._build_topology(temporary_store)
 
+        # Use cheaper variants if we know that cells are parallelepipeds (i.e. linearly transformed)
+        # (Cells only, not as much difference for sides)
+        self.cell_position = (
+            self._cell_position_parallelepiped if assume_parallelepiped_cells else self._cell_position_generic
+        )
+        self.cell_deformation_gradient = (
+            self._cell_deformation_gradient_parallelepiped
+            if assume_parallelepiped_cells
+            else self._cell_deformation_gradient_generic
+        )
+
         self._make_default_dependent_implementations()
+        self.cell_coordinates = self._make_cell_coordinates(assume_linear=assume_parallelepiped_cells)
+        self.side_coordinates = self._make_side_coordinates(assume_linear=assume_parallelepiped_cells)
+        self.cell_closest_point = self._make_cell_closest_point(assume_linear=assume_parallelepiped_cells)
+        self.side_closest_point = self._make_side_closest_point(assume_linear=assume_parallelepiped_cells)
+
+        if build_bvh:
+            self.build_bvh(self.positions.device)
 
     def cell_count(self):
         return self.hex_vertex_indices.shape[0]
@@ -203,19 +228,18 @@ class Hexmesh(Geometry):
 
     # Geometry device interface
 
-    @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, args: CellArg, device):
         args.hex_vertex_indices = self.hex_vertex_indices.to(device)
         args.positions = self.positions.to(device)
-        args.vertex_hex_offsets = self._vertex_hex_offsets.to(device)
-        args.vertex_hex_indices = self._vertex_hex_indices.to(device)
-
-        return args
+        args.hex_bvh = self.bvh_id(device)
 
     @wp.func
-    def cell_position(args: CellArg, s: Sample):
+    def _cell_position_generic(args: CellArg, s: Sample):
         hex_idx = args.hex_vertex_indices[s.element_index]
 
         w_p = s.element_coords
@@ -242,9 +266,18 @@ class Hexmesh(Geometry):
         )
 
     @wp.func
-    def cell_deformation_gradient(cell_arg: CellArg, s: Sample):
+    def _cell_position_parallelepiped(args: CellArg, s: Sample):
+        hex_idx = args.hex_vertex_indices[s.element_index]
+        w = s.element_coords
+        p0 = args.positions[hex_idx[0]]
+        p1 = args.positions[hex_idx[1]]
+        p2 = args.positions[hex_idx[3]]
+        p3 = args.positions[hex_idx[4]]
+        return w[0] * p1 + w[1] * p2 + w[2] * p3 + (1.0 - w[0] - w[1] - w[2]) * p0
+
+    @wp.func
+    def _cell_deformation_gradient_generic(cell_arg: CellArg, s: Sample):
         """Deformation gradient at `coords`"""
-        """Transposed deformation gradient at `coords`"""
         hex_idx = cell_arg.hex_vertex_indices[s.element_index]
 
         w_p = s.element_coords
@@ -261,31 +294,43 @@ class Hexmesh(Geometry):
             + wp.outer(cell_arg.positions[hex_idx[7]], wp.vec3(-w_p[1] * w_p[2], w_m[0] * w_p[2], w_m[0] * w_p[1]))
         )
 
+    @wp.func
+    def _cell_deformation_gradient_parallelepiped(cell_arg: CellArg, s: Sample):
+        """Deformation gradient at `coords`"""
+        hex_idx = cell_arg.hex_vertex_indices[s.element_index]
+
+        p0 = cell_arg.positions[hex_idx[0]]
+        p1 = cell_arg.positions[hex_idx[1]]
+        p2 = cell_arg.positions[hex_idx[3]]
+        p3 = cell_arg.positions[hex_idx[4]]
+        return wp.matrix_from_cols(p1 - p0, p2 - p0, p3 - p0)
+
     @cached_arg_value
     def side_index_arg_value(self, device) -> SideIndexArg:
         args = self.SideIndexArg()
+        self.fill_side_index_arg(args, device)
+        return args
 
+    def fill_side_index_arg(self, args: SideIndexArg, device):
         args.boundary_face_indices = self._boundary_face_indices.to(device)
 
-        return args
-
     @wp.func
     def boundary_side_index(args: SideIndexArg, boundary_side_index: int):
         """Boundary side to side index"""
 
         return args.boundary_face_indices[boundary_side_index]
 
-    @cached_arg_value
     def side_arg_value(self, device) -> CellArg:
         args = self.SideArg()
+        self.fill_side_arg(args, device)
+        return args
 
-        args.cell_arg = self.cell_arg_value(device)
+    def fill_side_arg(self, args: SideArg, device):
+        self.fill_cell_arg(args.cell_arg, device)
         args.face_vertex_indices = self._face_vertex_indices.to(device)
         args.face_hex_indices = self._face_hex_indices.to(device)
         args.face_hex_face_orientation = self._face_hex_face_orientation.to(device)
 
-        return args
-
     @wp.func
     def side_position(args: SideArg, s: Sample):
         face_idx = args.face_vertex_indices[s.element_index]
@@ -332,7 +377,7 @@ class Hexmesh(Geometry):
 
     @wp.func
     def _hex_local_face_coords(hex_coords: Coords, face_index: int):
-        # Coordinatex in local face coordinates system
+        # Coordinates in local face coordinates system
         # Sign of last coordinate (out of face)
 
         face_coords = wp.vec2(
@@ -909,3 +954,24 @@ class Hexmesh(Geometry):
                         + unique_beg
                     )
                     hex_edge_indices[t][k] = edge_id
+
+    @wp.func
+    def cell_bvh_id(cell_arg: HexmeshCellArg):
+        return cell_arg.hex_bvh
+
+    @wp.func
+    def cell_bounds(cell_arg: HexmeshCellArg, cell_index: ElementIndex):
+        vidx = cell_arg.hex_vertex_indices[cell_index]
+        p0 = cell_arg.positions[vidx[0]]
+        p1 = cell_arg.positions[vidx[1]]
+        p2 = cell_arg.positions[vidx[2]]
+        p3 = cell_arg.positions[vidx[3]]
+        lo0, up0 = wp.min(wp.min(p0, p1), wp.min(p2, p3)), wp.max(wp.max(p0, p1), wp.max(p2, p3))
+
+        p4 = cell_arg.positions[vidx[4]]
+        p5 = cell_arg.positions[vidx[5]]
+        p6 = cell_arg.positions[vidx[6]]
+        p7 = cell_arg.positions[vidx[7]]
+        lo1, up1 = wp.min(wp.min(p4, p5), wp.min(p6, p7)), wp.max(wp.max(p4, p5), wp.max(p6, p7))
+
+        return wp.min(lo0, lo1), wp.max(up0, up1)

warp/fem/geometry/nanogrid.py

@@ -13,13 +13,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional
+from typing import Any, Optional
 
 import numpy as np
 
 import warp as wp
 from warp.fem import cache, utils
-from warp.fem.linalg import basis_element
 from warp.fem.types import NULL_ELEMENT_INDEX, OUTSIDE, Coords, ElementIndex, Sample, make_free_sample
 
 from .element import Cube, Square
@@ -111,6 +110,13 @@ class Nanogrid(Geometry):
         self._edge_grid = None
         self._edge_count = 0
 
+        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))
+        )
+
     @property
     def cell_grid(self) -> wp.Volume:
         return self._cell_grid
@@ -158,14 +164,15 @@ class Nanogrid(Geometry):
     @cache.cached_arg_value
     def cell_arg_value(self, device) -> CellArg:
         args = self.CellArg()
-        args.cell_grid = self._cell_grid.id
-        args.cell_ijk = self._cell_ijk
+        self.fill_cell_arg(args, device)
+        return args
 
-        transform = np.array(self._cell_grid_info.transform_matrix).reshape(3, 3)
-        args.inverse_transform = wp.mat33f(np.linalg.inv(transform))
-        args.cell_volume = abs(np.linalg.det(transform))
+    def fill_cell_arg(self, arg, device):
+        arg.cell_grid = self._cell_grid.id
+        arg.cell_ijk = self._cell_ijk
 
-        return args
+        arg.inverse_transform = self._inverse_transform
+        arg.cell_volume = self._cell_volume
 
     @wp.func
     def cell_position(args: CellArg, s: Sample):
@@ -180,66 +187,106 @@ class Nanogrid(Geometry):
     def cell_inverse_deformation_gradient(args: CellArg, s: Sample):
         return args.inverse_transform
 
-    @wp.func
-    def cell_lookup(args: CellArg, pos: wp.vec3):
-        uvw = wp.volume_world_to_index(args.cell_grid, pos) + wp.vec3(0.5)
-        ijk = wp.vec3i(int(wp.floor(uvw[0])), int(wp.floor(uvw[1])), int(wp.floor(uvw[2])))
-        cell_index = wp.volume_lookup_index(args.cell_grid, ijk[0], ijk[1], ijk[2])
-
-        coords = uvw - wp.vec3(ijk)
-        if cell_index == -1:
-            if wp.min(coords) == 0.0 or wp.max(coords) == 1.0:
-                il = wp.where(coords[0] > 0.5, 0, -1)
-                jl = wp.where(coords[1] > 0.5, 0, -1)
-                kl = wp.where(coords[2] > 0.5, 0, -1)
-
-                for n in range(8):
-                    ni = n >> 2
-                    nj = (n & 2) >> 1
-                    nk = n & 1
-                    nijk = ijk + wp.vec3i(ni + il, nj + jl, nk + kl)
-
-                    coords = uvw - wp.vec3(nijk)
-                    if wp.min(coords) >= 0.0 and wp.max(coords) <= 1.0:
-                        cell_index = wp.volume_lookup_index(args.cell_grid, nijk[0], nijk[1], nijk[2])
-                        if cell_index != -1:
-                            return make_free_sample(cell_index, coords)
-
-            return make_free_sample(NULL_ELEMENT_INDEX, Coords(OUTSIDE))
-
-        return make_free_sample(cell_index, coords)
+    def supports_cell_lookup(self, device):
+        return True
 
     @wp.func
-    def _project_on_voxel_at_origin(coords: wp.vec3):
-        proj_coords = wp.min(wp.max(coords, wp.vec3(0.0)), wp.vec3(1.0))
-        return wp.length_sq(coords - proj_coords), proj_coords
+    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_lookup(args: CellArg, pos: wp.vec3, guess: Sample):
-        s_global = Nanogrid.cell_lookup(args, pos)
-
-        if s_global.element_index != NULL_ELEMENT_INDEX:
-            return s_global
+    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
 
-        closest_voxel = int(NULL_ELEMENT_INDEX)
-        closest_coords = Coords(OUTSIDE)
-        closest_dist = float(1.0e8)
+    @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
 
-        # project to closest in stencil
+    @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)
-        cell_ijk = args.cell_ijk[guess.element_index]
-        for ni in range(-1, 2):
-            for nj in range(-1, 2):
-                for nk in range(-1, 2):
-                    nijk = cell_ijk + wp.vec3i(ni, nj, nk)
-                    cell_idx = wp.volume_lookup_index(args.cell_grid, nijk[0], nijk[1], nijk[2])
-                    dist, coords = Nanogrid._project_on_voxel_at_origin(uvw - wp.vec3(nijk))
-                    if cell_idx != -1 and dist <= closest_dist:
-                        closest_dist = dist
-                        closest_voxel = cell_idx
-                        closest_coords = coords
+        return Nanogrid._cell_coordinates_local(args, cell_index, uvw)
 
-        return make_free_sample(closest_voxel, closest_coords)
+    @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
+
+    @staticmethod
+    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)
+        def cell_lookup(args: grid_geo.CellArg, pos: wp.vec3, max_dist: float, filter_data: Any, filter_target: Any):
+            grid = args.cell_grid
+
+            # Start at corresponding voxel
+            uvw = wp.volume_world_to_index(grid, pos) + wp.vec3(0.5)
+            i, j, k = int(wp.floor(uvw[0])), int(wp.floor(uvw[1])), int(wp.floor(uvw[2]))
+            cell_index = grid_geo._lookup_cell_index(args, i, j, k)
+
+            if cell_index != -1:
+                coords = grid_geo._cell_coordinates_local(args, cell_index, uvw)
+                if wp.static(filter_func is None):
+                    return make_free_sample(cell_index, coords)
+                else:
+                    if filter_func(filter_data, cell_index) == filter_target:
+                        return make_free_sample(cell_index, coords)
+
+            # Iterate over increasingly larger neighborhoods
+            cell_size = wp.vec3(
+                wp.length(wp.volume_index_to_world_dir(grid, wp.vec3(1.0, 0.0, 0.0))),
+                wp.length(wp.volume_index_to_world_dir(grid, wp.vec3(0.0, 1.0, 0.0))),
+                wp.length(wp.volume_index_to_world_dir(grid, wp.vec3(0.0, 0.0, 1.0))),
+            )
+
+            offset = float(0.5)
+            min_cell_size = wp.min(cell_size)
+            max_offset = wp.ceil(max_dist / min_cell_size)
+            scales = wp.cw_div(wp.vec3(min_cell_size), wp.vec3(cell_size))
+
+            closest_cell = NULL_ELEMENT_INDEX
+            closest_coords = Coords()
+
+            while closest_cell == NULL_ELEMENT_INDEX:
+                uvw_min = wp.vec3i(uvw - offset * scales)
+                uvw_max = wp.vec3i(uvw + offset * scales) + wp.vec3i(1)
+
+                closest_dist = min_cell_size * min_cell_size * float(offset * offset)
+
+                for i in range(uvw_min[0], uvw_max[0]):
+                    for j in range(uvw_min[1], uvw_max[1]):
+                        for k in range(uvw_min[2], uvw_max[2]):
+                            cell_index = grid_geo._lookup_cell_index(args, i, j, k)
+                            if cell_index == -1:
+                                continue
+
+                            if wp.static(filter_func is not None):
+                                if filter_func(filter_data, cell_index) != filter_target:
+                                    continue
+                            dist, coords = grid_geo._cell_closest_point_local(args, cell_index, uvw)
+
+                            if dist <= closest_dist:
+                                closest_dist = dist
+                                closest_coords = coords
+                                closest_cell = cell_index
+
+                if offset >= max_offset:
+                    break
+                offset = wp.min(3.0 * offset, max_offset)
+
+            return make_free_sample(closest_cell, closest_coords)
+
+        return cell_lookup
+
+    def make_filtered_cell_lookup(self, filter_func):
+        return Nanogrid._make_filtered_cell_lookup(self, filter_func)
 
     @wp.func
     def cell_measure(args: CellArg, s: Sample):
@@ -253,43 +300,47 @@ class Nanogrid(Geometry):
 
     @cache.cached_arg_value
     def side_arg_value(self, device) -> SideArg:
-        self._ensure_face_grid()
-
         args = self.SideArg()
-        args.cell_arg = self.cell_arg_value(device)
-        args.face_ijk = self._face_ijk.to(device)
-        args.face_flags = self._face_flags.to(device)
-        transform = np.array(self._cell_grid_info.transform_matrix).reshape(3, 3)
-        args.face_areas = wp.vec3(
-            tuple(np.linalg.norm(np.cross(transform[:, k - 2], transform[:, k - 1])) for k in range(3))
-        )
-
+        self.fill_side_arg(args, device)
         return args
 
+    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.struct
     class SideIndexArg:
         boundary_face_indices: wp.array(dtype=int)
 
     @cache.cached_arg_value
     def side_index_arg_value(self, device) -> SideIndexArg:
-        self._ensure_face_grid()
-
         args = self.SideIndexArg()
-        args.boundary_face_indices = self._boundary_face_indices.to(device)
+        self.fill_side_index_arg(args, device)
         return args
 
+    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 boundary_side_index(args: SideIndexArg, boundary_side_index: int):
         return args.boundary_face_indices[boundary_side_index]
 
     @wp.func
-    def _side_to_cell_coords(axis: int, inner: float, side_coords: Coords):
+    def _side_to_cell_coords(axis: int, flip: int, inner: float, side_coords: Coords):
         uvw = wp.vec3()
         uvw[axis] = inner
-        uvw[(axis + 1) % 3] = side_coords[0]
-        uvw[(axis + 2) % 3] = side_coords[1]
+        uvw[(axis + 1 + flip) % 3] = side_coords[0]
+        uvw[(axis + 2 - flip) % 3] = side_coords[1]
         return uvw
 
+    @wp.func
+    def _cell_to_side_coords(axis: int, flip: int, cell_coords: Coords):
+        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)
@@ -305,18 +356,21 @@ class Nanogrid(Geometry):
     @wp.func
     def side_position(args: SideArg, s: Sample):
         ijk = args.face_ijk[s.element_index]
-        axis = Nanogrid._get_face_axis(args.face_flags[s.element_index])
+        flags = args.face_flags[s.element_index]
+        axis = Nanogrid._get_face_axis(flags)
+        flip = Nanogrid._get_face_inner_offset(flags)
 
-        uvw = wp.vec3(ijk) + Nanogrid._side_to_cell_coords(axis, 0.0, s.element_coords)
+        uvw = wp.vec3(ijk) + Nanogrid._side_to_cell_coords(axis, flip, 0.0, s.element_coords)
 
         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 = basis_element(wp.vec3(), (axis + 1 + flip) % 3)
-        v_axis = basis_element(wp.vec3(), (axis + 2 - flip) % 3)
-
+        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
@@ -382,15 +436,17 @@ class Nanogrid(Geometry):
     def side_inner_cell_coords(args: SideArg, side_index: ElementIndex, side_coords: Coords):
         flags = args.face_flags[side_index]
         axis = Nanogrid._get_face_axis(flags)
+        flip = Nanogrid._get_face_inner_offset(flags)
         offset = float(Nanogrid._get_face_inner_offset(flags))
-        return Nanogrid._side_to_cell_coords(axis, 1.0 - offset, side_coords)
+        return Nanogrid._side_to_cell_coords(axis, flip, 1.0 - offset, side_coords)
 
     @wp.func
     def side_outer_cell_coords(args: SideArg, side_index: ElementIndex, side_coords: Coords):
         flags = args.face_flags[side_index]
         axis = Nanogrid._get_face_axis(flags)
+        flip = Nanogrid._get_face_inner_offset(flags)
         offset = float(Nanogrid._get_face_outer_offset(flags))
-        return Nanogrid._side_to_cell_coords(axis, offset, side_coords)
+        return Nanogrid._side_to_cell_coords(axis, flip, offset, side_coords)
 
     @wp.func
     def side_from_cell_coords(
@@ -401,20 +457,44 @@ class Nanogrid(Geometry):
     ):
         flags = args.face_flags[side_index]
         axis = Nanogrid._get_face_axis(flags)
+        flip = Nanogrid._get_face_inner_offset(flags)
 
         cell_ijk = args.cell_arg.cell_ijk[element_index]
         side_ijk = args.face_ijk[side_index]
 
         on_side = float(side_ijk[axis] - cell_ijk[axis]) == element_coords[axis]
 
-        return wp.where(
-            on_side, Coords(element_coords[(axis + 1) % 3], element_coords[(axis + 2) % 3], 0.0), Coords(OUTSIDE)
-        )
+        return wp.where(on_side, Nanogrid._cell_to_side_coords(axis, flip, element_coords), Coords(OUTSIDE))
 
     @wp.func
     def side_to_cell_arg(side_arg: SideArg):
         return side_arg.cell_arg
 
+    @wp.func
+    def side_coordinates(args: SideArg, side_index: int, pos: wp.vec3):
+        cell_arg = args.cell_arg
+
+        ijk = args.face_ijk[side_index]
+        cell_coords = wp.volume_world_to_index(cell_arg.cell_grid, pos) + wp.vec3(0.5) - wp.vec3(ijk)
+
+        flags = args.face_flags[side_index]
+        axis = Nanogrid._get_face_axis(flags)
+        flip = Nanogrid._get_face_inner_offset(flags)
+        return Nanogrid._cell_to_side_coords(axis, flip, cell_coords)
+
+    @wp.func
+    def side_closest_point(args: SideArg, side_index: int, pos: wp.vec3):
+        coords = Nanogrid.side_coordinates(args, side_index, pos)
+
+        proj_coords = Coords(wp.clamp(coords[0], 0.0, 1.0), wp.clamp(coords[1], 0.0, 1.0), 0.0)
+
+        flags = args.face_flags[side_index]
+        axis = Nanogrid._get_face_axis(flags)
+        flip = Nanogrid._get_face_inner_offset(flags)
+        cell_coord_offset = Nanogrid._side_to_cell_coords(axis, flip, 0, coords - proj_coords)
+
+        return proj_coords, wp.length_sq(wp.volume_index_to_world_dir(args.cell_grid, cell_coord_offset))
+
     def _build_face_grid(self, temporary_store: Optional[cache.TemporaryStore] = None):
         device = self._cell_grid.device
         self._face_grid = _build_face_grid(self._cell_ijk, self._cell_grid, temporary_store)