warp-lang 1.7.2rc1__py3-none-manylinux_2_34_aarch64.whl → 1.8.1__py3-none-manylinux_2_34_aarch64.whl

warp/fem/field/field.py

@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Dict, Optional, Set
+from typing import Any, ClassVar, Dict, Optional, Set
 
 import warp as wp
 from warp.fem import cache
@@ -22,6 +22,7 @@ from warp.fem.geometry import DeformedGeometry, Geometry
 from warp.fem.operator import Operator, integrand
 from warp.fem.space import FunctionSpace, SpacePartition
 from warp.fem.types import NULL_ELEMENT_INDEX, ElementKind, Sample
+from warp.fem.utils import type_zero_element
 
 
 class FieldLike:
@@ -37,6 +38,10 @@ class FieldLike:
         """Value of the field-level arguments to be passed to device functions"""
         raise NotImplementedError
 
+    def fill_eval_arg(self, arg: "FieldLike.EvalArg", device):
+        """Fill the field-level arguments to be passed to device functions"""
+        raise NotImplementedError
+
     @property
     def degree(self) -> int:
         """Polynomial degree of the field, used to estimate necessary quadrature order"""
@@ -50,10 +55,6 @@ class FieldLike:
     def __str__(self) -> str:
         return self.name
 
-    def eval_arg_value(self, device):
-        """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
@@ -139,6 +140,51 @@ class GeometryField(FieldLike):
         """
         return DeformedGeometry(self, relative=relative)
 
+    @property
+    def gradient_dtype(self):
+        """Return type of the (world space) gradient operator. Assumes self.gradient_valid()"""
+        if wp.types.type_is_matrix(self.dtype):
+            return None
+
+        if wp.types.type_is_vector(self.dtype):
+            return cache.cached_mat_type(
+                shape=(wp.types.type_size(self.dtype), self.geometry.dimension),
+                dtype=wp.types.type_scalar_type(self.dtype),
+            )
+        if wp.types.type_is_quaternion(self.dtype):
+            return cache.cached_mat_type(
+                shape=(4, 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))
+
+    @property
+    def reference_gradient_dtype(self):
+        """Return type of the reference space gradient operator. Assumes self.gradient_valid()"""
+        if wp.types.type_is_matrix(self.dtype):
+            return None
+
+        if wp.types.type_is_vector(self.dtype):
+            return cache.cached_mat_type(
+                shape=(wp.types.type_size(self.dtype), self.geometry.cell_dimension),
+                dtype=wp.types.type_scalar_type(self.dtype),
+            )
+        if wp.types.type_is_quaternion(self.dtype):
+            return cache.cached_mat_type(
+                shape=(4, self.geometry.cell_dimension),
+                dtype=wp.types.type_scalar_type(self.dtype),
+            )
+        return cache.cached_vec_type(length=self.geometry.cell_dimension, dtype=wp.types.type_scalar_type(self.dtype))
+
+    @property
+    def divergence_dtype(self):
+        """Return type of the divergence operator. Assumes self.divergence_valid()"""
+        if wp.types.type_is_vector(self.dtype):
+            return wp.types.type_scalar_type(self.dtype)
+        if wp.types.type_is_matrix(self.dtype):
+            return cache.cached_vec_type(length=self.dtype._shape_[1], dtype=wp.types.type_scalar_type(self.dtype))
+        return None
+
 
 class SpaceField(GeometryField):
     """Base class for fields defined over a function space"""
@@ -178,33 +224,6 @@ class SpaceField(GeometryField):
     def dof_dtype(self) -> type:
         return self.space.dof_dtype
 
-    @property
-    def gradient_dtype(self):
-        """Return type of the (world space) 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))
-
-    @property
-    def reference_gradient_dtype(self):
-        """Return type of the reference space 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.cell_dimension),
-                dtype=wp.types.type_scalar_type(self.dtype),
-            )
-        return cache.cached_vec_type(length=self.geometry.cell_dimension, dtype=wp.types.type_scalar_type(self.dtype))
-
-    @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_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
 
@@ -251,6 +270,19 @@ class ImplicitField(GeometryField):
         degree: Optional hint for automatic determination of quadrature orders when integrating this field
     """
 
+    _dynamic_attribute_constructors: ClassVar = {
+        "ElementEvalArg": lambda obj: obj._make_element_eval_arg(),
+        "eval_degree": lambda obj: obj._make_eval_degree(),
+        "eval_inner": lambda obj: obj._make_eval_func(obj._func),
+        "eval_grad_inner": lambda obj: obj._make_eval_func(obj._grad_func),
+        "eval_div_inner": lambda obj: obj._make_eval_func(obj._div_func),
+        "eval_reference_grad_inner": lambda obj: obj._make_eval_reference_grad(),
+        "eval_outer": lambda obj: obj.eval_inner,
+        "eval_grad_outer": lambda obj: obj.eval_grad_inner,
+        "eval_div_outer": lambda obj: obj.eval_div_inner,
+        "eval_reference_grad_outer": lambda obj: obj.eval_reference_grad_inner,
+    }
+
     def __init__(
         self,
         domain: GeometryDomain,
@@ -284,18 +316,7 @@ class ImplicitField(GeometryField):
         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
+        cache.setup_dynamic_attributes(self)
 
     @property
     def values(self):
@@ -303,6 +324,7 @@ class ImplicitField(GeometryField):
 
     @values.setter
     def values(self, v):
+        self._values = v
         self._func_arg = cache.populate_argument_struct(self.EvalArg, v, self._func.func.__name__)
 
     @property
@@ -316,6 +338,9 @@ class ImplicitField(GeometryField):
     def eval_arg_value(self, device):
         return self._func_arg
 
+    def fill_eval_arg(self, arg, device):
+        cache.populate_argument_struct(self.EvalArg, self._values, self._func.func.__name__, arg)
+
     @property
     def degree(self) -> int:
         return self._degree
@@ -324,6 +349,12 @@ class ImplicitField(GeometryField):
     def name(self) -> str:
         return f"Implicit_{self.domain.name}_{self.degree}_{self.EvalArg.key}"
 
+    def gradient_valid(self) -> bool:
+        return self._grad_func is not None
+
+    def divergence_valid(self) -> bool:
+        return self._div_func is not None
+
     def _make_eval_func(self, func):
         if func is None:
             return None
@@ -387,6 +418,20 @@ class UniformField(GeometryField):
         value: Uniform value over the domain
     """
 
