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

warp/fem/geometry/trimesh.py

@@ -23,15 +23,13 @@ from warp.fem.cache import (
     cached_arg_value,
 )
 from warp.fem.types import (
-    NULL_ELEMENT_INDEX,
     OUTSIDE,
     Coords,
     ElementIndex,
     Sample,
-    make_free_sample,
 )
 
-from .closest_point import project_on_tri_at_origin
+from .closest_point import project_on_seg_at_origin, project_on_tri_at_origin
 from .element import LinearEdge, Triangle
 from .geometry import Geometry
 
@@ -40,10 +38,6 @@ from .geometry import Geometry
 class TrimeshCellArg:
     tri_vertex_indices: wp.array2d(dtype=int)
 
-    # for neighbor cell lookup
-    vertex_tri_offsets: wp.array(dtype=int)
-    vertex_tri_indices: wp.array(dtype=int)
-
     # for global cell lookup
     tri_bvh: wp.uint64
 
@@ -55,19 +49,6 @@ class TrimeshSideArg:
     edge_tri_indices: wp.array(dtype=wp.vec2i)
 
 
-_NULL_BVH = wp.constant(wp.uint64(0))
-
-
-@wp.func
-def _bvh_vec(v: wp.vec3):
-    return v
-
-
-@wp.func
-def _bvh_vec(v: wp.vec2):
-    return wp.vec3(v[0], v[1], 0.0)
-
-
 class Trimesh(Geometry):
     """Triangular mesh geometry"""
 
@@ -93,14 +74,8 @@ class Trimesh(Geometry):
 
         self._edge_vertex_indices: wp.array = None
         self._edge_tri_indices: wp.array = None
-        self._vertex_tri_offsets: wp.array = None
-        self._vertex_tri_indices: wp.array = None
         self._build_topology(temporary_store)
 
-        self._tri_bvh: wp.Bvh = None
-        if build_bvh:
-            self._build_bvh()
-
         # Flip edges so that normals point away from inner cell
         wp.launch(
             kernel=self._orient_edges,
@@ -110,34 +85,11 @@ class Trimesh(Geometry):
         )
 
         self._make_default_dependent_implementations()
+        self.cell_coordinates = self._make_cell_coordinates(assume_linear=True)
+        self.side_coordinates = self._make_side_coordinates(assume_linear=True)
 
-    def update_bvh(self, force_rebuild: bool = False):
-        """
-        Refits the BVH, or rebuilds it from scratch if `force_rebuild` is ``True``.
-        """
-
-        if self._tri_bvh is None or force_rebuild:
-            return self.build_bvh()
-
-        wp.launch(
-            Trimesh._compute_tri_bounds,
-            self.tri_vertex_indices,
-            self.positions,
-            self._tri_bvh.lowers,
-            self._tri_bvh.uppers,
-        )
-        self._tri_bvh.refit()
-
-    def _build_bvh(self, temporary_store: Optional[TemporaryStore] = None):
-        lowers = wp.array(shape=self.cell_count(), dtype=wp.vec3, device=self.positions.device)
-        uppers = wp.array(shape=self.cell_count(), dtype=wp.vec3, device=self.positions.device)
-        wp.launch(
-            Trimesh._compute_tri_bounds,
-            device=self.positions.device,
-            dim=self.cell_count(),
-            inputs=[self.tri_vertex_indices, self.positions, lowers, uppers],
-        )
-        self._tri_bvh = wp.Bvh(lowers, uppers)
+        if build_bvh:
+            self.build_bvh(self.positions.device)
 
     def cell_count(self):
         return self.tri_vertex_indices.shape[0]
@@ -169,106 +121,42 @@ class Trimesh(Geometry):
     class SideIndexArg:
         boundary_edge_indices: wp.array(dtype=int)
 
-    @cached_arg_value
-    def _cell_topo_arg_value(self, device):
-        args = TrimeshCellArg()
-
+    def _fill_cell_topo_arg(self, args: TrimeshCellArg, device):
         args.tri_vertex_indices = self.tri_vertex_indices.to(device)
-        args.vertex_tri_offsets = self._vertex_tri_offsets.to(device)
-        args.vertex_tri_indices = self._vertex_tri_indices.to(device)
-
-        return args
-
-    @cached_arg_value
-    def _side_topo_arg_value(self, device):
-        args = TrimeshSideArg()
+        args.tri_bvh = self.bvh_id(device)
 
