warp-lang 1.9.0__py3-none-manylinux_2_34_aarch64.whl → 1.10.0__py3-none-manylinux_2_34_aarch64.whl

warp/{fem → _src/fem}/geometry/tetmesh.py

@@ -16,13 +16,12 @@
 from typing import Optional
 
 import warp as wp
-from warp.fem.cache import (
+from warp._src.fem.cache import (
     TemporaryStore,
     borrow_temporary,
     borrow_temporary_like,
-    cached_arg_value,
 )
-from warp.fem.types import (
+from warp._src.fem.types import (
     OUTSIDE,
     Coords,
     ElementIndex,
@@ -30,9 +29,11 @@ from warp.fem.types import (
 )
 
 from .closest_point import project_on_tet_at_origin, project_on_tri_at_origin
-from .element import Tetrahedron, Triangle
+from .element import Element
 from .geometry import Geometry
 
+_wp_module_name_ = "warp.fem.geometry.tetmesh"
+
 
 @wp.struct
 class TetmeshCellArg:
@@ -108,11 +109,11 @@ class Tetmesh(Geometry):
     def boundary_side_count(self):
         return self._boundary_face_indices.shape[0]
 
-    def reference_cell(self) -> Tetrahedron:
-        return Tetrahedron()
+    def reference_cell(self) -> Element:
+        return Element.TETRAHEDRON
 
-    def reference_side(self) -> Triangle:
-        return Triangle()
+    def reference_side(self) -> Element:
+        return Element.TRIANGLE
 
     @property
     def tet_edge_indices(self) -> wp.array:
@@ -137,11 +138,6 @@ class Tetmesh(Geometry):
 
     # Geometry device interface
 
-    def cell_arg_value(self, device) -> CellArg:
-        args = self.CellArg()
-        self.fill_cell_arg(args, device)
-        return args
-
     def fill_cell_arg(self, args: CellArg, device):
         args.tet_vertex_indices = self.tet_vertex_indices.to(device)
         args.positions = self.positions.to(device)
@@ -183,12 +179,6 @@ class Tetmesh(Geometry):
         dist, coords = project_on_tet_at_origin(q, e1, e2, e3)
         return coords, dist
 
-    @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_face_indices = self._boundary_face_indices.to(device)
 
@@ -198,11 +188,6 @@ class Tetmesh(Geometry):
 
         return args.boundary_face_indices[boundary_side_index]
 
-    def side_arg_value(self, device) -> CellArg:
-        args = self.SideArg()
-        self.fill_side_arg(args, device)
-        return args
-
     def fill_side_arg(self, args: SideArg, device):
         self.fill_cell_arg(args.cell_arg, device)
         args.face_vertex_indices = self._face_vertex_indices.to(device)
@@ -325,8 +310,8 @@ class Tetmesh(Geometry):
         return side_arg.cell_arg
 
     def _build_topology(self, temporary_store: TemporaryStore):
-        from warp.fem.utils import compress_node_indices, host_read_at_index, masked_indices
-        from warp.utils import array_scan
+        from warp._src.fem.utils import compress_node_indices, host_read_at_index, masked_indices
+        from warp._src.utils import array_scan
 
         device = self.tet_vertex_indices.device
 
@@ -337,7 +322,7 @@ class Tetmesh(Geometry):
         self._vertex_tet_indices = vertex_tet_indices.detach()
 
         vertex_start_face_count = borrow_temporary(temporary_store, dtype=int, device=device, shape=self.vertex_count())
-        vertex_start_face_count.array.zero_()
+        vertex_start_face_count.zero_()
         vertex_start_face_offsets = borrow_temporary_like(vertex_start_face_count, temporary_store=temporary_store)
 
         vertex_face_other_vs = borrow_temporary(
@@ -350,10 +335,10 @@ class Tetmesh(Geometry):
             kernel=Tetmesh._count_starting_faces_kernel,
             device=device,
             dim=self.cell_count(),
-            inputs=[self.tet_vertex_indices, vertex_start_face_count.array],
+            inputs=[self.tet_vertex_indices, vertex_start_face_count],
         )
 
-        array_scan(in_array=vertex_start_face_count.array, out_array=vertex_start_face_offsets.array, inclusive=False)
+        array_scan(in_array=vertex_start_face_count, out_array=vertex_start_face_offsets, inclusive=False)
 
         # Count number of unique edges (deduplicate across faces)
         vertex_unique_face_count = vertex_start_face_count
@@ -365,21 +350,19 @@ class Tetmesh(Geometry):
                 self._vertex_tet_offsets,
                 self._vertex_tet_indices,
                 self.tet_vertex_indices,
-                vertex_start_face_offsets.array,
-                vertex_unique_face_count.array,
-                vertex_face_other_vs.array,
-                vertex_face_tets.array,
+                vertex_start_face_offsets,
+                vertex_unique_face_count,
+                vertex_face_other_vs,
+                vertex_face_tets,
             ],
         )
 
         vertex_unique_face_offsets = borrow_temporary_like(vertex_start_face_offsets, temporary_store=temporary_store)
-        array_scan(in_array=vertex_start_face_count.array, out_array=vertex_unique_face_offsets.array, inclusive=False)
+        array_scan(in_array=vertex_start_face_count, out_array=vertex_unique_face_offsets, inclusive=False)
 
         # Get back edge count to host
         face_count = int(
-            host_read_at_index(
-                vertex_unique_face_offsets.array, self.vertex_count() - 1, temporary_store=temporary_store
-            )
+            host_read_at_index(vertex_unique_face_offsets, self.vertex_count() - 1, temporary_store=temporary_store)
         )
 
         self._face_vertex_indices = wp.empty(shape=(face_count,), dtype=wp.vec3i, device=device)
@@ -393,14 +376,14 @@ class Tetmesh(Geometry):
             device=device,
             dim=self.vertex_count(),
             inputs=[
-                vertex_start_face_offsets.array,
-                vertex_unique_face_offsets.array,
-                vertex_unique_face_count.array,
-                vertex_face_other_vs.array,
-                vertex_face_tets.array,
+                vertex_start_face_offsets,
+                vertex_unique_face_offsets,
+                vertex_unique_face_count,
+                vertex_face_other_vs,
+                vertex_face_tets,
                 self._face_vertex_indices,
                 self._face_tet_indices,
-                boundary_mask.array,
+                boundary_mask,
             ],
         )
 
