warp-lang 1.7.2__py3-none-macosx_10_13_universal2.whl → 1.8.0__py3-none-macosx_10_13_universal2.whl

warp/fem/field/nodal_field.py

@@ -13,12 +13,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any
+from typing import Any, ClassVar
 
 import warp as wp
 from warp.fem import cache
 from warp.fem.space import CollocatedFunctionSpace, SpacePartition
 from warp.fem.types import NULL_NODE_INDEX, ElementIndex, Sample
+from warp.fem.utils import type_zero_element
 
 from .field import DiscreteField
 
@@ -26,26 +27,29 @@ from .field import DiscreteField
 class NodalFieldBase(DiscreteField):
     """Base class for nodal field and nodal field traces. Does not hold values"""
 
+    _dynamic_attribute_constructors: ClassVar = {
+        "EvalArg": lambda obj: obj._make_eval_arg(),
+        "ElementEvalArg": lambda obj: obj._make_element_eval_arg(),
+        "eval_degree": lambda obj: DiscreteField._make_eval_degree(obj),
+        "_read_node_value": lambda obj: obj._make_read_node_value(),
+        "eval_inner": lambda obj: obj._make_eval_inner(),
+        "eval_outer": lambda obj: obj._make_eval_outer(),
+        "eval_grad_inner": lambda obj: obj._make_eval_grad_inner(world_space=True),
+        "eval_grad_outer": lambda obj: obj._make_eval_grad_outer(world_space=True),
+        "eval_reference_grad_inner": lambda obj: obj._make_eval_grad_inner(world_space=False),
+        "eval_reference_grad_outer": lambda obj: obj._make_eval_grad_outer(world_space=False),
+        "eval_div_inner": lambda obj: obj._make_eval_div_inner(),
+        "eval_div_outer": lambda obj: obj._make_eval_div_outer(),
+        "set_node_value": lambda obj: obj._make_set_node_value(),
+        "node_partition_index": lambda obj: obj._make_node_partition_index(),
+        "node_count": lambda obj: obj._make_node_count(),
+        "node_index": lambda obj: obj._make_node_index(),
+        "at_node": lambda obj: obj._make_at_node(),
+    }
+
     def __init__(self, space: CollocatedFunctionSpace, space_partition: SpacePartition):
         super().__init__(space, space_partition)
-
-        self.EvalArg = self._make_eval_arg()
-        self.ElementEvalArg = self._make_element_eval_arg()
-        self.eval_degree = DiscreteField._make_eval_degree(self)
-
-        self._read_node_value = self._make_read_node_value()
-
-        self.eval_inner = self._make_eval_inner()
-        self.eval_outer = self._make_eval_outer()
-        self.eval_grad_inner = self._make_eval_grad_inner(world_space=True)
-        self.eval_grad_outer = self._make_eval_grad_outer(world_space=True)
-        self.eval_reference_grad_inner = self._make_eval_grad_inner(world_space=False)
-        self.eval_reference_grad_outer = self._make_eval_grad_outer(world_space=False)
-        self.eval_div_inner = self._make_eval_div_inner()
-        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()
+        cache.setup_dynamic_attributes(self)
 
     def _make_eval_arg(self):
         @cache.dynamic_struct(suffix=self.name)
@@ -66,6 +70,8 @@ class NodalFieldBase(DiscreteField):
         return ElementEvalArg
 
     def _make_read_node_value(self):
+        zero_element = type_zero_element(self.dof_dtype)
+
         @cache.dynamic_func(suffix=self.name)
         def read_node_value(args: self.ElementEvalArg, geo_element_index: ElementIndex, node_index_in_elt: int):
             nidx = self.space.topology.element_node_index(
@@ -73,26 +79,29 @@ class NodalFieldBase(DiscreteField):
             )
             pidx = self.space_partition.partition_node_index(args.eval_arg.partition_arg, nidx)
             if pidx == NULL_NODE_INDEX:
