warp-lang 1.10.0__py3-none-macosx_11_0_arm64.whl

warp/_src/fem/geometry/closest_point.py

@@ -0,0 +1,99 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any
+
+import warp as wp
+from warp._src.fem.types import Coords
+
+_wp_module_name_ = "warp.fem.geometry.closest_point"
+
+
+@wp.func
+def project_on_seg_at_origin(q: Any, seg: Any, len_sq: float):
+    s = wp.clamp(wp.dot(q, seg) / len_sq, 0.0, 1.0)
+    return wp.length_sq(q - s * seg), s
+
+
+@wp.func
+def project_on_tri_at_origin(q: Any, e1: Any, e2: Any):
+    e1e1 = wp.dot(e1, e1)
+    e1e2 = wp.dot(e1, e2)
+    e2e2 = wp.dot(e2, e2)
+
+    det = e1e1 * e2e2 - e1e2 * e1e2
+
+    if det > e1e1 * e2e2 * 1.0e-6:
+        e1p = wp.dot(e1, q)
+        e2p = wp.dot(e2, q)
+
+        s = (e2e2 * e1p - e1e2 * e2p) / det
+        t = (e1e1 * e2p - e1e2 * e1p) / det
+
+        if s >= 0.0 and t >= 0.0 and s + t <= 1.0:
+            # point inside triangle (distance can be non-zero in 3D case)
+            return wp.length_sq(q - s * e1 - t * e2), Coords(1.0 - s - t, s, t)
+
+    d1, s1 = project_on_seg_at_origin(q, e1, e1e1)
+    d2, s2 = project_on_seg_at_origin(q, e2, e2e2)
+    d12, s12 = project_on_seg_at_origin(q - e1, e2 - e1, wp.length_sq(e2 - e1))
+
+    if d1 <= d2:
+        if d1 <= d12:
+            return d1, Coords(1.0 - s1, s1, 0.0)
+    elif d2 <= d12:
+        return d2, Coords(1.0 - s2, 0.0, s2)
+
+    return d12, Coords(0.0, 1.0 - s12, s12)
+
+
+@wp.func
+def project_on_tet_at_origin(q: wp.vec3, e1: wp.vec3, e2: wp.vec3, e3: wp.vec3):
+    mat = wp.inverse(wp.matrix_from_cols(e1, e2, e3))
+    coords = mat * q
+
+    if wp.min(coords) >= 0.0 and coords[0] + coords[1] + coords[2] <= 1.0:
+        return 0.0, coords
+
+    # Not inside tet, compare closest point on each tri
+
+    d12, s12 = project_on_tri_at_origin(q, e1, e2)
+    d23, s23 = project_on_tri_at_origin(q, e2, e3)
+    d31, s31 = project_on_tri_at_origin(q, e3, e1)
+    d123, s123 = project_on_tri_at_origin(q - e1, e2 - e1, e3 - e1)
+
+    dmin = wp.min(wp.vec4(d12, d23, d31, d123))
+
+    if dmin == d12:
+        return dmin, Coords(s12[1], s12[2], 0.0)
+    elif dmin == d23:
+        return dmin, Coords(0.0, s23[1], s23[2])
+    elif dmin == d31:
+        return dmin, Coords(s31[2], 0.0, s31[1])
+    else:
+        return dmin, s123
+
+
+@wp.func
+def project_on_box_at_origin(coords: wp.vec3, sizes: wp.vec3):
+    proj_coords = wp.min(wp.max(coords, wp.vec3(0.0)), sizes)
+    return wp.length_sq(coords - proj_coords), wp.cw_div(proj_coords, sizes)
+
+
+@wp.func
+def project_on_box_at_origin(coords: wp.vec2, sizes: wp.vec2):
+    proj_coords = wp.min(wp.max(coords, wp.vec2(0.0)), sizes)
+    norm_coords = wp.cw_div(proj_coords, sizes)
+    return wp.length_sq(coords - proj_coords), Coords(norm_coords[0], norm_coords[1], 0.0)

warp/_src/fem/geometry/deformed_geometry.py

