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/partition.py

@@ -13,10 +13,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from functools import cached_property
 from typing import Any
 
 import warp as wp
-from warp.fem.cache import TemporaryStore, borrow_temporary, cached_arg_value
+from warp.fem.cache import TemporaryStore, borrow_temporary, cached_arg_value, dynamic_struct
 from warp.fem.types import NULL_ELEMENT_INDEX, ElementIndex
 from warp.fem.utils import masked_indices
 
@@ -63,9 +64,15 @@ class GeometryPartition:
     def cell_arg_value(self, device):
         raise NotImplementedError()
 
+    def fill_cell_arg(self, args: CellArg, device):
+        raise NotImplementedError()
+
     def side_arg_value(self, device):
         raise NotImplementedError()
 
+    def fill_side_arg(self, args: SideArg, device):
+        raise NotImplementedError()
+
     @staticmethod
     def cell_index(args: CellArg, partition_cell_index: int):
         """Index in the geometry of a partition cell"""
@@ -103,6 +110,7 @@ class WholeGeometryPartition(GeometryPartition):
 
         self.SideArg = geometry.SideIndexArg
         self.side_arg_value = geometry.side_index_arg_value
+        self.fill_side_arg = geometry.fill_side_index_arg
 
         self.cell_index = WholeGeometryPartition._identity_element_index
         self.partition_cell_index = WholeGeometryPartition._identity_element_index
@@ -135,6 +143,9 @@ class WholeGeometryPartition(GeometryPartition):
         arg = WholeGeometryPartition.CellArg()
         return arg
 
+    def fill_cell_arg(self, args: CellArg, device):
+        pass
+
     @wp.func
     def _identity_element_index(args: Any, idx: ElementIndex):
         return idx
@@ -143,6 +154,10 @@ class WholeGeometryPartition(GeometryPartition):
     def name(self) -> str:
         return self.geometry.name
 
+    @wp.func
+    def side_to_cell_arg(side_arg: Any):
+        return WholeGeometryPartition.CellArg()
+
 
 class CellBasedGeometryPartition(GeometryPartition):
     """Geometry partition based on a subset of cells. Interior, boundary and frontier sides are automatically categorized."""
@@ -154,11 +169,19 @@ class CellBasedGeometryPartition(GeometryPartition):
     ):
         super().__init__(geometry)
 
-    @wp.struct
-    class SideArg:
-        partition_side_indices: wp.array(dtype=int)
-        boundary_side_indices: wp.array(dtype=int)
-        frontier_side_indices: wp.array(dtype=int)
+    @cached_property
+    def SideArg(self):
+        return self._make_side_arg()
+
+    def _make_side_arg(self):
+        @dynamic_struct(suffix=self.name)
+        class SideArg:
+            cell_arg: self.CellArg
+            partition_side_indices: wp.array(dtype=int)
+            boundary_side_indices: wp.array(dtype=int)
+            frontier_side_indices: wp.array(dtype=int)
+
+        return SideArg
 
     def side_count(self) -> int:
         return self._partition_side_indices.array.shape[0]
@@ -171,24 +194,28 @@ class CellBasedGeometryPartition(GeometryPartition):
 
     @cached_arg_value
     def side_arg_value(self, device):
-        arg = LinearGeometryPartition.SideArg()
-        arg.partition_side_indices = self._partition_side_indices.array.to(device)
-        arg.boundary_side_indices = self._boundary_side_indices.array.to(device)
-        arg.frontier_side_indices = self._frontier_side_indices.array.to(device)
+        arg = self.SideArg()
+        self.fill_side_arg(arg, device)
         return arg
 
+    def fill_side_arg(self, args: SideArg, device):
+        self.fill_cell_arg(args.cell_arg, device)
+        args.partition_side_indices = self._partition_side_indices.array.to(device)
+        args.boundary_side_indices = self._boundary_side_indices.array.to(device)
+        args.frontier_side_indices = self._frontier_side_indices.array.to(device)
+
     @wp.func
-    def side_index(args: SideArg, partition_side_index: int):
+    def side_index(args: Any, partition_side_index: int):
         """partition side to side index"""
         return args.partition_side_indices[partition_side_index]
 
     @wp.func