-                return self.space.dof_dtype(0.0)
+                return zero_element()
 
             return args.eval_arg.dof_values[pidx]
 
         return read_node_value
 
     def _make_eval_inner(self):
+        zero_element = type_zero_element(self.dtype)
+
         @cache.dynamic_func(suffix=self.name)
         def eval_inner(args: self.ElementEvalArg, s: Sample):
             local_value_map = self.space.local_value_map_inner(args.elt_arg, s.element_index, s.element_coords)
             node_count = self.space.topology.element_node_count(
                 args.elt_arg, args.eval_arg.topology_arg, s.element_index
             )
-            res = self.space.dtype(0.0)
+            res = zero_element()
             for k in range(node_count):
+                w = self.space.element_inner_weight(
+                    args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, k, s.qp_index
+                )
                 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, s.qp_index
-                    ),
+                    w,
                     local_value_map,
                 )
             return res
@@ -104,6 +113,7 @@ class NodalFieldBase(DiscreteField):
             return None
 
         gradient_dtype = self.gradient_dtype if world_space else self.reference_gradient_dtype
+        zero_element = type_zero_element(gradient_dtype)
 
         @cache.dynamic_func(suffix=f"{self.name}{world_space}")
         def eval_grad_inner(args: self.ElementEvalArg, s: Sample, grad_transform: Any):
@@ -112,7 +122,7 @@ class NodalFieldBase(DiscreteField):
                 args.elt_arg, args.eval_arg.topology_arg, s.element_index
             )
 
-            res = gradient_dtype(0.0)
+            res = zero_element()
             for k in range(node_count):
                 res += self.space.space_gradient(
                     self._read_node_value(args, s.element_index, k),
@@ -144,6 +154,7 @@ class NodalFieldBase(DiscreteField):
     def _make_eval_div_inner(self):
         if not self.divergence_valid():
             return None
+        zero_element = type_zero_element(self.divergence_dtype)
 
         @cache.dynamic_func(suffix=self.name)
         def eval_div_inner(args: self.ElementEvalArg, s: Sample):
@@ -153,7 +164,7 @@ class NodalFieldBase(DiscreteField):
                 args.elt_arg, args.eval_arg.topology_arg, s.element_index
             )
 
-            res = self.divergence_dtype(0.0)
+            res = zero_element()
             for k in range(node_count):
                 res += self.space.space_divergence(
                     self._read_node_value(args, s.element_index, k),
@@ -168,6 +179,8 @@ class NodalFieldBase(DiscreteField):
         return eval_div_inner
 
     def _make_eval_outer(self):
+        zero_element = type_zero_element(self.dtype)
+
         @cache.dynamic_func(suffix=self.name)
         def eval_outer(args: self.ElementEvalArg, s: Sample):
             local_value_map = self.space.local_value_map_outer(args.elt_arg, s.element_index, s.element_coords)
@@ -175,7 +188,7 @@ class NodalFieldBase(DiscreteField):
                 args.elt_arg, args.eval_arg.topology_arg, s.element_index
             )
 
-            res = self.dtype(0.0)
+            res = zero_element()
             for k in range(node_count):
                 res += self.space.space_value(
                     self._read_node_value(args, s.element_index, k),
@@ -193,6 +206,7 @@ class NodalFieldBase(DiscreteField):
             return None
 
         gradient_dtype = self.gradient_dtype if world_space else self.reference_gradient_dtype
+        zero_element = type_zero_element(gradient_dtype)
 
         @cache.dynamic_func(suffix=f"{self.name}{world_space}")
         def eval_grad_outer(args: self.ElementEvalArg, s: Sample, grad_transform: Any):
@@ -201,7 +215,7 @@ class NodalFieldBase(DiscreteField):
                 args.elt_arg, args.eval_arg.topology_arg, s.element_index
             )
 
