warp-lang 1.9.1__py3-none-manylinux_2_34_aarch64.whl → 1.10.0rc2__py3-none-manylinux_2_34_aarch64.whl

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

@@ -13,16 +13,38 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from enum import IntEnum
 from typing import List, Tuple
 
 import warp as wp
-from warp.fem.polynomial import Polynomial, quadrature_1d
-from warp.fem.types import Coords
+from warp._src.fem.polynomial import Polynomial, quadrature_1d
+from warp._src.fem.types import Coords
 
-_vec1 = wp.types.vector(length=1, dtype=float)
+_vec1 = wp.vec(length=1, dtype=float)
 
 
-class Element:
+class Element(IntEnum):
+    """Enumeration of reference element types"""
+
+    LINE_SEGMENT = 1
+    SQUARE = 2
+    CUBE = 3
+    TRIANGLE = 4
+    TETRAHEDRON = 5
+
+    @property
+    def prototype(self) -> "PrototypeElement":
+        """Prototype element for the given element type"""
+        return {
+            Element.LINE_SEGMENT: LinearEdge,
+            Element.SQUARE: Square,
+            Element.CUBE: Cube,
+            Element.TRIANGLE: Triangle,
+            Element.TETRAHEDRON: Tetrahedron,
+        }[self]
+
+
+class PrototypeElement:
     dimension = 0
     """Intrinsic dimension of the element"""
 
@@ -76,7 +98,7 @@ def _point_count_from_order(order: int, family: Polynomial):
     return point_count
 
 
-class Cube(Element):
+class Cube(PrototypeElement):
     dimension = 3
 
     @staticmethod
@@ -97,7 +119,7 @@ class Cube(Element):
         return coords, weights
 
 
-class Square(Element):
+class Square(PrototypeElement):
     dimension = 2
 
     @staticmethod
@@ -118,7 +140,7 @@ class Square(Element):
         return coords, weights
 
 
-class LinearEdge(Element):
+class LinearEdge(PrototypeElement):
     dimension = 1
 
     @staticmethod
@@ -137,7 +159,7 @@ class LinearEdge(Element):
         return coords, weights_1d
 
 
-class Triangle(Element):
+class Triangle(PrototypeElement):
     dimension = 2
 
     @staticmethod
@@ -492,7 +514,7 @@ class Triangle(Element):
         return Coords(-ref_delta[0] - ref_delta[1], ref_delta[0], ref_delta[1])
 
 
-class Tetrahedron(Element):
+class Tetrahedron(PrototypeElement):
     dimension = 3
 
     @staticmethod
@@ -827,4 +849,4 @@ class Tetrahedron(Element):
         c = c - Coords(n, 0.0, n)
 
         # project on cube
-        return Element.project(c)
+        return PrototypeElement.project(c)

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

@@ -17,8 +17,9 @@ from functools import cached_property
 from typing import Any, ClassVar
 
 import warp as wp
-from warp.fem import cache
-from warp.fem.types import NULL_ELEMENT_INDEX, OUTSIDE, Coords, ElementIndex, ElementKind, Sample, make_free_sample
+from warp._src.codegen import Struct
+from warp._src.fem import cache
+from warp._src.fem.types import NULL_ELEMENT_INDEX, OUTSIDE, Coords, ElementIndex, ElementKind, Sample, make_free_sample
 
 from .element import Element
 
@@ -66,7 +67,7 @@ class Geometry:
     @property
     def cell_dimension(self) -> int:
         """Manifold dimension of the geometry cells"""
-        return self.reference_cell().dimension
+        return self.reference_cell().prototype.dimension
 
     @property
     def base(self) -> "Geometry":
@@ -80,22 +81,27 @@ class Geometry:
     def __str__(self) -> str:
         return self.name
 
-    CellArg: wp.codegen.Struct
+    CellArg: Struct
     """Structure containing arguments to be passed to device functions evaluating cell-related quantities"""
 
-    SideArg: wp.codegen.Struct
+    SideArg: Struct
     """Structure containing arguments to be passed to device functions evaluating side-related quantities"""
 
