warp-lang 1.0.1__py3-none-manylinux2014_x86_64.whl → 1.1.0__py3-none-manylinux2014_x86_64.whl

warp/fem/field/trial.py

@@ -1,183 +1,183 @@
-import warp as wp
-from warp.fem.domain import GeometryDomain
-from warp.fem.space import FunctionSpace, SpacePartition
-from warp.fem.types import Sample, get_node_index_in_element
-from warp.fem import utils, cache
-
-from .field import SpaceField
-
-
-class TrialField(SpaceField):
-    """Field defined over a domain that can be used as a trial function"""
-
-    def __init__(
-        self,
-        space: FunctionSpace,
-        space_partition: SpacePartition,
-        domain: GeometryDomain,
-    ):
-        if domain.dimension == space.dimension - 1:
-            space = space.trace()
-
-        if domain.dimension != space.dimension:
-            raise ValueError("Incompatible space and domain dimensions")
-
-        if not space.topology.is_derived_from(space_partition.space_topology):
-            raise ValueError("Incompatible space and space partition topologies")
-
-        super().__init__(space, space_partition)
-
-        self.domain = domain
-
-        self.EvalArg = self.space.SpaceArg
-        self.ElementEvalArg = self._make_element_eval_arg()
-
-        self.eval_degree = self._make_eval_degree()
-        self.eval_inner = self._make_eval_inner()
-        self.eval_grad_inner = self._make_eval_grad_inner()
-        self.eval_div_inner = self._make_eval_div_inner()
-        self.eval_outer = self._make_eval_outer()
-        self.eval_grad_outer = self._make_eval_grad_outer()
-        self.eval_div_outer = self._make_eval_div_outer()
-        self.at_node = self._make_at_node()
-
-    def partition_node_count(self) -> int:
-        """Returns the number of nodes in the associated space topology partition"""
-        return self.space_partition.node_count()
-
-    @property
-    def name(self) -> str:
-        return self.space.name + "Trial"
-
-    def eval_arg_value(self, device) -> wp.codegen.StructInstance:
-        return self.space.space_arg_value(device)
-
-    def _make_element_eval_arg(self):
-        @cache.dynamic_struct(suffix=self.name)
-        class ElementEvalArg:
-            elt_arg: self.domain.ElementArg
-            eval_arg: self.EvalArg
-
-        return ElementEvalArg
-
-    def _make_eval_inner(self):
-        @cache.dynamic_func(suffix=self.name)
-        def eval_trial_inner(args: self.ElementEvalArg, s: Sample):
-            weight = self.space.element_inner_weight(
-                args.elt_arg,
-                args.eval_arg,
-                s.element_index,
-                s.element_coords,
-                get_node_index_in_element(s.trial_dof),
-            )
-            return weight * self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.trial_dof)
-
-        return eval_trial_inner
-
-    def _make_eval_grad_inner(self):
-        if not self.gradient_valid():
-            return None
-
-        @cache.dynamic_func(suffix=self.name)
-        def eval_nabla_trial_inner(args: self.ElementEvalArg, s: Sample):
-            nabla_weight = self.space.element_inner_weight_gradient(
-                args.elt_arg,
-                args.eval_arg,
-                s.element_index,
-                s.element_coords,
-                get_node_index_in_element(s.trial_dof),
-            )
-            grad_transform = self.space.element_inner_reference_gradient_transform(args.elt_arg, s)
-            return utils.generalized_outer(
-                self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.trial_dof),
-                utils.apply_right(nabla_weight, grad_transform),
-            )
-
-        return eval_nabla_trial_inner
-
-    def _make_eval_div_inner(self):
-        if not self.divergence_valid():
-            return None
-
-        @cache.dynamic_func(suffix=self.name)
-        def eval_div_trial_inner(args: self.ElementEvalArg, s: Sample):
-            nabla_weight = self.space.element_inner_weight_gradient(
-                args.elt_arg,
-                args.eval_arg,
-                s.element_index,
-                s.element_coords,
-                get_node_index_in_element(s.trial_dof),
-            )
-            grad_transform = self.space.element_inner_reference_gradient_transform(args.elt_arg, s)
-            return utils.generalized_inner(
-                self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.trial_dof),
-                utils.apply_right(nabla_weight, grad_transform),
-            )
-
-        return eval_div_trial_inner
-
-    def _make_eval_outer(self):
-        @cache.dynamic_func(suffix=self.name)
-        def eval_trial_outer(args: self.ElementEvalArg, s: Sample):
-            weight = self.space.element_outer_weight(
-                args.elt_arg,
-                args.eval_arg,
-                s.element_index,
-                s.element_coords,
-                get_node_index_in_element(s.trial_dof),
-            )
-            return weight * self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.trial_dof)
-
-        return eval_trial_outer
-
-    def _make_eval_grad_outer(self):
-        if not self.gradient_valid():
-            return None
-
-        @cache.dynamic_func(suffix=self.name)
-        def eval_nabla_trial_outer(args: self.ElementEvalArg, s: Sample):
-            nabla_weight = self.space.element_outer_weight_gradient(
-                args.elt_arg,
-                args.eval_arg,
-                s.element_index,
-                s.element_coords,
-                get_node_index_in_element(s.trial_dof),
-            )
-            grad_transform = self.space.element_outer_reference_gradient_transform(args.elt_arg, s)
-            return utils.generalized_outer(
-                self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.trial_dof),
-                utils.apply_right(nabla_weight, grad_transform),
-            )
-
-        return eval_nabla_trial_outer
-
-    def _make_eval_div_outer(self):
-        if not self.divergence_valid():
-            return None
-
-        @cache.dynamic_func(suffix=self.name)
-        def eval_div_trial_outer(args: self.ElementEvalArg, s: Sample):
-            nabla_weight = self.space.element_outer_weight_gradient(
-                args.elt_arg,
-                args.eval_arg,
-                s.element_index,
-                s.element_coords,
-                get_node_index_in_element(s.trial_dof),
-            )
-            grad_transform = self.space.element_outer_reference_gradient_transform(args.elt_arg, s)
-            return utils.generalized_inner(
-                self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.trial_dof),
-                utils.apply_right(nabla_weight, grad_transform),
-            )
-
-        return eval_div_trial_outer
-
-    def _make_at_node(self):
-        @cache.dynamic_func(suffix=self.name)
-        def at_node(args: self.ElementEvalArg, s: Sample):
-            node_coords = self.space.node_coords_in_element(
-                args.elt_arg, args.eval_arg, s.element_index, get_node_index_in_element(s.trial_dof)
-            )
-            return Sample(s.element_index, node_coords, s.qp_index, s.qp_weight, s.test_dof, s.trial_dof)
-
-        return at_node
+import warp as wp
+from warp.fem import cache, utils
+from warp.fem.domain import GeometryDomain
+from warp.fem.space import FunctionSpace, SpacePartition
+from warp.fem.types import Sample, get_node_index_in_element
+
+from .field import SpaceField
+
+
+class TrialField(SpaceField):
+    """Field defined over a domain that can be used as a trial function"""
+
+    def __init__(
+        self,
+        space: FunctionSpace,
+        space_partition: SpacePartition,
+        domain: GeometryDomain,
+    ):
+        if domain.dimension == space.dimension - 1:
+            space = space.trace()
+
+        if domain.dimension != space.dimension:
+            raise ValueError("Incompatible space and domain dimensions")
+
+        if not space.topology.is_derived_from(space_partition.space_topology):
+            raise ValueError("Incompatible space and space partition topologies")
+
+        super().__init__(space, space_partition)
+
+        self.domain = domain
+
+        self.EvalArg = self.space.SpaceArg
+        self.ElementEvalArg = self._make_element_eval_arg()
+
+        self.eval_degree = self._make_eval_degree()
+        self.eval_inner = self._make_eval_inner()
+        self.eval_grad_inner = self._make_eval_grad_inner()
+        self.eval_div_inner = self._make_eval_div_inner()
+        self.eval_outer = self._make_eval_outer()
+        self.eval_grad_outer = self._make_eval_grad_outer()
+        self.eval_div_outer = self._make_eval_div_outer()
+        self.at_node = self._make_at_node()
+
+    def partition_node_count(self) -> int:
+        """Returns the number of nodes in the associated space topology partition"""
+        return self.space_partition.node_count()
+
+    @property
+    def name(self) -> str:
+        return self.space.name + "Trial"
+
+    def eval_arg_value(self, device) -> wp.codegen.StructInstance:
+        return self.space.space_arg_value(device)
+
+    def _make_element_eval_arg(self):
+        @cache.dynamic_struct(suffix=self.name)
+        class ElementEvalArg:
+            elt_arg: self.domain.ElementArg
+            eval_arg: self.EvalArg
+
+        return ElementEvalArg
+
+    def _make_eval_inner(self):
+        @cache.dynamic_func(suffix=self.name)
+        def eval_trial_inner(args: self.ElementEvalArg, s: Sample):
+            weight = self.space.element_inner_weight(
+                args.elt_arg,
+                args.eval_arg,
+                s.element_index,
+                s.element_coords,
+                get_node_index_in_element(s.trial_dof),
+            )
+            return weight * self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.trial_dof)
+
+        return eval_trial_inner
+
+    def _make_eval_grad_inner(self):
+        if not self.gradient_valid():
+            return None
+
+        @cache.dynamic_func(suffix=self.name)
+        def eval_nabla_trial_inner(args: self.ElementEvalArg, s: Sample):
+            nabla_weight = self.space.element_inner_weight_gradient(
+                args.elt_arg,
+                args.eval_arg,
+                s.element_index,
+                s.element_coords,
+                get_node_index_in_element(s.trial_dof),
+            )
+            grad_transform = self.space.element_inner_reference_gradient_transform(args.elt_arg, s)
+            return utils.generalized_outer(
+                self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.trial_dof),
+                utils.apply_right(nabla_weight, grad_transform),
+            )
+
+        return eval_nabla_trial_inner
+
+    def _make_eval_div_inner(self):
+        if not self.divergence_valid():
+            return None
+
+        @cache.dynamic_func(suffix=self.name)
+        def eval_div_trial_inner(args: self.ElementEvalArg, s: Sample):
+            nabla_weight = self.space.element_inner_weight_gradient(
+                args.elt_arg,
+                args.eval_arg,
+                s.element_index,
+                s.element_coords,
+                get_node_index_in_element(s.trial_dof),
+            )
+            grad_transform = self.space.element_inner_reference_gradient_transform(args.elt_arg, s)
+            return utils.generalized_inner(
+                self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.trial_dof),
+                utils.apply_right(nabla_weight, grad_transform),
+            )
+
+        return eval_div_trial_inner
+
+    def _make_eval_outer(self):
+        @cache.dynamic_func(suffix=self.name)
+        def eval_trial_outer(args: self.ElementEvalArg, s: Sample):
+            weight = self.space.element_outer_weight(
+                args.elt_arg,
+                args.eval_arg,
+                s.element_index,
+                s.element_coords,
+                get_node_index_in_element(s.trial_dof),
+            )
+            return weight * self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.trial_dof)
+
+        return eval_trial_outer
+
+    def _make_eval_grad_outer(self):
+        if not self.gradient_valid():
+            return None
+
+        @cache.dynamic_func(suffix=self.name)
+        def eval_nabla_trial_outer(args: self.ElementEvalArg, s: Sample):
+            nabla_weight = self.space.element_outer_weight_gradient(
+                args.elt_arg,
+                args.eval_arg,
+                s.element_index,
+                s.element_coords,
+                get_node_index_in_element(s.trial_dof),
+            )
+            grad_transform = self.space.element_outer_reference_gradient_transform(args.elt_arg, s)
+            return utils.generalized_outer(
+                self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.trial_dof),
+                utils.apply_right(nabla_weight, grad_transform),
+            )
+
+        return eval_nabla_trial_outer
+
+    def _make_eval_div_outer(self):
+        if not self.divergence_valid():
+            return None
+
+        @cache.dynamic_func(suffix=self.name)
+        def eval_div_trial_outer(args: self.ElementEvalArg, s: Sample):
+            nabla_weight = self.space.element_outer_weight_gradient(
+                args.elt_arg,
+                args.eval_arg,
+                s.element_index,
+                s.element_coords,
+                get_node_index_in_element(s.trial_dof),
+            )
+            grad_transform = self.space.element_outer_reference_gradient_transform(args.elt_arg, s)
+            return utils.generalized_inner(
+                self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.trial_dof),
+                utils.apply_right(nabla_weight, grad_transform),
+            )
+
+        return eval_div_trial_outer
+
+    def _make_at_node(self):
+        @cache.dynamic_func(suffix=self.name)
+        def at_node(args: self.ElementEvalArg, s: Sample):
+            node_coords = self.space.node_coords_in_element(
+                args.elt_arg, args.eval_arg, s.element_index, get_node_index_in_element(s.trial_dof)
+            )
+            return Sample(s.element_index, node_coords, s.qp_index, s.qp_weight, s.test_dof, s.trial_dof)
+
+        return at_node

warp/fem/geometry/__init__.py

@@ -1,19 +1,15 @@
-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
+from .deformed_geometry import DeformedGeometry
+from .element import Element
+from .geometry import Geometry
+from .grid_2d import Grid2D
+from .grid_3d import Grid3D
+from .hexmesh import Hexmesh
+from .partition import (
+    ExplicitGeometryPartition,
+    GeometryPartition,
+    LinearGeometryPartition,
+    WholeGeometryPartition,
+)
+from .quadmesh_2d import Quadmesh2D
+from .tetmesh import Tetmesh
+from .trimesh_2d import Trimesh2D

warp/fem/geometry/closest_point.py

@@ -1,70 +1,69 @@
-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
+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