-            res = gradient_dtype(0.0)
+            res = zero_element()
             for k in range(node_count):
                 res += self.space.space_gradient(
                     self._read_node_value(args, s.element_index, k),
@@ -234,6 +248,8 @@ class NodalFieldBase(DiscreteField):
         if not self.divergence_valid():
             return None
 
+        zero_element = type_zero_element(self.divergence_dtype)
+
         @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)
@@ -242,7 +258,7 @@ class NodalFieldBase(DiscreteField):
                 args.elt_arg, args.eval_arg.topology_arg, s.element_index
             )
 
-            res = self.divergence_dtype(0.0)
+            res = zero_element()
             for k in range(node_count):
                 res += self.space.space_divergence(
                     self._read_node_value(args, s.element_index, k),
@@ -279,6 +295,32 @@ class NodalFieldBase(DiscreteField):
 
         return node_partition_index
 
+    def _make_node_count(self):
+        @cache.dynamic_func(suffix=self.name)
+        def node_count(args: self.ElementEvalArg, s: Sample):
+            return self.space.topology.element_node_count(args.elt_arg, args.eval_arg.topology_arg, s.element_index)
+
+        return node_count
+
+    def _make_at_node(self):
+        @cache.dynamic_func(suffix=self.name)
+        def at_node(args: self.ElementEvalArg, s: Sample, node_index_in_elt: int):
+            node_coords = self.space.node_coords_in_element(
+                args.elt_arg, args.eval_arg.space_arg, s.element_index, node_index_in_elt
+            )
+            return Sample(s.element_index, node_coords, s.qp_index, s.qp_weight, s.test_dof, s.trial_dof)
+
+        return at_node
+
+    def _make_node_index(self):
+        @cache.dynamic_func(suffix=self.name)
+        def node_index(args: self.ElementEvalArg, s: Sample, node_index_in_elt: int):
+            return self.space.topology.element_node_index(
+                args.elt_arg, args.eval_arg.topology_arg, s.element_index, node_index_in_elt
+            )
+
+        return node_index
+
 
 class NodalField(NodalFieldBase):
     """A field holding values for all degrees of freedom at each node of the underlying function space partition
@@ -296,13 +338,15 @@ class NodalField(NodalFieldBase):
 
     def eval_arg_value(self, device):
         arg = self.EvalArg()
-        arg.dof_values = self._dof_values.to(device)
-        arg.space_arg = self.space.space_arg_value(device)
-        arg.partition_arg = self.space_partition.partition_arg_value(device)
-        arg.topology_arg = self.space.topology.topo_arg_value(device)
-
+        self.fill_eval_arg(arg, device)
         return arg
 
+    def fill_eval_arg(self, arg, device):
+        arg.dof_values = self._dof_values.to(device)
+        self.space.fill_space_arg(arg.space_arg, device)
+        self.space_partition.fill_partition_arg(arg.partition_arg, device)
+        self.space.topology.fill_topo_arg(arg.topology_arg, device)
+
     @property
     def dof_values(self) -> wp.array:
         """Returns a warp array containing the values at all degrees of freedom of the underlying space partition"""
@@ -328,13 +372,15 @@ class NodalField(NodalFieldBase):
 
         def eval_arg_value(self, device):
             arg = self.EvalArg()
-            arg.dof_values = self._field.dof_values.to(device)
-            arg.space_arg = self.space.space_arg_value(device)
-            arg.partition_arg = self.space_partition.partition_arg_value(device)
-            arg.topology_arg = self.space.topology.topo_arg_value(device)
-
+            self.fill_eval_arg(arg, device)
             return arg
 
+        def fill_eval_arg(self, arg, device):
+            arg.dof_values = self._field.dof_values.to(device)
+            self.space.fill_space_arg(arg.space_arg, device)
+            self.space_partition.fill_partition_arg(arg.partition_arg, device)
+            self.space.topology.fill_topo_arg(arg.topology_arg, device)
+
     def trace(self) -> Trace:
         trace_field = NodalField.Trace(self)
         return trace_field

warp/fem/field/virtual.py

@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Set
+from typing import Any, ClassVar, Dict, Set
 
 import warp as wp
 import warp.fem.operator as operator
@@ -23,6 +23,7 @@ from warp.fem.linalg import basis_coefficient, generalized_inner, generalized_ou
 from warp.fem.quadrature import Quadrature
 from warp.fem.space import FunctionSpace, SpacePartition, SpaceRestriction
 from warp.fem.types import NULL_NODE_INDEX, DofIndex, Sample, get_node_coord, get_node_index_in_element
+from warp.fem.utils import type_zero_element
 
 from .field import SpaceField
 
@@ -30,33 +31,52 @@ from .field import SpaceField
 class AdjointField(SpaceField):
     """Adjoint of a discrete field with respect to its degrees of freedom"""
 
-    def __init__(self, space: FunctionSpace, space_partition: SpaceRestriction):
+    _dynamic_attribute_constructors: ClassVar = {
+        "EvalArg": lambda obj: obj._make_eval_arg(),
+        "ElementEvalArg": lambda obj: obj._make_element_eval_arg(),
+        "eval_degree": lambda obj: obj._make_eval_degree(),
+        "eval_inner": lambda obj: obj._make_eval_inner(),
+        "eval_grad_inner": lambda obj: obj._make_eval_grad_inner(),
+        "eval_div_inner": lambda obj: obj._make_eval_div_inner(),
+        "eval_outer": lambda obj: obj._make_eval_outer(),
+        "eval_grad_outer": lambda obj: obj._make_eval_grad_outer(),
+        "eval_div_outer": lambda obj: obj._make_eval_div_outer(),
+        "node_count": lambda obj: obj._make_node_count(),
+        "at_node": lambda obj: obj._make_at_node(),
+        "node_index": lambda obj: obj._make_node_index(),
+    }
+
+    def __init__(self, space: FunctionSpace, space_partition: SpacePartition):
         super().__init__(space, space_partition=space_partition)
 
         self.node_dof_count = self.space.NODE_DOF_COUNT
         self.value_dof_count = self.space.VALUE_DOF_COUNT
 
-        self.EvalArg = self.space.SpaceArg
-        self.ElementEvalArg = self._make_element_eval_arg()
-
-        self.eval_arg_value = self.space.space_arg_value
-
-        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()
+        cache.setup_dynamic_attributes(self)
 
     @property
     def name(self) -> str:
         return f"{self.__class__.__name__}{self.space.name}{self._space_partition.name}"
 
-    def _make_element_eval_arg(self):
-        from warp.fem import cache
+    @cache.cached_arg_value
+    def eval_arg_value(self, device):
+        arg = self.EvalArg()
+        self.fill_eval_arg(arg, device)
+        return arg
+
+    def fill_eval_arg(self, arg, device):
+        self.space.fill_space_arg(arg.space_arg, device)
+        self.space.topology.fill_topo_arg(arg.topo_arg, device)
 
+    def _make_eval_arg(self):
+        @cache.dynamic_struct(suffix=self.name)
+        class EvalArg:
+            space_arg: self.space.SpaceArg
+            topo_arg: self.space.topology.TopologyArg
+
+        return EvalArg
+
+    def _make_element_eval_arg(self):
         @cache.dynamic_struct(suffix=self.name)
         class ElementEvalArg:
             elt_arg: self.space.topology.ElementArg
@@ -70,7 +90,7 @@ class AdjointField(SpaceField):
             dof = self._get_dof(s)
             node_weight = self.space.element_inner_weight(
                 args.elt_arg,
-                args.eval_arg,
+                args.eval_arg.space_arg,
                 s.element_index,
                 s.element_coords,
                 get_node_index_in_element(dof),
@@ -91,7 +111,7 @@ class AdjointField(SpaceField):
             dof = self._get_dof(s)
             nabla_weight = self.space.element_inner_weight_gradient(
                 args.elt_arg,
-                args.eval_arg,
+                args.eval_arg.space_arg,
                 s.element_index,
                 s.element_coords,
                 get_node_index_in_element(dof),
@@ -113,7 +133,7 @@ class AdjointField(SpaceField):
             dof = self._get_dof(s)
             nabla_weight = self.space.element_inner_weight_gradient(
                 args.elt_arg,
-                args.eval_arg,
+                args.eval_arg.space_arg,
                 s.element_index,
                 s.element_coords,
                 get_node_index_in_element(dof),
@@ -132,7 +152,7 @@ class AdjointField(SpaceField):
             dof = self._get_dof(s)
             node_weight = self.space.element_outer_weight(
                 args.elt_arg,
-                args.eval_arg,
+                args.eval_arg.space_arg,
                 s.element_index,
                 s.element_coords,
                 get_node_index_in_element(dof),
@@ -153,7 +173,7 @@ class AdjointField(SpaceField):
             dof = self._get_dof(s)
             nabla_weight = self.space.element_outer_weight_gradient(
                 args.elt_arg,
-                args.eval_arg,
+                args.eval_arg.space_arg,
                 s.element_index,
                 s.element_coords,
                 get_node_index_in_element(dof),
@@ -175,7 +195,7 @@ class AdjointField(SpaceField):
             dof = self._get_dof(s)
             nabla_weight = self.space.element_outer_weight_gradient(
                 args.elt_arg,
-                args.eval_arg,
+                args.eval_arg.space_arg,
                 s.element_index,
                 s.element_coords,
                 get_node_index_in_element(dof),
@@ -193,12 +213,30 @@ class AdjointField(SpaceField):
         def at_node(args: self.ElementEvalArg, s: Sample):
             dof = self._get_dof(s)
             node_coords = self.space.node_coords_in_element(
-                args.elt_arg, args.eval_arg, s.element_index, get_node_index_in_element(dof)
+                args.elt_arg, args.eval_arg.space_arg, s.element_index, get_node_index_in_element(dof)
             )
             return Sample(s.element_index, node_coords, s.qp_index, s.qp_weight, s.test_dof, s.trial_dof)
 
         return at_node
 
+    def _make_node_index(self):
+        @cache.dynamic_func(suffix=self.name)
+        def node_index(args: self.ElementEvalArg, s: Sample):
+            dof = self._get_dof(s)
+            node_idx = self.space.topology.element_node_index(
+                args.elt_arg, args.eval_arg.topo_arg, s.element_index, get_node_index_in_element(dof)
+            )
+            return node_idx
+
+        return node_index
+
+    def _make_node_count(self):
+        @cache.dynamic_func(suffix=self.name)
+        def node_count(args: self.ElementEvalArg, s: Sample):
+            return self.space.topology.element_node_count(args.elt_arg, args.eval_arg.topo_arg, s.element_index)
+
+        return node_count
+
 
 class TestField(AdjointField):
     """Field defined over a space restriction that can be used as a test function.
@@ -269,16 +307,16 @@ class LocalAdjointField(SpaceField):
     OUTER_GRAD_DOF = wp.constant(3)
     DOF_TYPE_COUNT = wp.constant(4)
 
-    _OP_DOF_MAP_CONTINUOUS = {
+    _OP_DOF_MAP_CONTINUOUS: ClassVar[Dict[operator.Operator, int]] = {
         operator.inner: INNER_DOF,
         operator.outer: INNER_DOF,
         operator.grad: INNER_GRAD_DOF,
-        operator.grad_outer: INNER_GRAD_DOF,
+        operator.grad_outer: OUTER_GRAD_DOF,
         operator.div: INNER_GRAD_DOF,
-        operator.div_outer: INNER_GRAD_DOF,
+        operator.div_outer: OUTER_GRAD_DOF,
     }
 
-    _OP_DOF_MAP_DISCONTINUOUS = {
+    _OP_DOF_MAP_DISCONTINUOUS: ClassVar[Dict[operator.Operator, int]] = {
         operator.inner: INNER_DOF,
         operator.outer: OUTER_DOF,
         operator.grad: INNER_GRAD_DOF,
@@ -293,6 +331,18 @@ class LocalAdjointField(SpaceField):
     class EvalArg:
         pass
 
+    _dynamic_attribute_constructors: ClassVar = {
+        "ElementEvalArg": lambda obj: obj._make_element_eval_arg(),
+        "eval_degree": lambda obj: obj._make_eval_degree(),
+        "_split_dof": lambda obj: obj._make_split_dof(),
+        "eval_inner": lambda obj: obj._make_eval_inner(),
+        "eval_grad_inner": lambda obj: obj._make_eval_grad_inner(),
+        "eval_div_inner": lambda obj: obj._make_eval_div_inner(),
+        "eval_outer": lambda obj: obj._make_eval_outer(),
+        "eval_grad_outer": lambda obj: obj._make_eval_grad_outer(),
+        "eval_div_outer": lambda obj: obj._make_eval_div_outer(),
+    }
+
     def __init__(self, field: AdjointField):
         # if not isinstance(field.space, CollocatedFunctionSpace):
         #     raise NotImplementedError("Local assembly only implemented for collocated function spaces")
@@ -305,9 +355,6 @@ class LocalAdjointField(SpaceField):
         self.value_dof_count = self.space.VALUE_DOF_COUNT
 
         self._dof_suffix = ""
-
-        self.ElementEvalArg = self._make_element_eval_arg()
-        self.eval_degree = self._make_eval_degree()
         self.at_node = None
 
         self._is_discontinuous = (self.space.element_inner_weight != self.space.element_outer_weight) or (
@@ -346,21 +393,7 @@ class LocalAdjointField(SpaceField):
         self._TAYLOR_DOF_COUNTS = dof_counts
 
         self._dof_suffix = "".join(str(c) for c in dof_counts)
-
-        self._split_dof = self._make_split_dof()
-
-        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()
-
-        if self._is_discontinuous:
-            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()
-        else:
-            self.eval_outer = self.eval_inner
-            self.eval_grad_outer = self.eval_grad_inner
-            self.eval_div_outer = self.eval_div_inner
+        cache.setup_dynamic_attributes(self)
 
     @property
     def name(self) -> str:
@@ -369,6 +402,9 @@ class LocalAdjointField(SpaceField):
     def eval_arg_value(self, device):
         return LocalAdjointField.EvalArg()
 
+    def fill_eval_arg(self, arg, device):
+        pass
+
     def _make_element_eval_arg(self):
         from warp.fem import cache
 
@@ -392,6 +428,7 @@ class LocalAdjointField(SpaceField):
 
     def _make_eval_inner(self):
         DOF_BEGIN = wp.constant(self._TAYLOR_DOF_OFFSETS[LocalAdjointField.INNER_DOF])
+        zero_element = type_zero_element(self.dtype)
 
         @cache.dynamic_func(suffix=self.name)
         def eval_test_inner(args: self.ElementEvalArg, s: Sample):
@@ -399,7 +436,7 @@ class LocalAdjointField(SpaceField):
 
             local_value_map = self.space.local_value_map_inner(args.elt_arg, s.element_index, s.element_coords)
             dof_value = self.space.value_basis_element(value_dof, local_value_map)
-            return wp.where(taylor_dof == 0, dof_value, self.dtype(0.0))
+            return wp.where(taylor_dof == 0, dof_value, zero_element())
 
         return eval_test_inner
 
@@ -409,13 +446,14 @@ class LocalAdjointField(SpaceField):
 
         DOF_BEGIN = wp.constant(self._TAYLOR_DOF_OFFSETS[LocalAdjointField.INNER_GRAD_DOF])
         DOF_COUNT = wp.constant(self._TAYLOR_DOF_COUNTS[LocalAdjointField.INNER_GRAD_DOF])
+        zero_element = type_zero_element(self.gradient_dtype)
 
         @cache.dynamic_func(suffix=self.name)
         def eval_nabla_test_inner(args: self.ElementEvalArg, s: Sample):
             value_dof, taylor_dof = self._split_dof(self._get_dof(s), DOF_BEGIN)
 
             if taylor_dof < 0 or taylor_dof >= DOF_COUNT:
-                return self.gradient_dtype(0.0)
+                return zero_element()
 
             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)
@@ -430,13 +468,14 @@ class LocalAdjointField(SpaceField):
 
         DOF_BEGIN = wp.constant(self._TAYLOR_DOF_OFFSETS[LocalAdjointField.INNER_GRAD_DOF])
         DOF_COUNT = wp.constant(self._TAYLOR_DOF_COUNTS[LocalAdjointField.INNER_GRAD_DOF])
+        zero_element = type_zero_element(self.divergence_dtype)
 
         @cache.dynamic_func(suffix=self.name)
         def eval_div_test_inner(args: self.ElementEvalArg, s: Sample):
             value_dof, taylor_dof = self._split_dof(self._get_dof(s), DOF_BEGIN)
 
             if taylor_dof < 0 or taylor_dof >= DOF_COUNT:
-                return self.divergence_dtype(0.0)
+                return zero_element()
 
             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)
@@ -446,7 +485,11 @@ class LocalAdjointField(SpaceField):
         return eval_div_test_inner
 
     def _make_eval_outer(self):
+        if not self._is_discontinuous:
+            return self.eval_inner
+
         DOF_BEGIN = wp.constant(self._TAYLOR_DOF_OFFSETS[LocalAdjointField.OUTER_DOF])
+        zero_element = type_zero_element(self.dtype)
 
         @cache.dynamic_func(suffix=self.name)
         def eval_test_outer(args: self.ElementEvalArg, s: Sample):
@@ -454,7 +497,7 @@ class LocalAdjointField(SpaceField):
 
             local_value_map = self.space.local_value_map_outer(args.elt_arg, s.element_index, s.element_coords)
             dof_value = self.space.value_basis_element(value_dof, local_value_map)
-            return wp.where(taylor_dof == 0, dof_value, self.dtype(0.0))
+            return wp.where(taylor_dof == 0, dof_value, zero_element())
 
         return eval_test_outer
 
@@ -464,13 +507,14 @@ class LocalAdjointField(SpaceField):
 
         DOF_BEGIN = wp.constant(self._TAYLOR_DOF_OFFSETS[LocalAdjointField.OUTER_GRAD_DOF])
         DOF_COUNT = wp.constant(self._TAYLOR_DOF_COUNTS[LocalAdjointField.OUTER_GRAD_DOF])
+        zero_element = type_zero_element(self.gradient_dtype)
 
         @cache.dynamic_func(suffix=self.name)
         def eval_nabla_test_outer(args: self.ElementEvalArg, s: Sample):
             value_dof, taylor_dof = self._split_dof(self._get_dof(s), DOF_BEGIN)
 
             if taylor_dof < 0 or taylor_dof >= DOF_COUNT:
-                return self.gradient_dtype(0.0)
+                return zero_element()
 
             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)
@@ -485,13 +529,14 @@ class LocalAdjointField(SpaceField):
 
         DOF_BEGIN = wp.constant(self._TAYLOR_DOF_OFFSETS[LocalAdjointField.OUTER_GRAD_DOF])
         DOF_COUNT = wp.constant(self._TAYLOR_DOF_COUNTS[LocalAdjointField.OUTER_GRAD_DOF])
+        zero_element = type_zero_element(self.divergence_dtype)
 
         @cache.dynamic_func(suffix=self.name)
         def eval_div_test_outer(args: self.ElementEvalArg, s: Sample):
             value_dof, taylor_dof = self._split_dof(self._get_dof(s), DOF_BEGIN)
 
             if taylor_dof < 0 or taylor_dof >= DOF_COUNT:
-                return self.divergence_dtype(0.0)
+                return zero_element()
 
             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)