-        args.cell_arg = self._cell_topo_arg_value(device)
+    def _fill_side_topo_arg(self, args: TrimeshSideArg, device):
+        self._fill_cell_topo_arg(args.cell_arg, device)
         args.edge_vertex_indices = self._edge_vertex_indices.to(device)
         args.edge_tri_indices = self._edge_tri_indices.to(device)
 
-        return args
-
-    def _bvh_id(self, device):
-        if self._tri_bvh is None or self._tri_bvh.device != wp.get_device(device):
-            return _NULL_BVH
-        return self._tri_bvh.id
-
     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: TrimeshCellArg, device):
+        self._fill_cell_topo_arg(args.topology, device)
         args.positions = self.positions.to(device)
-        args.topology.tri_bvh = self._bvh_id(device)
-
-        return args
 
     def side_arg_value(self, device):
         args = self.SideArg()
-
-        args.topology = self._side_topo_arg_value(device)
-        args.positions = self.positions.to(device)
-        args.topology.cell_arg.tri_bvh = self._bvh_id(device)
-
+        self.fill_side_arg(args, device)
         return args
 
-    @wp.func
-    def _project_on_tri(args: TrimeshCellArg, positions: wp.array(dtype=Any), pos: Any, tri_index: int):
-        p0 = positions[args.tri_vertex_indices[tri_index, 0]]
-
-        q = pos - p0
-        e1 = positions[args.tri_vertex_indices[tri_index, 1]] - p0
-        e2 = positions[args.tri_vertex_indices[tri_index, 2]] - p0
-
-        dist, coords = project_on_tri_at_origin(q, e1, e2)
-        return dist, coords
-
-    @wp.func
-    def _bvh_lookup(args: TrimeshCellArg, positions: wp.array(dtype=Any), pos: Any):
-        closest_tri = int(NULL_ELEMENT_INDEX)
-        closest_coords = Coords(OUTSIDE)
-        closest_dist = float(1.0e8)
-
-        if args.tri_bvh != _NULL_BVH:
-            bvh_pos = _bvh_vec(pos)
-            query = wp.bvh_query_aabb(args.tri_bvh, bvh_pos, bvh_pos)
-            tri = int(0)
-            while wp.bvh_query_next(query, tri):
-                dist, coords = Trimesh._project_on_tri(args, positions, pos, tri)
-                if dist <= closest_dist:
-                    closest_dist = dist
-                    closest_tri = tri
-                    closest_coords = coords
-
-        return closest_dist, closest_tri, closest_coords
-
-    @wp.func
-    def _cell_neighbor_lookup(args: TrimeshCellArg, positions: wp.array(dtype=Any), pos: Any, cell_index: int):
-        closest_dist = float(1.0e8)
-
-        for v in range(3):
-            vtx = args.tri_vertex_indices[cell_index, v]
-            tri_beg = args.vertex_tri_offsets[vtx]
-            tri_end = args.vertex_tri_offsets[vtx + 1]
-
-            for t in range(tri_beg, tri_end):
-                tri = args.vertex_tri_indices[t]
-                dist, coords = Trimesh._project_on_tri(args, positions, pos, tri)
-                if dist <= closest_dist:
-                    closest_dist = dist
-                    closest_tri = tri
-                    closest_coords = coords
-
-        return closest_dist, closest_tri, closest_coords
+    def fill_side_arg(self, args: TrimeshSideArg, device):
+        self._fill_side_topo_arg(args.topology, device)
+        args.positions = self.positions.to(device)
 
     @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 boundary_side_index(args: SideIndexArg, boundary_side_index: int):
         """Boundary side to side index"""
@@ -516,43 +404,20 @@ class Trimesh(Geometry):
             else:
                 boundary_mask[edge_index] = 0
 
