warp-lang 1.0.0b2__py3-none-win_amd64.whl → 1.0.0b6__py3-none-win_amd64.whl

warp/fem/geometry/trimesh_2d.py

@@ -1,38 +1,61 @@
-import warp as wp
+from typing import Optional
 
-from warp.fem.types import ElementIndex, Coords, vec2i, Sample
-from warp.fem.types import NULL_ELEMENT_INDEX, OUTSIDE, NULL_DOF_INDEX, NULL_QP_INDEX
+import warp as wp
+from warp.fem.cache import (
+    TemporaryStore,
+    borrow_temporary,
+    borrow_temporary_like,
+    cached_arg_value,
+)
+from warp.fem.types import (
+    NULL_ELEMENT_INDEX,
+    OUTSIDE,
+    Coords,
+    ElementIndex,
+    Sample,
+    make_free_sample,
+)
 
-from .geometry import Geometry
-from .element import Triangle, LinearEdge
 from .closest_point import project_on_tri_at_origin
+from .element import LinearEdge, Triangle
+from .geometry import Geometry
 
 
 @wp.struct
-class Trimesh2DArg:
+class Trimesh2DCellArg:
     tri_vertex_indices: wp.array2d(dtype=int)
     positions: wp.array(dtype=wp.vec2)
 
+    # for neighbor cell lookup
     vertex_tri_offsets: wp.array(dtype=int)
     vertex_tri_indices: wp.array(dtype=int)
 
-    edge_vertex_indices: wp.array(dtype=vec2i)
-    edge_tri_indices: wp.array(dtype=vec2i)
+    deformation_gradients: wp.array(dtype=wp.mat22f)
+
+
+@wp.struct
+class Trimesh2DSideArg:
+    cell_arg: Trimesh2DCellArg
+    edge_vertex_indices: wp.array(dtype=wp.vec2i)
+    edge_tri_indices: wp.array(dtype=wp.vec2i)
 
 
 class Trimesh2D(Geometry):
     """Two-dimensional triangular mesh geometry"""
 
-    def __init__(self, tri_vertex_indices: wp.array, positions: wp.array):
+    dimension = 2
+
+    def __init__(
+        self, tri_vertex_indices: wp.array, positions: wp.array, temporary_store: Optional[TemporaryStore] = None
+    ):
         """
         Constructs a two-dimensional triangular mesh.
 
         Args:
             tri_vertex_indices: warp array of shape (num_tris, 3) containing vertex indices for each tri
             positions: warp array of shape (num_vertices, 2) containing 2d position for each vertex
-
+            temporary_store: shared pool from which to allocate temporary arrays
         """
-        self.dimension = 2
 
         self.tri_vertex_indices = tri_vertex_indices
         self.positions = positions
@@ -41,8 +64,10 @@ class Trimesh2D(Geometry):
         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._build_topology()
+        self._deformation_gradients: wp.array = None
+        self._compute_deformation_gradients()
 
     def cell_count(self):
         return self.tri_vertex_indices.shape[0]
@@ -62,8 +87,16 @@ class Trimesh2D(Geometry):
     def reference_side(self) -> LinearEdge:
         return LinearEdge()
 
-    CellArg = Trimesh2DArg
-    SideArg = Trimesh2DArg
+    @property
+    def edge_tri_indices(self) -> wp.array:
+        return self._edge_tri_indices
+
+    @property
+    def edge_vertex_indices(self) -> wp.array:
+        return self._edge_vertex_indices
+
+    CellArg = Trimesh2DCellArg
+    SideArg = Trimesh2DSideArg
 
     @wp.struct
     class SideIndexArg:
@@ -71,15 +104,15 @@ class Trimesh2D(Geometry):
 
     # Geometry device interface
 
+    @cached_arg_value
     def cell_arg_value(self, device) -> CellArg:
         args = self.CellArg()
 
         args.tri_vertex_indices = self.tri_vertex_indices.to(device)
         args.positions = self.positions.to(device)
-        args.edge_vertex_indices = self._edge_vertex_indices.to(device)
-        args.edge_tri_indices = self._edge_tri_indices.to(device)
         args.vertex_tri_offsets = self._vertex_tri_offsets.to(device)
         args.vertex_tri_indices = self._vertex_tri_indices.to(device)
+        args.deformation_gradients = self._deformation_gradients.to(device)
 
         return args
 
