warp-lang 1.2.2__py3-none-manylinux2014_aarch64.whl → 1.3.0__py3-none-manylinux2014_aarch64.whl

warp/fem/field/field.py

@@ -1,9 +1,12 @@
-from typing import Any
+from typing import Any, Dict, Optional
 
 import warp as wp
+from warp.fem import cache
+from warp.fem.domain import GeometryDomain, Sides
 from warp.fem.geometry import DeformedGeometry, Geometry
+from warp.fem.operator import integrand
 from warp.fem.space import FunctionSpace, SpacePartition
-from warp.fem.types import Sample
+from warp.fem.types import NULL_ELEMENT_INDEX, ElementKind, Sample
 
 
 class FieldLike:
@@ -36,6 +39,14 @@ class FieldLike:
         """Value of arguments to be passed to device functions"""
         raise NotImplementedError
 
+    def gradient_valid(self) -> bool:
+        """Whether the gradient operator is implemented for this field."""
+        return False
+
+    def divergence_valid(self) -> bool:
+        """Whether the divergence operator is implemented for this field."""
+        return False
+
     @staticmethod
     def eval_inner(args: "ElementEvalArg", s: "Sample"):  # noqa: F821
         """Device function evaluating the inner field value at a sample point"""
@@ -66,14 +77,64 @@ class FieldLike:
         """Device function evaluating the outer field divergence at a sample point"""
         raise NotImplementedError
 
+    @staticmethod
+    def eval_degree(args: "ElementEvalArg"):  # noqa: F821
+        """Polynomial degree of the field is applicable, or hint for determination of interpolation order"""
+        raise NotImplementedError
+
 
-class SpaceField(FieldLike):
+class GeometryField(FieldLike):
+    """Base class for fields defined over a geometry"""
+
+    @property
+    def geometry(self) -> Geometry:
+        """Geometry over which the field is expressed"""
+        raise NotImplementedError
+
+    @property
+    def element_kind(self) -> ElementKind:
+        """Kind of element over which the field is expressed"""
+        raise NotImplementedError
+
+    @staticmethod
+    def eval_reference_grad_inner(args: "ElementEvalArg", s: "Sample"):  # noqa: F821
+        """Device function evaluating the inner field gradient with respect to reference element coordinates at a sample point"""
+        raise NotImplementedError
+
+    @staticmethod
+    def eval_reference_grad_outer(args: "ElementEvalArg", s: "Sample"):  # noqa: F821
+        """Device function evaluating the outer field gradient with respect to reference element coordinates at a sample point"""
+        raise NotImplementedError
+
+    def trace(self) -> FieldLike:
+        """Trace of this field over lower-dimensional elements"""
+        raise NotImplementedError
+
+    def make_deformed_geometry(self, relative=True) -> Geometry:
+        """Returns a deformed version of the underlying geometry, with positions displaced according to this field's values.
+
+        Args:
+            relative: If ``True``, the field is intepreted as a relative displacement over the original geometry. If ``False``, the field values are intepreted as absolute positions.
+
+        """
+        return DeformedGeometry(self, relative=relative)
+
+
+class SpaceField(GeometryField):
     """Base class for fields defined over a function space"""
 
     def __init__(self, space: FunctionSpace, space_partition: SpacePartition):
         self._space = space
         self._space_partition = space_partition
 
+    @property
+    def geometry(self) -> Geometry:
+        return self._space.geometry
+
+    @property
+    def element_kind(self) -> ElementKind:
+        return self._space.element_kind
+
     @property
     def space(self) -> FunctionSpace:
         return self._space
@@ -108,7 +169,6 @@ class SpaceField(FieldLike):
 
     def _make_eval_degree(self):
         ORDER = self.space.ORDER
-        from warp.fem import cache
 
         @cache.dynamic_func(suffix=self.name)
         def degree(args: self.ElementEvalArg):
@@ -130,10 +190,6 @@ class DiscreteField(SpaceField):
         """Sets degrees of freedom values from an array"""
         raise NotImplementedError
 