@@ -0,0 +1,277 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import ClassVar
+
+import warp as wp
+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)
+
+
+class DeformedGeometry(Geometry):
+    _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.
+
+        See also: :meth:`warp.fem.GeometryField.make_deformed_geometry`
+        """
+
+        from warp._src.fem.field import GeometryField
+
+        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")
+
+        self._relative = relative
+
+        self.field: GeometryField = field
+        self.field_trace = field.trace()
+        self.dimension = self.base.dimension
+
+        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
+        self.boundary_side_count = self.base.boundary_side_count
+        self.reference_cell = self.base.reference_cell
+        self.reference_side = self.base.reference_side
+
+        self.side_index_arg_value = self.base.side_index_arg_value
+        self.fill_side_index_arg = self.base.fill_side_index_arg
+        self.boundary_side_index = self.base.boundary_side_index
+
+        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'}"
+
+    @property
+    def base(self) -> Geometry:
+        return self.field.geometry.base
+
+    # Geometry device interface
+
+    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
+
+        return CellArg
+
+    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):
+            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):
+            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):
+            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):
+            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
+
+    def _make_side_arg(self):
+        @cache.dynamic_struct(suffix=self.name)
+        class SideArg:
+            base_arg: self.base.SideArg
+            trace_arg: self.field_trace.EvalArg
+            field_arg: self.field.EvalArg
+            cell_bvh: wp.uint64
+
+        return SideArg
+
+    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):
+            trace_arg = self.field_trace.ElementEvalArg(args.base_arg, args.trace_arg)
+            return self.field_trace.eval_inner(trace_arg, s)
+
+        @cache.dynamic_func(suffix=self.name)
+        def side_position(args: self.SideArg, s: Sample):
+            trace_arg = self.field_trace.ElementEvalArg(args.base_arg, args.trace_arg)
+            return self.field_trace.eval_inner(trace_arg, s) + self.base.side_position(args.base_arg, s)
+
+        return side_position if self._relative else side_position_absolute
+
+    def _make_side_deformation_gradient(self):
+        @cache.dynamic_func(suffix=self.name)
+        def side_deformation_gradient_absolute(args: self.SideArg, s: Sample):
+            base_def_grad = self.base.side_deformation_gradient(args.base_arg, s)
+            trace_arg = self.field_trace.ElementEvalArg(args.base_arg, args.trace_arg)
+
+            Du = self.field_trace.eval_grad_inner(trace_arg, s)
+            return Du * base_def_grad
+
+        @cache.dynamic_func(suffix=self.name)
+        def side_deformation_gradient(args: self.SideArg, s: Sample):
+            base_def_grad = self.base.side_deformation_gradient(args.base_arg, s)
+            trace_arg = self.field_trace.ElementEvalArg(args.base_arg, args.trace_arg)
+
+            Du = self.field_trace.eval_grad_inner(trace_arg, s)
+            return base_def_grad + Du * base_def_grad
+
+        return side_deformation_gradient if self._relative else side_deformation_gradient_absolute
+
+    def _make_side_inner_cell_index(self):
+        @cache.dynamic_func(suffix=self.name)
+        def side_inner_cell_index(args: self.SideArg, side_index: ElementIndex):
+            return self.base.side_inner_cell_index(args.base_arg, side_index)
+
+        return side_inner_cell_index
+
+    def _make_side_outer_cell_index(self):
+        @cache.dynamic_func(suffix=self.name)
+        def side_outer_cell_index(args: self.SideArg, side_index: ElementIndex):
+            return self.base.side_outer_cell_index(args.base_arg, side_index)
+
+        return side_outer_cell_index
+
+    def _make_side_inner_cell_coords(self):
+        @cache.dynamic_func(suffix=self.name)
+        def side_inner_cell_coords(args: self.SideArg, side_index: ElementIndex, side_coords: Coords):
+            return self.base.side_inner_cell_coords(args.base_arg, side_index, side_coords)
+
+        return side_inner_cell_coords
+
+    def _make_side_outer_cell_coords(self):
+        @cache.dynamic_func(suffix=self.name)
+        def side_outer_cell_coords(args: self.SideArg, side_index: ElementIndex, side_coords: Coords):
+            return self.base.side_outer_cell_coords(args.base_arg, side_index, side_coords)
+
+        return side_outer_cell_coords
+
+    def _make_side_from_cell_coords(self):
+        @cache.dynamic_func(suffix=self.name)
+        def side_from_cell_coords(
+            args: self.SideArg,
+            side_index: ElementIndex,
+            cell_index: ElementIndex,
+            cell_coords: Coords,
+        ):
+            return self.base.side_from_cell_coords(args.base_arg, side_index, cell_index, cell_coords)
+
+        return side_from_cell_coords
+
+    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, 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().prototype.instantiate_quadrature(
+            order=self.field.degree, family=Polynomial.LOBATTO_GAUSS_LEGENDRE
+        )
+
+        points = wp.matrix_from_rows(*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