-    SideIndexArg: wp.codegen.Struct
+    SideIndexArg: Struct
     """Structure containing arguments to be passed to device functions for indexing sides"""
 
+    @cache.cached_arg_value
     def cell_arg_value(self, device) -> "Geometry.CellArg":
         """Value of the arguments to be passed to cell-related device functions"""
-        raise NotImplementedError
+        args = self.CellArg()
+        self.fill_cell_arg(args, device)
+        return args
 
     def fill_cell_arg(self, args: "Geometry.CellArg", device):
         """Fill the arguments to be passed to cell-related device functions"""
-        raise NotImplementedError
+        if self.cell_arg_value is __class__.cell_arg_value:
+            raise NotImplementedError()
+        args.assign(self.cell_arg_value(device))
 
     @staticmethod
     def cell_position(args: "Geometry.CellArg", s: "Sample"):
@@ -130,13 +136,31 @@ class Geometry:
         For elements with the same dimension as the embedding space, this will be zero."""
         raise NotImplementedError
 
+    @cache.cached_arg_value
     def side_arg_value(self, device) -> "Geometry.SideArg":
         """Value of the arguments to be passed to side-related device functions"""
-        raise NotImplementedError
+        args = self.SideArg()
+        self.fill_side_arg(args, device)
+        return args
 
     def fill_side_arg(self, args: "Geometry.SideArg", device):
         """Fill the arguments to be passed to side-related device functions"""
-        raise NotImplementedError
+        if self.side_arg_value is __class__.side_arg_value:
+            raise NotImplementedError()
+        args.assign(self.side_arg_value(device))
+
+    @cache.cached_arg_value
+    def side_index_arg_value(self, device) -> "Geometry.SideIndexArg":
+        """Value of the arguments to be passed to side-related device functions"""
+        args = self.SideIndexArg()
+        self.fill_side_index_arg(args, device)
+        return args
+
+    def fill_side_index_arg(self, args: "Geometry.SideIndexArg", device):
+        """Fill the arguments to be passed to side-related device functions"""
+        if self.side_index_arg_value is __class__.side_index_arg_value:
+            raise NotImplementedError()
+        args.assign(self.side_index_arg_value(device))
 
     @staticmethod
     def boundary_side_index(args: "Geometry.SideIndexArg", boundary_side_index: int):
@@ -269,7 +293,7 @@ class Geometry:
         return wp.normalize(Fcross)
 
     def _make_cell_measure(self):
-        REF_MEASURE = wp.constant(self.reference_cell().measure())
+        REF_MEASURE = wp.constant(self.reference_cell().prototype.measure())
 
         @cache.dynamic_func(suffix=self.name)
         def cell_measure(args: self.CellArg, s: Sample):
@@ -279,7 +303,7 @@ class Geometry:
         return cell_measure
 
     def _make_cell_normal(self):
-        cell_dim = self.reference_cell().dimension
+        cell_dim = self.reference_cell().prototype.dimension
         geo_dim = self.dimension
         normal_vec = wp.vec(length=geo_dim, dtype=float)
 
@@ -300,7 +324,7 @@ class Geometry:
         return None
 
     def _make_cell_inverse_deformation_gradient(self):
-        cell_dim = self.reference_cell().dimension
+        cell_dim = self.reference_cell().prototype.dimension
         geo_dim = self.dimension
 
         @cache.dynamic_func(suffix=self.name)
@@ -316,7 +340,7 @@ class Geometry:
         return cell_inverse_deformation_gradient if cell_dim == geo_dim else cell_pseudoinverse_deformation_gradient
 
     def _make_side_inverse_deformation_gradient(self):
-        side_dim = self.reference_side().dimension
+        side_dim = self.reference_side().prototype.dimension
         geo_dim = self.dimension
 
         if side_dim == geo_dim:
@@ -345,7 +369,7 @@ class Geometry:
         return side_pseudoinverse_deformation_gradient
 
     def _make_side_measure(self):
-        REF_MEASURE = wp.constant(self.reference_side().measure())
+        REF_MEASURE = wp.constant(self.reference_side().prototype.measure())
 
         @cache.dynamic_func(suffix=self.name)
         def side_measure(args: self.SideArg, s: Sample):
@@ -370,7 +394,7 @@ class Geometry:
         return side_measure_ratio
 
     def _make_side_normal(self):
-        side_dim = self.reference_side().dimension
+        side_dim = self.reference_side().prototype.dimension
         geo_dim = self.dimension
 
         @cache.dynamic_func(suffix=self.name)
@@ -407,12 +431,12 @@ class Geometry:
         pos_type = cache.cached_vec_type(self.dimension, dtype=float)
 
         if element_kind == ElementKind.CELL:
-            ref_elt = self.reference_cell()
+            ref_elt = self.reference_cell().prototype
             arg_type = self.CellArg
             elt_pos = self.cell_position
             elt_inv_grad = self.cell_inverse_deformation_gradient
         else:
-            ref_elt = self.reference_side()
+            ref_elt = self.reference_side().prototype
             arg_type = self.SideArg
             elt_pos = self.side_position
             elt_inv_grad = self.side_inverse_deformation_gradient
@@ -452,12 +476,12 @@ class Geometry:
         element_coordinates = self._make_element_coordinates(element_kind=element_kind, assume_linear=assume_linear)
 
         if element_kind == ElementKind.CELL:
-            ref_elt = self.reference_cell()
+            ref_elt = self.reference_cell().prototype
             arg_type = self.CellArg
             elt_pos = self.cell_position
             elt_def_grad = self.cell_deformation_gradient
         else:
-            ref_elt = self.reference_side()
+            ref_elt = self.reference_side().prototype
             arg_type = self.SideArg
             elt_pos = self.side_position
             elt_def_grad = self.side_deformation_gradient
@@ -636,8 +660,12 @@ class Geometry:
 
         if self._bvhs is None:
             self._bvhs = {}
+
         self._bvhs[device.ordinal] = wp.Bvh(lowers, uppers)
 
+        Geometry.cell_arg_value.invalidate(self, device)
+        Geometry.side_arg_value.invalidate(self, device)
+
     def bvh_id(self, device):
         if self._bvhs is None:
             return _NULL_BVH_ID

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

@@ -13,14 +13,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from functools import cached_property
 from typing import Any, Optional
 
 import warp as wp
-from warp.fem.cache import cached_arg_value, dynamic_func
-from warp.fem.types import NULL_ELEMENT_INDEX, OUTSIDE, Coords, ElementIndex, Sample, make_free_sample
+from warp._src.fem.cache import cached_arg_value, dynamic_func
+from warp._src.fem.types import NULL_ELEMENT_INDEX, OUTSIDE, Coords, ElementIndex, Sample, make_free_sample
 
 from .closest_point import project_on_box_at_origin
-from .element import LinearEdge, Square
+from .element import Element
 from .geometry import Geometry
 
 
@@ -59,7 +60,7 @@ class Grid2D(Geometry):
 
         self._res = res
 
-    @property
+    @cached_property
     def extents(self) -> wp.vec3:
         # Avoid using native sub due to higher over of calling builtins from Python
         return wp.vec2(
@@ -67,7 +68,7 @@ class Grid2D(Geometry):
             self.bounds_hi[1] - self.bounds_lo[1],
         )
 
-    @property
+    @cached_property
     def cell_size(self) -> wp.vec2:
         ex = self.extents
         return wp.vec2(
@@ -87,11 +88,11 @@ class Grid2D(Geometry):
     def boundary_side_count(self):
         return 2 * (self.res[0] + self.res[1])
 
-    def reference_cell(self) -> Square:
-        return Square()
+    def reference_cell(self) -> Element:
+        return Element.SQUARE
 
-    def reference_side(self) -> LinearEdge:
-        return LinearEdge()
+    def reference_side(self) -> Element:
+        return Element.LINE_SEGMENT
 
     @property
     def res(self):
@@ -101,7 +102,7 @@ class Grid2D(Geometry):
     def origin(self):
         return self.bounds_lo
 
-    @property
+    @cached_property
     def strides(self):
         return wp.vec2i(self.res[1], 1)
 
@@ -183,12 +184,6 @@ class Grid2D(Geometry):
 
     # Geometry device interface
 
-    @cached_arg_value
-    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.res = self.res
         args.cell_size = self.cell_size
@@ -303,23 +298,17 @@ class Grid2D(Geometry):
     @cached_arg_value
     def side_arg_value(self, device) -> SideArg:
         args = self.SideArg()
-        self.fill_side_arg(args, device)
-        return args
-
-    def fill_side_arg(self, args: SideArg, device):
         args.axis_offsets = wp.vec2i(
             0,
             self.res[1],
         )
         args.cell_count = self.cell_count()
         args.cell_arg = self.cell_arg_value(device)
+        return args
 
     def side_index_arg_value(self, device) -> SideIndexArg:
         return self.side_arg_value(device)
 
-    def fill_side_index_arg(self, args: SideIndexArg, device):
-        self.fill_side_arg(args, device)
-
     @wp.func
     def boundary_side_index(args: SideArg, boundary_side_index: int):
         """Boundary side to side index"""

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

@@ -13,14 +13,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from functools import cached_property
 from typing import Any, Optional
 
 import warp as wp
-from warp.fem.cache import cached_arg_value, dynamic_func
-from warp.fem.types import NULL_ELEMENT_INDEX, OUTSIDE, Coords, ElementIndex, Sample, make_free_sample
+from warp._src.fem.cache import cached_arg_value, dynamic_func
+from warp._src.fem.types import NULL_ELEMENT_INDEX, OUTSIDE, Coords, ElementIndex, Sample, make_free_sample
 
 from .closest_point import project_on_box_at_origin
-from .element import Cube, Square
+from .element import Element
 from .geometry import Geometry
 
 
@@ -56,7 +57,7 @@ class Grid3D(Geometry):
 
         self._res = res
 
-    @property
+    @cached_property
     def extents(self) -> wp.vec3:
         # Avoid using native sub due to higher over of calling builtins from Python
         return wp.vec3(
@@ -65,7 +66,7 @@ class Grid3D(Geometry):
             self.bounds_hi[2] - self.bounds_lo[2],
         )
 
-    @property
+    @cached_property
     def cell_size(self) -> wp.vec3:
         ex = self.extents
         return wp.vec3(
@@ -97,11 +98,11 @@ class Grid3D(Geometry):
     def boundary_side_count(self):
         return 2 * (self.res[1]) * (self.res[2]) + (self.res[0]) * 2 * (self.res[2]) + (self.res[0]) * (self.res[1]) * 2
 
-    def reference_cell(self) -> Cube:
-        return Cube()
+    def reference_cell(self) -> Element:
+        return Element.CUBE
 
-    def reference_side(self) -> Square:
-        return Square()
+    def reference_side(self) -> Element:
+        return Element.SQUARE
 
     @property
     def res(self):
@@ -111,7 +112,7 @@ class Grid3D(Geometry):
     def origin(self):
         return self.bounds_lo
 
-    @property
+    @cached_property
     def strides(self):
         return wp.vec3i(self.res[1] * self.res[2], self.res[2], 1)
 
@@ -223,12 +224,6 @@ class Grid3D(Geometry):
 
     # Geometry device interface
 
-    @cached_arg_value
-    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.res = self.res
         args.origin = self.bounds_lo
@@ -348,10 +343,6 @@ class Grid3D(Geometry):
     @cached_arg_value
     def side_arg_value(self, device) -> SideArg:
         args = self.SideArg()
-        self.fill_side_arg(args, device)
-        return args
-
-    def fill_side_arg(self, args: SideArg, device):
         axis_dims = wp.vec3i(
             self.res[1] * self.res[2],
             self.res[2] * self.res[0],
@@ -364,13 +355,11 @@ class Grid3D(Geometry):
         )
         args.cell_count = self.cell_count()
         args.cell_arg = self.cell_arg_value(device)
+        return args
 
     def side_index_arg_value(self, device) -> SideIndexArg:
         return self.side_arg_value(device)
 
-    def fill_side_index_arg(self, args: SideIndexArg, device):
-        self.fill_side_arg(args, device)
-
     @wp.func
     def boundary_side_index(args: SideArg, boundary_side_index: int):
         """Boundary side to side index"""