-    @wp.kernel
-    def _compute_tri_bounds(
-        tri_vertex_indices: wp.array2d(dtype=int),
-        positions: wp.array(dtype=Any),
-        lowers: wp.array(dtype=wp.vec3),
-        uppers: wp.array(dtype=wp.vec3),
-    ):
-        t = wp.tid()
-        p0 = _bvh_vec(positions[tri_vertex_indices[t, 0]])
-        p1 = _bvh_vec(positions[tri_vertex_indices[t, 1]])
-        p2 = _bvh_vec(positions[tri_vertex_indices[t, 2]])
-
-        lowers[t] = wp.min(wp.min(p0, p1), p2)
-        uppers[t] = wp.max(wp.max(p0, p1), p2)
-
-
-@wp.struct
-class Trimesh2DCellArg:
-    topology: TrimeshCellArg
-    positions: wp.array(dtype=wp.vec2)
-
-
-@wp.struct
-class Trimesh2DSideArg:
-    topology: TrimeshSideArg
-    positions: wp.array(dtype=wp.vec2)
-
+    @wp.func
+    def cell_bvh_id(cell_arg: Any):
+        return cell_arg.topology.tri_bvh
 
-class Trimesh2D(Trimesh):
-    """2D Triangular mesh geometry"""
+    @wp.func
+    def cell_bounds(cell_arg: Any, cell_index: ElementIndex):
+        p0 = cell_arg.positions[cell_arg.topology.tri_vertex_indices[cell_index, 0]]
+        p1 = cell_arg.positions[cell_arg.topology.tri_vertex_indices[cell_index, 1]]
+        p2 = cell_arg.positions[cell_arg.topology.tri_vertex_indices[cell_index, 2]]
 
-    dimension = 2
-    CellArg = Trimesh2DCellArg
-    SideArg = Trimesh2DSideArg
+        return wp.min(wp.min(p0, p1), p2), wp.max(wp.max(p0, p1), p2)
 
     @wp.func
-    def cell_position(args: CellArg, s: Sample):
+    def cell_position(args: Any, s: Sample):
         tri_idx = args.topology.tri_vertex_indices[s.element_index]
         return (
             s.element_coords[0] * args.positions[tri_idx[0]]
@@ -561,7 +426,7 @@ class Trimesh2D(Trimesh):
         )
 
     @wp.func
-    def cell_deformation_gradient(args: CellArg, s: Sample):
+    def cell_deformation_gradient(args: Any, s: Sample):
         tri_idx = args.topology.tri_vertex_indices[s.element_index]
         p0 = args.positions[tri_idx[0]]
         p1 = args.positions[tri_idx[1]]
@@ -569,72 +434,89 @@ class Trimesh2D(Trimesh):
         return wp.matrix_from_cols(p1 - p0, p2 - p0)
 
     @wp.func
-    def cell_lookup(args: CellArg, pos: wp.vec2):
-        closest_dist, closest_tri, closest_coords = Trimesh._bvh_lookup(args.topology, args.positions, pos)
+    def cell_closest_point(args: Any, tri_index: ElementIndex, pos: Any):
+        vidx = args.topology.tri_vertex_indices[tri_index]
+        p0 = args.positions[vidx[0]]
 
-        return make_free_sample(closest_tri, closest_coords)
-
-    @wp.func
-    def cell_lookup(args: CellArg, pos: wp.vec2, guess: Sample):
-        closest_dist, closest_tri, closest_coords = Trimesh._bvh_lookup(args.topology, args.positions, pos)
-
-        if closest_tri == NULL_ELEMENT_INDEX:
-            closest_dist, closest_tri, closest_coords = Trimesh._cell_neighbor_lookup(
-                args.topology, args.positions, pos, guess.element_index
-            )
+        q = pos - p0
+        e1 = args.positions[vidx[1]] - p0
+        e2 = args.positions[vidx[2]] - p0
 
-        return make_free_sample(closest_tri, closest_coords)
+        dist, coords = project_on_tri_at_origin(q, e1, e2)
+        return coords, dist
 
     @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_normal(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]]
-        e = v1 - v0
+    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
 
-        return wp.normalize(wp.vec2(-e[1], e[0]))
+        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: SideArg, side_index: ElementIndex):
+    def side_inner_cell_index(arg: Any, side_index: ElementIndex):
         return arg.topology.edge_tri_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_tri_indices[side_index][1]
 
     @wp.func
-    def side_inner_cell_coords(args: SideArg, side_index: ElementIndex, side_coords: Coords):
-        inner_cell_index = Trimesh2D.side_inner_cell_index(args, side_index)
+    def side_inner_cell_coords(args: Any, side_index: ElementIndex, side_coords: Coords):
+        inner_cell_index = Trimesh.side_inner_cell_index(args, side_index)
         return Trimesh._edge_to_tri_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 = Trimesh2D.side_outer_cell_index(args, side_index)
