warp-lang 0.9.0__py3-none-win_amd64.whl → 0.11.0__py3-none-win_amd64.whl

warp/fem/geometry/__init__.py

@@ -0,0 +1,19 @@
+from .element import Element
+from .geometry import Geometry
+from .deformed_geometry import DeformedGeometry
+
+from .partition import (
+    GeometryPartition,
+    WholeGeometryPartition,
+    LinearGeometryPartition,
+    ExplicitGeometryPartition,
+)
+
+
+from .grid_2d import Grid2D
+from .grid_3d import Grid3D
+
+from .trimesh_2d import Trimesh2D
+from .tetmesh import Tetmesh
+from .quadmesh_2d import Quadmesh2D
+from .hexmesh import Hexmesh

warp/fem/geometry/closest_point.py

@@ -0,0 +1,70 @@
+from typing import Any
+
+import warp as wp
+
+from warp.fem.types import Coords
+
+
+@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.mat33(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

warp/fem/geometry/deformed_geometry.py

@@ -0,0 +1,271 @@
+from typing import Any
+
+import warp as wp
+
+from .geometry import Geometry
+
+from warp.fem.types import Sample, ElementIndex, Coords, make_free_sample
+from warp.fem import cache
+
+_mat32 = wp.mat(shape=(3, 2), dtype=float)
+
+
+class DeformedGeometry(Geometry):
+    def __init__(self, field):
+        """Constructs a Deformed Geometry from a displacement field defined over a base geometry"""
+
+        from warp.fem.field import DiscreteField
+
+        self.field: DiscreteField = field
+        self.base = self.field.space.geometry
+        self.dimension = self.base.dimension
+
+        if not wp.types.type_is_vector(field.dtype) or wp.types.type_length(field.dtype) != self.dimension:
+            raise ValueError("Invalid value type for position field")
+
+        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
+        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.cell_position = self._make_cell_position()
+        self.cell_deformation_gradient = self._make_cell_deformation_gradient()
+        self.cell_inverse_deformation_gradient = self._make_cell_inverse_deformation_gradient()
+        self.cell_measure = self._make_cell_measure()
+
+        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()
+        self.side_inner_inverse_deformation_gradient = self._make_side_inner_inverse_deformation_gradient()
+        self.side_outer_inverse_deformation_gradient = self._make_side_outer_inverse_deformation_gradient()
+        self.side_measure = self._make_side_measure()
+        self.side_measure_ratio = self._make_side_measure_ratio()
+        self.side_normal = self._make_side_normal()
+
+    @property
+    def name(self):
+        return f"DefGeo_{self.field.name}"
+
+    # Geometry device interface
+
+    @cache.cached_arg_value
+    def cell_arg_value(self, device) -> "DeformedGeometry.CellArg":
+        args = self.CellArg()
+
+        args.elt_arg = self.base.cell_arg_value(device)
+        args.eval_arg = self.field.eval_arg_value(device)
+
+        return args
+
+    def _make_cell_position(self):
+        @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)
+
+        return cell_position
+
+    def _make_cell_deformation_gradient(self):
+        @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
+            )
+
+        return cell_deformation_gradient
+
+    def _make_cell_inverse_deformation_gradient(self):
+        @cache.dynamic_func(suffix=self.name)
+        def cell_inverse_deformation_gradient(cell_arg: self.CellArg, s: Sample):
+            return wp.inverse(self.cell_deformation_gradient(cell_arg, s))
+
+        return cell_inverse_deformation_gradient
+
+    def _make_cell_measure(self):
+        REF_MEASURE = wp.constant(self.reference_cell().measure())
+
+        @cache.dynamic_func(suffix=self.name)
+        def cell_measure(args: self.CellArg, s: Sample):
+            return wp.abs(wp.determinant(self.cell_deformation_gradient(args, s))) * REF_MEASURE
+
+        return cell_measure
+
+    @wp.func
+    def cell_normal(args: Any, s: Sample):
+        return wp.vec2(0.0)
+
+    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
+
+        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)
+
+        return args
+
+    def _make_side_position(self):
+        @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
+
+    def _make_side_deformation_gradient(self):
+        @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
+
+    def _make_side_inner_inverse_deformation_gradient(self):
+        @cache.dynamic_func(suffix=self.name)
+        def side_inner_inverse_deformation_gradient(args: self.SideArg, s: Sample):
+            cell_index = self.side_inner_cell_index(args, s.element_index)
+            cell_coords = self.side_inner_cell_coords(args, s.element_index, s.element_coords)
+            cell_arg = self.side_to_cell_arg(args)
+            return self.cell_inverse_deformation_gradient(cell_arg, make_free_sample(cell_index, cell_coords))
+
+    def _make_side_outer_inverse_deformation_gradient(self):
+        @cache.dynamic_func(suffix=self.name)
+        def side_outer_inverse_deformation_gradient(args: self.SideArg, s: Sample):
+            cell_index = self.side_outer_cell_index(args, s.element_index)
+            cell_coords = self.side_outer_cell_coords(args, s.element_index, s.element_coords)
+            cell_arg = self.side_to_cell_arg(args)
+            return self.cell_inverse_deformation_gradient(cell_arg, make_free_sample(cell_index, cell_coords))
+
+    @wp.func
+    def _side_measure(F: wp.vec2):
+        return wp.length(F)
+
+    @wp.func
+    def _side_measure(F: _mat32):
+        Fcross = wp.vec3(
+            F[1, 0] * F[2, 1] - F[2, 0] * F[1, 1],
+            F[2, 0] * F[0, 1] - F[0, 0] * F[2, 1],
+            F[0, 0] * F[1, 1] - F[1, 0] * F[0, 1],
+        )
+        return wp.length(Fcross)
+
+    @wp.func
+    def _side_normal(F: wp.vec2):
+        return wp.normalize(wp.vec2(-F[1], F[0]))
+
+    @wp.func
+    def _side_normal(F: _mat32):
+        Fcross = wp.vec3(
+            F[1, 0] * F[2, 1] - F[2, 0] * F[1, 1],
+            F[2, 0] * F[0, 1] - F[0, 0] * F[2, 1],
+            F[0, 0] * F[1, 1] - F[1, 0] * F[0, 1],
+        )
+        return wp.normalize(Fcross)
+
+    def _make_side_measure(self):
+        REF_MEASURE = wp.constant(self.reference_side().measure())
+
+        @cache.dynamic_func(suffix=self.name)
+        def side_measure(args: self.SideArg, s: Sample):
+            F = self.side_deformation_gradient(args, s)
+            return DeformedGeometry._side_measure(F) * REF_MEASURE
+
+        return side_measure
+
+    def _make_side_measure_ratio(self):
+        @cache.dynamic_func(suffix=self.name)
+        def side_measure_ratio(args: self.SideArg, s: Sample):
+            inner = self.side_inner_cell_index(args, s.element_index)
+            outer = self.side_outer_cell_index(args, s.element_index)
+            inner_coords = self.side_inner_cell_coords(args, s.element_index, s.element_coords)
+            outer_coords = self.side_outer_cell_coords(args, s.element_index, s.element_coords)
+            cell_arg = self.side_to_cell_arg(args)
+            return self.side_measure(args, s) / wp.min(
+                self.cell_measure(cell_arg, make_free_sample(inner, inner_coords)),
+                self.cell_measure(cell_arg, make_free_sample(outer, outer_coords)),
+            )
+
+        return side_measure_ratio
+
+    def _make_side_normal(self):
+        @cache.dynamic_func(suffix=self.name)
+        def side_normal(args: self.SideArg, s: Sample):
+            F = self.side_deformation_gradient(args, s)
+            return DeformedGeometry._side_normal(F)
+
+        return side_normal
+
+    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)
+
+        return side_to_cell_arg