-    def trace(self) -> "DiscreteField":
-        """Trace of this field over a lower-dimensional function space"""
-        raise NotImplementedError
-
     @staticmethod
     def set_node_value(args: "FieldLike.EvalArg", node_index: int, value: Any):
         """Device function setting the value at given node"""
@@ -143,6 +199,393 @@ class DiscreteField(SpaceField):
     def name(self) -> str:
         return f"{self.__class__.__qualname__}_{self.space.name}_{self.space_partition.name}"
 
-    def make_deformed_geometry(self) -> Geometry:
-        """Returns a deformed version of the underlying geometry using this field's values as displacement"""
-        return DeformedGeometry(self)
+
+class ImplicitField(GeometryField):
+    """Field defined from an arbitrary function over a domain.
+    Does not support autodiff yet, so if gradient/divergence evaluation is required corresponding functions must be provided.
+
+    Args:
+        domain: Domain over which the field is defined
+        func: Warp function evaluating the field at a given position. Must accept at least one argument, with the first argument being the evaluation position (``wp.vec2`` or ``wp.vec3``).
+        values: Optional dictionary of additional argument values to be passed to the evaluation function.
+        grad_func: Optional gradient evaluation function; must take same arguments as `func`
+        div_func: Optional divergence evaluation function; must take same arguments as `func`
+        degree: Optional hint for automatic determination of quadrature orders when integrating this field
+    """
+
+    def __init__(
+        self,
+        domain: GeometryDomain,
+        func: wp.Function,
+        values: Optional[Dict[str, Any]] = None,
+        grad_func: Optional[wp.Function] = None,
+        div_func: Optional[wp.Function] = None,
+        degree=0,
+    ):
+        self.domain = domain
+        self._degree = degree
+
+        if not isinstance(func, wp.Function):
+            raise ValueError("Implicit field function must be a warp Function (decorated with `wp.func`)")
+
+        self._func = func
+        self._grad_func = grad_func
+        self._div_func = div_func
+
+        argspec = integrand(func.func).argspec
+        arg_types = argspec.annotations
+
+        pos_arg_type = arg_types.pop(argspec.args[0]) if arg_types else None
+        if not pos_arg_type or not wp.types.types_equal(
+            pos_arg_type, wp.vec(length=domain.geometry.dimension, dtype=float), match_generic=True
+        ):
+            raise ValueError(
+                f"Implicit field function '{func.func.__name__}' must accept a position as its first argument"
+            )
+
+        self.EvalArg = cache.get_argument_struct(arg_types)
+        self.values = values
+
+        self.ElementEvalArg = self._make_element_eval_arg()
+        self.eval_degree = self._make_eval_degree()
+
+        self.eval_inner = self._make_eval_func(func)
+        self.eval_grad_inner = self._make_eval_func(grad_func)
+        self.eval_div_inner = self._make_eval_func(div_func)
+        self.eval_reference_grad_inner = self._make_eval_reference_grad()
+
+        self.eval_outer = self.eval_inner
+        self.eval_grad_outer = self.eval_grad_inner
+        self.eval_div_outer = self.eval_div_inner
+        self.eval_reference_grad_outer = self.eval_reference_grad_inner
+
+    @property
+    def values(self):
+        return self._func_arg
+
+    @values.setter
+    def values(self, v):
+        self._func_arg = cache.populate_argument_struct(self.EvalArg, v, self._func.func.__name__)
+
+    @property
+    def geometry(self) -> Geometry:
+        return self.domain.geometry
+
+    @property
+    def element_kind(self) -> ElementKind:
+        return self.domain.element_kind
+
+    def eval_arg_value(self, device):
+        return self._func_arg
+
+    @property
+    def degree(self) -> int:
+        return self._degree
+
+    @property
+    def name(self) -> str:
+        return f"Implicit_{self.domain.name}_{self.degree}_{self.EvalArg.key}"
+
+    def _make_eval_func(self, func):
+        if func is None:
+            return None
+
+        @cache.dynamic_func(
+            suffix=f"{self.name}_{func.key}",
+            code_transformers=[cache.ExpandStarredArgumentStruct({"args.eval_arg": self.EvalArg})],
+        )
+        def eval_inner(args: self.ElementEvalArg, s: Sample):
+            pos = self.domain.element_position(args.elt_arg, s)
+            return func(pos, *args.eval_arg)
+
+        return eval_inner
+
+    def _make_eval_reference_grad(self):
+        if self.eval_grad_inner is None:
+            return None
+
+        @cache.dynamic_func(suffix=f"{self.eval_grad_inner.key}")
+        def eval_reference_grad_inner(args: self.ElementEvalArg, s: Sample):
+            return self.eval_grad_inner(args, s) * self.domain.element_deformation_gradient(args.elt_arg, s)
+
+        return eval_reference_grad_inner
+
+    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_degree(self):
+        ORDER = wp.constant(self._degree)
+
+        @cache.dynamic_func(suffix=self.name)
+        def degree(args: self.ElementEvalArg):
+            return ORDER
+
+        return degree
+
+    def trace(self):
+        if self.element_kind == ElementKind.SIDE:
+            raise RuntimeError("Trace only available for field defined on cell elements")
+
+        return ImplicitField(
+            domain=Sides(self.domain.geometry_partition),
+            func=self._func,
+            values={name: getattr(self.values, name) for name in self.EvalArg.vars},
+            grad_func=self._grad_func,
+            div_func=self._div_func,
+            degree=self._degree,
+        )
+
+
+class UniformField(GeometryField):
+    """Field defined as a constant value over a domain.
+
+    Args:
+        domain: Domain over which the field is defined
+        value: Uniform value over the domain
+    """
+
+    def __init__(self, domain: GeometryDomain, value: Any):
+        self.domain = domain
+
+        if not wp.types.is_value(value):
+            raise ValueError("value must be a Warp scalar, vector or matrix")
+
+        self.dtype = wp.types.type_to_warp(type(value))
+        self._value = self.dtype(value)
+
+        scalar_type = wp.types.type_scalar_type(self.dtype)
+        if wp.types.type_is_vector(self.dtype):
+            grad_type = wp.mat(shape=(wp.types.type_length(self.dtype), self.geometry.dimension), dtype=scalar_type)
+            div_type = scalar_type
+        elif wp.types.type_is_matrix(self.dtype):
+            grad_type = None
+            div_type = wp.vec(length=(wp.types.type_length(self.dtype) // self.geometry.dimension), dtype=scalar_type)
+        else:
+            div_type = None
+            grad_type = wp.vec(length=self.geometry.dimension, dtype=scalar_type)
+
+        self.EvalArg = self._make_eval_arg()
+        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_zero(grad_type)
+        self.eval_div_inner = self._make_eval_zero(div_type)
+        self.eval_reference_grad_inner = self.eval_grad_inner
+
+        self.eval_outer = self.eval_inner
+        self.eval_grad_outer = self.eval_grad_inner
+        self.eval_div_outer = self.eval_div_inner
+        self.eval_reference_grad_outer = self.eval_reference_grad_inner
+
+    @property
+    def value(self):
+        return self._value
+
+    @value.setter
+    def value(self, v):
+        value_type = wp.types.type_to_warp(type(v))
+        assert wp.types.types_equal(value_type, self.dtype)
+        self._value = self.dtype(v)
+
+    @property
+    def geometry(self) -> Geometry:
+        return self.domain.geometry
+
+    @property
+    def element_kind(self) -> ElementKind:
+        return self.domain.element_kind
+
+    def eval_arg_value(self, device):
+        arg = self.EvalArg()
+        arg.value = self.value
+        return arg
+
+    @property
+    def degree(self) -> int:
+        return 0
+
+    @property
+    def name(self) -> str:
+        return f"Uniform{self.domain.name}_{wp.types.get_type_code(self.dtype)}"
+
+    def _make_eval_inner(self):
+        @cache.dynamic_func(suffix=self.name)
+        def eval_inner(args: self.ElementEvalArg, s: Sample):
+            return args.eval_arg.value
+
+        return eval_inner
+
+    def _make_eval_zero(self, dtype):
+        if dtype is None:
+            return None
+
+        scalar_type = wp.types.type_scalar_type(dtype)
+
+        @cache.dynamic_func(suffix=f"{self.name}_{wp.types.get_type_code(dtype)}")
+        def eval_zero(args: self.ElementEvalArg, s: Sample):
+            return dtype(scalar_type(0.0))
+
+        return eval_zero
+
+    def _make_eval_arg(self):
+        @cache.dynamic_struct(suffix=self.name)
+        class EvalArg:
+            value: self.dtype
+
+        return EvalArg
+
+    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_degree(self):
+        @cache.dynamic_func(suffix=self.name)
+        def degree(args: self.ElementEvalArg):
+            return 0
+
+        return degree
+
+    def trace(self):
+        if self.element_kind == ElementKind.SIDE:
+            raise RuntimeError("Trace only available for field defined on cell elements")
+
+        return UniformField(domain=Sides(self.domain.geometry_partition), value=self.value)
+
+
+class NonconformingField(GeometryField):
+    """Field defined as the map of a DiscreteField over a non-conforming geometry.
+
+    Args:
+        domain: The new domain over which the nonconforming field will be evaluated
+        field: Nonconforming discrete field
+        background: Uniform value or domain-conforming field determining the value outside of the geometry of definition of `field`
+    """
+
+    _LOOKUP_EPS = wp.constant(1.0e-6)
+
+    def __init__(self, domain: GeometryDomain, field: DiscreteField, background: Any = 0.0):
+        self.domain = domain
+
+        self.field = field
+        self.dtype = field.dtype
+
+        if not isinstance(background, GeometryField):
+            background = UniformField(domain, self.dtype(background))
+        elif background.geometry != domain.geometry or background.element_kind != domain.element_kind:
+            raise ValueError("Background field must be conforming to the domain")
+        self.background = background
+
+        self.EvalArg = self._make_eval_arg()
+        self.ElementEvalArg = self._make_element_eval_arg()
+        self.eval_degree = self._make_eval_degree()
+
+        self.eval_inner = self._make_nonconforming_eval("eval_inner")
+        self.eval_grad_inner = self._make_nonconforming_eval("eval_grad_inner")
+        self.eval_div_inner = self._make_nonconforming_eval("eval_div_inner")
+        self.eval_reference_grad_inner = self._make_eval_reference_grad()
+
+        # Nonconforming evaluation is position based, does not handle discontinuous fields
+        self.eval_outer = self.eval_inner
+        self.eval_grad_outer = self.eval_grad_inner
+        self.eval_div_outer = self.eval_div_inner
+        self.eval_reference_grad_outer = self.eval_reference_grad_inner
+
+    @property
+    def geometry(self) -> Geometry:
+        return self.domain.geometry
+
+    @property
+    def element_kind(self) -> ElementKind:
+        return self.domain.element_kind
+
+    @cache.cached_arg_value
+    def eval_arg_value(self, device):
+        arg = self.EvalArg()
+        arg.field_cell_eval_arg = self.field.ElementEvalArg()
+        arg.field_cell_eval_arg.elt_arg = self.field.geometry.cell_arg_value(device)
+        arg.field_cell_eval_arg.eval_arg = self.field.eval_arg_value(device)
+        arg.background_arg = self.background.eval_arg_value(device)
+        return arg
+
+    @property
+    def degree(self) -> int:
+        return self.field.degree
+
+    @property
+    def name(self) -> str:
+        return f"{self.domain.name}_{self.field.name}_{self.background.name}"
+
+    def _make_nonconforming_eval(self, eval_func_name):
+        field_eval = getattr(self.field, eval_func_name)
+        bg_eval = getattr(self.background, eval_func_name)
+
+        if field_eval is None or bg_eval is None:
+            return None
+
+        @cache.dynamic_func(suffix=f"{eval_func_name}_{self.name}")
+        def eval_nc(args: self.ElementEvalArg, s: Sample):
+            pos = self.domain.element_position(args.elt_arg, s)
+            cell_arg = args.eval_arg.field_cell_eval_arg.elt_arg
+            nonconforming_s = self.field.geometry.cell_lookup(cell_arg, pos)
+            if (
+                nonconforming_s.element_index == NULL_ELEMENT_INDEX
+                or wp.length_sq(pos - self.field.geometry.cell_position(cell_arg, nonconforming_s))
+                > NonconformingField._LOOKUP_EPS
+            ):
+                return bg_eval(self.background.ElementEvalArg(args.elt_arg, args.eval_arg.background_arg), s)
+            return field_eval(
+                self.field.ElementEvalArg(cell_arg, args.eval_arg.field_cell_eval_arg.eval_arg), nonconforming_s
+            )
+
+        return eval_nc
+
+    def _make_eval_reference_grad(self):
+        if self.eval_grad_inner is None:
+            return None
+
+        @cache.dynamic_func(suffix=f"{self.eval_grad_inner.key}")
+        def eval_reference_grad_inner(args: self.ElementEvalArg, s: Sample):
+            return self.eval_grad_inner(args, s) * self.domain.element_deformation_gradient(args.elt_arg, s)
+
+        return eval_reference_grad_inner
+
+    def _make_eval_arg(self):
+        @cache.dynamic_struct(suffix=self.name)
+        class EvalArg:
+            field_cell_eval_arg: self.field.ElementEvalArg
+            background_arg: self.background.EvalArg
+
+        return EvalArg
+
+    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_degree(self):
+        @cache.dynamic_func(suffix=self.name)
+        def degree(args: self.ElementEvalArg):
+            return self.field.eval_degree(args.eval_arg.field_cell_eval_arg)
+
+        return degree
+
+    def trace(self):
+        if self.element_kind == ElementKind.SIDE:
+            raise RuntimeError("Trace only available for field defined on cell elements")
+
+        return NonconformingField(
+            domain=Sides(self.domain.geometry_partition), field=self.field, background=self.background.trace()
+        )

warp/fem/field/nodal_field.py

@@ -28,6 +28,7 @@ class NodalFieldBase(DiscreteField):
         self.eval_div_outer = self._make_eval_div_outer()
 
         self.set_node_value = self._make_set_node_value()
+        self.node_partition_index = self._make_node_partition_index()
 
     def _make_eval_arg(self):
         @cache.dynamic_struct(suffix=self.name)
@@ -62,14 +63,15 @@ class NodalFieldBase(DiscreteField):
         return read_node_value
 
     def _make_eval_inner(self):
-        NODES_PER_ELEMENT = self.space.topology.NODES_PER_ELEMENT
-
         @cache.dynamic_func(suffix=self.name)
         def eval_inner(args: self.ElementEvalArg, s: Sample):
             res = self.space.element_inner_weight(
                 args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, 0
             ) * self._read_node_value(args, s.element_index, 0)
-            for k in range(1, NODES_PER_ELEMENT):
+            node_count = self.space.topology.element_node_count(
+                args.elt_arg, args.eval_arg.topology_arg, s.element_index
+            )
+            for k in range(1, node_count):
                 res += self.space.element_inner_weight(
                     args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, k
                 ) * self._read_node_value(args, s.element_index, k)
@@ -78,8 +80,6 @@ class NodalFieldBase(DiscreteField):
         return eval_inner
 
     def _make_eval_grad_inner(self, world_space: bool):
-        NODES_PER_ELEMENT = self.space.topology.NODES_PER_ELEMENT
-
         if not self.gradient_valid():
             return None
 
@@ -91,7 +91,10 @@ class NodalFieldBase(DiscreteField):
                     args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, 0
                 ),
             )