@@ -418,17 +401,17 @@ class Tetmesh(Geometry):
             inputs=[self._face_vertex_indices, self._face_tet_indices, self.tet_vertex_indices, self.positions],
         )
 
-        boundary_face_indices, _ = masked_indices(boundary_mask.array)
+        boundary_face_indices, _ = masked_indices(boundary_mask)
         self._boundary_face_indices = boundary_face_indices.detach()
 
     def _compute_tet_edges(self, temporary_store: Optional[TemporaryStore] = None):
-        from warp.fem.utils import host_read_at_index
-        from warp.utils import array_scan
+        from warp._src.fem.utils import host_read_at_index
+        from warp._src.utils import array_scan
 
         device = self.tet_vertex_indices.device
 
         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_count.zero_()
         vertex_start_edge_offsets = borrow_temporary_like(vertex_start_edge_count, temporary_store=temporary_store)
 
         vertex_edge_ends = borrow_temporary(temporary_store, dtype=int, device=device, shape=(6 * self.cell_count()))
@@ -438,10 +421,10 @@ class Tetmesh(Geometry):
             kernel=Tetmesh._count_starting_edges_kernel,
             device=device,
             dim=self.cell_count(),
-            inputs=[self.tet_vertex_indices, vertex_start_edge_count.array],
+            inputs=[self.tet_vertex_indices, vertex_start_edge_count],
         )
 
-        array_scan(in_array=vertex_start_edge_count.array, out_array=vertex_start_edge_offsets.array, inclusive=False)
+        array_scan(in_array=vertex_start_edge_count, out_array=vertex_start_edge_offsets, inclusive=False)
 
         # Count number of unique edges (deduplicate across faces)
         vertex_unique_edge_count = vertex_start_edge_count