@@ -92,6 +125,14 @@ class Trimesh2D(Geometry):
             + s.element_coords[2] * args.positions[tri_idx[2]]
         )
 
+    @wp.func
+    def cell_deformation_gradient(args: CellArg, s: Sample):
+        return args.deformation_gradients[s.element_index]
+
+    @wp.func
+    def cell_inverse_deformation_gradient(args: CellArg, s: Sample):
+        return wp.inverse(args.deformation_gradients[s.element_index])
+
     @wp.func
     def _project_on_tri(args: CellArg, pos: wp.vec2, tri_index: int):
         p0 = args.positions[args.tri_vertex_indices[tri_index, 0]]
@@ -122,36 +163,17 @@ class Trimesh2D(Geometry):
                     closest_tri = tri
                     closest_coords = coords
 
-        return Sample(closest_tri, closest_coords, NULL_QP_INDEX, 0.0, NULL_DOF_INDEX, NULL_DOF_INDEX)
-
-    @wp.func
-    def cell_measure(args: CellArg, cell_index: ElementIndex, coords: Coords):
-        tri_idx = args.tri_vertex_indices[cell_index]
-
-        v0 = args.positions[tri_idx[0]]
-        v1 = args.positions[tri_idx[1]]
-        v2 = args.positions[tri_idx[2]]
-
-        e1 = v1 - v0
-        e2 = v2 - v0
-
-        return 0.5 * wp.abs(e1[0] * e2[1] - e1[1] * e2[0])
+        return make_free_sample(closest_tri, closest_coords)
 
     @wp.func
     def cell_measure(args: CellArg, s: Sample):
-        return Trimesh2D.cell_measure(args, s.element_index, s.element_coords)
-
-    @wp.func
-    def cell_measure_ratio(args: CellArg, s: Sample):
-        return 1.0
+        return 0.5 * wp.abs(wp.determinant(args.deformation_gradients[s.element_index]))
 
     @wp.func
     def cell_normal(args: CellArg, s: Sample):
         return wp.vec2(0.0)
 
-    def side_arg_value(self, device) -> SideArg:
-        return self.cell_arg_value(device)
-
+    @cached_arg_value
     def side_index_arg_value(self, device) -> SideIndexArg:
         args = self.SideIndexArg()
 
@@ -165,38 +187,61 @@ class Trimesh2D(Geometry):
 
         return args.boundary_edge_indices[boundary_side_index]
 
+    @cached_arg_value
+    def side_arg_value(self, device) -> CellArg:
+        args = self.SideArg()
+
+        args.cell_arg = self.cell_arg_value(device)
+        args.edge_vertex_indices = self._edge_vertex_indices.to(device)
+        args.edge_tri_indices = self._edge_tri_indices.to(device)
+
+        return args
+
     @wp.func
     def side_position(args: SideArg, s: Sample):
         edge_idx = args.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]
-        ]
+        return (1.0 - s.element_coords[0]) * args.cell_arg.positions[edge_idx[0]] + s.element_coords[
+            0
+        ] * args.cell_arg.positions[edge_idx[1]]
 
     @wp.func
-    def side_measure(args: SideArg, side_index: ElementIndex, coords: Coords):
-        edge_idx = args.edge_vertex_indices[side_index]
-        v0 = args.positions[edge_idx[0]]
-        v1 = args.positions[edge_idx[1]]
-        return wp.length(v1 - v0)
+    def side_deformation_gradient(args: SideArg, s: Sample):
+        edge_idx = args.edge_vertex_indices[s.element_index]
+        v0 = args.cell_arg.positions[edge_idx[0]]
+        v1 = args.cell_arg.positions[edge_idx[1]]
+        return v1 - v0
+
+    @wp.func
+    def side_inner_inverse_deformation_gradient(args: SideArg, s: Sample):
+        cell_index = Trimesh2D.side_inner_cell_index(args, s.element_index)
+        return wp.inverse(args.cell_arg.deformation_gradients[cell_index])
+
+    @wp.func
+    def side_outer_inverse_deformation_gradient(args: SideArg, s: Sample):
+        cell_index = Trimesh2D.side_outer_cell_index(args, s.element_index)
+        return wp.inverse(args.cell_arg.deformation_gradients[cell_index])
 
     @wp.func
     def side_measure(args: SideArg, s: Sample):
