Lazy Replace

[JDBetteridge]

It would be very useful to have a lazy version of replace, which only gets applied when the expression tree is traversed.

Example:
Say I have an expression expr which contains a ufl.ConstantValue named c as a parameter. If I want to have the derivative of expr wrt. c I can do so by creating a ufl.Coefficient named f and doing:

ufl.replace(expr, {c: f})
ufl.derivative(expr, f)

However this traverses the expression tree twice.

(This issue would also be fixed by implementing derivative wrt. ufl.ConstantValue, but there are more examples where this is desirable). This came up as a discussion point when reviewing a recent Firedrake PR. Tagging @dham .

[wence-]

I think that you would still (in the general case) have to traverse the expression tree twice if you want your transformers to not have to handle every UFL node. In the specific case of replace, this is because differentiation doesn't commute through replacement: derivative(replace(expr, {c: f}), c) != replace(derivative(expr, c), {c: f}. In general it's because if you think as the expression transformation rules as fixed point "local" replacements with the recursive knot tied by map_expr_dag, then it is only some transformation rules that have the functorial property in the sense of map f o g == map f o map g.

I think specifically for replace you could solve it by having a Replace node, and then all transformers would be augmented to have a dispatch on Replace that traverses and replaces in the subtree before applying themselves to the result.