warp-lang 1.2.1__py3-none-manylinux2014_x86_64.whl → 1.3.0__py3-none-manylinux2014_x86_64.whl

warp/fem/geometry/deformed_geometry.py

@@ -10,18 +10,32 @@ _mat32 = wp.mat(shape=(3, 2), dtype=float)
 
 
 class DeformedGeometry(Geometry):
-    def __init__(self, field):
-        """Constructs a Deformed Geometry from a displacement field defined over a base geometry"""
-
-        from warp.fem.field import DiscreteField
-
-        self.field: DiscreteField = field
-        self.base = self.field.space.geometry
+    def __init__(self, field: "wp.fem.field.GeometryField", relative: bool = True):
+        """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.
+
+        See also: :meth:`warp.fem.DiscreteField.make_deformed_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
+            ):
+                raise ValueError(
+                    "Invalid value type for position field, must be vector-valued with same dimension as underlying geometry"
+                )
+        if field.eval_grad_inner is None:
+            raise ValueError("Gradient evaluation is not supported on the passed field")
+
+        self._relative = relative
+
+        self.field: GeometryField = field
+        self.base = self.field.geometry
         self.dimension = self.base.dimension
 
-        if not wp.types.type_is_vector(field.dtype) or wp.types.type_length(field.dtype) != self.dimension:
-            raise ValueError("Invalid value type for position field")
-
         self.CellArg = self.field.ElementEvalArg
 
         self.field_trace = field.trace()
@@ -60,7 +74,7 @@ class DeformedGeometry(Geometry):
 
     @property
     def name(self):
-        return f"DefGeo_{self.field.name}"
+        return f"DefGeo_{self.field.name}_{'rel' if self._relative else 'abs'}"
 
     # Geometry device interface
 
@@ -74,20 +88,28 @@ class DeformedGeometry(Geometry):
         return args
 
     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)
+
         @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)
 
-        return cell_position
+        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)
+
         @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
             )
 
-        return cell_deformation_gradient
+        return cell_deformation_gradient if self._relative else cell_deformation_gradient_absolute
 
     def _make_cell_inverse_deformation_gradient(self):
         @cache.dynamic_func(suffix=self.name)
@@ -129,14 +151,27 @@ class DeformedGeometry(Geometry):
         return args
 
     def _make_side_position(self):
+        @cache.dynamic_func(suffix=self.name)
+        def side_position_absolute(args: self.SideArg, s: Sample):
+            trace_arg = self.field_trace.ElementEvalArg(args.base_arg, args.trace_arg)
+            return self.field_trace.eval_inner(trace_arg, s)
+
         @cache.dynamic_func(suffix=self.name)
         def side_position(args: self.SideArg, s: Sample):
             trace_arg = self.field_trace.ElementEvalArg(args.base_arg, args.trace_arg)
             return self.field_trace.eval_inner(trace_arg, s) + self.base.side_position(args.base_arg, s)
 
-        return side_position
+        return side_position if self._relative else side_position_absolute
 
     def _make_side_deformation_gradient(self):
+        @cache.dynamic_func(suffix=self.name)
+        def side_deformation_gradient_absolute(args: self.SideArg, s: Sample):
+            base_def_grad = self.base.side_deformation_gradient(args.base_arg, s)
+            trace_arg = self.field_trace.ElementEvalArg(args.base_arg, args.trace_arg)
+
+            Du = self.field_trace.eval_grad_inner(trace_arg, s)
+            return Du * base_def_grad
+
         @cache.dynamic_func(suffix=self.name)
         def side_deformation_gradient(args: self.SideArg, s: Sample):
             base_def_grad = self.base.side_deformation_gradient(args.base_arg, s)
@@ -145,7 +180,7 @@ class DeformedGeometry(Geometry):
             Du = self.field_trace.eval_grad_inner(trace_arg, s)
             return base_def_grad + Du * base_def_grad
 
-        return side_deformation_gradient
+        return side_deformation_gradient if self._relative else side_deformation_gradient_absolute
 
     def _make_side_inner_inverse_deformation_gradient(self):
         @cache.dynamic_func(suffix=self.name)

warp/fem/geometry/hexmesh.py

@@ -65,7 +65,7 @@ EDGE_VERTEX_INDICES = wp.constant(
     )
 )
 
-# orthogal transform for face coordinates given first vertex + winding
+# orthogonal transform for face coordinates given first vertex + winding
 # (two rows per entry)
 
 FACE_ORIENTATION = [
@@ -439,12 +439,12 @@ class Hexmesh(Geometry):
         return side_arg.cell_arg
 
     def _build_topology(self, temporary_store: TemporaryStore):
-        from warp.fem.utils import compress_node_indices, masked_indices
+        from warp.fem.utils import compress_node_indices, host_read_at_index, masked_indices
         from warp.utils import array_scan
 
         device = self.hex_vertex_indices.device
 
-        vertex_hex_offsets, vertex_hex_indices, _, __ = compress_node_indices(
+        vertex_hex_offsets, vertex_hex_indices = compress_node_indices(
             self.vertex_count(), self.hex_vertex_indices, temporary_store=temporary_store
         )
         self._vertex_hex_offsets = vertex_hex_offsets.detach()
@@ -492,16 +492,11 @@ class Hexmesh(Geometry):
         array_scan(in_array=vertex_start_face_count.array, out_array=vertex_unique_face_offsets.array, inclusive=False)
 
         # Get back edge count to host
-        if device.is_cuda:
-            face_count = borrow_temporary(temporary_store, shape=(1,), dtype=int, device="cpu", pinned=True)
-            # Last vertex will not own any edge, so its count will be zero; just fetching last prefix count is ok
-            wp.copy(
-                dest=face_count.array, src=vertex_unique_face_offsets.array, src_offset=self.vertex_count() - 1, count=1
+        face_count = int(
+            host_read_at_index(
+                vertex_unique_face_offsets.array, self.vertex_count() - 1, temporary_store=temporary_store
             )
-            wp.synchronize_stream(wp.get_stream(device))
-            face_count = int(face_count.array.numpy()[0])
-        else:
-            face_count = int(vertex_unique_face_offsets.array.numpy()[self.vertex_count() - 1])
+        )
 
         self._face_vertex_indices = wp.empty(shape=(face_count,), dtype=wp.vec4i, device=device)
         self._face_hex_indices = wp.empty(shape=(face_count,), dtype=wp.vec2i, device=device)
@@ -557,6 +552,7 @@ class Hexmesh(Geometry):
         self._boundary_face_indices = boundary_face_indices.detach()
 
     def _compute_hex_edges(self, temporary_store: Optional[TemporaryStore] = None):
+        from warp.fem.utils import host_read_at_index
         from warp.utils import array_scan
 
         device = self.hex_vertex_indices.device
@@ -599,19 +595,11 @@ class Hexmesh(Geometry):
         array_scan(in_array=vertex_start_edge_count.array, out_array=vertex_unique_edge_offsets.array, inclusive=False)
 
         # Get back edge count to host
-        if device.is_cuda:
-            edge_count = borrow_temporary(temporary_store, shape=(1,), dtype=int, device="cpu", pinned=True)
-            # Last vertex will not own any edge, so its count will be zero; just fetching last prefix count is ok
-            wp.copy(
-                dest=edge_count.array,
-                src=vertex_unique_edge_offsets.array,
-                src_offset=self.vertex_count() - 1,
-                count=1,
+        self._edge_count = int(
+            host_read_at_index(
+                vertex_unique_edge_offsets.array, self.vertex_count() - 1, temporary_store=temporary_store
             )
-            wp.synchronize_stream(wp.get_stream(device))
-            self._edge_count = int(edge_count.array.numpy()[0])
-        else:
-            self._edge_count = int(vertex_unique_edge_offsets.array.numpy()[self.vertex_count() - 1])
+        )
 
         self._hex_edge_indices = wp.empty(
             dtype=int, device=self.hex_vertex_indices.device, shape=(self.cell_count(), 12)

warp/fem/geometry/nanogrid.py

@@ -60,9 +60,21 @@ class NanogridSideArg:
 
 
 class Nanogrid(Geometry):
+    """Sparse grid geometry"""
+
     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.
+
+        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()
 
@@ -77,31 +89,29 @@ class Nanogrid(Geometry):
         self._node_ijk = wp.array(shape=(node_count,), dtype=wp.vec3i, device=device)
         self._node_grid.get_voxels(out=self._node_ijk)
 
-        self._face_grid = _build_face_grid(self._cell_ijk, grid, temporary_store)
-        face_count = self._face_grid.get_voxel_count()
-        self._face_ijk = wp.array(shape=(face_count,), dtype=wp.vec3i, device=device)
-        self._face_grid.get_voxels(out=self._face_ijk)
-
-        self._face_flags = wp.array(shape=(face_count,), dtype=wp.uint8, device=device)
-        boundary_face_mask = cache.borrow_temporary(temporary_store, shape=(face_count,), dtype=wp.int32, device=device)
-
-        wp.launch(
-            _build_face_flags,
-            dim=face_count,
-            device=device,
-            inputs=[grid.id, self._face_ijk, self._face_flags, boundary_face_mask.array],
-        )
-        boundary_face_indices, _ = utils.masked_indices(boundary_face_mask.array)
-        self._boundary_face_indices = boundary_face_indices.detach()
+        self._face_grid = None
+        self._face_ijk = None
 
         self._edge_grid = None
         self._edge_ijk = None
 
-    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)
-        edge_count = self._edge_grid.get_voxel_count()
-        self._edge_ijk = wp.array(shape=(edge_count,), dtype=wp.vec3i, device=self._edge_grid.device)
-        self._edge_grid.get_voxels(out=self._edge_ijk)
+    @property
+    def cell_grid(self) -> wp.Volume:
+        return self._cell_grid
+
+    @property
+    def vertex_grid(self) -> wp.Volume:
+        return self._node_grid
+
+    @property
+    def face_grid(self) -> wp.Volume:
+        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]
@@ -110,17 +120,17 @@ class Nanogrid(Geometry):
         return self._node_ijk.shape[0]
 
     def side_count(self):
+        self._ensure_face_grid()
         return self._face_ijk.shape[0]
 
-    def edge_count(self):
-        if self._edge_ijk is None:
-            self._build_edge_grid()
-
-        return self._edge_ijk.shape[0]
-
     def boundary_side_count(self):
+        self._ensure_face_grid()
         return self._boundary_face_indices.shape[0]
 
+    def edge_count(self):
+        self._ensure_edge_grid()
+        return self._edge_ijk.shape[0]
+
     def reference_cell(self) -> Cube:
         return Cube()
 
@@ -144,7 +154,7 @@ class Nanogrid(Geometry):
     @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)
+        return wp.volume_index_to_world(args.cell_grid, uvw - wp.vec3(0.5))
 
     @wp.func
     def cell_deformation_gradient(args: CellArg, s: Sample):
@@ -156,7 +166,7 @@ class Nanogrid(Geometry):
 
     @wp.func
     def cell_lookup(args: CellArg, pos: wp.vec3):
-        uvw = wp.volume_world_to_index(args.cell_grid, pos)
+        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])))
         element_index = wp.volume_lookup_index(args.cell_grid, ijk[0], ijk[1], ijk[2])
 
@@ -166,9 +176,37 @@ class Nanogrid(Geometry):
             make_free_sample(NULL_ELEMENT_INDEX, Coords(OUTSIDE)),
         )
 
+    @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
+
     @wp.func
     def cell_lookup(args: CellArg, pos: wp.vec3, guess: Sample):
-        return Nanogrid.cell_lookup(args, pos)
+        s_global = Nanogrid.cell_lookup(args, pos)
+
+        if s_global.element_index != NULL_ELEMENT_INDEX:
+            return s_global
+
+        closest_voxel = int(NULL_ELEMENT_INDEX)
+        closest_coords = Coords(OUTSIDE)
+        closest_dist = float(1.0e8)
+
+        # project to closest in stencil
+        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 make_free_sample(closest_voxel, closest_coords)
 
     @wp.func
     def cell_measure(args: CellArg, s: Sample):
@@ -182,6 +220,8 @@ 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)
@@ -199,6 +239,8 @@ class Nanogrid(Geometry):
 
     @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)
         return args
@@ -235,7 +277,7 @@ class Nanogrid(Geometry):
         uvw = wp.vec3(ijk) + Nanogrid._side_to_cell_coords(axis, 0.0, s.element_coords)
 
         cell_grid = args.cell_arg.cell_grid
-        return wp.volume_index_to_world(cell_grid, uvw)
+        return wp.volume_index_to_world(cell_grid, uvw - wp.vec3(0.5))
 
     @wp.func
     def _face_tangent_vecs(args: SideArg, axis: int, flip: int):
@@ -342,6 +384,39 @@ class Nanogrid(Geometry):
     def side_to_cell_arg(side_arg: SideArg):
         return side_arg.cell_arg
 
+    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)
+        face_count = self._face_grid.get_voxel_count()
+        self._face_ijk = wp.array(shape=(face_count,), dtype=wp.vec3i, device=device)
+        self._face_grid.get_voxels(out=self._face_ijk)
+
+        self._face_flags = wp.array(shape=(face_count,), dtype=wp.uint8, device=device)
+        boundary_face_mask = cache.borrow_temporary(temporary_store, shape=(face_count,), dtype=wp.int32, device=device)
+
+        wp.launch(
+            _build_face_flags,
+            dim=face_count,
+            device=device,
+            inputs=[self._cell_grid.id, self._face_ijk, self._face_flags, boundary_face_mask.array],
+        )
+        boundary_face_indices, _ = utils.masked_indices(boundary_face_mask.array)
+        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)
+        edge_count = self._edge_grid.get_voxel_count()
+        self._edge_ijk = wp.array(shape=(edge_count,), dtype=wp.vec3i, device=self._edge_grid.device)
+        self._edge_grid.get_voxels(out=self._edge_ijk)
+
+    def _ensure_face_grid(self):
+        if self._face_ijk is None:
+            self._build_face_grid()
+
+    def _ensure_edge_grid(self):
+        if self._edge_ijk is None:
+            self._build_edge_grid()
+
 
 @wp.kernel
 def _cell_node_indices(

warp/fem/geometry/quadmesh_2d.py

@@ -299,12 +299,12 @@ class Quadmesh2D(Geometry):
         return side_arg.cell_arg
 
     def _build_topology(self, temporary_store: TemporaryStore):
-        from warp.fem.utils import compress_node_indices, masked_indices
+        from warp.fem.utils import compress_node_indices, host_read_at_index, masked_indices
         from warp.utils import array_scan
 
         device = self.quad_vertex_indices.device
 
-        vertex_quad_offsets, vertex_quad_indices, _, __ = compress_node_indices(
+        vertex_quad_offsets, vertex_quad_indices = compress_node_indices(
             self.vertex_count(), self.quad_vertex_indices, temporary_store=temporary_store
         )
         self._vertex_quad_offsets = vertex_quad_offsets.detach()
@@ -350,16 +350,11 @@ class Quadmesh2D(Geometry):
         array_scan(in_array=vertex_start_edge_count.array, out_array=vertex_unique_edge_offsets.array, inclusive=False)
 
         # Get back edge count to host
-        if device.is_cuda:
-            edge_count = borrow_temporary(temporary_store, shape=(1,), dtype=int, device="cpu", pinned=True)
-            # Last vertex will not own any edge, so its count will be zero; just fetching last prefix count is ok
-            wp.copy(
-                dest=edge_count.array, src=vertex_unique_edge_offsets.array, src_offset=self.vertex_count() - 1, count=1
+        edge_count = int(
+            host_read_at_index(
+                vertex_unique_edge_offsets.array, self.vertex_count() - 1, temporary_store=temporary_store
             )
-            wp.synchronize_stream(wp.get_stream(device))
-            edge_count = int(edge_count.array.numpy()[0])
-        else:
-            edge_count = int(vertex_unique_edge_offsets.array.numpy()[self.vertex_count() - 1])
+        )
 
         self._edge_vertex_indices = wp.empty(shape=(edge_count,), dtype=wp.vec2i, device=device)
         self._edge_quad_indices = wp.empty(shape=(edge_count,), dtype=wp.vec2i, device=device)