@@ -453,22 +436,18 @@ class Tetmesh(Geometry):
                 self._vertex_tet_offsets,
                 self._vertex_tet_indices,
                 self.tet_vertex_indices,
-                vertex_start_edge_offsets.array,
-                vertex_unique_edge_count.array,
-                vertex_edge_ends.array,
+                vertex_start_edge_offsets,
+                vertex_unique_edge_count,
+                vertex_edge_ends,
             ],
         )
 
-        vertex_unique_edge_offsets = borrow_temporary_like(
-            vertex_start_edge_offsets.array, temporary_store=temporary_store
-        )
-        array_scan(in_array=vertex_start_edge_count.array, out_array=vertex_unique_edge_offsets.array, 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, out_array=vertex_unique_edge_offsets, inclusive=False)
 
         # Get back edge count to host
         self._edge_count = int(
-            host_read_at_index(
-                vertex_unique_edge_offsets.array, self.vertex_count() - 1, temporary_store=temporary_store
-            )
+            host_read_at_index(vertex_unique_edge_offsets, self.vertex_count() - 1, temporary_store=temporary_store)
         )
 
         self._tet_edge_indices = wp.empty(
@@ -484,10 +463,10 @@ class Tetmesh(Geometry):
                 self._vertex_tet_offsets,
                 self._vertex_tet_indices,
                 self.tet_vertex_indices,
-                vertex_start_edge_offsets.array,
-                vertex_unique_edge_offsets.array,
-                vertex_unique_edge_count.array,
-                vertex_edge_ends.array,
+                vertex_start_edge_offsets,
+                vertex_unique_edge_offsets,
+                vertex_unique_edge_count,
+                vertex_edge_ends,
                 self._tet_edge_indices,
             ],
         )

warp/{fem → _src/fem}/geometry/trimesh.py

@@ -16,13 +16,12 @@
 from typing import Any, Optional
 
 import warp as wp
-from warp.fem.cache import (
+from warp._src.fem.cache import (
     TemporaryStore,
     borrow_temporary,
     borrow_temporary_like,
-    cached_arg_value,
 )
-from warp.fem.types import (
+from warp._src.fem.types import (
     OUTSIDE,
     Coords,
     ElementIndex,
@@ -30,9 +29,11 @@ from warp.fem.types import (
 )
 
 from .closest_point import project_on_seg_at_origin, project_on_tri_at_origin
-from .element import LinearEdge, Triangle
+from .element import Element
 from .geometry import Geometry
 
+_wp_module_name_ = "warp.fem.geometry.trimesh"
+
 
 @wp.struct
 class TrimeshCellArg:
@@ -103,11 +104,11 @@ class Trimesh(Geometry):
     def boundary_side_count(self):
         return self._boundary_edge_indices.shape[0]
 
-    def reference_cell(self) -> Triangle:
-        return Triangle()
+    def reference_cell(self) -> Element:
+        return Element.TRIANGLE
 
-    def reference_side(self) -> LinearEdge:
-        return LinearEdge()
+    def reference_side(self) -> Element:
+        return Element.LINE_SEGMENT
 
     @property
     def edge_tri_indices(self) -> wp.array:
@@ -130,30 +131,14 @@ class Trimesh(Geometry):
         args.edge_vertex_indices = self._edge_vertex_indices.to(device)
         args.edge_tri_indices = self._edge_tri_indices.to(device)
 
-    def cell_arg_value(self, device):
-        args = self.CellArg()
-        self.fill_cell_arg(args, device)
-        return args
-
     def fill_cell_arg(self, args: TrimeshCellArg, device):
         self._fill_cell_topo_arg(args.topology, device)
         args.positions = self.positions.to(device)
 
-    def side_arg_value(self, device):
-        args = self.SideArg()
-        self.fill_side_arg(args, device)
-        return args
-
     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)
 
@@ -216,8 +201,8 @@ class Trimesh(Geometry):
         return wp.where(tri_coords[start] + tri_coords[end] > 0.999, Coords(tri_coords[end], 0.0, 0.0), Coords(OUTSIDE))
 
     def _build_topology(self, temporary_store: TemporaryStore):
-        from warp.fem.utils import compress_node_indices, host_read_at_index, masked_indices
-        from warp.utils import array_scan
+        from warp._src.fem.utils import compress_node_indices, host_read_at_index, masked_indices
+        from warp._src.utils import array_scan
 
         device = self.tri_vertex_indices.device
 
@@ -228,7 +213,7 @@ class Trimesh(Geometry):
         self._vertex_tri_indices = vertex_tri_indices.detach()
 
         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_count.zero_()
         vertex_start_edge_offsets = borrow_temporary_like(vertex_start_edge_count, temporary_store=temporary_store)
 
         vertex_edge_ends = borrow_temporary(temporary_store, dtype=int, device=device, shape=(3 * self.cell_count()))
@@ -239,10 +224,10 @@ class Trimesh(Geometry):
             kernel=Trimesh._count_starting_edges_kernel,
             device=device,
             dim=self.cell_count(),
-            inputs=[self.tri_vertex_indices, vertex_start_edge_count.array],
+            inputs=[self.tri_vertex_indices, vertex_start_edge_count],
         )
 
-        array_scan(in_array=vertex_start_edge_count.array, out_array=vertex_start_edge_offsets.array, inclusive=False)
+        array_scan(in_array=vertex_start_edge_count, out_array=vertex_start_edge_offsets, inclusive=False)
 
         # Count number of unique edges (deduplicate across faces)
         vertex_unique_edge_count = vertex_start_edge_count
@@ -254,21 +239,19 @@ class Trimesh(Geometry):
                 self._vertex_tri_offsets,
                 self._vertex_tri_indices,
                 self.tri_vertex_indices,
-                vertex_start_edge_offsets.array,
-                vertex_unique_edge_count.array,
-                vertex_edge_ends.array,
-                vertex_edge_tris.array,
+                vertex_start_edge_offsets,
+                vertex_unique_edge_count,
+                vertex_edge_ends,
+                vertex_edge_tris,
             ],
         )
 
         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)
