warp-lang 1.7.2rc1__py3-none-win_amd64.whl → 1.8.1__py3-none-win_amd64.whl

warp/fem/field/virtual.py

@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Set
+from typing import Any, ClassVar, Dict, Set
 
 import warp as wp
 import warp.fem.operator as operator
@@ -23,6 +23,7 @@ from warp.fem.linalg import basis_coefficient, generalized_inner, generalized_ou
 from warp.fem.quadrature import Quadrature
 from warp.fem.space import FunctionSpace, SpacePartition, SpaceRestriction
 from warp.fem.types import NULL_NODE_INDEX, DofIndex, Sample, get_node_coord, get_node_index_in_element
+from warp.fem.utils import type_zero_element
 
 from .field import SpaceField
 
@@ -30,33 +31,52 @@ from .field import SpaceField
 class AdjointField(SpaceField):
     """Adjoint of a discrete field with respect to its degrees of freedom"""
 
-    def __init__(self, space: FunctionSpace, space_partition: SpaceRestriction):
+    _dynamic_attribute_constructors: ClassVar = {
+        "EvalArg": lambda obj: obj._make_eval_arg(),
+        "ElementEvalArg": lambda obj: obj._make_element_eval_arg(),
+        "eval_degree": lambda obj: obj._make_eval_degree(),
+        "eval_inner": lambda obj: obj._make_eval_inner(),
+        "eval_grad_inner": lambda obj: obj._make_eval_grad_inner(),
+        "eval_div_inner": lambda obj: obj._make_eval_div_inner(),
+        "eval_outer": lambda obj: obj._make_eval_outer(),
+        "eval_grad_outer": lambda obj: obj._make_eval_grad_outer(),
+        "eval_div_outer": lambda obj: obj._make_eval_div_outer(),
+        "node_count": lambda obj: obj._make_node_count(),
+        "at_node": lambda obj: obj._make_at_node(),
+        "node_index": lambda obj: obj._make_node_index(),
+    }
+
+    def __init__(self, space: FunctionSpace, space_partition: SpacePartition):
         super().__init__(space, space_partition=space_partition)
 
         self.node_dof_count = self.space.NODE_DOF_COUNT
         self.value_dof_count = self.space.VALUE_DOF_COUNT
 
-        self.EvalArg = self.space.SpaceArg
-        self.ElementEvalArg = self._make_element_eval_arg()
-
-        self.eval_arg_value = self.space.space_arg_value
-
-        self.eval_degree = self._make_eval_degree()
-        self.eval_inner = self._make_eval_inner()
-        self.eval_grad_inner = self._make_eval_grad_inner()
-        self.eval_div_inner = self._make_eval_div_inner()
-        self.eval_outer = self._make_eval_outer()
-        self.eval_grad_outer = self._make_eval_grad_outer()
-        self.eval_div_outer = self._make_eval_div_outer()
-        self.at_node = self._make_at_node()
+        cache.setup_dynamic_attributes(self)
 
     @property
     def name(self) -> str:
         return f"{self.__class__.__name__}{self.space.name}{self._space_partition.name}"
 
-    def _make_element_eval_arg(self):
-        from warp.fem import cache
+    @cache.cached_arg_value
+    def eval_arg_value(self, device):
+        arg = self.EvalArg()
+        self.fill_eval_arg(arg, device)
+        return arg
+
+    def fill_eval_arg(self, arg, device):
+        self.space.fill_space_arg(arg.space_arg, device)
+        self.space.topology.fill_topo_arg(arg.topo_arg, device)
 
+    def _make_eval_arg(self):
+        @cache.dynamic_struct(suffix=self.name)
+        class EvalArg:
+            space_arg: self.space.SpaceArg
+            topo_arg: self.space.topology.TopologyArg
+
+        return EvalArg
+
+    def _make_element_eval_arg(self):
         @cache.dynamic_struct(suffix=self.name)
         class ElementEvalArg:
             elt_arg: self.space.topology.ElementArg
