[BUG] fem.integrate failing with double precision polynomial space vectors

[cmflannery]

Bug Description

fem.integrate is failing with double precision vectors. Updating the fem.make_polynomial_space dtype from wp.vec2 to wp.vec2d causes fem.integrate to fail with the following error:

Module warp.fem.space.restriction.dyn.fill_element_node_indices_Grid2D_BoundarySides_GridBipolynomialSpaceTopology_9_Trace_Whole 25805a8 load on device 'cuda:0' took 0.39 ms  (cached)
Module warp.fem.utils 8b39cbb load on device 'cuda:0' took 0.28 ms  (cached)
Module warp.fem.space.restriction 4263cdf load on device 'cuda:0' took 0.29 ms  (cached)
Module __main__.mass_form e618e55 load on device 'cuda:0' took 30.84 ms
Traceback (most recent call last):
  File "/.../lib/python3.11/site-packages/warp/codegen.py", line 1006, in build
    adj.eval(adj.tree.body[0])
  File "/.../lib/python3.11/site-packages/warp/codegen.py", line 2700, in eval
    return emit_node(adj, node)
           ^^^^^^^^^^^^^^^^^^^^
  File "/.../lib/python3.11/site-packages/warp/codegen.py", line 1611, in emit_FunctionDef
    adj.eval(f)
  File "/.../lib/python3.11/site-packages/warp/codegen.py", line 2700, in eval
    return emit_node(adj, node)
           ^^^^^^^^^^^^^^^^^^^^
  File "/.../lib/python3.11/site-packages/warp/codegen.py", line 2558, in emit_Return
    var = (adj.eval(node.value),)
           ^^^^^^^^^^^^^^^^^^^^
  File "/.../lib/python3.11/site-packages/warp/codegen.py", line 2700, in eval
    return emit_node(adj, node)
           ^^^^^^^^^^^^^^^^^^^^
  File "/.../lib/python3.11/site-packages/warp/codegen.py", line 2238, in emit_Call
    args = tuple(adj.resolve_arg(x) for x in node.args)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/.../lib/python3.11/site-packages/warp/codegen.py", line 2238, in <genexpr>
    args = tuple(adj.resolve_arg(x) for x in node.args)
                 ^^^^^^^^^^^^^^^^^^
  File "/.../lib/python3.11/site-packages/warp/codegen.py", line 2141, in resolve_arg
    var = adj.eval(arg)
          ^^^^^^^^^^^^^
  File "/.../lib/python3.11/site-packages/warp/codegen.py", line 2700, in eval
    return emit_node(adj, node)
           ^^^^^^^^^^^^^^^^^^^^
  File "/.../lib/python3.11/site-packages/warp/codegen.py", line 2242, in emit_Call
    out = adj.add_call(func, args, kwargs, type_args, min_outputs=min_outputs)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/.../lib/python3.11/site-packages/warp/codegen.py", line 1324, in add_call
    return_type = func.value_func(
                  ^^^^^^^^^^^^^^^^
  File "/.../lib/python3.11/site-packages/warp/builtins.py", line 831, in vector_value_func
    raise RuntimeError()

.....
          raise RuntimeError(
RuntimeError: Error while parsing function "integrate_kernel_fn" at /.../lib/python3.11/site-packages/warp/fem/integrate.py:926:
            val = integrand_func(sample, fields, values)
;Error while parsing function "mass_form" at /.../scripts/warp_min.py:13:
    return wp.dot(u(s), v(s))
;Error while parsing function "eval_test_inner" at /.../lib/python3.11/site-packages/warp/fem/field/virtual.py:60:
    dof_value = self.space.node_basis_element(get_node_coord(dof))
;Error while parsing function "node_basis_element" at /.../lib/python3.11/site-packages/warp/fem/space/basis_function_space.py:90:
    return basis_element(self.dof_dtype(0.0), dof_coord)

Minimal Example

import warp as wp
import warp.fem as fem


@fem.integrand
def mass_form(
    s: fem.Sample,
    u: fem.Field,
    v: fem.Field,
) -> wp.float64:
    """Mass form for L2 projection of vector fields."""
    # For vector fields, we want the dot product of the vectors
    return wp.dot(u(s), v(s))


def main():
    geo = fem.Grid2D(res=wp.vec2i(3))
    boundary = fem.BoundarySides(geo)

    # Fails with dtype=wp.vec2d
    u_space = fem.make_polynomial_space(geo, degree=2, dtype=wp.vec2d)
    
    # # Passes with dtype=wp.vec2
    # u_space = fem.make_polynomial_space(geo, degree=2, dtype=wp.vec2)

    u_test = fem.make_test(space=u_space, domain=boundary)
    u_trial = fem.make_trial(space=u_space, domain=boundary)
    u_projector = fem.integrate(
        mass_form, fields={"u": u_trial, "v": u_test}, nodal=True
    )
    print("Successfully created projector")


if __name__ == "__main__":
    wp.init()
    main()

System Information

Python

• Python 3.11.10
• Poetry package manager
• warp 1.5.0

Hardware and CUDA Version

CUDA Toolkit 12.6, Driver 12.4
Devices:
"cpu" : "x86_64"
"cuda:0" : "NVIDIA L40S" (44 GiB, sm_89, mempool enabled)

[gdaviet]

Hi @cmflannery, indeed double-precision-valued spaces are not currently supported by fem.integrate.
Currently, the rationale has been to stay in single-precision for quantities evaluated at any single sample point, then switch to double precision to perform accumulation over the different samples.
While this strategy usually gives a good performance vs numerical stability compromise, I understand that there are cases where having the full evaluation done in full precision would be necessary, so support for double-precision fields and geometries is definitely on the roadmap (as well as complex-valued, eventually).
In any case, the error message should be made more indicative of the current limitation

[cmflannery]

Gotcha, thanks for the clarification. A more indicative error message would definitely be appreciated!