+        array_scan(in_array=vertex_start_edge_count, out_array=vertex_unique_edge_offsets, inclusive=False)
 
         # Get back edge count to host
         edge_count = int(
-            host_read_at_index(
-                vertex_unique_edge_offsets.array, self.vertex_count() - 1, temporary_store=temporary_store
-            )
+            host_read_at_index(vertex_unique_edge_offsets, self.vertex_count() - 1, temporary_store=temporary_store)
         )
 
         self._edge_vertex_indices = wp.empty(shape=(edge_count,), dtype=wp.vec2i, device=device)
@@ -282,14 +265,14 @@ class Trimesh(Geometry):
             device=device,
             dim=self.vertex_count(),
             inputs=[
-                vertex_start_edge_offsets.array,
-                vertex_unique_edge_offsets.array,
-                vertex_unique_edge_count.array,
-                vertex_edge_ends.array,
-                vertex_edge_tris.array,
+                vertex_start_edge_offsets,
+                vertex_unique_edge_offsets,
+                vertex_unique_edge_count,
+                vertex_edge_ends,
+                vertex_edge_tris,
                 self._edge_vertex_indices,
                 self._edge_tri_indices,
-                boundary_mask.array,
+                boundary_mask,
             ],
         )
 
@@ -299,7 +282,7 @@ class Trimesh(Geometry):
         vertex_edge_ends.release()
         vertex_edge_tris.release()
 
-        boundary_edge_indices, _ = masked_indices(boundary_mask.array, temporary_store=temporary_store)
+        boundary_edge_indices, _ = masked_indices(boundary_mask, temporary_store=temporary_store)
         self._boundary_edge_indices = boundary_edge_indices.detach()
 
         boundary_mask.release()
@@ -467,7 +450,7 @@ class Trimesh(Geometry):
         q = pos - p0
         e = args.positions[edge_idx[1]] - p0
 
-        dist, t = project_on_seg_at_origin(q, e, wp.lengh_sq(e))
+        dist, t = project_on_seg_at_origin(q, e, wp.length_sq(e))
         return Coords(t, 0.0, 0.0), dist
 
     @wp.func