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

@@ -16,12 +16,15 @@
 from typing import ClassVar
 
 import warp as wp
-from warp.fem import cache
-from warp.fem.polynomial import Polynomial
-from warp.fem.types import Coords, ElementIndex, Sample, make_free_sample
+from warp._src.fem import cache
+from warp._src.fem.polynomial import Polynomial
+from warp._src.fem.types import Coords, ElementIndex, Sample, make_free_sample
+from warp._src.types import type_is_vector, type_size
 
 from .geometry import Geometry
 
+_wp_module_name_ = "warp.fem.geometry.deformed_geometry"
+
 _mat32 = wp.mat(shape=(3, 2), dtype=float)
 
 
@@ -54,16 +57,17 @@ class DeformedGeometry(Geometry):
         """Constructs a Deformed Geometry from a displacement or absolute position field defined over a base geometry.
         The deformation field does not need to be isoparameteric.
 
-        See also: :meth:`warp.fem.DiscreteField.make_deformed_geometry`
+        See also: :meth:`warp.fem.GeometryField.make_deformed_geometry`
         """
 
-        from warp.fem.field import DiscreteField, GeometryField
+        from warp._src.fem.field import GeometryField
 
-        if isinstance(field, DiscreteField):
-            if not wp.types.type_is_vector(field.dtype) or wp.types.type_size(field.dtype) != field.geometry.dimension:
-                raise ValueError(
-                    "Invalid value type for position field, must be vector-valued with same dimension as underlying geometry"
-                )
+        if field.dtype is not None and (
+            not type_is_vector(field.dtype) or type_size(field.dtype) != field.geometry.dimension
+        ):
+            raise ValueError(
+                "Invalid value type for position field, must be vector-valued with same dimension as underlying geometry"
+            )
         if field.eval_grad_inner is None:
             raise ValueError("Gradient evaluation is not supported on the passed field")
 
@@ -113,11 +117,6 @@ class DeformedGeometry(Geometry):
 
         return CellArg
 
-    def cell_arg_value(self, device) -> "DeformedGeometry.CellArg":
-        args = self.CellArg()
-        self.fill_cell_arg(args, device)
-        return args
-
     def fill_cell_arg(self, args: "DeformedGeometry.CellArg", device):
         self.base.fill_cell_arg(args.base_arg, device)
         self.field.fill_eval_arg(args.field_arg, device)
@@ -161,11 +160,6 @@ class DeformedGeometry(Geometry):
 
         return SideArg
 
-    def side_arg_value(self, device) -> "DeformedGeometry.SideArg":
-        args = self.SideArg()
-        self.fill_side_arg(args, device)
-        return args
-
     def fill_side_arg(self, args: "DeformedGeometry.SideArg", device):
         self.base.fill_side_arg(args.base_arg, device)
         self.field.fill_eval_arg(args.field_arg, device)
@@ -259,11 +253,11 @@ class DeformedGeometry(Geometry):
         return cell_bvh_id
 
     def _make_cell_bounds(self):
-        points, _weights = self.reference_cell().instantiate_quadrature(
+        points, _weights = self.reference_cell().prototype.instantiate_quadrature(
             order=self.field.degree, family=Polynomial.LOBATTO_GAUSS_LEGENDRE
         )
 
-        points = cache.cached_mat_type((len(points), 3), dtype=float)(points)
+        points = wp.matrix_from_rows(*points)
         point_count = len(points)
 
         @cache.dynamic_func(suffix=self.name)

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

@@ -13,16 +13,40 @@
 # 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)
+_wp_module_name_ = "warp.fem.geometry.element"
 
+_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 +100,7 @@ def _point_count_from_order(order: int, family: Polynomial):
     return point_count
 
 
-class Cube(Element):
+class Cube(PrototypeElement):
     dimension = 3
 
     @staticmethod
@@ -97,7 +121,7 @@ class Cube(Element):
         return coords, weights
 
 
-class Square(Element):
+class Square(PrototypeElement):
     dimension = 2
 
     @staticmethod
@@ -118,7 +142,7 @@ class Square(Element):
         return coords, weights
 
 
-class LinearEdge(Element):
+class LinearEdge(PrototypeElement):
     dimension = 1
 
     @staticmethod
@@ -137,7 +161,7 @@ class LinearEdge(Element):
         return coords, weights_1d
 
 
-class Triangle(Element):
+class Triangle(PrototypeElement):
     dimension = 2
 
     @staticmethod
@@ -492,7 +516,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 +851,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,11 +17,14 @@ 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
 
+_wp_module_name_ = "warp.fem.geometry.geometry"
+
 _mat32 = wp.mat(shape=(3, 2), dtype=float)
 
 _NULL_BVH_ID = wp.uint64(0)
@@ -66,7 +69,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 +83,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 +138,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 +295,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 +305,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 +326,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 +342,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 +371,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 +396,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 +433,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 +478,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 +662,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,16 +13,19 @@
 # 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
 
+_wp_module_name_ = "warp.fem.geometry.grid_2d"
+
 
 @wp.struct
 class Grid2DCellArg:
@@ -59,7 +62,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 +70,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 +90,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 +104,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 +186,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 +300,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,16 +13,19 @@
 # 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
 
+_wp_module_name_ = "warp.fem.geometry.grid_3d"
+
 
 @wp.struct
 class Grid3DCellArg:
@@ -56,7 +59,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 +68,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 +100,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 +114,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 +226,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 +345,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 +357,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"""