+    def side_outer_cell_coords(args: Any, side_index: ElementIndex, side_coords: Coords):
+        outer_cell_index = Trimesh.side_outer_cell_index(args, side_index)
         return Trimesh._edge_to_tri_coords(args.topology, side_index, outer_cell_index, side_coords)
 
     @wp.func
     def side_from_cell_coords(
-        args: SideArg,
+        args: Any,
         side_index: ElementIndex,
         tri_index: ElementIndex,
         tri_coords: Coords,
     ):
         return Trimesh._tri_to_edge_coords(args.topology, side_index, tri_index, tri_coords)
 
+
+@wp.struct
+class Trimesh2DCellArg:
+    topology: TrimeshCellArg
+    positions: wp.array(dtype=wp.vec2)
+
+
+@wp.struct
+class Trimesh2DSideArg:
+    topology: TrimeshSideArg
+    positions: wp.array(dtype=wp.vec2)
+
+
+class Trimesh2D(Trimesh):
+    """2D Triangular mesh geometry"""
+
+    dimension = 2
+    CellArg = Trimesh2DCellArg
+    SideArg = Trimesh2DSideArg
+
     @wp.func
     def side_to_cell_arg(side_arg: SideArg):
         return Trimesh2DCellArg(side_arg.topology.cell_arg, side_arg.positions)
@@ -660,7 +542,7 @@ class Trimesh2D(Trimesh):
 
         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(tri_centroid - edge_center, edge_normal) > 0.0:
@@ -686,81 +568,6 @@ class Trimesh3D(Trimesh):
     CellArg = Trimesh3DCellArg
     SideArg = Trimesh3DSideArg
 
-    @wp.func
-    def cell_position(args: CellArg, s: Sample):
-        tri_idx = args.topology.tri_vertex_indices[s.element_index]
-        return (
-            s.element_coords[0] * args.positions[tri_idx[0]]
-            + s.element_coords[1] * args.positions[tri_idx[1]]
-            + s.element_coords[2] * args.positions[tri_idx[2]]
-        )
-
-    @wp.func
-    def cell_deformation_gradient(args: CellArg, s: Sample):
-        tri_idx = args.topology.tri_vertex_indices[s.element_index]
-        p0 = args.positions[tri_idx[0]]
-        p1 = args.positions[tri_idx[1]]
-        p2 = args.positions[tri_idx[2]]
-        return wp.matrix_from_cols(p1 - p0, p2 - p0)
-
-    @wp.func
-    def cell_lookup(args: CellArg, pos: wp.vec3):
-        closest_dist, closest_tri, closest_coords = Trimesh._bvh_lookup(args.topology, args.positions, pos)
-
-        return make_free_sample(closest_tri, closest_coords)
-
-    @wp.func
-    def cell_lookup(args: CellArg, pos: wp.vec3, guess: Sample):
-        closest_dist, closest_tri, closest_coords = Trimesh._bvh_lookup(args.topology, args.positions, pos)
-
-        if closest_tri == NULL_ELEMENT_INDEX:
-            closest_dist, closest_tri, closest_coords = Trimesh._cell_neighbor_lookup(
-                args.topology, args.positions, pos, guess.element_index
-            )
-
-        return make_free_sample(closest_tri, closest_coords)
-
-    @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_tri_indices[side_index][0]
-
-    @wp.func
-    def side_outer_cell_index(arg: SideArg, side_index: ElementIndex):
-        return arg.topology.edge_tri_indices[side_index][1]
-
-    @wp.func
-    def side_inner_cell_coords(args: SideArg, side_index: ElementIndex, side_coords: Coords):
-        inner_cell_index = Trimesh3D.side_inner_cell_index(args, side_index)
-        return Trimesh._edge_to_tri_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 = Trimesh3D.side_outer_cell_index(args, side_index)
-        return Trimesh._edge_to_tri_coords(args.topology, side_index, outer_cell_index, side_coords)
-
-    @wp.func
-    def side_from_cell_coords(
-        args: SideArg,
-        side_index: ElementIndex,
-        tri_index: ElementIndex,
-        tri_coords: Coords,
-    ):
-        return Trimesh._tri_to_edge_coords(args.topology, side_index, tri_index, tri_coords)
-
     @wp.func
     def side_to_cell_arg(side_arg: SideArg):
         return Trimesh3DCellArg(side_arg.topology.cell_arg, side_arg.positions)