-    def boundary_side_index(args: SideArg, boundary_side_index: int):
+    def boundary_side_index(args: Any, boundary_side_index: int):
         """Boundary side to side index"""
         return args.boundary_side_indices[boundary_side_index]
 
     @wp.func
-    def frontier_side_index(args: SideArg, frontier_side_index: int):
+    def frontier_side_index(args: Any, frontier_side_index: int):
         """Frontier side to side index"""
         return args.frontier_side_indices[frontier_side_index]
 
@@ -271,6 +298,10 @@ class CellBasedGeometryPartition(GeometryPartition):
         boundary_side_mask.release()
         frontier_side_mask.release()
 
+    @wp.func
+    def side_to_cell_arg(side_arg: Any):
+        return side_arg.cell_arg
+
 
 class LinearGeometryPartition(CellBasedGeometryPartition):
     def __init__(
@@ -314,10 +345,13 @@ class LinearGeometryPartition(CellBasedGeometryPartition):
 
     def cell_arg_value(self, device):
         arg = LinearGeometryPartition.CellArg()
-        arg.cell_begin = self.cell_begin
-        arg.cell_end = self.cell_end
+        self.fill_cell_arg(arg, device)
         return arg
 
+    def fill_cell_arg(self, args: CellArg, device):
+        args.cell_begin = self.cell_begin
+        args.cell_end = self.cell_end
+
     @wp.func
     def cell_index(args: CellArg, partition_cell_index: int):
         """Partition cell to cell index"""
@@ -326,14 +360,11 @@ class LinearGeometryPartition(CellBasedGeometryPartition):
     @wp.func
     def partition_cell_index(args: CellArg, cell_index: int):
         """Partition cell to cell index"""
-        if cell_index > args.cell_end:
-            return NULL_ELEMENT_INDEX
-
-        partition_cell_index = cell_index - args.cell_begin
-        if partition_cell_index < 0:
-            return NULL_ELEMENT_INDEX
-
-        return partition_cell_index
+        return wp.where(
+            cell_index >= args.cell_begin and cell_index < args.cell_end,
+            cell_index - args.cell_begin,
+            NULL_ELEMENT_INDEX,
+        )
 
     @wp.func
     def _cell_inclusion_test(arg: CellArg, cell_index: int):
@@ -372,10 +403,13 @@ class ExplicitGeometryPartition(CellBasedGeometryPartition):
     @cached_arg_value
     def cell_arg_value(self, device):
         arg = ExplicitGeometryPartition.CellArg()
-        arg.cell_index = self._cells.array.to(device)
-        arg.partition_cell_index = self._partition_cells.array.to(device)
+        self.fill_cell_arg(arg, device)
         return arg
 
+    def fill_cell_arg(self, args: CellArg, device):
+        args.cell_index = self._cells.array.to(device)
+        args.partition_cell_index = self._partition_cells.array.to(device)
+
     @wp.func
     def cell_index(args: CellArg, partition_cell_index: int):
         return args.cell_index[partition_cell_index]

warp/fem/geometry/quadmesh.py

@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional
+from typing import Any, Optional
 
 import warp as wp
 from warp.fem.cache import (
@@ -24,6 +24,7 @@ from warp.fem.cache import (
 )
 from warp.fem.types import OUTSIDE, Coords, ElementIndex, Sample
 
+from .closest_point import project_on_seg_at_origin
 from .element import LinearEdge, Square
 from .geometry import Geometry
 
@@ -32,9 +33,7 @@ from .geometry import Geometry
 class QuadmeshCellArg:
     quad_vertex_indices: wp.array2d(dtype=int)
 
-    # for neighbor cell lookup
-    vertex_quad_offsets: wp.array(dtype=int)
-    vertex_quad_indices: wp.array(dtype=int)
+    quad_bvh: wp.uint64
 
 
 @wp.struct
@@ -48,7 +47,11 @@ class Quadmesh(Geometry):
     """Quadrilateral mesh geometry"""
 
     def __init__(
-        self, quad_vertex_indices: wp.array, positions: wp.array, temporary_store: Optional[TemporaryStore] = None
+        self,
+        quad_vertex_indices: wp.array,
+        positions: wp.array,
+        build_bvh: bool = False,
+        temporary_store: Optional[TemporaryStore] = None,
     ):
         """
         Constructs a D-dimensional quadrilateral mesh.
@@ -77,6 +80,12 @@ class Quadmesh(Geometry):
         )
 
         self._make_default_dependent_implementations()
+        self.cell_closest_point = self._make_cell_closest_point()
+        self.cell_coordinates = self._make_cell_coordinates()
+        self.side_coordinates = self._make_side_coordinates(assume_linear=True)
+
+        if build_bvh:
+            self.build_bvh(self.positions.device)
 
     def cell_count(self):
         return self.quad_vertex_indices.shape[0]
@@ -108,50 +117,42 @@ class Quadmesh(Geometry):
     class SideIndexArg:
         boundary_edge_indices: wp.array(dtype=int)
 
-    @cached_arg_value
-    def _cell_topo_arg_value(self, device):
-        args = QuadmeshCellArg()
-
+    def fill_cell_topo_arg(self, args: QuadmeshCellArg, device):
         args.quad_vertex_indices = self.quad_vertex_indices.to(device)
-        args.vertex_quad_offsets = self._vertex_quad_offsets.to(device)
-        args.vertex_quad_indices = self._vertex_quad_indices.to(device)
-
-        return args
+        args.quad_bvh = self.bvh_id(device)
 
-    @cached_arg_value
-    def _side_topo_arg_value(self, device):
-        args = QuadmeshSideArg()
-
-        args.cell_arg = self._cell_topo_arg_value(device)
+    def fill_side_topo_arg(self, args: QuadmeshSideArg, device):
+        self.fill_cell_topo_arg(args.cell_arg, device)
         args.edge_vertex_indices = self._edge_vertex_indices.to(device)
         args.edge_quad_indices = self._edge_quad_indices.to(device)
 
-        return args
-
     def cell_arg_value(self, device):
         args = self.CellArg()
+        self.fill_cell_arg(args, device)
+        return args
 
-        args.topology = self._cell_topo_arg_value(device)
+    def fill_cell_arg(self, args: "Quadmesh.CellArg", device):
+        self.fill_cell_topo_arg(args.topology, device)
         args.positions = self.positions.to(device)
 
-        return args
-
     def side_arg_value(self, device):
         args = self.SideArg()
+        self.fill_side_arg(args, device)
+        return args
 
-        args.topology = self._side_topo_arg_value(device)
+    def fill_side_arg(self, args: "Quadmesh.SideArg", device):
+        self.fill_side_topo_arg(args.topology, device)
         args.positions = self.positions.to(device)
 
-        return args
-
     @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_edge_indices = self._boundary_edge_indices.to(device)
 
-        return args
-
     @wp.func
     def _edge_to_quad_coords(
         args: QuadmeshSideArg, side_index: ElementIndex, quad_index: ElementIndex, side_coords: Coords
@@ -400,28 +401,8 @@ class Quadmesh(Geometry):
             else:
                 boundary_mask[edge_index] = 0
 
-
-@wp.struct
-class Quadmesh2DCellArg:
-    topology: QuadmeshCellArg
-    positions: wp.array(dtype=wp.vec2)
-
-
-@wp.struct
-class Quadmesh2DSideArg:
-    topology: QuadmeshSideArg
-    positions: wp.array(dtype=wp.vec2)
-
-
-class Quadmesh2D(Quadmesh):
-    """Two-dimensional quadrilateral mesh"""
-
-    dimension = 2
-    CellArg = Quadmesh2DCellArg
-    SideArg = Quadmesh2DSideArg
-
     @wp.func
-    def cell_position(args: CellArg, s: Sample):
+    def cell_position(args: Any, s: Sample):
         quad_idx = args.topology.quad_vertex_indices[s.element_index]
 
         w_p = s.element_coords
@@ -440,7 +421,7 @@ class Quadmesh2D(Quadmesh):
         )
 
     @wp.func
-    def cell_deformation_gradient(cell_arg: CellArg, s: Sample):
+    def cell_deformation_gradient(cell_arg: Any, s: Sample):
         """Deformation gradient at `coords`"""
         quad_idx = cell_arg.topology.quad_vertex_indices[s.element_index]
 
@@ -455,46 +436,91 @@ class Quadmesh2D(Quadmesh):
         )
 
     @wp.func
-    def side_position(args: SideArg, s: Sample):
+    def side_position(args: Any, s: Sample):
         edge_idx = args.topology.edge_vertex_indices[s.element_index]
         return (1.0 - s.element_coords[0]) * args.positions[edge_idx[0]] + s.element_coords[0] * args.positions[
             edge_idx[1]
         ]
 
     @wp.func
-    def side_deformation_gradient(args: SideArg, s: Sample):
+    def side_deformation_gradient(args: Any, s: Sample):
         edge_idx = args.topology.edge_vertex_indices[s.element_index]
         v0 = args.positions[edge_idx[0]]
         v1 = args.positions[edge_idx[1]]
         return v1 - v0
 
     @wp.func
-    def side_inner_cell_index(arg: SideArg, side_index: ElementIndex):
+    def side_closest_point(args: Any, side_index: ElementIndex, pos: Any):
+        edge_idx = args.topology.edge_vertex_indices[side_index]
+        p0 = args.positions[edge_idx[0]]
+
+        q = pos - p0
+        e = args.positions[edge_idx[1]] - p0
+
+        dist, t = project_on_seg_at_origin(q, e, wp.lengh_sq(e))
+        return Coords(t, 0.0, 0.0), dist
+
+    @wp.func
+    def side_inner_cell_index(arg: Any, side_index: ElementIndex):
         return arg.topology.edge_quad_indices[side_index][0]
 
     @wp.func
-    def side_outer_cell_index(arg: SideArg, side_index: ElementIndex):
+    def side_outer_cell_index(arg: Any, side_index: ElementIndex):
         return arg.topology.edge_quad_indices[side_index][1]
 
     @wp.func
-    def side_inner_cell_coords(args: SideArg, side_index: ElementIndex, side_coords: Coords):
-        inner_cell_index = Quadmesh2D.side_inner_cell_index(args, side_index)
+    def side_inner_cell_coords(args: Any, side_index: ElementIndex, side_coords: Coords):
+        inner_cell_index = Quadmesh3D.side_inner_cell_index(args, side_index)
         return Quadmesh._edge_to_quad_coords(args.topology, side_index, inner_cell_index, side_coords)
 
     @wp.func
-    def side_outer_cell_coords(args: SideArg, side_index: ElementIndex, side_coords: Coords):
-        outer_cell_index = Quadmesh2D.side_outer_cell_index(args, side_index)
+    def side_outer_cell_coords(args: Any, side_index: ElementIndex, side_coords: Coords):
+        outer_cell_index = Quadmesh3D.side_outer_cell_index(args, side_index)
         return Quadmesh._edge_to_quad_coords(args.topology, side_index, outer_cell_index, side_coords)
 
     @wp.func
     def side_from_cell_coords(
-        args: SideArg,
+        args: Any,
         side_index: ElementIndex,
         quad_index: ElementIndex,
         quad_coords: Coords,
     ):
         return Quadmesh._quad_to_edge_coords(args.topology, side_index, quad_index, quad_coords)
 
+    @wp.func
+    def cell_bvh_id(cell_arg: Any):
+        return cell_arg.topology.quad_bvh
+
+    @wp.func
+    def cell_bounds(cell_arg: Any, cell_index: ElementIndex):
+        vidx = cell_arg.topology.quad_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]]
+
+        return wp.min(wp.min(p0, p1), wp.min(p2, p3)), wp.max(wp.max(p0, p1), wp.max(p2, p3))
+
+
+@wp.struct
+class Quadmesh2DCellArg:
+    topology: QuadmeshCellArg
+    positions: wp.array(dtype=wp.vec2)
+
+
+@wp.struct
+class Quadmesh2DSideArg:
+    topology: QuadmeshSideArg
+    positions: wp.array(dtype=wp.vec2)
+
+
+class Quadmesh2D(Quadmesh):
+    """Two-dimensional quadrilateral mesh"""
+
+    dimension = 2
+    CellArg = Quadmesh2DCellArg
+    SideArg = Quadmesh2DSideArg
+
     @wp.func
     def side_to_cell_arg(side_arg: SideArg):
         return Quadmesh2DCellArg(side_arg.topology.cell_arg, side_arg.positions)
@@ -522,7 +548,7 @@ class Quadmesh2D(Quadmesh):
 
         edge_center = 0.5 * (v1 + v0)
         edge_vec = v1 - v0
-        edge_normal = wp.vec2(-edge_vec[1], edge_vec[0])
+        edge_normal = Geometry._element_normal(edge_vec)
 
         # if edge normal points toward first triangle centroid, flip indices
         if wp.dot(quad_centroid - edge_center, edge_normal) > 0.0:
@@ -548,81 +574,6 @@ class Quadmesh3D(Quadmesh):
     CellArg = Quadmesh3DCellArg
     SideArg = Quadmesh3DSideArg
 
-    @wp.func
-    def cell_position(args: CellArg, s: Sample):
-        quad_idx = args.topology.quad_vertex_indices[s.element_index]
-
-        w_p = s.element_coords
-        w_m = Coords(1.0) - s.element_coords
-
-        # 0 : m m
-        # 1 : p m
-        # 2 : p p
-        # 3 : m p
-
-        return (
-            w_m[0] * w_m[1] * args.positions[quad_idx[0]]
-            + w_p[0] * w_m[1] * args.positions[quad_idx[1]]
-            + w_p[0] * w_p[1] * args.positions[quad_idx[2]]
-            + w_m[0] * w_p[1] * args.positions[quad_idx[3]]
-        )
-
-    @wp.func
-    def cell_deformation_gradient(cell_arg: CellArg, s: Sample):
-        """Deformation gradient at `coords`"""
-        quad_idx = cell_arg.topology.quad_vertex_indices[s.element_index]
-
-        w_p = s.element_coords
-        w_m = Coords(1.0) - s.element_coords
-
-        return (
-            wp.outer(cell_arg.positions[quad_idx[0]], wp.vec2(-w_m[1], -w_m[0]))
-            + wp.outer(cell_arg.positions[quad_idx[1]], wp.vec2(w_m[1], -w_p[0]))
-            + wp.outer(cell_arg.positions[quad_idx[2]], wp.vec2(w_p[1], w_p[0]))
-            + wp.outer(cell_arg.positions[quad_idx[3]], wp.vec2(-w_p[1], w_m[0]))
-        )
-
-    @wp.func
-    def side_position(args: SideArg, s: Sample):
-        edge_idx = args.topology.edge_vertex_indices[s.element_index]
-        return (1.0 - s.element_coords[0]) * args.positions[edge_idx[0]] + s.element_coords[0] * args.positions[
-            edge_idx[1]
-        ]
-
-    @wp.func
-    def side_deformation_gradient(args: SideArg, s: Sample):
-        edge_idx = args.topology.edge_vertex_indices[s.element_index]
-        v0 = args.positions[edge_idx[0]]
-        v1 = args.positions[edge_idx[1]]
-        return v1 - v0
-
-    @wp.func
-    def side_inner_cell_index(arg: SideArg, side_index: ElementIndex):
-        return arg.topology.edge_quad_indices[side_index][0]
-
-    @wp.func
-    def side_outer_cell_index(arg: SideArg, side_index: ElementIndex):
-        return arg.topology.edge_quad_indices[side_index][1]
-
-    @wp.func
-    def side_inner_cell_coords(args: SideArg, side_index: ElementIndex, side_coords: Coords):
-        inner_cell_index = Quadmesh3D.side_inner_cell_index(args, side_index)
-        return Quadmesh._edge_to_quad_coords(args.topology, side_index, inner_cell_index, side_coords)
-
-    @wp.func
-    def side_outer_cell_coords(args: SideArg, side_index: ElementIndex, side_coords: Coords):
-        outer_cell_index = Quadmesh3D.side_outer_cell_index(args, side_index)
-        return Quadmesh._edge_to_quad_coords(args.topology, side_index, outer_cell_index, side_coords)
-
-    @wp.func
-    def side_from_cell_coords(
-        args: SideArg,
-        side_index: ElementIndex,
-        quad_index: ElementIndex,
-        quad_coords: Coords,
-    ):
-        return Quadmesh._quad_to_edge_coords(args.topology, side_index, quad_index, quad_coords)
-
     @wp.func
     def side_to_cell_arg(side_arg: SideArg):
         return Quadmesh3DCellArg(side_arg.topology.cell_arg, side_arg.positions)