@@ -70,7 +90,7 @@ class AdjointField(SpaceField):
             dof = self._get_dof(s)
             node_weight = self.space.element_inner_weight(
                 args.elt_arg,
-                args.eval_arg,
+                args.eval_arg.space_arg,
                 s.element_index,
                 s.element_coords,
                 get_node_index_in_element(dof),
@@ -91,7 +111,7 @@ class AdjointField(SpaceField):
             dof = self._get_dof(s)
             nabla_weight = self.space.element_inner_weight_gradient(
                 args.elt_arg,
-                args.eval_arg,
+                args.eval_arg.space_arg,
                 s.element_index,
                 s.element_coords,
                 get_node_index_in_element(dof),
@@ -113,7 +133,7 @@ class AdjointField(SpaceField):
             dof = self._get_dof(s)
             nabla_weight = self.space.element_inner_weight_gradient(
                 args.elt_arg,
-                args.eval_arg,
+                args.eval_arg.space_arg,
                 s.element_index,
                 s.element_coords,
                 get_node_index_in_element(dof),
@@ -132,7 +152,7 @@ class AdjointField(SpaceField):
             dof = self._get_dof(s)
             node_weight = self.space.element_outer_weight(
                 args.elt_arg,
-                args.eval_arg,
+                args.eval_arg.space_arg,
                 s.element_index,
                 s.element_coords,
                 get_node_index_in_element(dof),
@@ -153,7 +173,7 @@ class AdjointField(SpaceField):
             dof = self._get_dof(s)
             nabla_weight = self.space.element_outer_weight_gradient(
                 args.elt_arg,
-                args.eval_arg,
+                args.eval_arg.space_arg,
                 s.element_index,
                 s.element_coords,
                 get_node_index_in_element(dof),
@@ -175,7 +195,7 @@ class AdjointField(SpaceField):
             dof = self._get_dof(s)
             nabla_weight = self.space.element_outer_weight_gradient(
                 args.elt_arg,
-                args.eval_arg,
+                args.eval_arg.space_arg,
                 s.element_index,
                 s.element_coords,
                 get_node_index_in_element(dof),
@@ -193,12 +213,30 @@ class AdjointField(SpaceField):
         def at_node(args: self.ElementEvalArg, s: Sample):
             dof = self._get_dof(s)
             node_coords = self.space.node_coords_in_element(
-                args.elt_arg, args.eval_arg, s.element_index, get_node_index_in_element(dof)
+                args.elt_arg, args.eval_arg.space_arg, s.element_index, get_node_index_in_element(dof)
             )
             return Sample(s.element_index, node_coords, s.qp_index, s.qp_weight, s.test_dof, s.trial_dof)
 
         return at_node
 
+    def _make_node_index(self):
+        @cache.dynamic_func(suffix=self.name)
+        def node_index(args: self.ElementEvalArg, s: Sample):
+            dof = self._get_dof(s)
+            node_idx = self.space.topology.element_node_index(
+                args.elt_arg, args.eval_arg.topo_arg, s.element_index, get_node_index_in_element(dof)
+            )
+            return node_idx
+
+        return node_index
+
+    def _make_node_count(self):
+        @cache.dynamic_func(suffix=self.name)
+        def node_count(args: self.ElementEvalArg, s: Sample):
+            return self.space.topology.element_node_count(args.elt_arg, args.eval_arg.topo_arg, s.element_index)
+
+        return node_count
+
 
 class TestField(AdjointField):
     """Field defined over a space restriction that can be used as a test function.
@@ -269,16 +307,16 @@ class LocalAdjointField(SpaceField):
     OUTER_GRAD_DOF = wp.constant(3)
     DOF_TYPE_COUNT = wp.constant(4)
 
-    _OP_DOF_MAP_CONTINUOUS = {
+    _OP_DOF_MAP_CONTINUOUS: ClassVar[Dict[operator.Operator, int]] = {
         operator.inner: INNER_DOF,
         operator.outer: INNER_DOF,
         operator.grad: INNER_GRAD_DOF,
-        operator.grad_outer: INNER_GRAD_DOF,
+        operator.grad_outer: OUTER_GRAD_DOF,
         operator.div: INNER_GRAD_DOF,
-        operator.div_outer: INNER_GRAD_DOF,
+        operator.div_outer: OUTER_GRAD_DOF,
     }
 
-    _OP_DOF_MAP_DISCONTINUOUS = {
+    _OP_DOF_MAP_DISCONTINUOUS: ClassVar[Dict[operator.Operator, int]] = {
         operator.inner: INNER_DOF,
         operator.outer: OUTER_DOF,
         operator.grad: INNER_GRAD_DOF,
@@ -293,6 +331,18 @@ class LocalAdjointField(SpaceField):
     class EvalArg:
         pass
 
+    _dynamic_attribute_constructors: ClassVar = {
+        "ElementEvalArg": lambda obj: obj._make_element_eval_arg(),
+        "eval_degree": lambda obj: obj._make_eval_degree(),
+        "_split_dof": lambda obj: obj._make_split_dof(),
+        "eval_inner": lambda obj: obj._make_eval_inner(),
+        "eval_grad_inner": lambda obj: obj._make_eval_grad_inner(),
+        "eval_div_inner": lambda obj: obj._make_eval_div_inner(),
+        "eval_outer": lambda obj: obj._make_eval_outer(),
+        "eval_grad_outer": lambda obj: obj._make_eval_grad_outer(),
+        "eval_div_outer": lambda obj: obj._make_eval_div_outer(),
+    }
+
     def __init__(self, field: AdjointField):
         # if not isinstance(field.space, CollocatedFunctionSpace):
         #     raise NotImplementedError("Local assembly only implemented for collocated function spaces")
@@ -305,9 +355,6 @@ class LocalAdjointField(SpaceField):
         self.value_dof_count = self.space.VALUE_DOF_COUNT
 
         self._dof_suffix = ""
-
-        self.ElementEvalArg = self._make_element_eval_arg()
-        self.eval_degree = self._make_eval_degree()
         self.at_node = None
 
         self._is_discontinuous = (self.space.element_inner_weight != self.space.element_outer_weight) or (
@@ -318,6 +365,8 @@ class LocalAdjointField(SpaceField):
         self._TAYLOR_DOF_COUNTS = LocalAdjointField.DofOffsets(0)
         self.TAYLOR_DOF_COUNT = 0
 
+        cache.setup_dynamic_attributes(self)
+
     def notify_operator_usage(self, ops: Set[operator.Operator]):
         # Rebuild degrees-of-freedom offsets based on used operators
 
@@ -346,21 +395,7 @@ class LocalAdjointField(SpaceField):
         self._TAYLOR_DOF_COUNTS = dof_counts
 
         self._dof_suffix = "".join(str(c) for c in dof_counts)
-
-        self._split_dof = self._make_split_dof()
-
-        self.eval_inner = self._make_eval_inner()
-        self.eval_grad_inner = self._make_eval_grad_inner()
-        self.eval_div_inner = self._make_eval_div_inner()
-
-        if self._is_discontinuous:
-            self.eval_outer = self._make_eval_outer()
-            self.eval_grad_outer = self._make_eval_grad_outer()
-            self.eval_div_outer = self._make_eval_div_outer()
-        else:
-            self.eval_outer = self.eval_inner
-            self.eval_grad_outer = self.eval_grad_inner
-            self.eval_div_outer = self.eval_div_inner
+        cache.setup_dynamic_attributes(self)
 
     @property
     def name(self) -> str:
@@ -369,6 +404,9 @@ class LocalAdjointField(SpaceField):
     def eval_arg_value(self, device):
         return LocalAdjointField.EvalArg()
 
+    def fill_eval_arg(self, arg, device):
+        pass
+
     def _make_element_eval_arg(self):
         from warp.fem import cache
 
@@ -392,6 +430,7 @@ class LocalAdjointField(SpaceField):
 
     def _make_eval_inner(self):
         DOF_BEGIN = wp.constant(self._TAYLOR_DOF_OFFSETS[LocalAdjointField.INNER_DOF])
+        zero_element = type_zero_element(self.dtype)
 
         @cache.dynamic_func(suffix=self.name)
         def eval_test_inner(args: self.ElementEvalArg, s: Sample):
@@ -399,7 +438,7 @@ class LocalAdjointField(SpaceField):
 
             local_value_map = self.space.local_value_map_inner(args.elt_arg, s.element_index, s.element_coords)
             dof_value = self.space.value_basis_element(value_dof, local_value_map)
-            return wp.where(taylor_dof == 0, dof_value, self.dtype(0.0))
+            return wp.where(taylor_dof == 0, dof_value, zero_element())
 
         return eval_test_inner
 
@@ -409,13 +448,14 @@ class LocalAdjointField(SpaceField):
 
         DOF_BEGIN = wp.constant(self._TAYLOR_DOF_OFFSETS[LocalAdjointField.INNER_GRAD_DOF])
         DOF_COUNT = wp.constant(self._TAYLOR_DOF_COUNTS[LocalAdjointField.INNER_GRAD_DOF])
+        zero_element = type_zero_element(self.gradient_dtype)
 
         @cache.dynamic_func(suffix=self.name)
         def eval_nabla_test_inner(args: self.ElementEvalArg, s: Sample):
             value_dof, taylor_dof = self._split_dof(self._get_dof(s), DOF_BEGIN)
 
             if taylor_dof < 0 or taylor_dof >= DOF_COUNT:
-                return self.gradient_dtype(0.0)
+                return zero_element()
 
             grad_transform = self.space.element_inner_reference_gradient_transform(args.elt_arg, s)
             local_value_map = self.space.local_value_map_inner(args.elt_arg, s.element_index, s.element_coords)
@@ -430,13 +470,14 @@ class LocalAdjointField(SpaceField):
 
         DOF_BEGIN = wp.constant(self._TAYLOR_DOF_OFFSETS[LocalAdjointField.INNER_GRAD_DOF])
         DOF_COUNT = wp.constant(self._TAYLOR_DOF_COUNTS[LocalAdjointField.INNER_GRAD_DOF])
+        zero_element = type_zero_element(self.divergence_dtype)
 
         @cache.dynamic_func(suffix=self.name)
         def eval_div_test_inner(args: self.ElementEvalArg, s: Sample):
             value_dof, taylor_dof = self._split_dof(self._get_dof(s), DOF_BEGIN)
 
             if taylor_dof < 0 or taylor_dof >= DOF_COUNT:
-                return self.divergence_dtype(0.0)
+                return zero_element()
 
             grad_transform = self.space.element_inner_reference_gradient_transform(args.elt_arg, s)
             local_value_map = self.space.local_value_map_inner(args.elt_arg, s.element_index, s.element_coords)
@@ -446,7 +487,11 @@ class LocalAdjointField(SpaceField):
         return eval_div_test_inner
 
     def _make_eval_outer(self):
+        if not self._is_discontinuous:
+            return self.eval_inner
+
         DOF_BEGIN = wp.constant(self._TAYLOR_DOF_OFFSETS[LocalAdjointField.OUTER_DOF])
+        zero_element = type_zero_element(self.dtype)
 
         @cache.dynamic_func(suffix=self.name)
         def eval_test_outer(args: self.ElementEvalArg, s: Sample):
@@ -454,7 +499,7 @@ class LocalAdjointField(SpaceField):
 
             local_value_map = self.space.local_value_map_outer(args.elt_arg, s.element_index, s.element_coords)
             dof_value = self.space.value_basis_element(value_dof, local_value_map)
-            return wp.where(taylor_dof == 0, dof_value, self.dtype(0.0))
+            return wp.where(taylor_dof == 0, dof_value, zero_element())
 
         return eval_test_outer
 
@@ -464,13 +509,14 @@ class LocalAdjointField(SpaceField):
 
         DOF_BEGIN = wp.constant(self._TAYLOR_DOF_OFFSETS[LocalAdjointField.OUTER_GRAD_DOF])
         DOF_COUNT = wp.constant(self._TAYLOR_DOF_COUNTS[LocalAdjointField.OUTER_GRAD_DOF])
+        zero_element = type_zero_element(self.gradient_dtype)
 
         @cache.dynamic_func(suffix=self.name)
         def eval_nabla_test_outer(args: self.ElementEvalArg, s: Sample):
             value_dof, taylor_dof = self._split_dof(self._get_dof(s), DOF_BEGIN)
 
             if taylor_dof < 0 or taylor_dof >= DOF_COUNT:
-                return self.gradient_dtype(0.0)
+                return zero_element()
 
             grad_transform = self.space.element_outer_reference_gradient_transform(args.elt_arg, s)
             local_value_map = self.space.local_value_map_outer(args.elt_arg, s.element_index, s.element_coords)
@@ -485,13 +531,14 @@ class LocalAdjointField(SpaceField):
 
         DOF_BEGIN = wp.constant(self._TAYLOR_DOF_OFFSETS[LocalAdjointField.OUTER_GRAD_DOF])
         DOF_COUNT = wp.constant(self._TAYLOR_DOF_COUNTS[LocalAdjointField.OUTER_GRAD_DOF])
+        zero_element = type_zero_element(self.divergence_dtype)
 
         @cache.dynamic_func(suffix=self.name)
         def eval_div_test_outer(args: self.ElementEvalArg, s: Sample):
             value_dof, taylor_dof = self._split_dof(self._get_dof(s), DOF_BEGIN)
 
             if taylor_dof < 0 or taylor_dof >= DOF_COUNT:
-                return self.divergence_dtype(0.0)
+                return zero_element()
 
             grad_transform = self.space.element_outer_reference_gradient_transform(args.elt_arg, s)
             local_value_map = self.space.local_value_map_outer(args.elt_arg, s.element_index, s.element_coords)

warp/fem/geometry/adaptive_nanogrid.py

@@ -19,7 +19,7 @@ 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.fem.types import OUTSIDE, Coords, ElementIndex, Sample, make_free_sample
 
 from .element import Cube, Square
 from .geometry import Geometry
@@ -107,6 +107,13 @@ class AdaptiveNanogrid(Geometry):
         self._stacked_face_grid = None
         self._stacked_face_count = 0
 
+        transform = self.transform
+        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
@@ -167,17 +174,17 @@ class AdaptiveNanogrid(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
-        args.cell_level = self._cell_level
-
-        transform = self.transform
-        args.inverse_transform = wp.mat33f(np.linalg.inv(transform))
-        args.cell_volume = abs(np.linalg.det(transform))
-        args.level_count = self.level_count
-
+        self.fill_cell_arg(args, device)
         return args
 
+    def fill_cell_arg(self, arg: CellArg, device):
+        arg.cell_grid = self._cell_grid.id
+        arg.cell_ijk = self._cell_ijk
+        arg.cell_level = self._cell_level
+        arg.inverse_transform = self._inverse_transform
+        arg.cell_volume = self._cell_volume
+        arg.level_count = self.level_count
+
     @wp.func
     def _cell_scale(args: CellArg, cell_index: int):
         return float(1 << int(args.cell_level[cell_index]))
@@ -202,72 +209,41 @@ class AdaptiveNanogrid(Geometry):
         scale = AdaptiveNanogrid._cell_scale(args, s.element_index)
         return args.inverse_transform / scale
 
-    @wp.func
-    def _make_sample(args: CellArg, cell_index: int, uvw: wp.vec3):
-        ijk = args.cell_ijk[cell_index]
-        return make_free_sample(cell_index, (uvw - wp.vec3(ijk)) / AdaptiveNanogrid._cell_scale(args, cell_index))
+    def supports_cell_lookup(self, device):
+        return True
 
     @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 = AdaptiveNanogrid.find_cell(args.cell_grid, ijk, args.level_count, args.cell_level)
-
-        if cell_index == -1:
-            coords = uvw - wp.vec3(ijk)
-
-            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 = AdaptiveNanogrid.find_cell(args.cell_grid, nijk, args.level_count, args.cell_level)
-                        if cell_index != -1:
-                            return AdaptiveNanogrid._make_sample(args, cell_index, uvw)
-
-            return make_free_sample(NULL_ELEMENT_INDEX, Coords(OUTSIDE))
-
-        return AdaptiveNanogrid._make_sample(args, cell_index, uvw)
+    def _lookup_cell_index(args: AdaptiveNanogridCellArg, i: int, j: int, k: int):
+        return AdaptiveNanogrid.find_cell(args.cell_grid, wp.vec3i(i, j, k), args.level_count, args.cell_level)
 
     @wp.func
-    def cell_lookup(args: CellArg, pos: wp.vec3, guess: Sample):
-        s_global = AdaptiveNanogrid.cell_lookup(args, pos)
+    def _cell_coordinates_local(args: AdaptiveNanogridCellArg, cell_index: int, uvw: wp.vec3):
+        ijk = wp.vec3(args.cell_ijk[cell_index])
+        rel_pos = uvw - ijk
+        scale = AdaptiveNanogrid._cell_scale(args, cell_index)
+        return rel_pos / scale
 
-        if s_global.element_index != NULL_ELEMENT_INDEX:
-            return s_global
+    @wp.func
+    def _cell_closest_point_local(args: AdaptiveNanogridCellArg, cell_index: int, uvw: wp.vec3):
+        ijk = wp.vec3(args.cell_ijk[cell_index])
+        rel_pos = uvw - ijk
+        scale = AdaptiveNanogrid._cell_scale(args, cell_index)
+        coords = wp.min(wp.max(rel_pos / scale, wp.vec3(0.0)), wp.vec3(1.0))
+        return wp.length_sq(wp.volume_index_to_world_dir(args.cell_grid, coords * scale - rel_pos)), coords
 
-        closest_voxel = int(NULL_ELEMENT_INDEX)
-        closest_coords = Coords(OUTSIDE)
-        closest_dist = float(1.0e8)
+    @wp.func
+    def cell_coordinates(args: AdaptiveNanogridCellArg, cell_index: int, pos: wp.vec3):
+        uvw = wp.volume_world_to_index(args.cell_grid, pos) + wp.vec3(0.5)
+        return AdaptiveNanogrid._cell_coordinates_local(args, cell_index, uvw)
 
-        # project to closest in stencil
+    @wp.func
+    def cell_closest_point(args: AdaptiveNanogridCellArg, 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 = AdaptiveNanogrid.find_cell(args.cell_grid, nijk, args.level_count, args.cell_level)
-                    if cell_idx != -1:
-                        nijk = args.cell_ijk[cell_idx]
-                        scale = AdaptiveNanogrid._cell_scale(args, cell_idx)
-                        coords = (uvw - wp.vec3(nijk)) / scale
-                        dist, proj_coords = Nanogrid._project_on_voxel_at_origin(coords)
-                        dist *= scale
-                        if dist <= closest_dist:
-                            closest_dist = dist
-                            closest_voxel = cell_idx
-                            closest_coords = proj_coords
-
-        return make_free_sample(closest_voxel, closest_coords)
+        dist, coords = AdaptiveNanogrid._cell_closest_point_local(args, cell_index, uvw)
+        return coords, dist
+
+    def make_filtered_cell_lookup(self, filter_func: wp.Function = None):
+        return Nanogrid._make_filtered_cell_lookup(self, filter_func)
 
     @wp.func
     def cell_measure(args: CellArg, s: Sample):
@@ -280,6 +256,7 @@ class AdaptiveNanogrid(Geometry):
 
     SideIndexArg = Nanogrid.SideIndexArg
     side_index_arg_value = Nanogrid.side_index_arg_value
+    fill_side_index_arg = Nanogrid.fill_side_index_arg
 
     SideArg = AdaptiveNanogridSideArg
 
@@ -289,20 +266,20 @@ class AdaptiveNanogrid(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)
-        args.face_cell_indices = self._face_cell_indices.to(device)
-        transform = self.transform
-        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_cell_indices = self._face_cell_indices.to(device)
+        arg.face_areas = self._face_areas
+
     @wp.func
     def boundary_side_index(args: SideIndexArg, boundary_side_index: int):
         return args.boundary_face_indices[boundary_side_index]
@@ -320,9 +297,10 @@ class AdaptiveNanogrid(Geometry):
         ijk = args.face_ijk[s.element_index]
         flags = args.face_flags[s.element_index]
         axis = Nanogrid._get_face_axis(flags)
+        flip = Nanogrid._get_face_inner_offset(flags)
         scale = AdaptiveNanogrid._get_face_scale(flags)
 
-        uvw = wp.vec3(ijk) + scale * Nanogrid._side_to_cell_coords(axis, 0.0, s.element_coords)
+        uvw = wp.vec3(ijk) + scale * 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))
@@ -391,9 +369,10 @@ class AdaptiveNanogrid(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 = Nanogrid._get_face_inner_offset(flags)
 
-        same_level_cell_coords = Nanogrid._side_to_cell_coords(axis, 1.0 - float(offset), side_coords)
+        same_level_cell_coords = Nanogrid._side_to_cell_coords(axis, flip, 1.0 - float(offset), side_coords)
         same_level_cell_ijk = args.face_ijk[side_index]
         side_level = AdaptiveNanogrid._get_face_level(flags)
         same_level_cell_ijk[axis] += (offset - 1) << side_level
@@ -410,9 +389,10 @@ class AdaptiveNanogrid(Geometry):
     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 = Nanogrid._get_face_outer_offset(flags)
 
-        same_level_cell_coords = Nanogrid._side_to_cell_coords(axis, float(offset), side_coords)
+        same_level_cell_coords = Nanogrid._side_to_cell_coords(axis, flip, float(offset), side_coords)
         same_level_cell_ijk = args.face_ijk[side_index]
         side_level = AdaptiveNanogrid._get_face_level(flags)
         same_level_cell_ijk[axis] -= offset << side_level
@@ -434,6 +414,7 @@ class AdaptiveNanogrid(Geometry):
     ):
         flags = args.face_flags[side_index]
         axis = Nanogrid._get_face_axis(flags)
+        flip = Nanogrid._get_face_inner_offset(flags)
         side_level = AdaptiveNanogrid._get_face_level(flags)
         cell_level = int(args.cell_arg.cell_level[element_index])
 
@@ -450,11 +431,35 @@ class AdaptiveNanogrid(Geometry):
             and wp.max(same_level_cell_coords) <= 1.0
         )
 
-        return wp.where(
-            on_side,
-            Coords(same_level_cell_coords[(axis + 1) % 3], same_level_cell_coords[(axis + 2) % 3], 0.0),
-            Coords(OUTSIDE),
-        )
+        return wp.where(on_side, Nanogrid._cell_to_side_coords(axis, flip, same_level_cell_coords), Coords(OUTSIDE))
+
+    @wp.func
+    def side_coordinates(args: SideArg, side_index: int, pos: wp.vec3):
+        cell_arg = args.cell_arg
+
+        ijk = args.face_ijk[side_index]
+        fine_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]
+        side_level = AdaptiveNanogrid._get_face_level(flags)
+        axis = Nanogrid._get_face_axis(flags)
+        flip = Nanogrid._get_face_inner_offset(flags)
+
+        return Nanogrid._cell_to_side_coords(axis, flip, fine_cell_coords / float(1 << side_level))
+
+    @wp.func
+    def side_closest_point(args: SideArg, side_index: int, pos: wp.vec3):
+        coords = AdaptiveNanogrid.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)
+        side_level = AdaptiveNanogrid._get_face_level(flags)
+        cell_coord_offset = Nanogrid._side_to_cell_coords(axis, flip, 0, coords - proj_coords) * float(1 << side_level)
+
+        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

warp/fem/geometry/closest_point.py

@@ -82,3 +82,16 @@ def project_on_tet_at_origin(q: wp.vec3, e1: wp.vec3, e2: wp.vec3, e3: wp.vec3):
         return dmin, Coords(s31[2], 0.0, s31[1])
     else:
         return dmin, s123
+
+
+@wp.func
+def project_on_box_at_origin(coords: wp.vec3, sizes: wp.vec3):
+    proj_coords = wp.min(wp.max(coords, wp.vec3(0.0)), sizes)
+    return wp.length_sq(coords - proj_coords), wp.cw_div(proj_coords, sizes)
+
+
+@wp.func
+def project_on_box_at_origin(coords: wp.vec2, sizes: wp.vec2):
+    proj_coords = wp.min(wp.max(coords, wp.vec2(0.0)), sizes)
+    norm_coords = wp.cw_div(proj_coords, sizes)
+    return wp.length_sq(coords - proj_coords), Coords(norm_coords[0], norm_coords[1], 0.0)