-        return Trimesh2D.side_measure(args, s.element_index, s.element_coords)
+        edge_idx = args.edge_vertex_indices[s.element_index]
+        v0 = args.cell_arg.positions[edge_idx[0]]
+        v1 = args.cell_arg.positions[edge_idx[1]]
+        return wp.length(v1 - v0)
 
     @wp.func
     def side_measure_ratio(args: SideArg, s: Sample):
         inner = Trimesh2D.side_inner_cell_index(args, s.element_index)
         outer = Trimesh2D.side_outer_cell_index(args, s.element_index)
         return Trimesh2D.side_measure(args, s) / wp.min(
-            Trimesh2D.cell_measure(args, inner, Coords()),
-            Trimesh2D.cell_measure(args, outer, Coords()),
+            Trimesh2D.cell_measure(args.cell_arg, make_free_sample(inner, Coords())),
+            Trimesh2D.cell_measure(args.cell_arg, make_free_sample(outer, Coords())),
         )
 
     @wp.func
     def side_normal(args: SideArg, s: Sample):
         edge_idx = args.edge_vertex_indices[s.element_index]
-        v0 = args.positions[edge_idx[0]]
-        v1 = args.positions[edge_idx[1]]
+        v0 = args.cell_arg.positions[edge_idx[0]]
+        v1 = args.cell_arg.positions[edge_idx[1]]
         e = v1 - v0
 
         return wp.normalize(wp.vec2(-e[1], e[0]))
@@ -212,7 +257,7 @@ class Trimesh2D(Geometry):
     @wp.func
     def edge_to_tri_coords(args: SideArg, side_index: ElementIndex, tri_index: ElementIndex, side_coords: Coords):
         edge_vidx = args.edge_vertex_indices[side_index]
-        tri_vidx = args.tri_vertex_indices[tri_index]
+        tri_vidx = args.cell_arg.tri_vertex_indices[tri_index]
 
         v0 = tri_vidx[0]
         v1 = tri_vidx[1]
@@ -238,9 +283,24 @@ class Trimesh2D(Geometry):
         return Coords(cx, cy, cz)
 
     @wp.func
-    def tri_to_edge_coords(args: SideArg, side_index: ElementIndex, tri_index: ElementIndex, tri_coords: Coords):
+    def side_inner_cell_coords(args: SideArg, side_index: ElementIndex, side_coords: Coords):
+        inner_cell_index = Trimesh2D.side_inner_cell_index(args, side_index)
+        return Trimesh2D.edge_to_tri_coords(args, 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)
+        return Trimesh2D.edge_to_tri_coords(args, 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,
+    ):
         edge_vidx = args.edge_vertex_indices[side_index]
-        tri_vidx = args.tri_vertex_indices[tri_index]
+        tri_vidx = args.cell_arg.tri_vertex_indices[tri_index]
 
         start = int(2)
         end = int(2)
@@ -256,31 +316,38 @@ class Trimesh2D(Geometry):
             tri_coords[start] + tri_coords[end] > 0.999, Coords(OUTSIDE), Coords(tri_coords[end], 0.0, 0.0)
         )
 
-    def _build_topology(self):
-        from warp.fem.utils import compress_node_indices, masked_indices, _get_pinned_temp_count_buffer
+    @wp.func
+    def side_to_cell_arg(side_arg: SideArg):
+        return side_arg.cell_arg
+
+    def _build_topology(self, temporary_store: TemporaryStore):
+        from warp.fem.utils import compress_node_indices, masked_indices
         from warp.utils import array_scan
 
         device = self.tri_vertex_indices.device
 
