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

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)

warp/fem/geometry/deformed_geometry.py

@@ -13,9 +13,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import ClassVar
+
 import warp as wp
 from warp.fem import cache
-from warp.fem.types import Coords, ElementIndex, Sample
+from warp.fem.polynomial import Polynomial
+from warp.fem.types import Coords, ElementIndex, Sample, make_free_sample
 
 from .geometry import Geometry
 
@@ -23,7 +26,31 @@ _mat32 = wp.mat(shape=(3, 2), dtype=float)
 
 
 class DeformedGeometry(Geometry):
-    def __init__(self, field: "wp.fem.field.GeometryField", relative: bool = True):
+    _dynamic_attribute_constructors_phase_1: ClassVar = {
+        "CellArg": lambda obj: obj._make_cell_arg(),
+        "SideArg": lambda obj: obj._make_side_arg(),
+        "cell_position": lambda obj: obj._make_cell_position(),
+        "cell_deformation_gradient": lambda obj: obj._make_cell_deformation_gradient(),
+        "side_to_cell_arg": lambda obj: obj._make_side_to_cell_arg(),
+        "side_position": lambda obj: obj._make_side_position(),
+        "side_deformation_gradient": lambda obj: obj._make_side_deformation_gradient(),
+        "side_inner_cell_index": lambda obj: obj._make_side_inner_cell_index(),
+        "side_outer_cell_index": lambda obj: obj._make_side_outer_cell_index(),
+        "side_inner_cell_coords": lambda obj: obj._make_side_inner_cell_coords(),
+        "side_outer_cell_coords": lambda obj: obj._make_side_outer_cell_coords(),
+        "side_from_cell_coords": lambda obj: obj._make_side_from_cell_coords(),
+        "cell_bvh_id": lambda obj: obj._make_cell_bvh_id(),
+        "cell_bounds": lambda obj: obj._make_cell_bounds(),
+    }
+
+    _dynamic_attribute_constructors_phase_2: ClassVar = {
+        "cell_closest_point": lambda obj: obj._make_cell_closest_point(),
+        "side_closest_point": lambda obj: obj._make_side_closest_point(),
+        "cell_coordinates": lambda obj: obj._make_cell_coordinates(),
+        "side_coordinates": lambda obj: obj._make_side_coordinates(),
+    }
+
+    def __init__(self, field: "wp.fem.field.GeometryField", relative: bool = True, build_bvh: bool = False):
         """Constructs a Deformed Geometry from a displacement or absolute position field defined over a base geometry.
         The deformation field does not need to be isoparameteric.
 
@@ -33,10 +60,7 @@ class DeformedGeometry(Geometry):
         from warp.fem.field import DiscreteField, GeometryField
 
         if isinstance(field, DiscreteField):
-            if (
-                not wp.types.type_is_vector(field.dtype)
-                or wp.types.type_length(field.dtype) != field.geometry.dimension
-            ):
+            if not wp.types.type_is_vector(field.dtype) or wp.types.type_size(field.dtype) != field.geometry.dimension:
                 raise ValueError(
                     "Invalid value type for position field, must be vector-valued with same dimension as underlying geometry"
                 )
@@ -46,14 +70,10 @@ class DeformedGeometry(Geometry):
         self._relative = relative
 
         self.field: GeometryField = field
+        self.field_trace = field.trace()
         self.dimension = self.base.dimension
 
-        self.CellArg = self.field.ElementEvalArg
-
-        self.field_trace = field.trace()
-        self.SideArg = self._make_side_arg()
         self.SideIndexArg = self.base.SideIndexArg
-
         self.cell_count = self.base.cell_count
         self.vertex_count = self.base.vertex_count
         self.side_count = self.base.side_count
@@ -62,23 +82,18 @@ class DeformedGeometry(Geometry):
         self.reference_side = self.base.reference_side
 
         self.side_index_arg_value = self.base.side_index_arg_value
-
-        self.cell_position = self._make_cell_position()
-        self.cell_deformation_gradient = self._make_cell_deformation_gradient()
-
+        self.fill_side_index_arg = self.base.fill_side_index_arg
         self.boundary_side_index = self.base.boundary_side_index
 
-        self.side_to_cell_arg = self._make_side_to_cell_arg()
-        self.side_position = self._make_side_position()
-        self.side_deformation_gradient = self._make_side_deformation_gradient()
-        self.side_inner_cell_index = self._make_side_inner_cell_index()
-        self.side_outer_cell_index = self._make_side_outer_cell_index()
-        self.side_inner_cell_coords = self._make_side_inner_cell_coords()
-        self.side_outer_cell_coords = self._make_side_outer_cell_coords()
-        self.side_from_cell_coords = self._make_side_from_cell_coords()
+        cache.setup_dynamic_attributes(self, constructors=self._dynamic_attribute_constructors_phase_1)
 
         self._make_default_dependent_implementations()
 
+        cache.setup_dynamic_attributes(self, constructors=self._dynamic_attribute_constructors_phase_2)
+
+        if build_bvh:
+            self.build_bvh(self.field.dof_values.device)
+
     @property
     def name(self) -> str:
         return f"DefGeo_{self.field.name}_{'rel' if self._relative else 'abs'}"
@@ -89,35 +104,49 @@ class DeformedGeometry(Geometry):
 
     # Geometry device interface
 
-    @cache.cached_arg_value
-    def cell_arg_value(self, device) -> "DeformedGeometry.CellArg":
-        args = self.CellArg()
+    def _make_cell_arg(self):
+        @cache.dynamic_struct(suffix=self.name)
+        class CellArg:
+            base_arg: self.base.CellArg
+            field_arg: self.field.EvalArg
+            cell_bvh: wp.uint64
 
-        args.elt_arg = self.base.cell_arg_value(device)
-        args.eval_arg = self.field.eval_arg_value(device)
+        return CellArg
 
+    def cell_arg_value(self, device) -> "DeformedGeometry.CellArg":
+        args = self.CellArg()
+        self.fill_cell_arg(args, device)
         return args
 
+    def fill_cell_arg(self, args: "DeformedGeometry.CellArg", device):
+        self.base.fill_cell_arg(args.base_arg, device)
+        self.field.fill_eval_arg(args.field_arg, device)
+        args.cell_bvh = self.bvh_id(device)
+
     def _make_cell_position(self):
         @cache.dynamic_func(suffix=self.name)
         def cell_position_absolute(cell_arg: self.CellArg, s: Sample):
-            return self.field.eval_inner(cell_arg, s)
+            field_arg = self.field.ElementEvalArg(cell_arg.base_arg, cell_arg.field_arg)
+            return self.field.eval_inner(field_arg, s)
 
         @cache.dynamic_func(suffix=self.name)
         def cell_position(cell_arg: self.CellArg, s: Sample):
-            return self.field.eval_inner(cell_arg, s) + self.base.cell_position(cell_arg.elt_arg, s)
+            field_arg = self.field.ElementEvalArg(cell_arg.base_arg, cell_arg.field_arg)
+            return self.field.eval_inner(field_arg, s) + self.base.cell_position(cell_arg.base_arg, s)
 
         return cell_position if self._relative else cell_position_absolute
 
     def _make_cell_deformation_gradient(self):
         @cache.dynamic_func(suffix=self.name)
         def cell_deformation_gradient_absolute(cell_arg: self.CellArg, s: Sample):
-            return self.field.eval_reference_grad_inner(cell_arg, s)
+            field_arg = self.field.ElementEvalArg(cell_arg.base_arg, cell_arg.field_arg)
+            return self.field.eval_reference_grad_inner(field_arg, s)
 
         @cache.dynamic_func(suffix=self.name)
         def cell_deformation_gradient(cell_arg: self.CellArg, s: Sample):
-            return self.field.eval_reference_grad_inner(cell_arg, s) + self.base.cell_deformation_gradient(
-                cell_arg.elt_arg, s
+            field_arg = self.field.ElementEvalArg(cell_arg.base_arg, cell_arg.field_arg)
+            return self.field.eval_reference_grad_inner(field_arg, s) + self.base.cell_deformation_gradient(
+                cell_arg.base_arg, s
             )
 
         return cell_deformation_gradient if self._relative else cell_deformation_gradient_absolute
@@ -128,19 +157,21 @@ class DeformedGeometry(Geometry):
             base_arg: self.base.SideArg
             trace_arg: self.field_trace.EvalArg
             field_arg: self.field.EvalArg
+            cell_bvh: wp.uint64
 
         return SideArg
 
-    @cache.cached_arg_value
     def side_arg_value(self, device) -> "DeformedGeometry.SideArg":
         args = self.SideArg()
-
-        args.base_arg = self.base.side_arg_value(device)
-        args.field_arg = self.field.eval_arg_value(device)
-        args.trace_arg = self.field_trace.eval_arg_value(device)
-
+        self.fill_side_arg(args, device)
         return args
 
+    def fill_side_arg(self, args: "DeformedGeometry.SideArg", device):
+        self.base.fill_side_arg(args.base_arg, device)
+        self.field.fill_eval_arg(args.field_arg, device)
+        self.field_trace.fill_eval_arg(args.trace_arg, device)
+        args.cell_bvh = self.bvh_id(device)
+
     def _make_side_position(self):
         @cache.dynamic_func(suffix=self.name)
         def side_position_absolute(args: self.SideArg, s: Sample):
@@ -216,6 +247,37 @@ class DeformedGeometry(Geometry):
     def _make_side_to_cell_arg(self):
         @cache.dynamic_func(suffix=self.name)
         def side_to_cell_arg(side_arg: self.SideArg):
-            return self.CellArg(self.base.side_to_cell_arg(side_arg.base_arg), side_arg.field_arg)
+            return self.CellArg(self.base.side_to_cell_arg(side_arg.base_arg), side_arg.field_arg, side_arg.cell_bvh)
 
         return side_to_cell_arg
+
+    def _make_cell_bvh_id(self):
+        @cache.dynamic_func(suffix=self.name)
+        def cell_bvh_id(cell_arg: self.CellArg):
+            return cell_arg.cell_bvh
+
+        return cell_bvh_id
+
+    def _make_cell_bounds(self):
+        points, _weights = self.reference_cell().instantiate_quadrature(
+            order=self.field.degree, family=Polynomial.LOBATTO_GAUSS_LEGENDRE
+        )
+
+        points = cache.cached_mat_type((len(points), 3), dtype=float)(points)
+        point_count = len(points)
+
+        @cache.dynamic_func(suffix=self.name)
+        def cell_bounds(cell_arg: self.CellArg, cell_index: ElementIndex):
+            lower = wp.vec3(1.0e8)
+            upper = wp.vec3(-1.0e8)
+            for k in range(point_count):
+                pos = self.cell_position(cell_arg, make_free_sample(cell_index, points[k]))
+                lower = wp.min(lower, pos)
+                upper = wp.max(upper, pos)
+
+            # pad the BBox to account for potential overflows
+            pad = 0.25 * (upper - lower)
+
+            return lower - pad, upper + pad
+
+        return cell_bounds

warp/fem/geometry/element.py

@@ -15,14 +15,18 @@
 
 from typing import List, Tuple
 
+import warp as wp
 from warp.fem.polynomial import Polynomial, quadrature_1d
 from warp.fem.types import Coords
 
+_vec1 = wp.types.vector(length=1, dtype=float)
+
 
 class Element:
     dimension = 0
     """Intrinsic dimension of the element"""
 
+    @staticmethod
     def measure() -> float:
         """Measure (area, volume, ...) of the reference element"""
         raise NotImplementedError
@@ -32,10 +36,32 @@ class Element:
         """Returns a quadrature of a given order for a prototypical element"""
         raise NotImplementedError
 
-    def center(self) -> Tuple[float]:
-        coords, _ = self.instantiate_quadrature(order=0, family=None)
+    @classmethod
+    def center(cls) -> Coords:
+        """Returns the coordinates for the center of the element"""
+        coords, _ = cls.instantiate_quadrature(order=0, family=None)
         return coords[0]
 
+    @wp.func
+    def project(v: Coords):
+        """project coordinates so that they belong to the element"""
+        return wp.min(wp.max(v, Coords(0.0)), Coords(1.0))
+
+    @wp.func
+    def coord_delta(ref_delta: wp.vec3):
+        """Transform a delta in reference space to element coords"""
+        return ref_delta
+
+    @wp.func
+    def coord_delta(ref_delta: wp.vec2):
+        """Transform a delta in reference space to element coords"""
+        return Coords(ref_delta[0], ref_delta[1], 0.0)
+
+    @wp.func
+    def coord_delta(ref_delta: _vec1):
+        """Transform a delta in reference space to element coords"""
+        return Coords(ref_delta[0], 0.0, 0.0)
+
 
 def _point_count_from_order(order: int, family: Polynomial):
     if family == Polynomial.GAUSS_LEGENDRE:
@@ -454,6 +480,17 @@ class Triangle(Element):
 
         return coords, weights
 
+    @wp.func
+    def project(v: Coords):
+        n = wp.max(0.0, (v[0] + v[1] - 1.0) * 0.5)
+        a = wp.clamp(v[0] - n, 0.0, 1.0)
+        b = wp.clamp(v[1] - n, 0.0, 1.0)
+        return Coords(a, b, 1.0 - (a + b))
+
+    @wp.func
+    def coord_delta(ref_delta: wp.vec2):
+        return Coords(-ref_delta[0] - ref_delta[1], ref_delta[0], ref_delta[1])
+
 
 class Tetrahedron(Element):
     dimension = 3
@@ -774,3 +811,20 @@ class Tetrahedron(Element):
             raise NotImplementedError
 
         return coords, weights
+
+    @wp.func
+    def project(v: Coords):
+        # project on 1-2-3 half-space
+        n = wp.max(0.0, (v[0] + v[1] + v[2] - 1.0) / 3.0)
+        c = v - Coords(n)
+
+        # project on 1-2, 2-3, 3-1 half-spaces
+        n = wp.max(0.0, (c[0] + c[1] - 1.0) * 0.5)
+        c = c - Coords(n, n, 0.0)
+        n = wp.max(0.0, (c[1] + c[2] - 1.0) * 0.5)
+        c = c - Coords(0.0, n, n)
+        n = wp.max(0.0, (c[2] + c[0] - 1.0) * 0.5)
+        c = c - Coords(n, 0.0, n)
+
+        # project on cube
+        return Element.project(c)