Allow specifying leaf class in recursive map and map-reduce

[majosm]

I think this could help simplify the logic in a couple of places in illinois-ceesd/mirgecom#566 (example).

cc @thomasgibson

[alexfikl]

One little comment, but code-wise this looks good to me!

What's the use case? Looking at illinois-ceesd/mirgecom#566, it seems like it would only be used with DOFArrays, but there's already a function for that in meshmode.dof_array. Would that do the trick?

[majosm]

What's the use case? Looking at illinois-ceesd/mirgecom#566, it seems like it would only be used with DOFArrays, but there's already a function for that in meshmode.dof_array. Would that do the trick?

TIL 🙂 Yes, I think it might. On the other hand, this change could remove the existing code duplication (i.e., rec_map_dof_array_container could be reimplemented in terms of rec_map_array_container). Either solution would seem to work for me.

[alexfikl]

On the other hand, this change could remove the existing code duplication (i.e., rec_map_dof_array_container could be reimplemented in terms of rec_map_array_container). Either solution would seem to work for me.

It actually was initially! Well, in terms of _map_array_container_impl. I re-implemented it directly in meshmode at some point in the hopes that it would be a bit faster (and not too much of a maintainence burden, since the code is straightforward).

I don't feel particularly strongly about it either, so we can let @inducer break the tie 😁

[inducer]

I'm gravitating towards having this implemented here and making the meshmode.dof_array use this via the leaf_class argument. I don't see too much benefit in keeping both virtually-identical implementations around:

• https://github.com/inducer/meshmode/blob/a7a8a060bb3103abcf1c4e509f72d6b324384d71/meshmode/dof_array.py#L359-L378

• arraycontext/arraycontext/container/traversal.py

  Lines 90 to 101 in 957be2f

  	def rec(_ary: ArrayOrContainerT) -> ArrayOrContainerT:
  	if type(_ary) is leaf_cls: # type(ary) is never None
  	return f(_ary)
  	try:
  	iterable = serialize_container(_ary)
  	except NotAnArrayContainerError:
  	return f(_ary)
  	else:
  	return deserialize_container(_ary, [
  	(key, frec(subary)) for key, subary in iterable
  	])

I also feel like the speed difference is probably non-existent. (pending benchmarks proving me wrong 😁 )

[alexfikl]

I also feel like the speed difference is probably non-existent. (pending benchmarks proving me wrong)

Made a simple benchmark using this patch and the version from meshmode is a bit slower!

INFO:__main__:[mm] min 6.767774e-03 mean 6.942166e-03 ± 1.912915e-04
INFO:__main__:[ac] min 5.201670e-03 mean 5.289515e-03 ± 7.113123e-05

so mostly just confused now 😕. Either way, it's definitely negligible!

[majosm]

Ready for another look @inducer.

[inducer]

LGTM, thanks! If we can ever get the (currently super flaky) CI to pass, this'll land.