warp-lang 1.7.2rc1__py3-none-manylinux_2_34_aarch64.whl → 1.8.1__py3-none-manylinux_2_34_aarch64.whl

warp/fem/geometry/deformed_geometry.py

@@ -13,9 +13,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import ClassVar
+
 import warp as wp
 from warp.fem import cache
-from warp.fem.types import Coords, ElementIndex, Sample
+from warp.fem.polynomial import Polynomial
+from warp.fem.types import Coords, ElementIndex, Sample, make_free_sample
 
 from .geometry import Geometry
 
@@ -23,7 +26,31 @@ _mat32 = wp.mat(shape=(3, 2), dtype=float)
 
 
 class DeformedGeometry(Geometry):
-    def __init__(self, field: "wp.fem.field.GeometryField", relative: bool = True):
+    _dynamic_attribute_constructors_phase_1: ClassVar = {
+        "CellArg": lambda obj: obj._make_cell_arg(),
+        "SideArg": lambda obj: obj._make_side_arg(),
+        "cell_position": lambda obj: obj._make_cell_position(),
+        "cell_deformation_gradient": lambda obj: obj._make_cell_deformation_gradient(),
+        "side_to_cell_arg": lambda obj: obj._make_side_to_cell_arg(),
+        "side_position": lambda obj: obj._make_side_position(),
+        "side_deformation_gradient": lambda obj: obj._make_side_deformation_gradient(),
+        "side_inner_cell_index": lambda obj: obj._make_side_inner_cell_index(),
+        "side_outer_cell_index": lambda obj: obj._make_side_outer_cell_index(),
+        "side_inner_cell_coords": lambda obj: obj._make_side_inner_cell_coords(),
+        "side_outer_cell_coords": lambda obj: obj._make_side_outer_cell_coords(),
+        "side_from_cell_coords": lambda obj: obj._make_side_from_cell_coords(),
+        "cell_bvh_id": lambda obj: obj._make_cell_bvh_id(),
+        "cell_bounds": lambda obj: obj._make_cell_bounds(),
+    }
+
+    _dynamic_attribute_constructors_phase_2: ClassVar = {
+        "cell_closest_point": lambda obj: obj._make_cell_closest_point(),
+        "side_closest_point": lambda obj: obj._make_side_closest_point(),
+        "cell_coordinates": lambda obj: obj._make_cell_coordinates(),
+        "side_coordinates": lambda obj: obj._make_side_coordinates(),
+    }
+
+    def __init__(self, field: "wp.fem.field.GeometryField", relative: bool = True, build_bvh: bool = False):
         """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.
 
@@ -33,10 +60,7 @@ class DeformedGeometry(Geometry):
         from warp.fem.field import DiscreteField, GeometryField
 
         if isinstance(field, DiscreteField):
-            if (
-                not wp.types.type_is_vector(field.dtype)
-                or wp.types.type_length(field.dtype) != field.geometry.dimension
-            ):
+            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"
                 )
@@ -46,14 +70,10 @@ class DeformedGeometry(Geometry):
         self._relative = relative
 
         self.field: GeometryField = field
+        self.field_trace = field.trace()
         self.dimension = self.base.dimension
 
-        self.CellArg = self.field.ElementEvalArg
-
-        self.field_trace = field.trace()
-        self.SideArg = self._make_side_arg()
         self.SideIndexArg = self.base.SideIndexArg
-
         self.cell_count = self.base.cell_count
         self.vertex_count = self.base.vertex_count
         self.side_count = self.base.side_count
@@ -62,23 +82,18 @@ class DeformedGeometry(Geometry):
         self.reference_side = self.base.reference_side
 
         self.side_index_arg_value = self.base.side_index_arg_value
-
-        self.cell_position = self._make_cell_position()
-        self.cell_deformation_gradient = self._make_cell_deformation_gradient()
-
+        self.fill_side_index_arg = self.base.fill_side_index_arg
         self.boundary_side_index = self.base.boundary_side_index
 
-        self.side_to_cell_arg = self._make_side_to_cell_arg()
-        self.side_position = self._make_side_position()
-        self.side_deformation_gradient = self._make_side_deformation_gradient()
-        self.side_inner_cell_index = self._make_side_inner_cell_index()
-        self.side_outer_cell_index = self._make_side_outer_cell_index()
-        self.side_inner_cell_coords = self._make_side_inner_cell_coords()
-        self.side_outer_cell_coords = self._make_side_outer_cell_coords()
-        self.side_from_cell_coords = self._make_side_from_cell_coords()
+        cache.setup_dynamic_attributes(self, constructors=self._dynamic_attribute_constructors_phase_1)
 
         self._make_default_dependent_implementations()
 
+        cache.setup_dynamic_attributes(self, constructors=self._dynamic_attribute_constructors_phase_2)
+
+        if build_bvh:
+            self.build_bvh(self.field.dof_values.device)
+
     @property
     def name(self) -> str:
         return f"DefGeo_{self.field.name}_{'rel' if self._relative else 'abs'}"
@@ -89,35 +104,49 @@ class DeformedGeometry(Geometry):
 
     # Geometry device interface
 
-    @cache.cached_arg_value
-    def cell_arg_value(self, device) -> "DeformedGeometry.CellArg":
-        args = self.CellArg()
+    def _make_cell_arg(self):
+        @cache.dynamic_struct(suffix=self.name)
+        class CellArg:
+            base_arg: self.base.CellArg
+            field_arg: self.field.EvalArg
+            cell_bvh: wp.uint64
 
-        args.elt_arg = self.base.cell_arg_value(device)
-        args.eval_arg = self.field.eval_arg_value(device)
+        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)
+        args.cell_bvh = self.bvh_id(device)
+
     def _make_cell_position(self):
         @cache.dynamic_func(suffix=self.name)
         def cell_position_absolute(cell_arg: self.CellArg, s: Sample):
-            return self.field.eval_inner(cell_arg, s)
+            field_arg = self.field.ElementEvalArg(cell_arg.base_arg, cell_arg.field_arg)
+            return self.field.eval_inner(field_arg, s)
 
         @cache.dynamic_func(suffix=self.name)
         def cell_position(cell_arg: self.CellArg, s: Sample):
-            return self.field.eval_inner(cell_arg, s) + self.base.cell_position(cell_arg.elt_arg, s)
+            field_arg = self.field.ElementEvalArg(cell_arg.base_arg, cell_arg.field_arg)
+            return self.field.eval_inner(field_arg, s) + self.base.cell_position(cell_arg.base_arg, s)
 
         return cell_position if self._relative else cell_position_absolute
 
     def _make_cell_deformation_gradient(self):
         @cache.dynamic_func(suffix=self.name)
         def cell_deformation_gradient_absolute(cell_arg: self.CellArg, s: Sample):
-            return self.field.eval_reference_grad_inner(cell_arg, s)
+            field_arg = self.field.ElementEvalArg(cell_arg.base_arg, cell_arg.field_arg)
+            return self.field.eval_reference_grad_inner(field_arg, s)
 
         @cache.dynamic_func(suffix=self.name)
         def cell_deformation_gradient(cell_arg: self.CellArg, s: Sample):
-            return self.field.eval_reference_grad_inner(cell_arg, s) + self.base.cell_deformation_gradient(
-                cell_arg.elt_arg, s
+            field_arg = self.field.ElementEvalArg(cell_arg.base_arg, cell_arg.field_arg)
+            return self.field.eval_reference_grad_inner(field_arg, s) + self.base.cell_deformation_gradient(
+                cell_arg.base_arg, s
             )
 
         return cell_deformation_gradient if self._relative else cell_deformation_gradient_absolute
@@ -128,19 +157,21 @@ class DeformedGeometry(Geometry):
             base_arg: self.base.SideArg
             trace_arg: self.field_trace.EvalArg
             field_arg: self.field.EvalArg
+            cell_bvh: wp.uint64
 
         return SideArg
 
-    @cache.cached_arg_value
     def side_arg_value(self, device) -> "DeformedGeometry.SideArg":
         args = self.SideArg()
-
-        args.base_arg = self.base.side_arg_value(device)
-        args.field_arg = self.field.eval_arg_value(device)
-        args.trace_arg = self.field_trace.eval_arg_value(device)
-
+        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)
+        self.field_trace.fill_eval_arg(args.trace_arg, device)
+        args.cell_bvh = self.bvh_id(device)
+
     def _make_side_position(self):
         @cache.dynamic_func(suffix=self.name)
         def side_position_absolute(args: self.SideArg, s: Sample):
@@ -216,6 +247,37 @@ class DeformedGeometry(Geometry):
     def _make_side_to_cell_arg(self):
         @cache.dynamic_func(suffix=self.name)
         def side_to_cell_arg(side_arg: self.SideArg):
-            return self.CellArg(self.base.side_to_cell_arg(side_arg.base_arg), side_arg.field_arg)
+            return self.CellArg(self.base.side_to_cell_arg(side_arg.base_arg), side_arg.field_arg, side_arg.cell_bvh)
 
         return side_to_cell_arg
+
+    def _make_cell_bvh_id(self):
+        @cache.dynamic_func(suffix=self.name)
+        def cell_bvh_id(cell_arg: self.CellArg):
+            return cell_arg.cell_bvh
+
+        return cell_bvh_id
+
+    def _make_cell_bounds(self):
+        points, _weights = self.reference_cell().instantiate_quadrature(
+            order=self.field.degree, family=Polynomial.LOBATTO_GAUSS_LEGENDRE
+        )
+
+        points = cache.cached_mat_type((len(points), 3), dtype=float)(points)
+        point_count = len(points)
+
+        @cache.dynamic_func(suffix=self.name)
+        def cell_bounds(cell_arg: self.CellArg, cell_index: ElementIndex):
+            lower = wp.vec3(1.0e8)
+            upper = wp.vec3(-1.0e8)
+            for k in range(point_count):
+                pos = self.cell_position(cell_arg, make_free_sample(cell_index, points[k]))
+                lower = wp.min(lower, pos)
+                upper = wp.max(upper, pos)
+
+            # pad the BBox to account for potential overflows
+            pad = 0.25 * (upper - lower)
+
+            return lower - pad, upper + pad
+
+        return cell_bounds

warp/fem/geometry/element.py

@@ -15,14 +15,18 @@
 
 from typing import List, Tuple
 
+import warp as wp
 from warp.fem.polynomial import Polynomial, quadrature_1d
 from warp.fem.types import Coords
 
+_vec1 = wp.types.vector(length=1, dtype=float)
+
 
 class Element:
     dimension = 0
     """Intrinsic dimension of the element"""
 
+    @staticmethod
     def measure() -> float:
         """Measure (area, volume, ...) of the reference element"""
         raise NotImplementedError
@@ -32,10 +36,32 @@ class Element:
         """Returns a quadrature of a given order for a prototypical element"""
         raise NotImplementedError
 
-    def center(self) -> Tuple[float]:
-        coords, _ = self.instantiate_quadrature(order=0, family=None)
+    @classmethod
+    def center(cls) -> Coords:
+        """Returns the coordinates for the center of the element"""
+        coords, _ = cls.instantiate_quadrature(order=0, family=None)
         return coords[0]
 
+    @wp.func
+    def project(v: Coords):
+        """project coordinates so that they belong to the element"""
+        return wp.min(wp.max(v, Coords(0.0)), Coords(1.0))
+
+    @wp.func
+    def coord_delta(ref_delta: wp.vec3):
+        """Transform a delta in reference space to element coords"""
+        return ref_delta
+
+    @wp.func
+    def coord_delta(ref_delta: wp.vec2):
+        """Transform a delta in reference space to element coords"""
+        return Coords(ref_delta[0], ref_delta[1], 0.0)
+
+    @wp.func
+    def coord_delta(ref_delta: _vec1):
+        """Transform a delta in reference space to element coords"""
+        return Coords(ref_delta[0], 0.0, 0.0)
+
 
 def _point_count_from_order(order: int, family: Polynomial):
     if family == Polynomial.GAUSS_LEGENDRE:
@@ -454,6 +480,17 @@ class Triangle(Element):
 
         return coords, weights
 
+    @wp.func
+    def project(v: Coords):
+        n = wp.max(0.0, (v[0] + v[1] - 1.0) * 0.5)
+        a = wp.clamp(v[0] - n, 0.0, 1.0)
+        b = wp.clamp(v[1] - n, 0.0, 1.0)
+        return Coords(a, b, 1.0 - (a + b))
+
+    @wp.func
+    def coord_delta(ref_delta: wp.vec2):
+        return Coords(-ref_delta[0] - ref_delta[1], ref_delta[0], ref_delta[1])
+
 
 class Tetrahedron(Element):
     dimension = 3
@@ -774,3 +811,20 @@ class Tetrahedron(Element):
             raise NotImplementedError
 
         return coords, weights
+
+    @wp.func
+    def project(v: Coords):
+        # project on 1-2-3 half-space
+        n = wp.max(0.0, (v[0] + v[1] + v[2] - 1.0) / 3.0)
+        c = v - Coords(n)
+
+        # project on 1-2, 2-3, 3-1 half-spaces
+        n = wp.max(0.0, (c[0] + c[1] - 1.0) * 0.5)
+        c = c - Coords(n, n, 0.0)
+        n = wp.max(0.0, (c[1] + c[2] - 1.0) * 0.5)
+        c = c - Coords(0.0, n, n)
+        n = wp.max(0.0, (c[2] + c[0] - 1.0) * 0.5)
+        c = c - Coords(n, 0.0, n)
+
+        # project on cube
+        return Element.project(c)