+    _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_zero(obj.gradient_dtype),
+        "eval_div_inner": lambda obj: obj._make_eval_zero(obj.divergence_dtype),
+        "eval_reference_grad_inner": lambda obj: obj._make_eval_zero(obj.reference_gradient_dtype),
+        "eval_outer": lambda obj: obj.eval_inner,
+        "eval_grad_outer": lambda obj: obj.eval_grad_inner,
+        "eval_div_outer": lambda obj: obj.eval_div_inner,
+        "eval_reference_grad_outer": lambda obj: obj.eval_reference_grad_inner,
+    }
+
     def __init__(self, domain: GeometryDomain, value: Any):
         self.domain = domain
 
@@ -396,30 +441,7 @@ class UniformField(GeometryField):
         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
+        cache.setup_dynamic_attributes(self)
 
     @property
     def value(self):
@@ -444,10 +466,19 @@ class UniformField(GeometryField):
         arg.value = self.value
         return arg
 
+    def fill_eval_arg(self, arg, device):
+        arg.value = self.value
+
     @property
     def degree(self) -> int:
         return 0
 
+    def gradient_valid(self) -> bool:
+        return self.gradient_dtype is not None
+
+    def divergence_valid(self) -> bool:
+        return self.divergence_dtype is not None
+
     @property
     def name(self) -> str:
         return f"Uniform{self.domain.name}_{wp.types.get_type_code(self.dtype)}"
@@ -463,11 +494,11 @@ class UniformField(GeometryField):
         if dtype is None:
             return None
 
-        scalar_type = wp.types.type_scalar_type(dtype)
+        zero_element = type_zero_element(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 zero_element()
 
         return eval_zero
 
@@ -511,6 +542,20 @@ class NonconformingField(GeometryField):
 
     _LOOKUP_EPS = wp.constant(1.0e-6)
 
+    _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_nonconforming_eval("eval_inner"),
+        "eval_grad_inner": lambda obj: obj._make_nonconforming_eval("eval_grad_inner"),
+        "eval_div_inner": lambda obj: obj._make_nonconforming_eval("eval_div_inner"),
+        "eval_reference_grad_inner": lambda obj: obj._make_eval_reference_grad(),
+        "eval_outer": lambda obj: obj.eval_inner,
+        "eval_grad_outer": lambda obj: obj.eval_grad_inner,
+        "eval_div_outer": lambda obj: obj.eval_div_inner,
+        "eval_reference_grad_outer": lambda obj: obj.eval_reference_grad_inner,
+    }
+
     def __init__(self, domain: GeometryDomain, field: DiscreteField, background: Any = 0.0):
         self.domain = domain
 
@@ -523,20 +568,7 @@ class NonconformingField(GeometryField):
             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
+        cache.setup_dynamic_attributes(self)
 
     @property
     def geometry(self) -> Geometry:
@@ -546,19 +578,26 @@ class NonconformingField(GeometryField):
     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)
+        self.fill_eval_arg(arg, device)
         return arg
 
+    def fill_eval_arg(self, arg, device):
+        self.field.fill_eval_arg(arg.field_cell_eval_arg.eval_arg, device)
+        self.field.geometry.fill_cell_arg(arg.field_cell_eval_arg.elt_arg, device)
+        self.background.fill_eval_arg(arg.background_arg, device)
+
     @property
     def degree(self) -> int:
         return self.field.degree
 
+    def gradient_valid(self) -> bool:
+        return self.field.gradient_valid() and self.background.gradient_valid()
+
+    def divergence_valid(self) -> bool:
+        return self.field.divergence_valid() and self.background.divergence_valid()
+
     @property
     def name(self) -> str:
         return f"{self.domain.name}_{self.field.name}_{self.background.name}"
@@ -570,20 +609,23 @@ class NonconformingField(GeometryField):
         if field_eval is None or bg_eval is None:
             return None
 
+        cell_lookup = self.field.geometry.cell_lookup
+
         @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
-            )
+            nonconforming_s = cell_lookup(cell_arg, pos, NonconformingField._LOOKUP_EPS)
+            if nonconforming_s.element_index != NULL_ELEMENT_INDEX:
+                if (
+                    wp.length_sq(pos - self.field.geometry.cell_position(cell_arg, nonconforming_s))
+                    <= NonconformingField._LOOKUP_EPS
+                ):
+                    return field_eval(
+                        self.field.ElementEvalArg(cell_arg, args.eval_arg.field_cell_eval_arg.eval_arg), nonconforming_s
+                    )
+
+            return bg_eval(self.background.ElementEvalArg(args.elt_arg, args.eval_arg.background_arg), s)
 
         return eval_nc
 

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