-            for k in range(1, NODES_PER_ELEMENT):
+            node_count = self.space.topology.element_node_count(
+                args.elt_arg, args.eval_arg.topology_arg, s.element_index
+            )
+            for k in range(1, node_count):
                 res += utils.generalized_outer(
                     self._read_node_value(args, s.element_index, k),
                     self.space.element_inner_weight_gradient(
@@ -109,8 +112,6 @@ class NodalFieldBase(DiscreteField):
         return eval_grad_inner_world_space if world_space else eval_grad_inner_ref_space
 
     def _make_eval_div_inner(self):
-        NODES_PER_ELEMENT = self.space.topology.NODES_PER_ELEMENT
-
         if not self.divergence_valid():
             return None
 
@@ -128,7 +129,10 @@ class NodalFieldBase(DiscreteField):
                 ),
             )
 
-            for k in range(1, NODES_PER_ELEMENT):
+            node_count = self.space.topology.element_node_count(
+                args.elt_arg, args.eval_arg.topology_arg, s.element_index
+            )
+            for k in range(1, node_count):
                 res += utils.generalized_inner(
                     self._read_node_value(args, s.element_index, k),
                     utils.apply_right(
@@ -143,8 +147,6 @@ class NodalFieldBase(DiscreteField):
         return eval_div_inner
 
     def _make_eval_outer(self):
-        NODES_PER_ELEMENT = self.space.topology.NODES_PER_ELEMENT
-
         @cache.dynamic_func(suffix=self.name)
         def eval_outer(args: self.ElementEvalArg, s: Sample):
             res = self.space.element_outer_weight(
@@ -154,7 +156,10 @@ class NodalFieldBase(DiscreteField):
                 s.element_coords,
                 0,
             ) * self._read_node_value(args, s.element_index, 0)
-            for k in range(1, NODES_PER_ELEMENT):
+            node_count = self.space.topology.element_node_count(
+                args.elt_arg, args.eval_arg.topology_arg, s.element_index
+            )
+            for k in range(1, node_count):
                 res += self.space.element_outer_weight(
                     args.elt_arg,
                     args.eval_arg.space_arg,
@@ -167,8 +172,6 @@ class NodalFieldBase(DiscreteField):
         return eval_outer
 
     def _make_eval_grad_outer(self, world_space: bool):
-        NODES_PER_ELEMENT = self.space.topology.NODES_PER_ELEMENT
-
         if not self.gradient_valid():
             return None
 
@@ -180,7 +183,10 @@ class NodalFieldBase(DiscreteField):
                     args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, 0
                 ),
             )
-            for k in range(1, NODES_PER_ELEMENT):
+            node_count = self.space.topology.element_node_count(
+                args.elt_arg, args.eval_arg.topology_arg, s.element_index
+            )
+            for k in range(1, node_count):
                 res += utils.generalized_outer(
                     self._read_node_value(args, s.element_index, k),
                     self.space.element_outer_weight_gradient(
@@ -198,8 +204,6 @@ class NodalFieldBase(DiscreteField):
         return eval_grad_outer_world_space if world_space else eval_grad_outer_ref_space
 
     def _make_eval_div_outer(self):
-        NODES_PER_ELEMENT = self.space.topology.NODES_PER_ELEMENT
-
         if not self.divergence_valid():
             return None
 
@@ -216,7 +220,11 @@ class NodalFieldBase(DiscreteField):
                     grad_transform,
                 ),
             )
-            for k in range(1, NODES_PER_ELEMENT):
+
+            node_count = self.space.topology.element_node_count(
+                args.elt_arg, args.eval_arg.topology_arg, s.element_index
+            )
+            for k in range(1, node_count):
                 res += utils.generalized_inner(
                     self._read_node_value(args, s.element_index, k),
                     utils.apply_right(
@@ -237,6 +245,13 @@ class NodalFieldBase(DiscreteField):
 
         return set_node_value
 
+    def _make_node_partition_index(self):
+        @cache.dynamic_func(suffix=self.name)
+        def node_partition_index(args: self.EvalArg, node_index: int):
+            return self.space_partition.partition_node_index(args.eval_arg.partition_arg, node_index)
+
+        return node_partition_index
+
 
 class NodalField(NodalFieldBase):
     """A field holding values for all degrees of freedom at each node of the underlying function space partition