warp-lang 1.0.0b2__py3-none-win_amd64.whl → 1.0.0b6__py3-none-win_amd64.whl

warp/fem/integrate.py

@@ -5,13 +5,12 @@ import warp as wp
 import re
 import ast
 
-from warp.sparse import BsrMatrix, bsr_zeros, bsr_set_from_triplets, bsr_copy, bsr_diag
+from warp.sparse import BsrMatrix, bsr_zeros, bsr_set_from_triplets, bsr_copy, bsr_assign
 from warp.types import type_length
 from warp.utils import array_cast
 from warp.codegen import get_annotations
 
 from warp.fem.domain import GeometryDomain
-from warp.fem.space import SpaceRestriction
 from warp.fem.field import (
     TestField,
     TrialField,
@@ -23,7 +22,7 @@ from warp.fem.field import (
 from warp.fem.quadrature import Quadrature, RegularQuadrature
 from warp.fem.operator import Operator, Integrand
 from warp.fem import cache
-from warp.fem.types import Domain, Field, Sample, DofIndex, NULL_DOF_INDEX, OUTSIDE
+from warp.fem.types import Domain, Field, Sample, DofIndex, NULL_DOF_INDEX, OUTSIDE, make_free_sample
 
 
 def _resolve_path(func, node):
@@ -98,7 +97,7 @@ class IntegrandTransformer(ast.NodeTransformer):
             operator = arg_type.call_operator
 
             call.func = ast.Attribute(
-                value=_path_to_ast_attribute(arg_type.__qualname__),
+                value=_path_to_ast_attribute(f"{arg_type.__module__}.{arg_type.__qualname__}"),
                 attr="call_operator",
                 ctx=ast.Load(),
             )
@@ -164,7 +163,7 @@ def _translate_integrand(integrand: Integrand, field_args: Dict[str, FieldLike])
     for arg in argspec.args:
         arg_type = argspec.annotations[arg]
         if arg_type == Field:
-            annotations[arg] = field_args[arg].EvalArg
+            annotations[arg] = field_args[arg].ElementEvalArg
         elif arg_type == Domain:
             annotations[arg] = field_args[arg].ElementArg
         else:
@@ -174,11 +173,9 @@ def _translate_integrand(integrand: Integrand, field_args: Dict[str, FieldLike])
     transformer = IntegrandTransformer(integrand, field_args)
 
     def is_field_like(f):
-        # WAR for isinstance not supporting Union in Python < 3.10
-        return any(isinstance(f, field_class) for field_class in FieldLike.__args__)
+        return isinstance(f, FieldLike)
 
     suffix = "_".join([f.name for f in field_args.values() if is_field_like(f)])
-    key = integrand.name + suffix
 
     func = cache.get_integrand_function(
         integrand=integrand,
@@ -265,18 +262,14 @@ def _gen_field_struct(field_args: Dict[str, FieldLike]):
         setattr(Fields, name, arg.EvalArg())
         annotations[name] = arg.EvalArg
 
-    Fields.__qualname__ = (
-        Fields.__name__
-        + "_"
-        + "_".join([f"{name}_{arg_struct.cls.__qualname__}" for name, arg_struct in annotations.items()])
-    )
-
     try:
         Fields.__annotations__ = annotations
     except AttributeError:
         setattr(Fields.__dict__, "__annotations__", annotations)
 
-    return cache.get_struct(Fields)
+    suffix = "_".join([f"{name}_{arg_struct.cls.__qualname__}" for name, arg_struct in annotations.items()])
+
+    return cache.get_struct(Fields, suffix=suffix)
 
 
 def _gen_value_struct(value_args: Dict[str, type]):
@@ -299,25 +292,34 @@ def _gen_value_struct(value_args: Dict[str, type]):
             return arg_type_name(arg_type.cls)
         return getattr(arg_type, "__name__", str(arg_type))
 
-    Values.__qualname__ = (
-        Values.__name__
-        + "_"
-        + "_".join([f"{name}_{arg_type_name(arg_type)}" for name, arg_type in annotations.items()])
-    )
-
     try:
         Values.__annotations__ = annotations
     except AttributeError:
         setattr(Values.__dict__, "__annotations__", annotations)
 
-    return cache.get_struct(Values)
+    suffix = "_".join([f"{name}_{arg_type_name(arg_type)}" for name, arg_type in annotations.items()])
+
+    return cache.get_struct(Values, suffix=suffix)
 
 
 def _get_trial_arg():
     pass
 
+
 def _get_test_arg():
     pass
+
+
+class _FieldWrappers:
+    pass
+
+
+def _register_integrand_field_wrappers(integrand_func: wp.Function, fields: Dict[str, FieldLike]):
+    integrand_func._field_wrappers = _FieldWrappers()
+    for name, field in fields.items():
+        setattr(integrand_func._field_wrappers, name, field.ElementEvalArg)
+
+
 class PassFieldArgsToIntegrand(ast.NodeTransformer):
     def __init__(
         self,
@@ -333,6 +335,7 @@ class PassFieldArgsToIntegrand(ast.NodeTransformer):
         values_var_name: str = "values",
         domain_var_name: str = "domain_arg",
         sample_var_name: str = "sample",
+        field_wrappers_attr: str = "_field_wrappers",
     ):
         self._arg_names = arg_names
         self._field_args = field_args
@@ -346,6 +349,7 @@ class PassFieldArgsToIntegrand(ast.NodeTransformer):
         self._values_var_name = values_var_name
         self._domain_var_name = domain_var_name
         self._sample_var_name = sample_var_name
+        self._field_wrappers_attr = field_wrappers_attr
 
     def visit_Call(self, call: ast.Call):
         call = self.generic_visit(call)
@@ -366,10 +370,25 @@ class PassFieldArgsToIntegrand(ast.NodeTransformer):
                     )
                 elif arg in self._field_args:
                     call.args.append(
-                        ast.Attribute(
-                            value=ast.Name(id=self._fields_var_name, ctx=ast.Load()),
-                            attr=arg,
-                            ctx=ast.Load(),
+                        ast.Call(
+                            func=ast.Attribute(
+                                value=ast.Attribute(
+                                    value=ast.Name(id=self._func_name, ctx=ast.Load()),
+                                    attr=self._field_wrappers_attr,
+                                    ctx=ast.Load(),
+                                ),
+                                attr=arg,
+                                ctx=ast.Load(),
+                            ),
+                            args=[
+                                ast.Name(id=self._domain_var_name, ctx=ast.Load()),
+                                ast.Attribute(
+                                    value=ast.Name(id=self._fields_var_name, ctx=ast.Load()),
+                                    attr=arg,
+                                    ctx=ast.Load(),
+                                ),
+                            ],
+                            keywords=[],
                         )
                     )
                 elif arg in self._value_args:
@@ -401,36 +420,6 @@ class PassFieldArgsToIntegrand(ast.NodeTransformer):
         return call
 
 
-def get_integrate_null_kernel(
-    integrand_func: wp.Function,
-    domain: GeometryDomain,
-    quadrature: Quadrature,
-    FieldStruct: wp.codegen.Struct,
-    ValueStruct: wp.codegen.Struct,
-):
-    def integrate_kernel_fn(
-        qp_arg: quadrature.Arg,
-        domain_arg: domain.ElementArg,
-        domain_index_arg: domain.ElementIndexArg,
-        fields: FieldStruct,
-        values: ValueStruct,
-    ):
-        element_index = domain.element_index(domain_index_arg, wp.tid())
-
-        test_dof_index = NULL_DOF_INDEX
-        trial_dof_index = NULL_DOF_INDEX
-
-        qp_point_count = quadrature.point_count(qp_arg, element_index)
-        for k in range(qp_point_count):
-            qp_index = quadrature.point_index(qp_arg, element_index, k)
-            qp_coords = quadrature.point_coords(qp_arg, element_index, k)
-            qp_weight = quadrature.point_weight(qp_arg, element_index, k)
-            sample = Sample(element_index, qp_coords, qp_index, qp_weight, test_dof_index, trial_dof_index)
-            integrand_func(sample, fields, values)
-
-    return integrate_kernel_fn
-
-
 def get_integrate_constant_kernel(
     integrand_func: wp.Function,
     domain: GeometryDomain,
@@ -453,14 +442,15 @@ def get_integrate_constant_kernel(
         test_dof_index = NULL_DOF_INDEX
         trial_dof_index = NULL_DOF_INDEX
 
-        qp_point_count = quadrature.point_count(qp_arg, element_index)
+        qp_point_count = quadrature.point_count(domain_arg, qp_arg, element_index)
         for k in range(qp_point_count):
-            qp_index = quadrature.point_index(qp_arg, element_index, k)
-            coords = quadrature.point_coords(qp_arg, element_index, k)
-            qp_weight = quadrature.point_weight(qp_arg, element_index, k)
-            vol = domain.element_measure(domain_arg, element_index, coords)
+            qp_index = quadrature.point_index(domain_arg, qp_arg, element_index, k)
+            coords = quadrature.point_coords(domain_arg, qp_arg, element_index, k)
+            qp_weight = quadrature.point_weight(domain_arg, qp_arg, element_index, k)
 
             sample = Sample(element_index, coords, qp_index, qp_weight, test_dof_index, trial_dof_index)
+            vol = domain.element_measure(domain_arg, sample)
+
             val = integrand_func(sample, fields, values)
 
             elem_sum += accumulate_dtype(qp_weight * vol * val)
@@ -476,42 +466,47 @@ def get_integrate_linear_kernel(
     quadrature: Quadrature,
     FieldStruct: wp.codegen.Struct,
     ValueStruct: wp.codegen.Struct,
-    test_space: SpaceRestriction,
+    test: TestField,
+    output_dtype,
     accumulate_dtype,
 ):
     def integrate_kernel_fn(
         qp_arg: quadrature.Arg,
         domain_arg: domain.ElementArg,
         domain_index_arg: domain.ElementIndexArg,
-        test_arg: test_space.NodeArg,
+        test_arg: test.space_restriction.NodeArg,
         fields: FieldStruct,
         values: ValueStruct,
-        result: wp.array2d(dtype=accumulate_dtype),
+        result: wp.array2d(dtype=output_dtype),
     ):
-        local_node_index = wp.tid()
-        node_index = test_space.node_partition_index(test_arg, local_node_index)
-        element_count = test_space.node_element_count(test_arg, local_node_index)
+        local_node_index, test_dof = wp.tid()
+        node_index = test.space_restriction.node_partition_index(test_arg, local_node_index)
+        element_count = test.space_restriction.node_element_count(test_arg, local_node_index)
 
         trial_dof_index = NULL_DOF_INDEX
 
+        val_sum = accumulate_dtype(0.0)
+
         for n in range(element_count):
-            node_element_index = test_space.node_element_index(test_arg, local_node_index, n)
+            node_element_index = test.space_restriction.node_element_index(test_arg, local_node_index, n)
             element_index = domain.element_index(domain_index_arg, node_element_index.domain_element_index)
 
-            qp_point_count = quadrature.point_count(qp_arg, element_index)
+            test_dof_index = DofIndex(node_element_index.node_index_in_element, test_dof)
+
+            qp_point_count = quadrature.point_count(domain_arg, qp_arg, element_index)
             for k in range(qp_point_count):
-                qp_index = quadrature.point_index(qp_arg, element_index, k)
-                coords = quadrature.point_coords(qp_arg, element_index, k)
+                qp_index = quadrature.point_index(domain_arg, qp_arg, element_index, k)
+                qp_coords = quadrature.point_coords(domain_arg, qp_arg, element_index, k)
+                qp_weight = quadrature.point_weight(domain_arg, qp_arg, element_index, k)
 
-                qp_weight = quadrature.point_weight(qp_arg, element_index, k)
-                vol = domain.element_measure(domain_arg, element_index, coords)
+                vol = domain.element_measure(domain_arg, make_free_sample(element_index, qp_coords))
 
-                for i in range(test_space.space.VALUE_DOF_COUNT):
-                    test_dof_index = DofIndex(node_element_index.node_index_in_element, i)
-                    sample = Sample(element_index, coords, qp_index, qp_weight, test_dof_index, trial_dof_index)
-                    val = integrand_func(sample, fields, values)
+                sample = Sample(element_index, qp_coords, qp_index, qp_weight, test_dof_index, trial_dof_index)
+                val = integrand_func(sample, fields, values)
+
+                val_sum += accumulate_dtype(qp_weight * vol * val)
 
-                    result[node_index, i] = result[node_index, i] + accumulate_dtype(qp_weight * vol * val)
+        result[node_index, test_dof] = output_dtype(val_sum)
 
     return integrate_kernel_fn
 
@@ -522,6 +517,7 @@ def get_integrate_linear_nodal_kernel(
     FieldStruct: wp.codegen.Struct,
     ValueStruct: wp.codegen.Struct,
     test: TestField,
+    output_dtype,
     accumulate_dtype,
 ):
     def integrate_kernel_fn(
@@ -530,7 +526,7 @@ def get_integrate_linear_nodal_kernel(
         test_restriction_arg: test.space_restriction.NodeArg,
         fields: FieldStruct,
         values: ValueStruct,
-        result: wp.array2d(dtype=accumulate_dtype),
+        result: wp.array2d(dtype=output_dtype),
     ):
         local_node_index, dof = wp.tid()
 
@@ -546,6 +542,7 @@ def get_integrate_linear_nodal_kernel(
             element_index = domain.element_index(domain_index_arg, node_element_index.domain_element_index)
 
             coords = test.space.node_coords_in_element(
+                domain_arg,
                 _get_test_arg(),
                 element_index,
                 node_element_index.node_index_in_element,
@@ -553,12 +550,12 @@ def get_integrate_linear_nodal_kernel(
 
             if coords[0] != OUTSIDE:
                 node_weight = test.space.node_quadrature_weight(
+                    domain_arg,
                     _get_test_arg(),
                     element_index,
                     node_element_index.node_index_in_element,
                 )
 
-                vol = domain.element_measure(domain_arg, element_index, coords)
                 test_dof_index = DofIndex(node_element_index.node_index_in_element, dof)
 
                 sample = Sample(
@@ -569,11 +566,12 @@ def get_integrate_linear_nodal_kernel(
                     test_dof_index,
                     trial_dof_index,
                 )
+                vol = domain.element_measure(domain_arg, sample)
                 val = integrand_func(sample, fields, values)
 
                 val_sum += accumulate_dtype(node_weight * vol * val)
 
-        result[node_index, dof] = val_sum
+        result[node_index, dof] = output_dtype(val_sum)
 
     return integrate_kernel_fn
 
@@ -584,80 +582,75 @@ def get_integrate_bilinear_kernel(
     quadrature: Quadrature,
     FieldStruct: wp.codegen.Struct,
     ValueStruct: wp.codegen.Struct,
-    test_space: SpaceRestriction,
+    test: TestField,
     trial: TrialField,
+    output_dtype,
     accumulate_dtype,
 ):
-    NODES_PER_ELEMENT = trial.space.NODES_PER_ELEMENT
+    NODES_PER_ELEMENT = trial.space.topology.NODES_PER_ELEMENT
 
     def integrate_kernel_fn(
         qp_arg: quadrature.Arg,
         domain_arg: domain.ElementArg,
         domain_index_arg: domain.ElementIndexArg,
-        test_arg: test_space.NodeArg,
+        test_arg: test.space_restriction.NodeArg,
         trial_partition_arg: trial.space_partition.PartitionArg,
+        trial_topology_arg: trial.space_partition.space_topology.TopologyArg,
         fields: FieldStruct,
         values: ValueStruct,
         row_offsets: wp.array(dtype=int),
         triplet_rows: wp.array(dtype=int),
         triplet_cols: wp.array(dtype=int),
-        triplet_values: wp.array3d(dtype=accumulate_dtype),
+        triplet_values: wp.array3d(dtype=output_dtype),
     ):
-        test_local_node_index = wp.tid()
+        test_local_node_index, trial_node, test_dof, trial_dof = wp.tid()
+
+        element_count = test.space_restriction.node_element_count(test_arg, test_local_node_index)
+        test_node_index = test.space_restriction.node_partition_index(test_arg, test_local_node_index)
 
-        element_count = test_space.node_element_count(test_arg, test_local_node_index)
-        test_node_index = test_space.node_partition_index(test_arg, test_local_node_index)
+        trial_dof_index = DofIndex(trial_node, trial_dof)
 
         for element in range(element_count):
-            test_element_index = test_space.node_element_index(test_arg, test_local_node_index, element)
+            test_element_index = test.space_restriction.node_element_index(test_arg, test_local_node_index, element)
             element_index = domain.element_index(domain_index_arg, test_element_index.domain_element_index)
-            qp_point_count = quadrature.point_count(qp_arg, element_index)
+            qp_point_count = quadrature.point_count(domain_arg, qp_arg, element_index)
 
-            start_offset = (row_offsets[test_node_index] + element) * NODES_PER_ELEMENT
+            test_dof_index = DofIndex(
+                test_element_index.node_index_in_element,
+                test_dof,
+            )
+
+            val_sum = accumulate_dtype(0.0)
 
             for k in range(qp_point_count):
-                qp_index = quadrature.point_index(qp_arg, element_index, k)
-                coords = quadrature.point_coords(qp_arg, element_index, k)
-
-                qp_weight = quadrature.point_weight(qp_arg, element_index, k)
-                vol = domain.element_measure(domain_arg, element_index, coords)
-
-                offset_cur = start_offset
-
-                for trial_n in range(NODES_PER_ELEMENT):
-                    for i in range(test_space.space.VALUE_DOF_COUNT):
-                        for j in range(trial.space.VALUE_DOF_COUNT):
-                            test_dof_index = DofIndex(
-                                test_element_index.node_index_in_element,
-                                i,
-                            )
-                            trial_dof_index = DofIndex(trial_n, j)
-                            sample = Sample(
-                                element_index,
-                                coords,
-                                qp_index,
-                                qp_weight,
-                                test_dof_index,
-                                trial_dof_index,
-                            )
-                            val = integrand_func(sample, fields, values)
-                            triplet_values[offset_cur, i, j] = triplet_values[offset_cur, i, j] + accumulate_dtype(
-                                qp_weight * vol * val
-                            )
-
-                    offset_cur += 1
-
-            # Set column indices
-            offset_cur = start_offset
-            for trial_n in range(NODES_PER_ELEMENT):
+                qp_index = quadrature.point_index(domain_arg, qp_arg, element_index, k)
+                coords = quadrature.point_coords(domain_arg, qp_arg, element_index, k)
+
+                qp_weight = quadrature.point_weight(domain_arg, qp_arg, element_index, k)
+                vol = domain.element_measure(domain_arg, make_free_sample(element_index, coords))
+
+                sample = Sample(
+                    element_index,
+                    coords,
+                    qp_index,
+                    qp_weight,
+                    test_dof_index,
+                    trial_dof_index,
+                )
+                val = integrand_func(sample, fields, values)
+                val_sum += accumulate_dtype(qp_weight * vol * val)
+
+            block_offset = (row_offsets[test_node_index] + element) * NODES_PER_ELEMENT + trial_node
+            triplet_values[block_offset, test_dof, trial_dof] = output_dtype(val_sum)
+
+            # Set row and column indices
+            if test_dof == 0 and trial_dof == 0:
                 trial_node_index = trial.space_partition.partition_node_index(
                     trial_partition_arg,
-                    trial.space.element_node_index(_get_trial_arg(), element_index, trial_n),
+                    trial.space.topology.element_node_index(domain_arg, trial_topology_arg, element_index, trial_node),
                 )
-
-                triplet_rows[offset_cur] = test_node_index
-                triplet_cols[offset_cur] = trial_node_index
-                offset_cur += 1
+                triplet_rows[block_offset] = test_node_index
+                triplet_cols[block_offset] = trial_node_index
 
     return integrate_kernel_fn
 
@@ -668,6 +661,7 @@ def get_integrate_bilinear_nodal_kernel(
     FieldStruct: wp.codegen.Struct,
     ValueStruct: wp.codegen.Struct,
     test: TestField,
+    output_dtype,
     accumulate_dtype,
 ):
     def integrate_kernel_fn(
@@ -678,7 +672,7 @@ def get_integrate_bilinear_nodal_kernel(
         values: ValueStruct,
         triplet_rows: wp.array(dtype=int),
         triplet_cols: wp.array(dtype=int),
-        triplet_values: wp.array3d(dtype=accumulate_dtype),
+        triplet_values: wp.array3d(dtype=output_dtype),
     ):
         local_node_index, test_dof, trial_dof = wp.tid()
 
@@ -692,6 +686,7 @@ def get_integrate_bilinear_nodal_kernel(
             element_index = domain.element_index(domain_index_arg, node_element_index.domain_element_index)
 
             coords = test.space.node_coords_in_element(
+                domain_arg,
                 _get_test_arg(),
                 element_index,
                 node_element_index.node_index_in_element,
@@ -699,13 +694,12 @@ def get_integrate_bilinear_nodal_kernel(
 
             if coords[0] != OUTSIDE:
                 node_weight = test.space.node_quadrature_weight(
+                    domain_arg,
                     _get_test_arg(),
                     element_index,
                     node_element_index.node_index_in_element,
                 )
 
-                vol = domain.element_measure(domain_arg, element_index, coords)
-
                 test_dof_index = DofIndex(node_element_index.node_index_in_element, test_dof)
                 trial_dof_index = DofIndex(node_element_index.node_index_in_element, trial_dof)
 
@@ -717,11 +711,12 @@ def get_integrate_bilinear_nodal_kernel(
                     test_dof_index,
                     trial_dof_index,
                 )
+                vol = domain.element_measure(domain_arg, sample)
                 val = integrand_func(sample, fields, values)
 
                 val_sum += accumulate_dtype(node_weight * vol * val)
 
-        triplet_values[local_node_index, test_dof, trial_dof] = val_sum
+        triplet_values[local_node_index, test_dof, trial_dof] = output_dtype(val_sum)
         triplet_rows[local_node_index] = node_index
         triplet_cols[local_node_index] = node_index
 
@@ -738,8 +733,12 @@ def _generate_integrate_kernel(
     trial: Optional[TrialField],
     trial_name: str,
     fields: Dict[str, FieldLike],
+    output_dtype: type,
     accumulate_dtype: type,
+    kernel_options: Dict[str, Any] = {},
 ) -> wp.Kernel:
+    output_dtype = wp.types.type_scalar_type(output_dtype)
+
     # Extract field arguments from integrand
     field_args, value_args, domain_name, sample_name = _get_integrand_field_arguments(
         integrand, fields=fields, domain=domain
@@ -749,7 +748,7 @@ def _generate_integrate_kernel(
     ValueStruct = _gen_value_struct(value_args)
 
     # Check if kernel exist in cache
-    kernel_suffix = f"_itg_{domain.name}_{FieldStruct.key}"
+    kernel_suffix = f"_itg_{wp.types.type_typestr(output_dtype)}{wp.types.type_typestr(accumulate_dtype)}_{domain.name}_{FieldStruct.key}"
     if nodal:
         kernel_suffix += "_nodal"
     else:
@@ -774,6 +773,8 @@ def _generate_integrate_kernel(
         field_args,
     )
 
+    _register_integrand_field_wrappers(integrand_func, fields)
+
     if test is None and trial is None:
         integrate_kernel_fn = get_integrate_constant_kernel(
             integrand_func,
@@ -791,6 +792,7 @@ def _generate_integrate_kernel(
                 FieldStruct,
                 ValueStruct,
                 test=test,
+                output_dtype=output_dtype,
                 accumulate_dtype=accumulate_dtype,
             )
         else:
@@ -800,7 +802,8 @@ def _generate_integrate_kernel(
                 quadrature,
                 FieldStruct,
                 ValueStruct,
-                test_space=test.space_restriction,
+                test=test,
+                output_dtype=output_dtype,
                 accumulate_dtype=accumulate_dtype,
             )
     else:
@@ -811,6 +814,7 @@ def _generate_integrate_kernel(
                 FieldStruct,
                 ValueStruct,
                 test=test,
+                output_dtype=output_dtype,
                 accumulate_dtype=accumulate_dtype,
             )
         else:
@@ -820,8 +824,9 @@ def _generate_integrate_kernel(
                 quadrature,
                 FieldStruct,
                 ValueStruct,
-                test_space=test.space_restriction,
+                test=test,
                 trial=trial,
+                output_dtype=output_dtype,
                 accumulate_dtype=accumulate_dtype,
             )
 
@@ -829,6 +834,7 @@ def _generate_integrate_kernel(
         integrand=integrand,
         kernel_fn=integrate_kernel_fn,
         suffix=kernel_suffix,
+        kernel_options=kernel_options,
         code_transformers=[
             PassFieldArgsToIntegrand(
                 arg_names=integrand.argspec.args,
@@ -837,7 +843,7 @@ def _generate_integrate_kernel(
                 sample_name=sample_name,
                 domain_name=domain_name,
                 test_name=test_name,
-                trial_name=trial_name
+                trial_name=trial_name,
             )
         ],
     )
@@ -846,7 +852,7 @@ def _generate_integrate_kernel(
 
 
 def _launch_integrate_kernel(
-    kernel: wp.kernel,
+    kernel: wp.Kernel,
     FieldStruct: wp.codegen.Struct,
     ValueStruct: wp.codegen.Struct,
     domain: GeometryDomain,
@@ -857,16 +863,11 @@ def _launch_integrate_kernel(
     fields: Dict[str, FieldLike],
     values: Dict[str, Any],
     accumulate_dtype: type,
+    temporary_store: Optional[cache.TemporaryStore],
     output_dtype: type,
     output: Optional[Union[wp.array, BsrMatrix]],
     device,
-) -> wp.Kernel:
-    if output_dtype is None:
-        if output is not None:
-            output_dtype = output.dtype
-        else:
-            output_dtype = accumulate_dtype
-
+):
     # Set-up launch arguments
     domain_elt_arg = domain.element_arg_value(device=device)
     domain_elt_index_arg = domain.element_index_arg_value(device=device)
@@ -882,14 +883,23 @@ def _launch_integrate_kernel(
     for k, v in values.items():
         setattr(value_struct_values, k, v)
 
-    # Constant
+    # Constant form
     if test is None and trial is None:
-        if output is None or output.dtype != accumulate_dtype:
-            result = wp.zeros(shape=(1), device=device, dtype=output_dtype)
+        if output is not None and output.dtype == accumulate_dtype:
+            if output.size < 1:
+                raise RuntimeError("Output array must be of size at least 1")
+            accumulate_array = output
         else:
-            result = output
-            result.zero_()
+            accumulate_temporary = cache.borrow_temporary(
+                shape=(1),
+                device=device,
+                dtype=accumulate_dtype,
+                temporary_store=temporary_store,
+                requires_grad=output is not None and output.requires_grad,
+            )
+            accumulate_array = accumulate_temporary.array
 
+        accumulate_array.zero_()
         wp.launch(
             kernel=kernel,
             dim=domain.element_count(),
@@ -899,43 +909,77 @@ def _launch_integrate_kernel(
                 domain_elt_index_arg,
                 field_arg_values,
                 value_struct_values,
-                result,
+                accumulate_array,
             ],
             device=device,
         )
 
-        if output is None:
-            return output_dtype(result.numpy()[0])
+        if output == accumulate_array:
+            return output
+        elif output is None:
+            return accumulate_array.numpy()[0]
         else:
-            if output != result:
-                array_cast(in_array=result, out_array=output)
+            array_cast(in_array=accumulate_array, out_array=output)
             return output
 
     test_arg = test.space_restriction.node_arg(device=device)
 
     # Linear form
     if trial is None:
-        if test.space.VALUE_DOF_COUNT == 1:
-            result_dtype = accumulate_dtype
+        # If an output array is provided with the correct type, accumulate directly into it
+        # Otherwise, grab a temporary array
+        if output is None:
+            if type_length(output_dtype) == test.space.VALUE_DOF_COUNT:
+                output_shape = (test.space_partition.node_count(),)
+            elif type_length(output_dtype) == 1:
+                output_shape = (test.space_partition.node_count(), test.space.VALUE_DOF_COUNT)
+            else:
+                raise RuntimeError(
+                    f"Incompatible output type {wp.types.type_repr(output_dtype)}, must be scalar or vector of length {test.space.VALUE_DOF_COUNT}"
+                )
+
+            output_temporary = cache.borrow_temporary(
+                temporary_store=temporary_store,
+                shape=output_shape,
+                dtype=output_dtype,
+                device=device,
+            )
+
+            output = output_temporary.array
+
         else:
-            result_dtype = wp.vec(length=test.space.VALUE_DOF_COUNT, dtype=accumulate_dtype)
+            output_temporary = None
 
-        result_array = wp.zeros(
-            shape=test.space_partition.node_count(),
-            dtype=result_dtype,
-            device=device,
-        )
+            if output.shape[0] < test.space_partition.node_count():
+                raise RuntimeError(f"Output array must have at least {test.space_partition.node_count()} rows")
+
+            output_dtype = output.dtype
+            if type_length(output_dtype) != test.space.VALUE_DOF_COUNT:
+                if type_length(output_dtype) != 1:
+                    raise RuntimeError(
+                        f"Incompatible output type {wp.types.type_repr(output_dtype)}, must be scalar or vector of length {test.space.VALUE_DOF_COUNT}"
+                    )
+                if output.ndim != 2 and output.shape[1] != test.space.VALUE_DOF_COUNT:
+                    raise RuntimeError(
+                        f"Incompatible output array shape, last dimension must be of size {test.space.VALUE_DOF_COUNT}"
+                    )
 
         # Launch the integration on the kernel on a 2d scalar view of the actual array
-        result_2d_view = wp.array(
-            data=None,
-            ptr=result_array.ptr,
-            capacity=result_array.capacity,
-            owner=False,
-            device=result_array.device,
-            shape=(test.space_partition.node_count(), test.space.VALUE_DOF_COUNT),
-            dtype=accumulate_dtype,
-        )
+        output.zero_()
+
+        def as_2d_array(array):
+            return wp.array(
+                data=None,
+                ptr=array.ptr,
+                capacity=array.capacity,
+                owner=False,
+                device=array.device,
+                shape=(test.space_partition.node_count(), test.space.VALUE_DOF_COUNT),
+                dtype=wp.types.type_scalar_type(output_dtype),
+                grad=None if array.grad is None else as_2d_array(array.grad),
+            )
+
+        output_view = output if output.ndim == 2 else as_2d_array(output)
 
         if nodal:
             wp.launch(
@@ -947,14 +991,14 @@ def _launch_integrate_kernel(
                     test_arg,
                     field_arg_values,
                     value_struct_values,
-                    result_2d_view,
+                    output_view,
                 ],
                 device=device,
             )
         else:
             wp.launch(
                 kernel=kernel,
-                dim=test.space_restriction.node_count(),
+                dim=(test.space_restriction.node_count(), test.space.VALUE_DOF_COUNT),
                 inputs=[
                     qp_arg,
                     domain_elt_arg,
@@ -962,55 +1006,47 @@ def _launch_integrate_kernel(
                     test_arg,
                     field_arg_values,
                     value_struct_values,
-                    result_2d_view,
+                    output_view,
                 ],
                 device=device,
             )
 
-        if output_dtype == result_array.dtype:
-            return result_array
-
-        output_type_length = type_length(output_dtype)
-        if output_type_length == test.space.VALUE_DOF_COUNT:
-            cast_result = wp.empty(dtype=output_dtype, shape=result_array.shape)
-        else:
-            cast_result = wp.empty(dtype=output_dtype, shape=result_2d_view.shape)
+        if output_temporary is not None:
+            return output_temporary.detach()
 
-        array_cast(in_array=result_array, out_array=cast_result)
-        return cast_result
+        return output
 
     # Bilinear form
 
     if test.space.VALUE_DOF_COUNT == 1 and trial.space.VALUE_DOF_COUNT == 1:
-        block_type = accumulate_dtype
+        block_type = output_dtype
     else:
-        block_type = wp.types.matrix(
-            shape=(test.space.VALUE_DOF_COUNT, trial.space.VALUE_DOF_COUNT), dtype=accumulate_dtype
+        block_type = cache.cached_mat_type(
+            shape=(test.space.VALUE_DOF_COUNT, trial.space.VALUE_DOF_COUNT), dtype=output_dtype
         )
 
-    bsr_matrix = bsr_zeros(
-        rows_of_blocks=test.space_partition.node_count(),
-        cols_of_blocks=trial.space_partition.node_count(),
-        block_type=block_type,
-        device=device,
-    )
-
     if nodal:
         nnz = test.space_restriction.node_count()
     else:
-        nnz = test.space_restriction.total_node_element_count() * trial.space.NODES_PER_ELEMENT
+        nnz = test.space_restriction.total_node_element_count() * trial.space.topology.NODES_PER_ELEMENT
 
-    triplet_rows = wp.empty(n=nnz, dtype=int, device=device)
-    triplet_cols = wp.empty(n=nnz, dtype=int, device=device)
-    triplet_values = wp.zeros(
+    triplet_rows_temp = cache.borrow_temporary(temporary_store, shape=(nnz,), dtype=int, device=device)
+    triplet_cols_temp = cache.borrow_temporary(temporary_store, shape=(nnz,), dtype=int, device=device)
+    triplet_values_temp = cache.borrow_temporary(
+        temporary_store,
         shape=(
             nnz,
             test.space.VALUE_DOF_COUNT,
             trial.space.VALUE_DOF_COUNT,
         ),
-        dtype=accumulate_dtype,
+        dtype=output_dtype,
         device=device,
     )
+    triplet_cols = triplet_cols_temp.array
+    triplet_rows = triplet_rows_temp.array
+    triplet_values = triplet_values_temp.array
+
+    triplet_values.zero_()
 
     if nodal:
         wp.launch(
@@ -1033,15 +1069,22 @@ def _launch_integrate_kernel(
         offsets = test.space_restriction.partition_element_offsets()
 
         trial_partition_arg = trial.space_partition.partition_arg_value(device)
+        trial_topology_arg = trial.space_partition.space_topology.topo_arg_value(device)
         wp.launch(
             kernel=kernel,
-            dim=test.space_restriction.node_count(),
+            dim=(
+                test.space_restriction.node_count(),
+                trial.space.topology.NODES_PER_ELEMENT,
+                test.space.VALUE_DOF_COUNT,
+                trial.space.VALUE_DOF_COUNT,
+            ),
             inputs=[
                 qp_arg,
                 domain_elt_arg,
                 domain_elt_index_arg,
                 test_arg,
                 trial_partition_arg,
+                trial_topology_arg,
                 field_arg_values,
                 value_struct_values,
                 offsets,
@@ -1052,38 +1095,63 @@ def _launch_integrate_kernel(
             device=device,
         )
 
-    bsr_set_from_triplets(bsr_matrix, triplet_rows, triplet_cols, triplet_values)
-    return bsr_matrix if output_dtype == accumulate_dtype else bsr_copy(bsr_matrix, scalar_type=output_dtype)
+    if output is not None:
+        if output.nrow != test.space_partition.node_count() or output.ncol != trial.space_partition.node_count():
+            raise RuntimeError(
+                f"Output matrix must have {test.space_partition.node_count()} rows and {trial.space_partition.node_count()} columns of blocks"
+            )
+
+    else:
+        output = bsr_zeros(
+            rows_of_blocks=test.space_partition.node_count(),
+            cols_of_blocks=trial.space_partition.node_count(),
+            block_type=block_type,
+            device=device,
+        )
+
+    bsr_set_from_triplets(output, triplet_rows, triplet_cols, triplet_values)
+
+    # Do not wait for garbage collection
+    triplet_values_temp.release()
+    triplet_rows_temp.release()
+    triplet_cols_temp.release()
+
+    return output
 
 
 def integrate(
     integrand: Integrand,
-    domain: GeometryDomain = None,
-    quadrature: Quadrature = None,
+    domain: Optional[GeometryDomain] = None,
+    quadrature: Optional[Quadrature] = None,
     nodal: bool = False,
-    fields={},
-    values={},
+    fields: Dict[str, FieldLike] = {},
+    values: Dict[str, Any] = {},
+    accumulate_dtype: type = wp.float64,
+    output_dtype: Optional[type] = None,
+    output: Optional[Union[BsrMatrix, wp.array]] = None,
     device=None,
-    accumulate_dtype=wp.float64,
-    output_dtype=None,
-    output=None,
+    temporary_store: Optional[cache.TemporaryStore] = None,
+    kernel_options: Dict[str, Any] = {},
 ):
     """
     Integrates a constant, linear or bilinear form, and returns a scalar, array, or sparse matrix, respectively.
 
     Args:
-        integrand: Form to be integrated, must have `wp.integrand` decorator
+        integrand: Form to be integrated, must have :func:`integrand` decorator
         domain: Integration domain. If None, deduced from fields
         quadrature: Quadrature formula. If None, deduced from domain and fields degree.
         nodal: For linear or bilinear form only, use the test function nodes as the quadrature points. Assumes Lagrange interpolation functions are used, and no differential or DG operator is evaluated on the test or trial functions.
         fields: Discrete, test, and trial fields to be passed to the integrand. Keys in the dictionary must match integrand parameter names.
-        values: Additional variable values to be passed to the integrand, can by of any type accepted by warp kernel launchs. Keys in the dictionary must match integrand parameter names.
-        device: Device on which to perform the integration
+        values: Additional variable values to be passed to the integrand, can be of any type accepted by warp kernel launchs. Keys in the dictionary must match integrand parameter names.
+        temporary_store: shared pool from which to allocate temporary arrays
         accumulate_dtype: Scalar type to be used for accumulating integration samples
-        output_dtype: Scalar type for returned results. If None, defaults to accumulate_dtype
+        output: Sparse matrix or warp array into which to store the result of the integration
+        output_dtype: Scalar type for returned results in `output` is notr provided. If None, defaults to `accumulate_dtype`
+        device: Device on which to perform the integration
+        kernel_options: Overloaded options to be passed to the kernel builder (e.g, ``{"enable_backward": True}``)
     """
     if not isinstance(integrand, Integrand):
-        raise ValueError("integrand must be tagged with @integrand decorator")
+        raise ValueError("integrand must be tagged with @warp.fem.integrand decorator")
 
     test, test_name, trial, trial_name = _get_test_and_trial_fields(fields)
 
@@ -1111,15 +1179,23 @@ def integrate(
             )
     else:
         if quadrature is None:
-            order = 0
-            if test is not None:
-                order += test.space.degree
-            if trial is not None:
-                order += trial.space.degree
+            order = sum(field.degree for field in fields.values())
             quadrature = RegularQuadrature(domain=domain, order=order)
         elif domain != quadrature.domain:
             raise ValueError("Incompatible integration and quadrature domain")
 
+    # Canonicalize types
+    accumulate_dtype = wp.types.type_to_warp(accumulate_dtype)
+    if output is not None:
+        if isinstance(output, BsrMatrix):
+            output_dtype = output.scalar_type
+        else:
+            output_dtype = output.dtype
+    elif output_dtype is None:
+        output_dtype = accumulate_dtype
+    else:
+        output_dtype = wp.types.type_to_warp(output_dtype)
+
     kernel, FieldStruct, ValueStruct = _generate_integrate_kernel(
         integrand=integrand,
         domain=domain,
@@ -1131,6 +1207,8 @@ def integrate(
         trial_name=trial_name,
         fields=fields,
         accumulate_dtype=accumulate_dtype,
+        output_dtype=output_dtype,
+        kernel_options=kernel_options,
     )
 
     return _launch_integrate_kernel(
@@ -1145,13 +1223,14 @@ def integrate(
         fields=fields,
         values=values,
         accumulate_dtype=accumulate_dtype,
+        temporary_store=temporary_store,
         output_dtype=output_dtype,
         output=output,
         device=device,
     )
 
 
-def get_interpolate_kernel(
+def get_interpolate_to_field_function(
     integrand_func: wp.Function,
     domain: GeometryDomain,
     FieldStruct: wp.codegen.Struct,
@@ -1160,7 +1239,8 @@ def get_interpolate_kernel(
 ):
     value_type = dest.space.dtype
 
-    def interpolate_kernel_fn(
+    def interpolate_to_field_fn(
+        local_node_index: int,
         domain_arg: domain.ElementArg,
         domain_index_arg: domain.ElementIndexArg,
         dest_node_arg: dest.space_restriction.NodeArg,
@@ -1168,19 +1248,15 @@ def get_interpolate_kernel(
         fields: FieldStruct,
         values: ValueStruct,
     ):
-        local_node_index = wp.tid()
         node_index = dest.space_restriction.node_partition_index(dest_node_arg, local_node_index)
-
         element_count = dest.space_restriction.node_element_count(dest_node_arg, local_node_index)
-        if element_count == 0:
-            return
 
         test_dof_index = NULL_DOF_INDEX
         trial_dof_index = NULL_DOF_INDEX
         node_weight = 1.0
 
-        # Volume-weighted average accross elements
-        # Superfluous if the function is continuous, but we might as well
+        # Volume-weighted average across elements
+        # Superfluous if the interpolated function is continuous, but helpful for visualizing discontinuous spaces
 
         val_sum = value_type(0.0)
         vol_sum = float(0.0)
@@ -1190,14 +1266,13 @@ def get_interpolate_kernel(
             element_index = domain.element_index(domain_index_arg, node_element_index.domain_element_index)
 
             coords = dest.space.node_coords_in_element(
+                domain_arg,
                 dest_eval_arg.space_arg,
                 element_index,
                 node_element_index.node_index_in_element,
             )
 
             if coords[0] != OUTSIDE:
-                vol = domain.element_measure(domain_arg, element_index, coords)
-
                 sample = Sample(
                     element_index,
                     coords,
@@ -1206,20 +1281,118 @@ def get_interpolate_kernel(
                     test_dof_index,
                     trial_dof_index,
                 )
+                vol = domain.element_measure(domain_arg, sample)
                 val = integrand_func(sample, fields, values)
 
                 vol_sum += vol
                 val_sum += vol * val
 
+        return val_sum, vol_sum
+
+    return interpolate_to_field_fn
+
+
+def get_interpolate_to_field_kernel(
+    interpolate_to_field_fn: wp.Function,
+    domain: GeometryDomain,
+    FieldStruct: wp.codegen.Struct,
+    ValueStruct: wp.codegen.Struct,
+    dest: FieldRestriction,
+):
+    def interpolate_to_field_kernel_fn(
+        domain_arg: domain.ElementArg,
+        domain_index_arg: domain.ElementIndexArg,
+        dest_node_arg: dest.space_restriction.NodeArg,
+        dest_eval_arg: dest.field.EvalArg,
+        fields: FieldStruct,
+        values: ValueStruct,
+    ):
+        local_node_index = wp.tid()
+
+        val_sum, vol_sum = interpolate_to_field_fn(
+            local_node_index, domain_arg, domain_index_arg, dest_node_arg, dest_eval_arg, fields, values
+        )
+
         if vol_sum > 0.0:
+            node_index = dest.space_restriction.node_partition_index(dest_node_arg, local_node_index)
             dest.field.set_node_value(dest_eval_arg, node_index, val_sum / vol_sum)
 
-    return interpolate_kernel_fn
+    return interpolate_to_field_kernel_fn
+
+
+def get_interpolate_to_array_kernel(
+    integrand_func: wp.Function,
+    domain: GeometryDomain,
+    quadrature: Quadrature,
+    FieldStruct: wp.codegen.Struct,
+    ValueStruct: wp.codegen.Struct,
+    value_type: type,
+):
+    def interpolate_to_array_kernel_fn(
+        qp_arg: quadrature.Arg,
+        domain_arg: quadrature.domain.ElementArg,
+        domain_index_arg: quadrature.domain.ElementIndexArg,
+        fields: FieldStruct,
+        values: ValueStruct,
+        result: wp.array(dtype=value_type),
+    ):
+        element_index = domain.element_index(domain_index_arg, wp.tid())
+
+        test_dof_index = NULL_DOF_INDEX
+        trial_dof_index = NULL_DOF_INDEX
+
+        qp_point_count = quadrature.point_count(domain_arg, qp_arg, element_index)
+        for k in range(qp_point_count):
+            qp_index = quadrature.point_index(domain_arg, qp_arg, element_index, k)
+            coords = quadrature.point_coords(domain_arg, qp_arg, element_index, k)
+            qp_weight = quadrature.point_weight(domain_arg, qp_arg, element_index, k)
+
+            sample = Sample(element_index, coords, qp_index, qp_weight, test_dof_index, trial_dof_index)
+
+            result[qp_index] = integrand_func(sample, fields, values)
+
+    return interpolate_to_array_kernel_fn
+
+
+def get_interpolate_nonvalued_kernel(
+    integrand_func: wp.Function,
+    domain: GeometryDomain,
+    quadrature: Quadrature,
+    FieldStruct: wp.codegen.Struct,
+    ValueStruct: wp.codegen.Struct,
+):
+    def interpolate_nonvalued_kernel_fn(
+        qp_arg: quadrature.Arg,
+        domain_arg: quadrature.domain.ElementArg,
+        domain_index_arg: quadrature.domain.ElementIndexArg,
+        fields: FieldStruct,
+        values: ValueStruct,
+    ):
+        element_index = domain.element_index(domain_index_arg, wp.tid())
+
+        test_dof_index = NULL_DOF_INDEX
+        trial_dof_index = NULL_DOF_INDEX
+
+        qp_point_count = quadrature.point_count(domain_arg, qp_arg, element_index)
+        for k in range(qp_point_count):
+            qp_index = quadrature.point_index(domain_arg, qp_arg, element_index, k)
+            coords = quadrature.point_coords(domain_arg, qp_arg, element_index, k)
+            qp_weight = quadrature.point_weight(domain_arg, qp_arg, element_index, k)
+
+            sample = Sample(element_index, coords, qp_index, qp_weight, test_dof_index, trial_dof_index)
+            integrand_func(sample, fields, values)
 
+    return interpolate_nonvalued_kernel_fn
 
-def _generate_interpolate_kernel(integrand: Integrand, dest: FieldLike, fields: Dict[str, FieldLike]) -> wp.Kernel:
-    domain = dest.domain
 
+def _generate_interpolate_kernel(
+    integrand: Integrand,
+    domain: GeometryDomain,
+    dest: Optional[Union[FieldLike, wp.array]],
+    quadrature: Optional[Quadrature],
+    fields: Dict[str, FieldLike],
+    kernel_options: Dict[str, Any] = {},
+) -> wp.Kernel:
     # Extract field arguments from integrand
     field_args, value_args, domain_name, sample_name = _get_integrand_field_arguments(
         integrand, fields=fields, domain=domain
@@ -1231,11 +1404,20 @@ def _generate_interpolate_kernel(integrand: Integrand, dest: FieldLike, fields:
         field_args,
     )
 
+    _register_integrand_field_wrappers(integrand_func, fields)
+
     FieldStruct = _gen_field_struct(field_args)
     ValueStruct = _gen_value_struct(value_args)
 
     # Check if kernel exist in cache
-    kernel_suffix = f"_itp_{FieldStruct.key}_{dest.domain.name}_{dest.space_restriction.space_partition.name}_{dest.space.name}"
+    if isinstance(dest, FieldRestriction):
+        kernel_suffix = (
+            f"_itp_{FieldStruct.key}_{dest.domain.name}_{dest.space_restriction.space_partition.name}_{dest.space.name}"
+        )
+    elif wp.types.is_array(dest):
+        kernel_suffix = f"_itp_{FieldStruct.key}_{quadrature.name}_{wp.types.type_repr(dest.dtype)}"
+    else:
+        kernel_suffix = f"_itp_{FieldStruct.key}_{quadrature.name}"
 
     kernel = cache.get_integrand_kernel(
         integrand=integrand,
@@ -1245,18 +1427,61 @@ def _generate_interpolate_kernel(integrand: Integrand, dest: FieldLike, fields:
         return kernel, FieldStruct, ValueStruct
 
     # Generate interpolation kernel
-    interpolate_kernel_fn = get_interpolate_kernel(
-        integrand_func,
-        domain,
-        dest=dest,
-        FieldStruct=FieldStruct,
-        ValueStruct=ValueStruct,
-    )
+    if isinstance(dest, FieldRestriction):
+        # need to split into kernel + function for diffferentiability
+        interpolate_fn = get_interpolate_to_field_function(
+            integrand_func,
+            domain,
+            dest=dest,
+            FieldStruct=FieldStruct,
+            ValueStruct=ValueStruct,
+        )
+
+        interpolate_fn = cache.get_integrand_function(
+            integrand=integrand,
+            func=interpolate_fn,
+            suffix=kernel_suffix,
+            code_transformers=[
+                PassFieldArgsToIntegrand(
+                    arg_names=integrand.argspec.args,
+                    field_args=field_args.keys(),
+                    value_args=value_args.keys(),
+                    sample_name=sample_name,
+                    domain_name=domain_name,
+                )
+            ],
+        )
+
+        interpolate_kernel_fn = get_interpolate_to_field_kernel(
+            interpolate_fn,
+            domain,
+            dest=dest,
+            FieldStruct=FieldStruct,
+            ValueStruct=ValueStruct,
+        )
+    elif wp.types.is_array(dest):
+        interpolate_kernel_fn = get_interpolate_to_array_kernel(
+            integrand_func,
+            domain=domain,
+            quadrature=quadrature,
+            value_type=dest.dtype,
+            FieldStruct=FieldStruct,
+            ValueStruct=ValueStruct,
+        )
+    else:
+        interpolate_kernel_fn = get_interpolate_nonvalued_kernel(
+            integrand_func,
+            domain=domain,
+            quadrature=quadrature,
+            FieldStruct=FieldStruct,
+            ValueStruct=ValueStruct,
+        )
 
     kernel = cache.get_integrand_kernel(
         integrand=integrand,
         kernel_fn=interpolate_kernel_fn,
         suffix=kernel_suffix,
+        kernel_options=kernel_options,
         code_transformers=[
             PassFieldArgsToIntegrand(
                 arg_names=integrand.argspec.args,
@@ -1275,16 +1500,16 @@ def _launch_interpolate_kernel(
     kernel: wp.kernel,
     FieldStruct: wp.codegen.Struct,
     ValueStruct: wp.codegen.Struct,
-    dest: FieldLike,
+    domain: GeometryDomain,
+    dest: Optional[Union[FieldRestriction, wp.array]],
+    quadrature: Optional[Quadrature],
     fields: Dict[str, FieldLike],
     values: Dict[str, Any],
     device,
 ) -> wp.Kernel:
     # Set-up launch arguments
-    elt_arg = dest.domain.element_arg_value(device=device)
-    elt_index_arg = dest.domain.element_index_arg_value(device=device)
-    dest_node_arg = dest.space_restriction.node_arg(device=device)
-    dest_eval_arg = dest.field.eval_arg_value(device=device)
+    elt_arg = domain.element_arg_value(device=device)
+    elt_index_arg = domain.element_index_arg_value(device=device)
 
     field_arg_values = FieldStruct()
     for k, v in fields.items():
@@ -1294,37 +1519,65 @@ def _launch_interpolate_kernel(
     for k, v in values.items():
         setattr(value_struct_values, k, v)
 
-    wp.launch(
-        kernel=kernel,
-        dim=dest.space_restriction.node_count(),
-        inputs=[
-            elt_arg,
-            elt_index_arg,
-            dest_node_arg,
-            dest_eval_arg,
-            field_arg_values,
-            value_struct_values,
-        ],
-        device=device,
-    )
+    if isinstance(dest, FieldRestriction):
+        dest_node_arg = dest.space_restriction.node_arg(device=device)
+        dest_eval_arg = dest.field.eval_arg_value(device=device)
+
+        wp.launch(
+            kernel=kernel,
+            dim=dest.space_restriction.node_count(),
+            inputs=[
+                elt_arg,
+                elt_index_arg,
+                dest_node_arg,
+                dest_eval_arg,
+                field_arg_values,
+                value_struct_values,
+            ],
+            device=device,
+        )
+    elif wp.types.is_array(dest):
+        qp_arg = quadrature.arg_value(device)
+        wp.launch(
+            kernel=kernel,
+            dim=domain.element_count(),
+            inputs=[qp_arg, elt_arg, elt_index_arg, field_arg_values, value_struct_values, dest],
+            device=device,
+        )
+    else:
+        qp_arg = quadrature.arg_value(device)
+        wp.launch(
+            kernel=kernel,
+            dim=domain.element_count(),
+            inputs=[qp_arg, elt_arg, elt_index_arg, field_arg_values, value_struct_values],
+            device=device,
+        )
 
 
 def interpolate(
     integrand: Integrand,
-    dest: Union[DiscreteField, FieldRestriction],
-    fields={},
-    values={},
+    dest: Optional[Union[DiscreteField, FieldRestriction, wp.array]] = None,
+    quadrature: Optional[Quadrature] = None,
+    fields: Dict[str, FieldLike] = {},
+    values: Dict[str, Any] = {},
     device=None,
+    kernel_options: Dict[str, Any] = {},
 ):
     """
-    Interpolates a function and assigns the result to a discrete field.
+    Interpolates a function at a finite set of sample points and optionally assigns the result to a discrete field or a raw warp array.
 
     Args:
-        integrand: Function to be interpolated, must have `wp.integrand` decorator
-        dest: Discrete field, or restriction of a discrete field to a domain, to which the interpolation result will be assigned
+        integrand: Function to be interpolated, must have :func:`integrand` decorator
+        dest: Where to store the interpolation result. Can be either
+
+         - a :class:`DiscreteField`, or restriction of a discrete field to a domain (from :func:`make_restriction`). In this case, interpolation will be performed at each node.
+         - a normal warp array. In this case, the `quadrature` argument defining the interpolation locations must be provided and the result of the `integrand` at each quadrature point will be assigned to the array.
+         - ``None``. In this case, the `quadrature` argument must also be provided and the `integrand` function is reponsible for dealing with the interpolation result.
+        quadrature: Quadrature formula defining the interpolation samples if `dest` is not a discrete field or field restriction.
         fields: Discrete fields to be passed to the integrand. Keys in the dictionary must match integrand parameters names.
-        values: Additional variable values to be passed to the integrand, can by of any type accepted by warp kernel launchs. Keys in the dictionary must match integrand parameter names.
+        values: Additional variable values to be passed to the integrand, can be of any type accepted by warp kernel launchs. Keys in the dictionary must match integrand parameter names.
         device: Device on which to perform the interpolation
+        kernel_options: Overloaded options to be passed to the kernel builder (e.g, ``{"enable_backward": True}``)
     """
     if not isinstance(integrand, Integrand):
         raise ValueError("integrand must be tagged with @integrand decorator")
@@ -1333,20 +1586,33 @@ def interpolate(
     if test is not None or trial is not None:
         raise ValueError("Test or Trial fields should not be used for interpolation")
 
-    if not isinstance(dest, FieldRestriction):
+    if isinstance(dest, DiscreteField):
         dest = make_restriction(dest)
 
+    if isinstance(dest, FieldRestriction):
+        domain = dest.domain
+    else:
+        if quadrature is None:
+            raise ValueError("When not interpolating to a field, a quadrature formula must be provided")
+
+        domain = quadrature.domain
+
     kernel, FieldStruct, ValueStruct = _generate_interpolate_kernel(
         integrand=integrand,
+        domain=domain,
         dest=dest,
+        quadrature=quadrature,
         fields=fields,
+        kernel_options=kernel_options,
     )
 
     return _launch_interpolate_kernel(
         kernel=kernel,
         FieldStruct=FieldStruct,
         ValueStruct=ValueStruct,
+        domain=domain,
         dest=dest,
+        quadrature=quadrature,
         fields=fields,
         values=values,
         device=device,