-        self._vertex_tri_offsets, self._vertex_tri_indices, _, __ = compress_node_indices(
-            self.vertex_count(), self.tri_vertex_indices
+        vertex_tri_offsets, vertex_tri_indices, _, __ = compress_node_indices(
+            self.vertex_count(), self.tri_vertex_indices, temporary_store=temporary_store
         )
+        self._vertex_tri_offsets = vertex_tri_offsets.detach()
+        self._vertex_tri_indices = vertex_tri_indices.detach()
 
-        vertex_start_edge_count = wp.zeros(dtype=int, device=device, shape=self.vertex_count())
-        vertex_start_edge_offsets = wp.empty_like(vertex_start_edge_count)
+        vertex_start_edge_count = borrow_temporary(temporary_store, dtype=int, device=device, shape=self.vertex_count())
+        vertex_start_edge_count.array.zero_()
+        vertex_start_edge_offsets = borrow_temporary_like(vertex_start_edge_count, temporary_store=temporary_store)
 
-        vertex_edge_ends = wp.empty(dtype=int, device=device, shape=(3 * self.cell_count()))
-        vertex_edge_tris = wp.empty(dtype=int, device=device, shape=(3 * self.cell_count(), 2))
+        vertex_edge_ends = borrow_temporary(temporary_store, dtype=int, device=device, shape=(3 * self.cell_count()))
+        vertex_edge_tris = borrow_temporary(temporary_store, dtype=int, device=device, shape=(3 * self.cell_count(), 2))
 
         # Count face edges starting at each vertex
         wp.launch(
             kernel=Trimesh2D._count_starting_edges_kernel,
             device=device,
             dim=self.cell_count(),
-            inputs=[self.tri_vertex_indices, vertex_start_edge_count],
+            inputs=[self.tri_vertex_indices, vertex_start_edge_count.array],
         )
 
-        array_scan(in_array=vertex_start_edge_count, out_array=vertex_start_edge_offsets, inclusive=False)
+        array_scan(in_array=vertex_start_edge_count.array, out_array=vertex_start_edge_offsets.array, inclusive=False)
 
         # Count number of unique edges (deduplicate across faces)
         vertex_unique_edge_count = vertex_start_edge_count
@@ -292,30 +359,32 @@ class Trimesh2D(Geometry):
                 self._vertex_tri_offsets,
                 self._vertex_tri_indices,
                 self.tri_vertex_indices,
-                vertex_start_edge_offsets,
-                vertex_unique_edge_count,
-                vertex_edge_ends,
-                vertex_edge_tris,
+                vertex_start_edge_offsets.array,
+                vertex_unique_edge_count.array,
+                vertex_edge_ends.array,
+                vertex_edge_tris.array,
             ],
         )
 
-        vertex_unique_edge_offsets = wp.empty_like(vertex_start_edge_offsets)
-        array_scan(in_array=vertex_start_edge_count, out_array=vertex_unique_edge_offsets, inclusive=False)
+        vertex_unique_edge_offsets = borrow_temporary_like(vertex_start_edge_offsets, temporary_store=temporary_store)
+        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 = _get_pinned_temp_count_buffer(device)
+            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, src=vertex_unique_edge_offsets, src_offset=self.vertex_count() - 1, count=1)
-            wp.synchronize_stream(wp.get_stream())
-            edge_count = int(edge_count.numpy()[0])
+            wp.copy(
+                dest=edge_count.array, src=vertex_unique_edge_offsets.array, src_offset=self.vertex_count() - 1, count=1
+            )
+            wp.synchronize_stream(wp.get_stream(device))
+            edge_count = int(edge_count.array.numpy()[0])
         else:
-            edge_count = int(vertex_unique_edge_offsets.numpy()[self.vertex_count() - 1])
+            edge_count = int(vertex_unique_edge_offsets.array.numpy()[self.vertex_count() - 1])
 
-        self._edge_vertex_indices = wp.empty(shape=(edge_count,), dtype=vec2i, device=device)
-        self._edge_tri_indices = wp.empty(shape=(edge_count,), dtype=vec2i, device=device)
+        self._edge_vertex_indices = wp.empty(shape=(edge_count,), dtype=wp.vec2i, device=device)
+        self._edge_tri_indices = wp.empty(shape=(edge_count,), dtype=wp.vec2i, device=device)
 
-        boundary_mask = wp.empty(shape=(edge_count,), dtype=int, device=device)
+        boundary_mask = borrow_temporary(temporary_store=temporary_store, shape=(edge_count,), dtype=int, device=device)
 
         # Compress edge data
         wp.launch(
@@ -323,17 +392,23 @@ class Trimesh2D(Geometry):
             device=device,
             dim=self.vertex_count(),
             inputs=[
-                vertex_start_edge_offsets,
-                vertex_unique_edge_offsets,
-                vertex_unique_edge_count,
-                vertex_edge_ends,
-                vertex_edge_tris,
+                vertex_start_edge_offsets.array,
+                vertex_unique_edge_offsets.array,
+                vertex_unique_edge_count.array,
+                vertex_edge_ends.array,
+                vertex_edge_tris.array,
                 self._edge_vertex_indices,
                 self._edge_tri_indices,
-                boundary_mask,
+                boundary_mask.array,
             ],
         )
 
