warp-lang 1.4.1__py3-none-manylinux2014_aarch64.whl → 1.5.0__py3-none-manylinux2014_aarch64.whl

warp/fem/cache.py

@@ -6,6 +6,7 @@ from copy import copy
 from typing import Any, Callable, Dict, Optional, Tuple, Union
 
 import warp as wp
+from warp.fem.operator import Integrand
 
 _kernel_cache = {}
 _struct_cache = {}
@@ -186,7 +187,7 @@ class ExpandStarredArgumentStruct(ast.NodeTransformer):
 
 
 def get_integrand_function(
-    integrand: "warp.fem.operator.Integrand",  # noqa: F821
+    integrand: Integrand,
     suffix: str,
     func=None,
     annotations=None,
@@ -208,27 +209,30 @@ def get_integrand_function(
 
 
 def get_integrand_kernel(
-    integrand: "warp.fem.operator.Integrand",  # noqa: F821
+    integrand: Integrand,
     suffix: str,
     kernel_fn: Optional[Callable] = None,
     kernel_options: Dict[str, Any] = None,
     code_transformers=None,
 ):
-    if kernel_options is None:
-        kernel_options = {}
+    options = integrand.module.options.copy()
+    options.update(integrand.kernel_options)
+    if kernel_options is not None:
+        options.update(kernel_options)
 
-    key = _make_key(integrand.func, suffix, use_qualified_name=True)
+    kernel_key = _make_key(integrand.func, suffix, use_qualified_name=True)
+    opts_key = "".join([f"{k}:{v}" for k, v in sorted(options.items())])
+    cache_key = kernel_key + opts_key
 
-    if key not in _kernel_cache:
+    if cache_key not in _kernel_cache:
         if kernel_fn is None:
             return None
 
         module = wp.get_module(f"{integrand.module.name}.{integrand.name}")
-        module.options = copy(integrand.module.options)
-        module.options.update(kernel_options)
-
-        _kernel_cache[key] = wp.Kernel(func=kernel_fn, key=key, module=module, code_transformers=code_transformers)
-    return _kernel_cache[key]
+        _kernel_cache[cache_key] = wp.Kernel(
+            func=kernel_fn, key=kernel_key, module=module, code_transformers=code_transformers, options=options
+        )
+    return _kernel_cache[cache_key]
 
 
 def cached_arg_value(func: Callable):
@@ -478,7 +482,7 @@ def borrow_temporary(
     if temporary_store is None:
         temporary_store = TemporaryStore._default_store
 
-    if temporary_store is None:
+    if temporary_store is None or (requires_grad and wp.context.runtime.tape is not None):
         return Temporary(
             array=wp.empty(shape=shape, dtype=dtype, pinned=pinned, device=device, requires_grad=requires_grad)
         )

warp/fem/dirichlet.py

@@ -1,7 +1,7 @@
 from typing import Any, Optional
 
 import warp as wp
-from warp.fem.utils import array_axpy, symmetric_eigenvalues_qr
+from warp.fem.linalg import array_axpy, symmetric_eigenvalues_qr
 from warp.sparse import BsrMatrix, bsr_assign, bsr_axpy, bsr_copy, bsr_mm, bsr_mv
 from warp.types import type_is_matrix, type_length
 

warp/fem/domain.py

@@ -1,4 +1,4 @@
-from typing import Optional, Union
+from typing import Any, Optional, Set, Union
 
 import warp as wp
 import warp.codegen
@@ -11,6 +11,7 @@ from warp.fem.geometry import (
     GeometryPartition,
     WholeGeometryPartition,
 )
+from warp.fem.operator import Operator
 from warp.fem.types import ElementKind
 
 GeometryOrPartition = Union[Geometry, GeometryPartition]
@@ -94,6 +95,10 @@ class GeometryDomain:
     element_lookup: wp.Function
     """Device function returning the sample point corresponding to a world position"""
 
+    def notify_operator_usage(self, ops: Set[Operator]):
+        """Makes the Domain aware that the operators `ops` will be applied"""
+        pass
+
 
 class Cells(GeometryDomain):
     """A Domain containing all cells of the geometry or geometry partition"""
@@ -160,6 +165,17 @@ class Cells(GeometryDomain):
     def element_lookup(self) -> wp.Function:
         return self.geometry.cell_lookup
 
+    @property
+    def domain_cell_arg(self) -> wp.Function:
+        return Cells._identity_fn
+
+    def cell_domain(self):
+        return self
+
+    @wp.func
+    def _identity_fn(x: Any):
+        return x
+
 
 class Sides(GeometryDomain):
     """A Domain containing all (interior and boundary) sides of the geometry or geometry partition"""
@@ -225,6 +241,33 @@ class Sides(GeometryDomain):
     def element_normal(self) -> wp.Function:
         return self.geometry.side_normal
 
+    @property
+    def element_inner_cell_index(self) -> wp.Function:
+        return self.geometry.side_inner_cell_index
+
+    @property
+    def element_outer_cell_index(self) -> wp.Function:
+        return self.geometry.side_outer_cell_index
+
+    @property
+    def element_inner_cell_coords(self) -> wp.Function:
+        return self.geometry.side_inner_cell_coords
+
+    @property
+    def element_outer_cell_coords(self) -> wp.Function:
+        return self.geometry.side_outer_cell_coords
+
+    @property
+    def cell_to_element_coords(self) -> wp.Function:
+        return self.geometry.side_from_cell_coords
+
+    @property
+    def domain_cell_arg(self) -> wp.Function:
+        return self.geometry.side_to_cell_arg
+
+    def cell_domain(self):
+        return Cells(self.geometry_partition)
+
 
 class BoundarySides(Sides):
     """A Domain containing boundary sides of the geometry or geometry partition"""

warp/fem/field/__init__.py

@@ -6,8 +6,7 @@ from warp.fem.space import FunctionSpace, SpacePartition, SpaceRestriction, make
 from .field import DiscreteField, FieldLike, GeometryField, ImplicitField, NonconformingField, SpaceField, UniformField
 from .nodal_field import NodalField
 from .restriction import FieldRestriction
-from .test import TestField
-from .trial import TrialField
+from .virtual import LocalTestField, LocalTrialField, TestField, TrialField
 
 
 def make_restriction(

warp/fem/field/field.py

@@ -1,10 +1,10 @@
-from typing import Any, Dict, Optional
+from typing import Any, Dict, Optional, Set
 
 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.operator import Operator, integrand
 from warp.fem.space import FunctionSpace, SpacePartition
 from warp.fem.types import NULL_ELEMENT_INDEX, ElementKind, Sample
 
@@ -48,32 +48,32 @@ class FieldLike:
         return False
 
     @staticmethod
-    def eval_inner(args: "ElementEvalArg", s: "Sample"):  # noqa: F821
+    def eval_inner(args: "ElementEvalArg", s: Sample):  # noqa: F821
         """Device function evaluating the inner field value at a sample point"""
         raise NotImplementedError
 
     @staticmethod
-    def eval_grad_inner(args: "ElementEvalArg", s: "Sample"):  # noqa: F821
+    def eval_grad_inner(args: "ElementEvalArg", s: Sample):  # noqa: F821
         """Device function evaluating the inner field gradient at a sample point"""
         raise NotImplementedError
 
     @staticmethod
-    def eval_div_inner(args: "ElementEvalArg", s: "Sample"):  # noqa: F821
+    def eval_div_inner(args: "ElementEvalArg", s: Sample):  # noqa: F821
         """Device function evaluating the inner field divergence at a sample point"""
         raise NotImplementedError
 
     @staticmethod
-    def eval_outer(args: "ElementEvalArg", s: "Sample"):  # noqa: F821
+    def eval_outer(args: "ElementEvalArg", s: Sample):  # noqa: F821
         """Device function evaluating the outer field value at a sample point"""
         raise NotImplementedError
 
     @staticmethod
-    def eval_grad_outer(args: "ElementEvalArg", s: "Sample"):  # noqa: F821
+    def eval_grad_outer(args: "ElementEvalArg", s: Sample):  # noqa: F821
         """Device function evaluating the outer field gradient at a sample point"""
         raise NotImplementedError
 
     @staticmethod
-    def eval_div_outer(args: "ElementEvalArg", s: "Sample"):  # noqa: F821
+    def eval_div_outer(args: "ElementEvalArg", s: Sample):  # noqa: F821
         """Device function evaluating the outer field divergence at a sample point"""
         raise NotImplementedError
 
@@ -82,6 +82,10 @@ class FieldLike:
         """Polynomial degree of the field is applicable, or hint for determination of interpolation order"""
         raise NotImplementedError
 
+    def notify_operator_usage(self, ops: Set[Operator]):
+        """Makes the Domain aware that the operators `ops` will be applied"""
+        pass
+
 
 class GeometryField(FieldLike):
     """Base class for fields defined over a geometry"""
@@ -97,12 +101,12 @@ class GeometryField(FieldLike):
         raise NotImplementedError
 
     @staticmethod
-    def eval_reference_grad_inner(args: "ElementEvalArg", s: "Sample"):  # noqa: F821
+    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
+    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
 
@@ -128,6 +132,9 @@ class SpaceField(GeometryField):
         self._space = space
         self._space_partition = space_partition
 
+        self.gradient_valid = self.space.gradient_valid
+        self.divergence_valid = self.space.divergence_valid
+
     @property
     def geometry(self) -> Geometry:
         return self._space.geometry
@@ -156,17 +163,22 @@ class SpaceField(GeometryField):
     def dof_dtype(self) -> type:
         return self.space.dof_dtype
 
-    def gradient_valid(self) -> bool:
-        """Whether gradient operator can be computed. Only for scalar and vector fields as higher-order tensors are not support yet"""
-        return not wp.types.type_is_matrix(self.dtype)
+    @property
+    def gradient_dtype(self):
+        """Return type of the gradient operator. Assumes self.gradient_valid()"""
+        if wp.types.type_is_vector(self.dtype):
+            return cache.cached_mat_type(
+                shape=(wp.types.type_length(self.dtype), self.geometry.dimension),
+                dtype=wp.types.type_scalar_type(self.dtype),
+            )
+        return cache.cached_vec_type(length=self.geometry.dimension, dtype=wp.types.type_scalar_type(self.dtype))
 
-    def divergence_valid(self) -> bool:
-        """Whether divergence of this field can be computed. Only for vector and tensor fields with same dimension as embedding geometry"""
+    @property
+    def divergence_dtype(self):
+        """Return type of the divergence operator. Assumes self.gradient_valid()"""
         if wp.types.type_is_vector(self.dtype):
-            return wp.types.type_length(self.dtype) == self.space.geometry.dimension
-        if wp.types.type_is_matrix(self.dtype):
-            return self.dtype._shape_[0] == self.space.geometry.dimension
-        return False
+            return wp.types.type_scalar_type(self.dtype)
+        return cache.cached_vec_type(length=self.dtype._shape_[1], dtype=wp.types.type_scalar_type(self.dtype))
 
     def _make_eval_degree(self):
         ORDER = self.space.ORDER

warp/fem/field/nodal_field.py

@@ -1,5 +1,7 @@
+from typing import Any
+
 import warp as wp
-from warp.fem import cache, utils
+from warp.fem import cache
 from warp.fem.space import CollocatedFunctionSpace, SpacePartition
 from warp.fem.types import NULL_NODE_INDEX, ElementIndex, Sample
 
@@ -56,58 +58,79 @@ class NodalFieldBase(DiscreteField):
             )
             pidx = self.space_partition.partition_node_index(args.eval_arg.partition_arg, nidx)
             if pidx == NULL_NODE_INDEX:
-                return self.space.dtype(0.0)
+                return self.space.dof_dtype(0.0)
 
-            return self.space.dof_mapper.dof_to_value(args.eval_arg.dof_values[pidx])
+            return args.eval_arg.dof_values[pidx]
 
         return read_node_value
 
     def _make_eval_inner(self):
         @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)
+            local_value_map = self.space.local_value_map_inner(args.elt_arg, s.element_index, s.element_coords)
+            res = self.space.space_value(
+                self._read_node_value(args, s.element_index, 0),
+                self.space.element_inner_weight(
+                    args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, 0
+                ),
+                local_value_map,
+            )
+
             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)
+                res += self.space.space_value(
+                    self._read_node_value(args, s.element_index, k),
+                    self.space.element_inner_weight(
+                        args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, k
+                    ),
+                    local_value_map,
+                )
             return res
 
         return eval_inner
 
     def _make_eval_grad_inner(self, world_space: bool):
-        if not self.gradient_valid():
+        if not self.space.gradient_valid():
             return None
 
         @cache.dynamic_func(suffix=self.name)
-        def eval_grad_inner_ref_space(args: self.ElementEvalArg, s: Sample):
-            res = utils.generalized_outer(
+        def eval_grad_inner(args: self.ElementEvalArg, s: Sample, grad_transform: Any):
+            local_value_map = self.space.local_value_map_inner(args.elt_arg, s.element_index, s.element_coords)
+
+            res = self.space.space_gradient(
                 self._read_node_value(args, s.element_index, 0),
                 self.space.element_inner_weight_gradient(
                     args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, 0
                 ),
+                local_value_map,
+                grad_transform,
             )
+
             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(
+                res += self.space.space_gradient(
                     self._read_node_value(args, s.element_index, k),
                     self.space.element_inner_weight_gradient(
                         args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, k
                     ),
+                    local_value_map,
+                    grad_transform,
                 )
             return res
 
+        @cache.dynamic_func(suffix=self.name)
+        def eval_grad_inner_ref_space(args: self.ElementEvalArg, s: Sample):
+            grad_transform = 1.0
+            return eval_grad_inner(args, s, grad_transform)
+
         @cache.dynamic_func(suffix=self.name)
         def eval_grad_inner_world_space(args: self.ElementEvalArg, s: Sample):
             grad_transform = self.space.element_inner_reference_gradient_transform(args.elt_arg, s)
-            res = eval_grad_inner_ref_space(args, s)
-            return res * grad_transform
+            return eval_grad_inner(args, s, grad_transform)
 
         return eval_grad_inner_world_space if world_space else eval_grad_inner_ref_space
 
@@ -118,25 +141,28 @@ class NodalFieldBase(DiscreteField):
         @cache.dynamic_func(suffix=self.name)
         def eval_div_inner(args: self.ElementEvalArg, s: Sample):
             grad_transform = self.space.element_inner_reference_gradient_transform(args.elt_arg, s)
+            local_value_map = self.space.local_value_map_inner(args.elt_arg, s.element_index, s.element_coords)
 
-            res = utils.generalized_inner(
+            res = self.space.space_divergence(
                 self._read_node_value(args, s.element_index, 0),
                 self.space.element_inner_weight_gradient(
                     args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, 0
-                )
-                * grad_transform,
+                ),
+                local_value_map,
+                grad_transform,
             )
 
             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(
+                res += self.space.space_divergence(
                     self._read_node_value(args, s.element_index, k),
                     self.space.element_inner_weight_gradient(
                         args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, k
-                    )
-                    * grad_transform,
+                    ),
+                    local_value_map,
+                    grad_transform,
                 )
             return res
 
@@ -145,57 +171,71 @@ class NodalFieldBase(DiscreteField):
     def _make_eval_outer(self):
         @cache.dynamic_func(suffix=self.name)
         def eval_outer(args: self.ElementEvalArg, s: Sample):
-            res = self.space.element_outer_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)
+            local_value_map = self.space.local_value_map_outer(args.elt_arg, s.element_index, s.element_coords)
+            res = self.space.space_value(
+                self._read_node_value(args, s.element_index, 0),
+                self.space.element_outer_weight(
+                    args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, 0
+                ),
+                local_value_map,
+            )
+
             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,
-                    s.element_index,
-                    s.element_coords,
-                    k,
-                ) * self._read_node_value(args, s.element_index, k)
+                res += self.space.space_value(
+                    self._read_node_value(args, s.element_index, k),
+                    self.space.element_outer_weight(
+                        args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, k
+                    ),
+                    local_value_map,
+                )
             return res
 
         return eval_outer
 
     def _make_eval_grad_outer(self, world_space: bool):
-        if not self.gradient_valid():
+        if not self.space.gradient_valid():
             return None
 
         @cache.dynamic_func(suffix=self.name)
-        def eval_grad_outer_ref_space(args: self.ElementEvalArg, s: Sample):
-            res = utils.generalized_outer(
+        def eval_grad_outer(args: self.ElementEvalArg, s: Sample, grad_transform: Any):
+            local_value_map = self.space.local_value_map_outer(args.elt_arg, s.element_index, s.element_coords)
+
+            res = self.space.space_gradient(
                 self._read_node_value(args, s.element_index, 0),
                 self.space.element_outer_weight_gradient(
                     args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, 0
                 ),
+                local_value_map,
+                grad_transform,
             )
+
             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(
+                res += self.space.space_gradient(
                     self._read_node_value(args, s.element_index, k),
                     self.space.element_outer_weight_gradient(
                         args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, k
                     ),
+                    local_value_map,
+                    grad_transform,
                 )
             return res
 
+        @cache.dynamic_func(suffix=self.name)
+        def eval_grad_outer_ref_space(args: self.ElementEvalArg, s: Sample):
+            grad_transform = 1.0
+            return eval_grad_outer_ref_space(args, s, grad_transform)
+
         @cache.dynamic_func(suffix=self.name)
         def eval_grad_outer_world_space(args: self.ElementEvalArg, s: Sample):
             grad_transform = self.space.element_outer_reference_gradient_transform(args.elt_arg, s)
-            res = eval_grad_outer_ref_space(args, s)
-            return res * grad_transform
+            return eval_grad_outer_ref_space(args, s, grad_transform)
 
         return eval_grad_outer_world_space if world_space else eval_grad_outer_ref_space
 
@@ -206,25 +246,28 @@ class NodalFieldBase(DiscreteField):
         @cache.dynamic_func(suffix=self.name)
         def eval_div_outer(args: self.ElementEvalArg, s: Sample):
             grad_transform = self.space.element_outer_reference_gradient_transform(args.elt_arg, s)
+            local_value_map = self.space.local_value_map_outer(args.elt_arg, s.element_index, s.element_coords)
 
-            res = utils.generalized_inner(
+            res = self.space.space_divergence(
                 self._read_node_value(args, s.element_index, 0),
                 self.space.element_outer_weight_gradient(
                     args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, 0
-                )
-                * grad_transform,
+                ),
+                local_value_map,
+                grad_transform,
             )
 
             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(
+                res += self.space.space_divergence(
                     self._read_node_value(args, s.element_index, k),
                     self.space.element_outer_weight_gradient(
                         args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, k
-                    )
-                    * grad_transform,
+                    ),
+                    local_value_map,
+                    grad_transform,
                 )
             return res
 
@@ -232,8 +275,17 @@ class NodalFieldBase(DiscreteField):
 
     def _make_set_node_value(self):
         @cache.dynamic_func(suffix=self.name)
-        def set_node_value(args: self.EvalArg, partition_node_index: int, value: self.space.dtype):
-            args.dof_values[partition_node_index] = self.space.dof_mapper.value_to_dof(value)
+        def set_node_value(
+            elt_arg: self.space.ElementArg,
+            eval_arg: self.EvalArg,
+            element_index: ElementIndex,
+            node_index_in_element: int,
+            partition_node_index: int,
+            value: self.space.dtype,
+        ):
+            eval_arg.dof_values[partition_node_index] = self.space.node_dof_value(
+                elt_arg, eval_arg.space_arg, element_index, node_index_in_element, value
+            )
 
         return set_node_value