+        vertex_start_edge_offsets.release()
+        vertex_unique_edge_offsets.release()
+        vertex_unique_edge_count.release()
+        vertex_edge_ends.release()
+        vertex_edge_tris.release()
+
         # Flip normals if necessary
         wp.launch(
             kernel=Trimesh2D._flip_edge_normals,
@@ -342,7 +417,20 @@ class Trimesh2D(Geometry):
             inputs=[self._edge_vertex_indices, self._edge_tri_indices, self.tri_vertex_indices, self.positions],
         )
 
-        self._boundary_edge_indices, _ = masked_indices(boundary_mask)
+        boundary_edge_indices, _ = masked_indices(boundary_mask.array, temporary_store=temporary_store)
+        self._boundary_edge_indices = boundary_edge_indices.detach()
+
+        boundary_mask.release()
+
+    def _compute_deformation_gradients(self):
+        self._deformation_gradients = wp.empty(dtype=wp.mat22f, device=self.positions.device, shape=(self.cell_count()))
+
+        wp.launch(
+            kernel=Trimesh2D._compute_deformation_gradients_kernel,
+            dim=self._deformation_gradients.shape,
+            device=self._deformation_gradients.device,
+            inputs=[self.tri_vertex_indices, self.positions, self._deformation_gradients],
+        )
 
     @wp.kernel
     def _count_starting_edges_kernel(
@@ -420,8 +508,8 @@ class Trimesh2D(Geometry):
         vertex_unique_edge_count: wp.array(dtype=int),
         uncompressed_edge_ends: wp.array(dtype=int),
         uncompressed_edge_tris: wp.array2d(dtype=int),
-        edge_vertex_indices: wp.array(dtype=vec2i),
-        edge_tri_indices: wp.array(dtype=vec2i),
+        edge_vertex_indices: wp.array(dtype=wp.vec2i),
+        edge_tri_indices: wp.array(dtype=wp.vec2i),
         boundary_mask: wp.array(dtype=int),
     ):
         v = wp.tid()
@@ -434,11 +522,11 @@ class Trimesh2D(Geometry):
             src_index = start_beg + e
             edge_index = unique_beg + e
 
-            edge_vertex_indices[edge_index] = vec2i(v, uncompressed_edge_ends[src_index])
+            edge_vertex_indices[edge_index] = wp.vec2i(v, uncompressed_edge_ends[src_index])
 
             t0 = uncompressed_edge_tris[src_index, 0]
             t1 = uncompressed_edge_tris[src_index, 1]
-            edge_tri_indices[edge_index] = vec2i(t0, t1)
+            edge_tri_indices[edge_index] = wp.vec2i(t0, t1)
             if t0 == t1:
                 boundary_mask[edge_index] = 1
             else:
@@ -446,8 +534,8 @@ class Trimesh2D(Geometry):
 
     @wp.kernel
     def _flip_edge_normals(
-        edge_vertex_indices: wp.array(dtype=vec2i),
-        edge_tri_indices: wp.array(dtype=vec2i),
+        edge_vertex_indices: wp.array(dtype=wp.vec2i),
+        edge_tri_indices: wp.array(dtype=wp.vec2i),
         tri_vertex_indices: wp.array2d(dtype=int),
         positions: wp.array(dtype=wp.vec2),
     ):
@@ -469,4 +557,21 @@ class Trimesh2D(Geometry):
 
         # if edge normal points toward first triangle centroid, flip indices
         if wp.dot(tri_centroid - edge_center, edge_normal) > 0.0:
-            edge_vertex_indices[e] = vec2i(edge_vidx[1], edge_vidx[0])
+            edge_vertex_indices[e] = wp.vec2i(edge_vidx[1], edge_vidx[0])
+
+    @wp.kernel
+    def _compute_deformation_gradients_kernel(
+        tri_vertex_indices: wp.array2d(dtype=int),
+        positions: wp.array(dtype=wp.vec2f),
+        transforms: wp.array(dtype=wp.mat22f),
+    ):
+        t = wp.tid()
+
+        p0 = positions[tri_vertex_indices[t, 0]]
+        p1 = positions[tri_vertex_indices[t, 1]]
+        p2 = positions[tri_vertex_indices[t, 2]]
+
+        e1 = p1 - p0
+        e2 = p2 - p0
+
+        transforms[t] = wp.mat22(e1, e2)